scripts/trainer_v2.py

# ============================================================================
# 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()