DreamX-Creator 1.0 on ZeroGPU: vendored videox_fun + dreamx_inference from AMAP-ML upstream; generate(image, prompt)->(mp4, last-frame PNG, seed), neutral keyframe when image empty, DREAMX_CKPT_DIR for persistent checkpoints, diffusers 0.37.1 stack

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,6 +1,8 @@
 # Local model / cache artifacts
 *.gguf
 *.safetensors
+*.pth
+/checkpoints/
 /tmp/
 .cache/
 __pycache__/

README.md

@@ -1,14 +1,30 @@
 ---
-title: Video Creator
-emoji: 🔥
+title: DreamX Creator
+emoji: 🎬
 colorFrom: green
 colorTo: blue
 sdk: docker
 app_file: app.py
 pinned: false
+hardware:
+  gpu: zero
 ---
 
-Docker-hosted LTX-2.5 Space for the AI Shorts Factory backend. Four Gradio
-endpoints via `/call/{fn}` (text_to_video, text_to_video_av, image_to_video,
-image_to_video_av). Requires a GPU tier — the free CPU tier (2 vCPU/16GB)
-boots but generation OOMs.
+Docker-hosted DreamX-Creator 1.0 Space for the AI Shorts Factory backend
+(engine vendored from `AMAP-ML/DreamX-Creator` / `hugging-apps/gd-ml-dreamx-creator`).
+
+One endpoint via Gradio's `/call/generate` protocol:
+
+- `generate(image, prompt, seconds, num_inference_steps, resolution_tokens,
+  guidance_scale, negative_prompt, seed, randomize_seed)`
+  → `(mp4 with jointly-denoised audio, last-frame PNG, seed)`.
+
+`image` may be empty — a neutral keyframe is generated in-app. The last-frame
+PNG is returned so each scene can feed the next one (continuity chaining).
+
+Weights (~43 GB peak) download at first boot from the ungated
+`GD-ML/DreamX-Creator` repo into `./checkpoints` (override with
+`DREAMX_CKPT_DIR`, e.g. `/data/dreamx` when persistent storage is enabled).
+
+Requires **ZeroGPU** hardware (`hardware: gpu: zero` above and/or the Space
+Settings toggle) — the free CPU tier cannot denoise.

app.py

@@ -1,443 +1,588 @@
-"""LTX-2.5 Space — ZeroGPU inference API for the AI Shorts Factory backend.
+"""DreamX-Creator 1.0 — native joint audio-video generation on ZeroGPU.
 
-Runs the distilled LTX-2.5 GGUF (Q4_K_M) on zero GPU as a Gradio demo. Four
-endpoints are exposed through Gradio's `/call/{fn}` protocol and consumed by
-the backend provider `app.ai.visuals.ltx`:
+Image + prompt -> a video whose soundtrack is denoised jointly with the frames
+by the same model (gated A2V / V2A cross-attention), so speech, foley and
+ambience stay in sync with the picture.
 
-    text_to_video(prompt, width, height, num_frames, enhance_prompt)
-    text_to_video_av(...)            # video + synchronized 48kHz narration audio
-    image_to_video(prompt, image_url, width, height, num_frames, enhance_prompt)
-    image_to_video_av(...)           # video + synchronized narration audio
+The inference path is the authors' own release code (`videox_fun/` +
+`dreamx_inference.py`, copied verbatim from the reference Space
+`hugging-apps/gd-ml-dreamx-creator` / `AMAP-ML/DreamX-Creator`); this file only
+wires it into Gradio and stages the checkpoint download so the container never
+holds all 43 GB of fp32 weights on disk at once.
 
-Every function returns ``(video_path, first_frame_path)``. Scene 2..N of the
-pipeline call image_to_video with the last frame URL of the previous clip so
-the subject stays continuous (last-frame chaining).
+Adapted for the AI Shorts Factory backend (Space `text_amon_API`):
 
-The backend always sends an already-composed prompt (quoted dialogue for the
-exact script words) and passes ``enhance_prompt=False``; the prompt enhancer
-is a hard no-op here so no rewrite happens.
-"""
+- `image` is optional: an empty value yields a neutral keyframe (Option A), so
+  scenes without a first-frame image can still generate.
+- Returns ``(mp4, last-frame PNG, seed)`` instead of ``(mp4, seed)``: the last
+  frame is what scene 2..N sends back as the first frame for continuity.
+- Checkpoint root can be redirected to persistent storage with
+  ``DREAMX_CKPT_DIR`` (default: this repo's ``./checkpoints``, the proven
+  upstream layout).
 
-from __future__ import annotations
+Served through Gradio's `/call/generate` protocol (``api_name="generate"``).
+"""
 
-import glob
 import os
-import subprocess
-import threading
-from pathlib import Path
-
-# ZeroGPU rule #1: `import spaces` must precede any CUDA-touching import
-# (torch etc.) — it monkey-patches torch.cuda at import time.
-import spaces  # noqa: E402
-
-import numpy as np
-import torch
-
-# Requires the Space secret HF_TOKEN after accepting the LTX-2.x Community
-# License (free for entities under $10M annual revenue). Needed for every file
-# served from the gated Lightricks/LTX-2.5-Diffusers repo (config, Gemma-4-12B
-# text encoder, VAEs, connectors, audio_vae + vocoder).
-HF_TOKEN = os.environ.get("HF_TOKEN", "").strip()
-
-# Direct download URL of the distilled GGUF transformer (~15GB). The backend
-# only needs LTX_SPACE_URL; this key is configured in the Space. A GGUF the
-# user uploads into the repo `model/` folder is preferred over this URL.
-LTX_MODEL_URL = os.environ.get(
-    "LTX_MODEL_URL",
-    "https://huggingface.co/realrebelai/LTX-2.5_GGUFs/resolve/main/LTX-2.5-Distilled-Q4_K_M.gguf",
-)
-PIPELINE_ID = "Lightricks/LTX-2.5-Diffusers"
-# Overridable repo id (e.g. a locally mirrored copy: LTX_PIPELINE_ID=/app/ltx25).
-LTX_PIPELINE_ID = os.environ.get("LTX_PIPELINE_ID", PIPELINE_ID)
-SEED = int(os.environ.get("LTX_SEED", "0"))
-FPS = int(os.environ.get("LTX_FPS", "24"))
 
-# Distilled model: guidance is baked into the weights, few-step schedule.
-INFERENCE_STEPS = int(os.environ.get("LTX_STEPS", "8"))
-GUIDANCE_SCALE = float(os.environ.get("LTX_GUIDANCE", "1.0"))
-NEGATIVE_PROMPT = "low quality, blurry, distorted, watermark"
-# Honest worst-case GPU wall-time per generation. ZeroGPU validates this
-# after a 1.5x multiplier and sits a ~60-120s continuous execution cap per
-# call on the free tier — 120 keeps us inside both (120 * 1.5 = 180 < 300
-# cap, and matches the real per-call window after the model is resident).
-GPU_DURATION = int(os.environ.get("LTX_GPU_DURATION", "120"))
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+# videox_fun falls back to torch SDPA when flash-attn is absent; the numerics are
+# identical here because every sequence in the batch is full-length (no padding).
+os.environ.setdefault("VIDEOX_ATTENTION_TYPE", "FLASH_ATTENTION")
 
-_lock = threading.Lock()
-_model = None
+import spaces  # noqa: E402  (must precede torch)
 
+import math  # noqa: E402
+import random  # noqa: E402
+import shutil  # noqa: E402
+import subprocess  # noqa: E402
+import tempfile  # noqa: E402
+import time  # noqa: E402
+from pathlib import Path  # noqa: E402
+from types import SimpleNamespace  # noqa: E402
 
-# ----------------------------------------------------------------- model load
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
 
+# Trusted upstream .pth checkpoints (UMT5 encoder, Wan2.2 VAE) are plain tensor
+# dicts saved before torch 2.6 flipped `weights_only` to True.
+_torch_load = torch.load
 
-def _find_local_gguf() -> str | None:
-    """Pick a GGUF the user uploaded into the repo (web UI ``model/`` folder)."""
-    for folder in ("model", "models", "weights"):
-        hits = [p for p in glob.glob(os.path.join(folder, "**", "*.gguf"), recursive=True) if os.path.isfile(p)]
-        if hits:
-            return sorted(hits)[0]
-    return None
 
+def _torch_load_compat(*args, **kwargs):
+    kwargs.setdefault("weights_only", False)
+    return _torch_load(*args, **kwargs)
 
-def _gguf_fingerprint(path: str) -> None:
-    """Fail fast if the GGUF is not an LTX-2 video transformer.
-
-    diffusers 0.40 converts GGUFs that keep the ComfyUI/native
-    ``model.diffusion_model.*`` tensor names (llama.cpp-style ``blk.*`` renames
-    are not supported by the LTX-2 GGUF converter).
-    """
-    from gguf import GGUFReader
 
-    reader = GGUFReader(path)
-    names = [tensor.name for tensor in reader.tensors]
-    print(f"[ltx] gguf tensors={len(names)}")
-
-    def count(marker: str) -> int:
-        return sum(1 for name in names if marker in name)
-
-    blocks = count("transformer_blocks.")
-    ltx2_av_gate = count("model.diffusion_model.av_ca_a2v_gate_adaln_single")
-    llama_blk = count("blk.") + count(".attn_qkv.")
-    print(
-        f"[ltx] gguf fingerprint blocks={blocks} ltx2_av_gate={ltx2_av_gate} llama_blk={llama_blk}"
-    )
-    if blocks == 0 and llama_blk == 0:
-        raise RuntimeError(
-            "GGUF does not look like an LTX-2 video transformer (no transformer_blocks found)"
-        )
-    if blocks == 0 and llama_blk > 0:
-        raise RuntimeError(
-            "GGUF uses llama.cpp-renamed tensors (blk.*); diffusers 0.40's LTX-2 GGUF "
-            "converter expects the native model.diffusion_model.* names. Use a GGUF "
-            "converted from the ComfyUI/native checkpoint instead."
-        )
-
-
-def _download_gguf(dest: str) -> str:
-    import httpx
-
-    if LTX_MODEL_URL.startswith("http"):
-        if os.path.exists(dest):
-            return dest
-        with httpx.stream("GET", LTX_MODEL_URL, follow_redirects=True, timeout=600) as r:
-            r.raise_for_status()
-            with open(dest, "wb") as fh:
-                for chunk in r.iter_bytes(chunk_size=1 << 20):
-                    fh.write(chunk)
-        return dest
-    from huggingface_hub import hf_hub_download
-
-    # Support "repo_id:filename" shorthand.
-    repo, _, filename = LTX_MODEL_URL.partition(":")
-    return hf_hub_download(repo, filename or "LTX-2.5-Distilled-Q4_K_M.gguf", token=HF_TOKEN or None)
-
-
-def _ensure_gguf() -> str:
-    """Local vendored GGUF first, then network download."""
-    local = _find_local_gguf()
-    if local is not None:
-        print(f"[ltx] using vendored GGUF: {local}")
-        return local
-    print("[ltx] no local GGUF found; downloading…")
-    return _download_gguf(os.environ.get("LTX_GGUF_CACHE", "/tmp/ltx-model.gguf"))
+torch.load = _torch_load_compat
 
+import gradio as gr  # noqa: E402
+from diffusers import FlowMatchEulerDiscreteScheduler  # noqa: E402
+from huggingface_hub import snapshot_download  # noqa: E402
+from omegaconf import OmegaConf  # noqa: E402
+from PIL import Image  # noqa: E402
+from transformers import AutoTokenizer  # noqa: E402
+
+from dreamx_inference import (  # noqa: E402
+    DEFAULT_NEGATIVE_PROMPT,
+    DirectionalMultimodalCFGAdapter,
+    filter_kwargs,
+    generate_joint_audio_video,
+)
+from videox_fun.models import AutoencoderKLWan3_8, WanT5EncoderModel  # noqa: E402
+from videox_fun.models.creator_dac_vae import CreatorDACVAE  # noqa: E402
+from videox_fun.models.creator_gating import WanCreatorGatingAVModel  # noqa: E402
+
+REPO_ID = "GD-ML/DreamX-Creator"
+APP_DIR = Path(__file__).parent.resolve()
+CKPT_DIR = Path(os.environ.get("DREAMX_CKPT_DIR", str(APP_DIR / "checkpoints")))
+CONFIG_PATH = APP_DIR / "config" / "config.yaml"
+WEIGHT_DTYPE = torch.bfloat16
+FPS = 24
+CACHE_VERSION = 1
+
+# Authors' defaults (audio_video_generation/inference.py + inference.sh).
+VIDEO_BRIDGE_GUIDANCE = 3.5
+AUDIO_BRIDGE_GUIDANCE = 3.5
+VIDEO_SHIFT = 5.0
+AUDIO_SHIFT = 5.0
+
+RESOLUTION_CHOICES = [
+    ("Fast — ~360p", 220),
+    ("Balanced — ~480p", 440),
+    ("Sharp — ~600p", 660),
+]
+
+# Neutral keyframe (Option A) dimensions — portrait 9:16, shorts-first.
+KEYFRAME_WIDTH = int(os.environ.get("DREAMX_KEYFRAME_WIDTH", "720"))
+KEYFRAME_HEIGHT = int(os.environ.get("DREAMX_KEYFRAME_HEIGHT", "1280"))
+
+
+def _log(msg: str) -> None:
+    print(f"[dreamx] {msg}", flush=True)
+
+
+def _disk() -> str:
+    total, used, free = shutil.disk_usage("/")
+    return f"disk used={used / 2**30:.1f}GB free={free / 2**30:.1f}GB"
+
+
+# --------------------------------------------------------------------------- #
+# Model loading (module scope, eagerly on "cuda" — ZeroGPU packs from here)
+# --------------------------------------------------------------------------- #
+
+_cfg = OmegaConf.load(CONFIG_PATH)
+_video_kwargs = OmegaConf.to_container(_cfg["video_transformer_additional_kwargs"], resolve=True)
+_audio_kwargs = OmegaConf.to_container(_cfg["audio_transformer_additional_kwargs"], resolve=True)
+_gating_kwargs = OmegaConf.to_container(_cfg["creator_gating_kwargs"], resolve=True)
+_video_vae_kwargs = OmegaConf.to_container(_cfg["video_vae_kwargs"], resolve=True)
+_text_encoder_kwargs = OmegaConf.to_container(_cfg["text_encoder_kwargs"], resolve=True)
+_scheduler_kwargs = OmegaConf.to_container(_cfg["scheduler_kwargs"], resolve=True)
+MAX_SEQUENCE_LENGTH = int(_text_encoder_kwargs.get("text_length", 512))
+
+_log(f"checkpoints in {CKPT_DIR} ({_disk()})")
+_log(f"downloading joint AV generator ... ({_disk()})")
+snapshot_download(REPO_ID, local_dir=str(CKPT_DIR), allow_patterns=["creator/*", "creator/**/*"])
+_log(f"loading joint AV generator ... ({_disk()})")
+
+transformer = WanCreatorGatingAVModel.from_pretrained(
+    pretrained_model_path=str(CKPT_DIR / "creator"),
+    video_pretrained_model_path=str(CKPT_DIR / "creator" / "video_model"),
+    audio_pretrained_model_path=str(CKPT_DIR / "creator" / "audio_model"),
+    video_subfolder=_video_kwargs.get("transformer_low_noise_model_subpath", None),
+    audio_subfolder=_audio_kwargs.get("transformer_low_noise_model_subpath", None),
+    video_kwargs=_video_kwargs,
+    audio_kwargs=_audio_kwargs,
+    low_cpu_mem_usage=True,
+    torch_dtype=WEIGHT_DTYPE,
+    use_temporal_rope=_gating_kwargs.get("use_temporal_rope", True),
+    audio_fps=_gating_kwargs.get("audio_fps", 48000.0 / 960.0),
+    vae_temporal_stride=_gating_kwargs.get("vae_temporal_stride", 4),
+    a2v_cross_attn_layers=_gating_kwargs.get("a2v_cross_attn_layers", None),
+    v2a_cross_attn_layers=_gating_kwargs.get("v2a_cross_attn_layers", None),
+    use_gating=_gating_kwargs.get("use_gating", True),
+    zero_init_cross_attn=_gating_kwargs.get("zero_init_cross_attn", False),
+    zero_init_gating=_gating_kwargs.get("zero_init_gating", True),
+    gate_init_value=_gating_kwargs.get("gate_init_value", 0.0),
+    a2v_gate_alphas=_gating_kwargs.get("a2v_gate_alphas", None),
+    v2a_gate_alphas=_gating_kwargs.get("v2a_gate_alphas", None),
+)
+transformer.eval()
+# fp32 source shards are no longer needed once the bf16 model is in RAM.
+shutil.rmtree(CKPT_DIR / "creator", ignore_errors=True)
+_log(f"joint AV generator loaded ({_disk()})")
+
+_log("downloading VAEs + UMT5-xxl text encoder ...")
+snapshot_download(
+    REPO_ID,
+    local_dir=str(CKPT_DIR),
+    allow_patterns=["audio_vae/*", "wan2.2_ti2v_5b/*", "wan2.2_ti2v_5b/**/*"],
+)
+_log(f"loading VAEs + text encoder ... ({_disk()})")
 
-def _build_transformer(gguf_path: str):
-    try:
-        # diffusers 0.40 (pinned): the LTX-2 video transformer is loadable from
-        # a single GGUF file via FromOriginalModelMixin.
-        from diffusers import LTX2VideoTransformer3DModel as TransformerCls
-    except ImportError:  # older versions only expose the generic AutoModel
-        from diffusers import AutoModel as TransformerCls
+_wan_dir = CKPT_DIR / "wan2.2_ti2v_5b"
+_video_vae_path = _wan_dir / _video_vae_kwargs.get("vae_subpath", "Wan2.2_VAE.pth")
+video_vae = AutoencoderKLWan3_8.from_pretrained(
+    str(_video_vae_path), additional_kwargs=_video_vae_kwargs
+).eval()
+audio_vae = CreatorDACVAE.from_pretrained(str(CKPT_DIR / "audio_vae"), strict=False).eval()
 
+tokenizer = AutoTokenizer.from_pretrained(
+    str(_wan_dir / _text_encoder_kwargs.get("tokenizer_subpath", "google/umt5-xxl"))
+)
+_text_encoder_path = _wan_dir / _text_encoder_kwargs.get(
+    "text_encoder_subpath", "models_t5_umt5-xxl-enc-bf16.pth"
+)
+text_encoder = WanT5EncoderModel.from_pretrained(
+    str(_text_encoder_path),
+    additional_kwargs=_text_encoder_kwargs,
+    low_cpu_mem_usage=True,
+    torch_dtype=WEIGHT_DTYPE,
+).eval()
+
+for _stale in (_video_vae_path, _text_encoder_path):
     try:
-        # GGUF config lives in quantizers since 0.33; newer releases keep it there.
-        from diffusers.quantizers.quantization_config import GGUFQuantizationConfig
-    except ImportError:  # legacy export path
-        from diffusers.utils import GGUFQuantizationConfig
-
-    kwargs = dict(
-        quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
-        dtype=torch.bfloat16,
-    )
-    _gguf_fingerprint(gguf_path)
-    token = HF_TOKEN or None
-    try:
-        # `config=` pulls transformer/config.json from the gated repo (token).
-        return TransformerCls.from_single_file(gguf_path, config=LTX_PIPELINE_ID, token=token, **kwargs)
-    except Exception as exc:  # config may live elsewhere; trust the GGUF KV metadata
-        print(f"[ltx] from_single_file with config failed ({exc}), retrying without")
-        return TransformerCls.from_single_file(gguf_path, token=token, **kwargs)
-
-
-def _load_model() -> dict:
-    global _model
-    if _model is not None:
-        return _model
-
-    with _lock:
-        if _model is not None:
-            return _model
-
-        if not HF_TOKEN and not os.path.isdir(LTX_PIPELINE_ID):
-            raise RuntimeError(
-                "HF_TOKEN secret is not set. Open the Space → Settings → Variables and "
-                "secrets and create a secret named HF_TOKEN: a fine-grained token with "
-                "read access to Lightricks/LTX-2.5-Diffusers (after accepting its "
-                "LTX-2.x Community License on that repo)."
-            )
+        os.remove(_stale)
+    except OSError:
+        pass
+shutil.rmtree(CKPT_DIR / "audio_vae", ignore_errors=True)
 
-        print("[ltx] downloading GGUF transformer…")
-        gguf_path = _ensure_gguf()
+transformer = DirectionalMultimodalCFGAdapter(
+    transformer,
+    video_scale=VIDEO_BRIDGE_GUIDANCE,
+    audio_scale=AUDIO_BRIDGE_GUIDANCE,
+    enable_a2v=True,
+    enable_v2a=True,
+).eval()
 
-        print("[ltx] building quantized transformer…")
-        transformer = _build_transformer(gguf_path)
+transformer.to("cuda")
+text_encoder.to("cuda")
+video_vae.to("cuda")
+audio_vae.to("cuda")
+_log(f"all models resident on cuda ({_disk()})")
 
-        from diffusers import LTX2ImageToVideoPipeline, LTX2Pipeline
 
-        built = {
-            "transformer": transformer,
-        }
-        token = HF_TOKEN or None
-        try:
-            built["t2v"] = LTX2Pipeline.from_pretrained(
-                LTX_PIPELINE_ID, transformer=transformer, torch_dtype=torch.bfloat16, token=token
-            )
-            built["i2v"] = LTX2ImageToVideoPipeline.from_pretrained(
-                LTX_PIPELINE_ID, transformer=transformer, torch_dtype=torch.bfloat16, token=token
-            )
-        except Exception as exc:  # pragma: no cover - component layout differences
-            msg = str(exc)
-            hint = ""
-            if "401" in msg or "is not a valid model identifier" in msg or "gated" in msg.lower():
-                hint = " (hint: accept the LTX-2.x Community License on the repo + set the HF_TOKEN secret)"
-            raise RuntimeError(
-                f"[ltx] pipeline build failed - check Lightricks/LTX-2.5-Diffusers "
-                f"component layout. {msg}{hint}"
-            ) from exc
+# --------------------------------------------------------------------------- #
+# Helpers
+# --------------------------------------------------------------------------- #
 
-        for pipe in (built["t2v"], built["i2v"]):
-            pipe.enable_model_cpu_offload()
-        _model = built
-        print("[ltx] model ready")
-        return _model
+def _latent_frames(duration: float) -> int:
+    num_frames = int(duration * FPS)
+    num_frames = int((num_frames - 1) // 4 * 4) + 1
+    return (num_frames - 1) // 4 + 1
 
 
-# ZeroGPU rule #2: load the model at module scope so weights are packed once
-# at startup and resident for every forked worker, instead of per-call.
-# `_generate` still guards with `_load_model()` so early calls block on the
-# lock until this background warm-up finishes.
-_warmup = threading.Thread(target=_load_model, name="ltx-warmup", daemon=True)
-_warmup.start()
+DURATION_CAP = 300
 
 
-# ----------------------------------------------------------------- generation
+def _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens) -> float:
+    """Predicted GPU seconds, fit to two measured ZeroGPU runs.
 
+    Measured on RTX PRO 6000 (sm_120), bf16, 3-branch directional multimodal CFG:
+      2640 video tokens x 10 steps -> 10.77s denoise, 16.0s total
+      7560 video tokens x 30 steps -> 95.04s denoise, 104.7s total
+    """
+    tokens = _latent_frames(float(seconds)) * int(resolution_tokens)
+    # per denoising step: linear (FFN/proj) + quadratic (self-attention) terms
+    per_step = 4.02e-4 * tokens + 2.256e-9 * tokens * tokens
+    denoise = per_step * int(num_inference_steps)
+    overhead = 4.5 + 7.0e-4 * tokens  # T5 encode, first-frame VAE encode, decode, mux
+    return denoise + overhead
+
+
+def _estimate_duration(
+    image=None,
+    prompt="",
+    seconds=3.0,
+    num_inference_steps=30,
+    resolution_tokens=440,
+    *args,
+    **kwargs,
+):
+    """GPU seconds to reserve — calibrated against measured runs on ZeroGPU."""
+    raw = _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens)
+    return int(min(DURATION_CAP, math.ceil(raw * 1.15) + 5))
+
+
+def _write_mp4(frames: np.ndarray, audio: np.ndarray, sample_rate: int, fps: int) -> str:
+    """Mux uint8 RGB frames + mono float audio into a single H.264/AAC mp4."""
+    height, width = frames.shape[1], frames.shape[2]
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    wav_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
+
+    import soundfile as sf
+
+    sf.write(wav_path, np.clip(audio, -1.0, 1.0), sample_rate)
+
+    base = [
+        "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
+        "-f", "rawvideo", "-pix_fmt", "rgb24",
+        "-s", f"{width}x{height}", "-r", str(fps), "-i", "-",
+        "-i", wav_path,
+    ]
+    for vcodec in ("libx264", "mpeg4"):
+        cmd = base + [
+            "-c:v", vcodec, "-pix_fmt", "yuv420p", "-crf", "18",
+            "-c:a", "aac", "-b:a", "192k", "-shortest", out_path,
+        ]
+        if vcodec == "mpeg4":
+            cmd.remove("-crf")
+            cmd.remove("18")
+        proc = subprocess.run(cmd, input=frames.tobytes(), capture_output=True)
+        if proc.returncode == 0 and os.path.getsize(out_path) > 0:
+            os.remove(wav_path)
+            return out_path
+        _log(f"ffmpeg ({vcodec}) failed: {proc.stderr.decode()[-600:]}")
+    raise gr.Error("ffmpeg failed to encode the generated video.")
+
+
+def _neutral_keyframe(width: int = KEYFRAME_WIDTH, height: int = KEYFRAME_HEIGHT) -> str:
+    """Deterministic dark diagonal gradient — the "Option A" neutral first frame.
+
+    By default a 9:16 portrait so generated clips keep a shorts-friendly aspect
+    ratio when no real first-frame image is provided.
+    """
+    y, x = np.mgrid[0:height, 0:width]
+    norm = np.sqrt((x / max(width - 1, 1)) ** 2 + (y / max(height - 1, 1)) ** 2)
+    tone = (0.06 + 0.16 * norm[..., None] * np.array([0.86, 0.95, 1.0])).clip(0, 1)
+    frame = (tone * 255.0).astype(np.uint8)
+    path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
+    Image.fromarray(frame).save(path)
+    return path
+
+
+# --------------------------------------------------------------------------- #
+# Inference
+# --------------------------------------------------------------------------- #
+
+
+@spaces.GPU(duration=_estimate_duration)
+def generate(
+    image: str,
+    prompt: str,
+    seconds: float = 3.0,
+    num_inference_steps: int = 30,
+    resolution_tokens: int = 440,
+    guidance_scale: float = 5.0,
+    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
+    seed: int = 113,
+    randomize_seed: bool = False,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate a video with a jointly-denoised soundtrack from a first frame and a prompt.
+
+    Args:
+        image: path to the image used as the first frame of the video. Empty
+            means "no image": a neutral keyframe is generated in-app instead.
+        prompt: description of the action AND the sound to generate; put spoken
+            lines in quotes (e.g. `Man says, 'hello there.'`).
+        seconds: length of the clip in seconds (24 fps).
+        num_inference_steps: number of flow-matching denoising steps.
+        resolution_tokens: spatial token budget; higher means higher resolution.
+        guidance_scale: classifier-free guidance strength for text.
+        negative_prompt: what to avoid in both the video and the audio.
+        seed: RNG seed for reproducible sampling.
+        randomize_seed: draw a fresh random seed instead of using `seed`.
+
+    Returns:
+        A tuple of (path to the generated mp4 with audio, path to the last-frame
+        PNG for chaining into the next scene, the seed actually used).
+    """
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please provide a prompt describing the motion and the sound.")
+    if not image:
+        _log("empty first-frame image -> neutral keyframe")
+        image = _neutral_keyframe()
+
+    # ZeroGPU kills a task that outruns its reservation, and the reservation is
+    # capped, so refuse the few extreme knob combinations that cannot fit.
+    if _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens) * 1.15 + 5 > DURATION_CAP:
+        raise gr.Error(
+            f"That combination needs more than the {DURATION_CAP}s GPU budget. "
+            "Lower the duration, the resolution, or the number of steps."
+        )
 
-def _try_call(pipe, **kwargs):
-    """Call the pipeline, disabling the (Gemma-4 based) prompt enhancer."""
-    try:
-        return pipe(enhance_prompt=False, **kwargs)
-    except TypeError:
-        return pipe(**kwargs)
-
-
-def _extract_frames(output) -> list:
-    frames = output.frames
-    if isinstance(frames[0], (list, tuple)):
-        return list(frames[0])
-    return list(frames)
-
-
-def _extract_audio(output):
-    audio = getattr(output, "audio", None)
-    if audio is None:
-        return None
-    if isinstance(audio, (list, tuple)):
-        arr, sr = audio[0], audio[1] if len(audio) > 1 else 48000
-    else:
-        arr, sr = audio, 48000
-    if torch.is_tensor(arr):
-        arr = arr.detach().float().cpu().numpy()
-    arr = np.asarray(arr)
-    if arr.ndim == 3:  # (batch, channels, samples)
-        arr = arr[0]
-    if arr.ndim == 2:  # collapse to mono; channel axis is the smaller dimension
-        arr = arr.mean(axis=int(arr.shape[0] < arr.shape[1]))
-    return np.ascontiguousarray(arr.astype(np.float32)), int(sr)
-
-
-def _mux_audio(video_path, frames, fps, audio) -> str:
-    from diffusers.utils import export_to_video
-
-    export_to_video(frames, video_path, fps=fps)
-    if audio is None:
-        return video_path
-
-    import imageio_ffmpeg
-    from scipy.io import wavfile
-
-    arr, sr = audio
-    vp = Path(video_path)
-    wav_path = vp.with_suffix(".wav")
-    wavfile.write(str(wav_path), sr, arr)
-    ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
-    muxed = str(vp.with_name(vp.stem + "_mux.mp4"))
-    subprocess.run(
-        [
-            ffmpeg, "-y",
-            "-i", str(video_path),
-            "-i", str(wav_path),
-            "-map", "0:v", "-map", "1:a",
-            "-c:v", "copy", "-c:a", "aac", "-shortest",
-            muxed,
-        ],
-        check=True,
-        capture_output=True,
+    if randomize_seed:
+        seed = random.randint(0, 2**31 - 1)
+    seed = int(seed)
+
+    device = torch.device("cuda")
+    video_scheduler_kwargs = dict(_scheduler_kwargs, shift=VIDEO_SHIFT)
+    audio_scheduler_kwargs = dict(_scheduler_kwargs, shift=AUDIO_SHIFT)
+
+    models = {
+        "config": _cfg,
+        "transformer": transformer,
+        "video_vae": video_vae,
+        "audio_vae": audio_vae,
+        "tokenizer": tokenizer,
+        "text_encoder": text_encoder,
+        # fresh schedulers per request: they carry mutable step state
+        "video_scheduler": FlowMatchEulerDiscreteScheduler(
+            **filter_kwargs(FlowMatchEulerDiscreteScheduler, video_scheduler_kwargs)
+        ),
+        "audio_scheduler": FlowMatchEulerDiscreteScheduler(
+            **filter_kwargs(FlowMatchEulerDiscreteScheduler, audio_scheduler_kwargs)
+        ),
+        "max_sequence_length": MAX_SEQUENCE_LENGTH,
+    }
+
+    args = SimpleNamespace(
+        config_path=str(CONFIG_PATH),
+        image=image,
+        output="output.mp4",
+        negative_prompt=negative_prompt or DEFAULT_NEGATIVE_PROMPT,
+        duration=float(seconds),
+        target_spatial_tokens=int(resolution_tokens),
+        min_token_ratio=0.95,
+        fps=FPS,
+        num_inference_steps=int(num_inference_steps),
+        guidance_scale=float(guidance_scale),
+        cfg_mode="multimodal",
+        video_bridge_guidance_scale=VIDEO_BRIDGE_GUIDANCE,
+        audio_bridge_guidance_scale=AUDIO_BRIDGE_GUIDANCE,
+        seed=seed,
+        video_shift=VIDEO_SHIFT,
+        audio_shift=AUDIO_SHIFT,
+        flow_match_mu=None,
+        sampler_name="Flow",
+        weight_dtype="bfloat16",
+        GPU_memory_mode="model_full_load",
+        text_encoder_cpu_offload=False,
+        video_vae_cpu_offload=False,
+        audio_vae_cpu_offload=False,
+        vae_cpu_offload=False,
+        use_temporal_rope=True,
+        audio_fps=48000.0 / 960.0,
+        vae_temporal_stride=4,
+        disable_a2v_cross_attn=False,
+        disable_v2a_cross_attn=False,
+        suppress_aux_writes=True,
+        skip_output_decode=False,
+        disable_progress=False,
+        synchronize_noise=False,
+        ulysses_degree=1,
+        ring_degree=1,
+        fsdp_dit=False,
     )
-    return muxed
-
-
-def _load_image(image_url: str):
-    if not image_url:
-        return None
-    import base64
-    import io
-
-    from PIL import Image, ImageOps
-
-    if image_url.startswith("data:"):
-        encoded = image_url.partition(",")[2]
-        return Image.open(io.BytesIO(base64.b64decode(encoded))).convert("RGB")
-
-    import httpx
-
-    data = httpx.get(image_url, timeout=120).content
-    image = Image.open(io.BytesIO(data)).convert("RGB")
-    return ImageOps.exif_transpose(image)
-
-
-def _generate(prompt, image_url, width, height, num_frames, enhance_prompt, with_audio):
-    if enhance_prompt:
-        # Never let the enhancer rewrite the scripted dialogue.
-        enhance_prompt = False
 
-    num_frames = max(9, int(num_frames))
-    num_frames = 1 + (num_frames - 1 - ((num_frames - 1) % 8))  # 1 + 8n
-    width = int(width) // 32 * 32
-    height = int(height) // 32 * 32
-
-    model = _load_model()
-    generator = torch.Generator(device="cuda").manual_seed(SEED)
-
-    image = _load_image(image_url)
-    common = dict(
-        prompt=prompt,
-        negative_prompt=NEGATIVE_PROMPT,
-        width=width,
-        height=height,
-        num_frames=num_frames,
-        num_inference_steps=INFERENCE_STEPS,
-        guidance_scale=GUIDANCE_SCALE,
-        generator=generator,
+    item = {
+        "prompt": prompt.strip(),
+        "video_prompt": prompt.strip(),
+        "audio_prompt": prompt.strip(),
+        "negative_prompt": args.negative_prompt,
+        "audio_negative_prompt": args.negative_prompt,
+        "duration": args.duration,
+        "guidance_scale": args.guidance_scale,
+        "num_inference_steps": args.num_inference_steps,
+        "seed": seed,
+        "first_frame_path": image,
+        "name": "sample",
+    }
+
+    started = time.perf_counter()
+    video_decoded, audio_decoded, num_frames = generate_joint_audio_video(
+        args, models, device, WEIGHT_DTYPE, item
     )
-    if image is not None:
-        output = _try_call(model["i2v"], image=image, **common)
-    else:
-        output = _try_call(model["t2v"], **common)
-
-    frames = _extract_frames(output)
-    audio = _extract_audio(output) if with_audio else None
-
-    video_path = f"/tmp/gradio/ltx_{threading.get_ident()}.mp4"
-    muxed = _mux_audio(video_path, frames, FPS, audio)
-
-    first_frame = f"/tmp/gradio/ltx_{threading.get_ident()}_first.png"
-    frames[0].save(first_frame)
-    return muxed, first_frame
-
-
-# ------------------------------------------------------------------ gradio UI
-
+    gpu_seconds = time.perf_counter() - started
 
-def text_to_video(prompt, width, height, num_frames, enhance_prompt=False):
-    return _generate(prompt, "", width, height, num_frames, enhance_prompt, with_audio=False)
+    frames = (
+        video_decoded[0].permute(1, 2, 3, 0).clamp(0, 1).numpy() * 255.0
+    ).astype(np.uint8)
+    waveform = audio_decoded.detach().float().cpu()
+    while waveform.ndim > 1:
+        waveform = waveform[0]
+    audio = waveform.numpy()
 
+    out_path = _write_mp4(frames, audio, int(audio_vae.sample_rate), FPS)
 
-def text_to_video_av(prompt, width, height, num_frames, enhance_prompt=False):
-    return _generate(prompt, "", width, height, num_frames, enhance_prompt, with_audio=True)
+    last_frame_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
+    Image.fromarray(frames[-1]).save(last_frame_path)
 
-
-def image_to_video(prompt, image_url, width, height, num_frames, enhance_prompt=False):
-    return _generate(prompt, image_url, width, height, num_frames, enhance_prompt, with_audio=False)
-
-
-def image_to_video_av(prompt, image_url, width, height, num_frames, enhance_prompt=False):
-    return _generate(prompt, image_url, width, height, num_frames, enhance_prompt, with_audio=True)
-
-
-import gradio as gr  # noqa: E402
-
-t2v_fns = [spaces.GPU(duration=GPU_DURATION)(text_to_video), spaces.GPU(duration=GPU_DURATION)(text_to_video_av)]
-i2v_fns = [spaces.GPU(duration=GPU_DURATION)(image_to_video), spaces.GPU(duration=GPU_DURATION)(image_to_video_av)]
-
-with gr.Blocks(title="LTX-2.5 Shorts Space") as demo:
-    gr.Markdown(
-        "# LTX-2.5 Shorts Space\n"
-        "ZeroGPU inference for the AI Shorts Factory backend. "
-        "Portrait 9:16 clips (~4s) with optional synchronized narration audio. "
-        "The prompt enhancer stays **off** so quoted dialogue is spoken exactly."
+    _log(
+        f"done in {gpu_seconds:.1f}s | {num_frames} frames @ {frames.shape[2]}x{frames.shape[1]} "
+        f"| steps={args.num_inference_steps} tokens={args.target_spatial_tokens} "
+        f"| reserved={_estimate_duration(image, prompt, seconds, num_inference_steps, resolution_tokens)}s"
     )
+    return out_path, last_frame_path, seed
+
+
+# --------------------------------------------------------------------------- #
+# UI
+# --------------------------------------------------------------------------- #
+
+EXAMPLES = [
+    [
+        "examples/case1.jpg",
+        "A man in a dark suit and white shirt is seated on a yellow couch, speaking about "
+        "the language of Americans. He uses hand gestures to emphasize his points, and the "
+        "background shows a cityscape with illuminated buildings, suggesting an urban setting, "
+        "possibly a studio with a city view. Man says, 'The thing about Americans that I've "
+        "thought about the language is that they speak, they say they speak English.'.",
+    ],
+    [
+        "examples/case4.jpg",
+        "The video shows a wolf standing on its hind legs, howling with its mouth wide open, "
+        "showing its teeth and tongue. The wolf's ears are perked up, and its eyes are focused "
+        "on something in the distance. The background consists of trees with green and yellow "
+        "leaves, indicating it is autumn. The wolf's howl is loud and resonant, filling the air "
+        "with its powerful voice. The sound of a dog howls and howls can be heard.",
+    ],
+    [
+        "examples/case3.jpg",
+        "The video captures a dramatic night scene with a series of lightning strikes "
+        "illuminating the dark sky and revealing the city lights below. The clouds move across "
+        "the sky, and the lightning strikes again, followed by a thunderclap. The sound of a "
+        "thunderstorm and rain falling can be heard.",
+    ],
+    [
+        "examples/case2.jpg",
+        "A humanoid robot is cooking in a modern kitchen. The robot, with a white body and blue "
+        "eyes, is stirring food in a pan on the stove. The kitchen is equipped with dark cabinets "
+        "and various utensils. The robot's movements are smooth and precise as it stirs the food, "
+        "causing steam to rise. The ambient sound of cooking can be heard throughout the scene.",
+    ],
+    [
+        "examples/case5.jpg",
+        "A person in a red plaid shirt is typing on a white keyboard placed on a wooden table. "
+        "The scene is set in a room with a wooden floor and a part of a white blanket visible in "
+        "the background. The person's hands are actively moving across the keyboard, indicating "
+        "typing activity. The ambient sound is the distinct sound of keys being pressed.",
+    ],
+    [
+        "examples/case6.jpg",
+        "A man is playing an acoustic guitar in a modern kitchen setting. He is focused on his "
+        "playing, with his hands moving along the strings and fretboard. The background features "
+        "a well-lit kitchen with wooden cabinets and hanging lights.",
+    ],
+]
+
+CSS = """
+#col-container { max-width: 1180px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
 
-    width = gr.Slider(256, 768, value=544, step=32, label="Width (÷32)")
-    height = gr.Slider(256, 1408, value=960, step=32, label="Height (÷32)")
-    num_frames = gr.Slider(9, 801, value=97, step=8, label="Frames (8n+1)")
-    enhance = gr.Checkbox(value=False, label="Enhance prompt (kept off)")
-    video_out = gr.Video(label="Clip")
-    frame_out = gr.Image(label="First frame (for chaining)")
-
-    with gr.Tab("Text → Video"):
-        t2v_prompt = gr.Textbox(label="Prompt", lines=4)
-        t2v_btn = gr.Button("Generate (video only)")
-        t2v_btn.click(
-            t2v_fns[0],
-            inputs=[t2v_prompt, width, height, num_frames, enhance],
-            outputs=[video_out, frame_out],
+with gr.Blocks(title="DreamX-Creator 1.0") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            "# 🎬 DreamX-Creator 1.0\n"
+            "Turn **one image + one prompt** into a video whose **soundtrack is generated "
+            "jointly with the frames** — speech, foley and ambience come out of the same "
+            "denoiser as the picture, so they stay in sync.\n\n"
+            "Describe the *sound* as well as the action. For speech, quote the line: "
+            "`Man says, 'hello there.'`\n\n"
+            "[Model](https://huggingface.co/GD-ML/DreamX-Creator) · "
+            "[Code](https://github.com/AMAP-ML/DreamX-Creator)"
         )
 
-    with gr.Tab("Text → Video + Audio"):
-        t2va_prompt = gr.Textbox(label="Prompt", lines=4)
-        t2va_btn = gr.Button("Generate (with narration)")
-        t2va_btn.click(
-            t2v_fns[1],
-            inputs=[t2va_prompt, width, height, num_frames, enhance],
-            outputs=[video_out, frame_out],
-        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                image = gr.Image(label="First frame (optional — leave empty for a neutral keyframe)", type="filepath", height=320)
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Describe the motion and the sound you want to hear…",
+                    lines=4,
+                )
+                run = gr.Button("Generate audio + video", variant="primary")
+            with gr.Column(scale=1):
+                video_out = gr.Video(label="Result (video + generated audio)", height=380)
+                last_frame_out = gr.Image(label="Last frame (feeds the next scene)", height=380)
+
+        with gr.Accordion("Advanced settings", open=False):
+            with gr.Row():
+                seconds = gr.Slider(
+                    label="Duration (seconds)", minimum=2.0, maximum=5.0, step=0.5, value=3.0
+                )
+                num_inference_steps = gr.Slider(
+                    label="Denoising steps", minimum=10, maximum=50, step=1, value=30
+                )
+            with gr.Row():
+                resolution_tokens = gr.Dropdown(
+                    label="Resolution budget (spatial tokens)",
+                    choices=RESOLUTION_CHOICES,
+                    value=440,
+                )
+                guidance_scale = gr.Slider(
+                    label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=5.0
+                )
+            negative_prompt = gr.Textbox(
+                label="Negative prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=2
+            )
+            with gr.Row():
+                seed = gr.Number(label="Seed", value=113, precision=0)
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
 
-    with gr.Tab("Image → Video"):
-        i2v_prompt = gr.Textbox(label="Prompt", lines=4)
-        i2v_image = gr.Textbox(
-            label="First-frame image URL (last frame from the previous clip)",
-            value="",
-        )
-        i2v_btn = gr.Button("Generate (video only)")
-        i2v_btn.click(
-            i2v_fns[0],
-            inputs=[i2v_prompt, i2v_image, width, height, num_frames, enhance],
-            outputs=[video_out, frame_out],
+        gr.Markdown(
+            "Longer clips, more steps and a bigger resolution budget all cost GPU time "
+            "roughly linearly (and attention grows quadratically with tokens × frames)."
         )
 
-    with gr.Tab("Image → Video + Audio"):
-        i2va_prompt = gr.Textbox(label="Prompt", lines=4)
-        i2va_image = gr.Textbox(
-            label="First-frame image URL (last frame from the previous clip)",
-            value="",
-        )
-        i2va_btn = gr.Button("Generate (with narration)")
-        i2va_btn.click(
-            i2v_fns[1],
-            inputs=[i2va_prompt, i2va_image, width, height, num_frames, enhance],
-            outputs=[video_out, frame_out],
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[image, prompt],
+            outputs=[video_out, last_frame_out, seed],
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+            label=f"Official Verse-Bench cases from the DreamX-Creator repo (v{CACHE_VERSION})",
         )
 
+    run.click(
+        fn=generate,
+        inputs=[
+            image,
+            prompt,
+            seconds,
+            num_inference_steps,
+            resolution_tokens,
+            guidance_scale,
+            negative_prompt,
+            seed,
+            randomize_seed,
+        ],
+        outputs=[video_out, last_frame_out, seed],
+        api_name="generate",
+    )
+
 if __name__ == "__main__":
-    demo.queue(default_concurrency_limit=1).launch(server_name="0.0.0.0")
+    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)

config/config.yaml

@@ -0,0 +1,54 @@
+video_transformer_additional_kwargs:
+  transformer_low_noise_model_subpath: ./
+  transformer_combination_type: single
+  dict_mapping:
+    in_dim: in_channels
+    dim: hidden_size
+
+audio_transformer_additional_kwargs:
+  transformer_low_noise_model_subpath: .
+  patch_size: [1]
+  in_dim: 128
+  out_dim: 128
+  vae_type: dac
+
+video_vae_kwargs:
+  vae_type: AutoencoderKLWan3_8
+  vae_subpath: Wan2.2_VAE.pth
+  temporal_compression_ratio: 4
+  spatial_compression_ratio: 16
+
+text_encoder_kwargs:
+  text_encoder_subpath: models_t5_umt5-xxl-enc-bf16.pth
+  tokenizer_subpath: google/umt5-xxl
+  text_length: 512
+  vocab: 256384
+  dim: 4096
+  dim_attn: 4096
+  dim_ffn: 10240
+  num_heads: 64
+  num_layers: 24
+  num_buckets: 32
+  shared_pos: false
+  dropout: 0.0
+
+scheduler_kwargs:
+  scheduler_subpath: null
+  num_train_timesteps: 1000
+  shift: 5.0
+  use_dynamic_shifting: false
+  base_shift: 0.5
+  max_shift: 1.15
+  base_image_seq_len: 256
+  max_image_seq_len: 4096
+
+creator_gating_kwargs:
+  use_temporal_rope: true
+  audio_fps: 50.0
+  vae_temporal_stride: 4
+  a2v_cross_attn_layers: [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
+  v2a_cross_attn_layers: [15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
+  use_gating: true
+  zero_init_cross_attn: false
+  zero_init_gating: false
+  gate_init_value: 0.5

dreamx_inference.py

@@ -0,0 +1,1030 @@
+"""Single-image Creator audio-video inference.
+
+The supplied image is encoded as the first video frame. The prompt conditions
+the joint Creator video/audio denoiser, which writes a video, a WAV, and a merged
+MP4. This entrypoint intentionally has no dataset, Verse-Bench, training, or
+model-download workflow.
+"""
+
+import argparse
+import inspect
+import math
+import os
+import sys
+import time
+import subprocess
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
+from PIL import Image
+from transformers import AutoTokenizer as HFAutoTokenizer
+try:  # optional: only used by this module's standalone CLI entrypoint
+    from torchvision.io import write_video
+except Exception:  # pragma: no cover
+    write_video = None
+from torch import nn
+from tqdm.auto import tqdm
+
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+
+current_file_path = os.path.abspath(__file__)
+release_root = os.path.dirname(current_file_path)
+if release_root not in sys.path:
+    sys.path.insert(0, release_root)
+
+from videox_fun.models import AutoencoderKLWan3_8, WanT5EncoderModel
+from videox_fun.models.creator_gating import WanCreatorGatingAVModel
+from videox_fun.models.creator_dac_vae import CreatorDACVAE
+
+
+def filter_kwargs(cls, kwargs):
+    """Keep scheduler options compatible across diffusers versions."""
+    valid = set(inspect.signature(cls.__init__).parameters)
+    return {key: value for key, value in kwargs.items() if key in valid}
+
+
+class DirectionalMultimodalCFGAdapter(nn.Module):
+    """Preserve the experiment's three-branch multimodal CFG in one call."""
+
+    def __init__(self, model, video_scale, audio_scale, enable_a2v, enable_v2a):
+        super().__init__()
+        self.model = model
+        self.video_scale = float(video_scale)
+        self.audio_scale = float(audio_scale)
+        self.enable_a2v = bool(enable_a2v)
+        self.enable_v2a = bool(enable_v2a)
+
+    @property
+    def video_patch_size(self):
+        return self.model.video_patch_size
+
+    @property
+    def audio_patch_size(self):
+        return self.model.audio_patch_size
+
+    @staticmethod
+    def _expand(value, name):
+        if isinstance(value, torch.Tensor):
+            if value.ndim == 0 or value.shape[0] != 2:
+                raise ValueError(f"Multimodal CFG expects {name} batch size 2")
+            return torch.cat((value[0:1], value[0:1], value[1:2]), dim=0)
+        if isinstance(value, (list, tuple)) and len(value) == 2:
+            return [value[0], value[0], value[1]]
+        raise TypeError(f"Unsupported {name} batch value: {type(value)!r}")
+
+    @classmethod
+    def _expand_inputs(cls, inputs):
+        expanded = dict(inputs)
+        for key in ("x", "t", "context", "y", "clip_fea"):
+            if expanded.get(key) is not None:
+                expanded[key] = cls._expand(expanded[key], key)
+        return expanded
+
+    @staticmethod
+    def _collapse(prediction, bridge_scale):
+        if prediction.shape[0] != 3:
+            raise ValueError("Creator multimodal CFG expects three model predictions")
+        d00, d0b, dtb = prediction[0:1], prediction[1:2], prediction[2:3]
+        bridge = d00 + bridge_scale * (d0b - d00)
+        return torch.cat((bridge, bridge + (dtb - d0b)), dim=0)
+
+    def forward(self, video, audio, dtype=torch.bfloat16, return_dict=True, **kwargs):
+        if self.model.training:
+            raise RuntimeError("Inference adapter received a training model")
+        model_video = self._expand_inputs(video)
+        model_audio = self._expand_inputs(audio)
+        device = model_video["t"].device
+        bridge_mask = torch.tensor([False, True, True], device=device, dtype=torch.bool)
+        output = self.model(
+            video=model_video,
+            audio=model_audio,
+            dtype=dtype,
+            return_dict=True,
+            enable_a2v=bridge_mask & self.enable_a2v,
+            enable_v2a=bridge_mask & self.enable_v2a,
+        )
+        result = {
+            "video": self._collapse(output["video"], self.video_scale),
+            "audio": self._collapse(output["audio"], self.audio_scale),
+        }
+        return result if return_dict else (result["video"], result["audio"])
+
+try:
+    import soundfile as sf
+except ImportError:
+    sf = None
+
+
+# ============================================================================
+    # Single-process helpers
+# ============================================================================
+
+def print_info(*args, **kwargs):
+    print(*args, **kwargs)
+
+
+def synchronize_tensor(tensor: torch.Tensor, args) -> torch.Tensor:
+    """Broadcast a rank-0 tensor when the SP launcher requests exact inputs."""
+    if getattr(args, "synchronize_noise", False) and dist.is_initialized():
+        dist.broadcast(tensor, src=0)
+    return tensor
+
+
+# ============================================================================
+# Argument parsing
+# ============================================================================
+
+DEFAULT_NEGATIVE_PROMPT = (
+    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指"
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Single-image Creator audio-video inference")
+
+    # Model paths
+    parser.add_argument("--config_path", type=str, default=os.path.join(os.path.dirname(__file__), "config/config.yaml"))
+    parser.add_argument("--model_name", type=str, required=True,
+                        help="Wan2.2-TI2V-5B model directory")
+    parser.add_argument("--transformer_path", type=str, required=True,
+                        help="Creator checkpoint directory with video_model/ and audio_model/")
+    parser.add_argument("--audio_vae_path", type=str, required=True,
+                        help="Path to CreatorDACVAE audio VAE")
+
+    # Input / Output
+    parser.add_argument("--image", type=str, required=True,
+                        help="Input image used as the first video frame")
+    parser.add_argument("--prompt", type=str, required=True, help="Text prompt")
+    parser.add_argument("--output", type=str, default="./outputs/output.mp4",
+                        help="Output MP4; a WAV with the same stem is also written")
+    parser.add_argument("--negative_prompt", type=str, default=DEFAULT_NEGATIVE_PROMPT)
+    # Generation parameters
+    parser.add_argument("--duration", type=float, default=5.0,
+                        help="Target duration in seconds")
+    parser.add_argument("--target_spatial_tokens", type=int, default=880,
+                        help="Maximum spatial tokens per frame")
+    parser.add_argument("--min_token_ratio", type=float, default=0.95,
+                        help="Minimum fraction of target spatial tokens")
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--num_inference_steps", type=int, default=50)
+    parser.add_argument("--guidance_scale", type=float, default=5.0)
+    parser.add_argument("--cfg_mode", type=str, default="multimodal",
+                        choices=["text", "multimodal"],
+                        help="CFG mode. 'multimodal' uses Creator dual CFG with separate "
+                             "bridge and text guidance (3 model evaluations per step).")
+    parser.add_argument("--video_bridge_guidance_scale", type=float, default=3.5)
+    parser.add_argument("--audio_bridge_guidance_scale", type=float, default=3.5)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--video_shift", type=float, default=5.0,
+                        help="Noise schedule shift for video denoising")
+    parser.add_argument("--audio_shift", type=float, default=5.0,
+                        help="Noise schedule shift for audio denoising")
+
+    # Sampler
+    parser.add_argument("--sampler_name", type=str, default="Flow", choices=["Flow"])
+
+    # Memory & compute
+    parser.add_argument("--weight_dtype", type=str, default="bfloat16",
+                        choices=["float16", "bfloat16", "float32"])
+    parser.add_argument("--GPU_memory_mode", type=str, default="model_full_load",
+                        choices=["model_full_load", "model_cpu_offload"])
+    offload_action = argparse.BooleanOptionalAction
+    parser.add_argument(
+        "--text_encoder_cpu_offload",
+        action=offload_action,
+        default=None,
+        help="Keep the text encoder on CPU except while encoding prompts. "
+             "Defaults to the GPU_memory_mode setting.",
+    )
+    parser.add_argument(
+        "--video_vae_cpu_offload",
+        action=offload_action,
+        default=None,
+        help="Keep the video VAE on CPU except while encoding/decoding. "
+             "Defaults to the GPU_memory_mode setting.",
+    )
+    parser.add_argument(
+        "--audio_vae_cpu_offload",
+        action=offload_action,
+        default=None,
+        help="Keep the audio VAE on CPU except while decoding. "
+             "Defaults to the GPU_memory_mode setting.",
+    )
+    parser.add_argument(
+        "--vae_cpu_offload",
+        action=offload_action,
+        default=None,
+        help="Enable or disable CPU offload for both video and audio VAEs. "
+             "Individual VAE flags take precedence.",
+    )
+
+    # Temporal RoPE
+    parser.add_argument("--use_temporal_rope", type=bool, default=True)
+    parser.add_argument("--audio_fps", type=float, default=48000.0 / 960.0,
+                        help="Audio latent FPS (DAC: 48000/960=50)")
+    parser.add_argument("--vae_temporal_stride", type=int, default=4)
+
+    # Runtime cross-attention controls
+    parser.add_argument("--disable_a2v_cross_attn", "--disable-a2v-cross-attn",
+                        "--disable_a2v", "--disable-a2v", action="store_true",
+                        help="Disable audio-to-video cross attention at inference time.")
+    parser.add_argument("--disable_v2a_cross_attn", "--disable-v2a-cross-attn",
+                        "--disable_v2a", "--disable-v2a", action="store_true",
+                        help="Disable video-to-audio cross attention at inference time.")
+
+    args = parser.parse_args()
+    return args
+
+
+# ============================================================================
+# Device / dtype helpers
+# ============================================================================
+
+def init_device() -> torch.device:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if device.type == "cuda":
+        torch.cuda.set_device(0)
+    return device
+
+
+def resolve_weight_dtype(dtype_arg: str, device: torch.device) -> torch.dtype:
+    if device.type == "cpu":
+        return torch.float32
+    if dtype_arg == "float16":
+        return torch.float16
+    if dtype_arg == "bfloat16":
+        return torch.bfloat16
+    return torch.float32
+
+
+def resolve_cpu_offload_flags(args):
+    """Resolve dedicated offload flags, preserving the legacy memory mode."""
+    legacy_offload = args.GPU_memory_mode == "model_cpu_offload"
+
+    def resolve(value):
+        return legacy_offload if value is None else bool(value)
+
+    vae_override = getattr(args, "vae_cpu_offload", None)
+
+    def resolve_vae(value):
+        return resolve(vae_override if value is None else value)
+
+    return {
+        "transformer": legacy_offload,
+        "text_encoder": resolve(getattr(args, "text_encoder_cpu_offload", None)),
+        "video_vae": resolve_vae(getattr(args, "video_vae_cpu_offload", None)),
+        "audio_vae": resolve_vae(getattr(args, "audio_vae_cpu_offload", None)),
+    }
+
+
+def clear_cuda_cache(device: torch.device):
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+
+
+# ============================================================================
+# Text encoding
+# ============================================================================
+
+def _get_t5_prompt_embeds(tokenizer, text_encoder, prompt, max_sequence_length, device, dtype):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=max_sequence_length,
+        truncation=True,
+        add_special_tokens=True,
+        return_tensors="pt",
+    )
+    text_input_ids = text_inputs.input_ids
+    prompt_attention_mask = text_inputs.attention_mask
+    seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
+
+    prompt_embeds = text_encoder(
+        text_input_ids.to(device),
+        attention_mask=prompt_attention_mask.to(device),
+    )[0]
+    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+    return [embed[:seq_len] for embed, seq_len in zip(prompt_embeds, seq_lens.tolist())]
+
+
+def encode_prompt(tokenizer, text_encoder, prompt, negative_prompt, guidance_scale,
+                  max_sequence_length, device, dtype):
+    prompt_embeds = _get_t5_prompt_embeds(
+        tokenizer, text_encoder, prompt, max_sequence_length, device, dtype)
+
+    if guidance_scale <= 1.0:
+        return prompt_embeds
+
+    negative_prompt = negative_prompt or ""
+    negative_prompt_embeds = _get_t5_prompt_embeds(
+        tokenizer, text_encoder, negative_prompt, max_sequence_length, device, dtype)
+    return negative_prompt_embeds + prompt_embeds
+
+
+# ============================================================================
+# Latent shape helpers
+# ============================================================================
+
+def compute_audio_latent_length(duration: float, audio_vae: CreatorDACVAE) -> int:
+    """Compute audio latent time length from duration.
+
+    For DAC VAE: latent_T = ceil(duration * sample_rate / hop_length)
+    """
+    sample_rate = audio_vae.sample_rate
+    hop_length = audio_vae.hop_length
+    num_samples = int(duration * sample_rate)
+    latent_T = math.ceil(num_samples / hop_length)
+    return latent_T
+
+
+def get_audio_num_tokens(latent_T: int, patch_size: tuple) -> int:
+    """Compute number of audio tokens after patching."""
+    return latent_T // patch_size[0]
+
+
+def compute_video_latent_shape(duration, fps, vae_temporal_ratio=4, vae_spatial_ratio=8,
+                               height=480, width=832):
+    """Compute video latent dimensions from duration and resolution."""
+    num_frames = int(duration * fps)
+    num_frames = int((num_frames - 1) // vae_temporal_ratio * vae_temporal_ratio) + 1
+    latent_frames = (num_frames - 1) // vae_temporal_ratio + 1
+    latent_height = height // vae_spatial_ratio
+    latent_width = width // vae_spatial_ratio
+    return num_frames, latent_frames, latent_height, latent_width
+
+
+# ============================================================================
+# Audio saving
+# ============================================================================
+
+def save_audio_wav(waveform: torch.Tensor, sample_rate: int, output_path: str):
+    """Save waveform tensor to wav file.
+
+    Args:
+        waveform: [B, 1, T] or [1, T] or [T] tensor
+        sample_rate: audio sample rate
+        output_path: path to save wav file
+    """
+    waveform = waveform.cpu().float()
+    if waveform.ndim == 3:
+        waveform = waveform.squeeze(0)  # [1, T]
+    if waveform.ndim == 1:
+        waveform = waveform.unsqueeze(0)  # [1, T]
+
+    if sf is not None:
+        sf.write(output_path, waveform.transpose(0, 1).numpy(), sample_rate)
+        return
+
+    try:
+        import torchaudio
+        torchaudio.save(output_path, waveform, sample_rate)
+        return
+    except ImportError:
+        pass
+
+    import wave
+    waveform = waveform.clamp(-1.0, 1.0)
+    waveform_int16 = (waveform * 32767.0).to(torch.int16).transpose(0, 1).contiguous().numpy()
+    with wave.open(output_path, "wb") as wav_file:
+        wav_file.setnchannels(int(waveform_int16.shape[1]))
+        wav_file.setsampwidth(2)
+        wav_file.setframerate(sample_rate)
+        wav_file.writeframes(waveform_int16.tobytes())
+
+
+# ============================================================================
+# Scheduler helpers
+# ============================================================================
+
+def retrieve_timesteps(scheduler, num_inference_steps=None, device=None,
+                       timesteps=None, sigmas=None, **kwargs):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed.")
+    if timesteps is not None:
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        return scheduler.timesteps, len(scheduler.timesteps)
+    elif sigmas is not None:
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        return scheduler.timesteps, len(scheduler.timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        return scheduler.timesteps, num_inference_steps
+
+
+def prepare_extra_step_kwargs(scheduler, eta=0.0):
+    extra_step_kwargs = {}
+    if "eta" in set(inspect.signature(scheduler.step).parameters.keys()):
+        extra_step_kwargs["eta"] = eta
+    return extra_step_kwargs
+
+
+def resolve_scheduler_shifts(args):
+    """Resolve separate AV shifts while supporting callers that only define --shift."""
+    video_shift = getattr(args, "video_shift", 5.0)
+    audio_shift = getattr(args, "audio_shift", 5.0)
+    return float(video_shift), float(audio_shift)
+
+
+# ============================================================================
+# Image to latent encoding
+# ============================================================================
+
+def encode_first_frame(image: Image.Image, vae, device, dtype):
+    """Encode first frame image into VAE latent space.
+
+    Returns latent of shape [1, C, 1, H//8, W//8].
+    """
+    image = image.convert("RGB")
+    img_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0
+    img_tensor = img_tensor * 2.0 - 1.0  # normalize to [-1, 1]
+    img_tensor = img_tensor.unsqueeze(0).unsqueeze(2)  # [1, 3, 1, H, W]
+    img_tensor = img_tensor.to(device=device, dtype=dtype)
+    vae = vae.to(dtype=dtype, device=device)
+    with torch.no_grad():
+        posterior = vae.encode(img_tensor)[0]
+        latent = posterior.sample()
+    return latent  # [1, C, 1, H//8, W//8]
+
+
+def _budget_spatial_size(
+    target_spatial_tokens: int,
+    source_height: int,
+    source_width: int,
+    spatial_divisor_h: int,
+    spatial_divisor_w: int,
+    min_token_ratio: float = 0.95,
+):
+    """Pick a near-aspect-ratio size within the spatial-token budget."""
+    if target_spatial_tokens <= 0:
+        raise ValueError("target_spatial_tokens must be positive")
+    if source_height <= 0 or source_width <= 0:
+        raise ValueError(f"Invalid source size: {(source_height, source_width)}")
+    if not 0.0 < min_token_ratio <= 1.0:
+        raise ValueError("min_token_ratio must be in (0, 1]")
+
+    min_spatial_tokens = max(1, int(math.ceil(target_spatial_tokens * min_token_ratio)))
+    source_ratio = float(source_height) / float(source_width)
+    best = None
+    for token_height in range(1, target_spatial_tokens + 1):
+        max_token_width = target_spatial_tokens // token_height
+        if max_token_width < 1:
+            continue
+        ideal_token_width = (
+            source_width * token_height * spatial_divisor_h
+            / (source_height * spatial_divisor_w)
+        )
+        for token_width in {
+            1,
+            max_token_width,
+            int(math.floor(ideal_token_width)),
+            int(math.ceil(ideal_token_width)),
+        }:
+            if not 1 <= token_width <= max_token_width:
+                continue
+            spatial_tokens = token_height * token_width
+            height = token_height * spatial_divisor_h
+            width = token_width * spatial_divisor_w
+            aspect_error = abs(math.log((height / width) / source_ratio))
+            score = (
+                max(0, min_spatial_tokens - spatial_tokens),
+                aspect_error,
+                target_spatial_tokens - spatial_tokens,
+                height,
+                width,
+            )
+            if best is None or score < best[0]:
+                best = (score, height, width, spatial_tokens)
+
+    if best is None:
+        raise ValueError(f"Unable to resolve target_spatial_tokens={target_spatial_tokens}")
+    return best[1], best[2], best[3]
+
+
+def compute_dynamic_resolution(
+    image: Image.Image,
+    target_spatial_tokens: int = 880,
+    min_token_ratio: float = 0.95,
+    spatial_divisor_h: int = 32,
+    spatial_divisor_w: int = 32,
+):
+    """Preserve input aspect ratio while using 0.95-1.0 of the token budget."""
+    source_width, source_height = image.size
+    height, width, actual_tokens = _budget_spatial_size(
+        int(target_spatial_tokens),
+        source_height,
+        source_width,
+        spatial_divisor_h,
+        spatial_divisor_w,
+        min_token_ratio,
+    )
+    return height, width, int(target_spatial_tokens), actual_tokens
+
+
+# ============================================================================
+# Model setup
+# ============================================================================
+
+def setup_models(args, device, weight_dtype):
+    config = OmegaConf.load(args.config_path)
+
+    # Video transformer kwargs
+    video_transformer_kwargs = OmegaConf.to_container(
+        config.get("video_transformer_additional_kwargs",
+                   config.get("transformer_additional_kwargs", {})),
+        resolve=True,
+    )
+    # Audio transformer kwargs
+    audio_transformer_kwargs = OmegaConf.to_container(
+        config.get("audio_transformer_additional_kwargs",
+                   config.get("transformer_additional_kwargs", {})),
+        resolve=True,
+    )
+
+    # Resolve transformer paths
+    video_path = args.transformer_path
+    audio_path = args.transformer_path
+
+    # Check if joint checkpoint (has video_model/ and audio_model/ subdirs)
+    if args.transformer_path is not None:
+        video_sub = os.path.join(args.transformer_path, "video_model")
+        audio_sub = os.path.join(args.transformer_path, "audio_model")
+        if os.path.isdir(video_sub) and os.path.isdir(audio_sub):
+            video_path = video_sub
+            audio_path = audio_sub
+
+    # Creator gating kwargs
+    creator_gating_kwargs = OmegaConf.to_container(
+        config.get("creator_gating_kwargs", {}),
+        resolve=True,
+    )
+
+    print_info(f"Loading Creator gating AV transformer: video={video_path}, audio={audio_path}")
+    transformer = WanCreatorGatingAVModel.from_pretrained(
+        pretrained_model_path=args.transformer_path,
+        video_pretrained_model_path=video_path,
+        audio_pretrained_model_path=audio_path,
+        video_subfolder=video_transformer_kwargs.get("transformer_low_noise_model_subpath", None),
+        audio_subfolder=audio_transformer_kwargs.get("transformer_low_noise_model_subpath", None),
+        video_kwargs=video_transformer_kwargs,
+        audio_kwargs=audio_transformer_kwargs,
+        low_cpu_mem_usage=True,
+        torch_dtype=weight_dtype,
+        # Pass gating-specific kwargs from config
+        use_temporal_rope=creator_gating_kwargs.get("use_temporal_rope", args.use_temporal_rope),
+        audio_fps=creator_gating_kwargs.get("audio_fps", args.audio_fps),
+        vae_temporal_stride=creator_gating_kwargs.get("vae_temporal_stride", args.vae_temporal_stride),
+        a2v_cross_attn_layers=creator_gating_kwargs.get("a2v_cross_attn_layers", None),
+        v2a_cross_attn_layers=creator_gating_kwargs.get("v2a_cross_attn_layers", None),
+        use_gating=creator_gating_kwargs.get("use_gating", True),
+        zero_init_cross_attn=creator_gating_kwargs.get("zero_init_cross_attn", False),
+        zero_init_gating=creator_gating_kwargs.get("zero_init_gating", True),
+        gate_init_value=creator_gating_kwargs.get("gate_init_value", 0.0),
+        a2v_gate_alphas=creator_gating_kwargs.get("a2v_gate_alphas", None),
+        v2a_gate_alphas=creator_gating_kwargs.get("v2a_gate_alphas", None),
+    )
+
+    # Load video VAE
+    video_vae_path = os.path.join(
+        args.model_name, config.get("video_vae_kwargs", {}).get("vae_subpath", "vae"))
+    print_info(f"Loading video VAE from: {video_vae_path}")
+    vae_kwargs = OmegaConf.to_container(config.get("video_vae_kwargs", {}), resolve=True)
+    video_vae = AutoencoderKLWan3_8.from_pretrained(video_vae_path, additional_kwargs=vae_kwargs)
+
+    # Load audio VAE (CreatorDACVAE)
+    print_info(f"Loading audio VAE (CreatorDACVAE) from: {args.audio_vae_path}")
+    audio_vae = CreatorDACVAE.from_pretrained(args.audio_vae_path, strict=False)
+
+    # Load tokenizer and text encoder
+    text_encoder_kwargs = OmegaConf.to_container(config.get("text_encoder_kwargs", {}), resolve=True)
+    tokenizer_path = os.path.join(args.model_name, text_encoder_kwargs.get("tokenizer_subpath", "tokenizer"))
+    text_encoder_path = os.path.join(
+        args.model_name, text_encoder_kwargs.get("text_encoder_subpath", "text_encoder"))
+
+    print_info(f"Loading tokenizer from: {tokenizer_path}")
+    tokenizer = HFAutoTokenizer.from_pretrained(tokenizer_path)
+
+    print_info(f"Loading text encoder from: {text_encoder_path}")
+    text_encoder = WanT5EncoderModel.from_pretrained(
+        text_encoder_path,
+        additional_kwargs=text_encoder_kwargs,
+        low_cpu_mem_usage=True,
+        torch_dtype=weight_dtype,
+    )
+
+    # Setup schedulers
+    scheduler_dict = {"Flow": FlowMatchEulerDiscreteScheduler}
+    video_shift, audio_shift = resolve_scheduler_shifts(args)
+    scheduler_kwargs = OmegaConf.to_container(config.get("scheduler_kwargs", {}), resolve=True)
+    video_scheduler_kwargs = dict(scheduler_kwargs)
+    audio_scheduler_kwargs = dict(scheduler_kwargs)
+    video_scheduler_kwargs["shift"] = video_shift
+    audio_scheduler_kwargs["shift"] = audio_shift
+    Chosen_Scheduler = scheduler_dict[args.sampler_name]
+    video_scheduler = Chosen_Scheduler(**filter_kwargs(Chosen_Scheduler, video_scheduler_kwargs))
+    audio_scheduler = Chosen_Scheduler(**filter_kwargs(Chosen_Scheduler, audio_scheduler_kwargs))
+
+    # Set models to eval
+    transformer.eval()
+    if args.cfg_mode == "multimodal":
+        transformer = DirectionalMultimodalCFGAdapter(
+            transformer,
+            video_scale=args.video_bridge_guidance_scale,
+            audio_scale=args.audio_bridge_guidance_scale,
+            enable_a2v=not args.disable_a2v_cross_attn,
+            enable_v2a=not args.disable_v2a_cross_attn,
+        ).eval()
+    text_encoder.eval()
+    video_vae.eval()
+    audio_vae.eval()
+
+    offload_flags = resolve_cpu_offload_flags(args)
+
+    # Keep explicitly offloaded modules on CPU until their short GPU phase.
+    if not offload_flags["transformer"]:
+        transformer.to(device)
+    if not offload_flags["text_encoder"]:
+        text_encoder.to(device)
+    if not offload_flags["video_vae"]:
+        video_vae.to(device)
+    if not offload_flags["audio_vae"]:
+        audio_vae.to(device)
+
+    return {
+        "config": config,
+        "transformer": transformer,
+        "video_vae": video_vae,
+        "audio_vae": audio_vae,
+        "tokenizer": tokenizer,
+        "text_encoder": text_encoder,
+        "video_scheduler": video_scheduler,
+        "audio_scheduler": audio_scheduler,
+        "max_sequence_length": int(text_encoder_kwargs.get("text_length", 512)),
+    }
+
+
+# ============================================================================
+# Joint denoising loop
+# ============================================================================
+
+@torch.no_grad()
+def generate_joint_audio_video(args, models, device, weight_dtype, item):
+    """Generate audio and video jointly for a single item."""
+    transformer = models["transformer"]
+    video_vae = models["video_vae"]
+    audio_vae = models["audio_vae"]
+    tokenizer = models["tokenizer"]
+    text_encoder = models["text_encoder"]
+    video_scheduler = models["video_scheduler"]
+    audio_scheduler = models["audio_scheduler"]
+    max_sequence_length = models["max_sequence_length"]
+    offload_flags = resolve_cpu_offload_flags(args)
+
+    prompt = item["prompt"]
+    video_prompt = item.get("video_prompt", prompt)
+    audio_prompt = item.get("audio_prompt", prompt)
+    negative_prompt = item.get("negative_prompt", args.negative_prompt)
+    audio_negative_prompt = item.get("audio_negative_prompt", args.negative_prompt)
+    duration = float(item.get("duration", args.duration))
+    guidance_scale = float(item.get("guidance_scale", args.guidance_scale))
+    num_inference_steps = int(item.get("num_inference_steps", args.num_inference_steps))
+    seed = int(item.get("seed", args.seed))
+    # ---- Step 1: Load the supplied first frame ----
+    first_frame_path = item.get("first_frame_path", args.image)
+    if not os.path.isfile(first_frame_path):
+        raise FileNotFoundError(f"Input image not found: {first_frame_path}")
+    print_info(f"Loading first frame from: {first_frame_path}")
+    with Image.open(first_frame_path) as source_image:
+        source_image = source_image.convert("RGB")
+    video_patch_size = tuple(int(v) for v in transformer.video_patch_size)
+    spatial_compression = int(getattr(video_vae.config, "spatial_compression_ratio", 16))
+    height, width, target_tokens, actual_tokens = compute_dynamic_resolution(
+        source_image,
+        target_spatial_tokens=args.target_spatial_tokens,
+        min_token_ratio=args.min_token_ratio,
+        spatial_divisor_h=spatial_compression * video_patch_size[1],
+        spatial_divisor_w=spatial_compression * video_patch_size[2],
+    )
+    print_info(
+        f"Dynamic resolution: {height}x{width}, "
+        f"spatial_tokens={actual_tokens}/{target_tokens}"
+    )
+    first_frame_image = source_image.resize((width, height), Image.Resampling.BICUBIC)
+
+    do_classifier_free_guidance = guidance_scale > 1.0
+
+    # Save the resized frame once in distributed inference.
+    if not getattr(args, "suppress_aux_writes", False):
+        first_frame_save_path = os.path.splitext(args.output)[0] + "_first_frame.png"
+        first_frame_image.save(first_frame_save_path)
+        print_info(f"First frame saved to: {first_frame_save_path}")
+
+    # ---- Step 2: Encode first frame to video latent ----
+    if offload_flags["video_vae"]:
+        video_vae.to(device)
+
+    first_frame_latent = encode_first_frame(first_frame_image, video_vae, device, weight_dtype)
+    first_frame_latent = synchronize_tensor(first_frame_latent, args)
+    latent_channels = first_frame_latent.shape[1]
+    latent_height = first_frame_latent.shape[3]
+    latent_width = first_frame_latent.shape[4]
+
+    if offload_flags["video_vae"]:
+        video_vae.to("cpu")
+        clear_cuda_cache(device)
+
+    # ---- Step 3: Compute latent shapes ----
+    vae_temporal_ratio = int(getattr(video_vae.config, "temporal_compression_ratio", 4))
+    num_frames = int(duration * args.fps)
+    num_frames = int((num_frames - 1) // vae_temporal_ratio * vae_temporal_ratio) + 1
+    latent_frames = (num_frames - 1) // vae_temporal_ratio + 1
+
+    # Video sequence length
+    video_patch_size = tuple(int(v) for v in transformer.video_patch_size)
+    video_seq_len = (latent_frames * latent_height * latent_width) // math.prod(video_patch_size)
+
+    # Audio latent shape (DAC VAE: latent is [D, T])
+    audio_duration = num_frames / args.fps
+    audio_latent_T = compute_audio_latent_length(audio_duration, audio_vae)
+    audio_latent_dim = audio_vae.latent_dim
+    audio_patch_size = tuple(int(v) for v in transformer.audio_patch_size)
+    # Align latent_T to patch_size boundary
+    if audio_latent_T % audio_patch_size[0] != 0:
+        audio_latent_T = math.ceil(audio_latent_T / audio_patch_size[0]) * audio_patch_size[0]
+    audio_seq_len = get_audio_num_tokens(audio_latent_T, audio_patch_size)
+
+    print_info(f"Video: {num_frames} frames, latent [{latent_frames}, {latent_height}, {latent_width}], "
+               f"seq_len={video_seq_len}")
+    print_info(f"Audio: duration={audio_duration}s, latent [{audio_latent_dim}, {audio_latent_T}], "
+               f"seq_len={audio_seq_len}")
+
+    # ---- Step 4: Encode text prompts ----
+    if offload_flags["text_encoder"]:
+        text_encoder.to(device)
+
+    video_prompt_embeds = encode_prompt(
+        tokenizer, text_encoder, video_prompt, negative_prompt,
+        guidance_scale, max_sequence_length, device, weight_dtype)
+
+    if audio_prompt != video_prompt or audio_negative_prompt != negative_prompt:
+        audio_prompt_embeds = encode_prompt(
+            tokenizer, text_encoder, audio_prompt, audio_negative_prompt,
+            guidance_scale, max_sequence_length, device, weight_dtype)
+    else:
+        audio_prompt_embeds = video_prompt_embeds
+
+    if offload_flags["text_encoder"]:
+        text_encoder.to("cpu")
+        clear_cuda_cache(device)
+
+    # ---- Step 5: Initialize noise latents ----
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    # Video noise: [1, C, latent_frames, H_lat, W_lat]
+    video_latents = torch.randn(
+        (1, latent_channels, latent_frames, latent_height, latent_width),
+        generator=generator, device=device, dtype=weight_dtype,
+    )
+    # Apply first-frame conditioning
+    video_latents[:, :, 0:1, :, :] = first_frame_latent.to(weight_dtype)
+
+    # Audio noise: [1, D, T] (DAC latent format)
+    audio_latents = torch.randn(
+        (1, audio_latent_dim, audio_latent_T),
+        generator=generator, device=device, dtype=weight_dtype,
+    )
+    video_latents = synchronize_tensor(video_latents, args)
+    audio_latents = synchronize_tensor(audio_latents, args)
+
+    # ---- Step 6: Setup schedulers ----
+    video_shift, audio_shift = resolve_scheduler_shifts(args)
+    timestep_kwargs = {}
+    if getattr(args, "flow_match_mu", None) is not None:
+        timestep_kwargs["mu"] = float(args.flow_match_mu)
+    video_timesteps, num_inference_steps = retrieve_timesteps(
+        video_scheduler, num_inference_steps, device, **timestep_kwargs)
+    audio_timesteps, _ = retrieve_timesteps(
+        audio_scheduler, num_inference_steps, device, **timestep_kwargs)
+
+    timesteps = video_timesteps
+    if len(audio_timesteps) != len(video_timesteps):
+        raise ValueError(
+            f"Video/audio timestep count mismatch: {len(video_timesteps)} vs {len(audio_timesteps)}"
+        )
+    if hasattr(video_scheduler, "init_noise_sigma"):
+        video_latents[:, :, 1:, :, :] = video_latents[:, :, 1:, :, :] * video_scheduler.init_noise_sigma
+    if hasattr(audio_scheduler, "init_noise_sigma"):
+        audio_latents = audio_latents * audio_scheduler.init_noise_sigma
+
+    video_extra_step_kwargs = prepare_extra_step_kwargs(video_scheduler)
+    audio_extra_step_kwargs = prepare_extra_step_kwargs(audio_scheduler)
+
+    # ---- Step 7: Per-timestep video token count for first frame ----
+    patch_t, patch_h, patch_w = video_patch_size
+    first_frame_tokens = (latent_height * latent_width) // (patch_h * patch_w)
+
+    # ---- Step 8: Denoising loop ----
+    if offload_flags["transformer"]:
+        transformer.to(device)
+
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+    start_time = time.time()
+
+    enable_a2v = not getattr(args, "disable_a2v_cross_attn", False)
+    enable_v2a = not getattr(args, "disable_v2a_cross_attn", False)
+
+    for step_idx, t in enumerate(tqdm(
+        timesteps,
+        desc="Denoising",
+        unit="step",
+        disable=bool(getattr(args, "disable_progress", False)),
+    )):
+        audio_t = audio_timesteps[step_idx]
+        video_noisy_input = video_latents.clone()
+
+        # Per-token timestep for video (first frame = 0, rest = t)
+        t_scalar = t.item() if t.dim() == 0 else t
+        video_timestep_tokens = torch.ones(1, video_seq_len, device=device) * t_scalar
+        video_timestep_tokens[:, :first_frame_tokens] = 0
+
+        audio_timestep = audio_t.unsqueeze(0) if audio_t.dim() == 0 else audio_t
+        # video_timestep_tokens = (video_timestep_tokens / 1000) ** 0.25 * 1000
+        # Build model inputs
+        if do_classifier_free_guidance:
+            video_x_input = [video_noisy_input[0], video_noisy_input[0]]
+            video_t_input = torch.cat([video_timestep_tokens, video_timestep_tokens], dim=0)
+            video_context_input = video_prompt_embeds
+
+            audio_x_input = [audio_latents[0], audio_latents[0]]
+            audio_t_input = audio_timestep.expand(2)
+            audio_context_input = audio_prompt_embeds
+        else:
+            video_x_input = [video_noisy_input[0]]
+            video_t_input = video_timestep_tokens
+            video_context_input = video_prompt_embeds
+
+            audio_x_input = [audio_latents[0]]
+            audio_t_input = audio_timestep
+            audio_context_input = audio_prompt_embeds
+
+        video_inputs = {
+            "x": video_x_input,
+            "t": video_t_input,
+            "context": video_context_input,
+            "y": None,
+            "seq_len": video_seq_len,
+            "video_fps": float(args.fps),
+        }
+        audio_inputs = {
+            "x": audio_x_input,
+            "t": audio_t_input,
+            "context": audio_context_input,
+            "y": None,
+            "seq_len": audio_seq_len,
+        }
+
+        with torch.autocast("cuda", dtype=weight_dtype):
+            model_output = transformer(
+                video=video_inputs,
+                audio=audio_inputs,
+                dtype=weight_dtype,
+                enable_a2v=enable_a2v,
+                enable_v2a=enable_v2a,
+            )
+
+        video_pred = model_output["video"]
+        audio_pred = model_output["audio"]
+
+        # Apply CFG
+        if do_classifier_free_guidance:
+            video_pred_uncond = video_pred[0:1]
+            video_pred_text = video_pred[1:2]
+            video_noise_pred = video_pred_uncond + guidance_scale * (video_pred_text - video_pred_uncond)
+            audio_pred_uncond = audio_pred[0:1]
+            audio_pred_text = audio_pred[1:2]
+            audio_noise_pred = audio_pred_uncond + guidance_scale * (audio_pred_text - audio_pred_uncond)
+        else:
+            video_noise_pred = video_pred
+            audio_noise_pred = audio_pred
+
+        # Scheduler step
+        video_latents_denoised = video_scheduler.step(
+            video_noise_pred, t, video_latents, **video_extra_step_kwargs, return_dict=False)[0]
+        video_latents_denoised[:, :, 0:1, :, :] = first_frame_latent.to(weight_dtype)
+        video_latents = video_latents_denoised
+
+        audio_latents = audio_scheduler.step(
+            audio_noise_pred, audio_t, audio_latents, **audio_extra_step_kwargs, return_dict=False)[0]
+
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+    elapsed = time.time() - start_time
+    print_info(f"Denoising completed in {elapsed:.2f}s")
+
+    if offload_flags["transformer"]:
+        transformer.to("cpu")
+        clear_cuda_cache(device)
+
+    if getattr(args, "skip_output_decode", False):
+        return None, None, num_frames
+
+    # ---- Step 9: Decode video ----
+    if offload_flags["video_vae"]:
+        video_vae.to(device)
+
+    with torch.no_grad():
+        video_decoded = video_vae.decode(video_latents.to(video_vae.dtype))[0]
+    video_decoded = (video_decoded / 2.0 + 0.5).clamp(0, 1).float().cpu()
+
+    if offload_flags["video_vae"]:
+        video_vae.to("cpu")
+        clear_cuda_cache(device)
+
+    # ---- Step 10: Decode audio (CreatorDACVAE) ----
+    if offload_flags["audio_vae"]:
+        audio_vae.to(device)
+
+    with torch.no_grad():
+        # CreatorDACVAE.decode expects [B, D, T'] and returns [B, 1, T]
+        audio_decoded = audio_vae.decode(audio_latents.float().to(audio_vae.dac.device if hasattr(audio_vae, 'dac') else device))
+
+    if offload_flags["audio_vae"]:
+        audio_vae.to("cpu")
+        clear_cuda_cache(device)
+
+    return video_decoded, audio_decoded, num_frames
+
+
+def main():
+    args = parse_args()
+    device = init_device()
+    weight_dtype = resolve_weight_dtype(args.weight_dtype, device)
+
+    output_path = os.path.abspath(args.output)
+    if not os.path.splitext(output_path)[1]:
+        output_path += ".mp4"
+    args.output = output_path
+    output_dir = os.path.dirname(output_path) or "."
+    os.makedirs(output_dir, exist_ok=True)
+    args.output_dir = output_dir
+
+    print_info("=" * 80)
+    print_info("Creator single-image audio-video inference")
+    print_info("=" * 80)
+    print_info(f"Device: {device} | dtype: {weight_dtype}")
+    offload_flags = resolve_cpu_offload_flags(args)
+    print_info(
+        "CPU offload: "
+        f"text_encoder={'on' if offload_flags['text_encoder'] else 'off'}, "
+        f"video_vae={'on' if offload_flags['video_vae'] else 'off'}, "
+        f"audio_vae={'on' if offload_flags['audio_vae'] else 'off'}"
+    )
+    print_info(
+        "Cross attention: "
+        f"A2V={'disabled' if args.disable_a2v_cross_attn else 'enabled'}, "
+        f"V2A={'disabled' if args.disable_v2a_cross_attn else 'enabled'}"
+    )
+
+    # Load models
+    models = setup_models(args, device, weight_dtype)
+
+    item = {
+        "prompt": args.prompt,
+        "video_prompt": args.prompt,
+        "audio_prompt": args.prompt,
+        "negative_prompt": args.negative_prompt,
+        "audio_negative_prompt": args.negative_prompt,
+        "duration": args.duration,
+        "guidance_scale": args.guidance_scale,
+        "num_inference_steps": args.num_inference_steps,
+        "seed": args.seed,
+        "name": "sample",
+    }
+    video_decoded, audio_decoded, _ = generate_joint_audio_video(
+        args, models, device, weight_dtype, item)
+    video_path = os.path.splitext(output_path)[0] + ".video.mp4"
+    audio_path = os.path.splitext(output_path)[0] + ".wav"
+    frames = (video_decoded[0].permute(1, 2, 3, 0).clamp(0, 1).numpy() * 255).astype("uint8")
+    write_video(video_path, torch.from_numpy(frames), fps=args.fps, video_codec="h264")
+    save_audio_wav(audio_decoded, int(models["audio_vae"].sample_rate), audio_path)
+    mux_result = subprocess.run(
+        ["ffmpeg", "-y", "-i", video_path, "-i", audio_path,
+         "-c:v", "copy", "-c:a", "aac", "-shortest", output_path],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if mux_result.returncode != 0:
+        raise RuntimeError(
+            f"ffmpeg failed while muxing {output_path}:\n{mux_result.stderr.strip()}"
+        )
+    print_info(f"Saved: {output_path}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -1,18 +1,18 @@
-# For ZeroGPU, gradio/spaces/huggingface_hub/torch are preinstalled and
-# platform-managed; the Dockerfile installs them explicitly.
-# Pinned to 0.40: LTX2 video transformer + GGUF single-file loading both work
-# here (LTX2VideoTransformer3DModel + GGUFQuantizer). Earlier versions lack
-# AutoModel/from_single_file support for LTX-2 GGUF.
-diffusers==0.40.0
+# For ZeroGPU, gradio/spaces/huggingface_hub/torch are platform-managed; the
+# Dockerfile installs torch explicitly from the cu124 index and then this file.
+# DreamX-Creator (Wan2.2-TI2V-5B + gated A2V/V2A) engine pins mirror the upstream
+# Space (hugging-apps/gd-ml-dreamx-creator) adapted to our known-good base.
+diffusers==0.37.1
 transformers==5.16.1
 accelerate>=1.0
-gguf
 safetensors
+sentencepiece
+# .pth checkpoints are saved pre-torch-2.6; app.py wraps torch.load(weights_only=False).
 # torch 2.6.0+cu124 predates numpy 2.5 ABI; 2.1.x is the known-good pair.
 numpy==2.1.3
-scipy
+omegaconf
+einops
 pillow
-sentencepiece
-protobuf
-imageio-ffmpeg
-httpx
+soundfile
+tqdm
+protobuf

videox_fun/__init__.py

@@ -0,0 +1 @@
+"""Minimal runtime package for single-image Creator inference."""

videox_fun/dist/__init__.py

@@ -0,0 +1,3 @@
+from .sequence_parallel import all_gather_sequence, all_to_all, ulysses_attention
+
+__all__ = ["all_gather_sequence", "all_to_all", "ulysses_attention"]

videox_fun/dist/sequence_parallel.py

@@ -0,0 +1,42 @@
+"""Small raw-process-group helpers for Ulysses-style sequence parallelism."""
+
+from __future__ import annotations
+
+import torch
+import torch.distributed as dist
+
+
+def all_gather_sequence(tensor: torch.Tensor, group=None) -> torch.Tensor:
+    """Gather equally-sized sequence chunks along dimension 1."""
+    if not dist.is_initialized() or dist.get_world_size(group) == 1:
+        return tensor
+    tensor = tensor.contiguous()
+    chunks = [torch.empty_like(tensor) for _ in range(dist.get_world_size(group))]
+    dist.all_gather(chunks, tensor, group=group)
+    return torch.cat(chunks, dim=1).contiguous()
+
+
+def all_to_all(tensor: torch.Tensor, scatter_dim: int, gather_dim: int, group=None) -> torch.Tensor:
+    """Scatter one dimension and gather another, matching Wan2.2's Ulysses layout."""
+    if not dist.is_initialized() or dist.get_world_size(group) == 1:
+        return tensor
+
+    world_size = dist.get_world_size(group)
+    if tensor.shape[scatter_dim] % world_size != 0:
+        raise ValueError(
+            f"Dimension {scatter_dim} ({tensor.shape[scatter_dim]}) must be divisible by "
+            f"SP size {world_size}"
+        )
+    inputs = [part.contiguous() for part in tensor.chunk(world_size, dim=scatter_dim)]
+    outputs = [torch.empty_like(inputs[0]) for _ in range(world_size)]
+    dist.all_to_all(outputs, inputs, group=group)
+    return torch.cat(outputs, dim=gather_dim).contiguous()
+
+
+def ulysses_attention(q, k, v, attention_fn, *, k_lens=None, window_size=(-1, -1), group=None):
+    """Run attention on local sequence chunks with Ulysses head/sequence exchange."""
+    q = all_to_all(q, scatter_dim=2, gather_dim=1, group=group)
+    k = all_to_all(k, scatter_dim=2, gather_dim=1, group=group)
+    v = all_to_all(v, scatter_dim=2, gather_dim=1, group=group)
+    output = attention_fn(q, k, v, k_lens=k_lens, window_size=window_size)
+    return all_to_all(output, scatter_dim=1, gather_dim=2, group=group)

videox_fun/models/__init__.py

@@ -0,0 +1,7 @@
+"""Model definitions required by the release inference entrypoint."""
+
+from .creator_dac_vae import CreatorDACVAE
+from .wan_text_encoder import WanT5EncoderModel
+from .wan_vae3_8 import AutoencoderKLWan3_8
+
+__all__ = ["AutoencoderKLWan3_8", "CreatorDACVAE", "WanT5EncoderModel"]

videox_fun/models/attention_utils.py

@@ -0,0 +1,211 @@
+import os
+
+import torch
+import warnings
+
+try:
+    import flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+try:
+    major, minor = torch.cuda.get_device_capability(0)
+    if f"{major}.{minor}" == "8.0":
+        from sageattention_sm80 import sageattn
+        SAGE_ATTENTION_AVAILABLE = True
+    elif f"{major}.{minor}" == "8.6":
+        from sageattention_sm86 import sageattn
+        SAGE_ATTENTION_AVAILABLE = True
+    elif f"{major}.{minor}" == "8.9":
+        from sageattention_sm89 import sageattn
+        SAGE_ATTENTION_AVAILABLE = True
+    elif f"{major}.{minor}" == "9.0":
+        from sageattention_sm90 import sageattn
+        SAGE_ATTENTION_AVAILABLE = True
+    elif major>9:
+        from sageattention_sm120 import sageattn
+        SAGE_ATTENTION_AVAILABLE = True
+except:
+    try:
+        from sageattention import sageattn
+        SAGE_ATTENTION_AVAILABLE = True
+    except:
+        sageattn = None
+        SAGE_ATTENTION_AVAILABLE = False
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+    # params
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # preprocess query
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor(
+            [lq] * b, dtype=torch.int32).to(
+                device=q.device, non_blocking=True)
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+    # preprocess key, value
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor(
+            [lk] * b, dtype=torch.int32).to(
+                device=k.device, non_blocking=True)
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn(
+            'Flash attention 3 is not available, use flash attention 2 instead.'
+        )
+
+    # apply attention
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p, window_size are not supported in FA3 now.
+        x = flash_attn_interface.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            seqused_q=None,
+            seqused_k=None,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic)[0].unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic).unflatten(0, (b, lq))
+
+    # output
+    return x.type(out_dtype)
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+    attention_type=None,
+    attn_mask=None,
+):
+    attention_type = os.environ.get("VIDEOX_ATTENTION_TYPE", "FLASH_ATTENTION") if attention_type is None else attention_type
+    if torch.is_grad_enabled() and attention_type == "SAGE_ATTENTION":
+        attention_type = "FLASH_ATTENTION"
+
+    if attention_type == "SAGE_ATTENTION" and SAGE_ATTENTION_AVAILABLE:
+        if q_lens is not None or k_lens is not None:
+            warnings.warn(
+                'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+            )
+
+        out = sageattn(
+            q, k, v, attn_mask=attn_mask, tensor_layout="NHD", is_causal=causal, dropout_p=dropout_p)
+
+    elif attention_type == "FLASH_ATTENTION" and (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE):
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+    else:
+        if q_lens is not None or k_lens is not None:
+            warnings.warn(
+                'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+            )
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
+
+        out = out.transpose(1, 2).contiguous()
+    return out

videox_fun/models/creator/__init__.py

@@ -0,0 +1 @@
+"""Minimal Creator audio/VAE building blocks."""

videox_fun/models/creator/creator_audio_dit.py

@@ -0,0 +1,74 @@
+"""
+Code adapted from DiffSynth-Studio's Wan DiT implementation:
+https://github.com/modelscope/DiffSynth-Studio/blob/main/diffsynth/models/wan_video_dit.py
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+
+from .creator_video_dit import DiTBlock
+
+
+def sinusoidal_embedding_1d(dim, position):
+    sinusoid = torch.outer(position.type(torch.float64), torch.pow(
+        10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.to(position.dtype)
+
+
+def precompute_freqs_cis_1d(dim: int, end: int = 16384, theta: float = 10000.0):
+    f_freqs_cis = precompute_freqs_cis(dim, end, theta)
+    return f_freqs_cis.chunk(3, dim=-1)
+
+
+def precompute_freqs_cis(dim: int, end: int = 16384, theta: float = 10000.0, s: float = 1.0):
+    # 1d rope precompute
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)
+                   [: (dim // 2)].double() / dim))
+    pos = torch.arange(end, dtype=torch.float64, device=freqs.device) * s
+    freqs = torch.outer(pos, freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+class MLP(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, has_pos_emb=False):
+        super().__init__()
+        self.proj = torch.nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.Linear(in_dim, in_dim),
+            nn.GELU(),
+            nn.Linear(in_dim, out_dim),
+            nn.LayerNorm(out_dim)
+        )
+        self.has_pos_emb = has_pos_emb
+        if has_pos_emb:
+            self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280)))
+
+    def forward(self, x):
+        if self.has_pos_emb:
+            x = x + self.emb_pos.to(dtype=x.dtype, device=x.device)
+        return self.proj(x)
+
+
+class Head(nn.Module):
+    def __init__(self, dim: int, out_dim: int, patch_size, eps: float):
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.head = nn.Linear(dim, out_dim * math.prod(patch_size))
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, t_mod):
+        if len(t_mod.shape) == 3:
+            shift, scale = (self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2)).chunk(2, dim=2)
+            x = (self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2)))
+        else:
+            # NOTE: `t_mod` used to be [B, C]. When B=1 broadcasting works, but for B>1
+            # it does not align with [1, 2, C]. We therefore unsqueeze at dim=1 here.
+            shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1)).chunk(2, dim=1)
+            x = (self.head(self.norm(x) * (1 + scale) + shift))
+        return x

videox_fun/models/creator/creator_video_dit.py

@@ -0,0 +1,103 @@
+"""Shared DiT blocks used by the Creator audio transformer."""
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import RMSNorm
+
+from ..attention_utils import attention
+
+
+def sinusoidal_embedding_1d(dim, position):
+    sinusoid = torch.outer(
+        position.type(torch.float64),
+        torch.pow(10000, -torch.arange(dim // 2, dtype=torch.float64, device=position.device).div(dim // 2)),
+    )
+    return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1).to(position.dtype)
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    return (x * (1 + scale) + shift).to(shift.dtype)
+
+
+def rope_apply_head_dim(x, freqs, head_dim):
+    x = rearrange(x, "b s (n d) -> b s n d", d=head_dim)
+    x_complex = torch.view_as_complex(x.to(torch.float64).reshape(*x.shape[:3], -1, 2))
+    return torch.view_as_real(x_complex * freqs).flatten(2).to(x.dtype)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+    def forward(self, x, freqs, seq_lens=None):
+        q = rope_apply_head_dim(self.norm_q(self.q(x)), freqs, self.head_dim)
+        k = rope_apply_head_dim(self.norm_k(self.k(x)), freqs, self.head_dim)
+        v = self.v(x)
+        b, s = q.shape[:2]
+        q = q.view(b, s, self.num_heads, self.head_dim)
+        k = k.view(b, s, self.num_heads, self.head_dim)
+        v = v.view(b, s, self.num_heads, self.head_dim)
+        return self.o(attention(q, k, v, k_lens=seq_lens).flatten(2))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+    def forward(self, x: torch.Tensor, context: torch.Tensor):
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(context))
+        v = self.v(context)
+        b, s = q.shape[:2]
+        n = self.num_heads
+        d = q.shape[-1] // n
+        output = attention(
+            q.view(b, s, n, d),
+            k.view(b, k.shape[1], n, d),
+            v.view(b, v.shape[1], n, d),
+        )
+        return self.o(output.flatten(2))
+
+
+class GateModule(nn.Module):
+    def forward(self, x, gate, residual):
+        return x + gate * residual
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.self_attn = SelfAttention(dim, num_heads, eps)
+        self.cross_attn = CrossAttention(dim, num_heads, eps, has_image_input)
+        self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm3 = nn.LayerNorm(dim, eps=eps)
+        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim))
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.gate = GateModule()
+
+    def forward(self, x, context, t_mod, freqs, seq_lens=None):
+        chunk_dim = 2 if t_mod.ndim == 4 else 1
+        modulation = (self.modulation.to(t_mod) + t_mod).chunk(6, dim=chunk_dim)
+        if chunk_dim == 2:
+            modulation = tuple(part.squeeze(2) for part in modulation)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation
+        x = self.gate(x, gate_msa, self.self_attn(modulate(self.norm1(x), shift_msa, scale_msa), freqs, seq_lens))
+        x = x + self.cross_attn(self.norm3(x), context)
+        return self.gate(x, gate_mlp, self.ffn(modulate(self.norm2(x), shift_mlp, scale_mlp)))

videox_fun/models/creator/dac_vae.py

@@ -0,0 +1,878 @@
+"""
+Code adapted from HunyuanVideo-Foley's DAC VAE implementation:
+https://github.com/Tencent-Hunyuan/HunyuanVideo-Foley/tree/main/hunyuanvideo_foley/models/dac_vae
+
+We modified the original code to fit this project (e.g., integration and refactors).
+"""
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from torch import nn
+from torch.nn.utils import weight_norm
+
+SUPPORTED_VERSIONS = ["1.0.0"]
+
+
+@dataclass
+class DACFile:
+    codes: torch.Tensor
+
+    # Metadata
+    chunk_length: int
+    original_length: int
+    input_db: float
+    channels: int
+    sample_rate: int
+    padding: bool
+    dac_version: str
+
+    def save(self, path):
+        artifacts = {
+            "codes": self.codes.numpy().astype(np.uint16),
+            "metadata": {
+                "input_db": self.input_db.numpy().astype(np.float32),
+                "original_length": self.original_length,
+                "sample_rate": self.sample_rate,
+                "chunk_length": self.chunk_length,
+                "channels": self.channels,
+                "padding": self.padding,
+                "dac_version": SUPPORTED_VERSIONS[-1],
+            },
+        }
+        path = Path(path).with_suffix(".dac")
+        with open(path, "wb") as f:
+            np.save(f, artifacts)
+        return path
+
+    @classmethod
+    def load(cls, path):
+        artifacts = np.load(path, allow_pickle=True)[()]
+        codes = torch.from_numpy(artifacts["codes"].astype(int))
+        if artifacts["metadata"].get("dac_version", None) not in SUPPORTED_VERSIONS:
+            raise RuntimeError(
+                f"Given file {path} can't be loaded with this version of descript-audio-codec."
+            )
+        return cls(codes=codes, **artifacts["metadata"])
+
+
+class CodecMixin:
+    @property
+    def padding(self):
+        if not hasattr(self, "_padding"):
+            self._padding = True
+        return self._padding
+
+    @padding.setter
+    def padding(self, value):
+        assert isinstance(value, bool)
+
+        layers = [
+            l for l in self.modules() if isinstance(l, (nn.Conv1d, nn.ConvTranspose1d))
+        ]
+
+        for layer in layers:
+            if value:
+                if hasattr(layer, "original_padding"):
+                    layer.padding = layer.original_padding
+            else:
+                layer.original_padding = layer.padding
+                layer.padding = tuple(0 for _ in range(len(layer.padding)))
+
+        self._padding = value
+
+    def get_delay(self):
+        # Any number works here, delay is invariant to input length
+        l_out = self.get_output_length(0)
+        L = l_out
+
+        layers = []
+        for layer in self.modules():
+            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+                layers.append(layer)
+
+        for layer in reversed(layers):
+            d = layer.dilation[0]
+            k = layer.kernel_size[0]
+            s = layer.stride[0]
+
+            if isinstance(layer, nn.ConvTranspose1d):
+                L = ((L - d * (k - 1) - 1) / s) + 1
+            elif isinstance(layer, nn.Conv1d):
+                L = (L - 1) * s + d * (k - 1) + 1
+
+            L = math.ceil(L)
+
+        l_in = L
+
+        return (l_in - l_out) // 2
+
+    def get_output_length(self, input_length):
+        L = input_length
+        # Calculate output length
+        for layer in self.modules():
+            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+                d = layer.dilation[0]
+                k = layer.kernel_size[0]
+                s = layer.stride[0]
+
+                if isinstance(layer, nn.Conv1d):
+                    L = ((L - d * (k - 1) - 1) / s) + 1
+                elif isinstance(layer, nn.ConvTranspose1d):
+                    L = (L - 1) * s + d * (k - 1) + 1
+
+                L = math.floor(L)
+        return L
+
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+import torch.nn.functional as F
+from einops import rearrange
+
+
+class VectorQuantize(nn.Module):
+    """
+    Implementation of VQ similar to Karpathy's repo:
+    https://github.com/karpathy/deep-vector-quantization
+    Additionally uses following tricks from Improved VQGAN
+    (https://arxiv.org/pdf/2110.04627.pdf):
+        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
+            for improved codebook usage
+        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
+            improves training stability
+    """
+
+    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
+        self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def forward(self, z):
+        """Quantized the input tensor using a fixed codebook and returns
+        the corresponding codebook vectors
+
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized continuous representation of input
+        Tensor[1]
+            Commitment loss to train encoder to predict vectors closer to codebook
+            entries
+        Tensor[1]
+            Codebook loss to update the codebook
+        Tensor[B x T]
+            Codebook indices (quantized discrete representation of input)
+        Tensor[B x D x T]
+            Projected latents (continuous representation of input before quantization)
+        """
+
+        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
+        z_e = self.in_proj(z)  # z_e : (B x D x T)
+        z_q, indices = self.decode_latents(z_e)
+
+        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+        codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+
+        z_q = (
+            z_e + (z_q - z_e).detach()
+        )  # noop in forward pass, straight-through gradient estimator in backward pass
+
+        z_q = self.out_proj(z_q)
+
+        return z_q, commitment_loss, codebook_loss, indices, z_e
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight  # codebook: (N x D)
+
+        # L2 normalize encodings and codebook (ViT-VQGAN)
+        encodings = F.normalize(encodings)
+        codebook = F.normalize(codebook)
+
+        # Compute euclidean distance with codebook
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+        return z_q, indices
+
+
+class ResidualVectorQuantize(nn.Module):
+    """
+    Introduced in SoundStream: An end2end neural audio codec
+    https://arxiv.org/abs/2107.03312
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 512,
+        n_codebooks: int = 9,
+        codebook_size: int = 1024,
+        codebook_dim: Union[int, list] = 8,
+        quantizer_dropout: float = 0.0,
+    ):
+        super().__init__()
+        if isinstance(codebook_dim, int):
+            codebook_dim = [codebook_dim for _ in range(n_codebooks)]
+
+        self.n_codebooks = n_codebooks
+        self.codebook_dim = codebook_dim
+        self.codebook_size = codebook_size
+
+        self.quantizers = nn.ModuleList(
+            [
+                VectorQuantize(input_dim, codebook_size, codebook_dim[i])
+                for i in range(n_codebooks)
+            ]
+        )
+        self.quantizer_dropout = quantizer_dropout
+
+    def forward(self, z, n_quantizers: int = None):
+        """Quantized the input tensor using a fixed set of `n` codebooks and returns
+        the corresponding codebook vectors
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+        n_quantizers : int, optional
+            No. of quantizers to use
+            (n_quantizers < self.n_codebooks ex: for quantizer dropout)
+            Note: if `self.quantizer_dropout` is True, this argument is ignored
+                when in training mode, and a random number of quantizers is used.
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+        """
+        z_q = 0
+        residual = z
+        commitment_loss = 0
+        codebook_loss = 0
+
+        codebook_indices = []
+        latents = []
+
+        if n_quantizers is None:
+            n_quantizers = self.n_codebooks
+        if self.training:
+            n_quantizers = torch.ones((z.shape[0],)) * self.n_codebooks + 1
+            dropout = torch.randint(1, self.n_codebooks + 1, (z.shape[0],))
+            n_dropout = int(z.shape[0] * self.quantizer_dropout)
+            n_quantizers[:n_dropout] = dropout[:n_dropout]
+            n_quantizers = n_quantizers.to(z.device)
+
+        for i, quantizer in enumerate(self.quantizers):
+            if self.training is False and i >= n_quantizers:
+                break
+
+            z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+                residual
+            )
+
+            # Create mask to apply quantizer dropout
+            mask = (
+                torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
+            )
+            z_q = z_q + z_q_i * mask[:, None, None]
+            residual = residual - z_q_i
+
+            # Sum losses
+            commitment_loss += (commitment_loss_i * mask).mean()
+            codebook_loss += (codebook_loss_i * mask).mean()
+
+            codebook_indices.append(indices_i)
+            latents.append(z_e_i)
+
+        codes = torch.stack(codebook_indices, dim=1)
+        latents = torch.cat(latents, dim=1)
+
+        return z_q, codes, latents, commitment_loss, codebook_loss
+
+    def from_codes(self, codes: torch.Tensor):
+        """Given the quantized codes, reconstruct the continuous representation
+        Parameters
+        ----------
+        codes : Tensor[B x N x T]
+            Quantized discrete representation of input
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized continuous representation of input
+        """
+        z_q = 0.0
+        z_p = []
+        n_codebooks = codes.shape[1]
+        for i in range(n_codebooks):
+            z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
+            z_p.append(z_p_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+        return z_q, torch.cat(z_p, dim=1), codes
+
+    def from_latents(self, latents: torch.Tensor):
+        """Given the unquantized latents, reconstruct the
+        continuous representation after quantization.
+
+        Parameters
+        ----------
+        latents : Tensor[B x N x T]
+            Continuous representation of input after projection
+
+        Returns
+        -------
+        Tensor[B x D x T]
+            Quantized representation of full-projected space
+        Tensor[B x D x T]
+            Quantized representation of latent space
+        """
+        z_q = 0
+        z_p = []
+        codes = []
+        dims = np.cumsum([0] + [q.codebook_dim for q in self.quantizers])
+
+        n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[
+            0
+        ]
+        for i in range(n_codebooks):
+            j, k = dims[i], dims[i + 1]
+            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
+            z_p.append(z_p_i)
+            codes.append(codes_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+
+        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)
+
+
+class AbstractDistribution:
+    def sample(self):
+        raise NotImplementedError()
+
+    def mode(self):
+        raise NotImplementedError()
+
+
+class DiracDistribution(AbstractDistribution):
+    def __init__(self, value):
+        self.value = value
+
+    def sample(self):
+        return self.value
+
+    def mode(self):
+        return self.value
+
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        else:
+            if other is None:
+                return 0.5 * torch.mean(
+                    torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
+                    dim=[1, 2],
+                )
+            else:
+                return 0.5 * torch.mean(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var
+                    - 1.0
+                    - self.logvar
+                    + other.logvar,
+                    dim=[1, 2],
+                )
+
+    def nll(self, sample, dims=[1, 2]):
+        if self.deterministic:
+            return torch.Tensor([0.0])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims,
+        )
+
+    def mode(self):
+        return self.mean
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, torch.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for torch.exp().
+    logvar1, logvar2 = [x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor) for x in (logvar1, logvar2)]
+
+    return 0.5 * (
+        -1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
+    )
+
+
+def init_weights(m):
+    if isinstance(m, nn.Conv1d):
+        nn.init.trunc_normal_(m.weight, std=0.02)
+        nn.init.constant_(m.bias, 0)
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int = 16, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(dim // 2, dilation=1),
+            ResidualUnit(dim // 2, dilation=3),
+            ResidualUnit(dim // 2, dilation=9),
+            Snake1d(dim // 2),
+            WNConv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 64,
+        strides: list = [2, 4, 8, 8],
+        d_latent: int = 64,
+    ):
+        super().__init__()
+        # Create first convolution
+        self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for stride in strides:
+            d_model *= 2
+            self.block += [EncoderBlock(d_model, stride=stride)]
+
+        # Create last convolution
+        self.block += [
+            Snake1d(d_model),
+            WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
+        ]
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(input_dim),
+            WNConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+            ResidualUnit(output_dim, dilation=1),
+            ResidualUnit(output_dim, dilation=3),
+            ResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        input_channel,
+        channels,
+        rates,
+        d_out: int = 1,
+    ):
+        super().__init__()
+
+        # Add first conv layer
+        layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)]
+
+        # Add upsampling + MRF blocks
+        for i, stride in enumerate(rates):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            layers += [DecoderBlock(input_dim, output_dim, stride)]
+
+        # Add final conv layer
+        layers += [
+            Snake1d(output_dim),
+            WNConv1d(output_dim, d_out, kernel_size=7, padding=3),
+            nn.Tanh(),
+        ]
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class DAC(CodecMixin, ModelMixin, ConfigMixin):
+
+    @register_to_config
+    def __init__(
+        self,
+        encoder_dim: int = 128,
+        encoder_rates: List[int] = [2, 3, 4, 5, 8],
+        latent_dim: int = None, # 128
+        decoder_dim: int = 2048,
+        decoder_rates: List[int] = [8, 5, 4, 3, 2],
+        n_codebooks: int = 9,
+        codebook_size: int = 1024,
+        codebook_dim: Union[int, list] = 8,
+        quantizer_dropout: bool = False,
+        sample_rate: int = 48000,
+        continuous: bool = True,
+        use_weight_norm: bool = False,
+    ):
+        super().__init__()
+
+        self.encoder_dim = encoder_dim
+        self.encoder_rates = encoder_rates
+        self.decoder_dim = decoder_dim
+        self.decoder_rates = decoder_rates
+        self.sample_rate = sample_rate
+        self.continuous = continuous
+        self.use_weight_norm = use_weight_norm
+
+        if latent_dim is None:
+            latent_dim = encoder_dim * (2 ** len(encoder_rates))
+
+        self.latent_dim = latent_dim
+
+        self.hop_length = np.prod(encoder_rates)
+        self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
+
+        if not continuous:
+            self.n_codebooks = n_codebooks
+            self.codebook_size = codebook_size
+            self.codebook_dim = codebook_dim
+            self.quantizer = ResidualVectorQuantize(
+                input_dim=latent_dim,
+                n_codebooks=n_codebooks,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+                quantizer_dropout=quantizer_dropout,
+            )
+        else:
+            self.quant_conv = torch.nn.Conv1d(latent_dim, 2 * latent_dim, 1)
+            self.post_quant_conv = torch.nn.Conv1d(latent_dim, latent_dim, 1)
+
+        self.decoder = Decoder(
+            latent_dim,
+            decoder_dim,
+            decoder_rates,
+        )
+        self.sample_rate = sample_rate
+        self.apply(init_weights)
+
+        if not self.use_weight_norm:
+            self.remove_weight_norm()
+
+    @property
+    def dtype(self):
+        """Get the dtype of the model parameters."""
+        # Return the dtype of the first parameter found
+        for param in self.parameters():
+            return param.dtype
+        return torch.float32  # fallback
+
+    @property
+    def device(self):
+        """Get the device of the model parameters."""
+        # Return the device of the first parameter found
+        for param in self.parameters():
+            return param.device
+        return torch.device('cpu')  # fallback
+
+    def preprocess(self, audio_data, sample_rate):
+        if sample_rate is None:
+            sample_rate = self.sample_rate
+        assert sample_rate == self.sample_rate
+
+        length = audio_data.shape[-1]
+        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+        audio_data = nn.functional.pad(audio_data, (0, right_pad))
+
+        return audio_data
+
+    @apply_forward_hook
+    def encode(
+        self,
+        audio_data: torch.Tensor,
+        n_quantizers: int = None,
+    ):
+        """Encode given audio data and return quantized latent codes
+
+        Parameters
+        ----------
+        audio_data : Tensor[B x 1 x T]
+            Audio data to encode
+        n_quantizers : int, optional
+            Number of quantizers to use, by default None
+            If None, all quantizers are used.
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+            "length" : int
+                Number of samples in input audio
+        """
+        z = self.encoder(audio_data)  # [B x D x T]
+        if not self.continuous:
+            z, codes, latents, commitment_loss, codebook_loss = self.quantizer(z, n_quantizers)
+        else:
+            z = self.quant_conv(z)  # [B x 2D x T]
+            z = DiagonalGaussianDistribution(z)
+            codes, latents, commitment_loss, codebook_loss = None, None, 0, 0
+
+        return z, codes, latents, commitment_loss, codebook_loss
+
+    @apply_forward_hook
+    def decode(self, z: torch.Tensor):
+        """Decode given latent codes and return audio data
+
+        Parameters
+        ----------
+        z : Tensor[B x D x T]
+            Quantized continuous representation of input
+        length : int, optional
+            Number of samples in output audio, by default None
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "audio" : Tensor[B x 1 x length]
+                Decoded audio data.
+        """
+        if not self.continuous:
+            audio = self.decoder(z)
+        else:
+            z = self.post_quant_conv(z)
+            audio = self.decoder(z)
+
+        return audio
+
+    def forward(
+        self,
+        audio_data: torch.Tensor,
+        sample_rate: int = None,
+        n_quantizers: int = None,
+    ):
+        """Model forward pass
+
+        Parameters
+        ----------
+        audio_data : Tensor[B x 1 x T]
+            Audio data to encode
+        sample_rate : int, optional
+            Sample rate of audio data in Hz, by default None
+            If None, defaults to `self.sample_rate`
+        n_quantizers : int, optional
+            Number of quantizers to use, by default None.
+            If None, all quantizers are used.
+
+        Returns
+        -------
+        dict
+            A dictionary with the following keys:
+            "z" : Tensor[B x D x T]
+                Quantized continuous representation of input
+            "codes" : Tensor[B x N x T]
+                Codebook indices for each codebook
+                (quantized discrete representation of input)
+            "latents" : Tensor[B x N*D x T]
+                Projected latents (continuous representation of input before quantization)
+            "vq/commitment_loss" : Tensor[1]
+                Commitment loss to train encoder to predict vectors closer to codebook
+                entries
+            "vq/codebook_loss" : Tensor[1]
+                Codebook loss to update the codebook
+            "length" : int
+                Number of samples in input audio
+            "audio" : Tensor[B x 1 x length]
+                Decoded audio data.
+        """
+        length = audio_data.shape[-1]
+        audio_data = self.preprocess(audio_data, sample_rate)
+        if not self.continuous:
+            z, codes, latents, commitment_loss, codebook_loss = self.encode(audio_data, n_quantizers)
+
+            x = self.decode(z)
+            return {
+                "audio": x[..., :length],
+                "z": z,
+                "codes": codes,
+                "latents": latents,
+                "vq/commitment_loss": commitment_loss,
+                "vq/codebook_loss": codebook_loss,
+            }
+        else:
+            posterior, _, _, _, _ = self.encode(audio_data, n_quantizers)
+            z = posterior.sample()
+            x = self.decode(z)
+
+            kl_loss = posterior.kl()
+            kl_loss = kl_loss.mean()
+
+            return {
+                "audio": x[..., :length],
+                "z": z,
+                "kl_loss": kl_loss,
+            }
+
+    def remove_weight_norm(self):
+        """
+        Remove weight_norm from all modules in the model.
+        This fuses the weight_g and weight_v parameters into a single weight parameter.
+        Should be called before inference for better performance.
+        Returns:
+            self: The model with weight_norm removed
+        """
+        from torch.nn.utils import remove_weight_norm
+        for module in list(self.modules()):
+            if hasattr(module, "_forward_pre_hooks"):
+                for hook in list(module._forward_pre_hooks.values()):
+                    if "WeightNorm" in str(type(hook)):
+                        try:
+                            remove_weight_norm(module)
+                        except ValueError:
+                            continue
+        if self.use_weight_norm:
+            self.use_weight_norm = False
+            self.register_to_config(use_weight_norm=False)
+        return self

videox_fun/models/creator_audio.py

@@ -0,0 +1,353 @@
+"""Creator audio diffusion transformer used by the release inference path."""
+
+import glob
+import json
+import logging
+import os
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+from einops import rearrange
+
+from .creator.creator_video_dit import DiTBlock
+from .creator.creator_audio_dit import (
+    Head,
+    MLP,
+    sinusoidal_embedding_1d,
+    precompute_freqs_cis_1d,
+)
+
+
+class CreatorAudioModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+    """Creator Audio DiT with wan_audio2-compatible forward interface."""
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size=(1,),
+        text_len=512,
+        in_dim=128,
+        dim=2048,
+        ffn_dim=8192,
+        freq_dim=256,
+        text_dim=4096,
+        out_dim=128,
+        num_heads=16,
+        num_layers=32,
+        eps=1e-6,
+        has_image_input=False,
+        has_image_pos_emb=False,
+        has_ref_conv=False,
+        vae_type="dac",
+        **kwargs
+    ):
+        super().__init__()
+
+        self.patch_size = tuple(patch_size) if isinstance(patch_size, (list, tuple)) else (patch_size,)
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.eps = eps
+        self.has_image_input = has_image_input
+        self.vae_type = vae_type
+
+        # Patch embedding (1D conv)
+        self.patch_embedding = nn.Conv1d(
+            in_dim, dim, kernel_size=self.patch_size[0], stride=self.patch_size[0],
+        )
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(dim, dim),
+        )
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim),
+        )
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), nn.Linear(dim, dim * 6),
+        )
+
+        # DiTBlock stack
+        self.blocks = nn.ModuleList([
+            DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps)
+            for _ in range(num_layers)
+        ])
+
+        self.head = Head(dim, out_dim, self.patch_size, eps)
+
+        # RoPE precomputation
+        head_dim = dim // num_heads
+        self.freqs = precompute_freqs_cis_1d(head_dim)
+
+        # Optional image / reference embeddings
+        if has_image_input:
+            self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb)
+        if has_ref_conv:
+            self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2))
+        self.has_ref_conv = has_ref_conv
+
+    # -----------------------------------------------------------------
+    # Unpatchify
+    # -----------------------------------------------------------------
+
+    def unpatchify(self, x, grid_sizes, output_shapes):
+        """Unpatchify tokens back to audio latent shape.
+
+        Args:
+            x: [B, seq_len, out_dim * patch_size]
+            grid_sizes: [B, 1] token lengths
+            output_shapes: list of original audio shapes [(C, T), ...]
+
+        Returns:
+            list of restored tensors matching output_shapes.
+        """
+        output = []
+        for i, (grid_size, original_shape) in enumerate(zip(grid_sizes, output_shapes)):
+            f = int(grid_size[0].item())
+            restored = rearrange(
+                x[i, :f].unsqueeze(0),
+                "b f (p c) -> b c (f p)",
+                f=f, p=self.patch_size[0],
+            )
+            # Trim to original time length
+            orig_T = original_shape[-1]
+            restored = restored[:, :, :orig_T].squeeze(0)
+            output.append(restored)
+        return output
+
+    # -----------------------------------------------------------------
+    # RoPE frequencies
+    # -----------------------------------------------------------------
+    def _build_freqs(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        """Build RoPE complex frequencies [seq_len, 1, rope_dim]."""
+        freqs = torch.cat([
+            self.freqs[0][:seq_len].view(seq_len, -1),
+            self.freqs[1][:seq_len].view(seq_len, -1),
+            self.freqs[2][:seq_len].view(seq_len, -1),
+        ], dim=-1).reshape(seq_len, 1, -1).to(device)
+        return freqs
+
+    # -----------------------------------------------------------------
+    # Weight initialisation
+    # -----------------------------------------------------------------
+    def init_weights(self):
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        if self.patch_embedding.bias is not None:
+            nn.init.zeros_(self.patch_embedding.bias)
+        for module in self.text_embedding.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, std=0.02)
+        for module in self.time_embedding.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, std=0.02)
+
+        nn.init.zeros_(self.head.head.weight)
+        if self.head.head.bias is not None:
+            nn.init.zeros_(self.head.head.bias)
+
+    # -----------------------------------------------------------------
+    # forward (matches wan_audio2.WanAudioModel interface)
+    # -----------------------------------------------------------------
+    def forward(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        clip_fea=None,
+        y=None,
+        dtype=torch.bfloat16,
+        **kwargs,
+    ):
+        """
+        Args:
+            x: list of [C, T] audio latent samples, or [B, C, T] tensor.
+            t: [B] or [B, T] timesteps.
+            context: list of [S, text_dim] or [B, S, text_dim] tensor.
+            seq_len: max sequence length after patching.
+            clip_fea: optional CLIP features for image conditioning.
+            y: optional conditioning latent (e.g. for i2a), same format as x.
+        """
+        # --- Normalise inputs to lists ---
+        if isinstance(x, torch.Tensor):
+            x = [sample for sample in x]
+        if isinstance(context, torch.Tensor):
+            context = [sample for sample in context]
+        if y is not None and isinstance(y, torch.Tensor):
+            y = [sample for sample in y]
+
+        device = self.patch_embedding.weight.device
+        dtype = x[0].dtype if len(x) > 0 else self.patch_embedding.weight.dtype
+        batch_size = len(x)
+
+        # Concatenate conditioning y (e.g. for i2v/i2a)
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # Remember original shapes for unpatchify
+        original_audio_shapes = [tuple(sample.shape) for sample in x]
+
+        # --- Patchify per sample, then pad tokens ---
+        patchified_samples = []
+        grid_sizes_list = []
+        for sample in x:
+            tokens_i = self.patch_embedding(sample.unsqueeze(0).to(device))  # [1, dim, T']
+            tokens_i = rearrange(tokens_i, "1 c f -> f c").contiguous()
+            patchified_samples.append(tokens_i)
+            grid_sizes_list.append(tokens_i.shape[0])
+
+        grid_sizes = torch.tensor(
+            [[g] for g in grid_sizes_list], dtype=torch.long, device=device,
+        )
+        seq_lens = grid_sizes[:, 0]
+        assert seq_lens.max() <= seq_len, (
+            f"Max token length {seq_lens.max().item()} exceeds seq_len {seq_len}"
+        )
+
+        # Pad each sample's tokens to seq_len and stack
+        x = torch.stack([
+            torch.cat([u, u.new_zeros(seq_len - u.size(0), u.size(1))], dim=0)
+            for u in patchified_samples
+        ])
+
+        # --- Time embedding ---
+        with torch.amp.autocast("cuda", dtype=torch.float32):
+            if t.dim() != 1:
+                # Per-token timesteps [B, T]
+                if t.size(1) < seq_len:
+                    pad_size = seq_len - t.size(1)
+                    last_elements = t[:, -1].unsqueeze(1)
+                    t = torch.cat([t, last_elements.repeat(1, pad_size)], dim=1)
+                bt = t.size(0)
+                e = self.time_embedding(
+                    sinusoidal_embedding_1d(self.freq_dim, t.flatten())
+                    .unflatten(0, (bt, seq_len)).float()
+                )
+                e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+            else:
+                e = self.time_embedding(
+                    sinusoidal_embedding_1d(self.freq_dim, t).float()
+                )
+                e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+
+        # --- Context embedding ---
+        context = self.text_embedding(
+            torch.stack([
+                torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                for u in context
+            ])
+        )
+
+        # Image embedding
+        if self.has_image_input and clip_fea is not None:
+            clip_embedding = self.img_emb(clip_fea)
+            context = torch.cat([clip_embedding, context], dim=1)
+
+        # --- RoPE frequencies ---
+        freqs = self._build_freqs(seq_len, device)
+
+        for block in self.blocks:
+            x = block(x, context, e0, freqs, seq_lens=seq_lens)
+
+        x = self.head(x, e)
+        x = self.unpatchify(x, grid_sizes, original_audio_shapes)
+        return x
+
+    # -----------------------------------------------------------------
+    # from_pretrained
+    # -----------------------------------------------------------------
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_path,
+        subfolder=None,
+        transformer_additional_kwargs=None,
+        low_cpu_mem_usage=False,
+        in_dim=None,
+        out_dim=None,
+        patch_size=None,
+        torch_dtype=torch.bfloat16,
+    ):
+        transformer_additional_kwargs = dict(transformer_additional_kwargs or {})
+        if subfolder is not None:
+            pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
+
+        config_file = os.path.join(pretrained_model_path, "config.json")
+        if not os.path.isfile(config_file):
+            raise RuntimeError(f"{config_file} does not exist")
+
+        with open(config_file, "r") as fp:
+            config = json.load(fp)
+
+        from diffusers.utils import WEIGHTS_NAME
+
+        model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
+        model_file_safetensors = model_file.replace(".bin", ".safetensors")
+
+        # Explicit architecture overrides
+        if in_dim is not None:
+            transformer_additional_kwargs["in_dim"] = in_dim
+        if out_dim is not None:
+            transformer_additional_kwargs["out_dim"] = out_dim
+        if patch_size is not None:
+            transformer_additional_kwargs["patch_size"] = patch_size
+
+        if "dict_mapping" in transformer_additional_kwargs:
+            for key, value in transformer_additional_kwargs["dict_mapping"].items():
+                if value not in transformer_additional_kwargs:
+                    transformer_additional_kwargs[value] = config[key]
+
+        # Merge overrides into config
+        model_config = dict(config)
+        model_config.update(transformer_additional_kwargs)
+
+        model = cls.from_config(model_config, **transformer_additional_kwargs)
+
+        # --- Load state dict ---
+        if os.path.exists(model_file):
+            state_dict = torch.load(model_file, map_location="cpu")
+        elif os.path.exists(model_file_safetensors):
+            from safetensors.torch import load_file
+            state_dict = load_file(model_file_safetensors)
+        else:
+            from safetensors.torch import load_file
+            state_dict = {}
+            for shard in glob.glob(os.path.join(pretrained_model_path, "*.safetensors")):
+                state_dict.update(load_file(shard))
+
+        # Filter by shape match
+        model_sd = model.state_dict()
+        filtered_state_dict = {}
+        for key, value in state_dict.items():
+            if key in model_sd and model_sd[key].shape == value.shape:
+                filtered_state_dict[key] = value
+            else:
+                logging.info("Skipping key %s due to size mismatch or absence.", key)
+
+        missing, unexpected = model.load_state_dict(filtered_state_dict, strict=False)
+        logging.info(
+            "CreatorAudioModel missing keys: %d, unexpected keys: %d",
+            len(missing), len(unexpected),
+        )
+        return model.to(torch_dtype)
+
+# Convenience aliases
+CreatorAudioTransformerModel = CreatorAudioModel

videox_fun/models/creator_dac_vae.py

@@ -0,0 +1,151 @@
+import os
+from typing import Optional
+from pathlib import Path
+
+import torch
+
+from videox_fun.models.creator.dac_vae import DAC, DiagonalGaussianDistribution
+
+
+class CreatorDACVAE(torch.nn.Module):
+    """
+    High-level wrapper around the DAC (Descript Audio Codec) VAE in continuous mode.
+    Mirrors the LTXAudioVAE interface used by the inference pipeline.
+    """
+
+    def __init__(self, dac_model: DAC) -> None:
+        super().__init__()
+        self.dac = dac_model
+
+    @property
+    def sample_rate(self) -> int:
+        return self.dac.sample_rate
+
+    @property
+    def hop_length(self) -> int:
+        return self.dac.hop_length
+
+    @property
+    def latent_dim(self) -> int:
+        return self.dac.latent_dim
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_path: str | os.PathLike[str],
+        strict: bool = False,
+    ) -> "CreatorDACVAE":
+        pretrained_model_path = Path(pretrained_model_path)
+
+        if pretrained_model_path.is_dir():
+            dac_model = DAC.from_pretrained(pretrained_model_path)
+        else:
+            dac_model = DAC.from_pretrained(pretrained_model_path.parent)
+
+        return cls(dac_model=dac_model)
+
+    def _preprocess_waveform(self, waveform: torch.Tensor) -> torch.Tensor:
+        """Normalize waveform to [B, 1, T] mono and pad to hop_length boundary."""
+        """The Creator audio VAE currently supports mono audio."""
+        if waveform.ndim == 1:
+            waveform = waveform.unsqueeze(0).unsqueeze(0)
+        elif waveform.ndim == 2:
+            if waveform.size(0) > 1:
+                waveform = waveform.mean(dim=0, keepdim=True)
+            waveform = waveform.unsqueeze(0)
+        elif waveform.ndim == 3:
+            if waveform.size(1) > 1:
+                waveform = waveform.mean(dim=1, keepdim=True)
+
+        waveform = self.dac.preprocess(waveform, self.sample_rate)
+        return waveform
+
+    @torch.inference_mode()
+    def encode_posterior(
+        self,
+        audio: torch.Tensor,
+        sampling_rate: Optional[int] = None,
+        deterministic: bool | None = None,
+    ) -> DiagonalGaussianDistribution:
+        """Encode audio waveform and return the posterior distribution.
+
+        Parameters
+        ----------
+        audio : Tensor
+            Raw waveform tensor. Accepts shapes [T], [C, T], or [B, C, T].
+        sampling_rate : int, optional
+            Not used directly; kept for API compatibility with LTXAudioVAE.
+        deterministic : bool, optional
+            If True, std is zeroed so sampling returns the mean.
+
+        Returns
+        -------
+        DiagonalGaussianDistribution
+        """
+        waveform = self._preprocess_waveform(audio)
+        posterior, _, _, _, _ = self.dac.encode(waveform)
+        if deterministic is not None:
+            posterior.deterministic = deterministic
+            if deterministic:
+                posterior.std = torch.zeros_like(posterior.mean)
+                posterior.var = torch.zeros_like(posterior.mean)
+        return posterior
+
+    @torch.inference_mode()
+    def encode(
+        self,
+        audio: torch.Tensor,
+        sampling_rate: Optional[int] = None,
+        sample: bool = False,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        """Encode audio waveform to latent representation.
+
+        Parameters
+        ----------
+        audio : Tensor
+            Raw waveform tensor.
+        sampling_rate : int, optional
+            Kept for API compatibility.
+        sample : bool
+            If True, sample from the posterior; otherwise return the mean.
+        generator : torch.Generator, optional
+            RNG for reproducible sampling.
+
+        Returns
+        -------
+        Tensor [B, D, T']
+            Continuous latent codes.
+        """
+        posterior = self.encode_posterior(audio, sampling_rate=sampling_rate)
+        if sample:
+            return posterior.sample()
+        return posterior.mode()
+
+    @torch.inference_mode()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        """Decode latent codes back to waveform.
+
+        Parameters
+        ----------
+        latent : Tensor [B, D, T']
+            Continuous latent codes.
+
+        Returns
+        -------
+        Tensor [B, 1, T]
+            Reconstructed waveform.
+        """
+        return self.dac.decode(latent)
+
+    @torch.inference_mode()
+    def reconstruct(
+        self,
+        audio: torch.Tensor,
+        sampling_rate: Optional[int] = None,
+        sample: bool = False,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.Tensor:
+        """Encode then decode (round-trip reconstruction)."""
+        latent = self.encode(audio, sampling_rate=sampling_rate, sample=sample, generator=generator)
+        return self.decode(latent)

videox_fun/models/creator_gating.py

@@ -0,0 +1,1286 @@
+"""Configurable gated cross-attention joint audio-video model.
+
+This experimental variant follows the AV cross-attention design but allows
+audio-to-video (A2V) and video-to-audio (V2A) cross attention to be enabled
+independently for each layer. Cross-modal attention can optionally apply a
+fixed per-layer alpha times a per-head sigmoid gate to the attention context
+before the output projection.
+"""
+
+import logging
+import math
+import os
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from .attention_utils import attention
+from .creator_audio import CreatorAudioModel
+from .wan_transformer3d_prope import (
+    Wan2_2Transformer3DModel,
+    WanRMSNorm,
+    WanTransformer3DModel,
+    rope_apply_qk,
+)
+from .creator.creator_video_dit import sinusoidal_embedding_1d
+from .creator.creator_video_dit import rope_apply_head_dim
+from ..dist.sequence_parallel import all_gather_sequence, ulysses_attention
+
+
+LayerSelection = Optional[Union[bool, str, Sequence[bool], Sequence[int], torch.Tensor]]
+LayerAlphas = Optional[
+    Union[float, Sequence[float], Mapping[Union[int, str], float], torch.Tensor]
+]
+
+
+@torch.amp.autocast("cuda", enabled=False)
+def temporal_rope_1d(
+    x: torch.Tensor,
+    temporal_positions: torch.Tensor,
+    inv_freqs_1d: torch.Tensor,
+) -> torch.Tensor:
+    """Apply 1D temporal RoPE to [B, L, num_heads, head_dim] tensors."""
+    dtype = x.dtype
+    batch_size, seq_len, num_heads, head_dim = x.shape
+    half_dim = head_dim // 2
+
+    freqs = torch.einsum(
+        "bl,d->bld",
+        temporal_positions.to(torch.float64),
+        inv_freqs_1d.to(x.device, torch.float64),
+    )
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
+
+    x_complex = torch.view_as_complex(
+        x.to(torch.float64).reshape(batch_size, seq_len, num_heads, half_dim, 2)
+    )
+    x_out = torch.view_as_real(x_complex * freqs_cis.unsqueeze(2)).flatten(3)
+    return x_out.to(dtype)
+
+
+def compute_video_temporal_positions(
+    grid_sizes: torch.Tensor,
+    seq_len: int,
+    device: torch.device,
+    audio_fps: float = 48000.0 / 960.0,
+    video_fps: float = 16.0,
+    vae_temporal_stride: int = 4,
+) -> torch.Tensor:
+    """Compute video token temporal positions in audio-token time units."""
+    batch_size = grid_sizes.size(0)
+    positions = torch.zeros(batch_size, seq_len, device=device, dtype=torch.float64)
+    video_latent_fps = video_fps / vae_temporal_stride
+    scale = audio_fps / video_latent_fps
+
+    for sample_idx, (num_frames, height, width) in enumerate(grid_sizes.tolist()):
+        spatial_size = int(height * width)
+        num_tokens = int(num_frames * spatial_size)
+        frame_indices = torch.arange(num_tokens, device=device, dtype=torch.float64) // spatial_size
+        positions[sample_idx, :num_tokens] = frame_indices * scale
+    return positions
+
+
+def compute_audio_temporal_positions(
+    seq_lens: torch.Tensor,
+    seq_len: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Compute sequential temporal positions for audio tokens."""
+    batch_size = seq_lens.size(0)
+    positions = torch.zeros(batch_size, seq_len, device=device, dtype=torch.float64)
+    base_positions = torch.arange(seq_len, device=device, dtype=torch.float64)
+    for sample_idx in range(batch_size):
+        valid_len = int(seq_lens[sample_idx].item())
+        positions[sample_idx, :valid_len] = base_positions[:valid_len]
+    return positions
+
+
+def _apply_video_rope_local(
+    x: torch.Tensor,
+    grid_sizes: torch.Tensor,
+    freqs: torch.Tensor,
+    sp_rank: int,
+    sp_world_size: int,
+) -> torch.Tensor:
+    """Apply the Wan 3D RoPE slice belonging to this sequence-parallel rank."""
+    if sp_world_size <= 1:
+        return rope_apply_qk(x, x, grid_sizes, freqs)[0]
+
+    local_len, num_heads, complex_dim = x.size(1), x.size(2), x.size(3) // 2
+    freq_parts = freqs.split(
+        [complex_dim - 2 * (complex_dim // 3), complex_dim // 3, complex_dim // 3],
+        dim=1,
+    )
+    output = []
+    for sample_idx, (frames, height, width) in enumerate(grid_sizes.tolist()):
+        full_len = int(frames * height * width)
+        sample = x[sample_idx, :local_len].to(torch.float64)
+        sample_complex = torch.view_as_complex(
+            sample.reshape(local_len, num_heads, -1, 2)
+        )
+        full_freqs = torch.cat(
+            [
+                freq_parts[0][:frames].view(frames, 1, 1, -1).expand(frames, height, width, -1),
+                freq_parts[1][:height].view(1, height, 1, -1).expand(frames, height, width, -1),
+                freq_parts[2][:width].view(1, 1, width, -1).expand(frames, height, width, -1),
+            ],
+            dim=-1,
+        ).reshape(full_len, 1, -1)
+        if full_freqs.size(0) < local_len * sp_world_size:
+            full_freqs = torch.cat(
+                [
+                    full_freqs,
+                    torch.ones(
+                        local_len * sp_world_size - full_freqs.size(0),
+                        full_freqs.size(1),
+                        full_freqs.size(2),
+                        dtype=full_freqs.dtype,
+                        device=full_freqs.device,
+                    ),
+                ],
+                dim=0,
+            )
+        start = sp_rank * local_len
+        local_freqs = full_freqs[start : start + local_len]
+        rotated = torch.view_as_real(sample_complex * local_freqs).flatten(2)
+        if x.size(1) > local_len:
+            rotated = torch.cat([rotated, x[sample_idx, local_len:]], dim=0)
+        output.append(rotated)
+    return torch.stack(output).to(x.dtype)
+
+
+def _to_list(value) -> List[torch.Tensor]:
+    if isinstance(value, torch.Tensor):
+        return [sample for sample in value]
+    return list(value)
+
+
+def _build_time_embeddings(
+    time_embedding: nn.Module,
+    time_projection: nn.Module,
+    freq_dim: int,
+    dim: int,
+    timesteps: torch.Tensor,
+    seq_len: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Build time embeddings and projected modulations."""
+
+    with torch.amp.autocast("cuda", dtype=torch.float32):
+        if timesteps.dim() != 1:
+            if timesteps.size(1) < seq_len:
+                pad_size = seq_len - timesteps.size(1)
+                timesteps = torch.cat(
+                    [timesteps, timesteps[:, -1:].repeat(1, pad_size)], dim=1
+                )
+            batch_size = timesteps.size(0)
+            embedding = time_embedding(
+                sinusoidal_embedding_1d(freq_dim, timesteps.flatten())
+                .unflatten(0, (batch_size, seq_len))
+                .float()
+            )
+            modulation = time_projection(embedding).unflatten(2, (6, dim))
+        else:
+            embedding = time_embedding(sinusoidal_embedding_1d(freq_dim, timesteps).float())
+            modulation = time_projection(embedding).unflatten(1, (6, dim))
+    return embedding, modulation
+
+
+def _embed_context(text_embedding: nn.Module, context, text_len: int) -> torch.Tensor:
+    """Embed and right-pad text context to the model fixed text length."""
+    if isinstance(context, torch.Tensor):
+        samples = [sample for sample in context]
+    else:
+        samples = list(context)
+    return text_embedding(
+        torch.stack([
+            torch.cat([sample, sample.new_zeros(text_len - sample.size(0), sample.size(1))])
+            for sample in samples
+        ])
+    )
+
+
+def _logit_from_gate_value(gate_init_value: Optional[float]) -> float:
+    """Convert an initial gate value in [0, 1] to a sigmoid bias."""
+    if gate_init_value is None:
+        return 0.0
+    gate_value = min(max(float(gate_init_value), 1e-6), 1.0 - 1e-6)
+    return math.log(gate_value / (1.0 - gate_value))
+
+
+def _expand_layer_selection(
+    selection: LayerSelection,
+    num_layers: int,
+    name: str,
+) -> List[bool]:
+    """Expand a layer-selection config to a bool mask of length ``num_layers``.
+
+    Accepted forms:
+      - ``None`` or ``True``: enable every layer.
+      - ``False``: disable every layer.
+      - bool mask with length ``num_layers``.
+      - 0/1 mask with length ``num_layers``.
+      - list/tuple/tensor of layer indices to enable.
+      - strings: ``"all"``, ``"none"``, ``"0,2,5"``.
+    """
+    if selection is None:
+        return [False] * num_layers
+    if isinstance(selection, bool):
+        return [selection] * num_layers
+    if isinstance(selection, torch.Tensor):
+        selection = selection.cpu().tolist()
+    if isinstance(selection, str):
+        normalized = selection.strip().lower()
+        if normalized in {"", "none", "false", "off", "0"}:
+            return [False] * num_layers
+        if normalized in {"all", "true", "on", "1"}:
+            return [True] * num_layers
+        indices = [int(part.strip()) for part in selection.split(",") if part.strip()]
+        mask = [False] * num_layers
+        for layer_idx in indices:
+            if layer_idx < 0 or layer_idx >= num_layers:
+                raise ValueError(f"{name} layer index {layer_idx} out of range [0, {num_layers})")
+            mask[layer_idx] = True
+        return mask
+
+    values = list(selection)
+    if not values:
+        return [False] * num_layers
+
+    if all(isinstance(value, bool) for value in values):
+        if len(values) != num_layers:
+            raise ValueError(f"{name} bool mask must have length {num_layers}, got {len(values)}")
+        return [bool(value) for value in values]
+
+    if all(isinstance(value, int) for value in values):
+        if len(values) == num_layers and all(int(value) in {0, 1} for value in values):
+            return [bool(value) for value in values]
+        mask = [False] * num_layers
+        for layer_idx in values:
+            if layer_idx < 0 or layer_idx >= num_layers:
+                raise ValueError(f"{name} layer index {layer_idx} out of range [0, {num_layers})")
+            mask[int(layer_idx)] = True
+        return mask
+
+    raise TypeError(
+        f"{name} must be None, bool, string, bool mask, 0/1 mask, or layer-index sequence"
+    )
+
+
+def _expand_layer_alphas(
+    values: LayerAlphas,
+    num_layers: int,
+    name: str,
+) -> List[float]:
+    """Expand fixed per-layer gate multipliers, defaulting each layer to 1.0.
+
+    Accepted forms are a scalar shared by all layers, a full sequence with
+    ``num_layers`` entries, or a mapping of layer index to alpha. Unspecified
+    mapping entries retain the default value 1.0.
+    """
+    if values is None:
+        return [1.0] * num_layers
+    if isinstance(values, torch.Tensor):
+        values = values.cpu().tolist()
+
+    def validate(value, layer_label: str) -> float:
+        alpha = float(value)
+        if not math.isfinite(alpha) or alpha < 0.0:
+            raise ValueError(f"{name} {layer_label} must be finite and non-negative, got {value}")
+        return alpha
+
+    if isinstance(values, (int, float)) and not isinstance(values, bool):
+        return [validate(values, "scalar")] * num_layers
+
+    if isinstance(values, Mapping):
+        alphas = [1.0] * num_layers
+        for raw_layer_idx, value in values.items():
+            try:
+                layer_idx = int(raw_layer_idx)
+            except (TypeError, ValueError) as exc:
+                raise ValueError(
+                    f"{name} mapping key must be a layer index, got {raw_layer_idx!r}"
+                ) from exc
+            if layer_idx < 0 or layer_idx >= num_layers:
+                raise ValueError(f"{name} layer index {layer_idx} out of range [0, {num_layers})")
+            alphas[layer_idx] = validate(value, f"layer {layer_idx}")
+        return alphas
+
+    if isinstance(values, Sequence) and not isinstance(values, (str, bytes)):
+        if len(values) != num_layers:
+            raise ValueError(f"{name} sequence must have length {num_layers}, got {len(values)}")
+        return [validate(value, f"layer {layer_idx}") for layer_idx, value in enumerate(values)]
+
+    raise TypeError(f"{name} must be None, a scalar, a mapping, or a full layer sequence")
+
+
+class GatedCrossModalAttention(nn.Module):
+    """Cross-modal attention with optional alpha-scaled sigmoid context gating.
+
+    For A2V, ``x`` is video hidden states and ``y`` is audio hidden states.
+    For V2A, ``x`` is audio hidden states and ``y`` is video hidden states.
+    When ``use_gating=False``, this module follows the reference cross-attn
+    behavior: query from raw ``x`` and key/value from normalized ``y``.
+    """
+
+    def __init__(
+        self,
+        q_dim: int,
+        kv_dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        zero_init_output: bool = False,
+        use_gating: bool = True,
+        zero_init_gating: bool = False,
+        gate_init_value: Optional[float] = None,
+        gate_alpha: float = 1.0,
+    ):
+        super().__init__()
+        assert q_dim % num_heads == 0
+        self.q_dim = q_dim
+        self.kv_dim = kv_dim
+        self.num_heads = num_heads
+        self.head_dim = q_dim // num_heads
+        self.use_gating = bool(use_gating)
+        self.gate_alpha = _expand_layer_alphas(gate_alpha, 1, "gate_alpha")[0]
+
+        self.norm = nn.LayerNorm(kv_dim, eps=eps)
+        # if self.use_gating:
+        self.norm_x = nn.LayerNorm(q_dim, eps=eps)
+
+        self.q = nn.Linear(q_dim, q_dim)
+        self.k = nn.Linear(kv_dim, q_dim)
+        self.v = nn.Linear(kv_dim, q_dim)
+        self.o = nn.Linear(q_dim, q_dim)
+
+        self.norm_q = WanRMSNorm(q_dim, eps=eps)
+        self.norm_k = WanRMSNorm(q_dim, eps=eps)
+
+        if self.use_gating:
+            self.gate_hidden = nn.Linear(q_dim, num_heads, bias=False)
+            self.gate_context_norm = nn.LayerNorm(self.head_dim, eps=eps)
+            self.gate_context = nn.Linear(self.head_dim, 1, bias=False)
+            self.gate_bias = nn.Parameter(torch.empty(num_heads))
+
+        self._init_weights(
+            zero_init_output=zero_init_output,
+            zero_init_gating=zero_init_gating,
+            gate_init_value=gate_init_value,
+        )
+
+    def _init_weights(
+        self,
+        zero_init_output: bool,
+        zero_init_gating: bool,
+        gate_init_value: Optional[float],
+    ):
+        for module in [self.q, self.k, self.v, self.o]:
+            nn.init.xavier_uniform_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+        if zero_init_output:
+            nn.init.zeros_(self.o.weight)
+            nn.init.zeros_(self.o.bias)
+
+        if self.use_gating:
+            if zero_init_gating:
+                nn.init.zeros_(self.gate_hidden.weight)
+                nn.init.zeros_(self.gate_context.weight)
+            else:
+                nn.init.xavier_uniform_(self.gate_hidden.weight)
+                nn.init.xavier_uniform_(self.gate_context.weight)
+            nn.init.constant_(self.gate_bias, _logit_from_gate_value(gate_init_value))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        y_lens: Optional[torch.Tensor] = None,
+        dtype: torch.dtype = torch.bfloat16,
+        q_temporal_pos: Optional[torch.Tensor] = None,
+        k_temporal_pos: Optional[torch.Tensor] = None,
+        temporal_rope_inv_freq: Optional[torch.Tensor] = None,
+        sequence_parallel: bool = False,
+        sp_group=None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: primary hidden states [B, Lq, q_dim]
+            y: conditioning hidden states [B, Lk, kv_dim]
+            y_lens: valid lengths of y per sample [B]
+            q_temporal_pos: [B, Lq] temporal positions for query
+            k_temporal_pos: [B, Lk] temporal positions for key
+            temporal_rope_inv_freq: [head_dim // 2] inv frequencies for temporal RoPE
+        Returns:
+            Cross-attention output [B, Lq, q_dim].
+        """
+        batch_size = x.size(0)
+        num_heads = self.num_heads
+        head_dim = self.head_dim
+
+        # In SP mode x is a local query chunk while y is a local conditioning
+        # chunk. Cross-modal attention needs the complete conditioning sequence.
+        if sequence_parallel:
+            y = all_gather_sequence(y, group=sp_group)
+            if k_temporal_pos is not None:
+                k_temporal_pos = all_gather_sequence(
+                    k_temporal_pos.unsqueeze(-1), group=sp_group
+                ).squeeze(-1)
+
+        x_for_q = self.norm_x(x)
+        y_norm = self.norm(y)
+
+        query = self.norm_q(self.q(x_for_q.to(dtype))).view(batch_size, -1, num_heads, head_dim)
+        key = self.norm_k(self.k(y_norm.to(dtype))).view(batch_size, -1, num_heads, head_dim)
+        value = self.v(y_norm.to(dtype)).view(batch_size, -1, num_heads, head_dim)
+
+        if q_temporal_pos is not None and k_temporal_pos is not None and temporal_rope_inv_freq is not None:
+            query = temporal_rope_1d(query, q_temporal_pos, temporal_rope_inv_freq)
+            key = temporal_rope_1d(key, k_temporal_pos, temporal_rope_inv_freq)
+
+        context = attention(query.to(dtype), key.to(dtype), value.to(dtype), k_lens=y_lens)
+        context = context.to(dtype)
+
+        if self.use_gating:
+            hidden_gate = self.gate_hidden(x_for_q.to(dtype)).view(batch_size, -1, num_heads, 1)
+            context_gate = self.gate_context(
+                self.gate_context_norm(context)
+            )
+            gate = self.gate_alpha * torch.sigmoid(
+                hidden_gate + context_gate + self.gate_bias.view(1, 1, num_heads, 1)
+            )
+            context = gate * context
+        return self.o(context.flatten(2))
+
+
+class GatedJointBlock(nn.Module):
+    """One joint block with independently configurable A2V and V2A attention."""
+
+    def __init__(
+        self,
+        video_block: nn.Module,
+        audio_block: nn.Module,
+        video_dim: int,
+        audio_dim: int,
+        video_num_heads: int,
+        audio_num_heads: int,
+        enable_a2v_cross_attn: bool = True,
+        enable_v2a_cross_attn: bool = True,
+        zero_init_output: bool = False,
+        zero_init_video_cross_attn: bool | None = None,
+        zero_init_audio_cross_attn: bool | None = None,
+        use_a2v_gating: bool = True,
+        use_v2a_gating: bool = True,
+        zero_init_a2v_gating: bool = False,
+        zero_init_v2a_gating: bool = False,
+        a2v_gate_init_value: Optional[float] = None,
+        v2a_gate_init_value: Optional[float] = None,
+        a2v_gate_alpha: float = 1.0,
+        v2a_gate_alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.video_block = video_block
+        self.audio_block = audio_block
+        self.enable_a2v_cross_attn = bool(enable_a2v_cross_attn)
+        self.enable_v2a_cross_attn = bool(enable_v2a_cross_attn)
+
+        zero_init_video = zero_init_video_cross_attn if zero_init_video_cross_attn is not None else zero_init_output
+        zero_init_audio = zero_init_audio_cross_attn if zero_init_audio_cross_attn is not None else zero_init_output
+
+        if self.enable_a2v_cross_attn:
+            self.video_cross_attn_audio = GatedCrossModalAttention(
+                q_dim=video_dim,
+                kv_dim=audio_dim,
+                num_heads=video_num_heads,
+                zero_init_output=zero_init_video,
+                use_gating=use_a2v_gating,
+                zero_init_gating=zero_init_a2v_gating,
+                gate_init_value=a2v_gate_init_value,
+                gate_alpha=a2v_gate_alpha,
+            )
+        else:
+            self.video_cross_attn_audio = None
+
+        if self.enable_v2a_cross_attn:
+            self.audio_cross_attn_video = GatedCrossModalAttention(
+                q_dim=audio_dim,
+                kv_dim=video_dim,
+                num_heads=audio_num_heads,
+                zero_init_output=zero_init_audio,
+                use_gating=use_v2a_gating,
+                zero_init_gating=zero_init_v2a_gating,
+                gate_init_value=v2a_gate_init_value,
+                gate_alpha=v2a_gate_alpha,
+            )
+        else:
+            self.audio_cross_attn_video = None
+
+    def forward(
+        self,
+        video_x: torch.Tensor,
+        audio_x: torch.Tensor,
+        video_kwargs: Dict[str, Any],
+        audio_kwargs: Dict[str, Any],
+        dtype: torch.dtype = torch.bfloat16,
+        enable_a2v: Optional[Union[bool, torch.Tensor]] = None,
+        enable_v2a: Optional[Union[bool, torch.Tensor]] = None,
+        sequence_parallel: bool = False,
+        sp_group=None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        video_block = self.video_block
+        audio_block = self.audio_block
+
+        video_x = self._video_selfattn_and_text(
+            video_x, video_block, video_kwargs, dtype,
+            sequence_parallel=sequence_parallel,
+            sp_group=sp_group,
+        )
+        audio_x = self._audio_selfattn_and_text(
+            audio_x, audio_block, audio_kwargs, dtype,
+            sequence_parallel=sequence_parallel,
+            sp_group=sp_group,
+        )
+
+        temporal_rope_inv_freq = video_kwargs.get("temporal_rope_inv_freq")
+        video_temporal_pos = video_kwargs.get("temporal_positions")
+        audio_temporal_pos = audio_kwargs.get("temporal_positions")
+
+        # Runtime switches can be batch-wide booleans or per-branch CFG masks.
+        a2v_is_tensor = isinstance(enable_a2v, torch.Tensor)
+        v2a_is_tensor = isinstance(enable_v2a, torch.Tensor)
+        run_a2v = self.video_cross_attn_audio is not None and (
+            a2v_is_tensor or enable_a2v is None or enable_a2v
+        )
+        run_v2a = self.audio_cross_attn_video is not None and (
+            v2a_is_tensor or enable_v2a is None or enable_v2a
+        )
+
+        # Cache pre-cross-attention states so A2V and V2A are updated jointly:
+        # both directions must attend to the *same* pre-update snapshot, otherwise
+        # V2A would condition on the already-A2V-updated video (serial dependency).
+        video_x_pre = video_x
+        audio_x_pre = audio_x
+
+        if run_a2v:
+            a2v_result = self.video_cross_attn_audio(
+                x=video_x_pre,
+                y=audio_x_pre,
+                y_lens=audio_kwargs.get("seq_lens"),
+                dtype=dtype,
+                q_temporal_pos=video_temporal_pos,
+                k_temporal_pos=audio_temporal_pos,
+                temporal_rope_inv_freq=temporal_rope_inv_freq,
+                sequence_parallel=sequence_parallel,
+                sp_group=sp_group,
+            )
+            a2v_out = a2v_result
+            if a2v_is_tensor:
+                a2v_mask = enable_a2v.view(-1, *([1] * (a2v_out.dim() - 1))).to(
+                    device=a2v_out.device, dtype=a2v_out.dtype
+                )
+                a2v_out = a2v_out * a2v_mask
+            video_x = video_x + a2v_out
+
+        if run_v2a:
+            v2a_result = self.audio_cross_attn_video(
+                x=audio_x_pre,
+                y=video_x_pre,
+                y_lens=video_kwargs.get("seq_lens"),
+                dtype=dtype,
+                q_temporal_pos=audio_temporal_pos,
+                k_temporal_pos=video_temporal_pos,
+                temporal_rope_inv_freq=temporal_rope_inv_freq,
+                sequence_parallel=sequence_parallel,
+                sp_group=sp_group,
+            )
+            v2a_out = v2a_result
+            if v2a_is_tensor:
+                v2a_mask = enable_v2a.view(-1, *([1] * (v2a_out.dim() - 1))).to(
+                    device=v2a_out.device, dtype=v2a_out.dtype
+                )
+                v2a_out = v2a_out * v2a_mask
+            audio_x = audio_x + v2a_out
+
+        video_x = self._video_ffn(video_x, video_block, video_kwargs, dtype)
+        audio_x = self._audio_ffn(audio_x, audio_block, audio_kwargs, dtype)
+        return video_x, audio_x
+
+    def _video_selfattn_and_text(
+        self,
+        x: torch.Tensor,
+        block,
+        kwargs: Dict[str, Any],
+        dtype: torch.dtype,
+        sequence_parallel: bool = False,
+        sp_group=None,
+    ) -> torch.Tensor:
+        e0 = kwargs["e0"]
+        seq_lens = kwargs["seq_lens"]
+        grid_sizes = kwargs["grid_sizes"]
+        freqs = kwargs["freqs"]
+        context = kwargs["context"]
+        context_lens = kwargs.get("context_lens")
+
+        if e0.dim() > 3:
+            modulation = (block.modulation.unsqueeze(0) + e0).chunk(6, dim=2)
+            modulation = [part.squeeze(2) for part in modulation]
+        else:
+            modulation = (block.modulation + e0).chunk(6, dim=1)
+
+        kwargs["_video_e"] = modulation
+
+        temp_x = block.norm1(x) * (1 + modulation[1]) + modulation[0]
+        temp_x = temp_x.to(dtype)
+
+        self_attn = block.self_attn
+        batch_size, seq_len = temp_x.shape[:2]
+        num_heads, head_dim = self_attn.num_heads, self_attn.head_dim
+        query = self_attn.norm_q(self_attn.q(temp_x)).view(batch_size, seq_len, num_heads, head_dim)
+        key = self_attn.norm_k(self_attn.k(temp_x)).view(batch_size, seq_len, num_heads, head_dim)
+        value = self_attn.v(temp_x).view(batch_size, seq_len, num_heads, head_dim)
+        if sequence_parallel:
+            sp_rank = int(kwargs["sp_rank"])
+            sp_world_size = int(kwargs["sp_world_size"])
+            query = _apply_video_rope_local(query, grid_sizes, freqs, sp_rank, sp_world_size)
+            key = _apply_video_rope_local(key, grid_sizes, freqs, sp_rank, sp_world_size)
+            attn_output = ulysses_attention(
+                query.to(dtype),
+                key.to(dtype),
+                value.to(dtype),
+                attention,
+                k_lens=seq_lens,
+                window_size=getattr(self_attn, "window_size", (-1, -1)),
+                group=sp_group,
+            )
+        else:
+            query, key = rope_apply_qk(query, key, grid_sizes, freqs)
+            attn_output = attention(
+                query.to(dtype),
+                key.to(dtype),
+                v=value.to(dtype),
+                k_lens=seq_lens,
+                window_size=getattr(self_attn, "window_size", (-1, -1)),
+            )
+        attn_output = attn_output.to(dtype).flatten(2)
+        attn_output = self_attn.o(attn_output)
+
+        x = x + attn_output * modulation[2]
+        x = x + block.cross_attn(block.norm3(x), context, context_lens, dtype)
+        return x
+
+    def _audio_selfattn_and_text(
+        self,
+        x: torch.Tensor,
+        block,
+        kwargs: Dict[str, Any],
+        dtype: torch.dtype,
+        sequence_parallel: bool = False,
+        sp_group=None,
+    ) -> torch.Tensor:
+        time_mod = kwargs["e0"]
+        freqs = kwargs["freqs"]
+        context = kwargs["context"]
+
+        has_seq_mod = len(time_mod.shape) == 4
+        chunk_dim = 2 if has_seq_mod else 1
+        modulation = (
+            block.modulation.to(dtype=time_mod.dtype, device=time_mod.device) + time_mod
+        ).chunk(6, dim=chunk_dim)
+        if has_seq_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [
+                part.squeeze(2) for part in modulation
+            ]
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation
+
+        kwargs["_audio_mod"] = (shift_mlp, scale_mlp, gate_mlp)
+
+        # norm1 (LayerNorm) upcasts to fp32 and the modulation terms are fp32,
+        # so input_x is fp32 here. block.self_attn calls flash-attn, which only
+        # supports fp16/bf16 -> cast down first (mirrors the video self-attn path).
+        input_x = (block.norm1(x) * (1 + scale_msa) + shift_msa).to(dtype)
+        if sequence_parallel:
+            self_attn = block.self_attn
+            batch_size, seq_len = input_x.shape[:2]
+            num_heads, head_dim = self_attn.num_heads, self_attn.head_dim
+            query = self_attn.norm_q(self_attn.q(input_x)).view(batch_size, seq_len, num_heads, head_dim)
+            key = self_attn.norm_k(self_attn.k(input_x)).view(batch_size, seq_len, num_heads, head_dim)
+            value = self_attn.v(input_x).view(batch_size, seq_len, num_heads, head_dim)
+            query = rope_apply_head_dim(
+                query.flatten(2), freqs, head_dim
+            ).view(batch_size, seq_len, num_heads, head_dim)
+            key = rope_apply_head_dim(
+                key.flatten(2), freqs, head_dim
+            ).view(batch_size, seq_len, num_heads, head_dim)
+            self_attn_output = ulysses_attention(
+                query,
+                key,
+                value,
+                attention,
+                k_lens=kwargs.get("seq_lens"),
+                window_size=getattr(self_attn, "window_size", (-1, -1)),
+                group=sp_group,
+            ).flatten(2)
+            self_attn_output = self_attn.o(self_attn_output)
+        else:
+            self_attn_output = block.self_attn(input_x, freqs, seq_lens=kwargs.get("seq_lens"))
+        x = block.gate(x, gate_msa, self_attn_output)
+        # norm3 (LayerNorm) upcasts to fp32 and context (text embeds) may be fp32;
+        # audio cross_attn calls flash-attn (fp16/bf16 only) with no internal cast,
+        # so cast both q-source and kv-source down first.
+        x = x + block.cross_attn(block.norm3(x).to(dtype), context.to(dtype))
+        return x
+
+    def _video_ffn(
+        self, x: torch.Tensor, block, kwargs: Dict[str, Any], dtype: torch.dtype
+    ) -> torch.Tensor:
+        modulation = kwargs["_video_e"]
+        temp_x = block.norm2(x) * (1 + modulation[4]) + modulation[3]
+        temp_x = temp_x.to(dtype)
+        return x + block.ffn(temp_x) * modulation[5]
+
+    def _audio_ffn(
+        self, x: torch.Tensor, block, kwargs: Dict[str, Any], dtype: torch.dtype
+    ) -> torch.Tensor:
+        shift_mlp, scale_mlp, gate_mlp = kwargs["_audio_mod"]
+        input_x = block.norm2(x) * (1 + scale_mlp) + shift_mlp
+        return block.gate(x, gate_mlp, block.ffn(input_x))
+
+
+class WanCreatorGatingAVModel(nn.Module):
+    """Wan+Creator audio-video model with configurable gated cross attention.
+
+    A2V means the video branch attends to audio and updates video tokens.
+    V2A means the audio branch attends to video and updates audio tokens.
+    """
+
+    def __init__(
+        self,
+        video_model: WanTransformer3DModel,
+        audio_model: CreatorAudioModel,
+        use_temporal_rope: bool = True,
+        audio_fps: float = 48000 / 960,
+        vae_temporal_stride: int = 4,
+        zero_init_cross_attn: bool = False,
+        zero_init_video_cross_attn: bool | None = None,
+        zero_init_audio_cross_attn: bool | None = None,
+        a2v_cross_attn_layers: LayerSelection = None,
+        v2a_cross_attn_layers: LayerSelection = None,
+        use_gating: bool = True,
+        use_a2v_gating: bool | None = None,
+        use_v2a_gating: bool | None = None,
+        zero_init_gating: bool = False,
+        zero_init_a2v_gating: bool | None = None,
+        zero_init_v2a_gating: bool | None = None,
+        gate_init_value: Optional[float] = None,
+        a2v_gate_init_value: Optional[float] = None,
+        v2a_gate_init_value: Optional[float] = None,
+        a2v_gate_alphas: LayerAlphas = None,
+        v2a_gate_alphas: LayerAlphas = None,
+    ):
+        nn.Module.__init__(self)
+        self.video_model = video_model
+        self.audio_model = audio_model
+
+        self.video_dim = int(video_model.dim)
+        self.audio_dim = int(audio_model.dim)
+        self.video_num_heads = int(video_model.num_heads)
+        self.audio_num_heads = int(audio_model.num_heads)
+        self.num_layers = int(video_model.num_layers)
+        self.video_patch_size = tuple(int(value) for value in video_model.patch_size)
+        self.audio_patch_size = tuple(int(value) for value in audio_model.patch_size)
+
+        assert int(audio_model.num_layers) == self.num_layers, (
+            f"Video ({self.num_layers}) and audio ({audio_model.num_layers}) must have same number of layers"
+        )
+
+        a2v_enabled = _expand_layer_selection(
+            a2v_cross_attn_layers, self.num_layers, "a2v_cross_attn_layers"
+        )
+        v2a_enabled = _expand_layer_selection(
+            v2a_cross_attn_layers, self.num_layers, "v2a_cross_attn_layers"
+        )
+        self.a2v_cross_attn_layers = a2v_enabled
+        self.v2a_cross_attn_layers = v2a_enabled
+        resolved_a2v_gate_alphas = _expand_layer_alphas(
+            a2v_gate_alphas, self.num_layers, "a2v_gate_alphas"
+        )
+        resolved_v2a_gate_alphas = _expand_layer_alphas(
+            v2a_gate_alphas, self.num_layers, "v2a_gate_alphas"
+        )
+        self.a2v_gate_alphas = resolved_a2v_gate_alphas
+        self.v2a_gate_alphas = resolved_v2a_gate_alphas
+
+        resolved_use_a2v_gating = use_gating if use_a2v_gating is None else use_a2v_gating
+        resolved_use_v2a_gating = use_gating if use_v2a_gating is None else use_v2a_gating
+        resolved_zero_init_a2v_gating = (
+            zero_init_gating if zero_init_a2v_gating is None else zero_init_a2v_gating
+        )
+        resolved_zero_init_v2a_gating = (
+            zero_init_gating if zero_init_v2a_gating is None else zero_init_v2a_gating
+        )
+        resolved_a2v_gate_init_value = (
+            gate_init_value if a2v_gate_init_value is None else a2v_gate_init_value
+        )
+        resolved_v2a_gate_init_value = (
+            gate_init_value if v2a_gate_init_value is None else v2a_gate_init_value
+        )
+        video_blocks = list(video_model.blocks)
+        audio_blocks = list(audio_model.blocks)
+        self.joint_blocks = nn.ModuleList([
+            GatedJointBlock(
+                video_block=video_block,
+                audio_block=audio_block,
+                video_dim=self.video_dim,
+                audio_dim=self.audio_dim,
+                video_num_heads=self.video_num_heads,
+                audio_num_heads=self.audio_num_heads,
+                enable_a2v_cross_attn=a2v_enabled[layer_idx],
+                enable_v2a_cross_attn=v2a_enabled[layer_idx],
+                zero_init_output=zero_init_cross_attn,
+                zero_init_video_cross_attn=zero_init_video_cross_attn,
+                zero_init_audio_cross_attn=zero_init_audio_cross_attn,
+                use_a2v_gating=resolved_use_a2v_gating,
+                use_v2a_gating=resolved_use_v2a_gating,
+                zero_init_a2v_gating=resolved_zero_init_a2v_gating,
+                zero_init_v2a_gating=resolved_zero_init_v2a_gating,
+                a2v_gate_init_value=resolved_a2v_gate_init_value,
+                v2a_gate_init_value=resolved_v2a_gate_init_value,
+                a2v_gate_alpha=resolved_a2v_gate_alphas[layer_idx],
+                v2a_gate_alpha=resolved_v2a_gate_alphas[layer_idx],
+            )
+            for layer_idx, (video_block, audio_block) in enumerate(zip(video_blocks, audio_blocks))
+        ])
+
+        video_model.blocks = nn.ModuleList()
+        audio_model.blocks = nn.ModuleList()
+
+        self.use_temporal_rope = use_temporal_rope
+        self.audio_fps = audio_fps
+        self.vae_temporal_stride = vae_temporal_stride
+        if use_temporal_rope:
+            head_dim = self.video_dim // self.video_num_heads
+            self.temporal_rope_inv_freq = 1.0 / (
+                10000.0 ** (torch.arange(0, head_dim, 2, dtype=torch.float64) / head_dim)
+            )
+        else:
+            self.temporal_rope_inv_freq = None
+
+        # Ulysses-style sequence-parallel inference state. All ranks keep a
+        # complete copy of the weights and exchange token/head dimensions in
+        # attention, matching the Wan2.2 inference design.
+        self.sp_world_size = 1
+        self.sp_world_rank = 0
+        self.sp_group = None
+
+    def enable_multi_gpus_inference(self, group=None) -> None:
+        """Enable raw-process-group sequence parallelism for inference."""
+        import torch.distributed as dist
+
+        if not dist.is_initialized():
+            raise RuntimeError("Sequence-parallel inference requires an initialized process group")
+        self.sp_world_size = dist.get_world_size(group)
+        self.sp_world_rank = dist.get_rank(group)
+        self.sp_group = group
+        if self.video_num_heads % self.sp_world_size != 0:
+            raise ValueError(
+                f"Video attention heads ({self.video_num_heads}) must be divisible by "
+                f"SP size ({self.sp_world_size})"
+            )
+        if self.audio_num_heads % self.sp_world_size != 0:
+            raise ValueError(
+                f"Audio attention heads ({self.audio_num_heads}) must be divisible by "
+                f"SP size ({self.sp_world_size})"
+            )
+
+    def _apply_sequence_parallel(self, video_state: Dict[str, Any], audio_state: Dict[str, Any]):
+        """Shard prepared video/audio token states along their sequence axes."""
+        if self.sp_world_size <= 1:
+            video_state["sp_rank"] = 0
+            video_state["sp_world_size"] = 1
+            audio_state["sp_rank"] = 0
+            audio_state["sp_world_size"] = 1
+            return video_state, audio_state
+
+        def shard(state: Dict[str, Any], *, audio: bool):
+            local_len = state["x"].size(1) // self.sp_world_size
+            rank = self.sp_world_rank
+            state["x"] = torch.chunk(state["x"], self.sp_world_size, dim=1)[rank]
+            if not audio:
+                if state["e"].dim() >= 3:
+                    state["e"] = torch.chunk(state["e"], self.sp_world_size, dim=1)[rank]
+                if state["e0"].dim() >= 4:
+                    state["e0"] = torch.chunk(state["e0"], self.sp_world_size, dim=1)[rank]
+            state["freqs"] = (
+                torch.chunk(state["freqs"], self.sp_world_size, dim=0)[rank]
+                if audio else state["freqs"]
+            )
+            if state.get("temporal_positions") is not None:
+                state["temporal_positions"] = torch.chunk(
+                    state["temporal_positions"], self.sp_world_size, dim=1
+                )[rank]
+            state["local_seq_lens"] = (
+                state["seq_lens"] - rank * local_len
+            ).clamp(min=0, max=local_len)
+            state["sp_rank"] = rank
+            state["sp_world_size"] = self.sp_world_size
+            return state
+
+        return shard(video_state, audio=False), shard(audio_state, audio=True)
+
+    def _prepare_video(self, video_inputs: Dict[str, Any], dtype: torch.dtype) -> Dict[str, Any]:
+        video_model = self.video_model
+        device = video_model.patch_embedding.weight.device
+
+        if video_model.freqs.device != device:
+            video_model.freqs = video_model.freqs.to(device)
+
+        x_list = _to_list(video_inputs["x"])
+        y = video_inputs.get("y")
+        if y is not None:
+            y_list = _to_list(y)
+            x_list = [torch.cat([sample, condition], dim=0) for sample, condition in zip(x_list, y_list)]
+
+        x_list = [video_model.patch_embedding(sample.unsqueeze(0)) for sample in x_list]
+        grid_sizes = torch.stack([
+            torch.tensor(sample.shape[2:], dtype=torch.long, device=device) for sample in x_list
+        ])
+        x_list = [sample.flatten(2).transpose(1, 2) for sample in x_list]
+        seq_lens = torch.tensor([sample.size(1) for sample in x_list], dtype=torch.long, device=device)
+
+        seq_len = self._round_seq_len(int(video_inputs["seq_len"]))
+        assert int(seq_lens.max().item()) <= seq_len
+        x = torch.cat([
+            torch.cat([sample, sample.new_zeros(1, seq_len - sample.size(1), sample.size(2))], dim=1)
+            for sample in x_list
+        ])
+
+        timesteps = video_inputs["t"].to(device)
+        embedding, modulation = _build_time_embeddings(
+            video_model.time_embedding,
+            video_model.time_projection,
+            int(video_model.freq_dim),
+            int(video_model.dim),
+            timesteps,
+            seq_len,
+        )
+
+        context = _embed_context(video_model.text_embedding, video_inputs["context"], int(video_model.text_len))
+
+        return {
+            "x": x,
+            "e": embedding,
+            "e0": modulation,
+            "seq_lens": seq_lens,
+            "grid_sizes": grid_sizes,
+            "freqs": video_model.freqs,
+            "context": context,
+            "context_lens": None,
+            "seq_len": seq_len,
+        }
+
+    def _prepare_audio(self, audio_inputs: Dict[str, Any], dtype: torch.dtype) -> Dict[str, Any]:
+        audio_model = self.audio_model
+        device = audio_model.patch_embedding.weight.device
+
+        x_list = _to_list(audio_inputs["x"])
+        y = audio_inputs.get("y")
+        if y is not None:
+            y_list = _to_list(y)
+            x_list = [torch.cat([sample, condition], dim=0) for sample, condition in zip(x_list, y_list)]
+
+        original_audio_shapes = [tuple(sample.shape) for sample in x_list]
+
+        patchified = []
+        grid_sizes_list = []
+        for sample in x_list:
+            tokens = audio_model.patch_embedding(sample.unsqueeze(0).to(device))
+            tokens = rearrange(tokens, "1 c f -> f c").contiguous()
+            patchified.append(tokens)
+            grid_sizes_list.append(tokens.shape[0])
+
+        grid_sizes = torch.tensor([[grid_size] for grid_size in grid_sizes_list], dtype=torch.long, device=device)
+        seq_lens = grid_sizes[:, 0]
+        seq_len = self._round_seq_len(int(audio_inputs["seq_len"]))
+        assert int(seq_lens.max().item()) <= seq_len
+
+        x = torch.stack([
+            torch.cat([sample, sample.new_zeros(seq_len - sample.size(0), sample.size(1))], dim=0)
+            for sample in patchified
+        ])
+
+        timesteps = audio_inputs["t"].to(device)
+        embedding, modulation = _build_time_embeddings(
+            audio_model.time_embedding,
+            audio_model.time_projection,
+            int(audio_model.freq_dim),
+            int(audio_model.dim),
+            timesteps,
+            seq_len,
+        )
+
+        context = _embed_context(audio_model.text_embedding, audio_inputs["context"], int(audio_model.text_len))
+        freqs = audio_model._build_freqs(seq_len, device)
+
+        clip_fea = audio_inputs.get("clip_fea")
+        if audio_model.has_image_input and clip_fea is not None:
+            clip_embedding = audio_model.img_emb(clip_fea)
+            context = torch.cat([clip_embedding, context], dim=1)
+
+        return {
+            "x": x,
+            "e": embedding,
+            "e0": modulation,
+            "seq_lens": seq_lens,
+            "grid_sizes": grid_sizes,
+            "freqs": freqs,
+            "context": context,
+            "context_lens": None,
+            "seq_len": seq_len,
+            "original_audio_shapes": original_audio_shapes,
+        }
+
+    def forward(
+        self,
+        video: Dict[str, Any],
+        audio: Dict[str, Any],
+        dtype: torch.dtype = torch.bfloat16,
+        return_dict: bool = True,
+        enable_a2v: Optional[Union[bool, torch.Tensor]] = None,
+        enable_v2a: Optional[Union[bool, torch.Tensor]] = None,
+    ):
+        """Forward pass of the joint audio-video model.
+
+        Args:
+            video: video input dict with keys 'x', 't', 'context', 'seq_len', etc.
+            audio: audio input dict with keys 'x', 't', 'context', 'seq_len', etc.
+            dtype: computation dtype for attention ops (default bfloat16).
+            return_dict: if True, return dict with 'video'/'audio' keys; else tuple.
+            enable_a2v: gate A2V cross-attention (video attending to audio).
+                A bool tensor can select the enabled CFG branches per sample.
+            enable_v2a: gate V2A cross-attention (audio attending to video), same
+                semantics as enable_a2v.
+
+        Returns:
+            dict or tuple of (video_output, audio_output) tensors.
+        """
+        video_state = self._prepare_video(video, dtype)
+        audio_state = self._prepare_audio(audio, dtype)
+
+        device = video_state["x"].device
+
+        video_temporal_pos = None
+        audio_temporal_pos = None
+        temporal_rope_inv_freq = None
+        if self.use_temporal_rope and self.temporal_rope_inv_freq is not None:
+            video_fps = float(video.get("video_fps", 16.0))
+            temporal_rope_inv_freq = self.temporal_rope_inv_freq
+            video_temporal_pos = compute_video_temporal_positions(
+                video_state["grid_sizes"],
+                video_state["x"].size(1),
+                device,
+                audio_fps=self.audio_fps,
+                video_fps=video_fps,
+                vae_temporal_stride=self.vae_temporal_stride,
+            )
+            audio_temporal_pos = compute_audio_temporal_positions(
+                audio_state["seq_lens"],
+                audio_state["x"].size(1),
+                device,
+            )
+
+        video_state["temporal_positions"] = video_temporal_pos
+        audio_state["temporal_positions"] = audio_temporal_pos
+        video_state, audio_state = self._apply_sequence_parallel(video_state, audio_state)
+        video_x = video_state["x"]
+        audio_x = audio_state["x"]
+
+        video_kwargs = {
+            "e0": video_state["e0"],
+            "seq_lens": video_state["seq_lens"],
+            "grid_sizes": video_state["grid_sizes"],
+            "freqs": video_state["freqs"],
+            "context": video_state["context"],
+            "context_lens": video_state["context_lens"],
+            "temporal_positions": video_state["temporal_positions"],
+            "temporal_rope_inv_freq": temporal_rope_inv_freq,
+            "sp_rank": video_state["sp_rank"],
+            "sp_world_size": video_state["sp_world_size"],
+        }
+        audio_kwargs = {
+            "e0": audio_state["e0"],
+            "seq_lens": audio_state["seq_lens"],
+            "freqs": audio_state["freqs"],
+            "context": audio_state["context"],
+            "temporal_positions": audio_state["temporal_positions"],
+            "sp_rank": audio_state["sp_rank"],
+            "sp_world_size": audio_state["sp_world_size"],
+        }
+
+        runtime_cross_attn = {
+            "enable_a2v": enable_a2v,
+            "enable_v2a": enable_v2a,
+            "sequence_parallel": self.sp_world_size > 1,
+            "sp_group": self.sp_group,
+        }
+
+        for joint_block in self.joint_blocks:
+            video_x, audio_x = joint_block(
+                video_x, audio_x, video_kwargs, audio_kwargs, dtype,
+                **runtime_cross_attn,
+            )
+
+        video_output = self.video_model.head(video_x, video_state["e"])
+        audio_output = self.audio_model.head(audio_x, audio_state["e"])
+
+        if self.sp_world_size > 1:
+            video_output = all_gather_sequence(video_output, group=self.sp_group)
+            audio_output = all_gather_sequence(audio_output, group=self.sp_group)
+
+        video_output = torch.stack(
+            self.video_model.unpatchify(video_output, video_state["grid_sizes"])
+        )
+        audio_output = torch.stack(
+            self.audio_model.unpatchify(
+                audio_output,
+                audio_state["grid_sizes"],
+                audio_state["original_audio_shapes"],
+            )
+        )
+        result = {"video": video_output, "audio": audio_output}
+        return result if return_dict else (video_output, audio_output)
+
+    def _round_seq_len(self, seq_len: int) -> int:
+        if self.sp_world_size > 1:
+            return int(math.ceil(seq_len / self.sp_world_size) * self.sp_world_size)
+        return int(seq_len)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_path: Optional[str] = None,
+        video_pretrained_model_path: Optional[str] = None,
+        audio_pretrained_model_path: Optional[str] = None,
+        video_subfolder: Optional[str] = None,
+        audio_subfolder: Optional[str] = None,
+        video_kwargs: Optional[Dict] = None,
+        audio_kwargs: Optional[Dict] = None,
+        video_model_cls=Wan2_2Transformer3DModel,
+        audio_model_cls=CreatorAudioModel,
+        low_cpu_mem_usage: bool = False,
+        torch_dtype: torch.dtype = torch.bfloat16,
+        use_temporal_rope: bool = True,
+        audio_fps: float = 48000.0 / 960.0,
+        vae_temporal_stride: int = 4,
+        zero_init_cross_attn: bool = False,
+        zero_init_video_cross_attn: bool | None = None,
+        zero_init_audio_cross_attn: bool | None = None,
+        a2v_cross_attn_layers: LayerSelection = None,
+        v2a_cross_attn_layers: LayerSelection = None,
+        use_gating: bool = True,
+        use_a2v_gating: bool | None = None,
+        use_v2a_gating: bool | None = None,
+        zero_init_gating: bool = False,
+        zero_init_a2v_gating: bool | None = None,
+        zero_init_v2a_gating: bool | None = None,
+        gate_init_value: Optional[float] = None,
+        a2v_gate_init_value: Optional[float] = None,
+        v2a_gate_init_value: Optional[float] = None,
+        a2v_gate_alphas: LayerAlphas = None,
+        v2a_gate_alphas: LayerAlphas = None,
+    ):
+        video_kwargs = dict(video_kwargs or {})
+        audio_kwargs = dict(audio_kwargs or {})
+
+        if video_pretrained_model_path is not None and audio_pretrained_model_path is not None:
+            video_path = video_pretrained_model_path
+            audio_path = audio_pretrained_model_path
+        elif pretrained_model_path is not None:
+            video_path = os.path.join(pretrained_model_path, "video_model")
+            audio_path = os.path.join(pretrained_model_path, "audio_model")
+        else:
+            raise ValueError(
+                "Must provide either pretrained_model_path or both video_pretrained_model_path "
+                "and audio_pretrained_model_path"
+            )
+
+        logging.info("Loading video model from: %s", video_path)
+        video_model = video_model_cls.from_pretrained(
+            video_path,
+            subfolder=video_subfolder,
+            transformer_additional_kwargs=video_kwargs,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+            torch_dtype=torch_dtype,
+        )
+
+        logging.info("Loading audio model from: %s", audio_path)
+        audio_model = audio_model_cls.from_pretrained(
+            audio_path,
+            subfolder=audio_subfolder,
+            transformer_additional_kwargs=audio_kwargs,
+            low_cpu_mem_usage=low_cpu_mem_usage,
+            torch_dtype=torch_dtype,
+        )
+
+        model = cls(
+            video_model=video_model,
+            audio_model=audio_model,
+            use_temporal_rope=use_temporal_rope,
+            audio_fps=audio_fps,
+            vae_temporal_stride=vae_temporal_stride,
+            zero_init_cross_attn=zero_init_cross_attn,
+            zero_init_video_cross_attn=zero_init_video_cross_attn,
+            zero_init_audio_cross_attn=zero_init_audio_cross_attn,
+            a2v_cross_attn_layers=a2v_cross_attn_layers,
+            v2a_cross_attn_layers=v2a_cross_attn_layers,
+            use_gating=use_gating,
+            use_a2v_gating=use_a2v_gating,
+            use_v2a_gating=use_v2a_gating,
+            zero_init_gating=zero_init_gating,
+            zero_init_a2v_gating=zero_init_a2v_gating,
+            zero_init_v2a_gating=zero_init_v2a_gating,
+            gate_init_value=gate_init_value,
+            a2v_gate_init_value=a2v_gate_init_value,
+            v2a_gate_init_value=v2a_gate_init_value,
+            a2v_gate_alphas=a2v_gate_alphas,
+            v2a_gate_alphas=v2a_gate_alphas,
+        ).to(torch_dtype)
+        if pretrained_model_path is not None:
+            cross_attn_file = os.path.join(pretrained_model_path, "cross_attn_weights.safetensors")
+            cross_attn_file_bin = os.path.join(pretrained_model_path, "cross_attn_weights.bin")
+
+            if os.path.exists(cross_attn_file):
+                from safetensors.torch import load_file
+
+                cross_attn_state = load_file(cross_attn_file)
+                logging.info("Loading cross-attn weights from: %s (%d keys)", cross_attn_file, len(cross_attn_state))
+            elif os.path.exists(cross_attn_file_bin):
+                cross_attn_state = torch.load(cross_attn_file_bin, map_location="cpu")
+                logging.info("Loading cross-attn weights from: %s (%d keys)", cross_attn_file_bin, len(cross_attn_state))
+            else:
+                cross_attn_state = None
+                logging.warning("No cross_attn_weights found in %s, skipping.", pretrained_model_path)
+
+            if cross_attn_state is not None:
+                missing, unexpected = model.load_state_dict(cross_attn_state, strict=False)
+                logging.info("Cross-attn load: %d missing, %d unexpected keys", len(missing), len(unexpected))
+                if unexpected:
+                    logging.warning("Unexpected keys in cross_attn_weights: %s", unexpected[:10])
+
+        return model
+
+
+WanCreatorCrossAttnGatingAVModel = WanCreatorGatingAVModel
+WanCreatorGatedCrossAttnAVModel = WanCreatorGatingAVModel
+
+__all__ = [
+    "GatedCrossModalAttention",
+    "GatedJointBlock",
+    "WanCreatorGatingAVModel",
+    "WanCreatorCrossAttnGatingAVModel",
+    "WanCreatorGatedCrossAttnAVModel",
+    "compute_audio_temporal_positions",
+    "compute_video_temporal_positions",
+]

videox_fun/models/wan_text_encoder.py

@@ -0,0 +1,389 @@
+# Modified from https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/t5.py
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+
+
+def fp16_clamp(x):
+    if x.dtype == torch.float16 and torch.isinf(x).any():
+        clamp = torch.finfo(x.dtype).max - 1000
+        x = torch.clamp(x, min=-clamp, max=clamp)
+    return x
+
+
+def init_weights(m):
+    if isinstance(m, T5LayerNorm):
+        nn.init.ones_(m.weight)
+    elif isinstance(m, T5FeedForward):
+        nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+    elif isinstance(m, T5Attention):
+        nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)
+        nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)
+    elif isinstance(m, T5RelativeEmbedding):
+        nn.init.normal_(
+            m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)
+
+
+class GELU(nn.Module):
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.tanh(
+            math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super(T5LayerNorm, self).__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +
+                            self.eps)
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            x = x.type_as(self.weight)
+        return self.weight * x
+
+
+class T5Attention(nn.Module):
+    def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+        assert dim_attn % num_heads == 0
+        super(T5Attention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.num_heads = num_heads
+        self.head_dim = dim_attn // num_heads
+
+        # layers
+        self.q = nn.Linear(dim, dim_attn, bias=False)
+        self.k = nn.Linear(dim, dim_attn, bias=False)
+        self.v = nn.Linear(dim, dim_attn, bias=False)
+        self.o = nn.Linear(dim_attn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context=None, mask=None, pos_bias=None):
+        """
+        x:          [B, L1, C].
+        context:    [B, L2, C] or None.
+        mask:       [B, L2] or [B, L1, L2] or None.
+        """
+        # check inputs
+        context = x if context is None else context
+        b, n, c = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.q(x).view(b, -1, n, c)
+        k = self.k(context).view(b, -1, n, c)
+        v = self.v(context).view(b, -1, n, c)
+
+        # attention bias
+        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+        if pos_bias is not None:
+            attn_bias += pos_bias
+        if mask is not None:
+            assert mask.ndim in [2, 3]
+            mask = mask.view(b, 1, 1,
+                             -1) if mask.ndim == 2 else mask.unsqueeze(1)
+            attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+        # compute attention (T5 does not use scaling)
+        attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum('bnij,bjnc->binc', attn, v)
+
+        # output
+        x = x.reshape(b, -1, n * c)
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5FeedForward(nn.Module):
+
+    def __init__(self, dim, dim_ffn, dropout=0.1):
+        super(T5FeedForward, self).__init__()
+        self.dim = dim
+        self.dim_ffn = dim_ffn
+
+        # layers
+        self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+        self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+        self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x) * self.gate(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5SelfAttention(nn.Module):
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5SelfAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True)
+
+    def forward(self, x, mask=None, pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.ffn(self.norm2(x)))
+        return x
+
+
+class T5CrossAttention(nn.Module):
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5CrossAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm3 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=False)
+
+    def forward(self,
+                x,
+                mask=None,
+                encoder_states=None,
+                encoder_mask=None,
+                pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.cross_attn(
+            self.norm2(x), context=encoder_states, mask=encoder_mask))
+        x = fp16_clamp(x + self.ffn(self.norm3(x)))
+        return x
+
+
+class T5RelativeEmbedding(nn.Module):
+    def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+        super(T5RelativeEmbedding, self).__init__()
+        self.num_buckets = num_buckets
+        self.num_heads = num_heads
+        self.bidirectional = bidirectional
+        self.max_dist = max_dist
+
+        # layers
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, lq, lk):
+        device = self.embedding.weight.device
+        # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \
+        #     torch.arange(lq).unsqueeze(1).to(device)
+        if torch.device(type="meta") != device:
+            rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \
+                torch.arange(lq, device=device).unsqueeze(1)
+        else:
+            rel_pos = torch.arange(lk).unsqueeze(0) - \
+                torch.arange(lq).unsqueeze(1)
+        rel_pos = self._relative_position_bucket(rel_pos)
+        rel_pos_embeds = self.embedding(rel_pos)
+        rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(
+            0)  # [1, N, Lq, Lk]
+        return rel_pos_embeds.contiguous()
+
+    def _relative_position_bucket(self, rel_pos):
+        # preprocess
+        if self.bidirectional:
+            num_buckets = self.num_buckets // 2
+            rel_buckets = (rel_pos > 0).long() * num_buckets
+            rel_pos = torch.abs(rel_pos)
+        else:
+            num_buckets = self.num_buckets
+            rel_buckets = 0
+            rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+        # embeddings for small and large positions
+        max_exact = num_buckets // 2
+        rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /
+                                     math.log(self.max_dist / max_exact) *
+                                     (num_buckets - max_exact)).long()
+        rel_pos_large = torch.min(
+            rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+        rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+        return rel_buckets
+
+class WanT5EncoderModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+    def __init__(self,
+                 vocab,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(WanT5EncoderModel, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+            else nn.Embedding(vocab, dim)
+        self.pos_embedding = T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True) if shared_pos else None
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([
+            T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+                            shared_pos, dropout) for _ in range(num_layers)
+        ])
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(
+        self, 
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+    ):
+        x = self.token_embedding(input_ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1),
+                               x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, attention_mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return (x, )
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_path, additional_kwargs={}, low_cpu_mem_usage=False, torch_dtype=torch.bfloat16):
+        def filter_kwargs(cls, kwargs):
+            import inspect
+            sig = inspect.signature(cls.__init__)
+            valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+            return filtered_kwargs
+
+        if low_cpu_mem_usage:
+            try:
+                import re
+
+                from diffusers import __version__ as diffusers_version
+                if diffusers_version >= "0.33.0":
+                    from diffusers.models.model_loading_utils import \
+                        load_model_dict_into_meta
+                else:
+                    from diffusers.models.modeling_utils import \
+                        load_model_dict_into_meta
+                from diffusers.utils import is_accelerate_available
+                if is_accelerate_available():
+                    import accelerate
+                
+                # Instantiate model with empty weights
+                with accelerate.init_empty_weights():
+                    model = cls(**filter_kwargs(cls, additional_kwargs))
+
+                param_device = "cpu"
+                if pretrained_model_path.endswith(".safetensors"):
+                    from safetensors.torch import load_file
+                    state_dict = load_file(pretrained_model_path)
+                else:
+                    state_dict = torch.load(pretrained_model_path, map_location="cpu")
+
+                if diffusers_version >= "0.33.0":
+                    # Diffusers has refactored `load_model_dict_into_meta` since version 0.33.0 in this commit:
+                    # https://github.com/huggingface/diffusers/commit/f5929e03060d56063ff34b25a8308833bec7c785.
+                    load_model_dict_into_meta(
+                        model,
+                        state_dict,
+                        dtype=torch_dtype,
+                        model_name_or_path=pretrained_model_path,
+                    )
+                else:
+                    # move the params from meta device to cpu
+                    missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+                    if len(missing_keys) > 0:
+                        raise ValueError(
+                            f"Cannot load {cls} from {pretrained_model_path} because the following keys are"
+                            f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
+                            " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
+                            " those weights or else make sure your checkpoint file is correct."
+                        )
+
+                    unexpected_keys = load_model_dict_into_meta(
+                        model,
+                        state_dict,
+                        device=param_device,
+                        dtype=torch_dtype,
+                        model_name_or_path=pretrained_model_path,
+                    )
+
+                    if cls._keys_to_ignore_on_load_unexpected is not None:
+                        for pat in cls._keys_to_ignore_on_load_unexpected:
+                            unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
+                    if len(unexpected_keys) > 0:
+                        logging.warning("Unused text encoder keys: %d", len(unexpected_keys))
+                
+                return model
+            except Exception:
+                logging.warning("Falling back to regular text encoder loading")
+        
+        model = cls(**filter_kwargs(cls, additional_kwargs))
+        if pretrained_model_path.endswith(".safetensors"):
+            from safetensors.torch import load_file, safe_open
+            state_dict = load_file(pretrained_model_path)
+        else:
+            state_dict = torch.load(pretrained_model_path, map_location="cpu")
+        m, u = model.load_state_dict(state_dict, strict=False)
+        model = model.to(torch_dtype)
+        return model

videox_fun/models/wan_transformer3d_prope.py

@@ -0,0 +1,1035 @@
+# Modified from https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/model.py
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import glob
+import json
+import math
+import os
+import logging
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+from .attention_utils import attention
+
+
+def sinusoidal_embedding_1d(dim, position):
+    # preprocess
+    assert dim % 2 == 0
+    half = dim // 2
+    position = position.type(torch.float64)
+
+    # calculation
+    sinusoid = torch.outer(
+        position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x
+
+
+@amp.autocast(enabled=False)
+def rope_params(max_seq_len, dim, theta=10000):
+    assert dim % 2 == 0
+    freqs = torch.outer(
+        torch.arange(max_seq_len),
+        1.0 / torch.pow(theta,
+                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
+    freqs = torch.polar(torch.ones_like(freqs), freqs)
+    return freqs
+
+
+# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid
+def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
+    tw = tgt_width
+    th = tgt_height
+    h, w = src
+    r = h / w
+    if r > (th / tw):
+        resize_height = th
+        resize_width = int(round(th / h * w))
+    else:
+        resize_width = tw
+        resize_height = int(round(tw / w * h))
+
+    crop_top = int(round((th - resize_height) / 2.0))
+    crop_left = int(round((tw - resize_width) / 2.0))
+
+    return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)
+
+
+@amp.autocast(enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float32).reshape(
+            seq_len, n, -1, 2))
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+                            dim=-1).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[i, seq_len:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).to(x.dtype)
+
+
+def rope_apply_qk(q, k, grid_sizes, freqs):
+    q = rope_apply(q, grid_sizes, freqs)
+    k = rope_apply(k, grid_sizes, freqs)
+    return q, k
+
+
+class WanRMSNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class WanLayerNorm(nn.LayerNorm):
+
+    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class WanSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, seq_lens, grid_sizes, freqs, dtype=torch.bfloat16, t=0, **kwargs):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            q = self.norm_q(self.q(x.to(dtype))).view(b, s, n, d)
+            k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
+            v = self.v(x.to(dtype)).view(b, s, n, d)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+
+        q, k = rope_apply_qk(q, k, grid_sizes, freqs)
+
+        x = attention(
+            q.to(dtype), 
+            k.to(dtype), 
+            v=v.to(dtype),
+            k_lens=seq_lens,
+            window_size=self.window_size)
+        x = x.to(dtype)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+    def forward(self, x, context, context_lens, dtype=torch.bfloat16, t=0):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+        k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+        v = self.v(context.to(dtype)).view(b, -1, n, d)
+
+        # compute attention
+        x = attention(
+            q.to(dtype), 
+            k.to(dtype), 
+            v.to(dtype), 
+            k_lens=context_lens
+        )
+        x = x.to(dtype)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+        self.k_img = nn.Linear(dim, dim)
+        self.v_img = nn.Linear(dim, dim)
+        # self.alpha = nn.Parameter(torch.zeros((1, )))
+        self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, context, context_lens, dtype=torch.bfloat16, t=0):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        context_img = context[:, :257]
+        context = context[:, 257:]
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+        k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+        v = self.v(context.to(dtype)).view(b, -1, n, d)
+        k_img = self.norm_k_img(self.k_img(context_img.to(dtype))).view(b, -1, n, d)
+        v_img = self.v_img(context_img.to(dtype)).view(b, -1, n, d)
+
+        img_x = attention(
+            q.to(dtype), 
+            k_img.to(dtype), 
+            v_img.to(dtype), 
+            k_lens=None
+        )
+        img_x = img_x.to(dtype)
+        # compute attention
+        x = attention(
+            q.to(dtype), 
+            k.to(dtype), 
+            v.to(dtype), 
+            k_lens=context_lens
+        )
+        x = x.to(dtype)
+
+        # output
+        x = x.flatten(2)
+        img_x = img_x.flatten(2)
+        x = x + img_x
+        x = self.o(x)
+        return x
+
+
+class WanCrossAttention(WanSelfAttention):
+    def forward(self, x, context, context_lens, dtype=torch.bfloat16, t=0):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+        # compute query, key, value
+        q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+        k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+        v = self.v(context.to(dtype)).view(b, -1, n, d)
+        # compute attention
+        x = attention(q.to(dtype), k.to(dtype), v.to(dtype), k_lens=context_lens)
+        # output
+        x = x.flatten(2)
+        x = self.o(x.to(dtype))
+        return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+    't2v_cross_attn': WanT2VCrossAttention,
+    'i2v_cross_attn': WanI2VCrossAttention,
+    'cross_attn': WanCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+
+    def __init__(self,
+                 cross_attn_type,
+                 dim,
+                 ffn_dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=False,
+                 eps=1e-6,
+                 **kwargs):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # layers
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+                                          eps)
+        self.norm3 = WanLayerNorm(
+            dim, eps,
+            elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,
+                                                                      num_heads,
+                                                                      (-1, -1),
+                                                                      qk_norm,
+                                                                      eps)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+            nn.Linear(ffn_dim, dim))
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+        dtype=torch.bfloat16,
+        t=0,
+        group=None,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        if e.dim() > 3:
+            e = (self.modulation.unsqueeze(0) + e).chunk(6, dim=2)
+            e = [e.squeeze(2) for e in e]
+        else:        
+            e = (self.modulation + e).chunk(6, dim=1)
+
+        # self-attention (RoPE branch)
+        temp_x = self.norm1(x) * (1 + e[1]) + e[0]
+        temp_x = temp_x.to(dtype)
+
+        y = self.self_attn(temp_x, seq_lens, grid_sizes, freqs, dtype, group=group)
+
+        x = x + y * e[2]
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e):
+            # cross-attention
+            x = x + self.cross_attn(self.norm3(x), context, context_lens, dtype)
+
+            # ffn function
+            temp_x = self.norm2(x) * (1 + e[4]) + e[3]
+            temp_x = temp_x.to(dtype)
+            
+            y = self.ffn(temp_x)
+            x = x + y * e[5]
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+        return x
+
+
+class Head(nn.Module):
+
+    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+
+        # layers
+        out_dim = math.prod(patch_size) * out_dim
+        self.norm = WanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_dim)
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, e):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            e(Tensor): Shape [B, C]
+        """
+        if e.dim() > 2:
+            e = (self.modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2)
+            e = [e.squeeze(2) for e in e]
+        else:
+            e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+        
+        x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+        return x
+
+
+class MLPProj(torch.nn.Module):
+
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+
+        self.proj = torch.nn.Sequential(
+            torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
+            torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
+            torch.nn.LayerNorm(out_dim))
+
+    def forward(self, image_embeds):
+        clip_extra_context_tokens = self.proj(image_embeds)
+        return clip_extra_context_tokens
+
+
+
+class WanTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    # ignore_for_config = [
+    #     'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+    # ]
+    # _no_split_modules = ['WanAttentionBlock']
+    @register_to_config
+    def __init__(
+        self,
+        model_type='t2v',
+        patch_size=(1, 2, 2),
+        text_len=512,
+        in_dim=16,
+        dim=2048,
+        ffn_dim=8192,
+        freq_dim=256,
+        text_dim=4096,
+        out_dim=16,
+        num_heads=16,
+        num_layers=32,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=True,
+        eps=1e-6,
+        in_channels=16,
+        hidden_size=2048,
+        add_control_adapter=True,
+        control_adapter_type="baseline",
+        in_dim_control_adapter=24,
+        downscale_factor_control_adapter=8,
+        add_ref_conv=False,
+        in_dim_ref_conv=16,
+        cross_attn_type=None,
+        cam_method='prope',
+        attn_compress=2,
+        cam_self_attn_layers=None,
+        traning=False,
+    ):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        # assert model_type in ['t2v', 'i2v', 'ti2v']
+        self.model_type = model_type
+
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # embeddings
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim))
+
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks
+        if cross_attn_type is None:
+            if model_type == "t2v":
+                cross_attn_type = 't2v_cross_attn'
+            else:
+                cross_attn_type = 'i2v_cross_attn'
+            # else:
+                # cross_attn_type = "cross_attn"
+
+        attn_class = WanAttentionBlock # by default use WanAttentionBlock
+        self.control_adapter = None
+
+        self.blocks = nn.ModuleList([
+            attn_class(cross_attn_type, dim, ffn_dim, num_heads,
+                              window_size, qk_norm, cross_attn_norm, eps)
+                               
+            for _ in range(num_layers)
+        ])
+
+        # head
+        self.head = Head(dim, out_dim, patch_size, eps)
+
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+        self.d = d
+        self.dim = dim
+        self.freqs = torch.cat(
+            [
+                rope_params(1024, d - 4 * (d // 6)),
+                rope_params(1024, 2 * (d // 6)),
+                rope_params(1024, 2 * (d // 6))
+            ],
+            dim=1
+        )
+
+        if model_type == 'i2v':
+            self.img_emb = MLPProj(1280, dim)
+
+        if add_ref_conv:
+            self.ref_conv = nn.Conv2d(in_dim_ref_conv, dim, kernel_size=patch_size[1:], stride=patch_size[1:])
+        else:
+            self.ref_conv = None
+
+        self.customize_initialize_weights()
+    
+    def customize_initialize_weights(self):
+        self.init_weights()
+
+    def forward(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        clip_fea=None,
+        y=None,
+        y_camera=None,
+        full_ref=None,
+        subject_ref=None,
+        cond_flag=True,
+        dtype=torch.bfloat16,
+    ):
+        r"""
+        Forward pass through the diffusion model
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            clip_fea (Tensor, *optional*):
+                CLIP image features for image-to-video mode
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+            cond_flag (`bool`, *optional*, defaults to True):
+                Flag to indicate whether to forward the condition input
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+
+        device = self.patch_embedding.weight.device
+        dtype = x[0].dtype
+        if self.freqs.device != device and torch.device(type="meta") != device:
+            self.freqs = self.freqs.to(device)
+
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        grid_sizes = torch.stack(
+            [torch.tensor(u.shape[2:], dtype=torch.long, device=device) for u in x])
+
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long, device=device)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat([
+            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+                      dim=1) for u in x
+        ])
+
+        # time embeddings
+        with amp.autocast(dtype=torch.float32):
+            if t.dim() != 1:
+                if t.size(1) < seq_len:
+                    pad_size = seq_len - t.size(1)
+                    last_elements = t[:, -1].unsqueeze(1)
+                    padding = last_elements.repeat(1, pad_size)
+                    t = torch.cat([t, padding], dim=1)
+                bt = t.size(0)
+                ft = t.flatten()
+                e = self.time_embedding(
+                    sinusoidal_embedding_1d(self.freq_dim,
+                                            ft).unflatten(0, (bt, seq_len)).float())
+                e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+            else:
+                e = self.time_embedding(
+                    sinusoidal_embedding_1d(self.freq_dim, t).float())
+                e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([
+                torch.cat(
+                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                for u in context
+            ]))
+
+        # if clip_fea is not None:
+        #     context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
+        #     context = torch.concat([context_clip, context], dim=1)
+
+        for block in self.blocks:
+            x = block(
+                x, e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes,
+                freqs=self.freqs, context=context, context_lens=context_lens,
+                dtype=dtype, t=t, group=None,
+            )
+
+        x = self.head(x, e)
+
+        return torch.stack(self.unpatchify(x, grid_sizes))
+
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum('fhwpqrc->cfphqwr', u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out
+
+    def init_weights(self):
+        r"""
+        Initialize model parameters using Xavier initialization.
+        """
+
+        # basic init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # init embeddings
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.text_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+
+        # init output layer
+        nn.init.zeros_(self.head.head.weight)
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_path, subfolder=None, transformer_additional_kwargs={},
+        low_cpu_mem_usage=False, torch_dtype=torch.bfloat16
+    ):
+        if subfolder is not None:
+            pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
+        config_file = os.path.join(pretrained_model_path, 'config.json')
+        if not os.path.isfile(config_file):
+            raise RuntimeError(f"{config_file} does not exist")
+        with open(config_file, "r") as f:
+            config = json.load(f)
+
+        from diffusers.utils import WEIGHTS_NAME
+        model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
+        model_file_safetensors = model_file.replace(".bin", ".safetensors")
+
+        if "dict_mapping" in transformer_additional_kwargs.keys():
+            for key in transformer_additional_kwargs["dict_mapping"]:
+                transformer_additional_kwargs[transformer_additional_kwargs["dict_mapping"][key]] = config[key]
+
+        if low_cpu_mem_usage:
+            try:
+                import re
+
+                from diffusers import __version__ as diffusers_version
+                if diffusers_version >= "0.33.0":
+                    from diffusers.models.model_loading_utils import \
+                        load_model_dict_into_meta
+                else:
+                    from diffusers.models.modeling_utils import \
+                        load_model_dict_into_meta
+                from diffusers.utils import is_accelerate_available
+                if is_accelerate_available():
+                    import accelerate
+                
+                # Instantiate model with empty weights
+                with accelerate.init_empty_weights():
+                    model = cls.from_config(config, **transformer_additional_kwargs)
+
+                param_device = "cpu"
+                if os.path.exists(model_file):
+                    state_dict = torch.load(model_file, map_location="cpu")
+                elif os.path.exists(model_file_safetensors):
+                    from safetensors.torch import load_file, safe_open
+                    state_dict = load_file(model_file_safetensors)
+                else:
+                    from safetensors.torch import load_file, safe_open
+                    model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+                    state_dict = {}
+                    for _model_file_safetensors in model_files_safetensors:
+                        _state_dict = load_file(_model_file_safetensors)
+                        for key in _state_dict:
+                            state_dict[key] = _state_dict[key]
+
+                if model.state_dict()['patch_embedding.weight'].size() != state_dict['patch_embedding.weight'].size():
+                    model.state_dict()['patch_embedding.weight'][:, :state_dict['patch_embedding.weight'].size()[1], :, :] = state_dict['patch_embedding.weight'][:, :model.state_dict()['patch_embedding.weight'].size()[1], :, :]
+                    model.state_dict()['patch_embedding.weight'][:, state_dict['patch_embedding.weight'].size()[1]:, :, :] = 0
+                    state_dict['patch_embedding.weight'] = model.state_dict()['patch_embedding.weight']
+
+                filtered_state_dict = {}
+                for key in state_dict:
+                    if key in model.state_dict() and model.state_dict()[key].size() == state_dict[key].size():
+                        filtered_state_dict[key] = state_dict[key]
+                        
+                model_keys = set(model.state_dict().keys())
+                loaded_keys = set(filtered_state_dict.keys())
+                missing_keys = model_keys - loaded_keys
+
+                def initialize_missing_parameters(missing_keys, model_state_dict, torch_dtype=None):
+                    initialized_dict = {}
+                    
+                    with torch.no_grad():
+                        for key in missing_keys:
+                            param_shape = model_state_dict[key].shape
+                            param_dtype = torch_dtype if torch_dtype is not None else model_state_dict[key].dtype
+                            if 'weight' in key:
+                                if any(norm_type in key for norm_type in ['norm', 'ln_', 'layer_norm', 'group_norm', 'batch_norm']):
+                                    initialized_dict[key] = torch.ones(param_shape, dtype=param_dtype)
+                                elif 'embedding' in key or 'embed' in key:
+                                    initialized_dict[key] = torch.randn(param_shape, dtype=param_dtype) * 0.02
+                                elif 'head' in key or 'output' in key or 'proj_out' in key or 'project_out' in key or "out" in key or 'out_proj' in key:
+                                    logging.info(f"Initialize {key} with zeros")
+                                    initialized_dict[key] = torch.zeros(param_shape, dtype=param_dtype)
+                                elif len(param_shape) >= 2:
+                                    initialized_dict[key] = torch.empty(param_shape, dtype=param_dtype)
+                                    nn.init.xavier_uniform_(initialized_dict[key])
+                                else:
+                                    initialized_dict[key] = torch.randn(param_shape, dtype=param_dtype) * 0.02
+                            elif 'bias' in key:
+                                initialized_dict[key] = torch.zeros(param_shape, dtype=param_dtype)
+                            elif 'running_mean' in key:
+                                initialized_dict[key] = torch.zeros(param_shape, dtype=param_dtype)
+                            elif 'running_var' in key:
+                                initialized_dict[key] = torch.ones(param_shape, dtype=param_dtype)
+                            elif 'num_batches_tracked' in key:
+                                initialized_dict[key] = torch.zeros(param_shape, dtype=torch.long)
+                            else:
+                                initialized_dict[key] = torch.zeros(param_shape, dtype=param_dtype)
+
+                    return initialized_dict
+
+                if missing_keys:
+                    initialized_params = initialize_missing_parameters(
+                        missing_keys, 
+                        model.state_dict(), 
+                        torch_dtype
+                    )
+                    filtered_state_dict.update(initialized_params)
+                    # logging.info(f"Weights of Camera_adapter: {initialized_params}")
+
+                if diffusers_version >= "0.33.0":
+                    # Diffusers has refactored `load_model_dict_into_meta` since version 0.33.0 in this commit:
+                    # https://github.com/huggingface/diffusers/commit/f5929e03060d56063ff34b25a8308833bec7c785.
+                    load_model_dict_into_meta(
+                        model,
+                        filtered_state_dict,
+                        dtype=torch_dtype,
+                        model_name_or_path=pretrained_model_path,
+                    )
+                else:
+                    model._convert_deprecated_attention_blocks(filtered_state_dict)
+                    unexpected_keys = load_model_dict_into_meta(
+                        model,
+                        filtered_state_dict,
+                        device=param_device,
+                        dtype=torch_dtype,
+                        model_name_or_path=pretrained_model_path,
+                    )
+
+                    if cls._keys_to_ignore_on_load_unexpected is not None:
+                        for pat in cls._keys_to_ignore_on_load_unexpected:
+                            unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
+                    if len(unexpected_keys) > 0:
+                        logging.warning("Unused checkpoint keys: %d", len(unexpected_keys))
+                
+                return model
+            except Exception:
+                logging.warning("Falling back to regular transformer loading")
+
+        model = cls.from_config(config, **transformer_additional_kwargs)
+        if os.path.exists(model_file):
+            state_dict = torch.load(model_file, map_location="cpu")
+        elif os.path.exists(model_file_safetensors):
+            from safetensors.torch import load_file, safe_open
+            state_dict = load_file(model_file_safetensors)
+        else:
+            from safetensors.torch import load_file, safe_open
+            model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+            state_dict = {}
+            for _model_file_safetensors in model_files_safetensors:
+                _state_dict = load_file(_model_file_safetensors)
+                for key in _state_dict:
+                    state_dict[key] = _state_dict[key]
+        
+        if model.state_dict()['patch_embedding.weight'].size() != state_dict['patch_embedding.weight'].size():
+            model.state_dict()['patch_embedding.weight'][:, :state_dict['patch_embedding.weight'].size()[1], :, :] = state_dict['patch_embedding.weight'][:, :model.state_dict()['patch_embedding.weight'].size()[1], :, :]
+            model.state_dict()['patch_embedding.weight'][:, state_dict['patch_embedding.weight'].size()[1]:, :, :] = 0
+            state_dict['patch_embedding.weight'] = model.state_dict()['patch_embedding.weight']
+        
+        tmp_state_dict = {} 
+        for key in state_dict:
+            if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size():
+                tmp_state_dict[key] = state_dict[key]
+                
+        state_dict = tmp_state_dict
+
+        m, u = model.load_state_dict(state_dict, strict=False)
+        model = model.to(torch_dtype)
+        return model
+
+
+class Wan2_2Transformer3DModel(WanTransformer3DModel):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    # ignore_for_config = [
+    #     'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+    # ]
+    # _no_split_modules = ['WanAttentionBlock']
+    def __init__(
+        self,
+        model_type='t2v',
+        patch_size=(1, 2, 2),
+        text_len=512,
+        in_dim=16,
+        dim=2048,
+        ffn_dim=8192,
+        freq_dim=256,
+        text_dim=4096,
+        out_dim=16,
+        num_heads=16,
+        num_layers=32,
+        window_size=(-1, -1),
+        qk_norm=True,
+        cross_attn_norm=True,
+        eps=1e-6,
+        in_channels=16,
+        hidden_size=2048,
+        add_control_adapter=False,
+        in_dim_control_adapter=24,
+        downscale_factor_control_adapter=8,
+        add_ref_conv=False,
+        control_adapter_type="baseline",
+        in_dim_ref_conv=16,
+        cam_method='prope',
+        attn_compress=2,
+        cam_self_attn_layers=None,
+    ):
+        r"""
+        Initialize the diffusion model backbone.
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+        super().__init__(
+            model_type=model_type,
+            patch_size=patch_size,
+            text_len=text_len,
+            in_dim=in_dim,
+            dim=dim,
+            ffn_dim=ffn_dim,
+            freq_dim=freq_dim,
+            text_dim=text_dim,
+            out_dim=out_dim,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            window_size=window_size,
+            qk_norm=qk_norm,
+            cross_attn_norm=cross_attn_norm,
+            eps=eps,
+            in_channels=in_channels,
+            hidden_size=hidden_size,
+            add_control_adapter=add_control_adapter,
+            control_adapter_type=control_adapter_type,
+            in_dim_control_adapter=in_dim_control_adapter,
+            downscale_factor_control_adapter=downscale_factor_control_adapter,
+            add_ref_conv=add_ref_conv,
+            in_dim_ref_conv=in_dim_ref_conv,
+            cross_attn_type="cross_attn",
+            cam_method=cam_method,
+            attn_compress=attn_compress,
+            cam_self_attn_layers=cam_self_attn_layers,
+        )
+        
+        if hasattr(self, "img_emb"):
+            del self.img_emb

videox_fun/models/wan_vae3_8.py

@@ -0,0 +1,1248 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from typing import Tuple, Union
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.autoencoders.vae import (DecoderOutput,
+                                               DiagonalGaussianDistribution)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from einops import rearrange
+
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return (F.normalize(x, dim=(1 if self.channel_first else -1)) *
+                self.scale * self.gamma + self.bias)
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in (
+            "none",
+            "upsample2d",
+            "upsample3d",
+            "downsample2d",
+            "downsample3d",
+        )
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+                # nn.Conv2d(dim, dim//2, 3, padding=1)
+            )
+            self.time_conv = CausalConv3d(
+                dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(
+                dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
+                            feat_cache[idx] != "Rep"):
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat(
+                            [
+                                feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                    cache_x.device),
+                                cache_x,
+                            ],
+                            dim=2,
+                        )
+                    if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
+                            feat_cache[idx] == "Rep"):
+                        cache_x = torch.cat(
+                            [
+                                torch.zeros_like(cache_x).to(cache_x.device),
+                                cache_x
+                            ],
+                            dim=2,
+                        )
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+                                    3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(
+                        torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data.detach().clone()
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight = nn.Parameter(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = (
+            CausalConv3d(in_dim, out_dim, 1)
+            if in_dim != out_dim else nn.Identity())
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = (
+            self.to_qkv(x).reshape(b * t, 1, c * 3,
+                                   -1).permute(0, 1, 3,
+                                               2).contiguous().chunk(3, dim=-1))
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+def patchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        x = rearrange(
+            x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b c f (h q) (w r) -> b (c r q) f h w",
+            q=patch_size,
+            r=patch_size,
+        )
+    else:
+        raise ValueError(f"Invalid input shape: {x.shape}")
+
+    return x
+
+
+def unpatchify(x, patch_size):
+    if patch_size == 1:
+        return x
+
+    if x.dim() == 4:
+        x = rearrange(
+            x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
+    elif x.dim() == 5:
+        x = rearrange(
+            x,
+            "b (c r q) f h w -> b c f (h q) (w r)",
+            q=patch_size,
+            r=patch_size,
+        )
+    return x
+
+
+class AvgDown3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert in_channels * self.factor % out_channels == 0
+        self.group_size = in_channels * self.factor // out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+        pad = (0, 0, 0, 0, pad_t, 0)
+        x = F.pad(x, pad)
+        B, C, T, H, W = x.shape
+        x = x.view(
+            B,
+            C,
+            T // self.factor_t,
+            self.factor_t,
+            H // self.factor_s,
+            self.factor_s,
+            W // self.factor_s,
+            self.factor_s,
+        )
+        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+        x = x.view(
+            B,
+            C * self.factor,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.view(
+            B,
+            self.out_channels,
+            self.group_size,
+            T // self.factor_t,
+            H // self.factor_s,
+            W // self.factor_s,
+        )
+        x = x.mean(dim=2)
+        return x
+
+
+class DupUp3D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        factor_t,
+        factor_s=1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+
+        assert out_channels * self.factor % in_channels == 0
+        self.repeats = out_channels * self.factor // in_channels
+
+    def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            self.factor_t,
+            self.factor_s,
+            self.factor_s,
+            x.size(2),
+            x.size(3),
+            x.size(4),
+        )
+        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            x.size(2) * self.factor_t,
+            x.size(4) * self.factor_s,
+            x.size(6) * self.factor_s,
+        )
+        if first_chunk:
+            x = x[:, :, self.factor_t - 1:, :, :]
+        return x
+
+
+class Down_ResidualBlock(nn.Module):
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 dropout,
+                 mult,
+                 temperal_downsample=False,
+                 down_flag=False):
+        super().__init__()
+
+        # Shortcut path with downsample
+        self.avg_shortcut = AvgDown3D(
+            in_dim,
+            out_dim,
+            factor_t=2 if temperal_downsample else 1,
+            factor_s=2 if down_flag else 1,
+        )
+
+        # Main path with residual blocks and downsample
+        downsamples = []
+        for _ in range(mult):
+            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final downsample block
+        if down_flag:
+            mode = "downsample3d" if temperal_downsample else "downsample2d"
+            downsamples.append(Resample(out_dim, mode=mode))
+
+        self.downsamples = nn.Sequential(*downsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        x_copy = x.clone()
+        for module in self.downsamples:
+            x = module(x, feat_cache, feat_idx)
+
+        return x + self.avg_shortcut(x_copy)
+
+
+class Up_ResidualBlock(nn.Module):
+
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 dropout,
+                 mult,
+                 temperal_upsample=False,
+                 up_flag=False):
+        super().__init__()
+        # Shortcut path with upsample
+        if up_flag:
+            self.avg_shortcut = DupUp3D(
+                in_dim,
+                out_dim,
+                factor_t=2 if temperal_upsample else 1,
+                factor_s=2 if up_flag else 1,
+            )
+        else:
+            self.avg_shortcut = None
+
+        # Main path with residual blocks and upsample
+        upsamples = []
+        for _ in range(mult):
+            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+
+        # Add the final upsample block
+        if up_flag:
+            mode = "upsample3d" if temperal_upsample else "upsample2d"
+            upsamples.append(Resample(out_dim, mode=mode))
+
+        self.upsamples = nn.Sequential(*upsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        x_main = x.clone()
+        for module in self.upsamples:
+            x_main = module(x_main, feat_cache, feat_idx)
+        if self.avg_shortcut is not None:
+            x_shortcut = self.avg_shortcut(x, first_chunk)
+            return x_main + x_shortcut
+        else:
+            return x_main
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_down_flag = (
+                temperal_downsample[i]
+                if i < len(temperal_downsample) else False)
+            downsamples.append(
+                Down_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks,
+                    temperal_downsample=t_down_flag,
+                    down_flag=i != len(dim_mult) - 1,
+                ))
+            scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+
+        # # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device),
+                        cache_x,
+                    ],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[False, True, True],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_up_flag = temperal_upsample[i] if i < len(
+                temperal_upsample) else False
+            upsamples.append(
+                Up_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks + 1,
+                    temperal_upsample=t_up_flag,
+                    up_flag=i != len(dim_mult) - 1,
+                ))
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 12, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device),
+                        cache_x,
+                    ],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx, first_chunk)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class AutoencoderKLWan2_2_(nn.Module):
+
+    def __init__(
+        self,
+        dim=160,
+        dec_dim=256,
+        z_dim=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, False],
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(
+            dim,
+            z_dim * 2,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_downsample,
+            dropout,
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dec_dim,
+            z_dim,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_upsample,
+            dropout,
+        )
+
+    def forward(self, x, scale=[0, 1]):
+        mu = self.encode(x, scale)
+        x_recon = self.decode(mu, scale)
+        return x_recon, mu
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        # z: [b,c,t,h,w]
+        scale = [item.to(x.device, x.dtype) for item in scale]
+        x = patchify(x, patch_size=2)
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        x = torch.cat([mu, log_var], dim = 1)
+        self.clear_cache()
+        return x
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        # z: [b,c,t,h,w]
+        scale = [item.to(z.device, z.dtype) for item in scale]
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                    first_chunk=True,
+                )
+            else:
+                out_ = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        out = unpatchify(out, patch_size=2)
+        self.clear_cache()
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs):
+    # params
+    cfg = dict(
+        dim=dim,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[True, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    # init model
+    model = AutoencoderKLWan2_2_(**cfg)
+
+    return model
+
+
+class AutoencoderKLWan3_8(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+    @register_to_config
+    def __init__(
+        self,
+        latent_channels=48,
+        c_dim=160,
+        vae_pth=None,
+        dim_mult=[1, 2, 4, 4],
+        temperal_downsample=[False, True, True],
+        temporal_compression_ratio=4,
+        spatial_compression_ratio=8
+    ):
+        super().__init__()
+        mean = torch.tensor(
+            [
+                -0.2289,
+                -0.0052,
+                -0.1323,
+                -0.2339,
+                -0.2799,
+                0.0174,
+                0.1838,
+                0.1557,
+                -0.1382,
+                0.0542,
+                0.2813,
+                0.0891,
+                0.1570,
+                -0.0098,
+                0.0375,
+                -0.1825,
+                -0.2246,
+                -0.1207,
+                -0.0698,
+                0.5109,
+                0.2665,
+                -0.2108,
+                -0.2158,
+                0.2502,
+                -0.2055,
+                -0.0322,
+                0.1109,
+                0.1567,
+                -0.0729,
+                0.0899,
+                -0.2799,
+                -0.1230,
+                -0.0313,
+                -0.1649,
+                0.0117,
+                0.0723,
+                -0.2839,
+                -0.2083,
+                -0.0520,
+                0.3748,
+                0.0152,
+                0.1957,
+                0.1433,
+                -0.2944,
+                0.3573,
+                -0.0548,
+                -0.1681,
+                -0.0667,
+            ], dtype=torch.float32
+        )
+        std = torch.tensor(
+            [
+                0.4765,
+                1.0364,
+                0.4514,
+                1.1677,
+                0.5313,
+                0.4990,
+                0.4818,
+                0.5013,
+                0.8158,
+                1.0344,
+                0.5894,
+                1.0901,
+                0.6885,
+                0.6165,
+                0.8454,
+                0.4978,
+                0.5759,
+                0.3523,
+                0.7135,
+                0.6804,
+                0.5833,
+                1.4146,
+                0.8986,
+                0.5659,
+                0.7069,
+                0.5338,
+                0.4889,
+                0.4917,
+                0.4069,
+                0.4999,
+                0.6866,
+                0.4093,
+                0.5709,
+                0.6065,
+                0.6415,
+                0.4944,
+                0.5726,
+                1.2042,
+                0.5458,
+                1.6887,
+                0.3971,
+                1.0600,
+                0.3943,
+                0.5537,
+                0.5444,
+                0.4089,
+                0.7468,
+                0.7744,
+            ], dtype=torch.float32
+        )
+        self.scale = [mean, 1.0 / std]
+
+        # init model
+        self.model = _video_vae(
+                pretrained_path=vae_pth,
+                z_dim=latent_channels,
+                dim=c_dim,
+                dim_mult=dim_mult,
+                temperal_downsample=temperal_downsample,
+            ).eval().requires_grad_(False)
+
+        self.use_tiling = False
+        self.tile_size = (34, 34)
+        self.tile_stride = (18, 16)
+
+    def enable_tiling(self, tile_size=(34, 34), tile_stride=(18, 16)):
+        """Enable spatial tiling for encode/decode to reduce peak GPU memory."""
+        self.use_tiling = True
+        self.tile_size = tile_size
+        self.tile_stride = tile_stride
+
+    def disable_tiling(self):
+        """Disable spatial tiling."""
+        self.use_tiling = False
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        x = [
+            self.model.encode(u.unsqueeze(0), self.scale).squeeze(0)
+            for u in x
+        ]
+        x = torch.stack(x)
+        return x
+
+    @apply_forward_hook
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        if self.use_tiling:
+            return self.tiled_encode(x, tile_size=self.tile_size, tile_stride=self.tile_stride)
+
+        h = self._encode(x)
+
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, zs):
+        dec = [
+            self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0)
+            for u in zs
+        ]
+        dec = torch.stack(dec)
+
+        return DecoderOutput(sample=dec)
+
+    @apply_forward_hook
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_tiling:
+            return self.tiled_decode(z, tile_size=self.tile_size, tile_stride=self.tile_stride)
+
+        decoded = self._decode(z).sample
+
+        if not return_dict:
+            return (decoded,)
+        return DecoderOutput(sample=decoded)
+
+    @staticmethod
+    def _build_1d_mask(length, left_bound, right_bound, border_width):
+        """Build a 1D blending mask with linear ramp at non-boundary edges."""
+        mask = torch.ones((length,))
+        if not left_bound:
+            mask[:border_width] = (torch.arange(border_width) + 1) / border_width
+        if not right_bound:
+            mask[-border_width:] = torch.flip(
+                (torch.arange(border_width) + 1) / border_width, dims=(0,)
+            )
+        return mask
+
+    def _build_spatial_mask(self, data, is_bound, border_width):
+        """Build a 2D spatial blending mask from H and W 1D masks.
+
+        Args:
+            data: tensor of shape [B, C, T, H, W] (only H, W are used).
+            is_bound: (top, bottom, left, right) booleans.
+            border_width: (border_h, border_w) in pixel space.
+        """
+        _, _, _, height, width = data.shape
+        h_mask = self._build_1d_mask(height, is_bound[0], is_bound[1], border_width[0])
+        w_mask = self._build_1d_mask(width, is_bound[2], is_bound[3], border_width[1])
+        h_mask = h_mask.unsqueeze(1).expand(height, width)
+        w_mask = w_mask.unsqueeze(0).expand(height, width)
+        mask = torch.stack([h_mask, w_mask]).min(dim=0).values
+        return mask.reshape(1, 1, 1, height, width)
+
+    def tiled_decode(self, z: torch.Tensor, tile_size=(34, 34), tile_stride=(18, 16)):
+        """Decode latent with spatial tiling to reduce peak GPU memory.
+
+        Args:
+            z: Latent tensor [B, C, T, H_latent, W_latent].
+            tile_size: (tile_h, tile_w) in latent spatial units.
+            tile_stride: (stride_h, stride_w) in latent spatial units.
+
+        Returns:
+            DecoderOutput with decoded video [B, 3, T_out, H_out, W_out] clamped to [-1, 1].
+        """
+        upsampling_factor = self.config.spatial_compression_ratio
+        temporal_ratio = self.config.temporal_compression_ratio
+
+        _, _, latent_t, latent_h, latent_w = z.shape
+        size_h, size_w = tile_size
+        stride_h, stride_w = tile_stride
+
+        # Build tile task list (skip redundant trailing tiles)
+        tasks = []
+        for h in range(0, latent_h, stride_h):
+            if h - stride_h >= 0 and h - stride_h + size_h >= latent_h:
+                continue
+            for w in range(0, latent_w, stride_w):
+                if w - stride_w >= 0 and w - stride_w + size_w >= latent_w:
+                    continue
+                tasks.append((h, h + size_h, w, w + size_w))
+
+        data_device = "cpu"
+        computation_device = z.device
+
+        out_t = latent_t * temporal_ratio - (temporal_ratio - 1)
+        out_h = latent_h * upsampling_factor
+        out_w = latent_w * upsampling_factor
+
+        weight = torch.zeros((1, 1, out_t, out_h, out_w), dtype=z.dtype, device=data_device)
+        values = torch.zeros((1, 3, out_t, out_h, out_w), dtype=z.dtype, device=data_device)
+
+        border_h = (size_h - stride_h) * upsampling_factor
+        border_w = (size_w - stride_w) * upsampling_factor
+
+        for h_start, h_end, w_start, w_end in tasks:
+            tile_latent = z[:, :, :, h_start:h_end, w_start:w_end].to(computation_device)
+            tile_decoded = self.model.decode(tile_latent, self.scale).float().clamp_(-1, 1).to(data_device)
+
+            is_bound = (h_start == 0, h_end >= latent_h, w_start == 0, w_end >= latent_w)
+            mask = self._build_spatial_mask(
+                tile_decoded, is_bound=is_bound, border_width=(border_h, border_w)
+            ).to(dtype=z.dtype, device=data_device)
+
+            target_h = h_start * upsampling_factor
+            target_w = w_start * upsampling_factor
+            th = tile_decoded.shape[3]
+            tw = tile_decoded.shape[4]
+            values[:, :, :, target_h:target_h + th, target_w:target_w + tw] += tile_decoded * mask
+            weight[:, :, :, target_h:target_h + th, target_w:target_w + tw] += mask
+
+        decoded = (values / weight).clamp_(-1, 1)
+        return DecoderOutput(sample=decoded.to(z.device))
+
+    def tiled_encode(self, x: torch.Tensor, tile_size=(34, 34), tile_stride=(18, 16)):
+        """Encode video with spatial tiling to reduce peak GPU memory.
+
+        Args:
+            x: Pixel video tensor [B, C, T, H, W].
+            tile_size: (tile_h, tile_w) in latent spatial units.
+            tile_stride: (stride_h, stride_w) in latent spatial units.
+
+        Returns:
+            AutoencoderKLOutput with latent distribution.
+        """
+        upsampling_factor = self.config.spatial_compression_ratio
+        temporal_ratio = self.config.temporal_compression_ratio
+        latent_channels = self.config.latent_channels
+
+        _, _, pixel_t, pixel_h, pixel_w = x.shape
+        size_h = tile_size[0] * upsampling_factor
+        size_w = tile_size[1] * upsampling_factor
+        stride_h = tile_stride[0] * upsampling_factor
+        stride_w = tile_stride[1] * upsampling_factor
+
+        tasks = []
+        for h in range(0, pixel_h, stride_h):
+            if h - stride_h >= 0 and h - stride_h + size_h >= pixel_h:
+                continue
+            for w in range(0, pixel_w, stride_w):
+                if w - stride_w >= 0 and w - stride_w + size_w >= pixel_w:
+                    continue
+                tasks.append((h, h + size_h, w, w + size_w))
+
+        data_device = "cpu"
+        computation_device = x.device
+
+        latent_h = pixel_h // upsampling_factor
+        latent_w = pixel_w // upsampling_factor
+        latent_t = (pixel_t + temporal_ratio - 1) // temporal_ratio
+
+        weight = torch.zeros((1, 1, latent_t, latent_h, latent_w), dtype=x.dtype, device=data_device)
+        values = torch.zeros((1, latent_channels * 2, latent_t, latent_h, latent_w), dtype=x.dtype, device=data_device)
+
+        border_h = tile_size[0] - tile_stride[0]
+        border_w = tile_size[1] - tile_stride[1]
+
+        for h_start, h_end, w_start, w_end in tasks:
+            tile_pixel = x[:, :, :, h_start:h_end, w_start:w_end].to(computation_device)
+            tile_encoded = self.model.encode(tile_pixel, self.scale).float().to(data_device)
+
+            is_bound = (h_start == 0, h_end >= pixel_h, w_start == 0, w_end >= pixel_w)
+            mask = self._build_spatial_mask(
+                tile_encoded, is_bound=is_bound, border_width=(border_h, border_w)
+            ).to(dtype=x.dtype, device=data_device)
+
+            target_h = h_start // upsampling_factor
+            target_w = w_start // upsampling_factor
+            th = tile_encoded.shape[3]
+            tw = tile_encoded.shape[4]
+            values[:, :, :, target_h:target_h + th, target_w:target_w + tw] += tile_encoded * mask
+            weight[:, :, :, target_h:target_h + th, target_w:target_w + tw] += mask
+
+        latent = (values / weight).to(x.device)
+        posterior = DiagonalGaussianDistribution(latent)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_path, additional_kwargs={}):
+        def filter_kwargs(cls, kwargs):
+            import inspect
+            sig = inspect.signature(cls.__init__)
+            valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+            return filtered_kwargs
+
+        model = cls(**filter_kwargs(cls, additional_kwargs))
+        if pretrained_model_path.endswith(".safetensors"):
+            from safetensors.torch import load_file, safe_open
+            state_dict = load_file(pretrained_model_path)
+        else:
+            state_dict = torch.load(pretrained_model_path, map_location="cpu")
+        tmp_state_dict = {} 
+        for key in state_dict:
+            tmp_state_dict["model." + key] = state_dict[key]
+        state_dict = tmp_state_dict
+        m, u = model.load_state_dict(state_dict, strict=False)
+
+        return model