# ============================================================================ # TinyFlux-Deep Training Cell - Combined Dataset # ============================================================================ # Datasets: # - FFHQ portraits (40k × 3 prompts = ~120k) # - flux-schnell-teacher-latents (train_simple_512 + train_512 + train_2_512 = ~40k) # - SportFashion_512x512 (54.6k) - with background mask # - SynthMoCap_smpl_512 (106k) - with SMPL body mask # Total: ~320k samples # # All encoded with flan-t5-base (768 dim) # Masked loss for foreground-focused training on product/body datasets # # USAGE: Run model.py cell first, then this cell # This converts tiny-flux-deep into tiny-flux-deep-v2, which is a different variant. # WARNING: It will impact performance and weights, so be aware. # ============================================================================ import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset from datasets import load_dataset, concatenate_datasets from transformers import T5EncoderModel, T5Tokenizer, CLIPTextModel, CLIPTokenizer from huggingface_hub import HfApi, hf_hub_download from safetensors.torch import save_file, load_file from torch.utils.tensorboard import SummaryWriter from tqdm.auto import tqdm import numpy as np import math import json import random from typing import Tuple, Optional, Dict, List import os from datetime import datetime from PIL import Image # ============================================================================ # CUDA OPTIMIZATIONS # ============================================================================ torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.backends.cudnn.benchmark = True torch.set_float32_matmul_precision('high') import warnings warnings.filterwarnings('ignore', message='.*TF32.*') # ============================================================================ # CONFIG # ============================================================================ BATCH_SIZE = 16 GRAD_ACCUM = 2 LR = 3e-4 EPOCHS = 20 MAX_SEQ = 128 SHIFT = 3.0 DEVICE = "cuda" DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 ALLOW_WEIGHT_UPGRADE = True # Set to False to require exact weight match # HuggingFace Hub HF_REPO = "AbstractPhil/tiny-flux-deep" SAVE_EVERY = 625 UPLOAD_EVERY = 625 SAMPLE_EVERY = 312 LOG_EVERY = 10 LOG_UPLOAD_EVERY = 625 # Checkpoint loading LOAD_TARGET = "latest" # "hub", "latest", "best", "none" RESUME_STEP = None # ============================================================================ # DATASET CONFIG - Enable/disable datasets for this run # ============================================================================ ENABLE_PORTRAIT = False ENABLE_SCHNELL = True ENABLE_SPORTFASHION = False # Disabled for disk space ENABLE_SYNTHMOCAP = False # Disabled for disk space # Dataset repos PORTRAIT_REPO = "AbstractPhil/ffhq_flux_latents_repaired" PORTRAIT_NUM_SHARDS = 11 SCHNELL_REPO = "AbstractPhil/flux-schnell-teacher-latents" SCHNELL_CONFIGS = ["train_simple_512"] # Add "train_512", "train_2_512" as disk allows SPORTFASHION_REPO = "Pianokill/SportFashion_512x512" SYNTHMOCAP_REPO = "toyxyz/SynthMoCap_smpl_512" # Masked loss config # Weight foreground higher than background FG_LOSS_WEIGHT = 2.0 # Foreground multiplier BG_LOSS_WEIGHT = 0.5 # Background multiplier USE_MASKED_LOSS = False # Min-SNR weighting for flow matching MIN_SNR_GAMMA = 5.0 # Paths CHECKPOINT_DIR = "./tiny_flux_deep_checkpoints" LOG_DIR = "./tiny_flux_deep_logs" SAMPLE_DIR = "./tiny_flux_deep_samples" ENCODING_CACHE_DIR = "./encoding_cache" LATENT_CACHE_DIR = "./latent_cache" os.makedirs(CHECKPOINT_DIR, exist_ok=True) os.makedirs(LOG_DIR, exist_ok=True) os.makedirs(SAMPLE_DIR, exist_ok=True) os.makedirs(ENCODING_CACHE_DIR, exist_ok=True) os.makedirs(LATENT_CACHE_DIR, exist_ok=True) # ============================================================================ # REGULARIZATION CONFIG # ============================================================================ TEXT_DROPOUT = 0.1 GUIDANCE_DROPOUT = 0.1 EMA_DECAY = 0.9999 # ============================================================================ # EMA # ============================================================================ class EMA: def __init__(self, model, decay=0.9999): self.decay = decay self.shadow = {} self._backup = {} if hasattr(model, '_orig_mod'): state = model._orig_mod.state_dict() else: state = model.state_dict() for k, v in state.items(): self.shadow[k] = v.clone().detach() @torch.no_grad() def update(self, model): if hasattr(model, '_orig_mod'): state = model._orig_mod.state_dict() else: state = model.state_dict() for k, v in state.items(): if k in self.shadow: self.shadow[k].lerp_(v.to(self.shadow[k].dtype), 1 - self.decay) def apply_shadow_for_eval(self, model): if hasattr(model, '_orig_mod'): self._backup = {k: v.clone() for k, v in model._orig_mod.state_dict().items()} model._orig_mod.load_state_dict(self.shadow) else: self._backup = {k: v.clone() for k, v in model.state_dict().items()} model.load_state_dict(self.shadow) def restore(self, model): if hasattr(model, '_orig_mod'): model._orig_mod.load_state_dict(self._backup) else: model.load_state_dict(self._backup) self._backup = {} def state_dict(self): return {'shadow': self.shadow, 'decay': self.decay} def load_state_dict(self, state): self.shadow = {k: v.clone() for k, v in state['shadow'].items()} self.decay = state.get('decay', self.decay) # ============================================================================ # REGULARIZATION # ============================================================================ def apply_text_dropout(t5_embeds, clip_pooled, dropout_prob=0.1): B = t5_embeds.shape[0] mask = torch.rand(B, device=t5_embeds.device) < dropout_prob t5_embeds = t5_embeds.clone() clip_pooled = clip_pooled.clone() t5_embeds[mask] = 0 clip_pooled[mask] = 0 return t5_embeds, clip_pooled, mask # ============================================================================ # MASKING UTILITIES # ============================================================================ def detect_background_color(image: Image.Image, sample_size: int = 100) -> Tuple[int, int, int]: """Detect dominant background color by sampling corners.""" img = np.array(image) if len(img.shape) == 2: img = np.stack([img] * 3, axis=-1) h, w = img.shape[:2] corners = [ img[:sample_size, :sample_size], # Top-left img[:sample_size, -sample_size:], # Top-right img[-sample_size:, :sample_size], # Bottom-left img[-sample_size:, -sample_size:], # Bottom-right ] # Compute median color across corners corner_pixels = np.concatenate([c.reshape(-1, 3) for c in corners], axis=0) bg_color = np.median(corner_pixels, axis=0).astype(np.uint8) return tuple(bg_color) def create_product_mask(image: Image.Image, threshold: int = 30) -> np.ndarray: """Create foreground mask for product images (non-background pixels).""" img = np.array(image).astype(np.float32) if len(img.shape) == 2: img = np.stack([img] * 3, axis=-1) bg_color = detect_background_color(image) bg_color = np.array(bg_color, dtype=np.float32) # Distance from background color diff = np.sqrt(np.sum((img - bg_color) ** 2, axis=-1)) mask = (diff > threshold).astype(np.float32) return mask def create_smpl_mask(conditioning_image: Image.Image, threshold: int = 20) -> np.ndarray: """Create body mask from SMPL conditioning render. SynthMoCap uses green/teal background. Body is rendered as mesh. Non-green pixels = body. """ img = np.array(conditioning_image).astype(np.float32) if len(img.shape) == 2: return (img > threshold).astype(np.float32) # Green background detection (high G, low R and B relative to G) r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2] # Background is typically green/teal # Body pixels have different color distribution is_background = (g > r + 20) & (g > b + 20) mask = (~is_background).astype(np.float32) return mask def downsample_mask_to_latent(mask: np.ndarray, latent_h: int = 64, latent_w: int = 64) -> torch.Tensor: """Downsample pixel mask to latent space dimensions.""" # Use area averaging for downsampling mask_pil = Image.fromarray((mask * 255).astype(np.uint8)) mask_pil = mask_pil.resize((latent_w, latent_h), Image.Resampling.BILINEAR) mask_latent = np.array(mask_pil).astype(np.float32) / 255.0 return torch.from_numpy(mask_latent) # ============================================================================ # HF HUB SETUP # ============================================================================ print("Setting up HuggingFace Hub...") api = HfApi() # ============================================================================ # FLOW MATCHING HELPERS # ============================================================================ def flux_shift(t, s=SHIFT): return s * t / (1 + (s - 1) * t) def min_snr_weight(t, gamma=MIN_SNR_GAMMA): """Min-SNR weighting for flow matching to balance loss across timesteps.""" snr = (t / (1 - t).clamp(min=1e-5)).pow(2) return torch.clamp(snr, max=gamma) / snr.clamp(min=1e-5) # ============================================================================ # LOAD TEXT ENCODERS # ============================================================================ print("Loading text encoders...") t5_tok = T5Tokenizer.from_pretrained("google/flan-t5-base") t5_enc = T5EncoderModel.from_pretrained("google/flan-t5-base", torch_dtype=DTYPE).to(DEVICE).eval() for p in t5_enc.parameters(): p.requires_grad = False clip_tok = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") clip_enc = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=DTYPE).to(DEVICE).eval() for p in clip_enc.parameters(): p.requires_grad = False print("✓ Text encoders loaded") # ============================================================================ # LOAD VAE # ============================================================================ print("Loading VAE...") from diffusers import AutoencoderKL vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=DTYPE).to(DEVICE).eval() for p in vae.parameters(): p.requires_grad = False VAE_SCALE = vae.config.scaling_factor print(f"✓ VAE loaded (scale={VAE_SCALE})") # ============================================================================ # ENCODING FUNCTIONS # ============================================================================ @torch.no_grad() def encode_prompt(prompt: str) -> Tuple[torch.Tensor, torch.Tensor]: t5_inputs = t5_tok(prompt, return_tensors="pt", padding="max_length", max_length=MAX_SEQ, truncation=True).to(DEVICE) t5_out = t5_enc(**t5_inputs).last_hidden_state clip_inputs = clip_tok(prompt, return_tensors="pt", padding="max_length", max_length=77, truncation=True).to(DEVICE) clip_out = clip_enc(**clip_inputs).pooler_output return t5_out.squeeze(0), clip_out.squeeze(0) @torch.no_grad() def encode_prompts_batched(prompts: List[str], batch_size: int = 64) -> Tuple[torch.Tensor, torch.Tensor]: all_t5 = [] all_clip = [] for i in tqdm(range(0, len(prompts), batch_size), desc="Encoding", leave=False): batch = prompts[i:i+batch_size] t5_inputs = t5_tok(batch, return_tensors="pt", padding="max_length", max_length=MAX_SEQ, truncation=True).to(DEVICE) t5_out = t5_enc(**t5_inputs).last_hidden_state all_t5.append(t5_out.cpu()) clip_inputs = clip_tok(batch, return_tensors="pt", padding="max_length", max_length=77, truncation=True).to(DEVICE) clip_out = clip_enc(**clip_inputs).pooler_output all_clip.append(clip_out.cpu()) return torch.cat(all_t5, dim=0), torch.cat(all_clip, dim=0) @torch.no_grad() def encode_image_to_latent(image: Image.Image) -> torch.Tensor: """Encode PIL image to VAE latent.""" if image.mode != "RGB": image = image.convert("RGB") # Resize to 512x512 if needed if image.size != (512, 512): image = image.resize((512, 512), Image.Resampling.LANCZOS) # To tensor and normalize img_tensor = torch.from_numpy(np.array(image)).float() / 255.0 img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W] img_tensor = (img_tensor * 2.0 - 1.0).to(DEVICE, dtype=DTYPE) # Encode latent = vae.encode(img_tensor).latent_dist.sample() latent = latent * VAE_SCALE return latent.squeeze(0).cpu() # [16, 64, 64] # ============================================================================ # LOAD DATASETS # ============================================================================ # --- 1. Portrait Dataset (FFHQ) --- portrait_ds = None portrait_indices = [] portrait_prompts = [] if ENABLE_PORTRAIT: print(f"\n[1/4] Loading portrait dataset from {PORTRAIT_REPO}...") portrait_shards = [] for i in range(PORTRAIT_NUM_SHARDS): split_name = f"train_{i:02d}" print(f" Loading {split_name}...") shard = load_dataset(PORTRAIT_REPO, split=split_name) portrait_shards.append(shard) portrait_ds = concatenate_datasets(portrait_shards) print(f"✓ Portrait: {len(portrait_ds)} base samples") # Extract triplicated prompts - batch read columns then iterate print(" Extracting prompts (columnar)...") # Batch read all three columns at once (fast Arrow read) florence_list = list(portrait_ds["text_florence"]) llava_list = list(portrait_ds["text_llava"]) blip_list = list(portrait_ds["text_blip"]) # Build from Python lists (instant) for i, (f, l, b) in enumerate(zip(florence_list, llava_list, blip_list)): if f and f.strip(): portrait_indices.append(i) portrait_prompts.append(f) if l and l.strip(): portrait_indices.append(i) portrait_prompts.append(l) if b and b.strip(): portrait_indices.append(i) portrait_prompts.append(b) print(f" Expanded: {len(portrait_prompts)} samples (3 prompts/image)") else: print("\n[1/4] Portrait dataset DISABLED") # --- 2. Schnell Teacher Dataset --- schnell_ds = None schnell_prompts = [] if ENABLE_SCHNELL: print(f"\n[2/4] Loading schnell teacher dataset from {SCHNELL_REPO}...") schnell_datasets = [] for config in SCHNELL_CONFIGS: print(f" Loading {config}...") ds = load_dataset(SCHNELL_REPO, config, split="train") schnell_datasets.append(ds) print(f" {len(ds)} samples") schnell_ds = concatenate_datasets(schnell_datasets) schnell_prompts = list(schnell_ds["prompt"]) print(f"✓ Schnell: {len(schnell_ds)} samples") else: print("\n[2/4] Schnell dataset DISABLED") # --- 3. SportFashion Dataset --- sportfashion_ds = None sportfashion_prompts = [] if ENABLE_SPORTFASHION: print(f"\n[3/4] Loading SportFashion dataset from {SPORTFASHION_REPO}...") sportfashion_ds = load_dataset(SPORTFASHION_REPO, split="train") sportfashion_prompts = list(sportfashion_ds["text"]) print(f"✓ SportFashion: {len(sportfashion_ds)} samples") else: print("\n[3/4] SportFashion dataset DISABLED") # --- 4. SynthMoCap Dataset --- synthmocap_ds = None synthmocap_prompts = [] if ENABLE_SYNTHMOCAP: print(f"\n[4/4] Loading SynthMoCap dataset from {SYNTHMOCAP_REPO}...") synthmocap_ds = load_dataset(SYNTHMOCAP_REPO, split="train") synthmocap_prompts = list(synthmocap_ds["text"]) print(f"✓ SynthMoCap: {len(synthmocap_ds)} samples") else: print("\n[4/4] SynthMoCap dataset DISABLED") # ============================================================================ # ENCODE ALL PROMPTS # ============================================================================ total_samples = len(portrait_prompts) + len(schnell_prompts) + len(sportfashion_prompts) + len(synthmocap_prompts) print(f"\nTotal combined samples: {total_samples}") def load_or_encode(cache_path, prompts, name): if not prompts: return None, None if os.path.exists(cache_path): print(f"Loading cached {name} encodings...") cached = torch.load(cache_path) return cached["t5_embeds"], cached["clip_pooled"] else: print(f"Encoding {len(prompts)} {name} prompts...") t5, clip = encode_prompts_batched(prompts, batch_size=64) torch.save({"t5_embeds": t5, "clip_pooled": clip}, cache_path) print(f"✓ Cached to {cache_path}") return t5, clip # Cache paths and encoding portrait_t5, portrait_clip = None, None schnell_t5, schnell_clip = None, None sportfashion_t5, sportfashion_clip = None, None synthmocap_t5, synthmocap_clip = None, None if portrait_prompts: portrait_enc_cache = os.path.join(ENCODING_CACHE_DIR, f"portrait_encodings_{len(portrait_prompts)}.pt") portrait_t5, portrait_clip = load_or_encode(portrait_enc_cache, portrait_prompts, "portrait") if schnell_prompts: schnell_enc_cache = os.path.join(ENCODING_CACHE_DIR, f"schnell_encodings_{len(schnell_prompts)}.pt") schnell_t5, schnell_clip = load_or_encode(schnell_enc_cache, schnell_prompts, "schnell") if sportfashion_prompts: sportfashion_enc_cache = os.path.join(ENCODING_CACHE_DIR, f"sportfashion_encodings_{len(sportfashion_prompts)}.pt") sportfashion_t5, sportfashion_clip = load_or_encode(sportfashion_enc_cache, sportfashion_prompts, "sportfashion") if synthmocap_prompts: synthmocap_enc_cache = os.path.join(ENCODING_CACHE_DIR, f"synthmocap_encodings_{len(synthmocap_prompts)}.pt") synthmocap_t5, synthmocap_clip = load_or_encode(synthmocap_enc_cache, synthmocap_prompts, "synthmocap") # ============================================================================ # COMBINED DATASET CLASS WITH MASK SUPPORT # ============================================================================ class CombinedDataset(Dataset): """Combined dataset with mask support for weighted loss.""" def __init__( self, portrait_ds, portrait_indices, portrait_t5, portrait_clip, schnell_ds, schnell_t5, schnell_clip, sportfashion_ds, sportfashion_t5, sportfashion_clip, synthmocap_ds, synthmocap_t5, synthmocap_clip, vae, vae_scale, device, dtype, compute_masks=True ): self.portrait_ds = portrait_ds self.portrait_indices = portrait_indices self.portrait_t5 = portrait_t5 self.portrait_clip = portrait_clip self.schnell_ds = schnell_ds self.schnell_t5 = schnell_t5 self.schnell_clip = schnell_clip self.sportfashion_ds = sportfashion_ds self.sportfashion_t5 = sportfashion_t5 self.sportfashion_clip = sportfashion_clip self.synthmocap_ds = synthmocap_ds self.synthmocap_t5 = synthmocap_t5 self.synthmocap_clip = synthmocap_clip self.vae = vae self.vae_scale = vae_scale self.device = device self.dtype = dtype self.compute_masks = compute_masks # Dataset sizes (0 if disabled) self.n_portrait = len(portrait_indices) if portrait_indices else 0 self.n_schnell = len(schnell_ds) if schnell_ds else 0 self.n_sportfashion = len(sportfashion_ds) if sportfashion_ds else 0 self.n_synthmocap = len(synthmocap_ds) if synthmocap_ds else 0 # Cumulative indices for fast lookup self.c1 = self.n_portrait self.c2 = self.c1 + self.n_schnell self.c3 = self.c2 + self.n_sportfashion self.total = self.c3 + self.n_synthmocap def __len__(self): return self.total def _get_latent_from_array(self, latent_data): """Convert latent data to tensor.""" if isinstance(latent_data, torch.Tensor): return latent_data.to(self.dtype) return torch.tensor(np.array(latent_data), dtype=self.dtype) @torch.no_grad() def _encode_image(self, image): """Encode PIL image to VAE latent.""" if image.mode != "RGB": image = image.convert("RGB") if image.size != (512, 512): image = image.resize((512, 512), Image.Resampling.LANCZOS) img_tensor = torch.from_numpy(np.array(image)).float() / 255.0 img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0) img_tensor = (img_tensor * 2.0 - 1.0).to(self.device, dtype=self.dtype) latent = self.vae.encode(img_tensor).latent_dist.sample() latent = latent * self.vae_scale return latent.squeeze(0).cpu() def __getitem__(self, idx): mask = None # Default: no mask (uniform loss) if idx < self.c1: # Portrait sample (has pre-computed latent, no mask needed) orig_idx = self.portrait_indices[idx] item = self.portrait_ds[orig_idx] latent = self._get_latent_from_array(item["latent"]) t5 = self.portrait_t5[idx] clip = self.portrait_clip[idx] elif idx < self.c2: # Schnell sample (has pre-computed latent, no mask needed) schnell_idx = idx - self.c1 item = self.schnell_ds[schnell_idx] latent = self._get_latent_from_array(item["latent"]) t5 = self.schnell_t5[schnell_idx] clip = self.schnell_clip[schnell_idx] elif idx < self.c3: # SportFashion (needs VAE encoding + product mask) sf_idx = idx - self.c2 item = self.sportfashion_ds[sf_idx] image = item["image"] latent = self._encode_image(image) t5 = self.sportfashion_t5[sf_idx] clip = self.sportfashion_clip[sf_idx] if self.compute_masks: pixel_mask = create_product_mask(image) mask = downsample_mask_to_latent(pixel_mask, 64, 64) else: # SynthMoCap (needs VAE encoding + SMPL body mask) sm_idx = idx - self.c3 item = self.synthmocap_ds[sm_idx] image = item["image"] conditioning = item["conditioning_image"] latent = self._encode_image(image) t5 = self.synthmocap_t5[sm_idx] clip = self.synthmocap_clip[sm_idx] if self.compute_masks: pixel_mask = create_smpl_mask(conditioning) mask = downsample_mask_to_latent(pixel_mask, 64, 64) result = { "latent": latent, "t5_embed": t5.to(self.dtype), "clip_pooled": clip.to(self.dtype), } if mask is not None: result["mask"] = mask.to(self.dtype) return result # ============================================================================ # COLLATE FUNCTION # ============================================================================ def collate_fn(batch): latents = torch.stack([b["latent"] for b in batch]) t5_embeds = torch.stack([b["t5_embed"] for b in batch]) clip_pooled = torch.stack([b["clip_pooled"] for b in batch]) # Handle masks (some samples may not have masks) masks = None if any("mask" in b for b in batch): masks = [] for b in batch: if "mask" in b: masks.append(b["mask"]) else: # No mask = uniform weight (all 1s) masks.append(torch.ones(64, 64, dtype=latents.dtype)) masks = torch.stack(masks) return { "latents": latents, "t5_embeds": t5_embeds, "clip_pooled": clip_pooled, "masks": masks, } # ============================================================================ # MASKED LOSS FUNCTION # ============================================================================ def masked_mse_loss(pred, target, mask=None, fg_weight=2.0, bg_weight=0.5, snr_weights=None): """ Compute MSE loss with optional foreground/background weighting and min-SNR. Args: pred: [B, H*W, C] predicted velocity target: [B, H*W, C] target velocity mask: [B, H, W] foreground mask (1=foreground, 0=background) or None fg_weight: Weight for foreground pixels bg_weight: Weight for background pixels snr_weights: [B] min-SNR weights per sample or None Returns: Scalar loss value """ B, N, C = pred.shape if mask is None: # No spatial mask - compute per-sample loss loss_per_sample = ((pred - target) ** 2).mean(dim=[1, 2]) # [B] else: H = W = int(math.sqrt(N)) mask_flat = mask.view(B, H * W, 1).to(pred.device) sq_error = (pred - target) ** 2 weights = mask_flat * fg_weight + (1 - mask_flat) * bg_weight weighted_error = sq_error * weights loss_per_sample = weighted_error.mean(dim=[1, 2]) # [B] # Apply min-SNR weighting if provided if snr_weights is not None: loss_per_sample = loss_per_sample * snr_weights return loss_per_sample.mean() # ============================================================================ # CREATE DATASET # ============================================================================ print("\nCreating combined dataset...") combined_ds = CombinedDataset( portrait_ds, portrait_indices, portrait_t5, portrait_clip, schnell_ds, schnell_t5, schnell_clip, sportfashion_ds, sportfashion_t5, sportfashion_clip, synthmocap_ds, synthmocap_t5, synthmocap_clip, vae, VAE_SCALE, DEVICE, DTYPE, compute_masks=USE_MASKED_LOSS ) print(f"✓ Combined dataset: {len(combined_ds)} samples") print(f" - Portraits (3x): {combined_ds.n_portrait:,}") print(f" - Schnell teacher: {combined_ds.n_schnell:,}") print(f" - SportFashion: {combined_ds.n_sportfashion:,}") print(f" - SynthMoCap: {combined_ds.n_synthmocap:,}") # ============================================================================ # DATALOADER # ============================================================================ loader = DataLoader( combined_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=8, pin_memory=True, collate_fn=collate_fn, drop_last=True, ) print(f"✓ DataLoader: {len(loader)} batches/epoch") # ============================================================================ # SAMPLING FUNCTION # ============================================================================ @torch.inference_mode() def generate_samples(model, prompts, num_steps=28, guidance_scale=3.5, H=64, W=64, use_ema=True): was_training = model.training model.eval() if use_ema and 'ema' in globals() and ema is not None: ema.apply_shadow_for_eval(model) B = len(prompts) C = 16 t5_list, clip_list = [], [] for p in prompts: t5, clip = encode_prompt(p) t5_list.append(t5) clip_list.append(clip) t5_embeds = torch.stack(t5_list).to(DTYPE) clip_pooleds = torch.stack(clip_list).to(DTYPE) x = torch.randn(B, H * W, C, device=DEVICE, dtype=DTYPE) img_ids = TinyFluxDeep.create_img_ids(B, H, W, DEVICE) t_linear = torch.linspace(0, 1, num_steps + 1, device=DEVICE, dtype=DTYPE) timesteps = flux_shift(t_linear, s=SHIFT) for i in range(num_steps): t_curr = timesteps[i] t_next = timesteps[i + 1] dt = t_next - t_curr t_batch = t_curr.expand(B).to(DTYPE) guidance = torch.full((B,), guidance_scale, device=DEVICE, dtype=DTYPE) with torch.autocast("cuda", dtype=DTYPE): v_cond = model( hidden_states=x, encoder_hidden_states=t5_embeds, pooled_projections=clip_pooleds, timestep=t_batch, img_ids=img_ids, guidance=guidance, ) x = x + v_cond * dt latents = x.reshape(B, H, W, C).permute(0, 3, 1, 2) latents = latents / VAE_SCALE with torch.autocast("cuda", dtype=DTYPE): images = vae.decode(latents.to(vae.dtype)).sample images = (images / 2 + 0.5).clamp(0, 1) if use_ema and 'ema' in globals() and ema is not None: ema.restore(model) if was_training: model.train() return images def save_samples(images, prompts, step, output_dir): from torchvision.utils import save_image os.makedirs(output_dir, exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") grid_path = os.path.join(output_dir, f"samples_step_{step}.png") save_image(images, grid_path, nrow=2, padding=2) try: api.upload_file( path_or_fileobj=grid_path, path_in_repo=f"samples/{timestamp}_step_{step}.png", repo_id=HF_REPO, ) except: pass # ============================================================================ # CHECKPOINT LOADING WITH WEIGHT UPGRADE SUPPORT # ============================================================================ # Add this config flag near your other CONFIG section: # # ALLOW_WEIGHT_UPGRADE = True # Allow loading old checkpoints into new model # ============================================================================ def load_checkpoint(model, optimizer, scheduler, target): """ Load checkpoint with optional weight upgrade support. When ALLOW_WEIGHT_UPGRADE=True: - Missing Q/K norm weights are initialized to ones (identity transform) - Unexpected keys (e.g., old sin_basis caches) are ignored - Model behavior is identical to old weights at load time When ALLOW_WEIGHT_UPGRADE=False: - Requires exact weight match (strict=True) """ start_step = 0 start_epoch = 0 if target == "none": print("Starting fresh (no checkpoint)") return start_step, start_epoch ckpt_path = None weights_path = None if target == "latest": if os.path.exists(CHECKPOINT_DIR): ckpts = [f for f in os.listdir(CHECKPOINT_DIR) if f.startswith("step_") and f.endswith(".pt")] if ckpts: steps = [int(f.split("_")[1].split(".")[0]) for f in ckpts] latest_step = max(steps) ckpt_path = os.path.join(CHECKPOINT_DIR, f"step_{latest_step}.pt") weights_path = ckpt_path.replace(".pt", ".safetensors") elif target == "hub" or target.startswith("hub:"): try: from huggingface_hub import list_repo_files if target.startswith("hub:"): step_name = target.split(":")[1] weights_path = hf_hub_download(HF_REPO, f"checkpoints/{step_name}.safetensors") start_step = int(step_name.split("_")[1]) if "_" in step_name else 0 print(f"Downloaded {step_name} from hub") else: files = list_repo_files(HF_REPO) ckpts = [f for f in files if f.startswith("checkpoints/step_") and f.endswith(".safetensors") and "_ema" not in f] if ckpts: steps = [int(f.split("_")[1].split(".")[0]) for f in ckpts] latest = max(steps) weights_path = hf_hub_download(HF_REPO, f"checkpoints/step_{latest}.safetensors") start_step = latest print(f"Downloaded step_{latest} from hub") except Exception as e: print(f"Could not download from hub: {e}") return start_step, start_epoch elif target == "best": ckpt_path = os.path.join(CHECKPOINT_DIR, "best.pt") weights_path = ckpt_path.replace(".pt", ".safetensors") elif os.path.exists(target): # Direct path provided if target.endswith(".safetensors"): weights_path = target ckpt_path = target.replace(".safetensors", ".pt") else: ckpt_path = target weights_path = target.replace(".pt", ".safetensors") if weights_path and os.path.exists(weights_path): print(f"Loading weights from {weights_path}") state_dict = load_file(weights_path) state_dict = {k: v.to(DTYPE) if v.is_floating_point() else v for k, v in state_dict.items()} # Get model reference (handle torch.compile wrapper) model_ref = model._orig_mod if hasattr(model, '_orig_mod') else model if ALLOW_WEIGHT_UPGRADE: # Flexible loading with weight upgrade missing, unexpected = load_with_weight_upgrade(model_ref, state_dict) if missing: print(f" ℹ Initialized {len(missing)} new parameters (identity)") if unexpected: print(f" ℹ Ignored {len(unexpected)} deprecated parameters") else: # Strict loading - must match exactly model_ref.load_state_dict(state_dict, strict=True) print(f"✓ Loaded model weights") if ckpt_path and os.path.exists(ckpt_path): state = torch.load(ckpt_path, map_location="cpu") start_step = state.get("step", 0) start_epoch = state.get("epoch", 0) try: optimizer.load_state_dict(state["optimizer"]) scheduler.load_state_dict(state["scheduler"]) print(f"✓ Loaded optimizer/scheduler state") except: print(" ⚠ Could not load optimizer state (will use fresh optimizer)") print(f"Resuming from step {start_step}, epoch {start_epoch}") return start_step, start_epoch def load_with_weight_upgrade(model, state_dict): """ Load state dict with automatic handling of: - Missing Q/K norm weights → initialize to ones (identity) - Unexpected keys → ignore (e.g., old sin_basis caches) Returns: (missing_keys, unexpected_keys) - lists of handled keys """ model_state = model.state_dict() # Keys that are new in the repaired model (Q/K norms) QK_NORM_PATTERNS = [ '.norm_q.weight', '.norm_k.weight', '.norm_added_q.weight', '.norm_added_k.weight', ] # Keys that may exist in old checkpoints but not new model DEPRECATED_PATTERNS = [ '.sin_basis', # Old cached sin embeddings ] loaded_keys = [] missing_keys = [] unexpected_keys = [] initialized_keys = [] # First pass: load matching weights for key in state_dict.keys(): if key in model_state: if state_dict[key].shape == model_state[key].shape: model_state[key] = state_dict[key] loaded_keys.append(key) else: print(f" ⚠ Shape mismatch for {key}: checkpoint {state_dict[key].shape} vs model {model_state[key].shape}") unexpected_keys.append(key) else: # Check if it's a known deprecated key is_deprecated = any(pat in key for pat in DEPRECATED_PATTERNS) if is_deprecated: unexpected_keys.append(key) else: print(f" ⚠ Unexpected key (not in model): {key}") unexpected_keys.append(key) # Second pass: handle missing keys for key in model_state.keys(): if key not in loaded_keys: # Check if it's a Q/K norm that needs identity initialization is_qk_norm = any(pat in key for pat in QK_NORM_PATTERNS) if is_qk_norm: # Initialize to ones (identity transform for RMSNorm) model_state[key] = torch.ones_like(model_state[key]) initialized_keys.append(key) else: missing_keys.append(key) print(f" ⚠ Missing key (not in checkpoint): {key}") # Load the updated state model.load_state_dict(model_state, strict=False) # Report if initialized_keys: print(f" ✓ Initialized Q/K norms to identity ({len(initialized_keys)} params):") # Group by block for cleaner output blocks = set() for k in initialized_keys: if 'double_blocks' in k: block_num = k.split('.')[1] blocks.add(f"double_blocks.{block_num}") elif 'single_blocks' in k: block_num = k.split('.')[1] blocks.add(f"single_blocks.{block_num}") for block in sorted(blocks): print(f" - {block}.attn.norm_[q,k,added_q,added_k]") if unexpected_keys: deprecated = [k for k in unexpected_keys if any(p in k for p in DEPRECATED_PATTERNS)] if deprecated: print(f" ✓ Ignored deprecated keys: {deprecated}") return missing_keys, unexpected_keys # ============================================================================ # ALSO UPDATE save_checkpoint TO STRIP _orig_mod PREFIX # ============================================================================ def save_checkpoint(model, optimizer, scheduler, step, epoch, loss, path, ema_state=None): """Save checkpoint with proper handling of torch.compile wrapper.""" os.makedirs(os.path.dirname(path) if os.path.dirname(path) else ".", exist_ok=True) # Get state dict, handling torch.compile wrapper if hasattr(model, '_orig_mod'): state_dict = model._orig_mod.state_dict() else: state_dict = model.state_dict() # Ensure proper dtype for storage state_dict = {k: v.to(DTYPE) if v.is_floating_point() else v for k, v in state_dict.items()} # Save weights weights_path = path.replace(".pt", ".safetensors") save_file(state_dict, weights_path) # Save EMA weights if provided if ema_state is not None: ema_weights = {k: v.to(DTYPE) if v.is_floating_point() else v for k, v in ema_state['shadow'].items()} ema_weights_path = path.replace(".pt", "_ema.safetensors") save_file(ema_weights, ema_weights_path) # Save optimizer/scheduler state state = { "step": step, "epoch": epoch, "loss": loss, "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), } if ema_state is not None: state["ema_decay"] = ema_state.get('decay', EMA_DECAY) torch.save(state, path) print(f" ✓ Saved checkpoint: step {step}") return weights_path # ============================================================================ # CREATE MODEL # ============================================================================ print("\nCreating TinyFluxDeep model...") config = TinyFluxDeepConfig() model = TinyFluxDeep(config).to(device=DEVICE, dtype=DTYPE) total_params = sum(p.numel() for p in model.parameters()) print(f"Total parameters: {total_params:,}") trainable_params = [p for p in model.parameters() if p.requires_grad] print(f"Trainable parameters: {sum(p.numel() for p in trainable_params):,}") # ============================================================================ # OPTIMIZER # ============================================================================ opt = torch.optim.AdamW(trainable_params, lr=LR, betas=(0.9, 0.99), weight_decay=0.01, fused=True) total_steps = len(loader) * EPOCHS // GRAD_ACCUM warmup = min(1000, total_steps // 10) def lr_fn(step): if step < warmup: return step / warmup return 0.5 * (1 + math.cos(math.pi * (step - warmup) / (total_steps - warmup))) sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_fn) # ============================================================================ # LOAD CHECKPOINT # ============================================================================ start_step, start_epoch = load_checkpoint(model, opt, sched, LOAD_TARGET) if RESUME_STEP is not None: start_step = RESUME_STEP # ============================================================================ # COMPILE # ============================================================================ model = torch.compile(model, mode="default") # ============================================================================ # EMA # ============================================================================ print("Initializing EMA...") ema = EMA(model, decay=EMA_DECAY) # ============================================================================ # TENSORBOARD # ============================================================================ run_name = f"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}" writer = SummaryWriter(os.path.join(LOG_DIR, run_name)) # Sample prompts SAMPLE_PROMPTS = [ "a photo of a cat sitting on a windowsill", "a portrait of a woman with red hair", "a black backpack on white background", "a person standing in a t-pose", ] # ============================================================================ # TRAINING LOOP # ============================================================================ print(f"\n{'='*60}") print(f"Training TinyFlux-Deep") print(f"{'='*60}") print(f"Total: {len(combined_ds):,} samples") print(f"Epochs: {EPOCHS}, Steps/epoch: {len(loader)}, Total: {total_steps}") print(f"Batch: {BATCH_SIZE} x {GRAD_ACCUM} = {BATCH_SIZE * GRAD_ACCUM}") print(f"Masked loss: {USE_MASKED_LOSS} (fg={FG_LOSS_WEIGHT}, bg={BG_LOSS_WEIGHT})") print(f"Min-SNR gamma: {MIN_SNR_GAMMA}") print(f"Resume: step {start_step}, epoch {start_epoch}") model.train() step = start_step best = float("inf") for ep in range(start_epoch, EPOCHS): ep_loss = 0 ep_batches = 0 pbar = tqdm(loader, desc=f"E{ep + 1}") for i, batch in enumerate(pbar): latents = batch["latents"].to(DEVICE, non_blocking=True) t5 = batch["t5_embeds"].to(DEVICE, non_blocking=True) clip = batch["clip_pooled"].to(DEVICE, non_blocking=True) masks = batch["masks"] if masks is not None: masks = masks.to(DEVICE, non_blocking=True) B, C, H, W = latents.shape data = latents.permute(0, 2, 3, 1).reshape(B, H * W, C) noise = torch.randn_like(data) if TEXT_DROPOUT > 0: t5, clip, _ = apply_text_dropout(t5, clip, TEXT_DROPOUT) t = torch.sigmoid(torch.randn(B, device=DEVICE)) t = flux_shift(t, s=SHIFT).to(DTYPE).clamp(1e-4, 1 - 1e-4) t_expanded = t.view(B, 1, 1) x_t = (1 - t_expanded) * noise + t_expanded * data v_target = data - noise img_ids = TinyFluxDeep.create_img_ids(B, H, W, DEVICE) guidance = torch.rand(B, device=DEVICE, dtype=DTYPE) * 4 + 1 if GUIDANCE_DROPOUT > 0: guide_mask = torch.rand(B, device=DEVICE) < GUIDANCE_DROPOUT guidance[guide_mask] = 1.0 with torch.autocast("cuda", dtype=DTYPE): v_pred = model( hidden_states=x_t, encoder_hidden_states=t5, pooled_projections=clip, timestep=t, img_ids=img_ids, guidance=guidance, ) # Compute loss with min-SNR weighting snr_weights = min_snr_weight(t) # [B] # Unified loss: handles mask + SNR weighting loss = masked_mse_loss( v_pred, v_target, mask=masks if USE_MASKED_LOSS else None, fg_weight=FG_LOSS_WEIGHT, bg_weight=BG_LOSS_WEIGHT, snr_weights=snr_weights ) / GRAD_ACCUM loss.backward() if (i + 1) % GRAD_ACCUM == 0: grad_norm = torch.nn.utils.clip_grad_norm_(trainable_params, 1.0) opt.step() sched.step() opt.zero_grad(set_to_none=True) ema.update(model) step += 1 if step % LOG_EVERY == 0: writer.add_scalar("train/loss", loss.item() * GRAD_ACCUM, step) writer.add_scalar("train/lr", sched.get_last_lr()[0], step) writer.add_scalar("train/grad_norm", grad_norm.item(), step) if step % SAMPLE_EVERY == 0: print(f"\n Generating samples at step {step}...") images = generate_samples(model, SAMPLE_PROMPTS, num_steps=20, use_ema=True) save_samples(images, SAMPLE_PROMPTS, step, SAMPLE_DIR) if step % SAVE_EVERY == 0: ckpt_path = os.path.join(CHECKPOINT_DIR, f"step_{step}.pt") weights_path = save_checkpoint(model, opt, sched, step, ep, loss.item(), ckpt_path, ema_state=ema.state_dict()) if step % UPLOAD_EVERY == 0: upload_checkpoint(weights_path, step) ep_loss += loss.item() * GRAD_ACCUM ep_batches += 1 pbar.set_postfix(loss=f"{loss.item() * GRAD_ACCUM:.4f}", step=step) avg = ep_loss / max(ep_batches, 1) print(f"Epoch {ep + 1} loss: {avg:.4f}") if avg < best: best = avg weights_path = save_checkpoint(model, opt, sched, step, ep, avg, os.path.join(CHECKPOINT_DIR, "best.pt"), ema_state=ema.state_dict()) try: api.upload_file(path_or_fileobj=weights_path, path_in_repo="model.safetensors", repo_id=HF_REPO) except: pass print(f"\n✓ Training complete! Best loss: {best:.4f}") writer.close()