scripts/trainer_v3_expert_guidance.py

# ============================================================================
# TinyFlux-Deep Training Cell - With Expert Distillation (Precached)
# ============================================================================
# Integrates SD1.5-flow-lune as a frozen timestep expert.
# Expert features are PRECACHED at 10 timestep buckets for speed.
# The ExpertPredictor learns to emulate expert features from (t, CLIP).
# At inference, no expert needed - predictor runs standalone.
#
# USAGE: Run model.py cell first, then this cell
# ============================================================================

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 = 40
MAX_SEQ = 128
SHIFT = 3.0
DEVICE = "cuda"
DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16

ALLOW_WEIGHT_UPGRADE = True

# 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 = "hub:step_305000"
RESUME_STEP = None

# ============================================================================
# EXPERT DISTILLATION CONFIG
# ============================================================================
ENABLE_EXPERT_DISTILLATION = True
EXPERT_CHECKPOINT = "AbstractPhil/sd15-flow-lune-flux"
EXPERT_CHECKPOINT_PATH = "flux_t2_6_pose_t4_6_port_t1_4/checkpoint-00018765/unet/diffusion_pytorch_model.safetensors"
EXPERT_DIM = 1280
EXPERT_HIDDEN_DIM = 512
EXPERT_DROPOUT = 0.1  # Prob of forcing predictor (applied outside model)
DISTILL_LOSS_WEIGHT = 0.1
DISTILL_WARMUP_STEPS = 1000

# Timestep buckets for precaching
EXPERT_T_BUCKETS = torch.linspace(0.05, 0.95, 10)

# ============================================================================
# DATASET CONFIG
# ============================================================================
ENABLE_PORTRAIT = False
ENABLE_SCHNELL = True
ENABLE_SPORTFASHION = False
ENABLE_SYNTHMOCAP = False

PORTRAIT_REPO = "AbstractPhil/ffhq_flux_latents_repaired"
PORTRAIT_NUM_SHARDS = 11
SCHNELL_REPO = "AbstractPhil/flux-schnell-teacher-latents"
SCHNELL_CONFIGS = ["train_512"]
SPORTFASHION_REPO = "Pianokill/SportFashion_512x512"
SYNTHMOCAP_REPO = "toyxyz/SynthMoCap_smpl_512"

FG_LOSS_WEIGHT = 2.0
BG_LOSS_WEIGHT = 0.5
USE_MASKED_LOSS = False
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


# ============================================================================
# EXPERT FEATURE CACHE (precached, fast lookup + interpolation)
# ============================================================================

class ExpertFeatureCache:
    """
    Precached SD1.5-flow expert features with timestep interpolation.
    
    Features extracted at 10 timestep buckets [0.05, 0.15, ..., 0.95].
    At runtime, interpolates between nearest buckets.
    """
    
    def __init__(self, features: torch.Tensor, t_buckets: torch.Tensor, dtype=torch.float16):
        self.features = features.to(dtype)  # [N, 10, 1280]
        self.t_buckets = t_buckets
        self.t_min = t_buckets[0].item()
        self.t_max = t_buckets[-1].item()
        self.t_step = (t_buckets[1] - t_buckets[0]).item()
        self.n_buckets = len(t_buckets)
        self.dtype = dtype
        
    def get_features(self, indices: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
        """
        Get interpolated expert features.
        
        Args:
            indices: [B] sample indices into dataset
            timesteps: [B] timesteps in [0, 1]
        
        Returns:
            [B, 1280] interpolated features
        """
        device = timesteps.device
        
        # Clamp to valid range
        t_clamped = timesteps.float().clamp(self.t_min, self.t_max)
        
        # Find bucket indices
        t_idx_float = (t_clamped - self.t_min) / self.t_step
        t_idx_low = t_idx_float.long().clamp(0, self.n_buckets - 2)
        t_idx_high = (t_idx_low + 1).clamp(0, self.n_buckets - 1)
        
        # Interpolation alpha
        alpha = (t_idx_float - t_idx_low.float()).unsqueeze(-1)  # [B, 1]
        
        # Gather (on CPU for large caches)
        idx_cpu = indices.cpu()
        t_low_cpu = t_idx_low.cpu()
        t_high_cpu = t_idx_high.cpu()
        
        f_low = self.features[idx_cpu, t_low_cpu]   # [B, 1280]
        f_high = self.features[idx_cpu, t_high_cpu]  # [B, 1280]
        
        # Interpolate and move to device
        result = (1 - alpha.cpu()) * f_low + alpha.cpu() * f_high
        return result.to(device=device, dtype=self.dtype)


def load_or_extract_expert_features(cache_path: str, prompts: List[str], name: str, 
                                     clip_tok, clip_enc, t_buckets: torch.Tensor,
                                     batch_size: int = 32) -> Optional[ExpertFeatureCache]:
    """
    Load cached expert features or extract them from SD1.5-flow.
    Follows same pattern as load_or_encode for text embeddings.
    """
    if not prompts or not ENABLE_EXPERT_DISTILLATION:
        return None
    
    # Check cache
    if os.path.exists(cache_path):
        print(f"Loading cached {name} expert features...")
        cached = torch.load(cache_path, map_location="cpu")
        cache = ExpertFeatureCache(cached["features"], cached["t_buckets"], DTYPE)
        print(f"  ✓ Loaded {cache.features.shape[0]} samples × {cache.n_buckets} timesteps")
        return cache
    
    # Extract features
    print(f"Extracting {name} expert features ({len(prompts)} × {len(t_buckets)} timesteps)...")
    print(f"  This is a one-time operation, will be cached for future runs.")
    
    # Load expert model temporarily
    checkpoint_path = hf_hub_download(
        repo_id=EXPERT_CHECKPOINT,
        filename=EXPERT_CHECKPOINT_PATH,
    )
    
    from diffusers import UNet2DConditionModel
    unet = UNet2DConditionModel.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5",
        subfolder="unet",
        torch_dtype=DTYPE,
    ).to(DEVICE).eval()
    
    state_dict = load_file(checkpoint_path)
    unet.load_state_dict(state_dict, strict=False)
    
    for p in unet.parameters():
        p.requires_grad = False
    
    # Hook for mid-block features
    mid_features = [None]
    def hook_fn(module, inp, out):
        mid_features[0] = out.mean(dim=[2, 3])
    unet.mid_block.register_forward_hook(hook_fn)
    
    # Extract
    n_prompts = len(prompts)
    n_buckets = len(t_buckets)
    all_features = torch.zeros(n_prompts, n_buckets, EXPERT_DIM, dtype=torch.float16)
    
    with torch.no_grad():
        for start_idx in tqdm(range(0, n_prompts, batch_size), desc=f"Extracting {name}"):
            end_idx = min(start_idx + batch_size, n_prompts)
            batch_prompts = prompts[start_idx:end_idx]
            B = len(batch_prompts)
            
            # Encode CLIP hidden states
            clip_inputs = clip_tok(
                batch_prompts, return_tensors="pt", padding="max_length",
                max_length=77, truncation=True
            ).to(DEVICE)
            clip_hidden = clip_enc(**clip_inputs).last_hidden_state  # [B, 77, 768]
            
            # Extract at each timestep bucket
            for t_idx, t_val in enumerate(t_buckets):
                timesteps = torch.full((B,), t_val.item(), device=DEVICE)
                latents = torch.randn(B, 4, 64, 64, device=DEVICE, dtype=DTYPE)
                
                _ = unet(latents, timesteps * 1000, encoder_hidden_states=clip_hidden.to(DTYPE))
                
                all_features[start_idx:end_idx, t_idx] = mid_features[0].cpu().to(torch.float16)
    
    # Cleanup
    del unet
    torch.cuda.empty_cache()
    
    # Save cache
    torch.save({"features": all_features, "t_buckets": t_buckets}, cache_path)
    print(f"  ✓ Cached to {cache_path}")
    print(f"  Size: {all_features.numel() * 2 / 1e9:.2f} GB")
    
    return ExpertFeatureCache(all_features, t_buckets, DTYPE)


# ============================================================================
# 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)
    
    def load_shadow(self, shadow_state):
        """Load EMA shadow weights, handling architecture changes gracefully."""
        device = next(iter(self.shadow.values())).device if self.shadow else 'cuda'
        
        loaded = 0
        skipped_old = 0
        kept_new = 0
        
        for k, v in shadow_state.items():
            if k in self.shadow:
                # Key exists in current model - load it
                self.shadow[k] = v.clone().to(device)
                loaded += 1
            else:
                # Key doesn't exist (deprecated like guidance_in)
                skipped_old += 1
        
        # Count new keys not in checkpoint
        for k in self.shadow:
            if k not in shadow_state:
                kept_new += 1
        
        print(f"  ✓ Restored EMA: {loaded} loaded, {skipped_old} deprecated skipped, {kept_new} new (fresh init)")


# ============================================================================
# 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]:
    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],
        img[:sample_size, -sample_size:],
        img[-sample_size:, :sample_size],
        img[-sample_size:, -sample_size:],
    ]
    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:
    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)
    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:
    img = np.array(conditioning_image).astype(np.float32)
    if len(img.shape) == 2:
        return (img > threshold).astype(np.float32)
    r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
    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:
    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):
    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:
    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(DEVICE, dtype=DTYPE)
    latent = vae.encode(img_tensor).latent_dist.sample()
    latent = latent * VAE_SCALE
    return latent.squeeze(0).cpu()


# ============================================================================
# LOAD DATASETS
# ============================================================================

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")
    print("  Extracting prompts (columnar)...")
    florence_list = list(portrait_ds["text_florence"])
    llava_list = list(portrait_ds["text_llava"])
    blip_list = list(portrait_ds["text_blip"])
    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")

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")

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")

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


# Standard text encodings
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")


# ============================================================================
# EXTRACT/LOAD EXPERT FEATURES (precached)
# ============================================================================
print("\n" + "="*60)
print("Expert Feature Caching")
print("="*60)

schnell_expert_cache = None
portrait_expert_cache = None
sportfashion_expert_cache = None
synthmocap_expert_cache = None

if schnell_prompts and ENABLE_EXPERT_DISTILLATION:
    schnell_expert_path = os.path.join(ENCODING_CACHE_DIR, f"schnell_expert_{len(schnell_prompts)}.pt")
    schnell_expert_cache = load_or_extract_expert_features(
        schnell_expert_path, schnell_prompts, "schnell",
        clip_tok, clip_enc, EXPERT_T_BUCKETS
    )

if portrait_prompts and ENABLE_EXPERT_DISTILLATION:
    portrait_expert_path = os.path.join(ENCODING_CACHE_DIR, f"portrait_expert_{len(portrait_prompts)}.pt")
    portrait_expert_cache = load_or_extract_expert_features(
        portrait_expert_path, portrait_prompts, "portrait",
        clip_tok, clip_enc, EXPERT_T_BUCKETS
    )

if sportfashion_prompts and ENABLE_EXPERT_DISTILLATION:
    sportfashion_expert_path = os.path.join(ENCODING_CACHE_DIR, f"sportfashion_expert_{len(sportfashion_prompts)}.pt")
    sportfashion_expert_cache = load_or_extract_expert_features(
        sportfashion_expert_path, sportfashion_prompts, "sportfashion",
        clip_tok, clip_enc, EXPERT_T_BUCKETS
    )

if synthmocap_prompts and ENABLE_EXPERT_DISTILLATION:
    synthmocap_expert_path = os.path.join(ENCODING_CACHE_DIR, f"synthmocap_expert_{len(synthmocap_prompts)}.pt")
    synthmocap_expert_cache = load_or_extract_expert_features(
        synthmocap_expert_path, synthmocap_prompts, "synthmocap",
        clip_tok, clip_enc, EXPERT_T_BUCKETS
    )


# ============================================================================
# COMBINED DATASET CLASS (with sample_idx for expert lookup)
# ============================================================================
class CombinedDataset(Dataset):
    """Combined dataset returning sample index for expert feature lookup."""
    
    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
        
        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
        
        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):
        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):
        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
        
        # Determine which dataset and local index
        if idx < self.c1:
            # Portrait
            local_idx = idx
            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]
            dataset_id = 0
            
        elif idx < self.c2:
            # Schnell
            local_idx = idx - self.c1
            item = self.schnell_ds[local_idx]
            latent = self._get_latent_from_array(item["latent"])
            t5 = self.schnell_t5[local_idx]
            clip = self.schnell_clip[local_idx]
            dataset_id = 1
            
        elif idx < self.c3:
            # SportFashion
            local_idx = idx - self.c2
            item = self.sportfashion_ds[local_idx]
            image = item["image"]
            latent = self._encode_image(image)
            t5 = self.sportfashion_t5[local_idx]
            clip = self.sportfashion_clip[local_idx]
            dataset_id = 2
            if self.compute_masks:
                pixel_mask = create_product_mask(image)
                mask = downsample_mask_to_latent(pixel_mask, 64, 64)
            
        else:
            # SynthMoCap
            local_idx = idx - self.c3
            item = self.synthmocap_ds[local_idx]
            image = item["image"]
            conditioning = item["conditioning_image"]
            latent = self._encode_image(image)
            t5 = self.synthmocap_t5[local_idx]
            clip = self.synthmocap_clip[local_idx]
            dataset_id = 3
            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),
            "sample_idx": idx,  # Global index for expert cache lookup
            "local_idx": local_idx,  # Local index within dataset
            "dataset_id": dataset_id,  # Which dataset (0=portrait, 1=schnell, etc)
        }
        
        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])
    sample_indices = torch.tensor([b["sample_idx"] for b in batch], dtype=torch.long)
    local_indices = torch.tensor([b["local_idx"] for b in batch], dtype=torch.long)
    dataset_ids = torch.tensor([b["dataset_id"] for b in batch], dtype=torch.long)
    
    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:
                masks.append(torch.ones(64, 64, dtype=latents.dtype))
        masks = torch.stack(masks)
    
    return {
        "latents": latents,
        "t5_embeds": t5_embeds,
        "clip_pooled": clip_pooled,
        "sample_indices": sample_indices,
        "local_indices": local_indices,
        "dataset_ids": dataset_ids,
        "masks": masks,
    }


# ============================================================================
# EXPERT FEATURE LOOKUP (handles multiple datasets)
# ============================================================================
def get_expert_features_for_batch(
    local_indices: torch.Tensor,
    dataset_ids: torch.Tensor,
    timesteps: torch.Tensor,
    portrait_cache: Optional[ExpertFeatureCache],
    schnell_cache: Optional[ExpertFeatureCache],
    sportfashion_cache: Optional[ExpertFeatureCache],
    synthmocap_cache: Optional[ExpertFeatureCache],
) -> Optional[torch.Tensor]:
    """Get expert features from the appropriate cache for each sample."""
    
    caches = [portrait_cache, schnell_cache, sportfashion_cache, synthmocap_cache]
    
    # Check if any cache is available
    if not any(c is not None for c in caches):
        return None
    
    B = local_indices.shape[0]
    device = timesteps.device
    features = torch.zeros(B, EXPERT_DIM, device=device, dtype=DTYPE)
    
    for ds_id, cache in enumerate(caches):
        if cache is None:
            continue
        
        # Find samples from this dataset
        mask = dataset_ids == ds_id
        if not mask.any():
            continue
        
        # Get features for these samples
        ds_local_indices = local_indices[mask]
        ds_timesteps = timesteps[mask]
        ds_features = cache.get_features(ds_local_indices, ds_timesteps)
        features[mask] = ds_features
    
    return features


# ============================================================================
# MASKED LOSS FUNCTION
# ============================================================================
def masked_mse_loss(pred, target, mask=None, fg_weight=2.0, bg_weight=0.5, snr_weights=None):
    B, N, C = pred.shape
    if mask is None:
        loss_per_sample = ((pred - target) ** 2).mean(dim=[1, 2])
    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])
    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:,}")
print(f"  - Expert distillation: {ENABLE_EXPERT_DISTILLATION}")


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

        with torch.autocast("cuda", dtype=DTYPE):
            # No expert_features at inference - predictor runs standalone
            v_cond = model(
                hidden_states=x,
                encoder_hidden_states=t5_embeds,
                pooled_projections=clip_pooleds,
                timestep=t_batch,
                img_ids=img_ids,
            )
        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 FUNCTIONS
# ============================================================================
def save_checkpoint(model, optimizer, scheduler, step, epoch, loss, path, ema=None):
    os.makedirs(os.path.dirname(path) if os.path.dirname(path) else ".", exist_ok=True)
    if hasattr(model, '_orig_mod'):
        state_dict = model._orig_mod.state_dict()
    else:
        state_dict = model.state_dict()
    state_dict = {k: v.to(DTYPE) if v.is_floating_point() else v for k, v in state_dict.items()}
    weights_path = path.replace(".pt", ".safetensors")
    save_file(state_dict, weights_path)
    if ema is not None:
        ema_weights = {k: v.to(DTYPE) if v.is_floating_point() else v for k, v in ema.shadow.items()}
        ema_weights_path = path.replace(".pt", "_ema.safetensors")
        save_file(ema_weights, ema_weights_path)
    state = {
        "step": step,
        "epoch": epoch,
        "loss": loss,
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict(),
    }
    if ema is not None:
        state["ema_decay"] = ema.decay
    torch.save(state, path)
    print(f"  ✓ Saved checkpoint: step {step}")
    return weights_path


def upload_checkpoint(weights_path, step):
    try:
        api.upload_file(
            path_or_fileobj=weights_path,
            path_in_repo=f"checkpoints/step_{step}.safetensors",
            repo_id=HF_REPO,
        )
        ema_path = weights_path.replace(".safetensors", "_ema.safetensors")
        if os.path.exists(ema_path):
            api.upload_file(
                path_or_fileobj=ema_path,
                path_in_repo=f"checkpoints/step_{step}_ema.safetensors",
                repo_id=HF_REPO,
            )
        print(f"  ✓ Uploaded checkpoint to {HF_REPO}")
    except Exception as e:
        print(f"  ⚠ Upload failed: {e}")


def load_with_weight_upgrade(model, state_dict):
    """
    Load state dict with automatic handling of:
      - Missing ExpertPredictor weights → initialize fresh
      - Missing Q/K norm weights → initialize to ones (identity)
      - Unexpected keys → ignore (e.g., old guidance_in, sin_basis caches)
    """
    model_state = model.state_dict()
    
    # Patterns for new weights that may not exist in old checkpoints
    NEW_WEIGHT_PATTERNS = [
        'expert_predictor.',  # New ExpertPredictor module
        '.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 = [
        'guidance_in.',  # Replaced by expert_predictor
        '.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:
            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:
            is_new = any(pat in key for pat in NEW_WEIGHT_PATTERNS)
            
            if is_new:
                # Keep default initialization for new modules
                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:
        # Group by module for cleaner output
        modules = set()
        for k in initialized_keys:
            parts = k.split('.')
            if len(parts) >= 2:
                modules.add(parts[0] + '.' + parts[1] if parts[0] == 'expert_predictor' else parts[0])
        print(f"  ✓ Initialized new modules (fresh): {sorted(modules)}")
    
    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: {len(deprecated)} (guidance_in, etc)")
    
    return missing_keys, unexpected_keys


def load_checkpoint(model, optimizer, scheduler, target):
    """
    Load checkpoint with weight upgrade support for ExpertPredictor.
    
    When ALLOW_WEIGHT_UPGRADE=True:
      - Missing ExpertPredictor weights are initialized fresh
      - Old guidance_in weights are ignored
      - Model continues training with new architecture
    """
    start_step = 0
    start_epoch = 0
    ema_state = None
    
    if target == "none":
        print("Starting fresh (no checkpoint)")
        return start_step, start_epoch, None
    
    ckpt_path = None
    weights_path = None
    ema_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")
                ema_weights_path = ckpt_path.replace(".pt", "_ema.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")
                try:
                    ema_weights_path = hf_hub_download(HF_REPO, f"checkpoints/{step_name}_ema.safetensors")
                    print(f"  Found EMA weights on hub")
                except:
                    ema_weights_path = None
                    print(f"  No EMA weights on hub (will start fresh)")
                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")
                    try:
                        ema_weights_path = hf_hub_download(HF_REPO, f"checkpoints/step_{latest}_ema.safetensors")
                        print(f"  Found EMA weights on hub")
                    except:
                        ema_weights_path = None
                        print(f"  No EMA weights on hub (will start fresh)")
                    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, None
            
    elif target == "best":
        ckpt_path = os.path.join(CHECKPOINT_DIR, "best.pt")
        weights_path = ckpt_path.replace(".pt", ".safetensors")
        ema_weights_path = ckpt_path.replace(".pt", "_ema.safetensors")
        
    elif os.path.exists(target):
        if target.endswith(".safetensors"):
            weights_path = target
            ckpt_path = target.replace(".safetensors", ".pt")
            ema_weights_path = target.replace(".safetensors", "_ema.safetensors")
        else:
            ckpt_path = target
            weights_path = target.replace(".pt", ".safetensors")
            ema_weights_path = target.replace(".pt", "_ema.safetensors")
    
    # Load main model weights
    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"  ⚠ {len(missing)} truly missing parameters (may need attention)")
        else:
            # Strict loading - must match exactly
            model_ref.load_state_dict(state_dict, strict=True)
        
        print(f"✓ Loaded model weights")
        
        # Load EMA weights if they exist
        if ema_weights_path and os.path.exists(ema_weights_path):
            ema_state = load_file(ema_weights_path)
            ema_state = {k: v.to(DTYPE) if v.is_floating_point() else v for k, v in ema_state.items()}
            print(f"✓ Loaded EMA weights ({len(ema_state)} params)")
        else:
            print(f"  ℹ No EMA weights found (will initialize fresh)")
    else:
        print(f"  ⚠ Weights file not found: {weights_path}")
        print(f"  Starting with fresh model")
        return start_step, start_epoch, None
    
    # Load optimizer/scheduler state
    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 Exception as e:
            print(f"  ⚠ Could not load optimizer state: {e}")
            print(f"  Will use fresh optimizer (this is fine for architecture changes)")
        print(f"Resuming from step {start_step}, epoch {start_epoch}")
    
    return start_step, start_epoch, ema_state


# ============================================================================
# CREATE MODEL
# ============================================================================
print("\nCreating TinyFluxDeep model with ExpertPredictor...")

config = TinyFluxDeepConfig(
    use_expert_predictor=ENABLE_EXPERT_DISTILLATION,
    expert_dim=EXPERT_DIM,
    expert_hidden_dim=EXPERT_HIDDEN_DIM,
    expert_dropout=EXPERT_DROPOUT,
    guidance_embeds=False,
)
model = TinyFluxDeep(config).to(device=DEVICE, dtype=DTYPE)

total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params:,}")

if hasattr(model, 'expert_predictor') and model.expert_predictor is not None:
    expert_params = sum(p.numel() for p in model.expert_predictor.parameters())
    print(f"Expert predictor parameters: {expert_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, ema_state = 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)
if ema_state is not None:
    ema.load_shadow(ema_state)
else:
    print("  Starting fresh EMA from current weights")


# ============================================================================
# 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 = [
    "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",
]


# ============================================================================
# DISTILLATION WEIGHT SCHEDULE
# ============================================================================
def get_distill_weight(step):
    if step < DISTILL_WARMUP_STEPS:
        return DISTILL_LOSS_WEIGHT * (step / DISTILL_WARMUP_STEPS)
    return DISTILL_LOSS_WEIGHT


# ============================================================================
# TRAINING LOOP
# ============================================================================
print(f"\n{'='*60}")
print(f"Training TinyFlux-Deep with Expert Distillation (Precached)")
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"Expert distillation: {ENABLE_EXPERT_DISTILLATION} (PRECACHED)")
if ENABLE_EXPERT_DISTILLATION:
    print(f"  - Expert: {EXPERT_CHECKPOINT}")
    print(f"  - Timestep buckets: {len(EXPERT_T_BUCKETS)}")
    print(f"  - Distill weight: {DISTILL_LOSS_WEIGHT} (warmup: {DISTILL_WARMUP_STEPS} steps)")
    print(f"  - Expert dropout: {EXPERT_DROPOUT}")
print(f"Masked loss: {USE_MASKED_LOSS}")
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_main_loss = 0
    ep_distill_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)
        local_indices = batch["local_indices"]
        dataset_ids = batch["dataset_ids"]
        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)

        # Get expert features from CACHE (fast!)
        expert_features = None
        if ENABLE_EXPERT_DISTILLATION:
            expert_features = get_expert_features_for_batch(
                local_indices, dataset_ids, t,
                portrait_expert_cache, schnell_expert_cache,
                sportfashion_expert_cache, synthmocap_expert_cache,
            )
            
            # Apply dropout OUTSIDE model (no graph break)
            if expert_features is not None and random.random() < EXPERT_DROPOUT:
                expert_features = None

        with torch.autocast("cuda", dtype=DTYPE):
            v_pred, expert_info = model(
                hidden_states=x_t,
                encoder_hidden_states=t5,
                pooled_projections=clip,
                timestep=t,
                img_ids=img_ids,
                expert_features=expert_features,
                return_expert_pred=True,
            )

        # Compute losses
        snr_weights = min_snr_weight(t)
        
        main_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
        )
        
        # Distillation loss
        distill_loss = torch.tensor(0.0, device=DEVICE)
        if expert_features is not None and expert_info is not None and 'expert_pred' in expert_info:
            distill_weight = get_distill_weight(step)
            distill_loss = F.mse_loss(expert_info['expert_pred'], expert_features)
            total_loss = main_loss + distill_weight * distill_loss
        else:
            total_loss = main_loss
        
        loss = total_loss / 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", total_loss.item(), step)
                writer.add_scalar("train/main_loss", main_loss.item(), step)
                if ENABLE_EXPERT_DISTILLATION:
                    writer.add_scalar("train/distill_loss", distill_loss.item(), step)
                    writer.add_scalar("train/distill_weight", get_distill_weight(step), 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, total_loss.item(), ckpt_path, ema=ema)
                if step % UPLOAD_EVERY == 0:
                    upload_checkpoint(weights_path, step)

        ep_loss += total_loss.item()
        ep_main_loss += main_loss.item()
        ep_distill_loss += distill_loss.item()
        ep_batches += 1
        
        pbar.set_postfix(
            loss=f"{total_loss.item():.4f}",
            main=f"{main_loss.item():.4f}",
            dist=f"{distill_loss.item():.4f}" if ENABLE_EXPERT_DISTILLATION else "off",
            step=step
        )

    avg = ep_loss / max(ep_batches, 1)
    avg_main = ep_main_loss / max(ep_batches, 1)
    avg_distill = ep_distill_loss / max(ep_batches, 1)
    
    print(f"Epoch {ep + 1} - total: {avg:.4f}, main: {avg_main:.4f}, distill: {avg_distill:.4f}")

    if avg < best:
        best = avg
        weights_path = save_checkpoint(model, opt, sched, step, ep, avg, os.path.join(CHECKPOINT_DIR, "best.pt"), ema=ema)
        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()