scripts/distill/dmd2_trainer.py

"""
DMD2 trainer for Anima (DMDR pattern + R3GAN + TSCD)
====================================================

DMDR paper の "2 モデル + LoRA on/off で real/fake 切替" 方式を踏襲。
ここに R3GAN discriminator と TSCD consistency loss を統合する。

Models:
  gen_model       : Anima DiT (full trainable, no LoRA)
  guidance_model  : Anima DiT (Q,V LoRA, scale 切替で real/fake)
  discriminator   : R3GAN projection D (text-conditional)
  ema_gen         : gen の EMA (TSCD target)

Update schedule per outer step:
  - guidance × N_GUIDANCE   (default 5)
  - generator × 1
  - discriminator × 1
"""
from __future__ import annotations
import copy
import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
import peft

from .anima_loader import AnimaBundle
from .r3gan_disc import R3GANDiscriminator, d_loss_r3gan, g_loss_r3gan


@dataclass
class DMD2Args:
    # Training
    batch_size: int = 4
    lr_gen: float = 2e-5
    lr_guidance: float = 2e-5
    lr_disc: float = 2e-4
    grad_clip: float = 1.0

    # DMD2 / DMDR
    guidance_updates: int = 5     # 1 outer step あたりの guidance update 回数
    lora_rank: int = 32
    lora_scale_f: float = 2.0     # fake-score (LoRA ON)
    lora_scale_r: float = 0.75    # real-score start (cosine decay → 0)
    cfg_r: float = 2.0            # real-score 側 CFG (Augmentation の Spear)
    dynamic_decay_steps: int = 2000  # lora_scale_r を cosine で 0 にする step

    # Timestep sampling
    gui_alpha: float = 4.0        # Beta(α, β) for continuous t
    gui_beta: float = 1.5
    num_inference_steps: int = 8  # gen の discrete grid (Phase A=8, B=4, C=2)

    # TSCD
    tscd_weight: float = 0.5
    ema_decay: float = 0.999

    # R3GAN
    adv_weight: float = 0.1
    r3gan_gamma: float = 50.0

    # CFG-Aug schedule
    cfg_aug_cold_steps: int = 200  # 最初は cfg_r=0 で wait


def attach_qv_lora(transformer: nn.Module, rank: int = 32) -> nn.Module:
    """attention Q/V projection に LoRA を attach。MiniTrainDIT の実 module 名を
    inspection で見つける。"""
    Q_KEYS = ("to_q", "q_proj", "wq", "q")
    V_KEYS = ("to_v", "v_proj", "wv", "v")
    target_modules = []
    for name, module in transformer.named_modules():
        if not isinstance(module, nn.Linear):
            continue
        leaf = name.rsplit(".", 1)[-1]
        if leaf in Q_KEYS or leaf in V_KEYS:
            target_modules.append(name)

    if not target_modules:
        # Fallback 1: 名前に 'attn' が入る Linear すべて
        for name, module in transformer.named_modules():
            if isinstance(module, nn.Linear) and (
                "attn" in name.lower() or "attention" in name.lower()
            ):
                target_modules.append(name)

    if not target_modules:
        print("[lora] WARN: no attention Q/V found, falling back to all-linear")
        target_modules = "all-linear"
    else:
        print(f"[lora] target_modules={len(target_modules)} layers (Q/V or attn)")

    cfg = peft.LoraConfig(
        r=rank,
        lora_alpha=rank,
        lora_dropout=0.0,
        bias="none",
        target_modules=target_modules,
    )
    return peft.get_peft_model(transformer, cfg)


def attach_wide_lora(transformer: nn.Module, rank: int = 32) -> nn.Module:
    """全 nn.Linear に LoRA を attach (LLM adapter 内部は除外)。
    Anima/Cosmos の AdaLN modulation・attention・MLP すべてを学習対象にする。
    timestep 解釈を含む蒸留タスクには Q,V LoRA では不足、これが必要。"""
    target_modules = []
    for name, module in transformer.named_modules():
        if not isinstance(module, nn.Linear):
            continue
        # LLM adapter (Qwen3→T5 bridge) は触らない (壊れやすい)
        if "llm_adapter" in name:
            continue
        target_modules.append(name)

    if not target_modules:
        raise RuntimeError("No nn.Linear found in transformer for wide LoRA")

    # 概要
    cats = {"attn": 0, "adaln": 0, "mlp": 0, "other": 0}
    for n in target_modules:
        nl = n.lower()
        if "attn" in nl or "attention" in nl:
            cats["attn"] += 1
        elif "adaln" in nl or "modulation" in nl:
            cats["adaln"] += 1
        elif "mlp" in nl or "ffn" in nl or "feed" in nl:
            cats["mlp"] += 1
        else:
            cats["other"] += 1
    print(f"[lora] wide target: total={len(target_modules)} "
          f"(attn={cats['attn']}, adaln={cats['adaln']}, mlp={cats['mlp']}, other={cats['other']})")

    cfg = peft.LoraConfig(
        r=rank,
        lora_alpha=rank,
        lora_dropout=0.0,
        bias="none",
        target_modules=target_modules,
    )
    return peft.get_peft_model(transformer, cfg)


def set_lora_scale(model: nn.Module, scale: float) -> None:
    """PEFT LoRA の scaling を runtime で書き換え。real/fake 切替用。"""
    for module in model.modules():
        if hasattr(module, "scaling") and isinstance(module.scaling, dict):
            for key in module.scaling:
                # alpha/r = 1.0 を基準に scale 倍する
                module.scaling[key] = scale


def sample_continuous_t(
    batch_size: int, alpha: float, beta: float, device: torch.device
) -> torch.Tensor:
    """logit-normal Beta(α, β) サンプリング。返り値 t ∈ (0, 1)"""
    # Beta -> logit-normal 風の重み (DMDR の sample_continue 簡略版)
    u = torch.rand(batch_size, device=device)
    # Beta(α, β) inverse CDF を近似:正規 logit に変換するシンプル版
    # ここでは log-normal で代用 (実装簡易のため)
    t = torch.distributions.Beta(alpha, beta).sample((batch_size,)).to(device)
    return t.clamp(1e-3, 1 - 1e-3)


def sample_discrete_t(
    batch_size: int, num_steps: int, device: torch.device
) -> torch.Tensor:
    """generator 側: 離散 grid {1/N, 2/N, ..., (N-1)/N} から uniform サンプル。"""
    grid = torch.linspace(1.0 / num_steps, 1.0 - 1.0 / num_steps, num_steps - 1, device=device)
    idx = torch.randint(0, num_steps - 1, (batch_size,), device=device)
    return grid[idx]


def cosine_decay(step: int, total: int) -> float:
    if step >= total:
        return 0.0
    return 0.5 * (1 + math.cos(math.pi * step / total))


class EMA:
    """生成器の EMA を別 model copy として保持。TSCD target として直接 forward 可能。

    only_trainable=True (default) なら trainable param のみ EMA 更新する。
    LoRA-only モードでは base は frozen なので大半が skip され、EMA cost が激減。"""

    def __init__(self, model: nn.Module, decay: float = 0.999,
                 only_trainable: bool = True):
        self.decay = decay
        self.only_trainable = only_trainable
        # 同 device/dtype の eval-only copy
        self.ema_model = copy.deepcopy(model).eval()
        for p in self.ema_model.parameters():
            p.requires_grad = False

    @torch.no_grad()
    def update(self, model: nn.Module) -> None:
        for ep, p in zip(self.ema_model.parameters(), model.parameters()):
            if self.only_trainable and not p.requires_grad:
                continue
            ep.mul_(self.decay).add_(p.detach(), alpha=1 - self.decay)
        # buffers (RoPE 等) は最新を採用
        for eb, b in zip(self.ema_model.buffers(), model.buffers()):
            eb.copy_(b.detach())

    def __call__(self, **kwargs):
        return self.ema_model(**kwargs)


# ============================================================================
# Trainer
# ============================================================================

class DMD2Trainer:
    def __init__(
        self,
        bundle: AnimaBundle,
        gen_transformer: nn.Module,
        guidance_transformer: nn.Module,  # PEFT-wrapped (LoRA attached)
        discriminator: R3GANDiscriminator,
        args: DMD2Args,
    ):
        self.bundle = bundle
        self.gen = gen_transformer
        self.guidance = guidance_transformer
        self.disc = discriminator
        self.args = args
        self.device = bundle.device
        self.step = 0

        # Optimizers (gen も frozen 部分があるなら filter)
        self.opt_gen = torch.optim.AdamW(
            [p for p in self.gen.parameters() if p.requires_grad],
            lr=args.lr_gen, betas=(0.9, 0.999),
            weight_decay=0.01, eps=1e-8,
        )
        self.opt_guidance = torch.optim.AdamW(
            [p for p in self.guidance.parameters() if p.requires_grad],
            lr=args.lr_guidance, betas=(0.9, 0.999), weight_decay=0.0, eps=1e-8,
        )
        self.opt_disc = torch.optim.Adam(
            self.disc.parameters(), lr=args.lr_disc, betas=(0.0, 0.99), eps=1e-8,
        )

        # EMA for TSCD
        self.ema = EMA(self.gen, decay=args.ema_decay)

    # -- helpers ------------------------------------------------------------
    def _set_train(self, gen: bool, guidance: bool, disc: bool):
        self.gen.train(gen); self.guidance.train(guidance); self.disc.train(disc)

    def _t_continuous(self, B: int) -> torch.Tensor:
        return sample_continuous_t(B, self.args.gui_alpha, self.args.gui_beta, self.device)

    def _t_discrete(self, B: int) -> torch.Tensor:
        return sample_discrete_t(B, self.args.num_inference_steps, self.device)

    def _cfg_r_current(self) -> float:
        """cold start 中は cfg_r=0、それ以降は args.cfg_r。"""
        return 0.0 if self.step < self.args.cfg_aug_cold_steps else self.args.cfg_r

    def _lora_scale_r_current(self) -> float:
        decay = cosine_decay(self.step, self.args.dynamic_decay_steps)
        return self.args.lora_scale_r * decay

    # -- guidance update ----------------------------------------------------
    def step_guidance(self, real_latents: torch.Tensor, cond: torch.Tensor) -> dict:
        """fake-score (LoRA ON) を v-prediction MSE で更新"""
        self._set_train(gen=False, guidance=True, disc=False)
        set_lora_scale(self.guidance, self.args.lora_scale_f)

        B = real_latents.size(0)
        noise = torch.randn_like(real_latents)
        t = self._t_continuous(B)
        noisy = AnimaBundle.add_noise(real_latents, noise, t)
        v_target = AnimaBundle.velocity_target(real_latents, noise)

        v_pred = AnimaBundle.dit_forward(self.guidance, noisy, t, cond)
        loss = F.mse_loss(v_pred, v_target)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            [p for p in self.guidance.parameters() if p.requires_grad],
            self.args.grad_clip)
        self.opt_guidance.step(); self.opt_guidance.zero_grad()
        return {"l_guidance": loss.detach()}

    # -- CFG forward (real-score side) -------------------------------------
    def _real_score_forward(
        self, noisy: torch.Tensor, t: torch.Tensor, cond: torch.Tensor, cfg: float,
    ) -> torch.Tensor:
        """LoRA OFF / decay 状態の guidance_model で real-side velocity を取得。
        CFG > 1 のときは null condition と blend。"""
        set_lora_scale(self.guidance, self._lora_scale_r_current())
        v_cond = AnimaBundle.dit_forward(self.guidance, noisy, t, cond)
        if cfg > 1.0:
            null_cond = torch.zeros_like(cond)
            v_uncond = AnimaBundle.dit_forward(self.guidance, noisy, t, null_cond)
            v = v_uncond + cfg * (v_cond - v_uncond)
        else:
            v = v_cond
        return v

    # -- TSCD consistency loss ---------------------------------------------
    def _tscd_loss(
        self, x0_pred: torch.Tensor, cond: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        """segment-wise consistency vs EMA gen。
        TSCD 簡易版: 同じ x0 ground 周辺で noisy 化 → EMA で x0 推定 → MSE。"""
        with torch.no_grad():
            noise = torch.randn_like(x0_pred)
            noisy = AnimaBundle.add_noise(x0_pred.detach(), noise, t)
            v_ema = AnimaBundle.dit_forward(self.ema.ema_model, noisy, t, cond)
            x0_ema = AnimaBundle.x0_from_velocity(noisy, v_ema, t)
        return F.mse_loss(x0_pred, x0_ema.detach())

    # -- generator update ---------------------------------------------------
    def step_generator(self, real_latents: torch.Tensor, cond: torch.Tensor) -> dict:
        self._set_train(gen=True, guidance=False, disc=False)
        B = real_latents.size(0)
        noise = torch.randn_like(real_latents)
        t = self._t_discrete(B)
        noisy = AnimaBundle.add_noise(real_latents, noise, t)

        # Generator forward
        v_gen = AnimaBundle.dit_forward(self.gen, noisy, t, cond)
        x0_pred = AnimaBundle.x0_from_velocity(noisy, v_gen, t)

        # === DMD gradient ===
        with torch.no_grad():
            cfg = self._cfg_r_current()
            v_real = self._real_score_forward(noisy, t, cond, cfg=cfg)
            set_lora_scale(self.guidance, self.args.lora_scale_f)
            v_fake = AnimaBundle.dit_forward(self.guidance, noisy, t, cond)
        x0_real = AnimaBundle.x0_from_velocity(noisy, v_real, t)
        x0_fake = AnimaBundle.x0_from_velocity(noisy, v_fake, t)
        p_real = x0_pred - x0_real
        p_fake = x0_pred - x0_fake
        denom = p_real.abs().mean(dim=list(range(1, p_real.dim())), keepdim=True) + 1e-8
        grad = (p_real - p_fake) / denom
        dmd_loss = 0.5 * F.mse_loss(x0_pred.float(), (x0_pred - grad).detach().float())

        # === TSCD consistency ===
        consist_loss = self._tscd_loss(x0_pred, cond, t)

        # === R3GAN adversarial (generator side) — adv_weight>0 のときだけ ===
        if self.args.adv_weight > 0:
            d_fake = self.disc(x0_pred, cond)
            with torch.no_grad():
                d_real_for_g = self.disc(real_latents, cond)
            adv_loss = g_loss_r3gan(d_real_for_g, d_fake)
        else:
            adv_loss = torch.zeros((), device=self.device)

        loss = (
            dmd_loss
            + self.args.tscd_weight * consist_loss
            + self.args.adv_weight * adv_loss
        )
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.gen.parameters(), self.args.grad_clip)
        self.opt_gen.step(); self.opt_gen.zero_grad()

        self.ema.update(self.gen)
        return {
            "l_dmd": dmd_loss.detach(),
            "l_consist": consist_loss.detach(),
            "l_adv_g": adv_loss.detach(),
            "loss_gen_total": loss.detach(),
        }

    # -- discriminator update ----------------------------------------------
    def step_discriminator(self, real_latents: torch.Tensor, cond: torch.Tensor) -> dict:
        self._set_train(gen=False, guidance=False, disc=True)
        B = real_latents.size(0)

        # 生成器で fake samples を作成 (no_grad)
        with torch.no_grad():
            noise = torch.randn_like(real_latents)
            t_init = torch.ones(B, device=self.device)
            v = AnimaBundle.dit_forward(self.gen, noise, t_init, cond)
            fake_latents = AnimaBundle.x0_from_velocity(noise, v, t_init)

        real_samples = real_latents.detach().clone().requires_grad_(True)
        fake_samples = fake_latents.detach().clone().requires_grad_(True)

        d_real = self.disc(real_samples, cond)
        d_fake = self.disc(fake_samples, cond)
        loss, metrics = d_loss_r3gan(
            d_real, d_fake, real_samples, fake_samples, gamma=self.args.r3gan_gamma,
        )

        # GUARD: R1/R2 が異常に大きくなったら abort
        # latent 空間 (16ch × 128×128 = 262k) では正常時でも R1/R2 ~ 10-100 になる。
        # Phase A は 700 -> 700k で爆発。1000 を超えたら abort。
        r1_val = float(metrics["d_r1"])
        r2_val = float(metrics["d_r2"])
        if r1_val > 1000.0 or r2_val > 1000.0:
            raise RuntimeError(
                f"R3GAN penalty exploded (r1={r1_val:.1f}, r2={r2_val:.1f}). "
                f"Reduce r3gan_gamma (current={self.args.r3gan_gamma}). "
                f"Aborting to prevent gen damage."
            )

        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.disc.parameters(), self.args.grad_clip)
        self.opt_disc.step(); self.opt_disc.zero_grad()
        metrics["l_disc"] = loss.detach()
        return metrics

    # -- main step ----------------------------------------------------------
    def train_step(self, batch: dict) -> dict:
        # 1) text encode + VAE encode (no grad)
        with torch.no_grad():
            cond = self.bundle.text_encode(batch["captions"])
            real_latents = self.bundle.vae_encode(batch["pixels"].to(self.device))

        # 2) guidance × N
        log = {}
        for _ in range(self.args.guidance_updates):
            log.update(self.step_guidance(real_latents, cond))

        # 3) generator × 1
        log.update(self.step_generator(real_latents, cond))

        # 4) discriminator × 1 (adv_weight=0 のときは skip して R3GAN を完全 disable)
        if self.args.adv_weight > 0:
            log.update(self.step_discriminator(real_latents, cond))

        self.step += 1
        return log