Introduced new trainer with improved systems and an included timestep

trainer_v2.py

@@ -0,0 +1,862 @@
+# =====================================================================================
+# SD1.5 Flow-Matching Trainer — David-Driven Adaptive Timestep Sampling
+# Quartermaster: Mirel
+# NEW: David-weighted timesteps + SD3 shift + adaptive chaos
+# =====================================================================================
+from __future__ import annotations
+import os, json, math, random, re, shutil
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+# Diffusers
+from diffusers import StableDiffusionPipeline, DDPMScheduler
+from diffusers.models.unets.unet_2d_condition import UNet2DConditionModel
+
+# Repo deps
+from geovocab2.train.model.core.geo_david_collective import GeoDavidCollective
+from geovocab2.data.prompt.symbolic_tree import SynthesisSystem
+
+# HF / safetensors
+from huggingface_hub import snapshot_download, HfApi, create_repo, hf_hub_download
+from safetensors.torch import load_file
+
+
+# =====================================================================================
+# 1) CONFIG
+# =====================================================================================
+@dataclass
+class BaseConfig:
+    run_name: str = "sd15_flowmatch_david_weighted"
+    out_dir: str = "./runs/sd15_flowmatch_david_weighted"
+    ckpt_dir: str = "./checkpoints_sd15_flow_david_weighted"
+    save_every: int = 1
+
+    # Data
+    num_samples: int = 200_000
+    batch_size: int = 32
+    num_workers: int = 2
+    seed: int = 42
+
+    # Models / Blocks
+    model_id: str = "runwayml/stable-diffusion-v1-5"
+    active_blocks: Tuple[str, ...] = ("down_0","down_1","down_2","down_3","mid","up_0","up_1","up_2","up_3")
+    pooling: str = "mean"
+
+    # Flow training
+    epochs: int = 10
+    lr: float = 1e-4
+    weight_decay: float = 1e-3
+    grad_clip: float = 1.0
+    amp: bool = True
+    
+    global_flow_weight: float = 1.0
+    block_penalty_weight: float = 0.2
+    use_local_flow_heads: bool = False
+    local_flow_weight: float = 1.0
+
+    # KD
+    use_kd: bool = True
+    kd_weight: float = 0.25
+
+    # David
+    david_repo_id: str = "AbstractPhil/geo-david-collective-sd15-base-e40"
+    david_cache_dir: str = "./_hf_david_cache"
+    david_state_key: Optional[str] = None
+
+    # Fusion
+    alpha_timestep: float = 0.5
+    beta_pattern: float = 0.25
+    delta_incoherence: float = 0.25
+    lambda_min: float = 0.5
+    lambda_max: float = 3.0
+
+    block_weights: Dict[str, float] = None
+    
+    # Timestep Weighting (David-guided adaptive sampling)
+    use_timestep_weighting: bool = True
+    use_david_weights: bool = True
+    timestep_shift: float = 3.0  # SD3-style shift (higher = bias toward clean)
+    base_jitter: int = 5  # Base ±jitter around bin center
+    adaptive_chaos: bool = True  # Scale jitter by pattern difficulty
+    profile_samples: int = 500  # Samples to profile David's difficulty
+
+    # Scheduler
+    num_train_timesteps: int = 1000
+
+    # Inference
+    sample_steps: int = 30
+    guidance_scale: float = 7.5
+    
+    # HuggingFace
+    hf_repo_id: Optional[str] = "AbstractPhil/sd15-flow-matching"
+    upload_every_epoch: bool = True
+    continue_training: bool = True
+
+    def __post_init__(self):
+        Path(self.out_dir).mkdir(parents=True, exist_ok=True)
+        Path(self.ckpt_dir).mkdir(parents=True, exist_ok=True)
+        Path(self.david_cache_dir).mkdir(parents=True, exist_ok=True)
+        if self.block_weights is None:
+            self.block_weights = {'down_0':0.7,'down_1':0.9,'down_2':1.0,'down_3':1.1,'mid':1.2,'up_0':1.1,'up_1':1.0,'up_2':0.9,'up_3':0.7}
+
+
+# =====================================================================================
+# 2) DAVID-WEIGHTED TIMESTEP SAMPLER
+# =====================================================================================
+class DavidWeightedTimestepSampler:
+    """
+    Samples timesteps weighted by David's inherent difficulty + SD3 shift + adaptive chaos.
+    """
+    def __init__(self, num_timesteps=1000, num_bins=100, shift=3.0, base_jitter=5, adaptive_chaos=True):
+        self.num_timesteps = num_timesteps
+        self.num_bins = num_bins
+        self.shift = shift
+        self.base_jitter = base_jitter
+        self.adaptive_chaos = adaptive_chaos
+        
+        self.difficulty_weights = None  # Timestep difficulty
+        self.pattern_difficulty = None   # Pattern confusion per bin
+    
+    def _apply_shift(self, t: float) -> float:
+        """Apply SD3-style timestep shift (operates on normalized t ∈ [0,1])."""
+        if self.shift <= 0:
+            return t
+        return self.shift * t / (1.0 + (self.shift - 1.0) * t)
+    
+    def compute_difficulty_from_david(self, david, teacher, device, num_samples=500):
+        """Profile David's confusion patterns to create difficulty map."""
+        print("🔍 Profiling David's timestep & pattern difficulty...")
+        
+        david.eval()
+        teacher.eval()
+        
+        # Track David's accuracy and pattern entropy per bin
+        correct_per_bin = torch.zeros(self.num_bins)
+        total_per_bin = torch.zeros(self.num_bins)
+        entropy_per_bin = torch.zeros(self.num_bins)
+        entropy_count_per_bin = torch.zeros(self.num_bins)
+        
+        with torch.no_grad():
+            for _ in tqdm(range(num_samples // 32), desc="Profiling David", leave=False):
+                # Random latents and timesteps
+                x = torch.randn(32, 4, 64, 64, device=device, dtype=torch.float16)
+                t = torch.randint(0, self.num_timesteps, (32,), device=device)
+                t_bins = (t // 10)
+                
+                # Dummy conditioning
+                ehs = torch.randn(32, 77, 768, device=device, dtype=torch.float16)
+                
+                # Get teacher features
+                teacher.hooks.clear()
+                _ = teacher.unet(x, t, encoder_hidden_states=ehs)
+                feats = {k: v.float() for k, v in teacher.hooks.bank.items()}
+                
+                # Pool features
+                pooled = {name: f.mean(dim=(2, 3)) for name, f in feats.items()}
+                
+                # Get David's outputs
+                outputs = david(pooled, t.float())
+                
+                # 1. Timestep difficulty (from classification error)
+                ts_key = None
+                for key in ["timestep_logits", "logits_timestep", "timestep_head_logits"]:
+                    if key in outputs:
+                        ts_key = key
+                        break
+                
+                if ts_key:
+                    ts_logits = outputs[ts_key]
+                    if isinstance(ts_logits, dict):
+                        ts_logits = torch.stack(list(ts_logits.values())).mean(0)
+                    
+                    preds = ts_logits.argmax(dim=-1)
+                    for pred, true_bin in zip(preds, t_bins):
+                        bin_idx = true_bin.item()
+                        correct_per_bin[bin_idx] += (pred == true_bin).float().item()
+                        total_per_bin[bin_idx] += 1
+                
+                # 2. Pattern difficulty (from entropy)
+                pt_key = None
+                for key in ["pattern_logits", "logits_pattern", "pattern_head_logits"]:
+                    if key in outputs:
+                        pt_key = key
+                        break
+                
+                if pt_key:
+                    pt_logits = outputs[pt_key]
+                    if isinstance(pt_logits, dict):
+                        pt_logits = torch.stack(list(pt_logits.values())).mean(0)
+                    
+                    P = pt_logits.softmax(-1)
+                    ent = -(P * P.clamp_min(1e-9).log()).sum(-1)
+                    norm_ent = ent / math.log(P.shape[-1])  # Normalize by max entropy
+                    
+                    for i, true_bin in enumerate(t_bins):
+                        bin_idx = true_bin.item()
+                        entropy_per_bin[bin_idx] += norm_ent[i].item()
+                        entropy_count_per_bin[bin_idx] += 1
+        
+        # Compute timestep difficulty (inverse of accuracy)
+        accuracy_per_bin = correct_per_bin / (total_per_bin.clamp(min=1))
+        timestep_difficulty = (1.0 - accuracy_per_bin) + 0.1  # Higher = harder
+        self.difficulty_weights = timestep_difficulty / timestep_difficulty.sum()
+        
+        # Compute pattern difficulty (average entropy per bin)
+        self.pattern_difficulty = entropy_per_bin / (entropy_count_per_bin.clamp(min=1))
+        self.pattern_difficulty = self.pattern_difficulty.clamp(min=0.1, max=1.0)
+        
+        print(f"✓ David difficulty map computed:")
+        print(f"  Avg timestep accuracy: {accuracy_per_bin.mean():.2%}")
+        print(f"  Hardest timestep bin: {accuracy_per_bin.argmin().item()} ({accuracy_per_bin.min():.2%} acc)")
+        print(f"  Easiest timestep bin: {accuracy_per_bin.argmax().item()} ({accuracy_per_bin.max():.2%} acc)")
+        print(f"  Avg pattern entropy: {self.pattern_difficulty.mean():.3f}")
+        
+        return self.difficulty_weights
+    
+    def sample(self, batch_size: int) -> List[int]:
+        """Sample timesteps with David weighting + shift + adaptive chaos."""
+        if self.difficulty_weights is None:
+            # Fallback to uniform
+            return [random.randint(0, self.num_timesteps - 1) for _ in range(batch_size)]
+        
+        timesteps = []
+        for _ in range(batch_size):
+            # 1. Sample bin weighted by David's difficulty
+            bin_idx = torch.multinomial(self.difficulty_weights, 1).item()
+            
+            # 2. Get bin center as normalized t
+            bin_center_raw = bin_idx * (self.num_timesteps // self.num_bins) + (self.num_timesteps // self.num_bins) // 2
+            t_normalized = bin_center_raw / self.num_timesteps
+            
+            # 3. Apply SD3 shift
+            t_shifted = self._apply_shift(t_normalized)
+            
+            # 4. Add adaptive chaos (jitter scaled by pattern difficulty)
+            if self.adaptive_chaos:
+                chaos_scale = self.pattern_difficulty[bin_idx].item()
+                jitter = int(self.base_jitter * (0.5 + chaos_scale))  # 0.5-1.5x base jitter
+            else:
+                jitter = self.base_jitter
+            
+            # 5. Convert back to raw timestep with jitter
+            t_raw = int(t_shifted * self.num_timesteps)
+            t_raw += random.randint(-jitter, jitter)
+            t_raw = max(0, min(self.num_timesteps - 1, t_raw))
+            
+            timesteps.append(t_raw)
+        
+        return timesteps
+
+
+# =====================================================================================
+# 3) DATA
+# =====================================================================================
+class SymbolicPromptDataset(Dataset):
+    def __init__(self, n:int, seed:int=42, timestep_sampler=None):
+        self.n = n
+        self.timestep_sampler = timestep_sampler
+        random.seed(seed)
+        self.sys = SynthesisSystem(seed=seed)
+
+    def __len__(self): return self.n
+
+    def __getitem__(self, idx):
+        r = self.sys.synthesize(complexity=random.choice([1,2,3,4,5]))
+        prompt = r['text']
+        
+        if self.timestep_sampler:
+            t = self.timestep_sampler.sample(1)[0]
+        else:
+            t = random.randint(0, 999)
+        
+        return {"prompt": prompt, "t": t}
+
+def collate(batch: List[dict]):
+    prompts = [b["prompt"] for b in batch]
+    t = torch.tensor([b["t"] for b in batch], dtype=torch.long)
+    t_bins = t // 10
+    return {"prompts": prompts, "t": t, "t_bins": t_bins}
+
+
+# =====================================================================================
+# 4) HOOKS + POOLING
+# =====================================================================================
+class HookBank:
+    def __init__(self, unet: UNet2DConditionModel, active: Tuple[str, ...]):
+        self.active = set(active)
+        self.bank: Dict[str, torch.Tensor] = {}
+        self.hooks: List[torch.utils.hooks.RemovableHandle] = []
+        self._register(unet)
+
+    def _register(self, unet: UNet2DConditionModel):
+        def mk(name):
+            def h(m, i, o):
+                out = o[0] if isinstance(o,(tuple,list)) else o
+                self.bank[name] = out
+            return h
+        for i, blk in enumerate(unet.down_blocks):
+            nm = f"down_{i}"
+            if nm in self.active: self.hooks.append(blk.register_forward_hook(mk(nm)))
+        if "mid" in self.active:
+            self.hooks.append(unet.mid_block.register_forward_hook(mk("mid")))
+        for i, blk in enumerate(unet.up_blocks):
+            nm = f"up_{i}"
+            if nm in self.active: self.hooks.append(blk.register_forward_hook(mk(nm)))
+
+    def clear(self): self.bank.clear()
+    def close(self):
+        for h in self.hooks: h.remove()
+        self.hooks.clear()
+
+def spatial_pool(x: torch.Tensor, name: str, policy: str) -> torch.Tensor:
+    if policy == "mean": return x.mean(dim=(2,3))
+    if policy == "max":  return x.amax(dim=(2,3))
+    if policy == "adaptive":
+        return x.mean(dim=(2,3)) if (name.startswith("down") or name=="mid") else x.amax(dim=(2,3))
+    raise ValueError(f"Unknown pooling: {policy}")
+
+
+# =====================================================================================
+# 5) TEACHER
+# =====================================================================================
+class SD15Teacher(nn.Module):
+    def __init__(self, cfg: BaseConfig, device: str):
+        super().__init__()
+        self.pipe = StableDiffusionPipeline.from_pretrained(cfg.model_id, torch_dtype=torch.float16, safety_checker=None).to(device)
+        self.unet: UNet2DConditionModel = self.pipe.unet
+        self.text_encoder = self.pipe.text_encoder
+        self.tokenizer = self.pipe.tokenizer
+        self.hooks = HookBank(self.unet, cfg.active_blocks)
+        self.sched = DDPMScheduler(num_train_timesteps=cfg.num_train_timesteps)
+        self.device = device
+        for p in self.parameters(): p.requires_grad_(False)
+
+    @torch.no_grad()
+    def encode(self, prompts: List[str]) -> torch.Tensor:
+        tok = self.tokenizer(prompts, padding="max_length", max_length=self.tokenizer.model_max_length,
+                             truncation=True, return_tensors="pt")
+        return self.text_encoder(tok.input_ids.to(self.device))[0]
+
+    @torch.no_grad()
+    def forward_eps_and_feats(self, x_t: torch.Tensor, t: torch.LongTensor, ehs: torch.Tensor):
+        self.hooks.clear()
+        eps_hat = self.unet(x_t, t, encoder_hidden_states=ehs).sample
+        feats = {k: v.detach().float() for k, v in self.hooks.bank.items()}
+        return eps_hat.float(), feats
+
+    def alpha_sigma(self, t: torch.LongTensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        ac = self.sched.alphas_cumprod.to(self.device)[t]
+        alpha = ac.sqrt().view(-1,1,1,1).float()
+        sigma = (1.0 - ac).sqrt().view(-1,1,1,1).float()
+        return alpha, sigma
+
+
+# =====================================================================================
+# 6) STUDENT
+# =====================================================================================
+class StudentUNet(nn.Module):
+    def __init__(self, teacher_unet: UNet2DConditionModel, active_blocks: Tuple[str,...], use_local_heads: bool):
+        super().__init__()
+        self.unet = UNet2DConditionModel.from_config(teacher_unet.config)
+        self.unet.load_state_dict(teacher_unet.state_dict(), strict=True)
+        self.hooks = HookBank(self.unet, active_blocks)
+        self.use_local_heads = use_local_heads
+        self.local_heads = nn.ModuleDict()
+
+    def _ensure_heads(self, feats: Dict[str, torch.Tensor]):
+        if not self.use_local_heads: return
+        if len(self.local_heads) == len(feats): return
+        
+        target_dtype = next(self.unet.parameters()).dtype
+        
+        for name, f in feats.items():
+            c = f.shape[1]
+            if name not in self.local_heads:
+                head = nn.Conv2d(c, 4, kernel_size=1)
+                head = head.to(dtype=target_dtype, device=f.device)
+                self.local_heads[name] = head
+
+    def forward(self, x_t: torch.Tensor, t: torch.LongTensor, ehs: torch.Tensor):
+        self.hooks.clear()
+        v_hat = self.unet(x_t, t, encoder_hidden_states=ehs).sample
+        feats = {k: v for k, v in self.hooks.bank.items()}
+        self._ensure_heads(feats)
+        return v_hat, feats
+
+
+# =====================================================================================
+# 7) DAVID + ASSESSOR + FUSION
+# =====================================================================================
+class DavidLoader:
+    def __init__(self, cfg: BaseConfig, device: str):
+        self.cfg = cfg
+        self.device = device
+        self.repo_dir = snapshot_download(repo_id=cfg.david_repo_id, local_dir=cfg.david_cache_dir, local_dir_use_symlinks=False)
+        self.config_path = os.path.join(self.repo_dir, "config.json")
+        self.weights_path = os.path.join(self.repo_dir, "model.safetensors")
+        with open(self.config_path, "r") as f:
+            self.hf_config = json.load(f)
+        
+        self.gdc = GeoDavidCollective(
+            block_configs=self.hf_config["block_configs"],
+            num_timestep_bins=int(self.hf_config["num_timestep_bins"]),
+            num_patterns_per_bin=int(self.hf_config["num_patterns_per_bin"]),
+            block_weights=self.hf_config.get("block_weights", {k:1.0 for k in self.hf_config["block_configs"].keys()}),
+            loss_config=self.hf_config.get("loss_config", {})
+        ).to(device).eval()
+        
+        state = load_file(self.weights_path)
+        self.gdc.load_state_dict(state, strict=False)
+        for p in self.gdc.parameters(): p.requires_grad_(False)
+        
+        print(f"✓ David loaded from HF: {self.repo_dir}")
+        print(f"  blocks={len(self.hf_config['block_configs'])} bins={self.hf_config['num_timestep_bins']} patterns={self.hf_config['num_patterns_per_bin']}")
+        
+        if "block_weights" in self.hf_config:
+            cfg.block_weights = self.hf_config["block_weights"]
+
+class DavidAssessor(nn.Module):
+    def __init__(self, gdc: GeoDavidCollective, pooling: str):
+        super().__init__()
+        self.gdc = gdc
+        self.pooling = pooling
+
+    def _pool(self, feats: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        return {k: spatial_pool(v, k, self.pooling) for k, v in feats.items()}
+
+    @torch.no_grad()
+    def forward(self, feats_student: Dict[str, torch.Tensor], t: torch.LongTensor
+               ) -> Tuple[Dict[str,float], Dict[str,float], Dict[str,float]]:
+        Zs = self._pool(feats_student)
+        outs = self.gdc(Zs, t.float())
+        e_t, e_p, coh = {}, {}, {}
+
+        ts_key = None
+        for key in ["timestep_logits", "logits_timestep", "timestep_head_logits"]:
+            if key in outs: ts_key = key; break
+        
+        pt_key = None
+        for key in ["pattern_logits", "logits_pattern", "pattern_head_logits"]:
+            if key in outs: pt_key = key; break
+
+        t_bins = (t // 10).to(next(self.gdc.parameters()).device)
+        if ts_key is not None:
+            ts_logits = outs[ts_key]
+            if isinstance(ts_logits, dict):
+                for name, L in ts_logits.items():
+                    ce = F.cross_entropy(L, t_bins, reduction="mean")
+                    e_t[name] = float(ce.item())
+            else:
+                ce = F.cross_entropy(ts_logits, t_bins, reduction="mean")
+                for name in Zs.keys():
+                    e_t[name] = float(ce.item())
+        else:
+            for name in Zs.keys(): e_t[name] = 0.0
+
+        if pt_key is not None:
+            pt_logits = outs[pt_key]
+            if isinstance(pt_logits, dict):
+                for name, L in pt_logits.items():
+                    P = L.softmax(-1)
+                    ent = -(P * (P.clamp_min(1e-9)).log()).sum(-1).mean()
+                    e_p[name] = float(ent.item() / math.log(P.shape[-1]))
+            else:
+                P = pt_logits.softmax(-1)
+                ent = -(P * (P.clamp_min(1e-9)).log()).sum(-1).mean()
+                for name in Zs.keys():
+                    e_p[name] = float(ent.item() / math.log(P.shape[-1]))
+        else:
+            for name in Zs.keys(): e_p[name] = 0.0
+
+        alphas = {}
+        try:
+            alphas = self.gdc.get_cantor_alphas()
+        except Exception:
+            alphas = {}
+        avg_alpha = float(sum(alphas.values())/max(len(alphas),1)) if alphas else 1.0
+        for name in Zs.keys():
+            coh[name] = avg_alpha
+
+        return e_t, e_p, coh
+
+class BlockPenaltyFusion:
+    def __init__(self, cfg: BaseConfig): self.cfg = cfg
+    def lambdas(self, e_t:Dict[str,float], e_p:Dict[str,float], coh:Dict[str,float]) -> Dict[str,float]:
+        lam = {}
+        for name, base in self.cfg.block_weights.items():
+            val = base * (1.0
+                          + self.cfg.alpha_timestep * float(e_t.get(name,0.0))
+                          + self.cfg.beta_pattern   * float(e_p.get(name,0.0))
+                          + self.cfg.delta_incoherence * (1.0 - float(coh.get(name,1.0))))
+            lam[name] = float(max(self.cfg.lambda_min, min(self.cfg.lambda_max, val)))
+        return lam
+
+
+# =====================================================================================
+# 8) TRAINER
+# =====================================================================================
+class FlowMatchDavidTrainer:
+    def __init__(self, cfg: BaseConfig, device: str = "cuda"):
+        self.cfg = cfg
+        self.device = device
+        self.start_epoch = 0
+        self.start_gstep = 0
+
+        # Initialize David first (needed for timestep sampler)
+        self.david_loader = DavidLoader(cfg, device)
+        self.david = self.david_loader.gdc
+        self.assessor = DavidAssessor(self.david, cfg.pooling)
+        self.fusion = BlockPenaltyFusion(cfg)
+        
+        # Initialize teacher (needed for David profiling)
+        self.teacher = SD15Teacher(cfg, device).eval()
+        
+        # Initialize timestep sampler
+        self.timestep_sampler = None
+        if cfg.use_timestep_weighting:
+            print("\n" + "="*70)
+            print("🎯 ADAPTIVE TIMESTEP SAMPLING ENABLED")
+            print(f"   David weighting: {cfg.use_david_weights}")
+            print(f"   SD3 shift: {cfg.timestep_shift}")
+            print(f"   Base jitter: ±{cfg.base_jitter}")
+            print(f"   Adaptive chaos: {cfg.adaptive_chaos}")
+            
+            self.timestep_sampler = DavidWeightedTimestepSampler(
+                num_timesteps=cfg.num_train_timesteps,
+                num_bins=100,
+                shift=cfg.timestep_shift if cfg.use_david_weights else 0.0,
+                base_jitter=cfg.base_jitter,
+                adaptive_chaos=cfg.adaptive_chaos
+            )
+            
+            if cfg.use_david_weights:
+                self.timestep_sampler.compute_difficulty_from_david(
+                    david=self.david,
+                    teacher=self.teacher,
+                    device=device,
+                    num_samples=cfg.profile_samples
+                )
+            print("="*70 + "\n")
+        
+        # Initialize dataset with sampler
+        self.dataset = SymbolicPromptDataset(cfg.num_samples, cfg.seed, self.timestep_sampler)
+        self.loader = DataLoader(self.dataset, batch_size=cfg.batch_size, shuffle=True,
+                                 num_workers=cfg.num_workers, pin_memory=True, collate_fn=collate)
+
+        # Initialize student
+        self.student = StudentUNet(self.teacher.unet, cfg.active_blocks, cfg.use_local_flow_heads).to(device)
+
+        self.opt = torch.optim.AdamW(self.student.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+        self.sched = torch.optim.lr_scheduler.CosineAnnealingLR(self.opt, T_max=cfg.epochs * len(self.loader))
+        self.scaler = torch.cuda.amp.GradScaler(enabled=cfg.amp)
+
+        # Load checkpoints
+        emergency_path = Path("./EMERGENCY_SAVE_SUCCESS.pt")
+        if not emergency_path.exists():
+            print("\n🔍 Emergency checkpoint not found locally, checking HuggingFace...")
+            emergency_path = self._download_emergency_checkpoint()
+        
+        if emergency_path and emergency_path.exists():
+            self._load_emergency_checkpoint(emergency_path)
+        elif cfg.continue_training:
+            self._load_latest_from_hf()
+
+        self.writer = SummaryWriter(log_dir=os.path.join(cfg.out_dir, cfg.run_name))
+
+    def _download_emergency_checkpoint(self) -> Optional[Path]:
+        """Download emergency checkpoint from HuggingFace backup repo."""
+        emergency_repo = "AbstractPhil/sd15-flow-emergency-backup"
+        emergency_file = "EMERGENCY_SAVE_SUCCESS.pt"
+        
+        try:
+            print(f"📥 Downloading emergency checkpoint from {emergency_repo}...")
+            local_path = hf_hub_download(
+                repo_id=emergency_repo,
+                filename=emergency_file,
+                repo_type="model",
+                cache_dir="./_emergency_cache"
+            )
+            
+            target_path = Path("./EMERGENCY_SAVE_SUCCESS.pt")
+            shutil.copy(local_path, target_path)
+            
+            size_mb = target_path.stat().st_size / 1e6
+            print(f"✅ Downloaded emergency checkpoint ({size_mb:.1f} MB)")
+            return target_path
+            
+        except Exception as e:
+            print(f"⚠️ Could not download emergency checkpoint: {e}")
+            return None
+
+    def _load_emergency_checkpoint(self, path: Path):
+        """Load emergency checkpoint with student_unet structure."""
+        try:
+            print(f"\n🚨 Found emergency checkpoint: {path}")
+            checkpoint = torch.load(path, map_location='cpu')
+            
+            if 'student_unet' in checkpoint:
+                print("📦 Loading emergency checkpoint format...")
+                missing, unexpected = self.student.unet.load_state_dict(checkpoint['student_unet'], strict=False)
+                print(f"✓ Loaded student UNet")
+                
+                if 'opt' in checkpoint:
+                    self.opt.load_state_dict(checkpoint['opt'])
+                    print("✓ Loaded optimizer state")
+                
+                if 'sched' in checkpoint:
+                    self.sched.load_state_dict(checkpoint['sched'])
+                    print("✓ Loaded scheduler state")
+                
+                if 'gstep' in checkpoint:
+                    self.start_gstep = checkpoint['gstep']
+                    self.start_epoch = self.start_gstep // len(self.loader)
+                    print(f"✓ Resuming from global step {self.start_gstep} (epoch ~{self.start_epoch})")
+                
+                print("✅ Emergency checkpoint loaded successfully!")
+                
+        except Exception as e:
+            print(f"⚠️ Failed to load emergency checkpoint: {e}")
+
+    def _load_latest_from_hf(self):
+        if not self.cfg.hf_repo_id:
+            return
+        
+        try:
+            api = HfApi()
+            print(f"\n🔍 Searching for latest checkpoint in {self.cfg.hf_repo_id}...")
+            
+            files = api.list_repo_files(repo_id=self.cfg.hf_repo_id, repo_type="model")
+            epochs = []
+            for f in files:
+                if f.endswith('.pt'):
+                    match = re.search(r'_e(\d+)\.pt$', f)
+                    if match:
+                        epochs.append((int(match.group(1)), f))
+            
+            if not epochs:
+                return
+            
+            latest_epoch, latest_file = max(epochs, key=lambda x: x[0])
+            print(f"📥 Downloading: {latest_file}")
+            
+            local_path = hf_hub_download(
+                repo_id=self.cfg.hf_repo_id,
+                filename=latest_file,
+                repo_type="model",
+                cache_dir=self.cfg.ckpt_dir
+            )
+            
+            checkpoint = torch.load(local_path, map_location='cpu')
+            
+            if 'student_unet' in checkpoint:
+                self.student.unet.load_state_dict(checkpoint['student_unet'], strict=False)
+            elif 'student' in checkpoint:
+                self.student.load_state_dict(checkpoint['student'], strict=False)
+            
+            if 'opt' in checkpoint:
+                self.opt.load_state_dict(checkpoint['opt'])
+            if 'sched' in checkpoint:
+                self.sched.load_state_dict(checkpoint['sched'])
+            
+            self.start_epoch = latest_epoch
+            self.start_gstep = latest_epoch * len(self.loader)
+            
+            print(f"✅ Resuming from epoch {self.start_epoch + 1}")
+            del checkpoint
+            torch.cuda.empty_cache()
+            
+        except Exception as e:
+            print(f"⚠️ Failed to load from HF: {e}")
+
+    def _v_star(self, x_t, t, eps_hat):
+        alpha, sigma = self.teacher.alpha_sigma(t)
+        x0_hat = (x_t - sigma * eps_hat) / (alpha + 1e-8)
+        return alpha * eps_hat - sigma * x0_hat
+
+    def _down_like(self, tgt: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
+        return F.interpolate(tgt, size=ref.shape[-2:], mode="bilinear", align_corners=False)
+
+    def _kd_cos(self, s: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        s = F.normalize(s, dim=-1); t = F.normalize(t, dim=-1)
+        return 1.0 - (s*t).sum(-1).mean()
+
+    def train(self):
+        cfg = self.cfg
+        gstep = self.start_gstep
+        
+        for ep in range(self.start_epoch, cfg.epochs):
+            self.student.train()
+            pbar = tqdm(self.loader, desc=f"Epoch {ep+1}/{cfg.epochs}", 
+                        dynamic_ncols=True, leave=True, position=0)
+            acc = {"L":0.0, "Lf":0.0, "Lb":0.0}
+
+            for it, batch in enumerate(pbar):
+                prompts = batch["prompts"]
+                t = batch["t"].to(self.device)
+
+                with torch.no_grad():
+                    ehs = self.teacher.encode(prompts)
+
+                x_t = torch.randn(len(prompts), 4, 64, 64, device=self.device, dtype=torch.float16)
+
+                with torch.no_grad():
+                    eps_hat, t_feats_spatial = self.teacher.forward_eps_and_feats(x_t.half(), t, ehs)
+                    v_star = self._v_star(x_t, t, eps_hat)
+
+                with torch.cuda.amp.autocast(enabled=cfg.amp):
+                    v_hat, s_feats_spatial = self.student(x_t, t, ehs)
+                    L_flow = F.mse_loss(v_hat, v_star)
+
+                    e_t, e_p, coh = self.assessor(s_feats_spatial, t)
+                    lam = self.fusion.lambdas(e_t, e_p, coh)
+
+                    L_blocks = torch.zeros((), device=self.device)
+                    for name, s_feat in s_feats_spatial.items():
+                        L_kd = torch.zeros((), device=self.device)
+                        if cfg.use_kd:
+                            s_pool = spatial_pool(s_feat, name, cfg.pooling)
+                            t_pool = spatial_pool(t_feats_spatial[name], name, cfg.pooling)
+                            L_kd = self._kd_cos(s_pool, t_pool)
+                        
+                        L_lf = torch.zeros((), device=self.device)
+                        if cfg.use_local_flow_heads and name in self.student.local_heads:
+                            v_loc = self.student.local_heads[name](s_feat)
+                            v_ds  = self._down_like(v_star, v_loc)
+                            L_lf  = F.mse_loss(v_loc, v_ds)
+                        L_blocks = L_blocks + lam.get(name,1.0) * (cfg.kd_weight * L_kd + cfg.local_flow_weight * L_lf)
+
+                    L_total = cfg.global_flow_weight*L_flow + cfg.block_penalty_weight*L_blocks
+
+                self.opt.zero_grad(set_to_none=True)
+                if cfg.amp:
+                    self.scaler.scale(L_total).backward()
+                    nn.utils.clip_grad_norm_(self.student.parameters(), cfg.grad_clip)
+                    self.scaler.step(self.opt); self.scaler.update()
+                else:
+                    L_total.backward()
+                    nn.utils.clip_grad_norm_(self.student.parameters(), cfg.grad_clip)
+                    self.opt.step()
+                self.sched.step(); gstep += 1
+
+                acc["L"]  += float(L_total.item())
+                acc["Lf"] += float(L_flow.item())
+                acc["Lb"] += float(L_blocks.item())
+
+                if it % 50 == 0:
+                    self.writer.add_scalar("train/total", float(L_total.item()), gstep)
+                    self.writer.add_scalar("train/flow",  float(L_flow.item()),  gstep)
+                    self.writer.add_scalar("train/blocks",float(L_blocks.item()), gstep)
+                    for k in list(lam.keys())[:4]:
+                        self.writer.add_scalar(f"lambda/{k}", lam[k], gstep)
+
+                if it % 10 == 0 or it == len(self.loader) - 1:
+                    pbar.set_postfix({
+                        "L": f"{float(L_total.item()):.4f}", 
+                        "Lf": f"{float(L_flow.item()):.4f}", 
+                        "Lb": f"{float(L_blocks.item()):.4f}"
+                    }, refresh=False)
+                
+                del x_t, eps_hat, v_star, v_hat, s_feats_spatial, t_feats_spatial
+
+            pbar.close()
+            
+            n = len(self.loader)
+            print(f"\n[Epoch {ep+1}] L={acc['L']/n:.4f} | L_flow={acc['Lf']/n:.4f} | L_blocks={acc['Lb']/n:.4f}")
+            self.writer.add_scalar("epoch/total", acc['L']/n, ep+1)
+            self.writer.add_scalar("epoch/flow",  acc['Lf']/n, ep+1)
+            self.writer.add_scalar("epoch/blocks",acc['Lb']/n, ep+1)
+
+            if (ep+1) % cfg.save_every == 0:
+                self._save(ep+1, gstep)
+
+        self._save("final", gstep)
+        self.writer.close()
+
+    def _save(self, tag, gstep):
+        """Save checkpoint and upload to HuggingFace."""
+        pt_path = Path(self.cfg.ckpt_dir) / f"{self.cfg.run_name}_e{tag}.pt"
+        torch.save({
+            "cfg": asdict(self.cfg),
+            "student": self.student.state_dict(),
+            "opt": self.opt.state_dict(),
+            "sched": self.sched.state_dict(),
+            "gstep": gstep
+        }, pt_path)
+        
+        size_mb = pt_path.stat().st_size / 1e6
+        print(f"✓ Saved checkpoint: {pt_path.name} ({size_mb:.1f} MB)")
+        
+        if self.cfg.upload_every_epoch and self.cfg.hf_repo_id:
+            self._upload_to_hf(pt_path, tag)
+
+    def _upload_to_hf(self, path: Path, tag):
+        """Upload checkpoint to HuggingFace."""
+        try:
+            api = HfApi()
+            create_repo(self.cfg.hf_repo_id, exist_ok=True, private=False, repo_type="model")
+            
+            print(f"📤 Uploading {path.name} to {self.cfg.hf_repo_id}...")
+            api.upload_file(
+                path_or_fileobj=str(path),
+                path_in_repo=path.name,
+                repo_id=self.cfg.hf_repo_id,
+                repo_type="model",
+                commit_message=f"Epoch {tag}"
+            )
+            print(f"✅ Uploaded: https://huggingface.co/{self.cfg.hf_repo_id}/{path.name}")
+            
+        except Exception as e:
+            print(f"⚠️ Upload failed: {e}")
+
+    @torch.no_grad()
+    def sample(self, prompts: List[str], steps: Optional[int]=None, guidance: Optional[float]=None) -> torch.Tensor:
+        steps = steps or self.cfg.sample_steps
+        guidance = guidance if guidance is not None else self.cfg.guidance_scale
+        cond_e = self.teacher.encode(prompts)
+        uncond_e = self.teacher.encode([""]*len(prompts))
+        sched = self.teacher.sched
+        sched.set_timesteps(steps, device=self.device)
+        x_t = torch.randn(len(prompts), 4, 64, 64, device=self.device)
+
+        for t_scalar in sched.timesteps:
+            t = torch.full((x_t.shape[0],), t_scalar, device=self.device, dtype=torch.long)
+            v_u, _ = self.student(x_t, t, uncond_e)
+            v_c, _ = self.student(x_t, t, cond_e)
+            v_hat = v_u + guidance*(v_c - v_u)
+
+            alpha, sigma = self.teacher.alpha_sigma(t)
+            denom = (alpha**2 + sigma**2)
+            x0_hat = (alpha * x_t - sigma * v_hat) / (denom + 1e-8)
+            eps_hat = (x_t - alpha * x0_hat) / (sigma + 1e-8)
+
+            step = sched.step(model_output=eps_hat, timestep=t_scalar, sample=x_t)
+            x_t = step.prev_sample
+
+        imgs = self.teacher.pipe.vae.decode(x_t / 0.18215).sample
+        return imgs.clamp(-1,1)
+
+
+# =====================================================================================
+# 9) MAIN
+# =====================================================================================
+def main():
+    cfg = BaseConfig()
+    print(json.dumps(asdict(cfg), indent=2))
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    if device != "cuda":
+        print("⚠️ A100 strongly recommended.")
+    trainer = FlowMatchDavidTrainer(cfg, device=device)
+    trainer.train()
+    _ = trainer.sample(["a castle at sunset"], steps=10, guidance=7.0)
+    print("✓ Training complete.")
+
+if __name__ == "__main__":
+    main()