Upload sdxl/train_sdxl_adapter.py with huggingface_hub

sdxl/train_sdxl_adapter.py

@@ -0,0 +1,263 @@
+"""
+Train SDXL adapter: Qwen3-4B activations -> SDXL prompt embeddings.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import json
+import os
+from pathlib import Path
+from datetime import datetime
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+from sdxl_adapter import SDXLCrossAttentionAdapter, LayerWeightedInput
+
+
+class ActivationEmbeddingDataset(Dataset):
+    """Dataset of (Qwen activation, SDXL embedding) pairs."""
+    def __init__(self, activation_dir, embedding_dir, metadata_path):
+        with open(metadata_path) as f:
+            self.metadata = json.load(f)
+        self.activation_dir = Path(activation_dir)
+        self.embedding_dir = Path(embedding_dir)
+
+        # Validate all files exist
+        valid = []
+        for item in self.metadata:
+            emotion = item['emotion']
+            idx = item['index']
+            act_file = self.activation_dir / f"{emotion}_{idx:02d}.npy"
+            emb_file = self.embedding_dir / f"{emotion}_{idx:02d}.npz"
+            if act_file.exists() and emb_file.exists():
+                valid.append(item)
+        self.metadata = valid
+        print(f"Dataset: {len(self.metadata)} valid pairs")
+
+    def __len__(self):
+        return len(self.metadata)
+
+    def __getitem__(self, idx):
+        item = self.metadata[idx]
+        emotion = item['emotion']
+        i = item['index']
+
+        # Load Qwen activation [7680]
+        act = np.load(self.activation_dir / f"{emotion}_{i:02d}.npy")
+        act = torch.from_numpy(act).float()
+
+        # Load SDXL embeddings
+        emb = np.load(self.embedding_dir / f"{emotion}_{i:02d}.npz")
+        prompt_embeds = torch.from_numpy(emb['prompt_embeds']).float().squeeze(0)  # [77, 2048]
+        pooled_embeds = torch.from_numpy(emb['pooled_prompt_embeds']).float().squeeze(0)  # [1280]
+
+        return act, prompt_embeds, pooled_embeds, emotion
+
+
+def compute_normalization(dataset):
+    """Compute mean/std of activations and targets for normalization."""
+    all_acts = []
+    all_main = []
+    all_pooled = []
+    for i in range(len(dataset)):
+        act, main, pooled, _ = dataset[i]
+        all_acts.append(act)
+        all_main.append(main)
+        all_pooled.append(pooled)
+
+    acts = torch.stack(all_acts)
+    mains = torch.stack(all_main)
+    pooleds = torch.stack(all_pooled)
+
+    return {
+        'act_mean': acts.mean(dim=0),
+        'act_std': acts.std(dim=0).clamp(min=1e-6),
+        'main_mean': mains.mean(dim=(0, 1)),  # [2048] - mean across batch and tokens
+        'main_std': mains.std(dim=(0, 1)).clamp(min=1e-6),
+        'pooled_mean': pooleds.mean(dim=0),  # [1280]
+        'pooled_std': pooleds.std(dim=0).clamp(min=1e-6),
+    }
+
+
+def train():
+    device = torch.device('cuda')
+    base_dir = Path('/home/beta1/milady-training')
+
+    print(f"[{datetime.now()}] Loading dataset...")
+    dataset = ActivationEmbeddingDataset(
+        activation_dir=base_dir / 'qwen_activations',
+        embedding_dir=base_dir / 'sdxl_embeddings',
+        metadata_path=base_dir / 'sdxl_emotions' / 'metadata.json',
+    )
+
+    # Compute normalization stats
+    print(f"[{datetime.now()}] Computing normalization statistics...")
+    norm = compute_normalization(dataset)
+    for k, v in norm.items():
+        print(f"  {k}: shape={v.shape}, mean={v.mean():.4f}, std={v.std():.4f}")
+
+    # Train/val split (80/20)
+    n = len(dataset)
+    n_val = max(1, n // 5)
+    n_train = n - n_val
+    train_ds, val_ds = torch.utils.data.random_split(
+        dataset, [n_train, n_val],
+        generator=torch.Generator().manual_seed(42)
+    )
+    print(f"Train: {n_train}, Val: {n_val}")
+
+    train_loader = DataLoader(train_ds, batch_size=16, shuffle=True, drop_last=True)
+    val_loader = DataLoader(val_ds, batch_size=16, shuffle=False)
+
+    # Models
+    layer_weight = LayerWeightedInput(n_layers=3, layer_dim=2560).to(device)
+    adapter = SDXLCrossAttentionAdapter(
+        in_dim=2560, rank=256, n_input_tokens=8,
+        n_heads=8, n_layers=3,
+        n_output_tokens=77, main_dim=2048, pooled_dim=1280,
+    ).to(device)
+
+    print(f"LayerWeight params: {sum(p.numel() for p in layer_weight.parameters()):,}")
+    print(f"Adapter params: {sum(p.numel() for p in adapter.parameters()):,}")
+
+    # Move norm stats to device
+    act_mean = norm['act_mean'].to(device)
+    act_std = norm['act_std'].to(device)
+
+    # Optimizer
+    params = list(adapter.parameters()) + list(layer_weight.parameters())
+    optimizer = torch.optim.AdamW(params, lr=1e-4, weight_decay=1e-2)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=500, eta_min=1e-6)
+
+    # Training loop
+    n_epochs = 500
+    best_val_loss = float('inf')
+    patience = 50
+    patience_counter = 0
+    save_dir = base_dir / 'sdxl_adapter_checkpoints'
+    save_dir.mkdir(exist_ok=True)
+
+    print(f"\n[{datetime.now()}] Starting training for {n_epochs} epochs...")
+
+    for epoch in range(n_epochs):
+        # Train
+        adapter.train()
+        layer_weight.train()
+        train_loss_sum = 0
+        train_main_loss_sum = 0
+        train_pooled_loss_sum = 0
+        n_batches = 0
+
+        for acts, target_main, target_pooled, _ in train_loader:
+            acts = acts.to(device)
+            target_main = target_main.to(device)
+            target_pooled = target_pooled.to(device)
+
+            # Normalize activations
+            acts_norm = (acts - act_mean) / act_std
+
+            # Layer weighting
+            x = layer_weight(acts_norm)  # [B, 2560]
+
+            # Forward
+            pred_main, pred_pooled = adapter(x)  # [B, 77, 2048], [B, 1280]
+
+            # Losses
+            main_loss = F.mse_loss(pred_main, target_main)
+            pooled_loss = F.mse_loss(pred_pooled, target_pooled)
+            loss = main_loss + 0.5 * pooled_loss
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(params, 1.0)
+            optimizer.step()
+
+            train_loss_sum += loss.item()
+            train_main_loss_sum += main_loss.item()
+            train_pooled_loss_sum += pooled_loss.item()
+            n_batches += 1
+
+        scheduler.step()
+
+        # Validate
+        adapter.eval()
+        layer_weight.eval()
+        val_loss_sum = 0
+        val_main_sum = 0
+        val_pooled_sum = 0
+        v_batches = 0
+
+        with torch.no_grad():
+            for acts, target_main, target_pooled, _ in val_loader:
+                acts = acts.to(device)
+                target_main = target_main.to(device)
+                target_pooled = target_pooled.to(device)
+
+                acts_norm = (acts - act_mean) / act_std
+                x = layer_weight(acts_norm)
+                pred_main, pred_pooled = adapter(x)
+
+                main_loss = F.mse_loss(pred_main, target_main)
+                pooled_loss = F.mse_loss(pred_pooled, target_pooled)
+                loss = main_loss + 0.5 * pooled_loss
+
+                val_loss_sum += loss.item()
+                val_main_sum += main_loss.item()
+                val_pooled_sum += pooled_loss.item()
+                v_batches += 1
+
+        avg_train = train_loss_sum / max(n_batches, 1)
+        avg_val = val_loss_sum / max(v_batches, 1)
+        avg_train_main = train_main_loss_sum / max(n_batches, 1)
+        avg_train_pooled = train_pooled_loss_sum / max(n_batches, 1)
+        avg_val_main = val_main_sum / max(v_batches, 1)
+        avg_val_pooled = val_pooled_sum / max(v_batches, 1)
+
+        if (epoch + 1) % 10 == 0 or epoch == 0:
+            lr = scheduler.get_last_lr()[0]
+            print(f"  Epoch {epoch+1:3d} | Train: {avg_train:.6f} (main={avg_train_main:.6f} pool={avg_train_pooled:.6f}) | Val: {avg_val:.6f} (main={avg_val_main:.6f} pool={avg_val_pooled:.6f}) | LR: {lr:.2e}")
+
+        # Early stopping & best model
+        if avg_val < best_val_loss:
+            best_val_loss = avg_val
+            patience_counter = 0
+            checkpoint = {
+                'epoch': epoch + 1,
+                'model_type': 'sdxl_cross_attention',
+                'adapter_state_dict': adapter.state_dict(),
+                'layer_weight_state_dict': layer_weight.state_dict(),
+                'in_dim': 2560,
+                'out_dim_main': 2048,
+                'out_dim_pooled': 1280,
+                'n_tokens': 77,
+                'rank': 256,
+                'n_input_tokens': 8,
+                'n_heads': 8,
+                'n_layers': 3,
+                'input_layers': 'learned_weight',
+                'hook_layers': [9, 18, 27],
+                'act_mean': act_mean.cpu(),
+                'act_std': act_std.cpu(),
+                'train_loss': avg_train,
+                'val_loss': avg_val,
+            }
+            torch.save(checkpoint, save_dir / 'best_sdxl_adapter.pt')
+        else:
+            patience_counter += 1
+            if patience_counter >= patience:
+                print(f"  Early stopping at epoch {epoch+1}")
+                break
+
+        # Periodic save
+        if (epoch + 1) % 100 == 0:
+            torch.save(checkpoint, save_dir / f'sdxl_adapter_epoch{epoch+1}.pt')
+
+    print(f"\n[{datetime.now()}] Training complete!")
+    print(f"Best val loss: {best_val_loss:.6f}")
+    print(f"Best model saved to: {save_dir / 'best_sdxl_adapter.pt'}")
+
+
+if __name__ == '__main__':
+    train()