Create trainer.py

trainer.py

@@ -0,0 +1,458 @@
+"""
+    SD15 Flow-Matching trainer
+    Author: AbstractPhil
+
+    Loads the current format pt and ensures through multiple validations that the process is correct for training.
+
+    Trains flow matching for sd15.
+
+    License: MIT
+    If you use my work, a cite wouldnt hurt.
+
+"""
+
+import os
+import json
+import datetime
+from dataclasses import dataclass, asdict
+from tqdm.auto import tqdm
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn.functional as F
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import DataLoader
+
+import datasets
+from diffusers import UNet2DConditionModel
+from huggingface_hub import HfApi, create_repo, hf_hub_download
+
+
+@dataclass
+class TrainConfig:
+    output_dir: str = "./outputs"
+    model_repo: str = "AbstractPhil/sd15-flow-matching-try2"
+    checkpoint_filename: str = "sd15_flowmatch_david_weighted_2_e34.pt"
+    dataset_name: str = "AbstractPhil/sd15-latent-distillation-500k"
+    
+    # HuggingFace upload settings
+    hf_repo_id: str = "AbstractPhil/sd15-flow-lune"
+    upload_to_hub: bool = True
+    
+    seed: int = 42
+    batch_size: int = 16
+    base_lr: float = 2e-6
+    shift: float = 2.0
+    dropout: float = 0.1
+    
+    max_train_steps: int = 50_000
+    checkpointing_steps: int = 1000
+    num_workers: int = 0
+    
+    # VAE scaling factor - multiply raw latents
+    vae_scale: float = 0.18215
+
+
+def load_student_unet(repo_id: str, filename: str, device="cuda") -> UNet2DConditionModel:
+    """Load UNet from .pt checkpoint containing student state_dict"""
+    # Download checkpoint from HuggingFace
+    print(f"Downloading checkpoint from {repo_id}/{filename}...")
+    checkpoint_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        repo_type="model"
+    )
+    print(f"✓ Downloaded to: {checkpoint_path}")
+    
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    
+    # Initialize UNet with SD1.5 config in fp32
+    print("Loading SD1.5 UNet architecture...")
+    unet = UNet2DConditionModel.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        subfolder="unet",
+        torch_dtype=torch.float32
+    )
+    
+    # Get original state for comparison
+    original_state_dict = {k: v.clone() for k, v in unet.state_dict().items()}
+    
+    # Load student weights and strip "unet." prefix
+    student_state_dict = checkpoint["student"]
+    
+    # Strip prefix if present
+    cleaned_student_dict = {}
+    for key, value in student_state_dict.items():
+        if key.startswith("unet."):
+            cleaned_key = key[5:]  # Remove "unet." prefix
+            cleaned_student_dict[cleaned_key] = value
+        else:
+            cleaned_student_dict[key] = value
+    
+    print(f"\n{'='*70}")
+    print("WEIGHT VERIFICATION")
+    print(f"{'='*70}")
+    
+    # 1. Compare keys
+    original_keys = set(original_state_dict.keys())
+    student_keys = set(cleaned_student_dict.keys())
+    
+    matching_keys = original_keys & student_keys
+    
+    print(f"Original UNet keys: {len(original_keys)}")
+    print(f"Student checkpoint keys: {len(student_keys)}")
+    print(f"Matching keys: {len(matching_keys)}")
+    
+    # 2. Compare student weights vs original BEFORE loading
+    total_params = 0
+    different_params = 0
+    mean_diff_sum = 0.0
+    max_diff = 0.0
+    
+    for key in matching_keys:
+        if key not in original_state_dict or key not in cleaned_student_dict:
+            continue
+            
+        orig = original_state_dict[key]
+        student = cleaned_student_dict[key].float()  # Convert to fp32 for comparison
+        
+        if orig.shape != student.shape:
+            print(f"⚠ Shape mismatch for {key}: {orig.shape} vs {student.shape}")
+            continue
+        
+        total_params += orig.numel()
+        
+        # Check if weights are different
+        diff = (orig - student).abs()
+        if diff.max() > 1e-6:
+            different_params += orig.numel()
+            mean_diff_sum += diff.sum().item()
+            max_diff = max(max_diff, diff.max().item())
+    
+    pct_different = (different_params / total_params * 100) if total_params > 0 else 0
+    avg_diff = mean_diff_sum / different_params if different_params > 0 else 0
+    
+    print(f"\nStudent vs Original (BEFORE loading):")
+    print(f"  Total parameters: {total_params:,}")
+    print(f"  Parameters different: {different_params:,} ({pct_different:.1f}%)")
+    print(f"  Average difference: {avg_diff:.6f}")
+    print(f"  Max difference: {max_diff:.6f}")
+    
+    # 3. Load weights
+    load_result = unet.load_state_dict(cleaned_student_dict, strict=False)
+    
+    if load_result.missing_keys:
+        print(f"\n⚠ Missing keys during load: {len(load_result.missing_keys)}")
+        for key in load_result.missing_keys[:3]:
+            print(f"  - {key}")
+    
+    if load_result.unexpected_keys:
+        print(f"⚠ Unexpected keys during load: {len(load_result.unexpected_keys)}")
+        for key in load_result.unexpected_keys[:3]:
+            print(f"  - {key}")
+    
+    # 4. Verify weights actually changed after loading
+    loaded_state_dict = unet.state_dict()
+    
+    total_params_after = 0
+    changed_params = 0
+    mean_diff_after = 0.0
+    max_diff_after = 0.0
+    
+    for key in matching_keys:
+        if key not in original_state_dict or key not in loaded_state_dict:
+            continue
+            
+        orig = original_state_dict[key]
+        loaded = loaded_state_dict[key]
+        
+        total_params_after += orig.numel()
+        
+        diff = (orig - loaded).abs()
+        if diff.max() > 1e-6:
+            changed_params += orig.numel()
+            mean_diff_after += diff.sum().item()
+            max_diff_after = max(max_diff_after, diff.max().item())
+    
+    pct_changed = (changed_params / total_params_after * 100) if total_params_after > 0 else 0
+    avg_diff_after = mean_diff_after / changed_params if changed_params > 0 else 0
+    
+    print(f"\nOriginal vs Loaded (AFTER loading):")
+    print(f"  Parameters changed: {changed_params:,} ({pct_changed:.1f}%)")
+    print(f"  Average difference: {avg_diff_after:.6f}")
+    print(f"  Max difference: {max_diff_after:.6f}")
+    
+    print(f"\n{'='*70}")
+    # Verification checks
+    if pct_different < 50:
+        print(f"⚠️ WARNING: Student weights only {pct_different:.1f}% different from base!")
+        print("   This checkpoint may not be trained.")
+    elif pct_changed < 90:
+        print(f"⚠️ WARNING: Only {pct_changed:.1f}% of weights changed after loading!")
+        print("   The load may have failed.")
+    else:
+        print(f"✅ Weights loaded successfully!")
+        print(f"   Checkpoint step: {checkpoint.get('gstep', 'unknown')}")
+        print(f"   {pct_different:.1f}% of weights differ from base SD1.5")
+    
+    print(f"{'='*70}\n")
+    
+    return unet.to(device)
+
+
+def train(config: TrainConfig):
+    device = "cuda"
+    torch.backends.cuda.matmul.allow_tf32 = True
+    
+    torch.manual_seed(config.seed)
+    torch.cuda.manual_seed(config.seed)
+    
+    # Setup output directory
+    date_time = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    real_output_dir = os.path.join(config.output_dir, date_time)
+    os.makedirs(real_output_dir, exist_ok=True)
+    t_writer = SummaryWriter(log_dir=real_output_dir, flush_secs=60)
+    
+    # Initialize HuggingFace API
+    hf_api = None
+    if config.upload_to_hub:
+        try:
+            hf_api = HfApi()
+            create_repo(
+                repo_id=config.hf_repo_id,
+                repo_type="model",
+                exist_ok=True,
+                private=False
+            )
+            print(f"✓ HuggingFace repo ready: {config.hf_repo_id}")
+        except Exception as e:
+            print(f"⚠ Hub upload disabled: {e}")
+            config.upload_to_hub = False
+    
+    # Save config locally and to hub
+    config_path = os.path.join(real_output_dir, "config.json")
+    with open(config_path, "w") as f:
+        json.dump(asdict(config), f, indent=2)
+    
+    if config.upload_to_hub:
+        hf_api.upload_file(
+            path_or_fileobj=config_path,
+            path_in_repo="config.json",
+            repo_id=config.hf_repo_id,
+            repo_type="model"
+        )
+    
+    # Load dataset in streaming mode
+    print(f"\nLoading dataset (streaming): {config.dataset_name}")
+    train_dataset = datasets.load_dataset(
+        config.dataset_name,
+        split="train",
+        streaming=True,
+        trust_remote_code=True
+    )
+    train_dataset = train_dataset.shuffle(seed=config.seed, buffer_size=1000)
+    print(f"✓ Dataset loaded in streaming mode")
+    
+    def collate_fn(examples):
+        # Latents are RAW from VAE - need to scale them
+        latents = torch.stack([torch.tensor(ex["latent"]) for ex in examples])
+        latents = latents * config.vae_scale  # Scale: ~[-6, 6] -> ~[-1, 1]
+        
+        clip_embeddings = torch.stack([torch.tensor(ex["clip_embedding"]) for ex in examples])
+        ids = [ex["id"] for ex in examples]
+        prompts = [ex["prompt"] for ex in examples]
+        
+        return latents, clip_embeddings, ids, prompts
+    
+    train_dataloader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        collate_fn=collate_fn,
+        num_workers=config.num_workers,
+    )
+    
+    # Verify first batch latent range (on GPU for speed)
+    print("\nVerifying latent scaling on first batch...")
+    first_batch = next(iter(train_dataloader))
+    latents_check, _, _, _ = first_batch
+    print(f"Raw latent range: [{latents_check.min():.3f}, {latents_check.max():.3f}]")
+    latents_check = latents_check.to(device)
+    print(f"After GPU transfer: [{latents_check.min():.3f}, {latents_check.max():.3f}]")
+    print(f"Expected: ~[-1, 1] for properly scaled latents")
+    del latents_check
+    
+    # Load pretrained student UNet
+    print(f"\nLoading model from HuggingFace...")
+    unet = load_student_unet(config.model_repo, config.checkpoint_filename, device=device)
+    unet.requires_grad_(True)
+    unet.enable_gradient_checkpointing()
+    unet.train()
+    
+    optimizer = torch.optim.Adam(
+        unet.parameters(),
+        lr=config.base_lr * (config.batch_size ** 0.5),
+    )
+    
+    global_step = 0
+    train_logs = {
+        "train_step": [],
+        "train_loss": [],
+        "train_timestep": [],
+        "trained_images": []
+    }
+    
+    def get_prediction(batch, log_to=None):
+        latents, encoder_hidden_states, ids, prompts = batch
+        
+        # Everything in fp32
+        latents = latents.to(dtype=torch.float32, device=device)
+        encoder_hidden_states = encoder_hidden_states.to(dtype=torch.float32, device=device)
+        
+        batch_size = latents.shape[0]
+        
+        # Apply dropout to conditioning for CFG support
+        dropout_mask = torch.rand(batch_size, device=device) < config.dropout
+        encoder_hidden_states = encoder_hidden_states.clone()
+        encoder_hidden_states[dropout_mask] = 0
+        
+        # Sample timesteps with shift
+        sigmas = torch.rand(batch_size, device=device)
+        sigmas = (config.shift * sigmas) / (1 + (config.shift - 1) * sigmas)
+        timesteps = sigmas * 1000
+        sigmas = sigmas[:, None, None, None]
+        
+        # Flow matching forward process
+        noise = torch.randn_like(latents)
+        noisy_latents = noise * sigmas + latents * (1 - sigmas)
+        target = noise - latents
+        
+        # Predict velocity
+        pred = unet(noisy_latents, timesteps, encoder_hidden_states, return_dict=False)[0]
+        
+        loss = F.mse_loss(pred, target, reduction="none")
+        loss = loss.mean(dim=list(range(1, len(loss.shape))))
+        
+        if log_to is not None:
+            for i in range(batch_size):
+                log_to["train_step"].append(global_step)
+                log_to["train_loss"].append(loss[i].item())
+                log_to["train_timestep"].append(timesteps[i].item())
+                log_to["trained_images"].append({
+                    "step": global_step,
+                    "id": ids[i],
+                    "prompt": prompts[i]
+                })
+        
+        return loss.mean()
+    
+    def plot_logs(log_dict):
+        plt.figure(figsize=(10, 6))
+        plt.scatter(
+            log_dict["train_timestep"],
+            log_dict["train_loss"],
+            s=3,
+            c=log_dict["train_step"],
+            marker=".",
+            cmap='cool'
+        )
+        plt.xlabel("timestep")
+        plt.ylabel("loss")
+        plt.yscale("log")
+        plt.colorbar(label="step")
+    
+    def save_checkpoint(step):
+        checkpoint_path = os.path.join(real_output_dir, f"checkpoint-{step:08}")
+        os.makedirs(checkpoint_path, exist_ok=True)
+        
+        # Save UNet weights as diffusers format
+        unet.save_pretrained(
+            os.path.join(checkpoint_path, "unet"),
+            safe_serialization=True
+        )
+        
+        # Save complete checkpoint in .pt format
+        pt_filename = f"sd15_flow_lune_e{step//1000}_s{step}.pt"
+        pt_path = os.path.join(checkpoint_path, pt_filename)
+        
+        torch.save({
+            "cfg": asdict(config),
+            "student": unet.state_dict(),
+            "opt": optimizer.state_dict(),
+            "gstep": step
+        }, pt_path)
+        
+        # Save training metadata
+        metadata = {
+            "step": step,
+            "trained_images": train_logs["trained_images"]
+        }
+        metadata_path = os.path.join(checkpoint_path, "trained_images.json")
+        with open(metadata_path, "w") as f:
+            json.dump(metadata, f, indent=2)
+        
+        print(f"✓ Checkpoint saved at step {step}")
+        
+        # Upload to HuggingFace Hub
+        if config.upload_to_hub and hf_api is not None:
+            try:
+                hf_api.upload_file(
+                    path_or_fileobj=pt_path,
+                    path_in_repo=pt_filename,
+                    repo_id=config.hf_repo_id,
+                    repo_type="model"
+                )
+                
+                hf_api.upload_folder(
+                    folder_path=os.path.join(checkpoint_path, "unet"),
+                    path_in_repo=f"checkpoint-{step:08}/unet",
+                    repo_id=config.hf_repo_id,
+                    repo_type="model"
+                )
+                
+                hf_api.upload_file(
+                    path_or_fileobj=metadata_path,
+                    path_in_repo=f"checkpoint-{step:08}/trained_images.json",
+                    repo_id=config.hf_repo_id,
+                    repo_type="model"
+                )
+                
+                print(f"✓ Uploaded to hub: {config.hf_repo_id}")
+            except Exception as e:
+                print(f"⚠ Upload failed: {e}")
+    
+    print("\nStarting training...")
+    progress_bar = tqdm(range(0, config.max_train_steps))
+    
+    for batch in train_dataloader:
+        loss = get_prediction(batch, log_to=train_logs)
+        t_writer.add_scalar("train/loss", loss.detach().item(), global_step)
+        
+        loss.backward()
+        
+        grad_norm = torch.nn.utils.clip_grad_norm_(unet.parameters(), 2.0)
+        t_writer.add_scalar("train/grad_norm", grad_norm.detach().item(), global_step)
+        
+        optimizer.step()
+        optimizer.zero_grad()
+        
+        progress_bar.update(1)
+        progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})
+        global_step += 1
+        
+        if global_step % 100 == 0:
+            plot_logs(train_logs)
+            t_writer.add_figure("train_loss", plt.gcf(), global_step)
+            plt.close()
+        
+        if global_step % config.checkpointing_steps == 0:
+            save_checkpoint(global_step)
+        
+        if global_step >= config.max_train_steps:
+            save_checkpoint(global_step)
+            print("\n✅ Training complete!")
+            return
+
+
+if __name__ == "__main__":
+    config = TrainConfig()
+    train(config)