scripts/distill/shortcut_module.py

"""
Shortcut Models 用の d-head injection。

設計:
  - transformer.d_head: nn.Linear(1 → cond_dim) (zero-init、初期は影響ゼロ)
  - transformer の t_embedder 出力に hook で d_head(d) を加算
  - 推論/訓練時に transformer._current_d = d を set してから forward
"""
from __future__ import annotations
import torch
import torch.nn as nn


def find_t_embedder(transformer: nn.Module) -> tuple[str, nn.Module]:
    """transformer 内部の t_embedder を発見。短い path を優先。"""
    candidates = []
    for name, module in transformer.named_modules():
        n = name.lower()
        if ("t_embedder" in n or "time_embedder" in n or "timestep_embedder" in n) and \
           ("llm_adapter" not in name):  # llm_adapter 内の embedder は除外
            candidates.append((name, module))
    if not candidates:
        # fallback: find by "time" in name (more inclusive)
        for name, module in transformer.named_modules():
            if "time" in name.lower() and "llm_adapter" not in name and \
               isinstance(module, (nn.Linear, nn.Sequential, nn.Module)) and \
               len(list(module.parameters())) > 0:
                candidates.append((name, module))
    if not candidates:
        raise RuntimeError("Could not find t_embedder module")
    # shortest path
    candidates.sort(key=lambda x: len(x[0].split(".")))
    name, module = candidates[0]
    print(f"[shortcut] found t_embedder at: {name}")
    return name, module


def _detect_t_embed_dim(transformer: nn.Module, t_embedder: nn.Module) -> int:
    """t_embedder の output dim を dummy forward で推定。"""
    device = next(t_embedder.parameters()).device
    dtype = next(t_embedder.parameters()).dtype
    with torch.no_grad():
        # 多くの実装は (B,) or (B,1) の timestep を受ける
        for shape in [(1,), (1, 1)]:
            try:
                t_in = torch.tensor([[0.5]] if len(shape) == 2 else [0.5],
                                    device=device, dtype=dtype)
                out = t_embedder(t_in)
                if isinstance(out, tuple):
                    out = out[0]
                return out.shape[-1]
            except Exception:
                continue
    raise RuntimeError("Could not detect t_embedder output dim")


def attach_shortcut_d_head(transformer: nn.Module) -> nn.Module:
    """transformer に d_head を attach、t_embedder 出力に hook で加算する。
    Returns: transformer (mutated)。`transformer._current_d` を set してから forward。"""
    name, t_emb = find_t_embedder(transformer)
    cond_dim = _detect_t_embed_dim(transformer, t_emb)
    print(f"[shortcut] t_embedder output dim: {cond_dim}")

    # d_head: 1 → cond_dim → cond_dim (small MLP, zero-init last layer)
    device = next(t_emb.parameters()).device
    dtype = next(t_emb.parameters()).dtype
    d_head = nn.Sequential(
        nn.Linear(1, cond_dim),
        nn.SiLU(),
        nn.Linear(cond_dim, cond_dim),
    ).to(device=device, dtype=dtype)
    # zero-init last linear so initial output is 0
    nn.init.zeros_(d_head[-1].weight)
    nn.init.zeros_(d_head[-1].bias)

    transformer.d_head = d_head
    transformer._current_d = None  # set per-forward

    def _hook(module, input, output):
        d = getattr(transformer, "_current_d", None)
        if d is None:
            return output
        # output shape: (B, cond_dim) or (B, 1, cond_dim) etc
        d_in = d.view(-1, 1).to(device=output.device if hasattr(output, 'device') else next(d_head.parameters()).device,
                                dtype=next(d_head.parameters()).dtype)
        d_emb = d_head(d_in)  # (B, cond_dim)
        # broadcast to match output shape
        if isinstance(output, tuple):
            base = output[0]
            d_emb_b = d_emb.view(base.size(0), *([1] * (base.dim() - 2)), base.size(-1)) if base.dim() > 2 else d_emb
            return (base + d_emb_b.to(dtype=base.dtype),) + output[1:]
        else:
            d_emb_b = d_emb.view(output.size(0), *([1] * (output.dim() - 2)), output.size(-1)) if output.dim() > 2 else d_emb
            return output + d_emb_b.to(dtype=output.dtype)

    handle = t_emb.register_forward_hook(_hook)
    transformer._d_head_hook_handle = handle

    return transformer


def set_shortcut_d(transformer: nn.Module, d: torch.Tensor | None):
    """forward 前に d を set。d=None で d-head 効果無効。"""
    transformer._current_d = d


def shortcut_d_head_params(transformer: nn.Module) -> list[nn.Parameter]:
    """d_head の trainable params (LoRA とは別の optimizer に渡す用)。"""
    if hasattr(transformer, "d_head"):
        return list(transformer.d_head.parameters())
    return []