scripts/distill/anima_loader.py

"""
Anima loader: diffusion-pipe の cosmos_predict2 Pipeline を活用して
DiT / VAE / Qwen3 text encoder / LLM adapter をロードする薄いラッパ。

diffusion-pipe 本体には sampling/inference 関数が無いので、ここで
- text_encode():  prompts -> crossattn_emb (LLM adapter 通過後)
- vae_encode():   pixels -> latents
- vae_decode():   latents -> pixels
- velocity():     DiT forward (rectified flow velocity prediction)
- add_noise():    noisy = (1-t)*latents + t*noise  (rectified flow forward)
- euler_step():   inference 用 1-step 前進
までを実装する。
"""
from __future__ import annotations
import os
import sys
from pathlib import Path
from dataclasses import dataclass

import torch
import torch.nn.functional as F


# diffusion-pipe を import path に強制配置。
# 重要: site-packages に top-level `utils` を持つ別パッケージがある場合、
# それが namespace package 化して shadow するため、(1) /workspace/diffusion-pipe を
# 先頭に固定 (2) cached `utils` モジュールを sys.modules から除外 する。
_DPIPE = Path("/workspace/diffusion-pipe")
if _DPIPE.exists():
    _dp = str(_DPIPE)
    sys.path = [_dp] + [p for p in sys.path if p != _dp]
    # 既にロード済みの top-level utils / models をクリア(他パッケージの shadow を排除)
    for _mod in list(sys.modules.keys()):
        if _mod == "utils" or _mod.startswith("utils.") \
           or _mod == "models" or _mod.startswith("models."):
            del sys.modules[_mod]


@dataclass
class AnimaPaths:
    transformer: str = "/models/checkpoints/anima-base-v1.0.safetensors"
    vae: str = "/models/checkpoints/qwen_image_vae.safetensors"
    llm: str = "/models/checkpoints/qwen_3_06b_base.safetensors"


class AnimaBundle:
    """Anima のモデル一式を保持。.transformer / .vae / .text_encoder / .tokenizer 等を直接触る。"""

    def __init__(self, pipeline, device: str | torch.device = "cuda"):
        self.pipeline = pipeline  # diffusion-pipe の CosmosPredict2Pipeline
        self.device = torch.device(device)

        self.transformer = pipeline.transformer
        self.vae = pipeline.vae
        self.text_encoder = pipeline.text_encoder  # Qwen3Model (inner)
        self.qwen_tokenizer = pipeline.tokenizer
        # LLM adapter は transformer 内部にある (Anima のみ)
        self.llm_adapter = getattr(pipeline.transformer, "llm_adapter", None)
        self.t5_tokenizer = getattr(pipeline, "t5_tokenizer", None)
        self.is_generic_llm = getattr(pipeline, "is_generic_llm", False)
        self.vae_scale = pipeline.vae.scale

    # ---- text encoding (Qwen3 → LLM adapter) ------------------------------
    @torch.no_grad()
    def text_encode(self, prompts: list[str]) -> torch.Tensor:
        """prompts -> crossattn_emb (B, 512, 1024) (DiT 用、LLM adapter 通過後)"""
        # diffusion-pipe の _tokenize を再現
        def _tok(tokenizer, prompts):
            return tokenizer(prompts, return_tensors="pt", truncation=True,
                              padding="max_length", max_length=512)

        qwen_enc = _tok(self.qwen_tokenizer, prompts)
        input_ids = qwen_enc.input_ids.to(self.device)
        attn_mask = qwen_enc.attention_mask.to(self.device)

        outputs = self.text_encoder(input_ids=input_ids, attention_mask=attn_mask)
        encoded = outputs.last_hidden_state
        encoded = encoded.masked_fill(~attn_mask.bool().unsqueeze(-1), 0.0)

        if self.llm_adapter is None or self.t5_tokenizer is None:
            return encoded

        t5_enc = _tok(self.t5_tokenizer, prompts)
        t5_ids = t5_enc.input_ids.to(self.device)
        t5_mask = t5_enc.attention_mask.to(self.device)

        crossattn = self.llm_adapter(
            source_hidden_states=encoded,
            target_input_ids=t5_ids,
            target_attention_mask=t5_mask,
            source_attention_mask=attn_mask,
        )
        crossattn = crossattn.masked_fill(~t5_mask.bool().unsqueeze(-1), 0.0)
        return crossattn

    # ---- VAE --------------------------------------------------------------
    @torch.no_grad()
    def vae_encode(self, pixels: torch.Tensor) -> torch.Tensor:
        """pixels (B, 3, H, W) in [-1, 1] -> latents (B, 16, 1, H/8, W/8)
        ※ Anima は静止画でも T=1 の 3D latent を扱う"""
        if pixels.dim() == 4:
            pixels = pixels.unsqueeze(2)  # (B, 3, 1, H, W)
        # VAE の weights dtype に合わせる
        vae_dtype = next(self.vae.model.parameters()).dtype
        pixels = pixels.to(device=self.device, dtype=vae_dtype)
        return self.vae.model.encode(pixels, self.vae_scale)

    @torch.no_grad()
    def vae_decode(self, latents: torch.Tensor) -> torch.Tensor:
        """latents (B, 16, 1, H_lat, W_lat) -> pixels (B, 3, 1, H, W) in [-1, 1]"""
        vae_dtype = next(self.vae.model.parameters()).dtype
        latents = latents.to(device=self.device, dtype=vae_dtype)
        return self.vae.model.decode(latents, self.vae_scale)

    # ---- DiT (rectified flow velocity prediction) -------------------------
    def velocity(
        self,
        latents: torch.Tensor,           # (B, 16, T, H_lat, W_lat) - noisy
        timesteps: torch.Tensor,         # (B,) in [0, 1]
        crossattn_emb: torch.Tensor,     # (B, 512, 1024)
        padding_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Anima DiT forward. 返り値は velocity (B, 16, T, H_lat, W_lat)。"""
        if padding_mask is None:
            padding_mask = self.zero_padding_mask(latents)
        return self.transformer(
            x_B_C_T_H_W=latents,
            timesteps_B_T=timesteps,
            crossattn_emb=crossattn_emb,
            padding_mask=padding_mask,
        )

    @staticmethod
    def zero_padding_mask(latents: torch.Tensor) -> torch.Tensor:
        """MiniTrainDIT.concat_padding_mask=True なので必須。shape=(B, 1, H_lat, W_lat)。"""
        B, _, _, H, W = latents.shape
        return torch.zeros(B, 1, H, W, dtype=latents.dtype, device=latents.device)

    @staticmethod
    def dit_forward(
        transformer: "torch.nn.Module",
        latents: torch.Tensor,
        t: torch.Tensor,
        crossattn_emb: torch.Tensor,
    ) -> torch.Tensor:
        """trainer から任意の transformer (gen / guidance) を呼ぶための薄い wrapper。
        全入力を transformer の weights dtype に揃え、padding_mask も自動生成。"""
        # weight dtype に揃える (t は float なので broadcast で他がアップキャストされやすい)
        w_dtype = next(transformer.parameters()).dtype
        latents = latents.to(dtype=w_dtype)
        t = t.to(dtype=w_dtype)
        crossattn_emb = crossattn_emb.to(dtype=w_dtype)
        padding_mask = AnimaBundle.zero_padding_mask(latents)
        return transformer(
            x_B_C_T_H_W=latents,
            timesteps_B_T=t,
            crossattn_emb=crossattn_emb,
            padding_mask=padding_mask,
        )

    # ---- rectified-flow utilities ----------------------------------------
    @staticmethod
    def add_noise(
        latents: torch.Tensor, noise: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        """forward process: x_t = (1-t)*x_0 + t*noise"""
        t_ = t.view(-1, *([1] * (latents.dim() - 1)))
        return (1 - t_) * latents + t_ * noise

    @staticmethod
    def velocity_target(latents: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
        """target velocity = noise - latents (rectified flow の training target)"""
        return noise - latents

    @staticmethod
    def x0_from_velocity(
        latents_t: torch.Tensor, v_pred: torch.Tensor, t: torch.Tensor
    ) -> torch.Tensor:
        """x_0 estimate = x_t - t * v_pred (linear FM のため)"""
        t_ = t.view(-1, *([1] * (latents_t.dim() - 1)))
        return latents_t - t_ * v_pred

    @staticmethod
    def euler_step(
        latents_t: torch.Tensor,
        v_pred: torch.Tensor,
        t: torch.Tensor,
        t_next: torch.Tensor,
    ) -> torch.Tensor:
        """Euler step: x_{t_next} = x_t + (t_next - t) * v_pred"""
        dt = (t_next - t).view(-1, *([1] * (latents_t.dim() - 1)))
        return latents_t + dt * v_pred


def build_anima(
    paths: AnimaPaths | None = None,
    device: str | torch.device = "cuda",
    dtype: torch.dtype = torch.bfloat16,
) -> AnimaBundle:
    """diffusion-pipe の CosmosPredict2Pipeline を初期化して AnimaBundle で返す。

    内部的には toml config を最小限作って Pipeline に渡す。
    """
    import importlib
    import importlib.util
    _dp = str(_DPIPE)

    # diffusion-pipe の utils/ と models/ に空 __init__.py を作成して regular package 化
    for sub in ("utils", "models", "optimizers"):
        sub_dir = os.path.join(_dp, sub)
        init_file = os.path.join(sub_dir, "__init__.py")
        if os.path.isdir(sub_dir) and not os.path.exists(init_file):
            with open(init_file, "w"):
                pass

    # sys.path 強制 + cache クリア
    sys.path = [_dp] + [p for p in sys.path if p != _dp]
    for _mod in list(sys.modules.keys()):
        if _mod in ("utils", "models", "optimizers") \
           or _mod.startswith(("utils.", "models.", "optimizers.")):
            del sys.modules[_mod]

    # importlib で **明示的に** 各 package を /workspace/diffusion-pipe 配下から
    # 読み込み、sys.modules に固定して shadow を完全排除する。
    def _load_pkg(name: str):
        pkg_dir = os.path.join(_dp, name)
        init_file = os.path.join(pkg_dir, "__init__.py")
        spec = importlib.util.spec_from_file_location(
            name, init_file,
            submodule_search_locations=[pkg_dir],
        )
        mod = importlib.util.module_from_spec(spec)
        sys.modules[name] = mod
        spec.loader.exec_module(mod)
        print(f"[setup] forced load {name} from {pkg_dir}")

    for _pkg in ("utils", "models", "optimizers"):
        _load_pkg(_pkg)

    # 動作確認: utils.common が解決できることを試す
    import utils.common  # noqa: F401
    print("[setup] utils.common import OK")

    from models import cosmos_predict2  # diffusion-pipe 側

    paths = paths or AnimaPaths()

    # CosmosPredict2Pipeline は config dict から **torch.dtype object** を期待する
    # (TOML 読み込み時に diffusion-pipe 側で str→dtype 変換するが、ここでは直接渡す)
    cfg = {
        "model": {
            "type": "anima",
            "transformer_path": str(paths.transformer),
            "vae_path": str(paths.vae),
            "llm_path": str(paths.llm),
            "dtype": dtype,
            "transformer_dtype": dtype,
            "timestep_sample_method": "logit_normal",
        },
        # adapter / optimizer 等は trainer 側で attach するので空でOK
        "adapter": {"type": "lora", "rank": 1, "alpha": 1, "dropout": 0.0, "dtype": dtype},
    }

    # __init__ で VAE + text encoder + tokenizers が CPU 上に load 済み
    pipeline = cosmos_predict2.CosmosPredict2Pipeline(cfg)
    # DiT をロード (transformer.llm_adapter も内部で attach される)
    pipeline.load_diffusion_model()

    # device / dtype に配置
    pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
    pipeline.text_encoder = pipeline.text_encoder.to(device=device, dtype=dtype)
    pipeline.vae.model = pipeline.vae.model.to(device=device, dtype=dtype)
    # mean/std は __init__ で 'cuda' に moved 済 (cosmos_predict2.py:200-201)

    return AnimaBundle(pipeline, device=device)