Create flow_leco_trainer.py

flow_leco_trainer.py

@@ -0,0 +1,530 @@
+"""
+Lune LECO Trainer - Fixed
+"""
+
+import os
+import json
+import datetime
+from dataclasses import dataclass, asdict, field
+from typing import List, Literal
+from tqdm.auto import tqdm
+from enum import Enum
+
+import torch
+import torch.nn.functional as F
+from torch.utils.tensorboard import SummaryWriter
+from safetensors.torch import save_file
+
+from diffusers import UNet2DConditionModel
+from transformers import CLIPTextModel, CLIPTokenizer
+from huggingface_hub import hf_hub_download
+
+
+class ActionType(str, Enum):
+    """LECO action types"""
+    ERASE = "erase"
+    ENHANCE = "enhance"
+    REPLACE = "replace"
+    SUPPRESS = "suppress"
+
+
+@dataclass
+class ConceptPair:
+    """
+    Single concept transformation pair.
+    
+    The LoRA learns: pred(concept) - pred(anchor)
+    
+    Examples:
+        Erase: ConceptPair("anime style", "")
+        Enhance: ConceptPair("masterpiece", "")
+        Replace: ConceptPair("van gogh", "monet")
+        Suppress: ConceptPair("nsfw", "sfw")
+    """
+    concept: str
+    anchor: str = ""
+    weight: float = 1.0
+    inference_weight: float = -1.0
+
+
+@dataclass
+class PreservationSet:
+    """Prompts that should remain unchanged"""
+    prompts: List[str] = field(default_factory=list)
+    weight: float = 0.3
+
+
+@dataclass
+class LECOConfig:
+    # Model paths
+    output_dir: str = "./leco_outputs"
+    base_model_repo: str = "AbstractPhil/sd15-flow-lune-flux"
+    base_checkpoint: str = "sd15_flow_flux_t2_6_pose_t4_6_port_t1_4_s18765.pt"
+    
+    # HuggingFace
+    hf_repo_id: str = "AbstractPhil/lune-leco-adapters"
+    upload_to_hub: bool = False
+    
+    # Training data
+    action: ActionType = ActionType.ERASE
+    concept_pairs: List[ConceptPair] = field(default_factory=list)
+    preservation: PreservationSet = field(default_factory=PreservationSet)
+    
+    # LoRA architecture
+    lora_rank: int = 4
+    lora_alpha: float = 1.0
+    lora_dropout: float = 0.0
+    training_method: Literal["full", "selfattn", "xattn", "noxattn", "innoxattn"] = "full"
+    
+    # Training
+    seed: int = 42
+    iterations: int = 1000
+    lr: float = 1e-4
+    pairs_per_step: int = 1
+    
+    # Flow-matching
+    shift: float = 2.5
+    min_timestep: float = 0.0
+    max_timestep: float = 1000.0
+    
+    # Resolution
+    resolution: int = 512
+
+
+def get_target_modules(training_method: str) -> List[str]:
+    """Get layer names for LoRA injection"""
+    attn1 = ["attn1.to_q", "attn1.to_k", "attn1.to_v", "attn1.to_out.0"]
+    attn2 = ["attn2.to_q", "attn2.to_k", "attn2.to_v", "attn2.to_out.0"]
+    
+    method_map = {
+        "full": attn1 + attn2,
+        "selfattn": attn1,
+        "xattn": attn2,
+        "noxattn": attn1,
+        "innoxattn": attn2
+    }
+    return method_map.get(training_method, attn1 + attn2)
+
+
+def create_lora_layers(unet: torch.nn.Module, config: LECOConfig):
+    """Create LoRA layers in ComfyUI/A1111 format"""
+    target_modules = get_target_modules(config.training_method)
+    lora_state = {}
+    trainable_params = []
+    
+    def get_lora_key(module_path: str) -> str:
+        return f"lora_unet_{module_path.replace('.', '_')}"
+    
+    print(f"Creating LoRA layers (method: {config.training_method})...")
+    layer_count = 0
+    
+    for name, module in unet.named_modules():
+        if not any(target in name for target in target_modules):
+            continue
+        
+        if not isinstance(module, torch.nn.Linear):
+            continue
+        
+        lora_key = get_lora_key(name)
+        in_dim = module.in_features
+        out_dim = module.out_features
+        rank = config.lora_rank
+        
+        # LoRA matrices
+        # down: [rank, in_features]
+        # up: [out_features, rank]
+        lora_down = torch.nn.Parameter(torch.zeros(rank, in_dim))
+        lora_up = torch.nn.Parameter(torch.zeros(out_dim, rank))
+        
+        torch.nn.init.kaiming_uniform_(lora_down, a=1.0)
+        torch.nn.init.zeros_(lora_up)
+        
+        lora_state[f"{lora_key}.lora_down.weight"] = lora_down
+        lora_state[f"{lora_key}.lora_up.weight"] = lora_up
+        lora_state[f"{lora_key}.alpha"] = torch.tensor(config.lora_alpha)
+        lora_state[f"{lora_key}._module"] = module
+        
+        trainable_params.extend([lora_down, lora_up])
+        layer_count += 1
+    
+    print(f"✓ Created {layer_count} LoRA layers ({len(trainable_params)} parameters)")
+    return lora_state, trainable_params
+
+
+def apply_lora_hooks(unet: torch.nn.Module, lora_state: dict, scale: float = 1.0) -> list:
+    """
+    Apply LoRA using forward hooks.
+    
+    LoRA computation: out = out + scale * (x @ down.T @ up.T)
+    Using F.linear: F.linear(x, W) computes x @ W.T
+    
+    So: F.linear(F.linear(x, down), up) gives x @ down.T @ up.T ✓
+    """
+    handles = []
+    
+    for key in lora_state:
+        if not key.endswith(".lora_down.weight"):
+            continue
+        
+        base_key = key.replace(".lora_down.weight", "")
+        module = lora_state[f"{base_key}._module"]
+        lora_down = lora_state[f"{base_key}.lora_down.weight"]
+        lora_up = lora_state[f"{base_key}.lora_up.weight"]
+        alpha = lora_state[f"{base_key}.alpha"].item()
+        rank = lora_down.shape[0]
+        
+        scaling = (alpha / rank) * scale
+        
+        def make_hook(down, up, s):
+            def forward_hook(mod, inp, out):
+                x = inp[0]
+                # F.linear handles transpose internally
+                # down is [rank, in_features], F.linear does x @ down.T
+                # up is [out_features, rank], F.linear does result @ up.T
+                lora_out = F.linear(F.linear(x, down), up)
+                return out + lora_out * s
+            return forward_hook
+        
+        handle = module.register_forward_hook(make_hook(lora_down, lora_up, scaling))
+        handles.append(handle)
+    
+    return handles
+
+
+def remove_lora_hooks(handles: list):
+    """Remove all LoRA hooks"""
+    for handle in handles:
+        handle.remove()
+
+
+@torch.no_grad()
+def encode_text(prompt: str, tokenizer, text_encoder, device) -> torch.Tensor:
+    """Encode text to CLIP embeddings"""
+    tokens = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        truncation=True,
+        return_tensors="pt"
+    ).input_ids.to(device)
+    
+    return text_encoder(tokens)[0]
+
+
+def compute_leco_loss(
+    unet: torch.nn.Module,
+    lora_state: dict,
+    pair: ConceptPair,
+    tokenizer,
+    text_encoder,
+    config: LECOConfig,
+    device: str = "cuda"
+):
+    """
+    Compute LECO loss for a concept pair.
+    
+    Trains LoRA to reproduce: pred(concept) - pred(anchor)
+    """
+    # Sample timestep
+    min_sigma = config.min_timestep / 1000.0
+    max_sigma = config.max_timestep / 1000.0
+    sigma = min_sigma + torch.rand(1, device=device) * (max_sigma - min_sigma)
+    sigma = (config.shift * sigma) / (1 + (config.shift - 1) * sigma)
+    timestep = sigma * 1000.0
+    sigma = sigma.view(1, 1, 1, 1)
+    
+    # Random noise
+    noise = torch.randn(1, 4, config.resolution // 8, config.resolution // 8, device=device)
+    noisy_input = sigma * noise
+    
+    # Encode prompts
+    concept_emb = encode_text(pair.concept, tokenizer, text_encoder, device)
+    anchor_emb = encode_text(pair.anchor, tokenizer, text_encoder, device)
+    
+    # Compute target direction (without LoRA)
+    with torch.no_grad():
+        pred_concept = unet(
+            noisy_input, timestep,
+            encoder_hidden_states=concept_emb,
+            return_dict=False
+        )[0]
+        
+        pred_anchor = unet(
+            noisy_input, timestep,
+            encoder_hidden_states=anchor_emb,
+            return_dict=False
+        )[0]
+        
+        target_delta = pred_concept - pred_anchor
+    
+    # Apply LoRA and get its contribution
+    handles = apply_lora_hooks(unet, lora_state, scale=1.0)
+    
+    try:
+        pred_with_lora = unet(
+            noisy_input, timestep,
+            encoder_hidden_states=concept_emb,
+            return_dict=False
+        )[0]
+        
+        lora_delta = pred_with_lora - pred_concept
+        loss = F.mse_loss(lora_delta, target_delta)
+        
+    finally:
+        remove_lora_hooks(handles)
+    
+    return loss, {
+        "timestep": timestep.item(),
+        "sigma": sigma.item(),
+        "target_norm": target_delta.norm().item(),
+        "lora_norm": lora_delta.norm().item()
+    }
+
+
+def compute_preservation_loss(
+    unet: torch.nn.Module,
+    lora_state: dict,
+    preservation: PreservationSet,
+    tokenizer,
+    text_encoder,
+    config: LECOConfig,
+    device: str = "cuda"
+):
+    """Penalize LoRA changes to preservation prompts"""
+    if not preservation.prompts:
+        return 0.0, {}
+    
+    min_sigma = config.min_timestep / 1000.0
+    max_sigma = config.max_timestep / 1000.0
+    sigma = min_sigma + torch.rand(1, device=device) * (max_sigma - min_sigma)
+    sigma = (config.shift * sigma) / (1 + (config.shift - 1) * sigma)
+    timestep = sigma * 1000.0
+    sigma = sigma.view(1, 1, 1, 1)
+    
+    total_loss = 0
+    
+    for prompt in preservation.prompts:
+        noise = torch.randn(1, 4, config.resolution // 8, config.resolution // 8, device=device)
+        noisy_input = sigma * noise
+        prompt_emb = encode_text(prompt, tokenizer, text_encoder, device)
+        
+        with torch.no_grad():
+            pred_base = unet(
+                noisy_input, timestep,
+                encoder_hidden_states=prompt_emb,
+                return_dict=False
+            )[0]
+        
+        handles = apply_lora_hooks(unet, lora_state, scale=1.0)
+        
+        try:
+            pred_with_lora = unet(
+                noisy_input, timestep,
+                encoder_hidden_states=prompt_emb,
+                return_dict=False
+            )[0]
+        finally:
+            remove_lora_hooks(handles)
+        
+        total_loss += F.mse_loss(pred_with_lora, pred_base)
+    
+    avg_loss = total_loss / len(preservation.prompts)
+    return avg_loss, {"count": len(preservation.prompts), "avg": avg_loss.item()}
+
+
+def train_leco(config: LECOConfig):
+    """Main training loop"""
+    device = "cuda"
+    torch.manual_seed(config.seed)
+    
+    if not config.concept_pairs:
+        raise ValueError("No concept pairs specified!")
+    
+    # Setup output
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    concept_names = "_".join([
+        p.concept.replace(" ", "")[:12] 
+        for p in config.concept_pairs[:2]
+    ])
+    if len(config.concept_pairs) > 2:
+        concept_names += f"_plus{len(config.concept_pairs)-2}"
+    
+    run_name = f"{config.action.value}_{concept_names}_{timestamp}"
+    output_dir = os.path.join(config.output_dir, run_name)
+    os.makedirs(output_dir, exist_ok=True)
+    
+    writer = SummaryWriter(log_dir=output_dir, flush_secs=60)
+    
+    with open(os.path.join(output_dir, "config.json"), "w") as f:
+        json.dump(asdict(config), f, indent=2)
+    
+    print("="*80)
+    print(f"LECO Training: {config.action.value.upper()}")
+    print("="*80)
+    
+    # Load model
+    print("\nLoading base model...")
+    checkpoint_path = hf_hub_download(
+        repo_id=config.base_model_repo,
+        filename=config.base_checkpoint,
+        repo_type="model"
+    )
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    
+    unet = UNet2DConditionModel.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        subfolder="unet",
+        torch_dtype=torch.float32
+    )
+    
+    student_dict = checkpoint["student"]
+    cleaned_dict = {k[5:] if k.startswith("unet.") else k: v for k, v in student_dict.items()}
+    unet.load_state_dict(cleaned_dict, strict=False)
+    unet = unet.to(device)
+    unet.requires_grad_(False)
+    unet.eval()
+    print("✓ Loaded UNet")
+    
+    # Load CLIP
+    print("Loading CLIP text encoder...")
+    tokenizer = CLIPTokenizer.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", subfolder="tokenizer"
+    )
+    text_encoder = CLIPTextModel.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", subfolder="text_encoder",
+        torch_dtype=torch.float32
+    ).to(device)
+    text_encoder.requires_grad_(False)
+    text_encoder.eval()
+    print("✓ Loaded CLIP")
+    
+    # Create LoRA
+    print(f"\nInjecting LoRA (rank={config.lora_rank}, alpha={config.lora_alpha})...")
+    lora_state, trainable_params = create_lora_layers(unet, config)
+    
+    for key in lora_state:
+        if isinstance(lora_state[key], torch.Tensor):
+            lora_state[key] = lora_state[key].to(device)
+    
+    optimizer = torch.optim.AdamW(trainable_params, lr=config.lr, weight_decay=0.01)
+    
+    # Print config
+    print(f"\nTraining Configuration:")
+    print(f"  Action: {config.action.value}")
+    print(f"  Concept pairs: {len(config.concept_pairs)}")
+    for i, pair in enumerate(config.concept_pairs, 1):
+        anchor_str = f"→ '{pair.anchor}'" if pair.anchor else "(none)"
+        print(f"    {i}. '{pair.concept}' {anchor_str} (weight: {pair.weight})")
+    
+    if config.preservation.prompts:
+        print(f"  Preservation: {len(config.preservation.prompts)} prompts")
+    
+    print(f"\n  Iterations: {config.iterations}")
+    print(f"  Learning rate: {config.lr}")
+    print(f"  Training method: {config.training_method}")
+    print("="*80 + "\n")
+    
+    # Training loop
+    progress = tqdm(range(config.iterations), desc="Training")
+    
+    for step in progress:
+        import random
+        if config.pairs_per_step >= len(config.concept_pairs):
+            active_pairs = config.concept_pairs
+        else:
+            active_pairs = random.sample(config.concept_pairs, config.pairs_per_step)
+        
+        total_loss = 0
+        all_metrics = []
+        
+        for pair in active_pairs:
+            loss, metrics = compute_leco_loss(
+                unet, lora_state, pair,
+                tokenizer, text_encoder, config, device
+            )
+            total_loss += loss * pair.weight
+            all_metrics.append(metrics)
+        
+        if config.preservation.prompts:
+            pres_loss, pres_metrics = compute_preservation_loss(
+                unet, lora_state, config.preservation,
+                tokenizer, text_encoder, config, device
+            )
+            total_loss += pres_loss * config.preservation.weight
+        else:
+            pres_loss = 0
+        
+        total_loss.backward()
+        grad_norm = torch.nn.utils.clip_grad_norm_(trainable_params, max_norm=1.0)
+        optimizer.step()
+        optimizer.zero_grad()
+        
+        # Logging
+        writer.add_scalar("loss/total", total_loss.item(), step)
+        writer.add_scalar("loss/preservation", pres_loss if isinstance(pres_loss, (float, int)) else pres_loss.item(), step)
+        writer.add_scalar("grad_norm", grad_norm.item(), step)
+        
+        avg_target = sum(m["target_norm"] for m in all_metrics) / len(all_metrics)
+        progress.set_postfix({
+            "loss": f"{total_loss.item():.4f}",
+            "grad": f"{grad_norm.item():.3f}",
+            "target": f"{avg_target:.3f}"
+        })
+        
+        if (step + 1) % 200 == 0 or step == config.iterations - 1:
+            save_checkpoint(lora_state, config, output_dir, step + 1, concept_names)
+    
+    writer.close()
+    
+    print("\n" + "="*80)
+    print("✅ Training complete!")
+    print(f"Output: {output_dir}")
+    print("="*80)
+    
+    return output_dir
+
+
+def save_checkpoint(lora_state, config, output_dir, step, name_suffix):
+    """Save LoRA in SafeTensors format"""
+    save_dict = {}
+    
+    for key, value in lora_state.items():
+        if isinstance(value, torch.Tensor) and not key.endswith("._module"):
+            save_dict[key] = value.detach().cpu()
+    
+    concepts_str = ", ".join([p.concept for p in config.concept_pairs])
+    anchors_str = ", ".join([p.anchor for p in config.concept_pairs if p.anchor])
+    
+    metadata = {
+        "ss_network_module": "networks.lora",
+        "ss_network_dim": str(config.lora_rank),
+        "ss_network_alpha": str(config.lora_alpha),
+        "ss_base_model": "runwayml/stable-diffusion-v1-5",
+        "ss_training_method": config.training_method,
+        "leco_action": config.action.value,
+        "leco_concepts": concepts_str,
+        "leco_anchors": anchors_str,
+        "leco_step": str(step)
+    }
+    
+    filename = f"leco_{name_suffix}_r{config.lora_rank}_s{step}.safetensors"
+    filepath = os.path.join(output_dir, filename)
+    
+    save_file(save_dict, filepath, metadata=metadata)
+    print(f"\n✓ Saved: {filename}")
+
+
+if __name__ == "__main__":
+    config = LECOConfig(
+        action=ActionType.ENHANCE,
+        concept_pairs=[
+            ConceptPair("masterpiece", "", weight=1.0),
+            ConceptPair("best quality", "", weight=1.0),
+            ConceptPair("highly detailed", "", weight=0.8),
+        ],
+        iterations=600,
+        lora_rank=4,
+        training_method="selfattn"
+    )
+    
+    train_leco(config)