""" 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)