SolitaryThinker
4a33540c08
The torch adapter boundary marshalled the latent host<->device on EVERY denoise step: _t uploaded the latent (and re-uploaded the text embeds) and _n downloaded the velocity with a forced CUDA sync — 2*N PCIe copies + N syncs per generation, buying nothing, since the latent could stay resident on the GPU the whole loop. Root cause was a numpy loop surface. But the loop MATH is already array-agnostic (CFG combine + flow-match Euler are pure arithmetic; the solver kernel already passes torch through). So introduce a per-platform array namespace (v2/platform/ array_ns): numpy on CPU (torch-free — the parity mini is unchanged), torch-on- device on cuda. The latent is seeded with numpy and uploaded ONCE; it then stays resident through forward -> CFG combine -> solver -> next step; a single host marshal happens at the request/output boundary (engine._to_artifact). Opt-in per recipe via ModelCard.device_io (set on the Wan cards). When set on a GPU box, build_component flips the components' TorchComponent.device_io so _out keeps tensors on-device (in fp32, matching the old _to_numpy cast so the combine dtype is unchanged). Un-migrated families and the CPU toy keep numpy in/out. Also: PrecisionPolicy.cast is array-preserving; _t accepts resident tensors. Verified BIT-IDENTICAL on real Wan2.1-1.3B / H100: the on-device latent equals the pre-change numpy-path latent exactly (max_abs_diff 0.0, np.array_equal True). CPU mini holds 237 passed / 5 pre-existing; pre-commit clean. Other WanDenoiseLoop families can flip device_io next (per-family GPU re-verify); non-Wan loops migrate to the xp namespace later.
647 lines
33 KiB
Python
647 lines
33 KiB
Python
"""Real torch component adapters for the GPU backend. See v2/README.md ("Running the real models").
|
|
|
|
A ``TorchComponent`` base centralizes the shared mechanics — ``.to(device,dtype).eval()``, the
|
|
numpy<->torch marshalling at the loop boundary, the ``set_forward_context`` wrap every fastvideo
|
|
forward needs, and the weight surface (copy_from / blend_from / clone). Thin per-model subclasses
|
|
carry only the forward semantics (Wan ``sigma*1000`` -> velocity; LTX-2 per-token sigma,
|
|
x0 -> ``(x_t-x0)/sigma``, joint A/V; VAE normalization; T5 padding; Gemma dual-projection; upsampler;
|
|
audio decode->vocoder). A single ``build_component(spec, instance, platform)`` dispatches by
|
|
``spec.kind`` through ``_MAKERS``.
|
|
|
|
The loop math (CFG combine, flow-match Euler) is array-agnostic, so the boundary is the ONLY place
|
|
host<->device crossing happens (``TorchComponent._t``/``_out``). When the card sets ``device_io`` (a GPU
|
|
box), egress keeps tensors on-device and the latent stays resident for the whole denoise loop — no
|
|
per-step round-trip; otherwise the default marshals to host numpy (un-migrated families / CPU toy). See
|
|
``v2/platform/array_ns`` for the loop-side namespace. Component construction goes through
|
|
``load_component``, the v2-owned loader seam (currently delegates to fastvideo's component_loader; a
|
|
per-module vendored cutover replaces it later, no caller changes).
|
|
|
|
Imported ONLY lazily — inside ``torch_cuda.py``'s registered builder — so importing the platform package
|
|
on a CPU box never imports torch and the mini stays green. Wan2.1 + LTX-2 A/V GPU-verified.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
# Flow-match timestep convention: the loop hands the raw sigma (1->0); the diffusers/FastVideo Wan
|
|
# convention embeds ``timestep = sigma * num_train_timesteps`` (BRINGUP risk B, confirmed on box).
|
|
NUM_TRAIN_TIMESTEPS = 1000
|
|
|
|
_RUNTIME_READY = False
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Construction helpers (FastVideoArgs / loaders / dist-init / marshalling) #
|
|
# --------------------------------------------------------------------------- #
|
|
def _ensure_fastvideo_runtime() -> None:
|
|
"""Single-process distributed env the real fastvideo loaders require (mirrors
|
|
``fastvideo/worker/gpu_worker.py:init_device``): the loaders call ``get_local_torch_device()`` and
|
|
build a 1x1 device mesh, which needs an initialized process group. Idempotent. BRINGUP risk A."""
|
|
global _RUNTIME_READY
|
|
if _RUNTIME_READY:
|
|
return
|
|
os.environ.setdefault("MASTER_ADDR", "localhost")
|
|
os.environ.setdefault("MASTER_PORT", "29500")
|
|
os.environ.setdefault("RANK", "0")
|
|
os.environ.setdefault("WORLD_SIZE", "1")
|
|
os.environ.setdefault("LOCAL_RANK", "0")
|
|
if torch.cuda.is_available():
|
|
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
|
|
from v2.distributed import maybe_init_distributed_environment_and_model_parallel
|
|
maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
|
|
_RUNTIME_READY = True
|
|
|
|
|
|
def _require_checkpoint(spec) -> str:
|
|
ckpt = getattr(spec, "checkpoint", "") or ""
|
|
if not ckpt:
|
|
raise RuntimeError(f"component {spec.component_id!r}: set ComponentSpec.checkpoint to the weights path / HF id "
|
|
f"for the GPU backend (load_id={spec.load_id!r}). See GPU_BRINGUP.md, risk A.")
|
|
return ckpt
|
|
|
|
|
|
def _model_root(spec) -> str:
|
|
"""Model root = the dir holding ``model_index.json`` + component subfolders. ``spec.checkpoint`` is
|
|
the component subfolder (e.g. ``<root>/transformer``); its parent is the registry root. BRINGUP A."""
|
|
return os.path.dirname(os.path.normpath(_require_checkpoint(spec)))
|
|
|
|
|
|
def _fastvideo_args(spec: Any) -> Any:
|
|
"""Build the FastVideoArgs the real loaders need (BRINGUP risk A). ``from_kwargs`` populates
|
|
``pipeline_config`` (dit/vae/text-encoder configs + precisions) from the model root. Single-GPU,
|
|
all offload/FSDP OFF — weights resident on one device, the simplest correct bring-up."""
|
|
from v2.fastvideo_args import FastVideoArgs
|
|
return FastVideoArgs.from_kwargs(model_path=_model_root(spec),
|
|
num_gpus=1,
|
|
tp_size=1,
|
|
sp_size=1,
|
|
dit_cpu_offload=False,
|
|
text_encoder_cpu_offload=False,
|
|
vae_cpu_offload=False,
|
|
image_encoder_cpu_offload=False,
|
|
dit_layerwise_offload=False,
|
|
use_fsdp_inference=False,
|
|
pin_cpu_memory=False)
|
|
|
|
|
|
def load_component(loader_attr: str, path: str, args):
|
|
"""v2-owned loader seam (``v2/loader``). Constructs a real component from a checkpoint via the
|
|
component loaders. ``WanTransformer3DModel`` / ``AutoencoderKLWan`` have NO ``from_pretrained``; the
|
|
loader reads the checkpoint config, resolves the class (UMT5 vs T5 from config), and loads weights.
|
|
Currently delegates to fastvideo's loader; a vendored cutover swaps the body, not the callers."""
|
|
from v2.loader import component_loader as _cl
|
|
return getattr(_cl, loader_attr)().load(path, args)
|
|
|
|
|
|
def _device(platform) -> str:
|
|
return "cuda" if platform.device == "cuda" else platform.device
|
|
|
|
|
|
def _native_dtype(module: Any) -> Any:
|
|
"""Keep each component at the precision its loader produced (Wan DiT bf16, VAE/text fp32) rather than
|
|
the card's uniform fp32 — faithful to fastvideo and ~2x faster/smaller for the DiT."""
|
|
try:
|
|
return next(module.parameters()).dtype
|
|
except StopIteration:
|
|
return torch.float32
|
|
|
|
|
|
def _to_torch(a: Any, *, device: Any, dtype: Any) -> torch.Tensor | None:
|
|
return None if a is None else torch.as_tensor(np.asarray(a), dtype=dtype, device=device)
|
|
|
|
|
|
def _to_numpy(t: Any) -> np.ndarray:
|
|
if hasattr(t, "sample"): # some heads wrap output in an object with .sample
|
|
t = t.sample
|
|
return t.detach().to("cpu", torch.float32).numpy()
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# TorchComponent — the shared base (eval/marshalling/forward-context/weights) #
|
|
# --------------------------------------------------------------------------- #
|
|
class TorchComponent:
|
|
"""Wraps a real ``fastvideo.models.*`` module to the mini's numpy duck-typed surface. Subclasses
|
|
override only the forward semantics (``__call__`` / ``encode`` / ``decode`` / ``upsample`` / ...)."""
|
|
|
|
def __init__(self, module: Any, *, device: Any, dtype: Any, eager: bool = True) -> None:
|
|
self.device, self.dtype = device, dtype
|
|
self.module = module.to(device=device, dtype=dtype).eval() if eager else module
|
|
# When True (set by build_component from the card's device_io flag on a GPU box), the loop
|
|
# boundary keeps tensors on-device — no per-step host<->device copy. Default False = today's
|
|
# numpy in/out, so un-migrated families and the CPU toy are unchanged.
|
|
self.device_io = False
|
|
|
|
# host<->device marshalling at the loop boundary (ONE place) ------------------------------------- #
|
|
def _t(self, a: Any, *, batch: bool = True) -> Any:
|
|
if a is None:
|
|
return None
|
|
# Ingress: an already-on-device tensor (the loop kept it resident) is moved to our dtype, never
|
|
# round-tripped through numpy; a host array is uploaded here.
|
|
t = a.to(device=self.device, dtype=self.dtype) if torch.is_tensor(a) else _to_torch(
|
|
a, device=self.device, dtype=self.dtype)
|
|
return t.unsqueeze(0) if batch else t
|
|
|
|
def _out(self, t: Any) -> Any:
|
|
"""Egress: keep the tensor on-device when ``device_io`` (it stays resident for the next step),
|
|
else marshal to host numpy (the default — un-migrated families / the CPU toy are unchanged).
|
|
Either way it lands in fp32 — matching the host path's ``_to_numpy`` (``.to(cpu, float32)``) so
|
|
the loop's CFG combine runs in fp32 regardless of the module's native dtype (e.g. Wan DiT bf16)."""
|
|
if not self.device_io:
|
|
return _to_numpy(t)
|
|
if hasattr(t, "sample"): # unwrap heads that wrap output in an object with .sample
|
|
t = t.sample
|
|
return t.float()
|
|
|
|
def _n(self, t: Any) -> Any:
|
|
return self._out(t.squeeze(0))
|
|
|
|
@staticmethod
|
|
def _ctx(current_timestep: Any = 0) -> Any:
|
|
"""The FastVideo attention layer reads attn_metadata via get_forward_context(), so every forward
|
|
runs inside set_forward_context(...). attn_metadata=None selects the dense SDPA path."""
|
|
from v2.forward_context import set_forward_context
|
|
return set_forward_context(current_timestep=current_timestep, attn_metadata=None)
|
|
|
|
# weight surface used by serving weight-sync + training ----------------------------------------- #
|
|
def copy_from(self, other) -> None:
|
|
self.module.load_state_dict(other.module.state_dict())
|
|
|
|
def blend_from(self, other, decay: float) -> None: # EMA / decayed-old-policy
|
|
with torch.no_grad():
|
|
for p, q in zip(self.module.parameters(), other.module.parameters(), strict=False):
|
|
p.mul_(decay).add_(q, alpha=1.0 - decay)
|
|
|
|
def clone(self) -> TorchComponent:
|
|
import copy
|
|
c = self.__class__.__new__(self.__class__)
|
|
c.__dict__.update(self.__dict__) # tokenizer / shapes / sibling refs shared
|
|
c.module = copy.deepcopy(self.module) # independent weights
|
|
return c
|
|
|
|
def mse_grad_step(self, *a, **k):
|
|
raise NotImplementedError(
|
|
"GPU training surface (mse_grad_step) is a separate workstream — see GPU_BRINGUP.md risk F")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# DiT adapters #
|
|
# --------------------------------------------------------------------------- #
|
|
class WanDiT(TorchComponent):
|
|
"""dit(latent[C,T,H,W], text_embed[seq,dim], sigma) -> velocity[C,T,H,W]. Real forward
|
|
(wanvideo.py): forward(hidden_states[B,C,T,H,W], encoder_hidden_states, timestep,
|
|
encoder_hidden_states_image=None) -> velocity (bare tensor)."""
|
|
|
|
def __init__(self,
|
|
module: Any,
|
|
*,
|
|
device: Any,
|
|
dtype: Any,
|
|
offload_group: Any = None,
|
|
component_id: str = "transformer") -> None:
|
|
# Wan2.2 MoE (A14B): two 14B experts don't both fit one 80GB GPU. With ``offload_group`` set, keep
|
|
# this expert on CPU and bring only the *active* one onto the GPU on demand (single swap at the
|
|
# boundary, not per-step thrash). Single-expert Wan stays resident (offload_group=None).
|
|
self.offload_group = offload_group
|
|
self.component_id = component_id
|
|
super().__init__(module, device=device, dtype=dtype, eager=(offload_group is None))
|
|
if offload_group is not None:
|
|
self.module = self.module.to(device="cpu", dtype=dtype).eval()
|
|
self._on_gpu = False
|
|
offload_group[component_id] = self
|
|
else:
|
|
self._on_gpu = True
|
|
# CausalWanTransformer3DModel (self-forcing student) conditions across chunks via an internal
|
|
# kv_cache, not a forward arg, so the chunk_rollout loop's latent ``context`` is ignored here.
|
|
self.causal = "Causal" in type(module).__name__
|
|
|
|
def _ensure_resident(self) -> None:
|
|
if self.offload_group is None or self._on_gpu:
|
|
return
|
|
for other in self.offload_group.values():
|
|
if other is not self and other._on_gpu:
|
|
other.module.to("cpu")
|
|
other._on_gpu = False
|
|
torch.cuda.empty_cache()
|
|
self.module.to(self.device)
|
|
self._on_gpu = True
|
|
|
|
def alloc_causal_caches(self, frame_seqlen: int, *, batch: int = 1, max_text_len: int = 512) -> Any:
|
|
"""Allocate the persistent per-block KV + cross-attn caches the causal student needs for
|
|
cross-chunk conditioning (mirrors CausalDenoisingStage._initialize_kv_cache). The chunk_rollout
|
|
loop owns these in its LoopState (per request, never on this shared adapter) and threads them
|
|
through __call__ so the model routes to _forward_inference."""
|
|
self._ensure_resident()
|
|
m = self.module
|
|
nblocks = len(m.blocks)
|
|
nheads, hdim = int(m.num_attention_heads), int(m.attention_head_dim)
|
|
local = int(getattr(m, "local_attn_size", -1) or -1)
|
|
if local != -1:
|
|
kv_size = local * frame_seqlen
|
|
else: # global window: model keeps the last GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES (21) frames
|
|
arch = getattr(getattr(m, "config", None), "arch_config", None)
|
|
sliding = int(getattr(arch, "sliding_window_num_frames", 0) or 0)
|
|
kv_size = max(sliding, 21) * frame_seqlen
|
|
dev, dt = self.device, self.dtype
|
|
kv = [{
|
|
"k": torch.zeros(batch, kv_size, nheads, hdim, device=dev, dtype=dt),
|
|
"v": torch.zeros(batch, kv_size, nheads, hdim, device=dev, dtype=dt),
|
|
"global_end_index": torch.tensor([0], dtype=torch.long, device=dev),
|
|
"local_end_index": torch.tensor([0], dtype=torch.long, device=dev),
|
|
} for _ in range(nblocks)]
|
|
ca = [{
|
|
"k": torch.zeros(batch, max_text_len, nheads, hdim, device=dev, dtype=dt),
|
|
"v": torch.zeros(batch, max_text_len, nheads, hdim, device=dev, dtype=dt),
|
|
"is_init": False,
|
|
} for _ in range(nblocks)]
|
|
return {"kv": kv, "crossattn": ca}
|
|
|
|
@torch.no_grad()
|
|
def __call__(self,
|
|
latent,
|
|
text_embed,
|
|
sigma,
|
|
context=None,
|
|
*,
|
|
cond=None,
|
|
kv_cache=None,
|
|
crossattn_cache=None,
|
|
current_start=0,
|
|
start_frame=0,
|
|
frame_seqlen=1560):
|
|
self._ensure_resident()
|
|
hs = self._t(latent)
|
|
if cond is not None: # i2v: concat [noise (16ch) ; mask+cond_latent (20ch)] -> 36ch DiT input
|
|
hs = torch.cat([hs, self._t(cond)], dim=1)
|
|
ehs = self._t(text_embed)
|
|
# timestep = sigma*1000 (BRINGUP risk B). Causal model needs per-latent-frame [B, num_frames].
|
|
ts = float(sigma) * NUM_TRAIN_TIMESTEPS
|
|
timestep = (torch.full(
|
|
(1, hs.shape[2]), ts, device=self.device) if self.causal else torch.tensor([ts], device=self.device))
|
|
img = None if self.causal else self._t(context) # i2v image embedding for standard Wan
|
|
fwd = dict(hidden_states=hs, encoder_hidden_states=ehs, timestep=timestep, encoder_hidden_states_image=img)
|
|
if self.causal and kv_cache is not None:
|
|
# Pass the persistent caches + frame offset so the model runs _forward_inference (CausVid
|
|
# Algorithm 2) and conditions this chunk on prior clean chunks. Without it the model falls
|
|
# through to _forward_train (no cross-chunk KV) -> chunks denoise blind -> discontinuity.
|
|
fwd.update(kv_cache=kv_cache,
|
|
crossattn_cache=crossattn_cache,
|
|
current_start=int(current_start),
|
|
cache_start=int(current_start),
|
|
start_frame=int(start_frame),
|
|
frame_seqlen=int(frame_seqlen))
|
|
with self._ctx():
|
|
velocity = self.module(**fwd)
|
|
return self._n(velocity) # rectified-flow velocity (BRINGUP risk C)
|
|
|
|
def clone(self) -> WanDiT:
|
|
c = super().clone()
|
|
assert isinstance(c, WanDiT)
|
|
c.offload_group, c._on_gpu = None, True # standalone resident copy
|
|
return c
|
|
|
|
|
|
class LTX2DiT(TorchComponent):
|
|
"""LTX-2 DiT: patchifies internally, takes a PER-TOKEN timestep ``ones(B, token_count, 1) * sigma``
|
|
(sigma direct, 0..1) + a per-sample ``video_sigma``; predicts ``denoised`` (x0), so the adapter
|
|
returns the flow-match velocity ``(x_t - x0)/sigma`` the loop integrates. Joint A/V (audio_latent
|
|
given) -> (video_velocity, audio_velocity) in one forward."""
|
|
|
|
def __init__(self, module: Any, *, device: Any, dtype: Any) -> None:
|
|
super().__init__(module, device=device, dtype=dtype)
|
|
from v2.models.dits.ltx2 import VideoLatentShape
|
|
self._VideoLatentShape = VideoLatentShape
|
|
|
|
@torch.no_grad()
|
|
def __call__(self, latent, text_embed, sigma, context=None, *, audio_latent=None, audio_text=None):
|
|
hs = self._t(latent)
|
|
ehs = self._t(text_embed)
|
|
s = float(sigma)
|
|
token_count = self.module.patchifier.get_token_count(self._VideoLatentShape.from_torch_shape(tuple(hs.shape)))
|
|
timestep = torch.full((1, token_count, 1), s, device=self.device, dtype=torch.float32)
|
|
video_sigma = torch.tensor([s], device=self.device, dtype=torch.float32)
|
|
av_kwargs: dict = {}
|
|
au = None
|
|
if audio_latent is not None: # joint A/V: audio latent [1,8,T,16] + audio text
|
|
from v2.models.audio.ltx2_audio_vae import AudioLatentShape
|
|
au = self._t(audio_latent)
|
|
aeh = self._t(audio_text)
|
|
atok = self.module.audio_patchifier.get_token_count(AudioLatentShape.from_torch_shape(tuple(au.shape)))
|
|
av_kwargs = dict(audio_hidden_states=au,
|
|
audio_encoder_hidden_states=aeh,
|
|
audio_timestep=torch.full((1, atok, 1), s, device=self.device, dtype=torch.float32),
|
|
audio_sigma=torch.tensor([s], device=self.device, dtype=torch.float32))
|
|
with self._ctx(current_timestep=s):
|
|
out = self.module(hidden_states=hs,
|
|
encoder_hidden_states=ehs,
|
|
timestep=timestep,
|
|
video_sigma=video_sigma,
|
|
encoder_attention_mask=None,
|
|
**av_kwargs)
|
|
if au is not None: # LTX-2 DiT predicts x0; integrate velocity
|
|
denoised_v, denoised_a = out
|
|
vel_v = (hs.float() - denoised_v.float()) / max(s, 1e-6)
|
|
vel_a = (au.float() - denoised_a.float()) / max(s, 1e-6)
|
|
return self._n(vel_v), self._n(vel_a)
|
|
denoised = out[0] if isinstance(out, tuple) else out
|
|
return self._n((hs.float() - denoised.float()) / max(s, 1e-6))
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# VAE adapters #
|
|
# --------------------------------------------------------------------------- #
|
|
class WanVAE(TorchComponent):
|
|
"""The DiT operates in NORMALIZED latent space, so encode applies ``(z - mean)/std`` and decode
|
|
inverts it; ``AutoencoderKLWan.encode/decode`` operate in RAW latent space (BRINGUP risk D)."""
|
|
|
|
def _mean_invstd(self, like: Any) -> tuple[torch.Tensor, torch.Tensor]:
|
|
mean = torch.tensor(self.module.latents_mean, device=like.device, dtype=torch.float32).view(1, -1, 1, 1, 1)
|
|
inv_std = (1.0 / torch.tensor(self.module.latents_std, device=like.device, dtype=torch.float32)).view(
|
|
1, -1, 1, 1, 1)
|
|
return mean, inv_std
|
|
|
|
@torch.no_grad()
|
|
def encode(self, video):
|
|
x = self._t(video)
|
|
dist = self.module.encode(x)
|
|
z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist)
|
|
mean, inv_std = self._mean_invstd(z)
|
|
return self._n((z.float() - mean) * inv_std) # -> the normalized latent the DiT expects
|
|
|
|
@torch.no_grad()
|
|
def decode(self, latent):
|
|
z = self._t(latent).float()
|
|
mean, inv_std = self._mean_invstd(z)
|
|
z = z / inv_std + mean # invert encode normalization -> raw latent
|
|
video = self.module.decode(z.to(self.dtype)) # -> video [B,3,T,H,W] in [-1,1]
|
|
return self._n(video)
|
|
|
|
|
|
class LTX2VAE(TorchComponent):
|
|
"""LTX-2 VideoDecoder un-normalizes internally (per-channel stats); the adapter just marshals."""
|
|
|
|
def __init__(self, module: Any, *, device: Any, dtype: Any) -> None:
|
|
super().__init__(module, device=device, dtype=dtype)
|
|
# Tiled decode (spatial tiles concatenated): the LTX-2 CausalVideoAutoencoder decode of a full-res,
|
|
# full-length (121f / 1536x1024) latent allocates >80GB of conv activations otherwise. fastvideo
|
|
# defaults vae_tiling=True for LTX-2; mirror it so the full clip fits on one GPU.
|
|
if hasattr(self.module, "enable_tiling"):
|
|
self.module.enable_tiling()
|
|
|
|
@torch.no_grad()
|
|
def decode(self, latent):
|
|
video = self.module.decode(self._t(latent))
|
|
if hasattr(video, "sample"):
|
|
video = video.sample
|
|
return self._n(video)
|
|
|
|
@torch.no_grad()
|
|
def encode(self, video):
|
|
dist = self.module.encode(self._t(video))
|
|
z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist)
|
|
return self._n(z)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Text encoders #
|
|
# --------------------------------------------------------------------------- #
|
|
class T5Encoder(TorchComponent):
|
|
"""(U)MT5 text_encoder.encode(text) -> embedding[text_len, dim] (numpy out)."""
|
|
|
|
def __init__(self, module: Any, tokenizer: Any, *, device: Any, dtype: Any, max_length: int = 512) -> None:
|
|
super().__init__(module, device=device, dtype=dtype)
|
|
self.tokenizer = tokenizer
|
|
self.max_length = max_length
|
|
|
|
@torch.no_grad()
|
|
def encode(self, text):
|
|
toks = self.tokenizer(text or "", return_tensors="pt", max_length=self.max_length, truncation=True)
|
|
ids = toks.input_ids.to(self.device)
|
|
mask = toks.attention_mask.to(self.device)
|
|
with self._ctx():
|
|
out = self.module(input_ids=ids, attention_mask=mask)
|
|
hidden = (out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]).squeeze(0)
|
|
# Wan convention: the DiT cross-attends over a FIXED-length text sequence — real-token rows then
|
|
# a ZERO-padded tail to text_len (a short unpadded sequence mis-conditions the DiT).
|
|
if hidden.shape[0] < self.max_length:
|
|
pad = hidden.new_zeros(self.max_length - hidden.shape[0], hidden.shape[1])
|
|
hidden = torch.cat([hidden, pad], dim=0)
|
|
return self._out(hidden)
|
|
|
|
|
|
class Gemma(TorchComponent):
|
|
"""LTX-2 Gemma text encoder. ``encode`` -> video projection; ``encode_av`` -> (video, audio): the 2.3
|
|
connector emits SEPARATE projections (audio in ``hidden_states[0]`` only when output_hidden_states)."""
|
|
|
|
def __init__(self, module: Any, tokenizer: Any, *, device: Any, dtype: Any, max_length: int = 256) -> None:
|
|
super().__init__(module, device=device, dtype=dtype)
|
|
self.tokenizer = tokenizer
|
|
self.max_length = max_length
|
|
|
|
def _tok(self, text: Any) -> tuple[torch.Tensor, torch.Tensor]:
|
|
toks = self.tokenizer(text or "",
|
|
return_tensors="pt",
|
|
max_length=self.max_length,
|
|
truncation=True,
|
|
padding="max_length")
|
|
return toks.input_ids.to(self.device), toks.attention_mask.to(self.device)
|
|
|
|
@torch.no_grad()
|
|
def encode(self, text):
|
|
ids, mask = self._tok(text)
|
|
with self._ctx():
|
|
out = self.module(input_ids=ids, attention_mask=mask)
|
|
hidden = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
|
|
return _to_numpy(hidden.squeeze(0))
|
|
|
|
@torch.no_grad()
|
|
def encode_av(self, text):
|
|
ids, mask = self._tok(text)
|
|
with self._ctx():
|
|
out = self.module(input_ids=ids, attention_mask=mask, output_hidden_states=True)
|
|
video = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
|
|
hs = getattr(out, "hidden_states", None)
|
|
audio = hs[0] if hs else video # separate audio connector projection
|
|
return _to_numpy(video.squeeze(0)), _to_numpy(audio.squeeze(0))
|
|
|
|
|
|
class CLIPImageEncoder(TorchComponent):
|
|
"""CLIP-vision image encoder for Wan i2v: ``encode_image(image) -> image_embeds`` (the DiT's
|
|
``encoder_hidden_states_image``). Mirrors fastvideo's ImageEncodingStage — the HF image processor
|
|
preprocesses, the encoder returns ``last_hidden_state``. BRINGUP: written-not-run; GPU-verify the
|
|
processor/encoder subfolders + dtype against a real i2v checkpoint (see GPU_BRINGUP.md)."""
|
|
|
|
def __init__(self, module: Any, processor: Any, *, device: Any, dtype: Any) -> None:
|
|
super().__init__(module, device=device, dtype=dtype)
|
|
self.processor = processor
|
|
|
|
@torch.no_grad()
|
|
def encode_image(self, image):
|
|
inputs = self.processor(images=image, return_tensors="pt").to(self.device)
|
|
with self._ctx():
|
|
out = self.module(**inputs)
|
|
embeds = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
|
|
return self._out(embeds.squeeze(0))
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Upsampler + audio (LTX-2) #
|
|
# --------------------------------------------------------------------------- #
|
|
class LTX2Upsampler(TorchComponent):
|
|
"""upsample(latent[C,T,H,W]) -> [C,T,2H,2W]. Applies the repo's ``upsample_video``: un_normalize (via
|
|
the video VAE's per_channel_statistics) -> learned 2x upsample -> normalize."""
|
|
|
|
def __init__(self, module: Any, stats_owner: Any, *, device: Any, dtype: Any) -> None:
|
|
super().__init__(module, device=device, dtype=dtype)
|
|
self.stats_owner = stats_owner # exposes per_channel_statistics (the VAE decoder/encoder)
|
|
|
|
@torch.no_grad()
|
|
def upsample(self, latent):
|
|
from v2.models.upsamplers.ltx2_upsampler import upsample_video
|
|
out = upsample_video(self._t(latent), self.stats_owner, self.module)
|
|
return self._n(out)
|
|
|
|
|
|
class LTX2Vocoder(TorchComponent):
|
|
"""Thin wrapper over the LTX-2 Vocoder (mel -> waveform @24kHz); chained by LTX2AudioVAE."""
|
|
|
|
|
|
class LTX2AudioVAE(TorchComponent):
|
|
"""audio_vae.decode(audio_latent[8,T,16]) -> waveform. Runs AudioDecoder (latent->mel) then the
|
|
Vocoder (mel->waveform). ``sample_rate`` = the vocoder's output rate (so a saver writes the right wav)."""
|
|
|
|
def __init__(self, module: Any, vocoder: Any, *, device: Any, dtype: Any) -> None:
|
|
super().__init__(module, device=device, dtype=dtype)
|
|
self.vocoder = vocoder # LTX2Vocoder | None
|
|
self.sample_rate = int(getattr(getattr(vocoder, "module", None), "output_sample_rate", 24000))
|
|
|
|
@torch.no_grad()
|
|
def decode(self, audio_latent):
|
|
mel = self.module(self._t(audio_latent))
|
|
if hasattr(mel, "sample"):
|
|
mel = mel.sample
|
|
if self.vocoder is not None:
|
|
wav = self.vocoder.module(mel)
|
|
if hasattr(wav, "sample"):
|
|
wav = wav.sample
|
|
else:
|
|
wav = mel
|
|
return self._n(wav)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# build_component — one dispatch (replaces the six build_torch_* + trampolines) #
|
|
# --------------------------------------------------------------------------- #
|
|
def _explicit_adapter(spec: Any, module: Any, device: Any, dtype: Any, *extra: Any) -> Any:
|
|
"""Build the card's explicitly-declared adapter (``ComponentSpec.adapter='module:Class'``) if set —
|
|
the seam that lets a NEW architecture's recipe carry its own TorchComponent subclass without editing
|
|
the dispatch here (so a port is a self-contained recipe package). Constructed as
|
|
``cls(module, *extra, device=device, dtype=dtype)``. Returns None when unset (-> built-in dispatch)."""
|
|
ref = getattr(spec, "adapter", "") or ""
|
|
if not ref:
|
|
return None
|
|
import importlib
|
|
mod, _, cls = ref.partition(":")
|
|
return getattr(importlib.import_module(mod), cls)(module, *extra, device=device, dtype=dtype)
|
|
|
|
|
|
def _make_dit(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
|
module = load_component("TransformerLoader", spec.checkpoint, args)
|
|
device, dtype = _device(platform), _native_dtype(module)
|
|
explicit = _explicit_adapter(spec, module, device, dtype)
|
|
if explicit is not None:
|
|
return explicit
|
|
if "LTX2" in type(module).__name__:
|
|
return LTX2DiT(module, device=device, dtype=dtype)
|
|
# Wan2.2 MoE: >1 DiT expert -> CPU-offload all but the active one (swapped at the boundary).
|
|
n_experts = sum(1 for c in instance.card.components.values() if getattr(c, "kind", None) == "dit")
|
|
grp = None
|
|
if n_experts > 1:
|
|
grp = getattr(instance, "_dit_offload_group", None)
|
|
if grp is None:
|
|
grp = instance._dit_offload_group = {}
|
|
return WanDiT(module, device=device, dtype=dtype, offload_group=grp, component_id=spec.component_id)
|
|
|
|
|
|
def _make_vae(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
|
module = load_component("VAELoader", spec.checkpoint, args)
|
|
device, dtype = _device(platform), _native_dtype(module)
|
|
explicit = _explicit_adapter(spec, module, device, dtype)
|
|
if explicit is not None:
|
|
return explicit
|
|
cls = LTX2VAE if "LTX2" in type(module).__name__ else WanVAE
|
|
return cls(module, device=device, dtype=dtype)
|
|
|
|
|
|
def _make_text_encoder(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
|
module = load_component("TextEncoderLoader", spec.checkpoint, args)
|
|
tokenizer = load_component("TokenizerLoader", os.path.join(_model_root(spec), "tokenizer"), args)
|
|
device, dtype = _device(platform), _native_dtype(module)
|
|
explicit = _explicit_adapter(spec, module, device, dtype, tokenizer) # adapter cls(module, tokenizer, ...)
|
|
if explicit is not None:
|
|
return explicit
|
|
cls = Gemma if "Gemma" in type(module).__name__ else T5Encoder
|
|
return cls(module, tokenizer, device=device, dtype=dtype)
|
|
|
|
|
|
def _make_image_encoder(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
|
module = load_component("ImageEncoderLoader", spec.checkpoint, args) # CLIP vision
|
|
# the HF image processor is a sibling subfolder (BRINGUP: confirm path on a real i2v checkpoint)
|
|
processor = load_component("ImageProcessorLoader", os.path.join(_model_root(spec), "image_processor"), args)
|
|
return CLIPImageEncoder(module, processor, device=_device(platform), dtype=_native_dtype(module))
|
|
|
|
|
|
def _make_upsampler(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
|
module = load_component("UpsamplerLoader", spec.checkpoint, args) # LTX2LatentUpsampler
|
|
vae = instance.component("vae")
|
|
dtype = getattr(vae, "dtype", _native_dtype(module))
|
|
vae_module = getattr(vae, "module", None)
|
|
# per_channel_statistics lives on the AE's decoder/encoder sub-module, not the top-level autoencoder.
|
|
stats_owner = getattr(vae_module, "decoder", None) or getattr(vae_module, "encoder", None) or vae_module
|
|
return LTX2Upsampler(module, stats_owner, device=_device(platform), dtype=dtype)
|
|
|
|
|
|
def _make_audio_vae(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
|
module = load_component("AudioDecoderLoader", spec.checkpoint, args) # LTX2AudioDecoder
|
|
voc = instance.component("vocoder") # chained: decoder -> vocoder
|
|
return LTX2AudioVAE(module, voc, device=_device(platform), dtype=_native_dtype(module))
|
|
|
|
|
|
def _make_vocoder(spec: Any, instance: Any, platform: Any, args: Any) -> TorchComponent:
|
|
module = load_component("VocoderLoader", spec.checkpoint, args) # LTX2Vocoder
|
|
return LTX2Vocoder(module, device=_device(platform), dtype=_native_dtype(module))
|
|
|
|
|
|
_MAKERS = {
|
|
"dit": _make_dit,
|
|
"vae": _make_vae,
|
|
"text_encoder": _make_text_encoder,
|
|
"image_encoder": _make_image_encoder,
|
|
"upsampler": _make_upsampler,
|
|
"audio_vae": _make_audio_vae,
|
|
"vocoder": _make_vocoder,
|
|
}
|
|
|
|
|
|
def build_component(spec: Any, instance: Any, platform: Any) -> TorchComponent:
|
|
"""The single cuda component builder (registered for every kind in ``torch_cuda.py``). Shared prefix
|
|
— checkpoint check, dist-init, FastVideoArgs — then dispatch by ``spec.kind`` to its maker."""
|
|
_require_checkpoint(spec) # fail fast on a mis-stamped card, before any dist init
|
|
_ensure_fastvideo_runtime()
|
|
args = _fastvideo_args(spec)
|
|
try:
|
|
maker = _MAKERS[spec.kind]
|
|
except KeyError:
|
|
raise RuntimeError(f"torch backend: no builder for component kind {spec.kind!r} "
|
|
f"(have {sorted(_MAKERS)})") from None
|
|
comp = maker(spec, instance, platform, args)
|
|
# On-device I/O opt-in: keep this component's loop boundary on-device when the card declares it
|
|
# (and we're actually on cuda). Off by default -> numpy in/out (un-migrated families unchanged).
|
|
if isinstance(comp, TorchComponent):
|
|
comp.device_io = bool(platform.device == "cuda"
|
|
and getattr(getattr(instance, "card", None), "device_io", False))
|
|
return comp
|