scripts/trainer_v4_testing.py

# ============================================================================
# TinyFlux-Deep v4.1 Training Cell - Dual Expert Distillation (Lune + Sol)
# ============================================================================
# Integrates:
#   - Lune: SD1.5-flow trajectory guidance (mid-block features)
#   - Sol: Geometric attention prior (attention statistics + spatial importance)
#
# Both expert features are PRECACHED at 10 timestep buckets for speed.
# At inference, predictors run standalone - no teachers needed.
#
# USAGE: Run model_v4.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 = 8
GRAD_ACCUM = 4
LR = 3e-4
EPOCHS = 10
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 = 1562
UPLOAD_EVERY = 1562
SAMPLE_EVERY = 781
LOG_EVERY = 200
LOG_UPLOAD_EVERY = 1562

# Checkpoint loading
# v4.1 init checkpoint (converted from v3 step_401434)
# Options:
#   "hub:checkpoint_runs/v4_init/lailah_401434_v4_init" - v4.1 init (no EMA, fresh Sol)
#   "hub:step_401434" - v3 checkpoint (will auto-remap expert_predictor -> lune_predictor)
#   "latest" - latest local checkpoint
#   "none" - start fresh
LOAD_TARGET = "hub:checkpoint_runs/v4_init/lailah_401434_v4_init"
RESUME_STEP = 401434

# ============================================================================
# EXPERT REPOSITORY (both Lune and Sol)
# ============================================================================
EXPERTS_REPO = "AbstractPhil/tinyflux-experts"

# ============================================================================
# LUNE EXPERT DISTILLATION CONFIG (trajectory guidance)
# ============================================================================
ENABLE_LUNE_DISTILLATION = True
LUNE_FILENAME = "sd15-flow-lune-unet.safetensors"
LUNE_DIM = 1280  # SD1.5 mid-block dimension
LUNE_HIDDEN_DIM = 512
LUNE_DROPOUT = 0.1
LUNE_LOSS_WEIGHT = 0.1
LUNE_WARMUP_STEPS = 1000
LUNE_DISTILL_MODE = "cosine"  # "hard", "soft", "cosine", "huber"

# ============================================================================
# SOL ATTENTION PRIOR CONFIG (structural guidance)
# ============================================================================
ENABLE_SOL_DISTILLATION = True
SOL_FILENAME = "sd15-flow-sol-unet.safetensors"
SOL_HIDDEN_DIM = 256
SOL_SPATIAL_SIZE = 8  # 8x8 spatial importance map
SOL_GEOMETRIC_WEIGHT = 0.7  # 70% geometric, 30% learned
SOL_LOSS_WEIGHT = 0.05
SOL_WARMUP_STEPS = 2000  # Start Sol later than Lune

# Timestep buckets for precaching (shared by Lune and Sol)
EXPERT_T_BUCKETS = torch.linspace(0.05, 0.95, 10)

# ============================================================================
# LOSS CONFIG
# ============================================================================
USE_HUBER_LOSS = True
HUBER_DELTA = 0.1
USE_SPATIAL_WEIGHTING = False  # Weight main loss by Sol spatial importance

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

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"
IMAGENET_REPO = "AbstractPhil/synthetic-imagenet-1k"
IMAGENET_SUBSET = "schnell_512"
OBJECT_RELATIONS_REPO = "AbstractPhil/synthetic-object-relations"

# Confidence threshold for misprediction filtering
IMAGENET_CONFIDENCE_THRESHOLD = 0.5  # If confident but wrong, remove label

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


# ============================================================================
# LUNE FEATURE CACHE (SD1.5 mid-block features)
# ============================================================================
class LuneFeatureCache:
    """
    Precached SD1.5-flow Lune features with timestep interpolation.
    Features extracted at 10 timestep buckets [0.05, 0.15, ..., 0.95].
    """

    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:
        device = timesteps.device
        t_clamped = timesteps.float().clamp(self.t_min, self.t_max)
        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)
        alpha = (t_idx_float - t_idx_low.float()).unsqueeze(-1)

        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]
        f_high = self.features[idx_cpu, t_high_cpu]

        result = (1 - alpha.cpu()) * f_low + alpha.cpu() * f_high
        return result.to(device=device, dtype=self.dtype)


# ============================================================================
# SOL FEATURE CACHE (attention statistics + spatial importance)
# ============================================================================
class SolFeatureCache:
    """
    Precached Sol attention statistics with timestep interpolation.
    
    Statistics per sample per timestep:
      - stats: [N, 10, 4] - locality, entropy, clustering, sparsity
      - spatial: [N, 10, 8, 8] - spatial importance map
    """

    def __init__(self, stats: torch.Tensor, spatial: torch.Tensor, 
                 t_buckets: torch.Tensor, dtype=torch.float16):
        self.stats = stats.to(dtype)  # [N, 10, 4]
        self.spatial = spatial.to(dtype)  # [N, 10, 8, 8]
        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) -> Tuple[torch.Tensor, torch.Tensor]:
        device = timesteps.device
        t_clamped = timesteps.float().clamp(self.t_min, self.t_max)
        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)
        
        alpha_stats = (t_idx_float - t_idx_low.float()).unsqueeze(-1)
        alpha_spatial = alpha_stats.unsqueeze(-1)

        idx_cpu = indices.cpu()
        t_low_cpu = t_idx_low.cpu()
        t_high_cpu = t_idx_high.cpu()

        s_low = self.stats[idx_cpu, t_low_cpu]
        s_high = self.stats[idx_cpu, t_high_cpu]
        stats_result = (1 - alpha_stats.cpu()) * s_low + alpha_stats.cpu() * s_high

        sp_low = self.spatial[idx_cpu, t_low_cpu]
        sp_high = self.spatial[idx_cpu, t_high_cpu]
        spatial_result = (1 - alpha_spatial.cpu()) * sp_low + alpha_spatial.cpu() * sp_high

        return (
            stats_result.to(device=device, dtype=self.dtype),
            spatial_result.to(device=device, dtype=self.dtype)
        )


def load_or_extract_lune_features(cache_path: str, prompts: List[str], name: str,
                                   clip_tok, clip_enc, t_buckets: torch.Tensor,
                                   batch_size: int = 32) -> Optional[LuneFeatureCache]:
    """Load cached Lune features or extract from SD1.5-flow teacher."""
    if not prompts or not ENABLE_LUNE_DISTILLATION:
        return None

    if os.path.exists(cache_path):
        print(f"Loading cached {name} Lune features...")
        cached = torch.load(cache_path, map_location="cpu")
        cache = LuneFeatureCache(cached["features"], cached["t_buckets"], DTYPE)
        print(f"  ✓ Loaded {cache.features.shape[0]} samples × {cache.n_buckets} timesteps")
        return cache

    print(f"Extracting {name} Lune features ({len(prompts)} × {len(t_buckets)} timesteps)...")
    print(f"  This is a one-time operation, will be cached.")

    checkpoint_path = hf_hub_download(
        repo_id=EXPERTS_REPO,
        filename=LUNE_FILENAME,
    )
    print(f"  Loaded Lune from {EXPERTS_REPO}/{LUNE_FILENAME}")

    from diffusers import UNet2DConditionModel
    unet = UNet2DConditionModel.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5",
        subfolder="unet",
        torch_dtype=torch.float16,
    ).to(DEVICE).eval()

    state_dict = load_file(checkpoint_path)
    unet.load_state_dict(state_dict, strict=False)
    
    # Convert to fp16 and compile for speed
    unet = unet.half()
    unet = torch.compile(unet, mode="reduce-overhead")
    print(f"  ✓ Lune UNet compiled (fp16)")

    for p in unet.parameters():
        p.requires_grad = False

    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)

    n_prompts = len(prompts)
    n_buckets = len(t_buckets)
    all_features = torch.zeros(n_prompts, n_buckets, LUNE_DIM, dtype=torch.float16)

    # A100 can handle large batches - 64 prompts × 10 timesteps = 640 UNet forward passes batched
    # SD1.5 UNet at 64x64 latents uses ~2GB for batch of 64, so 640 samples ~10-15GB
    LUNE_BATCH_PROMPTS = 64  # Number of prompts per iteration
    
    with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.float16):
        for start_idx in tqdm(range(0, n_prompts, LUNE_BATCH_PROMPTS), desc=f"Extracting {name} Lune"):
            end_idx = min(start_idx + LUNE_BATCH_PROMPTS, n_prompts)
            batch_prompts = prompts[start_idx:end_idx]
            B = len(batch_prompts)

            # Encode CLIP once per prompt batch
            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]

            # Expand for all timesteps: [B * n_buckets, 77, 768]
            clip_expanded = clip_hidden.unsqueeze(1).expand(-1, n_buckets, -1, -1)
            clip_expanded = clip_expanded.reshape(B * n_buckets, 77, -1)

            # Create timesteps for all buckets: [B * n_buckets]
            t_expanded = t_buckets.unsqueeze(0).expand(B, -1).reshape(-1).to(DEVICE)

            # Random latents: [B * n_buckets, 4, 64, 64]
            latents = torch.randn(B * n_buckets, 4, 64, 64, device=DEVICE, dtype=DTYPE)

            # Single batched UNet forward pass
            _ = unet(latents, t_expanded * 1000, encoder_hidden_states=clip_expanded.to(DTYPE))

            # Reshape features back to [B, n_buckets, LUNE_DIM]
            features = mid_features[0].reshape(B, n_buckets, -1)
            all_features[start_idx:end_idx] = features.cpu().to(torch.float16)

    del unet
    torch.cuda.empty_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 LuneFeatureCache(all_features, t_buckets, DTYPE)


def load_or_extract_sol_features(cache_path: str, prompts: List[str], name: str,
                                  clip_tok, clip_enc, t_buckets: torch.Tensor,
                                  spatial_size: int = 8,
                                  batch_size: int = 32) -> Optional[SolFeatureCache]:
    """Load cached Sol features or generate geometric heuristics."""
    if not prompts or not ENABLE_SOL_DISTILLATION:
        return None

    if os.path.exists(cache_path):
        print(f"Loading cached {name} Sol features...")
        cached = torch.load(cache_path, map_location="cpu")
        cache = SolFeatureCache(
            cached["stats"], cached["spatial"], cached["t_buckets"], DTYPE
        )
        print(f"  ✓ Loaded {cache.stats.shape[0]} samples × {cache.n_buckets} timesteps")
        return cache

    print(f"Generating {name} Sol features ({len(prompts)} × {len(t_buckets)} timesteps)...")
    print(f"  Using geometric heuristics (no teacher needed)")

    n_prompts = len(prompts)
    n_buckets = len(t_buckets)
    
    # Vectorized generation - no loops needed
    # Stats: [n_buckets, 4] then broadcast to [n_prompts, n_buckets, 4]
    t_vals = t_buckets.float()  # [n_buckets]
    
    locality = 1 - t_vals  # [n_buckets]
    entropy = t_vals
    clustering = 0.5 - 0.3 * (t_vals - 0.5).abs()
    sparsity = 1 - t_vals
    
    stats_per_t = torch.stack([locality, entropy, clustering, sparsity], dim=-1)  # [n_buckets, 4]
    all_stats = stats_per_t.unsqueeze(0).expand(n_prompts, -1, -1).to(torch.float16)  # [n_prompts, n_buckets, 4]
    
    # Spatial: [n_buckets, spatial_size, spatial_size] then broadcast
    y, x = torch.meshgrid(
        torch.linspace(-1, 1, spatial_size),
        torch.linspace(-1, 1, spatial_size),
        indexing='ij'
    )
    center_dist = torch.sqrt(x**2 + y**2)  # [spatial_size, spatial_size]
    
    # Vectorized across timesteps: [n_buckets, spatial_size, spatial_size]
    t_weight = (1 - t_vals).view(-1, 1, 1)  # [n_buckets, 1, 1]
    center_bias = 1 - center_dist.unsqueeze(0) * t_weight  # [n_buckets, spatial_size, spatial_size]
    center_bias = center_bias / center_bias.sum(dim=[-2, -1], keepdim=True)  # Normalize per timestep
    
    all_spatial = center_bias.unsqueeze(0).expand(n_prompts, -1, -1, -1).to(torch.float16)  # [n_prompts, n_buckets, 8, 8]

    torch.save({
        "stats": all_stats,
        "spatial": all_spatial,
        "t_buckets": t_buckets
    }, cache_path)
    print(f"  ✓ Cached to {cache_path}")

    return SolFeatureCache(all_stats, all_spatial, 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 sync_from_model(self, model, pattern=None):
        if hasattr(model, '_orig_mod'):
            model_state = model._orig_mod.state_dict()
        else:
            model_state = model.state_dict()
        
        synced = 0
        for k, v in model_state.items():
            if pattern is None or pattern in k:
                if k in self.shadow:
                    self.shadow[k] = v.clone().to(self.shadow[k].device)
                    synced += 1
        
        print(f"  ✓ Synced EMA: {synced} weights" + (f" matching '{pattern}'" if pattern else ""))

    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, model=None):
        device = next(iter(self.shadow.values())).device if self.shadow else 'cuda'

        loaded = 0
        skipped_old = 0
        initialized_from_model = 0

        for k, v in shadow_state.items():
            if k in self.shadow:
                self.shadow[k] = v.clone().to(device)
                loaded += 1
            else:
                skipped_old += 1

        if model is not None:
            if hasattr(model, '_orig_mod'):
                model_state = model._orig_mod.state_dict()
            else:
                model_state = model.state_dict()
            
            for k in self.shadow:
                if k not in shadow_state and k in model_state:
                    self.shadow[k] = model_state[k].clone().to(device)
                    initialized_from_model += 1

        print(f"  ✓ Restored EMA: {loaded} loaded, {skipped_old} deprecated, {initialized_from_model} new (from model)")


# ============================================================================
# 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()
@torch.no_grad()
def encode_prompts_batched(prompts: List[str], batch_size: int = 128) -> Tuple[torch.Tensor, torch.Tensor]:
    """Batch encode prompts with T5 and CLIP."""
    all_t5 = []
    all_clip = []
    for i in tqdm(range(0, len(prompts), batch_size), desc="Encoding prompts", 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/6] 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/6] Portrait dataset DISABLED")

schnell_ds = None
schnell_prompts = []

if ENABLE_SCHNELL:
    print(f"\n[2/6] 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/6] Schnell dataset DISABLED")

sportfashion_ds = None
sportfashion_prompts = []
sportfashion_latents = None
sportfashion_masks = None

if ENABLE_SPORTFASHION:
    print(f"\n[3/6] 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")
    
    # Precache latents and masks
    sportfashion_latent_cache = os.path.join(LATENT_CACHE_DIR, f"sportfashion_latents_{len(sportfashion_ds)}.pt")
    sportfashion_mask_cache = os.path.join(LATENT_CACHE_DIR, f"sportfashion_masks_{len(sportfashion_ds)}.pt")
    
    if os.path.exists(sportfashion_latent_cache):
        print(f"  Loading cached SportFashion latents...")
        sportfashion_latents = torch.load(sportfashion_latent_cache)
        print(f"  ✓ Loaded {len(sportfashion_latents)} latents")
        if os.path.exists(sportfashion_mask_cache):
            sportfashion_masks = torch.load(sportfashion_mask_cache)
            print(f"  ✓ Loaded {len(sportfashion_masks)} masks")
    else:
        print(f"  Encoding SportFashion images to latents (one-time)...")
        VAE_BATCH_SIZE = 64  # A100 can handle large batches
        sportfashion_latents = []
        sportfashion_masks = []
        with torch.no_grad():
            for start_idx in tqdm(range(0, len(sportfashion_ds), VAE_BATCH_SIZE), desc="Encoding latents"):
                end_idx = min(start_idx + VAE_BATCH_SIZE, len(sportfashion_ds))
                batch_images = []
                batch_masks = []
                for i in range(start_idx, end_idx):
                    image = sportfashion_ds[i]["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)
                    batch_images.append(img_tensor)
                    # Create mask
                    pixel_mask = create_product_mask(image)
                    mask = downsample_mask_to_latent(pixel_mask, 64, 64)
                    batch_masks.append(mask)
                batch_tensor = torch.stack(batch_images)
                batch_tensor = (batch_tensor * 2.0 - 1.0).to(DEVICE, dtype=DTYPE)
                latents = vae.encode(batch_tensor).latent_dist.sample()
                latents = latents * VAE_SCALE
                sportfashion_latents.append(latents.cpu())
                sportfashion_masks.extend(batch_masks)
        sportfashion_latents = torch.cat(sportfashion_latents, dim=0)
        sportfashion_masks = torch.stack(sportfashion_masks)
        torch.save(sportfashion_latents, sportfashion_latent_cache)
        torch.save(sportfashion_masks, sportfashion_mask_cache)
        print(f"  ✓ Cached to {sportfashion_latent_cache}")
else:
    print("\n[3/6] SportFashion dataset DISABLED")

synthmocap_ds = None
synthmocap_prompts = []
synthmocap_latents = None
synthmocap_masks = None

if ENABLE_SYNTHMOCAP:
    print(f"\n[4/6] 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")
    
    # Precache latents and masks
    synthmocap_latent_cache = os.path.join(LATENT_CACHE_DIR, f"synthmocap_latents_{len(synthmocap_ds)}.pt")
    synthmocap_mask_cache = os.path.join(LATENT_CACHE_DIR, f"synthmocap_masks_{len(synthmocap_ds)}.pt")
    
    if os.path.exists(synthmocap_latent_cache):
        print(f"  Loading cached SynthMoCap latents...")
        synthmocap_latents = torch.load(synthmocap_latent_cache)
        print(f"  ✓ Loaded {len(synthmocap_latents)} latents")
        if os.path.exists(synthmocap_mask_cache):
            synthmocap_masks = torch.load(synthmocap_mask_cache)
            print(f"  ✓ Loaded {len(synthmocap_masks)} masks")
    else:
        print(f"  Encoding SynthMoCap images to latents (one-time)...")
        VAE_BATCH_SIZE = 64  # A100 can handle large batches
        synthmocap_latents = []
        synthmocap_masks = []
        with torch.no_grad():
            for start_idx in tqdm(range(0, len(synthmocap_ds), VAE_BATCH_SIZE), desc="Encoding latents"):
                end_idx = min(start_idx + VAE_BATCH_SIZE, len(synthmocap_ds))
                batch_images = []
                batch_masks = []
                for i in range(start_idx, end_idx):
                    image = synthmocap_ds[i]["image"]
                    conditioning = synthmocap_ds[i]["conditioning_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)
                    batch_images.append(img_tensor)
                    # Create mask from conditioning image
                    pixel_mask = create_smpl_mask(conditioning)
                    mask = downsample_mask_to_latent(pixel_mask, 64, 64)
                    batch_masks.append(mask)
                batch_tensor = torch.stack(batch_images)
                batch_tensor = (batch_tensor * 2.0 - 1.0).to(DEVICE, dtype=DTYPE)
                latents = vae.encode(batch_tensor).latent_dist.sample()
                latents = latents * VAE_SCALE
                synthmocap_latents.append(latents.cpu())
                synthmocap_masks.extend(batch_masks)
        synthmocap_latents = torch.cat(synthmocap_latents, dim=0)
        synthmocap_masks = torch.stack(synthmocap_masks)
        torch.save(synthmocap_latents, synthmocap_latent_cache)
        torch.save(synthmocap_masks, synthmocap_mask_cache)
        print(f"  ✓ Cached to {synthmocap_latent_cache}")
else:
    print("\n[4/6] SynthMoCap dataset DISABLED")

# ============================================================================
# IMAGENET DATASET WITH SMART PROMPT FILTERING
# ============================================================================
imagenet_ds = None
imagenet_prompts = []


def build_imagenet_prompt(item):
    semantic_class = item.get("semantic_class", "object")
    semantic_subclass = item.get("semantic_subclass", "")
    label = item.get("label", "").replace("_", " ")
    base_prompt = item.get("prompt", "")
    synset_id = item.get("synset_id", "")
    
    pred_confidence = item.get("pred_confidence", 0.0)
    top1_correct = item.get("top1_correct", False)
    top5_correct = item.get("top5_correct", False)
    
    confident_but_wrong = (
        pred_confidence >= IMAGENET_CONFIDENCE_THRESHOLD and 
        not top1_correct and 
        not top5_correct
    )
    
    if confident_but_wrong:
        parts = ["subject", semantic_class]
        if semantic_subclass:
            parts.append(semantic_subclass)
        parts.append(base_prompt)
        parts.append(synset_id)
        parts.append("imagenet")
    else:
        parts = ["subject", semantic_class]
        if semantic_subclass:
            parts.append(semantic_subclass)
        if label:
            parts.append(label)
        parts.append(base_prompt)
        parts.append(synset_id)
        parts.append("imagenet")
    
    return ", ".join(p for p in parts if p)


if ENABLE_IMAGENET:
    print(f"\n[5/6] Loading Synthetic ImageNet from {IMAGENET_REPO}...")
    imagenet_ds = load_dataset(IMAGENET_REPO, IMAGENET_SUBSET, split="train")
    print(f"  Raw samples: {len(imagenet_ds)}")
    
    # Use columnar access - MUCH faster than row iteration
    print(f"  Building prompts...")
    semantic_classes = imagenet_ds["semantic_class"]
    semantic_subclasses = imagenet_ds.get("semantic_subclass", [""] * len(imagenet_ds)) if "semantic_subclass" in imagenet_ds.features else [""] * len(imagenet_ds)
    labels = imagenet_ds["label"]
    base_prompts = imagenet_ds["prompt"]
    synset_ids = imagenet_ds["synset_id"]
    pred_confidences = imagenet_ds.get("pred_confidence", [0.0] * len(imagenet_ds)) if "pred_confidence" in imagenet_ds.features else [0.0] * len(imagenet_ds)
    top1_corrects = imagenet_ds.get("top1_correct", [False] * len(imagenet_ds)) if "top1_correct" in imagenet_ds.features else [False] * len(imagenet_ds)
    top5_corrects = imagenet_ds.get("top5_correct", [False] * len(imagenet_ds)) if "top5_correct" in imagenet_ds.features else [False] * len(imagenet_ds)
    
    # Handle case where columns might not exist
    if not isinstance(semantic_subclasses, list):
        semantic_subclasses = list(semantic_subclasses) if semantic_subclasses else [""] * len(imagenet_ds)
    if not isinstance(pred_confidences, list):
        pred_confidences = list(pred_confidences) if pred_confidences else [0.0] * len(imagenet_ds)
    if not isinstance(top1_corrects, list):
        top1_corrects = list(top1_corrects) if top1_corrects else [False] * len(imagenet_ds)
    if not isinstance(top5_corrects, list):
        top5_corrects = list(top5_corrects) if top5_corrects else [False] * len(imagenet_ds)
    
    confident_wrong = 0
    for i in range(len(imagenet_ds)):
        semantic_class = semantic_classes[i] if semantic_classes[i] else "object"
        semantic_subclass = semantic_subclasses[i] if i < len(semantic_subclasses) else ""
        label = labels[i].replace("_", " ") if labels[i] else ""
        base_prompt = base_prompts[i] if base_prompts[i] else ""
        synset_id = synset_ids[i] if synset_ids[i] else ""
        pred_confidence = pred_confidences[i] if i < len(pred_confidences) else 0.0
        top1_correct = top1_corrects[i] if i < len(top1_corrects) else False
        top5_correct = top5_corrects[i] if i < len(top5_corrects) else False
        
        confident_but_wrong = (
            pred_confidence >= IMAGENET_CONFIDENCE_THRESHOLD and 
            not top1_correct and 
            not top5_correct
        )
        
        if confident_but_wrong:
            parts = ["subject", semantic_class]
            if semantic_subclass:
                parts.append(semantic_subclass)
            parts.append(base_prompt)
            parts.append(synset_id)
            parts.append("imagenet")
            confident_wrong += 1
        else:
            parts = ["subject", semantic_class]
            if semantic_subclass:
                parts.append(semantic_subclass)
            if label:
                parts.append(label)
            parts.append(base_prompt)
            parts.append(synset_id)
            parts.append("imagenet")
        
        imagenet_prompts.append(", ".join(p for p in parts if p))
    
    print(f"✓ ImageNet: {len(imagenet_ds)} samples")
    print(f"  Confident mispredictions (label removed): {confident_wrong}")
    
    imagenet_latent_cache = os.path.join(LATENT_CACHE_DIR, f"imagenet_latents_{len(imagenet_ds)}.pt")
    if os.path.exists(imagenet_latent_cache):
        print(f"  Loading cached ImageNet latents...")
        imagenet_latents = torch.load(imagenet_latent_cache)
        print(f"  ✓ Loaded {len(imagenet_latents)} latents")
    else:
        print(f"  Encoding ImageNet images to latents (one-time)...")
        VAE_BATCH_SIZE = 64  # A100 can handle large batches
        imagenet_latents = []
        with torch.no_grad():
            for start_idx in tqdm(range(0, len(imagenet_ds), VAE_BATCH_SIZE), desc="Encoding latents"):
                end_idx = min(start_idx + VAE_BATCH_SIZE, len(imagenet_ds))
                batch_images = []
                for i in range(start_idx, end_idx):
                    image = imagenet_ds[i]["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)
                    batch_images.append(img_tensor)
                batch_tensor = torch.stack(batch_images)
                batch_tensor = (batch_tensor * 2.0 - 1.0).to(DEVICE, dtype=DTYPE)
                latents = vae.encode(batch_tensor).latent_dist.sample()
                latents = latents * VAE_SCALE
                imagenet_latents.append(latents.cpu())
        imagenet_latents = torch.cat(imagenet_latents, dim=0)
        torch.save(imagenet_latents, imagenet_latent_cache)
        print(f"  ✓ Cached to {imagenet_latent_cache}")
else:
    print("\n[5/6] ImageNet dataset DISABLED")
    imagenet_latents = None

# ============================================================================
# OBJECT RELATIONS DATASET WITH SUBJECT PREFIX
# ============================================================================
object_relations_ds = None
object_relations_prompts = []
object_relations_latents = None


def build_object_relations_prompt(item):
    prompt = item.get("prompt", "")
    if random.random() < 0.5:
        return f"subject, object, {prompt}"
    else:
        return f"subject, {prompt}"


if ENABLE_OBJECT_RELATIONS:
    print(f"\n[6/6] Loading Object Relations from {OBJECT_RELATIONS_REPO}...")
    object_relations_ds = load_dataset(OBJECT_RELATIONS_REPO, "schnell_512_1", split="train")
    print(f"  Raw samples: {len(object_relations_ds)}")
    
    # Use columnar access - MUCH faster than row iteration
    print(f"  Building prompts...")
    all_prompts = object_relations_ds["prompt"]  # Get entire column at once
    
    random.seed(42)
    object_relations_prompts = []
    for prompt in all_prompts:
        if random.random() < 0.5:
            object_relations_prompts.append(f"subject, object, {prompt}")
        else:
            object_relations_prompts.append(f"subject, {prompt}")
    random.seed()
    
    subject_object_count = sum(1 for p in object_relations_prompts if p.startswith("subject, object,"))
    subject_only_count = len(object_relations_prompts) - subject_object_count
    print(f"✓ Object Relations: {len(object_relations_ds)} samples")
    print(f"  'subject, object, ...' prefix: {subject_object_count}")
    print(f"  'subject, ...' prefix: {subject_only_count}")
    
    object_relations_latent_cache = os.path.join(LATENT_CACHE_DIR, f"object_relations_latents_{len(object_relations_ds)}.pt")
    if os.path.exists(object_relations_latent_cache):
        print(f"  Loading cached Object Relations latents...")
        object_relations_latents = torch.load(object_relations_latent_cache)
        print(f"  ✓ Loaded {len(object_relations_latents)} latents")
    else:
        print(f"  Encoding Object Relations images to latents (one-time)...")
        VAE_BATCH_SIZE = 64  # A100 can handle large batches
        object_relations_latents = []
        with torch.no_grad():
            for start_idx in tqdm(range(0, len(object_relations_ds), VAE_BATCH_SIZE), desc="Encoding latents"):
                end_idx = min(start_idx + VAE_BATCH_SIZE, len(object_relations_ds))
                batch_images = []
                for i in range(start_idx, end_idx):
                    image = object_relations_ds[i]["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)
                    batch_images.append(img_tensor)
                batch_tensor = torch.stack(batch_images)
                batch_tensor = (batch_tensor * 2.0 - 1.0).to(DEVICE, dtype=DTYPE)
                latents = vae.encode(batch_tensor).latent_dist.sample()
                latents = latents * VAE_SCALE
                object_relations_latents.append(latents.cpu())
        object_relations_latents = torch.cat(object_relations_latents, dim=0)
        torch.save(object_relations_latents, object_relations_latent_cache)
        print(f"  ✓ Cached to {object_relations_latent_cache}")
else:
    print("\n[6/6] Object Relations dataset DISABLED")

# ============================================================================
# ENCODE ALL PROMPTS
# ============================================================================
total_samples = len(portrait_prompts) + len(schnell_prompts) + len(sportfashion_prompts) + len(synthmocap_prompts) + len(imagenet_prompts) + len(object_relations_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


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

imagenet_t5, imagenet_clip = None, None
if imagenet_prompts:
    imagenet_enc_cache = os.path.join(ENCODING_CACHE_DIR, f"imagenet_encodings_{len(imagenet_prompts)}.pt")
    imagenet_t5, imagenet_clip = load_or_encode(imagenet_enc_cache, imagenet_prompts, "imagenet")

object_relations_t5, object_relations_clip = None, None
if object_relations_prompts:
    object_relations_enc_cache = os.path.join(ENCODING_CACHE_DIR, f"object_relations_encodings_{len(object_relations_prompts)}.pt")
    object_relations_t5, object_relations_clip = load_or_encode(object_relations_enc_cache, object_relations_prompts, "object_relations")



# ============================================================================
# EXTRACT/LOAD LUNE AND SOL FEATURES (precached)
# ============================================================================
print("\n" + "=" * 60)
print("Expert Feature Caching (Lune + Sol)")
print("=" * 60)

# Lune caches
schnell_lune_cache = None
portrait_lune_cache = None
sportfashion_lune_cache = None
synthmocap_lune_cache = None
imagenet_lune_cache = None
object_relations_lune_cache = None

# Sol caches
schnell_sol_cache = None
portrait_sol_cache = None
sportfashion_sol_cache = None
synthmocap_sol_cache = None
imagenet_sol_cache = None
object_relations_sol_cache = None

if schnell_prompts:
    if ENABLE_LUNE_DISTILLATION:
        schnell_lune_path = os.path.join(ENCODING_CACHE_DIR, f"schnell_lune_{len(schnell_prompts)}.pt")
        schnell_lune_cache = load_or_extract_lune_features(
            schnell_lune_path, schnell_prompts, "schnell",
            clip_tok, clip_enc, EXPERT_T_BUCKETS
        )
    if ENABLE_SOL_DISTILLATION:
        schnell_sol_path = os.path.join(ENCODING_CACHE_DIR, f"schnell_sol_{len(schnell_prompts)}.pt")
        schnell_sol_cache = load_or_extract_sol_features(
            schnell_sol_path, schnell_prompts, "schnell",
            clip_tok, clip_enc, EXPERT_T_BUCKETS, SOL_SPATIAL_SIZE
        )

if portrait_prompts:
    if ENABLE_LUNE_DISTILLATION:
        portrait_lune_path = os.path.join(ENCODING_CACHE_DIR, f"portrait_lune_{len(portrait_prompts)}.pt")
        portrait_lune_cache = load_or_extract_lune_features(
            portrait_lune_path, portrait_prompts, "portrait",
            clip_tok, clip_enc, EXPERT_T_BUCKETS
        )
    if ENABLE_SOL_DISTILLATION:
        portrait_sol_path = os.path.join(ENCODING_CACHE_DIR, f"portrait_sol_{len(portrait_prompts)}.pt")
        portrait_sol_cache = load_or_extract_sol_features(
            portrait_sol_path, portrait_prompts, "portrait",
            clip_tok, clip_enc, EXPERT_T_BUCKETS, SOL_SPATIAL_SIZE
        )

if sportfashion_prompts:
    if ENABLE_LUNE_DISTILLATION:
        sportfashion_lune_path = os.path.join(ENCODING_CACHE_DIR, f"sportfashion_lune_{len(sportfashion_prompts)}.pt")
        sportfashion_lune_cache = load_or_extract_lune_features(
            sportfashion_lune_path, sportfashion_prompts, "sportfashion",
            clip_tok, clip_enc, EXPERT_T_BUCKETS
        )
    if ENABLE_SOL_DISTILLATION:
        sportfashion_sol_path = os.path.join(ENCODING_CACHE_DIR, f"sportfashion_sol_{len(sportfashion_prompts)}.pt")
        sportfashion_sol_cache = load_or_extract_sol_features(
            sportfashion_sol_path, sportfashion_prompts, "sportfashion",
            clip_tok, clip_enc, EXPERT_T_BUCKETS, SOL_SPATIAL_SIZE
        )

if synthmocap_prompts:
    if ENABLE_LUNE_DISTILLATION:
        synthmocap_lune_path = os.path.join(ENCODING_CACHE_DIR, f"synthmocap_lune_{len(synthmocap_prompts)}.pt")
        synthmocap_lune_cache = load_or_extract_lune_features(
            synthmocap_lune_path, synthmocap_prompts, "synthmocap",
            clip_tok, clip_enc, EXPERT_T_BUCKETS
        )
    if ENABLE_SOL_DISTILLATION:
        synthmocap_sol_path = os.path.join(ENCODING_CACHE_DIR, f"synthmocap_sol_{len(synthmocap_prompts)}.pt")
        synthmocap_sol_cache = load_or_extract_sol_features(
            synthmocap_sol_path, synthmocap_prompts, "synthmocap",
            clip_tok, clip_enc, EXPERT_T_BUCKETS, SOL_SPATIAL_SIZE
        )

if imagenet_prompts:
    if ENABLE_LUNE_DISTILLATION:
        imagenet_lune_path = os.path.join(ENCODING_CACHE_DIR, f"imagenet_lune_{len(imagenet_prompts)}.pt")
        imagenet_lune_cache = load_or_extract_lune_features(
            imagenet_lune_path, imagenet_prompts, "imagenet",
            clip_tok, clip_enc, EXPERT_T_BUCKETS
        )
    if ENABLE_SOL_DISTILLATION:
        imagenet_sol_path = os.path.join(ENCODING_CACHE_DIR, f"imagenet_sol_{len(imagenet_prompts)}.pt")
        imagenet_sol_cache = load_or_extract_sol_features(
            imagenet_sol_path, imagenet_prompts, "imagenet",
            clip_tok, clip_enc, EXPERT_T_BUCKETS, SOL_SPATIAL_SIZE
        )

if object_relations_prompts:
    if ENABLE_LUNE_DISTILLATION:
        object_relations_lune_path = os.path.join(ENCODING_CACHE_DIR, f"object_relations_lune_{len(object_relations_prompts)}.pt")
        object_relations_lune_cache = load_or_extract_lune_features(
            object_relations_lune_path, object_relations_prompts, "object_relations",
            clip_tok, clip_enc, EXPERT_T_BUCKETS
        )
    if ENABLE_SOL_DISTILLATION:
        object_relations_sol_path = os.path.join(ENCODING_CACHE_DIR, f"object_relations_sol_{len(object_relations_prompts)}.pt")
        object_relations_sol_cache = load_or_extract_sol_features(
            object_relations_sol_path, object_relations_prompts, "object_relations",
            clip_tok, clip_enc, EXPERT_T_BUCKETS, SOL_SPATIAL_SIZE
        )


# ============================================================================
# COMBINED DATASET CLASS
# ============================================================================
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_latents, sportfashion_masks, sportfashion_t5, sportfashion_clip,
            synthmocap_ds, synthmocap_latents, synthmocap_masks, synthmocap_t5, synthmocap_clip,
            imagenet_ds, imagenet_latents, imagenet_t5, imagenet_clip,
            object_relations_ds, object_relations_latents, object_relations_t5, object_relations_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_latents = sportfashion_latents
        self.sportfashion_masks = sportfashion_masks
        self.sportfashion_t5 = sportfashion_t5
        self.sportfashion_clip = sportfashion_clip

        self.synthmocap_ds = synthmocap_ds
        self.synthmocap_latents = synthmocap_latents
        self.synthmocap_masks = synthmocap_masks
        self.synthmocap_t5 = synthmocap_t5
        self.synthmocap_clip = synthmocap_clip

        self.imagenet_ds = imagenet_ds
        self.imagenet_latents = imagenet_latents
        self.imagenet_t5 = imagenet_t5
        self.imagenet_clip = imagenet_clip

        self.object_relations_ds = object_relations_ds
        self.object_relations_latents = object_relations_latents
        self.object_relations_t5 = object_relations_t5
        self.object_relations_clip = object_relations_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_latents) if sportfashion_latents is not None else 0
        self.n_synthmocap = len(synthmocap_latents) if synthmocap_latents is not None else 0
        self.n_imagenet = len(imagenet_latents) if imagenet_latents is not None else 0
        self.n_object_relations = len(object_relations_latents) if object_relations_latents is not None else 0

        self.c1 = self.n_portrait
        self.c2 = self.c1 + self.n_schnell
        self.c3 = self.c2 + self.n_sportfashion
        self.c4 = self.c3 + self.n_synthmocap
        self.c5 = self.c4 + self.n_imagenet
        self.total = self.c5 + self.n_object_relations

    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)

    def __getitem__(self, idx):
        mask = None

        if idx < self.c1:
            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:
            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:
            local_idx = idx - self.c2
            latent = self.sportfashion_latents[local_idx].to(self.dtype)
            t5 = self.sportfashion_t5[local_idx]
            clip = self.sportfashion_clip[local_idx]
            dataset_id = 2
            if self.compute_masks and self.sportfashion_masks is not None:
                mask = self.sportfashion_masks[local_idx].to(self.dtype)

        elif idx < self.c4:
            local_idx = idx - self.c3
            latent = self.synthmocap_latents[local_idx].to(self.dtype)
            t5 = self.synthmocap_t5[local_idx]
            clip = self.synthmocap_clip[local_idx]
            dataset_id = 3
            if self.compute_masks and self.synthmocap_masks is not None:
                mask = self.synthmocap_masks[local_idx].to(self.dtype)

        elif idx < self.c5:
            local_idx = idx - self.c4
            latent = self.imagenet_latents[local_idx].to(self.dtype)
            t5 = self.imagenet_t5[local_idx]
            clip = self.imagenet_clip[local_idx]
            dataset_id = 4

        else:
            local_idx = idx - self.c5
            latent = self.object_relations_latents[local_idx].to(self.dtype)
            t5 = self.object_relations_t5[local_idx]
            clip = self.object_relations_clip[local_idx]
            dataset_id = 5

        result = {
            "latent": latent,
            "t5_embed": t5.to(self.dtype),
            "clip_pooled": clip.to(self.dtype),
            "sample_idx": idx,
            "local_idx": local_idx,
            "dataset_id": dataset_id,
        }

        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, dual experts)
# ============================================================================
def get_lune_features_for_batch(
        local_indices: torch.Tensor,
        dataset_ids: torch.Tensor,
        timesteps: torch.Tensor,
) -> Optional[torch.Tensor]:
    """Get Lune features from the appropriate cache for each sample."""
    caches = [
        portrait_lune_cache, schnell_lune_cache, sportfashion_lune_cache,
        synthmocap_lune_cache, imagenet_lune_cache, object_relations_lune_cache
    ]

    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, LUNE_DIM, device=device, dtype=DTYPE)

    for ds_id, cache in enumerate(caches):
        if cache is None:
            continue
        mask = dataset_ids == ds_id
        if not mask.any():
            continue
        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


def get_sol_features_for_batch(
        local_indices: torch.Tensor,
        dataset_ids: torch.Tensor,
        timesteps: torch.Tensor,
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
    """Get Sol features (stats + spatial) from the appropriate cache."""
    caches = [
        portrait_sol_cache, schnell_sol_cache, sportfashion_sol_cache,
        synthmocap_sol_cache, imagenet_sol_cache, object_relations_sol_cache
    ]

    if not any(c is not None for c in caches):
        return None, None

    B = local_indices.shape[0]
    device = timesteps.device
    stats = torch.zeros(B, 3, device=device, dtype=DTYPE)  # 3 stats: locality, entropy, clustering
    spatial = torch.zeros(B, SOL_SPATIAL_SIZE, SOL_SPATIAL_SIZE, device=device, dtype=DTYPE)

    for ds_id, cache in enumerate(caches):
        if cache is None:
            continue
        mask = dataset_ids == ds_id
        if not mask.any():
            continue
        ds_local_indices = local_indices[mask]
        ds_timesteps = timesteps[mask]
        ds_stats, ds_spatial = cache.get_features(ds_local_indices, ds_timesteps)
        stats[mask] = ds_stats[:, :3]  # Drop redundant sparsity (was copy of locality)
        spatial[mask] = ds_spatial

    return stats, spatial


# ============================================================================
# LOSS FUNCTIONS
# ============================================================================
def huber_loss(pred, target, delta=0.1):
    """Huber loss - L2 for small errors, L1 for large."""
    diff = pred - target
    abs_diff = diff.abs()
    quadratic = torch.clamp(abs_diff, max=delta)
    linear = abs_diff - quadratic
    return 0.5 * quadratic ** 2 + delta * linear


def compute_main_loss(pred, target, mask=None, spatial_weights=None, 
                      fg_weight=2.0, bg_weight=0.5, snr_weights=None):
    """Compute main prediction loss with optional spatial weighting."""
    B, N, C = pred.shape
    
    if USE_HUBER_LOSS:
        loss_per_elem = huber_loss(pred, target, HUBER_DELTA)
    else:
        loss_per_elem = (pred - target) ** 2
    
    # Apply spatial weights from Sol if enabled
    if spatial_weights is not None and USE_SPATIAL_WEIGHTING:
        H = W = int(math.sqrt(N))
        # Upsample spatial weights from 8x8 to HxW
        spatial_upsampled = F.interpolate(
            spatial_weights.unsqueeze(1),  # [B, 1, 8, 8]
            size=(H, W),
            mode='bilinear',
            align_corners=False
        ).squeeze(1)  # [B, H, W]
        # Normalize so mean = 1
        spatial_upsampled = spatial_upsampled / (spatial_upsampled.mean(dim=[1, 2], keepdim=True) + 1e-6)
        spatial_flat = spatial_upsampled.view(B, N, 1)
        loss_per_elem = loss_per_elem * spatial_flat
    
    # Apply foreground/background mask
    if mask is not None:
        H = W = int(math.sqrt(N))
        mask_flat = mask.view(B, H * W, 1).to(pred.device)
        weights = mask_flat * fg_weight + (1 - mask_flat) * bg_weight
        loss_per_elem = loss_per_elem * weights
    
    loss_per_sample = loss_per_elem.mean(dim=[1, 2])
    
    if snr_weights is not None:
        loss_per_sample = loss_per_sample * snr_weights
    
    return loss_per_sample.mean()


def compute_lune_loss(pred, target, mode="cosine"):
    """Compute Lune distillation loss."""
    if mode == "cosine":
        # Cosine similarity loss (1 - cos_sim)
        pred_norm = F.normalize(pred, dim=-1)
        target_norm = F.normalize(target, dim=-1)
        return (1 - (pred_norm * target_norm).sum(dim=-1)).mean()
    elif mode == "huber":
        return huber_loss(pred, target, HUBER_DELTA).mean()
    elif mode == "soft":
        # Soft L2 with temperature
        return F.mse_loss(pred / 10.0, target / 10.0)
    else:  # hard
        return F.mse_loss(pred, target)


def compute_sol_loss(pred_stats, pred_spatial, target_stats, target_spatial):
    """Compute Sol distillation loss (stats + spatial)."""
    stats_loss = F.mse_loss(pred_stats, target_stats)
    spatial_loss = F.mse_loss(pred_spatial, target_spatial)
    return stats_loss + spatial_loss


# ============================================================================
# WEIGHT SCHEDULES
# ============================================================================
def get_lune_weight(step):
    if step < LUNE_WARMUP_STEPS:
        return LUNE_LOSS_WEIGHT * (step / LUNE_WARMUP_STEPS)
    return LUNE_LOSS_WEIGHT


def get_sol_weight(step):
    if step < SOL_WARMUP_STEPS:
        return SOL_LOSS_WEIGHT * (step / SOL_WARMUP_STEPS)
    return SOL_LOSS_WEIGHT


# ============================================================================
# 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_latents, sportfashion_masks, sportfashion_t5, sportfashion_clip,
    synthmocap_ds, synthmocap_latents, synthmocap_masks, synthmocap_t5, synthmocap_clip,
    imagenet_ds, imagenet_latents, imagenet_t5, imagenet_clip,
    object_relations_ds, object_relations_latents, object_relations_t5, object_relations_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"  - ImageNet:          {combined_ds.n_imagenet:,}")
print(f"  - Object Relations:  {combined_ds.n_object_relations:,}")
print(f"  - Lune distillation: {ENABLE_LUNE_DISTILLATION}")
print(f"  - Sol distillation:  {ENABLE_SOL_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=5.0, H=64, W=64,
                     use_ema=True, seed=None,
                     negative_prompt="blurry, distorted, low quality"):
    """Generate samples during training with proper CFG support."""
    was_training = model.training
    model.eval()

    if seed is not None:
        torch.manual_seed(seed)

    model_ref = model._orig_mod if hasattr(model, '_orig_mod') else model

    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_cond = torch.stack(t5_list).to(DTYPE)
    clip_cond = torch.stack(clip_list).to(DTYPE)

    if guidance_scale > 1.0:
        t5_uncond, clip_uncond = encode_prompt(negative_prompt)
        t5_uncond = t5_uncond.expand(B, -1, -1)
        clip_uncond = clip_uncond.expand(B, -1)
    else:
        t5_uncond, clip_uncond = None, None

    x = torch.randn(B, H * W, C, device=DEVICE, dtype=DTYPE)
    img_ids = model_ref.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):
            v_cond = model_ref(
                hidden_states=x,
                encoder_hidden_states=t5_cond,
                pooled_projections=clip_cond,
                timestep=t_batch,
                img_ids=img_ids,
            )
            if isinstance(v_cond, tuple):
                v_cond = v_cond[0]

            if guidance_scale > 1.0 and t5_uncond is not None:
                v_uncond = model_ref(
                    hidden_states=x,
                    encoder_hidden_states=t5_uncond,
                    pooled_projections=clip_uncond,
                    timestep=t_batch,
                    img_ids=img_ids,
                )
                if isinstance(v_uncond, tuple):
                    v_uncond = v_uncond[0]
                v = v_uncond + guidance_scale * (v_cond - v_uncond)
            else:
                v = v_cond

        x = x + v * 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 upload_logs():
    try:
        for root, dirs, files in os.walk(LOG_DIR):
            for f in files:
                if f.startswith("events.out.tfevents"):
                    local_path = os.path.join(root, f)
                    rel_path = os.path.relpath(local_path, LOG_DIR)
                    repo_path = f"logs/{rel_path}"
                    api.upload_file(
                        path_or_fileobj=local_path,
                        path_in_repo=repo_path,
                        repo_id=HF_REPO,
                    )
        print(f"  ✓ Uploaded logs to {HF_REPO}")
    except Exception as e:
        print(f"  ⚠ Log upload failed: {e}")



# ============================================================================
# WEIGHT UPGRADE LOADING (v3 -> v4.1)
# ============================================================================


def load_with_weight_upgrade(model, state_dict):
    """Load state dict with bidirectional remapping support.
    
    Handles:
    - v3 checkpoint (expert_predictor) -> v4 model (lune_predictor)  
    - v4 checkpoint (lune_predictor) -> model with (expert_predictor)
    """
    model_state = model.state_dict()
    
    # Detect which naming the MODEL uses
    model_has_expert = any('expert_predictor' in k for k in model_state.keys())
    model_has_lune = any('lune_predictor' in k for k in model_state.keys())
    
    # Detect which naming the CHECKPOINT uses
    ckpt_has_expert = any('expert_predictor' in k for k in state_dict.keys())
    ckpt_has_lune = any('lune_predictor' in k for k in state_dict.keys())
    
    # Build remap based on mismatch
    REMAP = {}
    if model_has_expert and ckpt_has_lune:
        # Checkpoint has lune_predictor, model expects expert_predictor
        print("  Remapping: lune_predictor -> expert_predictor")
        REMAP = {'lune_predictor.': 'expert_predictor.'}
    elif model_has_lune and ckpt_has_expert:
        # Checkpoint has expert_predictor, model expects lune_predictor
        print("  Remapping: expert_predictor -> lune_predictor")
        REMAP = {'expert_predictor.': 'lune_predictor.'}

    # New modules that may not exist in checkpoint
    NEW_WEIGHT_PATTERNS = [
        'expert_predictor.',
        'lune_predictor.',
        'sol_prior.',
        't5_vec_proj.',
        '.norm_q.weight',
        '.norm_k.weight',
        '.norm_added_q.weight',
        '.norm_added_k.weight',
    ]

    # Deprecated keys
    DEPRECATED_PATTERNS = [
        'guidance_in.',
        '.sin_basis',
    ]

    loaded_keys = []
    missing_keys = []
    unexpected_keys = []
    initialized_keys = []
    remapped_keys = []

    # First pass: remap checkpoint keys to match model
    remapped_state = {}
    for k, v in state_dict.items():
        new_k = k
        for old_pat, new_pat in REMAP.items():
            if old_pat in k:
                new_k = k.replace(old_pat, new_pat)
                remapped_keys.append(f"{k} -> {new_k}")
                break
        remapped_state[new_k] = v

    # Second pass: load matching weights
    for key, v in remapped_state.items():
        if key in model_state:
            if v.shape == model_state[key].shape:
                model_state[key] = v
                loaded_keys.append(key)
            else:
                print(f"  ⚠ Shape mismatch for {key}: checkpoint {v.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)

    # Third 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:
                initialized_keys.append(key)
            else:
                missing_keys.append(key)
                print(f"  ⚠ Missing key (not in checkpoint): {key}")

    model.load_state_dict(model_state, strict=False)

    # Report
    if remapped_keys:
        print(f"  ✓ Remapped v3->v4: {len(remapped_keys)} keys")
        for rk in remapped_keys[:5]:
            print(f"      {rk}")
        if len(remapped_keys) > 5:
            print(f"      ... and {len(remapped_keys) - 5} more")

    if initialized_keys:
        modules = set()
        for k in initialized_keys:
            parts = k.split('.')
            if len(parts) >= 2:
                modules.add(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)}")

    return missing_keys, unexpected_keys


def load_checkpoint(model, optimizer, scheduler, target):
    """Load checkpoint with weight upgrade support for v4.1."""
    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:"):
                path_or_name = target.split(":", 1)[1]
                
                # Check if it's a full path (contains /) or just a step name
                if "/" in path_or_name:
                    # Full path like checkpoint_runs/v4_init/lailah_401434_v4_init
                    weights_path = hf_hub_download(HF_REPO, f"{path_or_name}.safetensors")
                    try:
                        ema_weights_path = hf_hub_download(HF_REPO, f"{path_or_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)")
                    print(f"Downloaded {path_or_name} from hub")
                else:
                    # Simple step name like step_401434
                    step_name = path_or_name
                    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()}

        model_ref = model._orig_mod if hasattr(model, '_orig_mod') else model

        if ALLOW_WEIGHT_UPGRADE:
            missing, unexpected = load_with_weight_upgrade(model_ref, state_dict)
            if missing:
                print(f"  ⚠ {len(missing)} truly missing parameters")
        else:
            model_ref.load_state_dict(state_dict, strict=True)

        print(f"✓ Loaded model weights")

        # Load EMA weights
        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")
        print(f"Resuming from step {start_step}, epoch {start_epoch}")

    return start_step, start_epoch, ema_state



# ============================================================================
# CREATE MODEL (v4.1 with dual experts)
# ============================================================================
print("\nCreating TinyFlux v4.1 model with Lune + Sol...")

# Import model - expects model_v4.py to define TinyFluxConfig and TinyFlux
# If running as a notebook cell, ensure model_v4.py cell was run first
# If running as a script, uncomment the import below:
# from model_v4 import TinyFluxConfig, TinyFlux

config = TinyFluxConfig(
    hidden_size=512,
    num_attention_heads=4,
    attention_head_dim=128,
    num_double_layers=15,
    num_single_layers=25,
    
    # Lune expert (trajectory guidance)
    use_lune_expert=ENABLE_LUNE_DISTILLATION,
    lune_expert_dim=LUNE_DIM,
    lune_hidden_dim=LUNE_HIDDEN_DIM,
    lune_dropout=LUNE_DROPOUT,
    
    # Sol prior (structural guidance)
    use_sol_prior=ENABLE_SOL_DISTILLATION,
    sol_spatial_size=SOL_SPATIAL_SIZE,
    sol_hidden_dim=SOL_HIDDEN_DIM,
    sol_geometric_weight=SOL_GEOMETRIC_WEIGHT,
    
    # Other settings
    use_t5_vec=True,
    lune_distill_mode=LUNE_DISTILL_MODE,
    use_huber_loss=USE_HUBER_LOSS,
    huber_delta=HUBER_DELTA,
    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, 'lune_predictor') and model.lune_predictor is not None:
    lune_params = sum(p.numel() for p in model.lune_predictor.parameters())
    print(f"Lune predictor parameters: {lune_params:,}")

if hasattr(model, 'sol_prior') and model.sol_prior is not None:
    sol_params = sum(p.numel() for p in model.sol_prior.parameters())
    print(f"Sol prior parameters: {sol_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:
    # Remap v3 EMA keys to v4
    remapped_ema = {}
    for k, v in ema_state.items():
        #if k in V3_TO_V4_REMAP:
        #    remapped_ema[V3_TO_V4_REMAP[k]] = v
        #else:
            remapped_ema[k] = v
    ema.load_shadow(remapped_ema, model=model)
    
    # Sync new modules from model
    has_lune_in_ema = any('lune_predictor' in k for k in ema_state.keys())
    has_sol_in_ema = any('sol_prior' in k for k in ema_state.keys())
    
    if ENABLE_LUNE_DISTILLATION and not has_lune_in_ema:
        # Check if expert_predictor was in the v3 checkpoint (remapped to lune_predictor)
        has_expert_in_v3 = any('expert_predictor' in k for k in ema_state.keys())
        if not has_expert_in_v3:
            ema.sync_from_model(model, pattern='lune_predictor')
            print("  ✓ Force-synced lune_predictor (new weights)")
        else:
            print("  ✓ lune_predictor loaded from remapped v3 checkpoint")
    
    if ENABLE_SOL_DISTILLATION and not has_sol_in_ema:
        ema.sync_from_model(model, pattern='sol_prior')
        print("  ✓ Force-synced sol_prior (new weights)")
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",
]


# ============================================================================
# TRAINING LOOP
# ============================================================================
print(f"\n{'=' * 60}")
print(f"Training TinyFlux v4.1 with Dual Expert Distillation")
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"Lune distillation: {ENABLE_LUNE_DISTILLATION}")
if ENABLE_LUNE_DISTILLATION:
    print(f"  - Mode: {LUNE_DISTILL_MODE}")
    print(f"  - Weight: {LUNE_LOSS_WEIGHT} (warmup: {LUNE_WARMUP_STEPS} steps)")
print(f"Sol distillation: {ENABLE_SOL_DISTILLATION}")
if ENABLE_SOL_DISTILLATION:
    print(f"  - Weight: {SOL_LOSS_WEIGHT} (warmup: {SOL_WARMUP_STEPS} steps)")
print(f"Huber loss: {USE_HUBER_LOSS} (delta={HUBER_DELTA})")
print(f"Spatial weighting: {USE_SPATIAL_WEIGHTING}")
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_lune_loss = 0
    ep_sol_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
        lune_features = None
        sol_stats = None
        sol_spatial = None
        
        if ENABLE_LUNE_DISTILLATION:
            lune_features = get_lune_features_for_batch(local_indices, dataset_ids, t)
            if lune_features is not None and random.random() < LUNE_DROPOUT:
                lune_features = None
        
        if ENABLE_SOL_DISTILLATION:
            sol_stats, sol_spatial = get_sol_features_for_batch(local_indices, dataset_ids, t)

        with torch.autocast("cuda", dtype=DTYPE):
            result = model(
                hidden_states=x_t,
                encoder_hidden_states=t5,
                pooled_projections=clip,
                timestep=t,
                img_ids=img_ids,
                lune_features=lune_features,
                sol_stats=sol_stats,
                sol_spatial=sol_spatial,
                return_expert_pred=True,
            )
            
            if isinstance(result, tuple):
                v_pred, expert_info = result
            else:
                v_pred = result
                expert_info = {}

        # Compute losses
        snr_weights = min_snr_weight(t)

        # Main loss with optional spatial weighting from Sol
        spatial_weights = sol_spatial if USE_SPATIAL_WEIGHTING else None
        main_loss = compute_main_loss(
            v_pred, v_target,
            mask=masks if USE_MASKED_LOSS else None,
            spatial_weights=spatial_weights,
            fg_weight=FG_LOSS_WEIGHT,
            bg_weight=BG_LOSS_WEIGHT,
            snr_weights=snr_weights
        )

        # Lune distillation loss
        lune_loss = torch.tensor(0.0, device=DEVICE)
        if lune_features is not None and expert_info.get('lune') is not None:
            lune_loss = compute_lune_loss(
                expert_info['lune']['expert_pred'], lune_features, mode=LUNE_DISTILL_MODE
            )

        # Sol distillation loss
        sol_loss = torch.tensor(0.0, device=DEVICE)
        if sol_stats is not None and expert_info.get('sol') is not None:
          sol_loss = compute_sol_loss(
              expert_info['sol']['pred_stats'], expert_info['sol']['pred_spatial'],
              sol_stats, sol_spatial
          )
        # Total loss with warmup weights
        total_loss = main_loss
        total_loss = total_loss + get_lune_weight(step) * lune_loss
        total_loss = total_loss + get_sol_weight(step) * sol_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_LUNE_DISTILLATION:
                    writer.add_scalar("train/lune_loss", lune_loss.item(), step)
                    writer.add_scalar("train/lune_weight", get_lune_weight(step), step)
                if ENABLE_SOL_DISTILLATION:
                    writer.add_scalar("train/sol_loss", sol_loss.item(), step)
                    writer.add_scalar("train/sol_weight", get_sol_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=28,
                    guidance_scale=5.0,
                    use_ema=True,
                    negative_prompt="blurry, distorted, low quality, deformed",
                )
                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)
                if step % LOG_UPLOAD_EVERY == 0:
                    writer.flush()
                    upload_logs()

        ep_loss += total_loss.item()
        ep_main_loss += main_loss.item()
        ep_lune_loss += lune_loss.item()
        ep_sol_loss += sol_loss.item()
        ep_batches += 1

        pbar.set_postfix(
            loss=f"{total_loss.item():.4f}",
            main=f"{main_loss.item():.4f}",
            lune=f"{lune_loss.item():.4f}" if ENABLE_LUNE_DISTILLATION else "-",
            sol=f"{sol_loss.item():.4f}" if ENABLE_SOL_DISTILLATION else "-",
            step=step
        )

    avg = ep_loss / max(ep_batches, 1)
    avg_main = ep_main_loss / max(ep_batches, 1)
    avg_lune = ep_lune_loss / max(ep_batches, 1)
    avg_sol = ep_sol_loss / max(ep_batches, 1)

    print(f"Epoch {ep + 1} - total: {avg:.4f}, main: {avg_main:.4f}, lune: {avg_lune:.4f}, sol: {avg_sol:.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()