Create trainer.py

trainer.py

@@ -0,0 +1,653 @@
+# ============================================================================
+# TinyFlux Training Cell - Full Featured
+# ============================================================================
+# Run the model cell before this one (defines TinyFlux, TinyFluxConfig)
+# Dataset: AbstractPhil/flux-schnell-teacher-latents
+# Uploads checkpoints to: AbstractPhil/tiny-flux
+# ============================================================================
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from datasets import load_dataset
+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 os
+import json
+from datetime import datetime
+
+# ============================================================================
+# CONFIG
+# ============================================================================
+BATCH_SIZE = 4
+GRAD_ACCUM = 2
+LR = 1e-4
+EPOCHS = 10
+MAX_SEQ = 128
+MIN_SNR = 5.0
+SHIFT = 3.0
+DEVICE = "cuda"
+DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+
+# HuggingFace Hub
+HF_REPO = "AbstractPhil/tiny-flux"
+SAVE_EVERY = 500      # steps - local save
+UPLOAD_EVERY = 500    # steps - hub upload
+SAMPLE_EVERY = 250    # steps - generate samples
+LOG_EVERY = 10        # steps - tensorboard
+
+# Checkpoint loading target
+# Options:
+#   None or "latest" - load most recent checkpoint
+#   "best" - load best model
+#   int (e.g. 1500) - load specific step
+#   "hub:step_1000" - load specific checkpoint from hub
+#   "local:path/to/checkpoint.safetensors" - load specific local file
+#   "none" - start fresh, ignore existing checkpoints
+LOAD_TARGET = "latest"
+
+# Local paths
+CHECKPOINT_DIR = "./tiny_flux_checkpoints"
+LOG_DIR = "./tiny_flux_logs"
+SAMPLE_DIR = "./tiny_flux_samples"
+
+os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+os.makedirs(LOG_DIR, exist_ok=True)
+os.makedirs(SAMPLE_DIR, exist_ok=True)
+
+# ============================================================================
+# HF HUB SETUP
+# ============================================================================
+print("Setting up HuggingFace Hub...")
+api = HfApi()
+
+try:
+    api.create_repo(repo_id=HF_REPO, exist_ok=True, repo_type="model")
+    print(f"✓ Repo ready: {HF_REPO}")
+except Exception as e:
+    print(f"Note: {e}")
+
+# ============================================================================
+# TENSORBOARD
+# ============================================================================
+run_name = datetime.now().strftime("%Y%m%d_%H%M%S")
+writer = SummaryWriter(log_dir=os.path.join(LOG_DIR, run_name))
+print(f"✓ Tensorboard: {LOG_DIR}/{run_name}")
+
+# ============================================================================
+# LOAD DATASET
+# ============================================================================
+print("\nLoading dataset...")
+ds = load_dataset("AbstractPhil/flux-schnell-teacher-latents", split="train")
+print(f"Samples: {len(ds)}")
+
+# ============================================================================
+# LOAD TEXT ENCODERS
+# ============================================================================
+print("\nLoading flan-t5-base (768 dim)...")
+t5_tok = T5Tokenizer.from_pretrained("google/flan-t5-base")
+t5_enc = T5EncoderModel.from_pretrained("google/flan-t5-base", torch_dtype=DTYPE).to(DEVICE).eval()
+
+print("Loading CLIP-L...")
+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 t5_enc.parameters(): p.requires_grad = False
+for p in clip_enc.parameters(): p.requires_grad = False
+
+# ============================================================================
+# LOAD VAE FOR SAMPLE GENERATION
+# ============================================================================
+print("Loading Flux VAE for samples...")
+from diffusers import AutoencoderKL
+vae = AutoencoderKL.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell", 
+    subfolder="vae",
+    torch_dtype=DTYPE
+).to(DEVICE).eval()
+for p in vae.parameters(): p.requires_grad = False
+
+# ============================================================================
+# ENCODING HELPERS
+# ============================================================================
+@torch.no_grad()
+def encode_prompt(prompt):
+    t5_in = t5_tok(prompt, max_length=MAX_SEQ, padding="max_length", truncation=True, return_tensors="pt").to(DEVICE)
+    t5_out = t5_enc(input_ids=t5_in.input_ids, attention_mask=t5_in.attention_mask).last_hidden_state
+    
+    clip_in = clip_tok(prompt, max_length=77, padding="max_length", truncation=True, return_tensors="pt").to(DEVICE)
+    clip_out = clip_enc(input_ids=clip_in.input_ids, attention_mask=clip_in.attention_mask)
+    return t5_out, clip_out.pooler_output
+
+# ============================================================================
+# FLOW MATCHING HELPERS
+# ============================================================================
+def flux_shift(t, s=SHIFT):
+    return s * t / (1 + (s - 1) * t)
+
+def flux_shift_inverse(t_shifted, s=SHIFT):
+    """Inverse of flux_shift for sampling."""
+    return t_shifted / (s - (s - 1) * t_shifted)
+
+def min_snr_weight(t, gamma=MIN_SNR):
+    snr = (t / (1 - t).clamp(min=1e-5)).pow(2)
+    return torch.clamp(snr, max=gamma) / snr.clamp(min=1e-5)
+
+# ============================================================================
+# SAMPLING FUNCTION
+# ============================================================================
+@torch.no_grad()
+def generate_samples(model, prompts, num_steps=20, guidance_scale=3.5, H=64, W=64):
+    """Generate sample images using Euler sampling."""
+    model.eval()
+    B = len(prompts)
+    C = 16  # VAE channels
+    
+    # Encode prompts
+    t5_embeds, clip_pooleds = [], []
+    for p in prompts:
+        t5_out, clip_pooled = encode_prompt(p)
+        t5_embeds.append(t5_out.squeeze(0))
+        clip_pooleds.append(clip_pooled.squeeze(0))
+    t5_embeds = torch.stack(t5_embeds)
+    clip_pooleds = torch.stack(clip_pooleds)
+    
+    # Start from noise
+    x = torch.randn(B, H * W, C, device=DEVICE, dtype=DTYPE)
+    
+    # Create image IDs
+    img_ids = TinyFlux.create_img_ids(B, H, W, DEVICE)
+    
+    # Euler sampling with uniform timesteps
+    timesteps = torch.linspace(1, 0, num_steps + 1, device=DEVICE)[:-1]
+    
+    for i, t in enumerate(timesteps):
+        t_batch = t.expand(B)
+        dt = 1.0 / num_steps
+        
+        # Conditional prediction
+        guidance = torch.full((B,), guidance_scale, device=DEVICE, dtype=DTYPE)
+        v_cond = model(
+            hidden_states=x,
+            encoder_hidden_states=t5_embeds,
+            pooled_projections=clip_pooleds,
+            timestep=t_batch,
+            img_ids=img_ids,
+            guidance=guidance,
+        )
+        
+        # Euler step: x = x + v * dt (going from noise to data)
+        x = x + v_cond * dt
+    
+    # Reshape to image format: (B, H*W, C) -> (B, C, H, W)
+    latents = x.reshape(B, H, W, C).permute(0, 3, 1, 2)
+    
+    # Decode with VAE
+    latents = latents / vae.config.scaling_factor
+    images = vae.decode(latents.float()).sample
+    images = (images / 2 + 0.5).clamp(0, 1)
+    
+    model.train()
+    return images
+
+def save_samples(images, prompts, step, save_dir):
+    """Save sample images and log to tensorboard."""
+    from torchvision.utils import make_grid, save_image
+    
+    # Save individual images
+    for i, (img, prompt) in enumerate(zip(images, prompts)):
+        safe_prompt = prompt[:50].replace(" ", "_").replace("/", "-")
+        path = os.path.join(save_dir, f"step{step}_{i}_{safe_prompt}.png")
+        save_image(img, path)
+    
+    # Log grid to tensorboard
+    grid = make_grid(images, nrow=2, normalize=False)
+    writer.add_image("samples", grid, step)
+    
+    # Log prompts
+    writer.add_text("sample_prompts", "\n".join(prompts), step)
+    
+    print(f"  ✓ Saved {len(images)} samples")
+
+# ============================================================================
+# COLLATE
+# ============================================================================
+def collate(batch):
+    latents, t5_embeds, clip_embeds, prompts = [], [], [], []
+    for b in batch:
+        latents.append(torch.tensor(np.array(b["latent"]), dtype=DTYPE))
+        t5_out, clip_pooled = encode_prompt(b["prompt"])
+        t5_embeds.append(t5_out.squeeze(0))
+        clip_embeds.append(clip_pooled.squeeze(0))
+        prompts.append(b["prompt"])
+    return {
+        "latents": torch.stack(latents).to(DEVICE),
+        "t5_embeds": torch.stack(t5_embeds),
+        "clip_pooled": torch.stack(clip_embeds),
+        "prompts": prompts,
+    }
+
+# ============================================================================
+# CHECKPOINT FUNCTIONS
+# ============================================================================
+def save_checkpoint(model, optimizer, scheduler, step, epoch, loss, path):
+    """Save checkpoint locally."""
+    os.makedirs(os.path.dirname(path) if os.path.dirname(path) else ".", exist_ok=True)
+    
+    weights_path = path.replace(".pt", ".safetensors")
+    save_file(model.state_dict(), weights_path)
+    
+    state = {
+        "step": step,
+        "epoch": epoch,
+        "loss": loss,
+        "optimizer": optimizer.state_dict(),
+        "scheduler": scheduler.state_dict(),
+    }
+    torch.save(state, path)
+    print(f"  ✓ Saved checkpoint: step {step}")
+    return weights_path
+
+def upload_checkpoint(weights_path, step, config, include_logs=True):
+    """Upload checkpoint to HuggingFace Hub."""
+    try:
+        # Upload weights
+        api.upload_file(
+            path_or_fileobj=weights_path,
+            path_in_repo=f"checkpoints/step_{step}.safetensors",
+            repo_id=HF_REPO,
+            commit_message=f"Checkpoint step {step}",
+        )
+        
+        # Upload config
+        config_path = os.path.join(CHECKPOINT_DIR, "config.json")
+        with open(config_path, "w") as f:
+            json.dump(config.__dict__, f, indent=2)
+        api.upload_file(
+            path_or_fileobj=config_path,
+            path_in_repo="config.json",
+            repo_id=HF_REPO,
+        )
+        
+        # Upload tensorboard logs
+        if include_logs and os.path.exists(LOG_DIR):
+            api.upload_folder(
+                folder_path=LOG_DIR,
+                path_in_repo="logs",
+                repo_id=HF_REPO,
+                commit_message=f"Logs at step {step}",
+            )
+        
+        # Upload samples
+        if os.path.exists(SAMPLE_DIR) and os.listdir(SAMPLE_DIR):
+            api.upload_folder(
+                folder_path=SAMPLE_DIR,
+                path_in_repo="samples",
+                repo_id=HF_REPO,
+                commit_message=f"Samples at step {step}",
+            )
+        
+        print(f"  ✓ Uploaded to {HF_REPO}")
+    except Exception as e:
+        print(f"  ⚠ Upload failed: {e}")
+
+def load_checkpoint(model, optimizer, scheduler, target):
+    """
+    Load checkpoint based on target specification.
+    
+    Args:
+        target: 
+            None, "latest" - most recent checkpoint
+            "best" - best model
+            int (1500) - specific step
+            "hub:step_1000" - specific hub checkpoint  
+            "local:/path/to/file.safetensors" - specific local file
+            "none" - skip loading, start fresh
+    """
+    if target == "none":
+        print("Starting fresh (no checkpoint loading)")
+        return 0, 0
+    
+    start_step, start_epoch = 0, 0
+    
+    # Parse target
+    if target is None or target == "latest":
+        load_mode = "latest"
+        load_path = None
+    elif target == "best":
+        load_mode = "best"
+        load_path = None
+    elif isinstance(target, int):
+        load_mode = "step"
+        load_path = target
+    elif target.startswith("hub:"):
+        load_mode = "hub"
+        load_path = target[4:]  # Remove "hub:" prefix
+    elif target.startswith("local:"):
+        load_mode = "local"
+        load_path = target[6:]  # Remove "local:" prefix
+    else:
+        print(f"Unknown target format: {target}, trying as step number")
+        try:
+            load_mode = "step"
+            load_path = int(target)
+        except:
+            load_mode = "latest"
+            load_path = None
+    
+    # Load based on mode
+    if load_mode == "local":
+        # Direct local file
+        if os.path.exists(load_path):
+            weights = load_file(load_path)
+            model.load_state_dict(weights)
+            # Try to find associated state file
+            state_path = load_path.replace(".safetensors", ".pt")
+            if os.path.exists(state_path):
+                state = torch.load(state_path, weights_only=False)
+                optimizer.load_state_dict(state["optimizer"])
+                scheduler.load_state_dict(state["scheduler"])
+                start_step = state.get("step", 0)
+                start_epoch = state.get("epoch", 0)
+            print(f"✓ Loaded local: {load_path} (step {start_step})")
+            return start_step, start_epoch
+        else:
+            print(f"⚠ Local file not found: {load_path}")
+    
+    elif load_mode == "hub":
+        # Specific hub checkpoint
+        try:
+            filename = f"checkpoints/{load_path}.safetensors" if not load_path.endswith(".safetensors") else load_path
+            local_path = hf_hub_download(repo_id=HF_REPO, filename=filename)
+            weights = load_file(local_path)
+            model.load_state_dict(weights)
+            # Extract step from filename
+            if "step_" in load_path:
+                start_step = int(load_path.split("step_")[-1].replace(".safetensors", ""))
+            print(f"✓ Loaded from Hub: {filename} (step {start_step})")
+            return start_step, start_epoch
+        except Exception as e:
+            print(f"⚠ Hub load failed: {e}")
+    
+    elif load_mode == "best":
+        # Try hub best first
+        try:
+            local_path = hf_hub_download(repo_id=HF_REPO, filename="model.safetensors")
+            weights = load_file(local_path)
+            model.load_state_dict(weights)
+            print(f"✓ Loaded best model from Hub")
+            return start_step, start_epoch
+        except:
+            pass
+        # Try local best
+        best_path = os.path.join(CHECKPOINT_DIR, "best.safetensors")
+        if os.path.exists(best_path):
+            weights = load_file(best_path)
+            model.load_state_dict(weights)
+            state_path = best_path.replace(".safetensors", ".pt")
+            if os.path.exists(state_path):
+                state = torch.load(state_path, weights_only=False)
+                start_step = state.get("step", 0)
+                start_epoch = state.get("epoch", 0)
+            print(f"✓ Loaded local best (step {start_step})")
+            return start_step, start_epoch
+    
+    elif load_mode == "step":
+        # Specific step number
+        step_num = load_path
+        # Try hub
+        try:
+            filename = f"checkpoints/step_{step_num}.safetensors"
+            local_path = hf_hub_download(repo_id=HF_REPO, filename=filename)
+            weights = load_file(local_path)
+            model.load_state_dict(weights)
+            start_step = step_num
+            print(f"✓ Loaded step {step_num} from Hub")
+            return start_step, start_epoch
+        except:
+            pass
+        # Try local
+        local_path = os.path.join(CHECKPOINT_DIR, f"step_{step_num}.safetensors")
+        if os.path.exists(local_path):
+            weights = load_file(local_path)
+            model.load_state_dict(weights)
+            state_path = local_path.replace(".safetensors", ".pt")
+            if os.path.exists(state_path):
+                state = torch.load(state_path, weights_only=False)
+                optimizer.load_state_dict(state["optimizer"])
+                scheduler.load_state_dict(state["scheduler"])
+                start_epoch = state.get("epoch", 0)
+            start_step = step_num
+            print(f"✓ Loaded local step {step_num}")
+            return start_step, start_epoch
+        print(f"⚠ Step {step_num} not found")
+    
+    # Default: latest
+    # Try Hub first
+    try:
+        files = api.list_repo_files(repo_id=HF_REPO)
+        checkpoints = [f for f in files if f.startswith("checkpoints/step_") and f.endswith(".safetensors")]
+        if checkpoints:
+            checkpoints.sort(key=lambda x: int(x.split("step_")[-1].replace(".safetensors", "")))
+            latest = checkpoints[-1]
+            step = int(latest.split("step_")[-1].replace(".safetensors", ""))
+            local_path = hf_hub_download(repo_id=HF_REPO, filename=latest)
+            weights = load_file(local_path)
+            model.load_state_dict(weights)
+            start_step = step
+            print(f"✓ Loaded latest from Hub: step {step}")
+            return start_step, start_epoch
+    except Exception as e:
+        print(f"Hub check: {e}")
+    
+    # Try local
+    if os.path.exists(CHECKPOINT_DIR):
+        local_ckpts = [f for f in os.listdir(CHECKPOINT_DIR) if f.startswith("step_") and f.endswith(".safetensors")]
+        if local_ckpts:
+            local_ckpts.sort(key=lambda x: int(x.split("step_")[-1].replace(".safetensors", "")))
+            latest = local_ckpts[-1]
+            step = int(latest.split("step_")[-1].replace(".safetensors", ""))
+            weights_path = os.path.join(CHECKPOINT_DIR, latest)
+            weights = load_file(weights_path)
+            model.load_state_dict(weights)
+            state_path = weights_path.replace(".safetensors", ".pt")
+            if os.path.exists(state_path):
+                state = torch.load(state_path, weights_only=False)
+                optimizer.load_state_dict(state["optimizer"])
+                scheduler.load_state_dict(state["scheduler"])
+                start_epoch = state.get("epoch", 0)
+            start_step = step
+            print(f"✓ Loaded latest local: step {step}")
+            return start_step, start_epoch
+    
+    print("No checkpoint found, starting fresh")
+    return 0, 0
+
+# ============================================================================
+# DATALOADER
+# ============================================================================
+loader = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate, num_workers=0)
+
+# ============================================================================
+# MODEL
+# ============================================================================
+config = TinyFluxConfig()
+model = TinyFlux(config).to(DEVICE).to(DTYPE)
+print(f"\nParams: {sum(p.numel() for p in model.parameters()):,}")
+
+# ============================================================================
+# OPTIMIZER & SCHEDULER
+# ============================================================================
+opt = torch.optim.AdamW(model.parameters(), lr=LR, betas=(0.9, 0.99), weight_decay=0.01)
+total_steps = len(loader) * EPOCHS // GRAD_ACCUM
+warmup = min(500, 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
+# ============================================================================
+print(f"\nLoad target: {LOAD_TARGET}")
+start_step, start_epoch = load_checkpoint(model, opt, sched, LOAD_TARGET)
+
+# Log config to tensorboard
+writer.add_text("config", json.dumps(config.__dict__, indent=2), 0)
+writer.add_text("training_config", json.dumps({
+    "batch_size": BATCH_SIZE,
+    "grad_accum": GRAD_ACCUM,
+    "lr": LR,
+    "epochs": EPOCHS,
+    "min_snr": MIN_SNR,
+    "shift": SHIFT,
+}, indent=2), 0)
+
+# ============================================================================
+# SAMPLE PROMPTS FOR PERIODIC GENERATION
+# ============================================================================
+SAMPLE_PROMPTS = [
+    "a photo of a cat sitting on a windowsill",
+    "a beautiful sunset over mountains",
+    "a portrait of a woman with red hair",
+    "a futuristic cityscape at night",
+]
+
+# ============================================================================
+# TRAINING
+# ============================================================================
+print(f"\nTraining {EPOCHS} epochs, {total_steps} total steps")
+print(f"Resuming from step {start_step}, epoch {start_epoch}")
+print(f"Save: {SAVE_EVERY}, Upload: {UPLOAD_EVERY}, Sample: {SAMPLE_EVERY}, Log: {LOG_EVERY}")
+
+model.train()
+step = start_step
+best = float("inf")
+
+for ep in range(start_epoch, EPOCHS):
+    ep_loss = 0
+    ep_batches = 0
+    pbar = tqdm(loader, desc=f"E{ep+1}")
+    
+    for i, batch in enumerate(pbar):
+        lat = batch["latents"]
+        t5 = batch["t5_embeds"]
+        clip = batch["clip_pooled"]
+        
+        B, C, H, W = lat.shape
+        x1 = lat.permute(0, 2, 3, 1).reshape(B, H*W, C)
+        x0 = torch.randn_like(x1)
+        
+        t = torch.sigmoid(torch.randn(B, device=DEVICE))
+        t = flux_shift(t).to(DTYPE).clamp(1e-4, 1-1e-4)
+        
+        t_exp = t.view(B, 1, 1)
+        x_t = (1 - t_exp) * x0 + t_exp * x1
+        v_target = x1 - x0
+        
+        img_ids = TinyFlux.create_img_ids(B, H, W, DEVICE)
+        guidance = torch.rand(B, device=DEVICE, dtype=DTYPE) * 4 + 1
+        
+        with torch.autocast("cuda", dtype=DTYPE):
+            pred = model(
+                hidden_states=x_t,
+                encoder_hidden_states=t5,
+                pooled_projections=clip,
+                timestep=t,
+                img_ids=img_ids,
+                guidance=guidance,
+            )
+        
+        loss_raw = F.mse_loss(pred, v_target, reduction="none").mean(dim=[1,2])
+        snr_weights = min_snr_weight(t)
+        loss = (loss_raw * snr_weights).mean() / GRAD_ACCUM
+        loss.backward()
+        
+        if (i + 1) % GRAD_ACCUM == 0:
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+            sched.step()
+            opt.zero_grad()
+            step += 1
+            
+            # Tensorboard logging
+            if step % LOG_EVERY == 0:
+                writer.add_scalar("train/loss", loss.item() * GRAD_ACCUM, step)
+                writer.add_scalar("train/lr", sched.get_last_lr()[0], step)
+                writer.add_scalar("train/grad_norm", grad_norm.item(), step)
+                writer.add_scalar("train/t_mean", t.mean().item(), step)
+                writer.add_scalar("train/snr_weight_mean", snr_weights.mean().item(), step)
+            
+            # Generate samples
+            if step % SAMPLE_EVERY == 0:
+                print(f"\n  Generating samples at step {step}...")
+                images = generate_samples(model, SAMPLE_PROMPTS, num_steps=20)
+                save_samples(images, SAMPLE_PROMPTS, step, SAMPLE_DIR)
+            
+            # Save checkpoint
+            if step % SAVE_EVERY == 0:
+                ckpt_path = os.path.join(CHECKPOINT_DIR, f"step_{step}.pt")
+                weights_path = save_checkpoint(model, opt, sched, step, ep, loss.item(), ckpt_path)
+                
+                # Upload
+                if step % UPLOAD_EVERY == 0:
+                    upload_checkpoint(weights_path, step, config, include_logs=True)
+        
+        ep_loss += loss.item() * GRAD_ACCUM
+        ep_batches += 1
+        pbar.set_postfix(loss=f"{loss.item()*GRAD_ACCUM:.4f}", lr=f"{sched.get_last_lr()[0]:.1e}", step=step)
+    
+    avg = ep_loss / max(ep_batches, 1)
+    print(f"Epoch {ep+1} loss: {avg:.4f}")
+    writer.add_scalar("train/epoch_loss", avg, ep + 1)
+    
+    if avg < best:
+        best = avg
+        best_path = os.path.join(CHECKPOINT_DIR, "best.pt")
+        weights_path = save_checkpoint(model, opt, sched, step, ep, avg, best_path)
+        
+        try:
+            api.upload_file(
+                path_or_fileobj=weights_path,
+                path_in_repo="model.safetensors",
+                repo_id=HF_REPO,
+                commit_message=f"Best model (epoch {ep+1}, loss {avg:.4f})",
+            )
+            print(f"  ✓ Uploaded best to {HF_REPO}")
+        except Exception as e:
+            print(f"  ⚠ Upload failed: {e}")
+
+# ============================================================================
+# FINAL
+# ============================================================================
+print("\nSaving final model...")
+final_path = os.path.join(CHECKPOINT_DIR, "final.pt")
+weights_path = save_checkpoint(model, opt, sched, step, EPOCHS, best, final_path)
+
+# Final samples
+print("Generating final samples...")
+images = generate_samples(model, SAMPLE_PROMPTS, num_steps=20)
+save_samples(images, SAMPLE_PROMPTS, step, SAMPLE_DIR)
+
+# Final upload
+try:
+    api.upload_file(path_or_fileobj=weights_path, path_in_repo="model.safetensors", repo_id=HF_REPO)
+    config_path = os.path.join(CHECKPOINT_DIR, "config.json")
+    with open(config_path, "w") as f:
+        json.dump(config.__dict__, f, indent=2)
+    api.upload_file(path_or_fileobj=config_path, path_in_repo="config.json", repo_id=HF_REPO)
+    api.upload_folder(folder_path=LOG_DIR, path_in_repo="logs", repo_id=HF_REPO)
+    api.upload_folder(folder_path=SAMPLE_DIR, path_in_repo="samples", repo_id=HF_REPO)
+    print(f"\n✓ Training complete! https://huggingface.co/{HF_REPO}")
+except Exception as e:
+    print(f"\n⚠ Final upload failed: {e}")
+
+writer.close()
+print(f"Best loss: {best:.4f}")