# Any4D pipeline that uses Wan 2.1 Fun 1.3B InP as the base video diffusion model. # # Design: Any4DWanPipeline inherits from Any4DPipeline (which itself inherits from # Video2WorldPipeline). The whole Any4D flow (denoise / generate / apply_conditioner / # pack_streams_from_entries / ...) works unchanged because our Any4DWanDiT exposes the # same forward contract as the Cosmos-based Any4DDiT. # # The only Wan-specific concern at the pipeline layer is checkpoint loading: # the Wan 2.1 Fun 1.3B InP checkpoint ships as # diffusion_pytorch_model*.safetensors (DiT) # Wan2.1_VAE.pth (VAE) # models_t5_umt5-xxl-enc-bf16.pth (text encoder) # rather than the single Cosmos-Predict2 .pt file. We remap keys so they land on the # `self.wan.*` submodule of `Any4DWanDiT`. import os from typing import Any, Dict, Iterable, List, Optional import torch from cosmos_predict2.pipelines.video2world import Video2WorldPipeline from custom.any4d.a4d_pipe import Any4DPipeline from custom.utils.utils import get_checkpoint, get_folder, get_root_s3, _DEFAULT_LOCAL_ROOT from custom.utils.aws import aws_s3_sync from imaginaire.utils import distributed from imaginaire.utils import log # Mapping from Wan-2.1 state-dict key prefix (what ships in the .safetensors) to the path # inside our Any4DWanDiT tree. Wan's published checkpoint uses the keys directly at the # top level (e.g. `patch_embedding.weight`, `blocks.0.self_attn.q.weight`, `head.head.weight`, # `time_embedding.0.weight`, …), so we simply prefix them with `wan.`. _WAN_PREFIX = 'wan.' def _load_state_dict_from_path(path: str) -> Dict[str, torch.Tensor]: path = get_checkpoint(path) if path.endswith('.safetensors'): from safetensors.torch import load_file return load_file(path) return torch.load(path, map_location='cpu', weights_only=True) def _load_sharded_state_dict(paths: Iterable[str]) -> Dict[str, torch.Tensor]: """Load and merge multiple safetensors/pth shards into one state dict.""" merged: Dict[str, torch.Tensor] = {} for p in paths: sd = _load_state_dict_from_path(p) # Guard against accidental duplicate keys across shards. dup = set(merged) & set(sd) if dup: log.warning(f'Duplicate keys across Wan shards: {sorted(dup)[:5]}...') merged.update(sd) return merged def _strip_common_prefixes(k: str) -> str: for strip in ('net.', 'model.', 'dit.', 'module.'): if k.startswith(strip): return k[len(strip):] return k def _remap_wan_dit_state_dict(sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: """Prefix all keys with `wan.` so they land on Any4DWanDiT.wan.<...>""" return {_WAN_PREFIX + _strip_common_prefixes(k): v for k, v in sd.items()} class Any4DWanPipeline(Any4DPipeline): """Pipeline for Any4D + Wan 2.1 Fun 1.3B InP. Inherits denoise / generate / apply_conditioner / broadcast_split_for_model_parallelism from Any4DPipeline. Only checkpoint loading changes. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.any4d_active = True # ------------------------------------------------------------------ # Variable-length conditioning (generalizes Wan canonical i2v mode) # ------------------------------------------------------------------ # Rewrite masks so the rgb slot (channels [0:16]) is treated as noisy at ALL frames. # Wan's conditioning flows via the separate ref slot (channels [18:34]), not via # clean-frame replacement in the rgb slot. Without this override, the default # Any4D assemble pins rgb slot to clean x0 at cond frames, which is out-of- # distribution for the pretrained Wan DiT. # # The imask channel [16:17] (variable-length: `[1]*k + [0]*(T-k)` along T from the # dataloader) is passed through unchanged — _to_wan_36ch broadcasts it 1→4 channels # into Wan's mask_slot, and the matching ref_latent is built with gray-padding at # every non-given frame by a4d_vae_wan._build_wan_ref_latent. def _rewrite_masks_for_wan(self, x0_streams, masks) -> None: cfg = self.any4d_config for v_idx in range(cfg.num_views): v_key = f'v{v_idx}' if v_key not in x0_streams: continue entries = cfg.video_entries[v_idx] rgb_k = f'rgb{v_idx}' imask_k = f'{rgb_k}_input_mask' ref_k = f'{rgb_k}_ref' # The ref slot is optional: the Fun-InP / Fun-Control backbones carry a # reference latent at rgb{v}_ref, but the pure T2V-1.3B backbone # (wan_variant='t2v') has no ref slot — text is the only conditioning. So require # rgb + input_mask, and treat ref as present-only-if-declared. has_ref = ref_k in entries if rgb_k not in entries or imask_k not in entries: raise AssertionError( f'Any4DWanPipeline view {v_idx} missing a required entry in video_entries: ' f'{rgb_k}/{imask_k}.' ) rgb_s, rgb_e = entries[rgb_k][0], entries[rgb_k][1] imask_s, imask_e = entries[imask_k][0], entries[imask_k][1] if has_ref: ref_s, ref_e = entries[ref_k][0], entries[ref_k][1] x0 = x0_streams[v_key] B, C, T, H, W = x0.shape # Rewrite the rgb-channel masks: rgb slot must be noisy at all frames. # input_mask[rgb] = 0 (no clean pinning in the rgb slot) # output_mask[rgb] = 1 (entire rgb slot is the noise-injection target) # supervise_mask[rgb] = 1 (supervise the denoising loss at every frame to # match Wan 2.1 Fun InP training — conditioning is delivered through the # ref slot's extra channels, so the rgb latents are always predicted and # always supervised). device, dtype = x0.device, x0.dtype rgb_shape = (B, rgb_e - rgb_s, T, H, W) ones = torch.ones(rgb_shape, device=device, dtype=dtype) zeros = torch.zeros(rgb_shape, device=device, dtype=dtype) masks['input'][v_key][:, rgb_s:rgb_e] = zeros masks['output'][v_key][:, rgb_s:rgb_e] = ones masks['supervise'][v_key][:, rgb_s:rgb_e] = ones # The ref slot must stay read-through (input_mask=1, output_mask=0) so the # clean Wan reference latent is preserved by assemble. pack_streams_from_entries # already sets these correctly because rgb{N}_ref has default imputed mask = 1; # but also zero the supervise for ref so we don't accidentally include it in loss. # (Skipped for the pure T2V backbone, which has no ref slot.) if has_ref: masks['supervise'][v_key][:, ref_s:ref_e] = 0.0 # Optional edge-control slot (used as a structured-noise phase source on the # pure T2V path, not as a DiT input): keep it read-through and OUT of the loss. control_k = f'{rgb_k}_edge_control' if control_k in entries: ctl_s, ctl_e = entries[control_k][0], entries[control_k][1] masks['supervise'][v_key][:, ctl_s:ctl_e] = 0.0 # Overrides Any4DPipeline.generate solely to fix per-stream noise initialization. # The inherited implementation calls `misc.arch_invariant_rand(..., seed=seed)` with the # SAME seed across streams, producing byte-identical initial noise for v0 / v1 / v2. # Wan's joint self-attention over three identical-content view token sequences then # collapses into a degenerate state and target-frame generation comes out as abstract # hallucinations instead of video. Here we draw per-stream noise with a per-stream seed # offset so v0 / v1 / v2 start from independent samples — matching the smoke-test path. def generate(self, data_batch, seed=0, num_sampling_steps=20, guidance=0.0, sigma_max=None, sigma_min=None, cfg_zero_entries=(), is_negative_prompt=False, phase='val', cond_aug_sigma=0.0, **kwargs): import hashlib # Mixin: patch misc.arch_invariant_rand for this single call so each stream gets # its own deterministic seed derived from (seed, stream_name). from imaginaire.utils import misc _orig_rand = misc.arch_invariant_rand # Track which call (first invocation for each stream key within this generate). _call_counter = {'i': 0} _stream_keys = [k for k in data_batch.get('a4d_latent', {}).keys() if k.startswith(('rgb', 'v')) and not k.endswith('_mask') and not k.endswith('_ref')] # Any4DPipeline.generate iterates `state_shapes` (which is the x0_streams dict keys # minus {'xattn','adaln'}) to build yt_noise_start_streams. Key order is # insertion-order deterministic, so using an incrementing counter gives each # stream a unique derived seed. def _per_stream_rand(shape, dtype, device, seed=seed): idx = _call_counter['i'] _call_counter['i'] += 1 # Fold stream index into the seed via a fast hash; adds reproducibility while # ensuring distinct noise per stream. derived = (int(seed or 0) * 1_000_003 + idx * 2_654_435_761) & 0xFFFF_FFFF return _orig_rand(shape, dtype, device, seed=int(derived)) misc.arch_invariant_rand = _per_stream_rand try: return super().generate( data_batch, seed=seed, num_sampling_steps=num_sampling_steps, guidance=guidance, sigma_max=sigma_max, sigma_min=sigma_min, cfg_zero_entries=list(cfg_zero_entries), is_negative_prompt=is_negative_prompt, phase=phase, cond_aug_sigma=cond_aug_sigma, **kwargs, ) finally: misc.arch_invariant_rand = _orig_rand def _customize_initial_noise(self, yt_noise_start_streams, x0_streams, state_shapes, device, num_sampling_steps, sdedit_init_streams=None, sdedit_strength=1.0, **kwargs): """vid2vid initial-noise customization (kept out of the base Any4DPipeline loop): (1) PPD structured noise at val — phase from the clean RGB latent (setup 1, cutoff) or from the rgb{v}_edge_control slot (setup 2, source_mix + full-frequency); (2) SDEdit — start the schedule partway down from a supplied clean latent. Both are config/kwarg-gated, so the interpolation pipe (no structured noise, no sdedit) gets the noise back unchanged. Returns (yt_noise_start_streams, i_start).""" import torch # ---- (1) PPD structured noise (weather-transfer eval) ---- if getattr(self.any4d_config, 'val_structured_noise', False): from einops import rearrange as _rearrange from custom.any4d.a4d_logistics import rgb_latent_span from custom.wan.structured_noise import generate_structured_noise_subspace sm = self.scheduler.config.sigma_max # Full-frequency (cutoff=None => phase at ALL freqs) mirrors training: setup 2 is # full-frequency; setup 1 uses the fixed val cutoff. Val follows train. if bool(getattr(self.any4d_config, 'structured_noise_full_frequency', False)): cr = None else: cr = float(getattr(self.any4d_config, 'val_structured_noise_cutoff', 30.0)) K = int(getattr(self.any4d_config, 'structured_noise_channels', 16)) seed = int(getattr(self.any4d_config, 'structured_noise_seed', 0)) # Phase source: the clean rgb latent (setup 1), or the rgb{v}_edge_control slot when a # source_mix is configured (setup 2, multimodal — the source was routed into the slot). edge_phase = getattr(self.any4d_config, 'structured_noise_source_mix', None) is not None # P is the learned projection (DiT.sn_P_raw), re-orthonormalized via QR — same as # training. Under ema_scope() at val, self.dit holds EMA weights => EMA P. None only # in full-channel mode (K >= C), where the subspace projection is unused. def _qr_named(name): rw = getattr(self.dit, name, None) if rw is None: return None from torch.distributed.tensor import DTensor as _DTensor rw_full = rw.full_tensor() if isinstance(rw, _DTensor) else rw _Q, _ = torch.linalg.qr(rw_full.float()) # (C, K) orthonormal columns return _Q.t().contiguous() # (K, C) orthonormal rows P_learn = _qr_named('sn_P_raw') # noise projection (P_rgb) # PER-MODALITY P (setup 2): read this source's phase via its own P_m, lift noise via # P_rgb. Source modality is structured_noise_active_modality (set by the standalone eval # per --setup), else the dataloader's last round-robin pick. P_phase_permod = None if getattr(self.any4d_config, 'structured_noise_per_modality_p', False) and P_learn is not None: _m = getattr(self.any4d_config, 'structured_noise_active_modality', None) if not _m: try: from custom.wan.control_maps import get_last_sn_pick _m = get_last_sn_pick() except Exception: _m = None from custom.wan.structured_noise import SN_MODALITY_ALIAS _m = SN_MODALITY_ALIAS.get(_m, _m) if _m and _m != 'rgb': P_phase_permod = _qr_named(f'sn_P_raw_{_m}') for k in list(yt_noise_start_streams.keys()): if not k.startswith('v'): continue x0 = x0_streams.get(k) nz = yt_noise_start_streams[k] / sm # back to unit-variance noise if x0 is None or x0.ndim != 5 or x0.shape != nz.shape: continue Bv, Cv, Tv, Hv, Wv = x0.shape v = int(k[1:]) entries = self.any4d_config.video_entries[v] rs, re_ = rgb_latent_span(entries) # noise is injected into the rgb-latent channels # Phase source span: edge-control slot if a source_mix is set, else the rgb span. if edge_phase: ek = f'rgb{v}_edge_control' if ek not in entries: raise AssertionError( f'structured_noise_source_mix is set but view {v} has no ' f'{ek!r} entry in video_entries.') ps, pe = int(entries[ek][0]), int(entries[ek][1]) else: ps, pe = rs, re_ struct = generate_structured_noise_subspace( _rearrange(x0[:, ps:pe].float(), 'b c t h w -> (b t) c h w'), # phase source _rearrange(nz[:, rs:re_].float(), 'b c t h w -> (b t) c h w'), # rgb noise cutoff_radius=cr, num_channels=K, seed=seed, P=P_learn, P_phase=P_phase_permod, ) struct = _rearrange(struct, '(b t) c h w -> b c t h w', b=Bv, t=Tv).to(nz.dtype) nz_new = nz.clone() nz_new[:, rs:re_] = struct yt_noise_start_streams[k] = nz_new * sm # ---- (2) SDEdit: start the schedule partway down from a supplied clean latent ---- # xt = (1 - s_start) * x0_init + s_start * eps, eps = this run's (structured) unit noise. # Preserves the init latent's appearance/tone while re-denoising structure. i_start skips # the high-noise steps above s_start. sdedit_init_streams=None -> normal sampling (i_start=0). i_start = 0 if sdedit_init_streams is not None and 0.0 < sdedit_strength < 1.0: sm = self.scheduler.config.sigma_max # sdedit_strength is the START SIGMA s0 in (0,1): fraction of noise in the init. Decoupled # from the shifted step schedule -> start at the first step whose sigma <= s0. sig = self.scheduler.sigmas.to(device, dtype=torch.float32) below = (sig <= sdedit_strength).nonzero() i_start = int(below[0]) if below.numel() else (num_sampling_steps - 1) i_start = max(0, min(i_start, num_sampling_steps - 1)) s_start = sig[i_start] for k in state_shapes: if k not in sdedit_init_streams: continue eps = (yt_noise_start_streams[k].to(torch.float32) / sm) # unit-variance noise x0i = sdedit_init_streams[k].to(device=device, dtype=torch.float32) yt_noise_start_streams[k] = (1.0 - s_start) * x0i + s_start * eps log.info(f'[SDEdit] strength={sdedit_strength} -> start at step {i_start}/' f'{num_sampling_steps} (s={float(s_start):.3f})') return yt_noise_start_streams, i_start # Overrides Any4DPipeline.denoise with Wan-native flow-matching math. # No RectifiedFlowScaling, no EDM↔s translation: sigmas are Wan s ∈ [0, 1] directly, # xt = (1-s)·x0 + s·noise, DiT predicts velocity v = noise − x0, and # y0_pred = xt − s·v. See also cosmos_predict2/schedulers/wan_flow_match_scheduler.py. def denoise(self, x0_streams, yt_streams, masks, sigmas, condition, iteration=-1, phase='train'): from custom.any4d.a4d_logistics import assemble_network_inputs, assemble_clean_predictions # Rewrite rgb masks so assemble treats the rgb slot as noisy-everywhere # (conditioning flows via ref slot instead). Generalizes across variable-length # `[1]*k + [0]*(T-k)` imask patterns emitted by the dataloader. self._rewrite_masks_for_wan(x0_streams, masks) # Assemble the DiT input: clean x0 where input_mask=1 (ref slot + imask channel), # noisy yt where output_mask=1 (rgb slot). No cs_in scaling — yt_streams ARE # already the Wan xt = (1-s)·x0 + s·noise (produced by the pipeline caller). xt_assemb_streams = assemble_network_inputs(x0_streams, yt_streams, masks) xt_assemb_streams = {k: v.to(**self.tensor_kwargs) for k, v in xt_assemb_streams.items()} # Build the DiT timestep per stream: Wan's time_embedding expects s * 1000. timesteps_for_dit = {k: (s.to(torch.float32) * 1000.0).to(**self.tensor_kwargs) for k, s in sigmas.items()} # Call Wan DiT — returns velocity prediction v = noise - x0 (per stream). v_pred_streams = self.dit( x_in_streams=xt_assemb_streams, timesteps=timesteps_for_dit, **condition.to_dict(), iteration=iteration, ) # y0_pred = xt - s · v (Wan's one-line relation; no preconditioning needed). y0_raw_pred_streams = {} eps_pred_streams = {} for k, v_pred in v_pred_streams.items(): v_pred = v_pred.to(torch.float32) s = sigmas[k].to(torch.float32) yt = yt_streams[k].to(torch.float32) y0_raw_pred = yt - s * v_pred # x0 estimate eps_pred = y0_raw_pred + v_pred # since v = noise - x0 y0_raw_pred_streams[k] = y0_raw_pred eps_pred_streams[k] = eps_pred # Clean prediction: clean x0 where input_mask=1, y0_raw_pred where output_mask=1. y0_pred_streams = assemble_clean_predictions(x0_streams, y0_raw_pred_streams, masks) return (xt_assemb_streams, y0_pred_streams, eps_pred_streams, None) # Overrides Any4DPipeline.load_checkpoint (which overrides Video2WorldPipeline.load_checkpoint). def load_checkpoint(self, ckpt: dict, any4d_config=None): """ :param ckpt (dict): state dict. Two formats accepted: (a) Raw Wan Fun InP release — top-level keys like `patch_embedding.weight`, `blocks.0.self_attn.q.weight`, `head.head.weight`. We add the `wan.` prefix so they match Any4DWanDiT's module tree. (b) Any4D-saved checkpoint — keys already carry `wan.` / `view_emb_existing.` etc. Loaded as-is. Any-format `net.` or `module.` / `model.` / `dit.` prefix is stripped first by _remap_wan_dit_state_dict before detection. """ # Normalize keys (strip common wrappers like `net.` / `model.` / `dit.` / `module.`). stripped = {_strip_common_prefixes(k): v for k, v in ckpt.items()} already_any4d_format = any(k.startswith('wan.') for k in stripped.keys()) remapped = stripped if already_any4d_format else _remap_wan_dit_state_dict(ckpt) missing, surplus = self.dit.load_state_dict(remapped, strict=False, assign=True) new_parameter_prefixes = ( 'view_embs_newviews', 'view_emb_existing', 'sattn_inproj', 'sattn_outproj', 'xattn_inproj', 'action_embedder', ) unexpected_missing = [k for k in missing if not any(k.startswith(p) for p in new_parameter_prefixes) and not k.startswith('wan.')] # Note: missing may legitimately include Wan blocks/params that we replaced (head, patch_embedding) # when Any4D ch count differs — those are handled by _adapt_wan_patch_embed / _adapt_wan_head. log.warning(f'[Any4DWanPipeline] missing (sample): {missing[:8]}') log.warning(f'[Any4DWanPipeline] surplus (sample): {surplus[:8]}') log.info(f'[Any4DWanPipeline] missing: {len(missing)}, surplus: {len(surplus)}, ' f'unexpected_missing: {len(unexpected_missing)}') # ------------------------------------------------------------------ # Checkpoint loading helpers used by the model config to assemble the Wan DiT from # the multi-shard safetensors release. # ------------------------------------------------------------------ @staticmethod def gather_wan_checkpoint(dit_path: str) -> Dict[str, torch.Tensor]: """Load the Wan 2.1 Fun 1.3B InP DiT state dict from either: - a single path (.pt / .pth / .safetensors), or - a directory containing sharded diffusion_pytorch_model*.safetensors files. S3 paths are materialized to local first: a file URL via `get_checkpoint` (single-file `aws s3 cp`), a directory URL via an unconditional `aws s3 sync` (we don't use `get_folder` because its skip-if-exists shortcut can mis-fire when partial caches contain unrelated files). """ _file_exts = ('.safetensors', '.bin', '.pt', '.pth') if dit_path.startswith('s3://') and not dit_path.endswith(_file_exts): path = dit_path.replace(get_root_s3(dit_path), _DEFAULT_LOCAL_ROOT) # Idempotent: aws s3 sync only transfers missing/changed files. if distributed.is_local_rank0(): os.makedirs(path, exist_ok=True) aws_s3_sync(dit_path, path) distributed.barrier() else: path = get_checkpoint(dit_path) if os.path.isdir(path): shards = sorted( os.path.join(path, fn) for fn in os.listdir(path) if fn.startswith('diffusion_pytorch_model') and ( fn.endswith('.safetensors') or fn.endswith('.bin') or fn.endswith('.pth') ) ) if not shards: raise FileNotFoundError( f'No diffusion_pytorch_model*.safetensors in {path} — is this a Wan dir?') log.info(f'Loading Wan DiT from {len(shards)} shards in {path}') return _load_sharded_state_dict(shards) log.info(f'Loading Wan DiT from {path}') return _load_state_dict_from_path(path) # ---------------------------------------------------------------------- # Hook: extend Video2WorldPipeline.from_config to accept a Wan-style `dit_path`. # When the DiT is Any4DWanDiT, the checkpoint format is different from Cosmos, so # the default loader (which expects a single `.pt` with `net.` prefixed keys) must # be bypassed. # # We implement this by overriding `Any4DWanPipeline.from_config` (static method shared # with Video2WorldPipeline) and routing through the Wan-aware gather helper above. # ---------------------------------------------------------------------- def build_wan_pipeline( config, dit_path: str = "", text_encoder_path: str = "", vae_path: str = "", device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16, any4d_config=None, ): """Drop-in replacement for Video2WorldPipeline.from_config, but for Wan. This builds an Any4DWanPipeline, instantiates the Wan DiT from our own source (see custom/wan/a4d_network_wan.py:WanModel), and loads Wan 2.1 weights. """ from imaginaire.lazy_config import instantiate from imaginaire.utils import misc from cosmos_predict2.schedulers.wan_flow_match_scheduler import WanFlowMatchScheduler from cosmos_predict2.models.utils import init_weights_on_device from custom.any4d import a4d_surgery any4d_config = a4d_surgery.validate_config(any4d_config) # Mirror cosmos-side: data_only_disable_model implies the Wan-side remove_dit gates too. if any4d_config is not None and getattr(any4d_config, 'data_only_disable_model', False): config.remove_dit = True pipe = Any4DWanPipeline(device=device, torch_dtype=torch_dtype) pipe.any4d_config = any4d_config pipe.config = config pipe.precision = { 'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16, }[config.precision] pipe.tensor_kwargs = {'device': 'cuda', 'dtype': pipe.precision} # Data keys + Wan-native flow-matching scheduler. # No preconditioning scaling is needed: Any4DWanPipeline.denoise implements Wan's # xt = (1-s)·x0 + s·noise, v = noise - x0, y0 = xt - s·v directly (s-native, no EDM). pipe.sigma_data = config.sigma_data pipe.setup_data_key() pipe.scheduler = WanFlowMatchScheduler( shift=5.0, sigma_min=0.003 / 1.002, sigma_max=1.0, num_train_timesteps=1000, extra_one_step=True, ) pipe.scaling = None # Wan pipe doesn't use preconditioning; denoise() does raw flow-match. # Tokenizer = Wan VAE adapter. Instantiated via LazyCall registered in model.py. # Resolve S3 vae_path to a local cached path before instantiating (mirrors what # Video2WorldPipeline.from_config does for the Cosmos tokenizer at video2world.py:350). if not config.remove_dit: config.tokenizer.vae_pth = get_checkpoint(vae_path) pipe.tokenizer = instantiate(config.tokenizer) assert pipe.tokenizer.latent_ch == pipe.config.state_ch, \ f'Wan VAE latent_ch {pipe.tokenizer.latent_ch} != state_ch {pipe.config.state_ch}' # Move the actual Wan VAE weights to the target dtype + device (the default is # float32 on CPU after load). Without this, forward with bf16 inputs hits # "Input type (BFloat16) and bias type (float) should be the same". pipe.tokenizer.reset_dtype() pipe.tokenizer.to(device=device) # Text encoder: UMT5-XXL for per-sample prompts. When text_encoder_path is set, # it should point to the weights .pth file (or S3 path). The tokenizer is resolved # from the sibling google/umt5-xxl/ directory relative to the weights. if text_encoder_path and len(text_encoder_path) > 10: from custom.wan.wan_text_encoder import WanTextEncoder # Resolve tokenizer path: try sibling google/umt5-xxl/ relative to weights weights_resolved = get_checkpoint(text_encoder_path) parent = os.path.dirname(weights_resolved) tokenizer_path = os.path.join(parent, 'google', 'umt5-xxl') # If the sibling dir isn't already local, sync it from the corresponding S3 prefix. if not os.path.isdir(tokenizer_path) and text_encoder_path.startswith('s3://'): s3_parent = os.path.dirname(text_encoder_path) tokenizer_path = get_folder(f'{s3_parent}/google/umt5-xxl') if not os.path.isdir(tokenizer_path): # Fallback: look relative to dit_path (which for Wan 2.2 14B is a directory prefix). if dit_path and dit_path.endswith(('.pt', '.pth', '.safetensors', '.bin')): dit_resolved = get_checkpoint(dit_path) dit_parent = os.path.dirname(dit_resolved) elif dit_path: dit_parent = os.path.dirname(get_folder(dit_path)) else: dit_parent = '' tokenizer_path = os.path.join(dit_parent, 'google', 'umt5-xxl') if not os.path.isdir(tokenizer_path) and dit_path and dit_path.startswith('s3://'): s3_dit_parent = os.path.dirname(dit_path) tokenizer_path = get_folder(f'{s3_dit_parent}/google/umt5-xxl') assert os.path.isdir(tokenizer_path), ( f'Cannot find UMT5 tokenizer directory. Tried google/umt5-xxl/ relative to ' f'{parent} and {dit_parent if dit_path else "N/A"}. ' f'Set text_encoder_path="" to use default_text_wan.pkl instead.') with misc.timer('load WanTextEncoder'): pipe.text_encoder = WanTextEncoder( weights_path=text_encoder_path, tokenizer_path=tokenizer_path, device=device, dtype=torch_dtype) log.success(f'[Wan] Loaded UMT5-XXL text encoder from {text_encoder_path}') else: log.warning('[Wan] Skipping UMT5 text encoder; will use default_text_wan.pkl') pipe.text_encoder = None # Conditioner is pure logistics (remaps keys, applies text dropout), works as-is. if not config.remove_dit: pipe.conditioner = instantiate(config.conditioner) pipe.prompt_refiner = None pipe.text_guardrail_runner = None pipe.video_guardrail_runner = None # DiT: instantiate Any4DWanDiT via LazyCall `config.net`. if not config.remove_dit: with init_weights_on_device(device=torch.device('cpu')): dit_config = config.net dit_config.any4d_config = any4d_config with misc.timer('instantiate Any4DWanDiT'): pipe.dit = instantiate(dit_config).eval() # Load the Wan 2.1 DiT weights. if dit_path: state_dict = Any4DWanPipeline.gather_wan_checkpoint(dit_path) with misc.timer('Any4DWanDiT load_checkpoint'): pipe.load_checkpoint(state_dict, any4d_config=any4d_config) del state_dict log.success(f'Loaded Wan 2.1 DiT from {dit_path}') # The DiT is built under init_weights_on_device(device='cpu') then populated with # assign=True load_state_dict — after which its params are on CPU. Move to GPU # (FSDP will reshard later if distributed; single-GPU inference just uses cuda). if not config.remove_dit: pipe.dit = pipe.dit.to(device=device, dtype=pipe.precision) pipe.dit.eval() # Activation checkpointing on each Wan DiTBlock (applied before FSDP wraps the # DiT in Predict2Video2WorldModel.__init__). Required to fit Wan 2.2 14B on # H100 80GB; ~30% slower step in exchange for re-computing block activations # during backward instead of storing them all. if getattr(any4d_config, 'wan_activation_checkpoint', False): from cosmos_predict2.utils.fsdp_helper import apply_fsdp_checkpointing from custom.wan.a4d_network_wan import DiTBlock apply_fsdp_checkpointing(pipe.dit, [DiTBlock]) log.info('[Wan] Activation checkpointing enabled on DiTBlock') # EMA: build a frozen copy of the DiT for the training loop to exponentially average # into (Karras EDM2 power-function EMA). This MUST exist before the inherited # Predict2Video2WorldModel.__init__ -> Video2WorldPipeline.apply_fsdp runs, so apply_fsdp # shards dit_ema and installs the DTensor EMA worker. on_before_zero_grad then updates it # each step, ema_scope() swaps it in for validation, and state_dict() persists it under # net_ema.* keys. The copy here only seeds dit_ema from the base DiT; the checkpointer # reseeds it from loaded weights on resume/init. Mirrors build_wan_control_pipeline so the # plain Wan path (Fun-InP / pure T2V) supports ema_enabled=True too. Off => dit_ema=None. if config.ema.enabled and not config.remove_dit: import numpy as np from imaginaire.utils.ema import FastEmaModelUpdater pipe.dit_ema = instantiate(config.net).eval() pipe.dit_ema.requires_grad_(False) pipe.dit_ema = pipe.dit_ema.to(device=device, dtype=pipe.precision) pipe.dit_ema_worker = FastEmaModelUpdater() # swapped to DTensor variant by apply_fsdp s = config.ema.rate pipe.ema_exp_coefficient = np.roots([1, 7, 16 - s**-2, 12 - s**-2]).real.max() pipe.dit_ema_worker.copy_to(src_model=pipe.dit, tgt_model=pipe.dit_ema) log.success(f'Built EMA DiT (rate={s}, ema_exp_coefficient={pipe.ema_exp_coefficient:.4f})') else: pipe.dit_ema = None # Data parallel size (copied from Video2WorldPipeline.from_config). from megatron.core import parallel_state if parallel_state.is_initialized(): pipe.data_parallel_size = parallel_state.get_data_parallel_world_size() else: pipe.data_parallel_size = 1 return pipe