Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ 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/girl.png filter=lfs diff=lfs merge=lfs -text
+examples/news.wav filter=lfs diff=lfs merge=lfs -text
+examples/podcast_sichuan.wav filter=lfs diff=lfs merge=lfs -text
+examples/portrait.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,57 @@
 ---
-title: Leaptalk Talking Head
-emoji: 👀
-colorFrom: green
-colorTo: purple
+title: LeapTalk
+emoji: 🗣️
+colorFrom: pink
+colorTo: yellow
 sdk: gradio
 sdk_version: 6.22.0
-python_version: '3.12'
 app_file: app.py
+short_description: One-step audio-driven talking head video generation
+python_version: "3.12"
+startup_duration_timeout: 45m
 pinned: false
+license: apache-2.0
+models:
+  - z-rx/leaptalk
+  - Soul-AILab/SoulX-FlashHead-1_3B
+  - facebook/wav2vec2-base-960h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# LeapTalk — real-time talking-head generation
+
+Demo for [**z-rx/leaptalk**](https://huggingface.co/z-rx/leaptalk): give it a portrait photo
+and a speech clip, and it generates a lip-synced talking-head video.
+
+LeapTalk reformulates talking-head generation as a *data-to-data* Brownian-bridge transport
+(**Bridge Forcing**), which lets the distilled model synthesize each video chunk with a single
+network evaluation (1 NFE) while keeping identity and style stable over long autoregressive
+rollouts.
+
+- Paper: [LeapTalk: Breaking the Latency-Quality Trade-off in Talking Head Generation](https://huggingface.co/papers/2608.00079)
+- Project page: https://zhangrongxiang.github.io/leaptalk-page/
+- Code: https://github.com/zhangrongxiang/LeapTalk
+
+## Pipeline
+
+This Space reproduces the authors' reference streaming pipeline
+(`inference.py --model_type pro --lite --num_inference_steps 1 --audio_encode_mode stream`):
+
+| component | weights |
+|---|---|
+| base video DiT | `Soul-AILab/SoulX-FlashHead-1_3B` (`Model_Pro`) |
+| LeapTalk LoRA (merged) + audio projector | `z-rx/leaptalk` |
+| audio encoder | `facebook/wav2vec2-base-960h` |
+| Lite TAE video autoencoder | `taew2_1.pth` from `z-rx/leaptalk` |
+| sampler | ViBT Brownian-bridge scheduler (1 step / chunk) |
+
+Video is produced chunk-by-chunk (33 frames per chunk, 5 overlapping motion frames) at
+512x512, 25 fps, with the SoulX-style VAE round-trip history update and reference-colour
+correction.
+
+## Notes
+
+- Runs on ZeroGPU; the multi-GPU USP (xFuser sequence-parallel) path from the reference
+  implementation is disabled since the Space is single-GPU.
+- Example assets come from the [LeapTalk](https://github.com/zhangrongxiang/LeapTalk) and
+  [SoulX-FlashHead](https://github.com/Soul-AILab/SoulX-FlashHead) repositories (Apache-2.0);
+  the audio clips were trimmed to 9 seconds.

app.py

@@ -0,0 +1,515 @@
+"""LeapTalk — real-time audio-driven talking-head generation on ZeroGPU.
+
+Faithful port of the official reference implementation
+(https://github.com/zhangrongxiang/LeapTalk, `inference.py` streaming path):
+
+    SoulX-FlashHead-1_3B (Model_Pro)  +  LeapTalk LoRA (merged)
+    + LeapTalk audio projector        +  wav2vec2-base-960h audio encoder
+    + Lite TAE (taew2_1) VAE          +  ViBT Brownian-bridge scheduler
+
+Everything (chunking, audio windowing, bridge sampling, motion-frame
+round-trip, colour correction) mirrors the authors' `--lite` / `--model_type pro`
+/ `--audio_encode_mode stream` defaults from `inf.sh`.
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: E402  — must precede any torch / CUDA-touching import
+
+import math  # noqa: E402
+import shutil  # noqa: E402
+import subprocess  # noqa: E402
+import sys  # noqa: E402
+import tempfile  # noqa: E402
+import time  # noqa: E402
+import wave  # noqa: E402
+from collections import deque  # noqa: E402
+
+import gradio as gr  # noqa: E402
+import imageio  # noqa: E402
+import librosa  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+from huggingface_hub import hf_hub_download, snapshot_download  # noqa: E402
+from loguru import logger  # noqa: E402
+from peft import PeftModel  # noqa: E402
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import flash_head.src.pipeline.flash_head_pipeline as fh_pipe_mod  # noqa: E402
+
+# torch.compile is disabled: the released LoRA was not saved from a compiled base
+# (`--compile off` in the reference `inf.sh`), and TAEHV uses Python-level loops.
+fh_pipe_mod.COMPILE_MODEL = False
+fh_pipe_mod.COMPILE_VAE = False
+
+from flash_head.src.pipeline.flash_head_pipeline import FlashHeadPipeline  # noqa: E402
+from leaptalk_inference import (  # noqa: E402
+    StreamParams,
+    _audio_context_from_embeddings_range,
+    _bridge_sample_one_chunk,
+    _build_infer_timesteps,
+    _decode_to_cthw,
+    _encode_motion_prefix_from_decoded,
+    _get_inner_flashhead_model,
+    _maybe_apply_color_correction,
+)
+from vibt.scheduler import ViBTScheduler  # noqa: E402
+
+# --------------------------------------------------------------------------------------
+# Fixed inference configuration (reference defaults)
+# --------------------------------------------------------------------------------------
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+HEIGHT = WIDTH = 512
+FPS = 25
+SAMPLE_RATE = 16000
+FRAME_NUM = 33
+MOTION_FRAMES_LATENT_NUM = 2
+CACHED_AUDIO_DURATION = 8
+SHIFT_GAMMA = 5.0
+NOISE_SCALE = 1.0
+COLOR_CORRECTION_STRENGTH = 1.0
+MAX_SECONDS_CAP = 20
+
+# --------------------------------------------------------------------------------------
+# Weights
+# --------------------------------------------------------------------------------------
+logger.info("Downloading weights…")
+CKPT_DIR = snapshot_download(
+    "Soul-AILab/SoulX-FlashHead-1_3B", allow_patterns=["Model_Pro/*"]
+)
+WAV2VEC_DIR = snapshot_download(
+    "facebook/wav2vec2-base-960h",
+    allow_patterns=["*.json", "*.txt", "*.safetensors", "pytorch_model.bin"],
+)
+LEAPTALK_DIR = snapshot_download("z-rx/leaptalk")
+LORA_DIR = os.path.join(LEAPTALK_DIR, "lora")
+TAE_PATH = os.path.join(LEAPTALK_DIR, "taew2_1.pth")
+AUDIO_PROJ_PATH = os.path.join(LEAPTALK_DIR, "audio_proj_step_10400.pt")
+
+# --------------------------------------------------------------------------------------
+# Pipeline assembly (module scope, eagerly moved to CUDA)
+# --------------------------------------------------------------------------------------
+logger.info("Building FlashHead pipeline…")
+# Built on CPU first so the LoRA merge / projector load happen on real tensors,
+# then the whole stack is moved to CUDA eagerly (ZeroGPU packs it from there).
+pipeline = FlashHeadPipeline(
+    checkpoint_dir=CKPT_DIR,
+    model_type="pro",
+    wav2vec_dir=WAV2VEC_DIR,
+    device="cpu",
+    param_dtype=DTYPE,
+    use_usp=False,
+    use_tae=True,
+    tae_path=TAE_PATH,
+    tae_model_type="wan21",
+)
+
+logger.info("Merging LeapTalk LoRA…")
+pipeline.model = PeftModel.from_pretrained(pipeline.model, LORA_DIR, is_trainable=False)
+pipeline.model = pipeline.model.merge_and_unload()
+pipeline.model.eval().requires_grad_(False)
+
+logger.info("Loading LeapTalk audio projector…")
+_audio_proj_state = torch.load(AUDIO_PROJ_PATH, map_location="cpu", weights_only=True)
+_get_inner_flashhead_model(pipeline.model).audio_proj.load_state_dict(
+    _audio_proj_state, strict=True
+)
+del _audio_proj_state
+
+pipeline.device = DEVICE
+pipeline.model.to(DEVICE)
+pipeline.vae.device = DEVICE
+pipeline.vae.model.to(DEVICE)
+pipeline.audio_encoder.to(DEVICE)
+pipeline.audio_encoder.eval().requires_grad_(False)
+
+STREAM = StreamParams(
+    frame_num=FRAME_NUM,
+    motion_frames_latent_num=MOTION_FRAMES_LATENT_NUM,
+    tgt_fps=FPS,
+    sample_rate=SAMPLE_RATE,
+    cached_audio_duration=CACHED_AUDIO_DURATION,
+).init_with_stride(int(pipeline.config.vae_stride[0]))
+SLICE_SAMPLES = STREAM.slice_len * SAMPLE_RATE // FPS
+logger.info(
+    f"Ready. frame_num={STREAM.frame_num} motion_frames={STREAM.motion_frames_num} "
+    f"slice_len={STREAM.slice_len} ({SLICE_SAMPLES} samples/chunk)"
+)
+
+
+# --------------------------------------------------------------------------------------
+# Video helpers
+# --------------------------------------------------------------------------------------
+def _ffmpeg_exe() -> str:
+    exe = shutil.which("ffmpeg")
+    if exe:
+        return exe
+    import imageio_ffmpeg
+
+    return imageio_ffmpeg.get_ffmpeg_exe()
+
+
+def _write_wav(path: str, audio: np.ndarray, sample_rate: int = SAMPLE_RATE) -> str:
+    pcm = (np.clip(audio, -1.0, 1.0) * 32767.0).astype(np.int16)
+    with wave.open(path, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        wf.writeframes(pcm.tobytes())
+    return path
+
+
+def _mux(video_path: str, audio_path: str, out_path: str) -> str:
+    cmd = [
+        _ffmpeg_exe(), "-y",
+        "-i", video_path,
+        "-i", audio_path,
+        "-c:v", "copy",
+        "-c:a", "aac", "-b:a", "128k",
+        "-shortest",
+        "-movflags", "+faststart",
+        out_path,
+    ]
+    proc = subprocess.run(cmd, capture_output=True)
+    if proc.returncode != 0 or not os.path.exists(out_path):
+        logger.warning(f"ffmpeg mux failed: {proc.stderr.decode()[-800:]}")
+        shutil.copy(video_path, out_path)
+    return out_path
+
+
+def _num_chunks_for(seconds: float) -> int:
+    samples = max(int(seconds * SAMPLE_RATE), FRAME_NUM * SAMPLE_RATE // FPS)
+    return max(1, math.ceil(samples / SLICE_SAMPLES))
+
+
+def _estimate_duration(
+    portrait_image=None,
+    speech_audio=None,
+    max_seconds: float = 9.0,
+    num_inference_steps: int = 1,
+    guidance_scale: float = 1.0,
+    *args,
+    **kwargs,
+) -> int:
+    """ZeroGPU time budget: weight streaming + per-chunk cost."""
+    try:
+        chunks = _num_chunks_for(float(max_seconds))
+        nfe = max(1, int(num_inference_steps)) * (2 if float(guidance_scale) != 1.0 else 1)
+    except Exception:
+        chunks, nfe = _num_chunks_for(MAX_SECONDS_CAP), 1
+    return int(min(180, 25 + chunks * (0.9 + 0.55 * nfe)))
+
+
+# --------------------------------------------------------------------------------------
+# Inference
+# --------------------------------------------------------------------------------------
+@spaces.GPU(duration=_estimate_duration)
+def generate(
+    portrait_image: str,
+    speech_audio: str,
+    max_seconds: float = 9.0,
+    num_inference_steps: int = 1,
+    guidance_scale: float = 1.0,
+    seed: int = 42,
+    auto_crop_face: bool = True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Animate a portrait photo so that it speaks the given audio.
+
+    Args:
+        portrait_image: Path to a portrait photo (a single, roughly front-facing face).
+        speech_audio: Path to a speech audio file that drives lip and head motion.
+        max_seconds: Maximum number of seconds of the audio to animate.
+        num_inference_steps: Bridge-sampler steps per chunk. LeapTalk is distilled for 1.
+        guidance_scale: Audio classifier-free guidance. 1.0 disables it (2x faster).
+        seed: Random seed for the Brownian-bridge noise.
+        auto_crop_face: Detect and crop around the face before generating.
+
+    Returns:
+        A tuple of (path to the generated talking-head mp4, a short speed report).
+    """
+    if not portrait_image:
+        raise gr.Error("Please provide a portrait image.")
+    if not speech_audio:
+        raise gr.Error("Please provide a speech audio file.")
+
+    num_inference_steps = max(1, int(num_inference_steps))
+    guidance_scale = float(guidance_scale)
+    seed = int(seed)
+    max_seconds = float(np.clip(max_seconds, 1.0, MAX_SECONDS_CAP))
+
+    workdir = tempfile.mkdtemp(prefix="leaptalk_")
+    progress(0.02, desc="Preparing reference portrait…")
+
+    # ---- reference image -> anchor latent X0 (same call as inference.py) --------------
+    pipeline.prepare_params(
+        cond_image_path_or_dir=portrait_image,
+        target_size=(HEIGHT, WIDTH),
+        frame_num=STREAM.frame_num,
+        motion_frames_num=0,
+        sampling_steps=num_inference_steps,
+        seed=seed,
+        shift=SHIFT_GAMMA,
+        color_correction_strength=COLOR_CORRECTION_STRENGTH,
+        use_face_crop=bool(auto_crop_face),
+    )
+    X0 = pipeline.ref_img_latent.to(device=DEVICE, dtype=DTYPE)
+
+    # ---- scheduler -------------------------------------------------------------------
+    scheduler = ViBTScheduler(num_train_timesteps=1000)
+    scheduler.timesteps = _build_infer_timesteps(
+        step_list=None,
+        num_inference_steps=num_inference_steps,
+        shift_gamma=SHIFT_GAMMA,
+        device=DEVICE,
+        num_timesteps=1000,
+    )
+    scheduler.num_inference_steps = int(scheduler.timesteps.numel())
+    scheduler.set_parameters(noise_scale=NOISE_SCALE, shift_gamma=SHIFT_GAMMA, seed=seed)
+
+    # ---- audio (streaming ring buffer, exactly as inference.py --audio_encode_mode stream)
+    progress(0.06, desc="Loading audio…")
+    audio_all, _ = librosa.load(speech_audio, sr=SAMPLE_RATE, mono=True)
+    audio_all = audio_all[: int(max_seconds * SAMPLE_RATE)]
+    if audio_all.size == 0:
+        raise gr.Error("The audio file appears to be empty.")
+
+    frame_window_samples = STREAM.frame_num * SAMPLE_RATE // FPS
+    remainder = len(audio_all) % SLICE_SAMPLES
+    if remainder > 0:
+        audio_all = np.concatenate(
+            [audio_all, np.zeros(SLICE_SAMPLES - remainder, dtype=audio_all.dtype)]
+        )
+    if len(audio_all) < frame_window_samples:
+        audio_all = np.concatenate(
+            [audio_all, np.zeros(frame_window_samples - len(audio_all), dtype=audio_all.dtype)]
+        )
+    remainder = len(audio_all) % SLICE_SAMPLES
+    if remainder != 0:
+        audio_all = np.concatenate(
+            [audio_all, np.zeros(SLICE_SAMPLES - remainder, dtype=audio_all.dtype)]
+        )
+
+    slices = audio_all.reshape(-1, SLICE_SAMPLES)
+    num_chunks = int(slices.shape[0])
+
+    cached_len = SAMPLE_RATE * STREAM.cached_audio_duration
+    audio_end_idx = STREAM.cached_audio_duration * FPS
+    audio_start_idx = audio_end_idx - STREAM.frame_num
+    audio_dq = deque([0.0] * cached_len, maxlen=cached_len)
+
+    latent_motion_frames = X0[:, :1].unsqueeze(0).clone()
+    clamp_latent_len = int(latent_motion_frames.shape[2])
+
+    generated: list[np.ndarray] = []
+    gen_seconds = 0.0
+    gen_frames = 0
+
+    for chunk_idx in range(num_chunks):
+        progress(
+            0.08 + 0.88 * chunk_idx / num_chunks,
+            desc=f"Generating chunk {chunk_idx + 1}/{num_chunks}…",
+        )
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+
+        audio_dq.extend(slices[chunk_idx].tolist())
+        audio_cache = np.array(audio_dq, dtype=np.float32)
+        audio_emb = pipeline.preprocess_audio(audio_cache, sr=SAMPLE_RATE, fps=FPS)
+        if audio_emb is None:
+            raise gr.Error("Failed to extract audio embeddings.")
+        audio_emb = audio_emb.to(device=DEVICE, dtype=DTYPE)
+        audio_ctx = _audio_context_from_embeddings_range(
+            audio_emb,
+            start_idx=audio_start_idx,
+            end_idx=audio_end_idx,
+            device=DEVICE,
+            dtype=DTYPE,
+        )
+
+        x_final = _bridge_sample_one_chunk(
+            pipeline,
+            scheduler=scheduler,
+            ref_latent=X0,
+            audio_context=audio_ctx,
+            guidance_scale=guidance_scale,
+            latent_motion_frames=latent_motion_frames,
+            clamp_latent_len=clamp_latent_len,
+            device=DEVICE,
+            dtype=DTYPE,
+        )
+        decoded_cthw = _decode_to_cthw(pipeline, x_final)
+        decoded_cthw = _maybe_apply_color_correction(pipeline, decoded_cthw)
+
+        # SoulX-style VAE round-trip history update (reference default)
+        latent_motion_frames = _encode_motion_prefix_from_decoded(
+            pipeline,
+            decoded_video_cthw=decoded_cthw,
+            motion_frames_num=STREAM.motion_frames_num,
+            device=DEVICE,
+            dtype=DTYPE,
+        ).unsqueeze(0)
+        clamp_latent_len = int(latent_motion_frames.shape[2])
+
+        decoded_cthw = decoded_cthw[:, STREAM.motion_frames_num:]
+        video_thwc = (
+            ((decoded_cthw + 1.0) / 2.0)
+            .permute(1, 2, 3, 0)
+            .clamp(0.0, 1.0)
+            .mul(255.0)
+            .contiguous()
+        )
+        torch.cuda.synchronize()
+        chunk_seconds = time.perf_counter() - t0
+
+        frames_np = video_thwc.to(torch.float32).cpu().numpy().astype(np.uint8)
+        generated.append(frames_np)
+        gen_frames += int(frames_np.shape[0])
+        gen_seconds += chunk_seconds
+        logger.info(
+            f"chunk {chunk_idx + 1}/{num_chunks}: {chunk_seconds:.3f}s "
+            f"({frames_np.shape[0] / max(chunk_seconds, 1e-6):.1f} FPS)"
+        )
+
+    progress(0.97, desc="Encoding video…")
+    silent_path = os.path.join(workdir, "silent.mp4")
+    with imageio.get_writer(
+        silent_path,
+        format="mp4",
+        mode="I",
+        fps=FPS,
+        codec="h264",
+        pixelformat="yuv420p",
+        ffmpeg_params=["-bf", "0"],
+    ) as writer:
+        for frames_np in generated:
+            for frame in frames_np:
+                writer.append_data(frame)
+
+    wav_path = _write_wav(os.path.join(workdir, "track.wav"), audio_all)
+    out_path = _mux(silent_path, wav_path, os.path.join(workdir, "leaptalk.mp4"))
+
+    video_seconds = gen_frames / FPS
+    report = (
+        f"**{gen_frames} frames** ({video_seconds:.1f}s of video) in "
+        f"**{gen_seconds:.2f}s** of GPU time — "
+        f"**{gen_frames / max(gen_seconds, 1e-6):.1f} FPS** generation throughput "
+        f"({gen_frames / max(gen_seconds, 1e-6) / FPS:.2f}× real time) over "
+        f"{num_chunks} streaming chunks at {num_inference_steps} step"
+        f"{'s' if num_inference_steps > 1 else ''}/chunk."
+    )
+    return out_path, report
+
+
+# --------------------------------------------------------------------------------------
+# UI
+# --------------------------------------------------------------------------------------
+CSS = """
+#col-container { margin: 0 auto; max-width: 1180px; }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="LeapTalk") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+# 🗣️ LeapTalk — real-time talking heads
+
+Animate a **portrait photo** with a **speech clip**. LeapTalk reformulates talking-head
+generation as a Brownian-bridge transport (*Bridge Forcing*), which lets it synthesize each
+video chunk in a **single sampling step** while keeping identity stable over long rollouts.
+
+[Model](https://huggingface.co/z-rx/leaptalk) · [Paper](https://huggingface.co/papers/2608.00079)
+· [Project page](https://zhangrongxiang.github.io/leaptalk-page/)
+· [Code](https://github.com/zhangrongxiang/LeapTalk)
+· built on [SoulX-FlashHead-1.3B](https://huggingface.co/Soul-AILab/SoulX-FlashHead-1_3B)
+            """
+        )
+
+        with gr.Row():
+            with gr.Column():
+                portrait_image = gr.Image(
+                    label="Portrait", type="filepath", height=320, sources=["upload", "webcam", "clipboard"]
+                )
+                speech_audio = gr.Audio(
+                    label="Speech audio", type="filepath", sources=["upload", "microphone"]
+                )
+                run_btn = gr.Button("Generate talking head", variant="primary")
+            with gr.Column():
+                video_out = gr.Video(
+                    label="Result", height=460, autoplay=True, show_download_button=True
+                )
+                report_out = gr.Markdown()
+
+        with gr.Accordion("Advanced options", open=False):
+            with gr.Row():
+                max_seconds = gr.Slider(
+                    label="Max audio length (seconds)",
+                    minimum=1,
+                    maximum=MAX_SECONDS_CAP,
+                    step=1,
+                    value=9,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Sampling steps per chunk",
+                    minimum=1,
+                    maximum=4,
+                    step=1,
+                    value=1,
+                    info="LeapTalk is distilled for 1-step (1 NFE) generation.",
+                )
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Audio guidance scale",
+                    minimum=1.0,
+                    maximum=3.0,
+                    step=0.1,
+                    value=1.0,
+                    info="1.0 disables audio CFG; higher strengthens lip motion but doubles compute.",
+                )
+                seed = gr.Number(label="Seed", value=42, precision=0)
+            auto_crop_face = gr.Checkbox(
+                label="Auto-crop to face",
+                value=True,
+                info="Detects the face and crops around it; falls back to a centre crop.",
+            )
+
+        gr.Examples(
+            examples=[
+                ["examples/portrait.jpg", "examples/news.wav"],
+                ["examples/girl.png", "examples/podcast_sichuan.wav"],
+            ],
+            inputs=[portrait_image, speech_audio],
+            outputs=[video_out, report_out],
+            fn=generate,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+        gr.Markdown(
+            "Example assets: portrait + news clip from the "
+            "[LeapTalk](https://github.com/zhangrongxiang/LeapTalk) repository, portrait + podcast "
+            "clip from [SoulX-FlashHead](https://github.com/Soul-AILab/SoulX-FlashHead) "
+            "(both Apache-2.0). Audio clips were trimmed to 9 seconds."
+        )
+
+    gr.on(
+        triggers=[run_btn.click],
+        fn=generate,
+        inputs=[
+            portrait_image,
+            speech_audio,
+            max_seconds,
+            num_inference_steps,
+            guidance_scale,
+            seed,
+            auto_crop_face,
+        ],
+        outputs=[video_out, report_out],
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

examples/news.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08e71b2901fb19c33d02dc4c96a118d432f476a155ac849bcbf121602fdb3829
+size 1587644

examples/podcast_sichuan.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2dc8315e8f2a6a616baa467706769ca3a563696d280829fad68310357db60614
+size 288044

flash_head/audio_analysis/torch_utils.py

@@ -0,0 +1,20 @@
+import torch
+import torch.nn.functional as F
+
+
+def get_mask_from_lengths(lengths, max_len=None):
+    lengths = lengths.to(torch.long)
+    if max_len is None:
+        max_len = torch.max(lengths).item()
+
+    ids = torch.arange(0, max_len).unsqueeze(0).expand(lengths.shape[0], -1).to(lengths.device)
+    mask = ids < lengths.unsqueeze(1).expand(-1, max_len)
+
+    return mask
+
+
+def linear_interpolation(features, seq_len):
+    features = features.transpose(1, 2)
+    output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear')
+    return output_features.transpose(1, 2)
+

flash_head/audio_analysis/wav2vec2.py

@@ -0,0 +1,125 @@
+from transformers import Wav2Vec2Config, Wav2Vec2Model
+from transformers.modeling_outputs import BaseModelOutput
+
+from .torch_utils import linear_interpolation
+
+# the implementation of Wav2Vec2Model is borrowed from
+# https://github.com/huggingface/transformers/blob/HEAD/src/transformers/models/wav2vec2/modeling_wav2vec2.py
+# initialize our encoder with the pre-trained wav2vec 2.0 weights.
+class Wav2Vec2Model(Wav2Vec2Model):
+    def __init__(self, config: Wav2Vec2Config):
+        super().__init__(config)
+
+    def forward(
+        self,
+        input_values,
+        seq_len,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        self.config.output_attentions = False
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+    def feature_extract(
+        self,
+        input_values,
+        seq_len,
+    ):
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        return extract_features
+
+    def encode(
+        self,
+        extract_features,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        self.config.output_attentions = False
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+            
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

flash_head/src/modules/flash_head_model.py

@@ -0,0 +1,589 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Tuple, Optional
+from einops import rearrange
+from diffusers import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+import torch.distributed as dist
+
+# `torch.cuda.amp` is deprecated; fall back to the equivalent `torch.amp` API on
+# newer torch releases where the alias has been removed.
+try:
+    import torch.cuda.amp as amp
+except (ImportError, AttributeError):  # pragma: no cover
+    from torch import amp as _torch_amp
+
+    class amp:  # noqa: N801 - drop-in shim for torch.cuda.amp
+        @staticmethod
+        def autocast(*args, **kwargs):
+            return _torch_amp.autocast("cuda", *args, **kwargs)
+
+# xFuser is only needed for the multi-GPU USP (sequence-parallel) path. This Space
+# always runs single-GPU, so the import is made optional to keep the environment slim.
+try:
+    from xfuser.core.distributed import (
+        get_sequence_parallel_rank,
+        get_sequence_parallel_world_size,
+        get_sp_group,
+    )
+    from xfuser.core.long_ctx_attention import xFuserLongContextAttention
+except ImportError:  # single-GPU: USP is never enabled (see _maybe_enable_usp below)
+    get_sequence_parallel_rank = None
+    get_sequence_parallel_world_size = None
+    get_sp_group = None
+    xFuserLongContextAttention = None
+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:
+    from sageattention import sageattn
+    SAGE_ATTN_AVAILABLE = True
+except ModuleNotFoundError:
+    SAGE_ATTN_AVAILABLE = False
+
+
+def _maybe_enable_usp() -> tuple[bool, int, int]:
+    """Return (use_usp, sp_size, sp_rank).
+
+    The model code supports xFuser's USP/sequence-parallel path, but this
+    requires xFuser model-parallel groups to be initialized.
+
+    In plain torchrun DDP training we typically only initialize the default
+    torch.distributed process group; in that case xFuser groups are absent.
+    To keep DDP usable, we gracefully fall back to non-USP behavior.
+    """
+
+    if not dist.is_initialized():
+        return False, 1, 0
+
+    try:
+        # These call into xFuser parallel_state and will assert if the group
+        # is not initialized.
+        sp_size = int(get_sequence_parallel_world_size())
+        sp_rank = int(get_sequence_parallel_rank())
+        if sp_size <= 1:
+            return False, 1, 0
+        return True, sp_size, sp_rank
+    except Exception:
+        return False, 1, 0
+    
+    
+def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False):
+    if compatibility_mode:
+        q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)
+    elif SAGE_ATTN_AVAILABLE:
+        q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
+        x = sageattn(q, k, v)
+        x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)
+    elif FLASH_ATTN_3_AVAILABLE:
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+        x = flash_attn_interface.flash_attn_func(q, k, v)
+        x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
+    elif FLASH_ATTN_2_AVAILABLE:
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+        x = flash_attn.flash_attn_func(q, k, v)
+        x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
+    else:
+        q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)
+    return x
+
+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_3d(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 3d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta)
+    h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    return torch.cat([f_freqs_cis, h_freqs_cis, w_freqs_cis], dim=1)
+
+
+def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 1d rope precompute
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)
+                   [: (dim // 2)].double() / dim))
+    freqs = torch.outer(torch.arange(end, device=freqs.device), freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+def pad_freqs(original_tensor, target_len):
+    seq_len, s1, s2 = original_tensor.shape
+    pad_size = target_len - seq_len
+    padding_tensor = torch.ones(
+        pad_size,
+        s1,
+        s2,
+        dtype=original_tensor.dtype,
+        device=original_tensor.device)
+    padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
+    return padded_tensor
+
+def rope_apply(x, freqs, grid_sizes, use_usp=False, sp_size=1, sp_rank=0):
+    """
+    x:          [B, L, N, C].
+    grid_sizes: [B, 3].
+    freqs:      [M, C // 2].
+    """
+    s, n, c = x.size(1), x.size(2), x.size(3) // 2
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) # [[N, head_dim/2], [N, head_dim/2], [N, head_dim/2]] # T H W 极坐标
+
+    # loop over samples
+
+    (f, h, w) = grid_sizes
+    seq_len = f * h * w
+
+    # precompute multipliers
+    x_i = torch.view_as_complex(x[0, :s].to(torch.float64).reshape(
+        s, n, -1, 2)) # [L, N, C/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) # seq_lens, 1,  3 * dim / 2 (T H W)
+
+    if use_usp:
+        # apply rotary embedding
+        freqs_i = pad_freqs(freqs_i, s * sp_size)
+        s_per_rank = s
+        freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *
+                                                        s_per_rank), :, :]
+        x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
+        x_i = torch.cat([x_i, x[0, s:]])
+    else:
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[0, seq_len:]])
+    return x_i.unsqueeze(0).to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        dtype = x.dtype
+        return self.norm(x.float()).to(dtype) * self.weight
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        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)
+
+        self.use_usp, self.sp_size, self.sp_rank = _maybe_enable_usp()
+
+    def forward(self, x, freqs, grid_sizes):
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x)
+
+        if self.use_usp:
+            from yunchang.kernels import AttnType
+            if SAGE_ATTN_AVAILABLE:
+                attn_type = AttnType.SAGE_AUTO
+            else:
+                attn_type = AttnType.FA
+
+            x = xFuserLongContextAttention(attn_type=attn_type)(
+                None,
+                query=rope_apply(q, freqs, grid_sizes, self.use_usp, self.sp_size, self.sp_rank),
+                key=rope_apply(k, freqs, grid_sizes, self.use_usp, self.sp_size, self.sp_rank),
+                value=v.view(b, s, n, d),
+            ).flatten(2)
+        else:
+            x = flash_attention(
+                q=rope_apply(q, freqs, grid_sizes).flatten(2),
+                k=rope_apply(k, freqs, grid_sizes).flatten(2),
+                v=v,
+                num_heads=self.num_heads
+            )
+        return self.o(x)
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False):
+        super().__init__()
+        self.dim = dim
+        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)
+        self.has_image_input = has_image_input
+        if has_image_input:
+            self.k_img = nn.Linear(dim, dim)
+            self.v_img = nn.Linear(dim, dim)
+            self.norm_k_img = RMSNorm(dim, eps=eps)
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        if self.has_image_input:
+            img = y[:, :257]
+            ctx = y[:, 257:]
+        else:
+            ctx = y
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(ctx))
+        v = self.v(ctx)
+        x = flash_attention(q, k, v, num_heads=self.num_heads)
+        if self.has_image_input:
+            k_img = self.norm_k_img(self.k_img(img))
+            v_img = self.v_img(img)
+            y = flash_attention(q, k_img, v_img, num_heads=self.num_heads)
+            x = x + y
+        return self.o(x)
+
+class DiTAudioBlock(nn.Module):
+    def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6, i=0, num_layers=0):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.ffn_dim = ffn_dim
+        self.i = i
+        self.num_layers = num_layers
+
+        self.self_attn = SelfAttention(dim, num_heads, eps)
+        self.cross_attn = CrossAttention(
+            dim, num_heads, eps, has_image_input=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.use_usp, self.sp_size, self.sp_rank = _maybe_enable_usp()
+
+    def forward(self, x, context, t_mod, freqs, grid_sizes):
+        e = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1)
+
+        y = self.self_attn(
+            self.norm1(x) * (1 + e[1]) + e[0], freqs, grid_sizes)
+
+        x = x + y * e[2]
+
+        x_1 = rearrange(self.norm3(x), 'b (f l) c -> (b f) l c', f=context.shape[1])
+        context_1 = context.squeeze(0)
+
+        if self.use_usp:
+            context_1 = context_1.unsqueeze(1).repeat(1, self.sp_size, 1, 1).flatten(0,1)
+            context_1 = torch.chunk(context_1, self.sp_size, dim=0)[self.sp_rank]
+
+        x = x + self.cross_attn(x_1, context_1).flatten(0, 1).unsqueeze(0)
+
+        y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3])
+        x = x + y * e[5]
+
+        return x
+
+class MLP(torch.nn.Module):
+    def __init__(self, in_dim, out_dim):
+        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)
+        )
+
+    def forward(self, x):
+        return self.proj(x)
+
+
+class Head(nn.Module):
+    def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], 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):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            t_mod(Tensor): Shape [B*21, C]
+        """
+        B, L, D = x.shape
+        F = t_mod.shape[0] // B
+        shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device).unsqueeze(1) + t_mod.unflatten(dim=0, sizes=(B, t_mod.shape[0]//B)).unsqueeze(2)).chunk(2, dim=2)
+
+        x = rearrange(x, 'b (f l) d -> b f l d', f=F)
+        x = (self.head(self.norm(x) * (1 + scale) + shift))
+        x = rearrange(x, 'b f l d -> b (f l) d')
+        return x
+
+class WanModelAudioProject(ModelMixin, ConfigMixin):
+    _no_split_modules = ['DiTAudioBlock']
+    @register_to_config
+    def __init__(
+        self,
+        dim: int,
+        in_dim: int,
+        ffn_dim: int,
+        out_dim: int,
+        text_dim: int,
+        freq_dim: int,
+        eps: float,
+        vae_stride: Tuple[int, int, int],
+        patch_size: Tuple[int, int, int],
+        num_heads: int,
+        num_layers: int,
+        has_image_input: bool,
+        **kwargs,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.freq_dim = freq_dim
+        self.has_image_input = has_image_input
+        self.patch_size = patch_size
+
+        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))
+        self.blocks = nn.ModuleList([
+            DiTAudioBlock(has_image_input, dim, num_heads, ffn_dim, eps, i, num_layers)
+            for i in range(num_layers)
+        ])
+        self.head = Head(dim, out_dim, patch_size, eps)
+        head_dim = dim // num_heads
+        self.freqs = precompute_freqs_cis_3d(head_dim)
+
+        self.audio_emb = MLP(768, dim)
+
+        if has_image_input:
+            self.img_emb = MLP(1280, dim)
+
+        # init audio adapter
+        audio_window = 5
+        vae_scale = vae_stride[0]
+        intermediate_dim = 512
+        output_dim = 1536
+        context_tokens = 32
+        norm_output_audio = True
+        self.audio_window = audio_window
+        self.vae_scale = vae_scale
+        self.audio_proj = AudioProjModel(
+                    seq_len=audio_window,
+                    seq_len_vf=audio_window+vae_scale-1,
+                    intermediate_dim=intermediate_dim,
+                    output_dim=output_dim,
+                    context_tokens=context_tokens,
+                    norm_output_audio=norm_output_audio,
+                )
+
+        self.use_usp, self.sp_size, self.sp_rank = _maybe_enable_usp()
+        
+    def patchify(self, x: torch.Tensor):
+        x = self.patch_embedding(x)
+        grid_size = x.shape[2:]
+        x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous()
+        return x, grid_size  # x, grid_size: (f, h, w)
+
+    def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor):
+        return rearrange(
+            x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)',
+            f=grid_size[0], h=grid_size[1], w=grid_size[2], 
+            x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2]
+        )
+
+    def forward(self,
+                x: torch.Tensor,  #(1, 16, 9, 64, 64))
+                timestep: torch.Tensor, #(9,)
+                context: torch.Tensor, #(5, 33, 12, 768)
+                y: Optional[torch.Tensor] = None, #(1, 16, 9, 64, 64)
+                use_gradient_checkpointing: bool = False,
+                use_gradient_checkpointing_offload: bool = False,
+                **kwargs,
+                ):
+
+        if self.freqs.device != x.device:
+            self.freqs = self.freqs.to(x.device)
+
+        x = torch.cat([x, y], dim=1) # (1, 32, 9, 64, 64)
+        x, grid_sizes = self.patchify(x)
+        t = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, timestep.to(dtype=x.dtype)))
+        t_mod = self.time_projection(t).unflatten(1, (6, self.dim))   # (bsz, 6, 1536) 
+
+        # ==================== 音频条件处理 ====================
+        # 输入: context (bsz, 81, 5, 12, 768)
+        # - 81 帧 = 1 (第一帧) + 80 (后续帧, 每4帧对应VAE压缩后的1帧)
+        # - 5 是音频窗口大小 (audio_window)
+        # - 12 是音频特征的 blocks
+        # - 768 是音频特征维度
+        
+        audio_cond = context.to(device=x.device, dtype=x.dtype)
+        
+        # 1. 第一帧：直接使用完整的5帧音频窗口
+        first_frame_audio = audio_cond[:, :1, ...]  # (bsz, 1, 5, 12, 768)
+        
+        # 2. 后续帧：需要根据帧位置选择不同的音频窗口
+        # 将 32 帧重排为 (8 个 VAE latent, 每个4帧)
+        latter_frames_audio = rearrange(
+            audio_cond[:, 1:, ...], 
+            "b (n_latent n_frame) w s c -> b n_latent n_frame w s c", 
+            n_frame=self.vae_scale  # vae_scale=4
+        )  # (bsz, 8, 4, 5, 12, 768)
+        
+        mid_idx = self.audio_window // 2  # 窗口中心索引: 5//2=2
+        
+        # 为每个 latent 的4帧选择合适的音频窗口:
+        # - 第1帧 (帧索引0): 无过去，取前3帧窗口 [:mid_idx+1] = [:3]
+        # - 中间帧 (帧索引1-2): 取中心1帧 [mid_idx:mid_idx+1] = [2:3]
+        # - 第4帧 (帧索引3): 无未来，取后3帧窗口 [mid_idx:] = [2:]
+        
+        first_of_group = latter_frames_audio[:, :, :1, :mid_idx+1, ...]  # (bsz, 8, 1, 3, 12, 768)
+        middle_of_group = latter_frames_audio[:, :, 1:-1, mid_idx:mid_idx+1, ...]  # (bsz, 8, 2, 1, 12, 768)
+        last_of_group = latter_frames_audio[:, :, -1:, mid_idx:, ...]  # (bsz, 8, 1, 3, 12, 768)
+        
+        # 合并并展平窗口维度: (n_frame, window) -> (n_frame * window)
+        latter_frames_audio_processed = torch.cat([
+            rearrange(first_of_group, "b n_latent n_f w s c -> b n_latent (n_f w) s c"),
+            rearrange(middle_of_group, "b n_latent n_f w s c -> b n_latent (n_f w) s c"),
+            rearrange(last_of_group, "b n_latent n_f w s c -> b n_latent (n_f w) s c"),
+        ], dim=2)  # (bsz, 8, 1*3 + 2*1 + 1*3, 12, 768) = (bsz, 8, 8, 12, 768)
+        
+        # 3. 通过 AudioProjModel 投影到 DiT 所需的特征空间
+        context = self.audio_proj(
+            first_frame_audio, 
+            latter_frames_audio_processed
+        ).to(x.dtype)  # (bsz, 9, 32, 1536)
+
+        if self.use_usp:
+            x = torch.chunk(x, self.sp_size, dim=1)[self.sp_rank]
+
+        for block in self.blocks:
+            x = block(x, context, t_mod, self.freqs, grid_sizes)
+        x = self.head(x, t)   # (bsz, 9*32*32, 64)
+        if self.use_usp:
+            x = get_sp_group().all_gather(x, dim=1)
+        x = self.unpatchify(x, grid_sizes)  # (bsz, 16, 21, 64, 64)
+        return x
+
+
+class AudioProjModel(ModelMixin, ConfigMixin):
+    def __init__(
+        self,
+        seq_len=5,
+        seq_len_vf=12,
+        blocks=12,  
+        channels=768, 
+        intermediate_dim=512,
+        output_dim=768,
+        context_tokens=32,
+        norm_output_audio=False,
+    ):
+        super().__init__()
+
+        self.seq_len = seq_len
+        self.blocks = blocks
+        self.channels = channels
+        self.input_dim = seq_len * blocks * channels  
+        self.input_dim_vf = seq_len_vf * blocks * channels
+        self.intermediate_dim = intermediate_dim
+        self.context_tokens = context_tokens
+        self.output_dim = output_dim
+
+        # define multiple linear layers
+        self.proj1 = nn.Linear(self.input_dim, intermediate_dim)
+        self.proj1_vf = nn.Linear(self.input_dim_vf, intermediate_dim)
+        self.proj2 = nn.Linear(intermediate_dim, intermediate_dim)
+        self.proj3 = nn.Linear(intermediate_dim, context_tokens * output_dim)
+        self.norm = nn.LayerNorm(output_dim) if norm_output_audio else nn.Identity()
+
+    def forward(self, audio_embeds, audio_embeds_vf):
+        video_length = audio_embeds.shape[1] + audio_embeds_vf.shape[1]
+        B, _, _, S, C = audio_embeds.shape
+
+        # process audio of first frame
+        audio_embeds = rearrange(audio_embeds, "bz f w b c -> (bz f) w b c")
+        batch_size, window_size, blocks, channels = audio_embeds.shape
+        audio_embeds = audio_embeds.view(batch_size, window_size * blocks * channels)
+
+        # process audio of latter frame
+        audio_embeds_vf = rearrange(audio_embeds_vf, "bz f w b c -> (bz f) w b c")
+        batch_size_vf, window_size_vf, blocks_vf, channels_vf = audio_embeds_vf.shape
+        audio_embeds_vf = audio_embeds_vf.view(batch_size_vf, window_size_vf * blocks_vf * channels_vf)
+
+        # first projection
+        audio_embeds = torch.relu(self.proj1(audio_embeds)) 
+        audio_embeds_vf = torch.relu(self.proj1_vf(audio_embeds_vf)) 
+        audio_embeds = rearrange(audio_embeds, "(bz f) c -> bz f c", bz=B)
+        audio_embeds_vf = rearrange(audio_embeds_vf, "(bz f) c -> bz f c", bz=B)
+        audio_embeds_c = torch.concat([audio_embeds, audio_embeds_vf], dim=1) 
+        batch_size_c, N_t, C_a = audio_embeds_c.shape
+        audio_embeds_c = audio_embeds_c.view(batch_size_c*N_t, C_a)
+
+        # second projection
+        audio_embeds_c = torch.relu(self.proj2(audio_embeds_c))
+
+        context_tokens = self.proj3(audio_embeds_c).reshape(batch_size_c*N_t, self.context_tokens, self.output_dim)
+
+        # normalization and reshape
+        with amp.autocast(dtype=torch.float32):
+            context_tokens = self.norm(context_tokens)
+        context_tokens = rearrange(context_tokens, "(bz f) m c -> bz f m c", f=video_length)
+
+        return context_tokens

flash_head/src/pipeline/flash_head_pipeline.py

@@ -0,0 +1,379 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import os
+from PIL import Image
+from loguru import logger
+import time
+import numpy as np
+import torch
+import torch.distributed as dist
+from einops import rearrange
+
+from transformers import Wav2Vec2FeatureExtractor
+
+from flash_head.src.modules.flash_head_model import WanModelAudioProject
+from flash_head.audio_analysis.wav2vec2 import Wav2Vec2Model
+from flash_head.utils.utils import match_and_blend_colors_torch, resize_and_centercrop
+from flash_head.utils.facecrop import process_image
+
+# compile models to speedup inference
+COMPILE_MODEL = True
+COMPILE_VAE = True
+# use parallel vae to speedup decode/encode, only support WanVAE
+USE_PARALLEL_VAE = True
+
+
+class _TAEAdapter:
+    """Adapter to make TAEHV look like WanVAE (encode/decode signatures).
+
+    FlashHead code expects:
+    - encode(video[B=1,C,T,H,W] in [-1,1]) -> latent[C_lat,T_lat,h,w]
+    - decode(latent[C_lat,T_lat,h,w]) -> video[B=1,C,T,H,W] in [-1,1]
+
+    TAEHV expects NTCHW layout and operates in [0,1] pixel space.
+    """
+
+    def __init__(self, tae, dtype, device):
+        self.model = tae.eval().requires_grad_(False).to(device=device, dtype=dtype)
+        self.dtype = dtype
+        self.device = device
+
+    @torch.no_grad()
+    def encode(self, video, world_size_h=None, world_size_w=None):
+        if video is None or video.dim() != 5:
+            raise ValueError(f"TAE encode expects [1,C,T,H,W], got {None if video is None else tuple(video.shape)}")
+
+        # [1,C,T,H,W] in [-1,1] -> [1,T,C,H,W] in [0,1]
+        x = video.to(device=self.device, dtype=self.dtype)
+        x = (x + 1.0) / 2.0
+        x = x.permute(0, 2, 1, 3, 4).contiguous()
+
+        z = self.model.encode_video(x, parallel=True, show_progress_bar=False)  # [1,T_lat,C_lat,h,w]
+        z = z.permute(0, 2, 1, 3, 4).contiguous().squeeze(0)  # [C_lat,T_lat,h,w]
+        return z
+
+    @torch.no_grad()
+    def decode(self, zs):
+        if zs is None or zs.dim() != 4:
+            raise ValueError(f"TAE decode expects [C_lat,T_lat,h,w], got {None if zs is None else tuple(zs.shape)}")
+
+        z = zs.to(device=self.device, dtype=self.dtype)
+        z = z.unsqueeze(0).permute(0, 2, 1, 3, 4).contiguous()  # [1,T_lat,C_lat,h,w]
+
+        x = self.model.decode_video(z, parallel=True, show_progress_bar=False)  # [1,T,C,H,W] in [0,1]
+        x = x.permute(0, 2, 1, 3, 4).contiguous()  # [1,C,T,H,W]
+        x = x * 2.0 - 1.0
+        return x.clamp_(-1, 1)
+
+def get_cond_image_dict(cond_image_path_or_dir, use_face_crop):
+    def get_image(cond_image_path, use_face_crop):
+        if use_face_crop:
+            try:
+                image = process_image(cond_image_path)
+                return image
+            except Exception as e:
+                logger.error(f"Error processing {cond_image_path}: {e}")
+        return Image.open(cond_image_path).convert("RGB")
+
+    if os.path.isdir(cond_image_path_or_dir):
+        import glob
+        cond_image_list = glob.glob(os.path.join(cond_image_path_or_dir, "*.png"))
+        cond_image_list.sort()
+        cond_image_dict = {cond_image.split("/")[-1].split(".")[0]: get_image(cond_image, use_face_crop) for cond_image in cond_image_list}
+    else:
+        cond_image_dict = {cond_image_path_or_dir.split("/")[-1].split(".")[0]: get_image(cond_image_path_or_dir, use_face_crop)}
+    return cond_image_dict
+
+def timestep_transform(
+    t,
+    shift=5.0,
+    num_timesteps=1000,
+):
+    t = t / num_timesteps
+    # shift the timestep based on ratio
+    new_t = shift * t / (1 + (shift - 1) * t)
+    new_t = new_t * num_timesteps
+    return new_t
+
+
+class FlashHeadPipeline:
+    def __init__(
+        self,
+        checkpoint_dir,
+        model_type,
+        wav2vec_dir,
+        device="cuda",
+        param_dtype=torch.bfloat16,
+        use_usp=False,
+        num_timesteps=1000,
+        use_timestep_transform=True,
+        use_tae=False,
+        tae_path=None,
+        tae_model_type="wan21",
+    ):
+        r"""
+        Initializes the image-to-video generation model components.
+        Args:
+            checkpoint_dir (`str`):
+                Path to directory containing model checkpoints
+            wav2vec_dir (`str`):
+                Path to directory containing wav2vec checkpoints
+            use_usp (`bool`, *optional*, defaults to False):
+                Enable distribution strategy of USP.
+        """
+        self.param_dtype = param_dtype
+        self.device = device
+        self.rank = dist.get_rank() if dist.is_initialized() else 0
+        self.use_usp = use_usp and dist.is_initialized()
+        self.model_type = model_type
+        self.use_ltx = model_type == "lite"
+
+        if self.use_ltx:
+            model_dir = os.path.join(checkpoint_dir, "Model_Lite")
+            vae_dir = os.path.join(checkpoint_dir, "VAE_LTX")
+
+            from flash_head.ltx_video.ltx_vae import LtxVAE
+            self.vae = LtxVAE(
+                pretrained_model_type_or_path=vae_dir,
+                dtype=self.param_dtype,
+                device=self.device,
+            )
+        else:
+            vae_path = os.path.join(checkpoint_dir, "VAE_Wan/Wan2.1_VAE.pth")
+
+            if use_tae:
+                if tae_path is None:
+                    tae_path = os.path.join(checkpoint_dir, "VAE_Wan/taew2_1.pth")
+                from flash_head.wan.modules.tae import TAEHV
+
+                logger.info(f"Using TAEHV for VAE encode/decode: {tae_path}")
+                if USE_PARALLEL_VAE and self.use_usp:
+                    logger.warning("USE_PARALLEL_VAE is enabled but TAEHV does not support VAE parallel sharding; falling back to single-rank VAE.")
+
+                tae = TAEHV(checkpoint_path=tae_path, model_type=str(tae_model_type))
+                self.vae = _TAEAdapter(tae, dtype=self.param_dtype, device=self.device)
+                self.use_tae = True
+            else:
+                from flash_head.wan.modules import WanVAE
+                self.vae = WanVAE(
+                    vae_path=vae_path,
+                    dtype=self.param_dtype,
+                    device=self.device,
+                    parallel=(USE_PARALLEL_VAE and self.use_usp),
+                )
+                self.use_tae = False
+
+            if self.model_type == "pretrained":
+                self.audio_guide_scale = 3.0
+                model_dir = os.path.join(checkpoint_dir, "teacher")
+            elif self.model_type == "pro":
+                model_dir = os.path.join(checkpoint_dir, "Model_Pro")
+        
+        self.model = WanModelAudioProject.from_pretrained(model_dir)
+        self.model.eval().requires_grad_(False)
+        self.model.to(device=self.device, dtype=self.param_dtype)
+
+        self.config = self.model.config
+
+        if use_usp:
+            from xfuser.core.distributed import get_sequence_parallel_world_size
+            self.sp_size = get_sequence_parallel_world_size()
+        else:
+            self.sp_size = 1
+
+        if dist.is_initialized():
+            dist.barrier()
+
+        self.num_timesteps = num_timesteps
+        self.use_timestep_transform = use_timestep_transform
+
+        if COMPILE_MODEL:
+            self.model = torch.compile(self.model)
+        if COMPILE_VAE:
+            if self.use_ltx:
+                self.vae.model.encode = torch.compile(self.vae.model.encode)
+                self.vae.model.decode = torch.compile(self.vae.model.decode)
+            else:
+                if getattr(self, "use_tae", False):
+                    logger.warning("COMPILE_VAE is enabled, but TAEHV backend is selected; skipping torch.compile for VAE.")
+                else:
+                    self.vae.encode = torch.compile(self.vae.encode)
+                    self.vae.decode = torch.compile(self.vae.decode)
+
+        self.audio_encoder = Wav2Vec2Model.from_pretrained(wav2vec_dir, local_files_only=True).to(self.device)
+        self.audio_encoder.feature_extractor._freeze_parameters()
+        self.wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_dir, local_files_only=True)
+
+    @torch.no_grad()
+    def prepare_params(self,
+                        cond_image_path_or_dir,
+                        target_size,
+                        frame_num,
+                        motion_frames_num,
+                        sampling_steps,
+                        seed=None,
+                        shift=5.0,
+                        color_correction_strength=0.0,
+                        use_face_crop=False,
+                        ):
+        self.cond_image_dict = get_cond_image_dict(cond_image_path_or_dir, use_face_crop)
+
+        self.frame_num = frame_num
+        self.motion_frames_num = motion_frames_num
+        self.color_correction_strength = color_correction_strength
+
+        self.target_h, self.target_w = target_size
+        self.lat_h, self.lat_w = self.target_h // self.config.vae_stride[1], self.target_w // self.config.vae_stride[2]
+
+        self.generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        # prepare timesteps
+        if sampling_steps == 2:
+            timesteps = [1000, 500]
+        elif sampling_steps == 4:
+            timesteps = [1000, 750, 500, 250]
+        else:
+            timesteps = list(np.linspace(self.num_timesteps, 1, sampling_steps, dtype=np.float32))
+            
+        timesteps.append(0.)
+        timesteps = [torch.tensor([t], device=self.device) for t in timesteps]
+        if self.use_timestep_transform:
+            timesteps = [timestep_transform(t, shift=shift, num_timesteps=self.num_timesteps) for t in timesteps]
+        self.timesteps = timesteps
+
+        self.cond_image_tensor_dict = {}
+        self.ref_img_latent_dict = {}
+        for i, (person_name, cond_image_pil) in enumerate(self.cond_image_dict.items()):
+            cond_image_tensor = resize_and_centercrop(cond_image_pil, (self.target_h, self.target_w)).to(self.device, dtype=self.param_dtype) # 1 C 1 H W
+            cond_image_tensor = (cond_image_tensor / 255 - 0.5) * 2
+
+            self.cond_image_tensor_dict[person_name] = cond_image_tensor
+
+            video_frames = cond_image_tensor.repeat(1, 1, self.frame_num, 1, 1)
+            self.ref_img_latent_dict[person_name] = self.vae.encode(video_frames) # (16, 9, 64, 64) / (128, 5, 16, 16)
+            if i == 0:
+                self.reset_person_name(person_name)
+
+        return
+    
+    @torch.no_grad()
+    def reset_person_name(self, person_name=None):
+        if person_name is None or person_name not in self.cond_image_dict:
+            pass
+        else:
+            self.person_name = person_name
+        self.original_color_reference = self.cond_image_tensor_dict[self.person_name]
+        self.ref_img_latent = self.ref_img_latent_dict[self.person_name]
+        self.latent_motion_frames = self.ref_img_latent[:, :1].clone()
+
+    @torch.no_grad()
+    def preprocess_audio(self, speech_array, sr=16000, fps=25):
+        video_length = len(speech_array) * fps / sr
+
+        # wav2vec_feature_extractor
+        audio_feature = np.squeeze(
+            self.wav2vec_feature_extractor(speech_array, sampling_rate=sr).input_values
+        )
+        audio_feature = torch.from_numpy(audio_feature).float().to(device=self.device)
+        audio_feature = audio_feature.unsqueeze(0)
+
+        # audio encoder
+        with torch.no_grad():
+            embeddings = self.audio_encoder(audio_feature, seq_len=int(video_length), output_hidden_states=True)
+
+        if len(embeddings) == 0:
+            logger.error("Fail to extract audio embedding")
+            return None
+
+        audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
+        audio_emb = rearrange(audio_emb, "b s d -> s b d")
+        return audio_emb
+
+    @torch.no_grad()
+    def generate(self, audio_embedding):
+        # evaluation mode
+        with torch.no_grad():
+
+            # sample videos
+            noise = torch.randn(
+                self.config.out_dim, 
+                (self.frame_num - 1) // self.config.vae_stride[0] + 1,
+                self.lat_h,
+                self.lat_w,
+                dtype=self.param_dtype,
+                device=self.device,
+                generator=self.generator)
+
+            for i in range(len(self.timesteps)-1):
+                torch.cuda.synchronize()
+                start_time = time.time()
+
+                noise[:, :self.latent_motion_frames.shape[1]] = self.latent_motion_frames
+
+                flow_pred = self.model(
+                    x=noise.unsqueeze(0),
+                    timestep=self.timesteps[i],
+                    context=audio_embedding,
+                    y=self.ref_img_latent.unsqueeze(0),
+                )[0]
+
+                if self.model_type == "pretrained":
+                    flow_pred_drop_audio = self.model(
+                        x=noise.unsqueeze(0),
+                        timestep=self.timesteps[i],
+                        context=torch.zeros_like(audio_embedding),
+                        y=self.ref_img_latent.unsqueeze(0),
+                    )[0]
+                    flow_pred = flow_pred_drop_audio + self.audio_guide_scale * (flow_pred - flow_pred_drop_audio)
+
+                    # update latent
+                    dt = self.timesteps[i] - self.timesteps[i + 1]
+                    dt = (dt / self.num_timesteps).to(self.param_dtype)
+                    noise = noise - flow_pred * dt[:, None, None, None]
+                
+                else:
+                    # update latent
+                    t_i = (self.timesteps[i][:, None, None, None] / self.num_timesteps).to(self.param_dtype)
+                    t_i_1 = (self.timesteps[i+1][:, None, None, None] / self.num_timesteps).to(self.param_dtype)
+                    x_0 = noise - flow_pred * t_i
+
+                    noise = (1 - t_i_1) * x_0 + t_i_1 * torch.randn(x_0.size(), dtype=x_0.dtype, device=self.device, generator=self.generator)
+
+                torch.cuda.synchronize()
+                end_time = time.time()
+                if self.rank == 0:
+                    print(f'[generate] model denoise per step: {end_time - start_time}s')
+
+            noise[:, :self.latent_motion_frames.shape[1]] = self.latent_motion_frames
+
+            torch.cuda.synchronize()
+            start_decode_time = time.time()
+
+            videos = self.vae.decode(noise)
+
+            torch.cuda.synchronize()
+            end_decode_time = time.time()
+            if self.rank == 0:
+                print(f'[generate] decode video frames: {end_decode_time - start_decode_time}s')
+        
+        torch.cuda.synchronize()
+        start_color_correction_time = time.time()
+        if self.color_correction_strength > 0.0:
+            videos = match_and_blend_colors_torch(videos, self.original_color_reference, self.color_correction_strength)
+
+        cond_frame = videos[:, :, -self.motion_frames_num:].to(self.device)
+        torch.cuda.synchronize()
+        end_color_correction_time = time.time()
+        if self.rank == 0:
+            print(f'[generate] color correction: {end_color_correction_time - start_color_correction_time}s')
+
+        torch.cuda.synchronize()
+        start_encode_time = time.time()
+        self.latent_motion_frames = self.vae.encode(cond_frame)
+        torch.cuda.synchronize()
+        end_encode_time = time.time()
+        if self.rank == 0:
+            print(f'[generate] encode motion frames: {end_encode_time - start_encode_time}s')
+
+        gen_video_samples = videos #[:, :, self.motion_frames_num:]
+
+        return gen_video_samples[0].to(torch.float32)

flash_head/utils/cpu_face_handler.py

@@ -0,0 +1,55 @@
+import mediapipe as mp
+import numpy as np
+from typing import Tuple, List
+
+
+class CPUFaceHandler:
+    """Handler for CPU-based face detection using MediaPipe.
+    (2 ms/frame)
+    This handler provides a simple interface for face detection using MediaPipe's
+    face detection model. It's optimized for CPU usage and provides basic face
+    detection functionality.
+    """
+
+    def __init__(self, model_selection: int = 1, min_detection_confidence: float = 0.0):
+        """Initialize the face detection handler."""
+        self.detector = mp.solutions.face_detection.FaceDetection(
+            model_selection=model_selection,
+            min_detection_confidence=min_detection_confidence,
+        )
+
+    def detect(self, image: np.ndarray) -> Tuple[int, List[int]]:
+        """Detect faces in the given image.
+
+        Args:
+            image (np.ndarray): RGB image array.
+
+        Returns:
+            Tuple[int, List[int]]: A tuple containing:
+                - Number of faces detected (int)
+                - Bounding box coordinates [x1, y1, x2, y2] if exactly one face is detected,
+                  empty list otherwise
+        """
+        bboxs, scores = [], []
+        results = self.detector.process(image)
+        detection_result = results.detections
+        if detection_result is None:
+            return bboxs, scores
+        for detection in detection_result:
+            bboxC = detection.location_data.relative_bounding_box
+            x, y, w, h = bboxC.xmin, bboxC.ymin, bboxC.width, bboxC.height
+            x1, y1, x2, y2 = x, y, x + w, y + h
+            bboxs.append([x1, y1, x2, y2])
+            scores.append(detection.score[0])
+        return bboxs, scores
+
+    def __call__(self, image: np.ndarray) -> Tuple[int, List[int]]:
+        """Make the handler callable.
+
+        Args:
+            image (np.ndarray): RGB image array.
+
+        Returns:
+            Tuple[int, List[int]]: Same as detect() method.
+        """
+        return self.detect(image)

flash_head/utils/facecrop.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""
+人脸裁剪处理脚本
+从单张图像中检测人脸，裁剪并调整大小到指定尺寸
+"""
+import os
+from PIL import Image
+import numpy as np
+
+from flash_head.utils.cpu_face_handler import CPUFaceHandler
+
+def get_scaled_bbox(
+    bbox, img_w, img_h, ratio: float = 1.0, face_image: Image.Image = None
+):
+    """
+    根据人脸边界框计算缩放后的裁剪区域
+    
+    Args:
+        bbox: 人脸边界框 [x1, y1, x2, y2]
+        img_w: 图像宽度
+        img_h: 图像高度
+        ratio: 缩放比例，数值越大，人脸在画面中的比例越小（周围留白越多）
+        face_image: PIL Image 对象
+    
+    Returns:
+        裁剪后的人脸图像
+    """
+    x1, y1, x2, y2 = bbox
+
+    # Calculate center point
+    center_x = (x1 + x2) / 2
+    center_y = (y1 + y2) / 2
+
+    # Calculate width and height
+    width = x2 - x1
+
+    # Scale width and height
+    new_width = width * ratio
+    new_height = new_width
+
+    # tile pix
+    dis_x_left = new_width * 0.5
+    dis_x_right = new_width - dis_x_left  # 0.5new_width
+    dis_y_up = new_height * 0.55
+    dis_y_down = new_height - dis_y_up  # 0.45new_height
+
+    # Calculate new coordinates
+    new_x1 = int(max(0, center_x - dis_x_left))
+    new_y1 = int(max(0, center_y - dis_y_up))
+    new_x2 = int(min(img_w, center_x + dis_x_right))
+    new_y2 = int(min(img_h, center_y + dis_y_down))
+    scaled_bbox = [new_x1, new_y1, new_x2, new_y2]
+    crop_face = face_image.crop(scaled_bbox)
+    return crop_face
+
+
+def process_image(
+    input_path,
+    face_ratio=2.0,
+    target_size=(512, 512),
+):
+    """
+    处理单张图像，进行人脸检测和裁剪
+    
+    Args:
+        input_path: 输入图像路径
+        face_ratio: 人脸缩放比例，建议范围：1.5-3.0，默认2.0
+        target_size: 输出图像尺寸，默认(512, 512)
+    
+    Returns:
+        imgae: 处理后的图像
+    """
+    # 初始化人脸检测器
+    face_detector = CPUFaceHandler()
+    
+    # 验证输入文件
+    if not os.path.isfile(input_path):
+        raise ValueError(f"File not found: {input_path}")
+    
+    try:
+        # 读取图像
+        image = Image.open(input_path)
+        image = image.convert("RGB")
+        image_rgb = np.array(image)
+        img_h, img_w = image_rgb.shape[:2]
+        
+        # 检测人脸
+        boxes, scores = face_detector(image_rgb)
+        
+        if len(boxes) == 0:
+            raise ValueError("No face detected")
+        
+        # 转换边界框坐标（从相对坐标转为绝对坐标）
+        boxes_abs = [
+            boxes[0][0] * img_w,
+            boxes[0][1] * img_h,
+            boxes[0][2] * img_w,
+            boxes[0][3] * img_h
+        ]
+        
+        # 裁剪人脸
+        crop_face = get_scaled_bbox(boxes_abs, img_w, img_h, face_ratio, image)
+        
+        # 调整大小
+        crop_face = crop_face.resize(target_size)
+        
+        return crop_face
+            
+    except Exception as e:
+        raise ValueError(f"Error processing {input_path}: {e}")

flash_head/utils/utils.py

@@ -0,0 +1,222 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+from PIL import Image
+import torchvision.transforms as transforms
+import torch.nn as nn
+import pyloudnorm as pyln
+
+def rgb_to_lab_torch(rgb: torch.Tensor) -> torch.Tensor:
+    """
+    PyTorch GPU版本：RGB转Lab颜色空间（输入范围[0,1]，张量形状任意，最后一维为通道数）
+    参考CIE 1931标准转换公式
+    """
+    # 转换为线性RGB（sRGB伽马校正逆过程）
+    linear_rgb = torch.where(
+        rgb > 0.04045,
+        ((rgb + 0.055) / 1.055) ** 2.4,
+        rgb / 12.92
+    )
+    
+    # 线性RGB转XYZ（使用sRGB标准白点D65）
+    xyz_from_rgb = torch.tensor([
+        [0.4124564, 0.3575761, 0.1804375],
+        [0.2126729, 0.7151522, 0.0721750],
+        [0.0193339, 0.1191920, 0.9503041]
+    ], dtype=rgb.dtype, device=rgb.device)
+    
+    # 维度适配：确保输入为(B, ..., C)，矩阵乘法后保持空间维度
+    shape = linear_rgb.shape
+    linear_rgb_flat = linear_rgb.reshape(-1, 3)  # (N, 3)，N=B*T*H*W
+    xyz_flat = linear_rgb_flat @ xyz_from_rgb.T  # (N, 3)
+    xyz = xyz_flat.reshape(shape)  # 恢复原形状
+    
+    # XYZ转Lab（使用D65白点参数）
+    xyz_ref = torch.tensor([0.95047, 1.0, 1.08883], dtype=rgb.dtype, device=rgb.device)
+    xyz_normalized = xyz / xyz_ref[None, None, None, None, :]  # 广播适配(B, C, T, H, W)
+    
+    # 应用Lab转换公式
+    epsilon = 0.008856
+    kappa = 903.3
+    xyz_normalized = torch.clamp(xyz_normalized, 1e-8, 1.0)  # 避免log(0)
+    
+    f_xyz = torch.where(
+        xyz_normalized > epsilon,
+        xyz_normalized ** (1/3),
+        (kappa * xyz_normalized + 16) / 116
+    )
+    
+    L = 116 * f_xyz[..., 1] - 16  # Y通道对应亮度
+    a = 500 * (f_xyz[..., 0] - f_xyz[..., 1])  # X-Y对应红绿
+    b = 200 * (f_xyz[..., 1] - f_xyz[..., 2])  # Y-Z对应蓝黄
+    
+    lab = torch.stack([L, a, b], dim=-1)  # 最后一维拼接为Lab通道
+    return lab
+
+def lab_to_rgb_torch(lab: torch.Tensor) -> torch.Tensor:
+    """
+    PyTorch GPU版本：Lab转RGB颜色空间（输出范围[0,1]，张量形状任意，最后一维为通道数）
+    """
+    # Lab分离通道
+    L = lab[..., 0]
+    a = lab[..., 1]
+    b = lab[..., 2]
+    
+    # Lab转XYZ
+    f_y = (L + 16) / 116
+    f_x = (a / 500) + f_y
+    f_z = f_y - (b / 200)
+    
+    epsilon = 0.008856
+    kappa = 903.3
+    
+    x = torch.where(f_x ** 3 > epsilon, f_x ** 3, (116 * f_x - 16) / kappa)
+    y = torch.where(L > kappa * epsilon, ((L + 16) / 116) ** 3, L / kappa)
+    z = torch.where(f_z ** 3 > epsilon, f_z ** 3, (116 * f_z - 16) / kappa)
+    
+    # 乘以D65白点参数
+    xyz_ref = torch.tensor([0.95047, 1.0, 1.08883], dtype=lab.dtype, device=lab.device)
+    xyz = torch.stack([x, y, z], dim=-1) * xyz_ref[None, None, None, None, :]
+    
+    # XYZ转线性RGB
+    rgb_from_xyz = torch.tensor([
+        [3.2404542, -1.5371385, -0.4985314],
+        [-0.9692660, 1.8760108, 0.0415560],
+        [0.0556434, -0.2040259, 1.0572252]
+    ], dtype=lab.dtype, device=lab.device)
+    
+    # 维度适配：矩阵乘法
+    shape = xyz.shape
+    xyz_flat = xyz.reshape(-1, 3)  # (N, 3)
+    linear_rgb_flat = xyz_flat @ rgb_from_xyz.T  # (N, 3)
+    linear_rgb = linear_rgb_flat.reshape(shape)  # 恢复原形状
+    
+    # 线性RGB转sRGB（伽马校正）
+    rgb = torch.where(
+        linear_rgb > 0.0031308,
+        1.055 * (linear_rgb ** (1/2.4)) - 0.055,
+        12.92 * linear_rgb
+    )
+    
+    # 确保输出在[0,1]范围内
+    rgb = torch.clamp(rgb, 0.0, 1.0)
+    return rgb
+
+def match_and_blend_colors_torch(
+    source_chunk: torch.Tensor, 
+    reference_image: torch.Tensor, 
+    strength: float
+) -> torch.Tensor:
+    """
+    全GPU批量运算版本：将视频chunk的颜色匹配到参考图像并混合（支持B>1、T帧并行）
+    
+    Args:
+        source_chunk (torch.Tensor): 视频chunk (B, C, T, H, W)，范围[-1, 1]
+        reference_image (torch.Tensor): 参考图像 (B, C, 1, H, W)，范围[-1, 1]（B需与source_chunk一致）
+        strength (float): 颜色校正强度 (0.0-1.0)，0.0无校正，1.0完全校正
+    
+    Returns:
+        torch.Tensor: 颜色校正后的视频chunk (B, C, T, H, W)，范围[-1, 1]
+    """
+    # 强度为0直接返回原图
+    if strength <= 0.0:
+        return source_chunk.clone()
+    
+    # 验证强度范围
+    if not 0.0 <= strength <= 1.0:
+        raise ValueError(f"Strength必须在0.0-1.0之间，当前值：{strength}")
+    
+    # 验证输入形状（确保B一致，参考图T=1）
+    B, C, T, H, W = source_chunk.shape
+    assert reference_image.shape == (B, C, 1, H, W), \
+        f"参考图像形状需为(B, C, 1, H, W)，当前为{reference_image.shape}"
+    assert C == 3, f"仅支持3通道RGB图像，当前通道数：{C}"
+    
+    # 保持设备和数据类型一致
+    device = source_chunk.device
+    dtype = source_chunk.dtype
+    reference_image = reference_image.to(device=device, dtype=dtype)
+    
+    # 1. 从[-1,1]转换到[0,1]（GPU上直接运算）
+    source_01 = (source_chunk + 1.0) / 2.0
+    ref_01 = (reference_image + 1.0) / 2.0
+    
+    # 2. 调整维度顺序：(B, C, T, H, W) → (B, T, H, W, C)（适配颜色空间转换）
+    # 参考图：(B, C, 1, H, W) → (B, 1, H, W, C)
+    source_permuted = source_01.permute(0, 2, 3, 4, 1)  # 通道移到最后一维
+    ref_permuted = ref_01.permute(0, 2, 3, 4, 1)
+    
+    # 3. RGB转Lab（批量处理所有帧）
+    source_lab = rgb_to_lab_torch(source_permuted)
+    ref_lab = rgb_to_lab_torch(ref_permuted)  # (B, 1, H, W, 3)
+    
+    # 4. 批量颜色迁移：匹配L/a/b通道的均值和标准差（核心逻辑）
+    # 计算参考图各通道的均值和标准差（对H、W维度求统计，保持B维度）
+    ref_mean = ref_lab.mean(dim=[2, 3], keepdim=True)  # (B, 1, 1, 1, 3)
+    ref_std = ref_lab.std(dim=[2, 3], keepdim=True, unbiased=False)  # (B, 1, 1, 1, 3)
+    
+    # 计算源视频各通道的均值和标准差（对H、W维度求统计，保持B、T维度）
+    source_mean = source_lab.mean(dim=[2, 3], keepdim=True)  # (B, T, 1, 1, 3)
+    source_std = source_lab.std(dim=[2, 3], keepdim=True, unbiased=False)  # (B, T, 1, 1, 3)
+    
+    # 避免标准差为0的除法错误（用1.0替代0）
+    source_std_safe = torch.where(source_std < 1e-8, torch.ones_like(source_std), source_std)
+    
+    # 颜色迁移公式：(源 - 源均值) * (参考标准差/源标准差) + 参考均值
+    corrected_lab = (source_lab - source_mean) * (ref_std / source_std_safe) + ref_mean
+    
+    # 5. Lab转RGB（批量转换所有校正后的帧）
+    corrected_rgb_01 = lab_to_rgb_torch(corrected_lab)
+    
+    # 6. 批量混合原始帧和校正帧（按强度加权）
+    blended_rgb_01 = (1 - strength) * source_permuted + strength * corrected_rgb_01
+    
+    # 7. 还原维度顺序和数值范围：(B, T, H, W, C) → (B, C, T, H, W)，范围[0,1]→[-1,1]
+    blended_rgb_01 = blended_rgb_01.permute(0, 4, 1, 2, 3)  # 通道移回第二维
+    blended_rgb_minus1_1 = (blended_rgb_01 * 2.0) - 1.0
+    
+    # 8. 确保输出格式正确（连续内存布局）
+    output = blended_rgb_minus1_1.contiguous().to(device=device, dtype=dtype)
+    
+    return output
+
+def resize_and_centercrop(cond_image, target_size):
+    """
+    Resize image or tensor to the target size without padding.
+    """
+
+    # Get the original size
+    if isinstance(cond_image, torch.Tensor):
+        _, orig_h, orig_w = cond_image.shape
+    else:
+        orig_h, orig_w = cond_image.height, cond_image.width
+
+    target_h, target_w = target_size
+    
+    # Calculate the scaling factor for resizing
+    scale_h = target_h / orig_h
+    scale_w = target_w / orig_w
+    
+    # Compute the final size
+    scale = max(scale_h, scale_w)
+    final_h = math.ceil(scale * orig_h)
+    final_w = math.ceil(scale * orig_w)
+    
+    # Resize
+    if isinstance(cond_image, torch.Tensor):
+        if len(cond_image.shape) == 3:
+            cond_image = cond_image[None]
+        resized_tensor = nn.functional.interpolate(cond_image, size=(final_h, final_w), mode='nearest').contiguous() 
+        # crop
+        cropped_tensor = transforms.functional.center_crop(resized_tensor, target_size) 
+        cropped_tensor = cropped_tensor.squeeze(0)
+    else:
+        resized_image = cond_image.resize((final_w, final_h), resample=Image.BILINEAR)
+        resized_image = np.array(resized_image)
+        # tensor and crop
+        resized_tensor = torch.from_numpy(resized_image)[None, ...].permute(0, 3, 1, 2).contiguous()
+        cropped_tensor = transforms.functional.center_crop(resized_tensor, target_size)
+        cropped_tensor = cropped_tensor[:, :, None, :, :] 
+
+    return cropped_tensor

flash_head/wan/modules/__init__.py

@@ -0,0 +1,7 @@
+from .vae import WanVAE
+from .tae import TAEHV
+
+__all__ = [
+    'WanVAE',
+    'TAEHV',
+]

flash_head/wan/modules/tae.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+"""
+Tiny AutoEncoder for Hunyuan Video
+(DNN for encoding / decoding videos to Hunyuan Video's latent space)
+"""
+
+import os
+from collections import namedtuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file
+from tqdm.auto import tqdm
+
+DecoderResult = namedtuple("DecoderResult", ("frame", "memory"))
+TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
+
+
+def conv(n_in, n_out, **kwargs):
+    return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
+
+
+class Clamp(nn.Module):
+    def forward(self, x):
+        return torch.tanh(x / 3) * 3
+
+
+class MemBlock(nn.Module):
+    def __init__(self, n_in, n_out, act_func):
+        super().__init__()
+        self.conv = nn.Sequential(conv(n_in * 2, n_out), act_func, conv(n_out, n_out), act_func, conv(n_out, n_out))
+        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
+        self.act = act_func
+
+    def forward(self, x, past):
+        return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x))
+
+
+class TPool(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(n_f * stride, n_f, 1, bias=False)
+
+    def forward(self, x):
+        _NT, C, H, W = x.shape
+        return self.conv(x.reshape(-1, self.stride * C, H, W))
+
+
+class TGrow(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(n_f, n_f * stride, 1, bias=False)
+
+    def forward(self, x):
+        _NT, C, H, W = x.shape
+        x = self.conv(x)
+        return x.reshape(-1, C, H, W)
+
+
+def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
+    """
+    Apply a sequential model with memblocks to the given input.
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - parallel: if True, parallelize over timesteps (fast but uses O(T) memory)
+        if False, each timestep will be processed sequentially (slow but uses O(1) memory)
+    - show_progress_bar: if True, enables tqdm progressbar display
+
+    Returns NTCHW tensor of output data.
+    """
+    assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
+    N, T, C, H, W = x.shape
+    if parallel:
+        x = x.reshape(N * T, C, H, W)
+        # parallel over input timesteps, iterate over blocks
+        for b in tqdm(model, disable=not show_progress_bar):
+            if isinstance(b, MemBlock):
+                NT, C, H, W = x.shape
+                T = NT // N
+                _x = x.reshape(N, T, C, H, W)
+                mem = F.pad(_x, (0, 0, 0, 0, 0, 0, 1, 0), value=0)[:, :T].reshape(x.shape)
+                x = b(x, mem)
+            else:
+                x = b(x)
+        NT, C, H, W = x.shape
+        T = NT // N
+        x = x.view(N, T, C, H, W)
+    else:
+        # TODO(oboerbohan): at least on macos this still gradually uses more memory during decode...
+        # need to fix :(
+        out = []
+        # iterate over input timesteps and also iterate over blocks.
+        # because of the cursed TPool/TGrow blocks, this is not a nested loop,
+        # it's actually a ***graph traversal*** problem! so let's make a queue
+        work_queue = [TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))]
+        # in addition to manually managing our queue, we also need to manually manage our progressbar.
+        # we'll update it for every source node that we consume.
+        progress_bar = tqdm(range(T), disable=not show_progress_bar)
+        # we'll also need a separate addressable memory per node as well
+        mem = [None] * len(model)
+        while work_queue:
+            xt, i = work_queue.pop(0)
+            if i == 0:
+                # new source node consumed
+                progress_bar.update(1)
+            if i == len(model):
+                # reached end of the graph, append result to output list
+                out.append(xt)
+            else:
+                # fetch the block to process
+                b = model[i]
+                if isinstance(b, MemBlock):
+                    # mem blocks are simple since we're visiting the graph in causal order
+                    if mem[i] is None:
+                        xt_new = b(xt, xt * 0)
+                        mem[i] = xt
+                    else:
+                        xt_new = b(xt, mem[i])
+                        mem[i].copy_(xt)  # inplace might reduce mysterious pytorch memory allocations? doesn't help though
+                    # add successor to work queue
+                    work_queue.insert(0, TWorkItem(xt_new, i + 1))
+                elif isinstance(b, TPool):
+                    # pool blocks are miserable
+                    if mem[i] is None:
+                        mem[i] = []  # pool memory is itself a queue of inputs to pool
+                    mem[i].append(xt)
+                    if len(mem[i]) > b.stride:
+                        # pool mem is in invalid state, we should have pooled before this
+                        raise ValueError("???")
+                    elif len(mem[i]) < b.stride:
+                        # pool mem is not yet full, go back to processing the work queue
+                        pass
+                    else:
+                        # pool mem is ready, run the pool block
+                        N, C, H, W = xt.shape
+                        xt = b(torch.cat(mem[i], 1).view(N * b.stride, C, H, W))
+                        # reset the pool mem
+                        mem[i] = []
+                        # add successor to work queue
+                        work_queue.insert(0, TWorkItem(xt, i + 1))
+                elif isinstance(b, TGrow):
+                    xt = b(xt)
+                    NT, C, H, W = xt.shape
+                    # each tgrow has multiple successor nodes
+                    for xt_next in reversed(xt.view(N, b.stride * C, H, W).chunk(b.stride, 1)):
+                        # add successor to work queue
+                        work_queue.insert(0, TWorkItem(xt_next, i + 1))
+                else:
+                    # normal block with no funny business
+                    xt = b(xt)
+                    # add successor to work queue
+                    work_queue.insert(0, TWorkItem(xt, i + 1))
+        progress_bar.close()
+        x = torch.stack(out, 1)
+    return x
+
+
+class TAEHV(nn.Module):
+    def __init__(self, checkpoint_path="taehv.pth", decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True), patch_size=1, latent_channels=16, model_type="wan21"):
+        """Initialize pretrained TAEHV from the given checkpoint.
+
+        Arg:
+            checkpoint_path: path to weight file to load. taehv.pth for Hunyuan, taew2_1.pth for Wan 2.1.
+            decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+            decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+            patch_size: input/output pixelshuffle patch-size for this model.
+            latent_channels: number of latent channels (z dim) for this model.
+        """
+        super().__init__()
+        self.patch_size = patch_size
+        self.latent_channels = latent_channels
+        self.image_channels = 3
+        self.is_cogvideox = checkpoint_path is not None and "taecvx" in checkpoint_path
+        # if checkpoint_path is not None and "taew2_2" in checkpoint_path:
+        #     self.patch_size, self.latent_channels = 2, 48
+        self.model_type = model_type
+        if model_type == "wan22":
+            self.patch_size, self.latent_channels = 2, 48
+        if model_type == "hy15":
+            act_func = nn.LeakyReLU(0.2, inplace=True)
+        else:
+            act_func = nn.ReLU(inplace=True)
+
+        self.encoder = nn.Sequential(
+            conv(self.image_channels * self.patch_size**2, 64),
+            act_func,
+            TPool(64, 2),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64, act_func),
+            MemBlock(64, 64, act_func),
+            MemBlock(64, 64, act_func),
+            TPool(64, 2),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64, act_func),
+            MemBlock(64, 64, act_func),
+            MemBlock(64, 64, act_func),
+            TPool(64, 1),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64, act_func),
+            MemBlock(64, 64, act_func),
+            MemBlock(64, 64, act_func),
+            conv(64, self.latent_channels),
+        )
+        n_f = [256, 128, 64, 64]
+        self.frames_to_trim = 2 ** sum(decoder_time_upscale) - 1
+        self.decoder = nn.Sequential(
+            Clamp(),
+            conv(self.latent_channels, n_f[0]),
+            act_func,
+            MemBlock(n_f[0], n_f[0], act_func),
+            MemBlock(n_f[0], n_f[0], act_func),
+            MemBlock(n_f[0], n_f[0], act_func),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1),
+            TGrow(n_f[0], 1),
+            conv(n_f[0], n_f[1], bias=False),
+            MemBlock(n_f[1], n_f[1], act_func),
+            MemBlock(n_f[1], n_f[1], act_func),
+            MemBlock(n_f[1], n_f[1], act_func),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1),
+            TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1),
+            conv(n_f[1], n_f[2], bias=False),
+            MemBlock(n_f[2], n_f[2], act_func),
+            MemBlock(n_f[2], n_f[2], act_func),
+            MemBlock(n_f[2], n_f[2], act_func),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1),
+            TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1),
+            conv(n_f[2], n_f[3], bias=False),
+            act_func,
+            conv(n_f[3], self.image_channels * self.patch_size**2),
+        )
+        if checkpoint_path is not None:
+            ext = os.path.splitext(checkpoint_path)[1].lower()
+
+            if ext == ".pth":
+                state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+            elif ext == ".safetensors":
+                state_dict = load_file(checkpoint_path, device="cpu")
+            else:
+                raise ValueError(f"Unsupported checkpoint format: {ext}. Supported formats: .pth, .safetensors")
+
+            self.load_state_dict(self.patch_tgrow_layers(state_dict))
+
+    def patch_tgrow_layers(self, sd):
+        """Patch TGrow layers to use a smaller kernel if needed.
+
+        Args:
+            sd: state dict to patch
+        """
+        new_sd = self.state_dict()
+        for i, layer in enumerate(self.decoder):
+            if isinstance(layer, TGrow):
+                key = f"decoder.{i}.conv.weight"
+                if sd[key].shape[0] > new_sd[key].shape[0]:
+                    # take the last-timestep output channels
+                    sd[key] = sd[key][-new_sd[key].shape[0] :]
+        return sd
+
+    def encode_video(self, x, parallel=True, show_progress_bar=True):
+        """Encode a sequence of frames.
+
+        Args:
+            x: input NTCHW RGB (C=3) tensor with values in [0, 1].
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW latent tensor with ~Gaussian values.
+        """
+        if self.patch_size > 1:
+            x = F.pixel_unshuffle(x, self.patch_size)
+        if x.shape[1] % 4 != 0:
+            # pad at end to multiple of 4
+            n_pad = 4 - x.shape[1] % 4
+            padding = x[:, -1:].repeat_interleave(n_pad, dim=1)
+            x = torch.cat([x, padding], 1)
+        return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar)
+
+    def decode_video(self, x, parallel=True, show_progress_bar=True):
+        """Decode a sequence of frames.
+
+        Args:
+            x: input NTCHW latent (C=12) tensor with ~Gaussian values.
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW RGB tensor with ~[0, 1] values.
+        """
+        skip_trim = self.is_cogvideox and x.shape[1] % 2 == 0
+        x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
+        if self.model_type == "hy15":
+            x = x.clamp_(-1, 1)
+        else:
+            x = x.clamp_(0, 1)
+        if self.patch_size > 1:
+            x = F.pixel_shuffle(x, self.patch_size)
+        if skip_trim:
+            # skip trimming for cogvideox to make frame counts match.
+            # this still doesn't have correct temporal alignment for certain frame counts
+            # (cogvideox seems to pad at the start?), but for multiple-of-4 it's fine.
+            return x
+        return x[:, self.frames_to_trim :]

flash_head/wan/modules/vae.py

@@ -0,0 +1,1598 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from loguru import logger
+
+__all__ = [
+    "WanVAE",
+]
+
+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)
+
+
+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 // 2, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                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()
+                    # 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)
+
+                    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
+        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,1,1] = init_matrix * 0.5
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix  # * 0.5
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        # init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
+        conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix
+        conv.weight.data.copy_(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
+
+
+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(3, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            for _ in range(num_res_blocks):
+                downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = "downsample3d" if temperal_downsample[i] else "downsample2d"
+                downsamples.append(Resample(out_dim, mode=mode))
+                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 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 = 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 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
+
+
+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:])):
+            # residual (+attention) blocks
+            if i == 1 or i == 2 or i == 3:
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # upsample block
+            if i != len(dim_mult) - 1:
+                mode = "upsample3d" if temperal_upsample[i] else "upsample2d"
+                upsamples.append(Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 3, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        ## conv1
+        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 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 = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(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)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if 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 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
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if isinstance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class WanVAE_(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
+        self.temperal_upsample = temperal_downsample[::-1]
+        self.spatial_compression_ratio = 2 ** len(self.temperal_downsample)
+
+        # The minimal tile height and width for spatial tiling to be used
+        self.tile_sample_min_height = 256
+        self.tile_sample_min_width = 256
+
+        # The minimal distance between two spatial tiles
+        self.tile_sample_stride_height = 192
+        self.tile_sample_stride_width = 192
+        # 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(
+            dim,
+            z_dim,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_upsample,
+            dropout,
+        )
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        x_recon = self.decode(z)
+        return x_recon, mu, log_var
+
+    def blend_v(self, a, b, blend_extent):
+        blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (
+                1 - y / blend_extent
+            ) + b[:, :, :, y, :] * (y / blend_extent)
+        return b
+
+    def blend_h(self, a, b, blend_extent):
+        blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (
+                1 - x / blend_extent
+            ) + b[:, :, :, :, x] * (x / blend_extent)
+        return b
+
+    def tiled_encode(self, x, scale):
+        _, _, num_frames, height, width = x.shape
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        # Split x into overlapping tiles and encode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            row = []
+            for j in range(0, width, self.tile_sample_stride_width):
+                self.clear_cache()
+                time = []
+                frame_range = 1 + (num_frames - 1) // 4
+                for k in range(frame_range):
+                    self._enc_conv_idx = [0]
+                    if k == 0:
+                        tile = x[
+                            :,
+                            :,
+                            :1,
+                            i : i + self.tile_sample_min_height,
+                            j : j + self.tile_sample_min_width,
+                        ]
+                    else:
+                        tile = x[
+                            :,
+                            :,
+                            1 + 4 * (k - 1) : 1 + 4 * k,
+                            i : i + self.tile_sample_min_height,
+                            j : j + self.tile_sample_min_width,
+                        ]
+                    tile = self.encoder(
+                        tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx
+                    )
+                    mu, log_var = self.conv1(tile).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]
+
+                    time.append(mu)
+
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
+        self.clear_cache()
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(
+                    tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
+        return enc
+
+    def tiled_decode(self, z, 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]
+
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = (
+            self.tile_sample_min_height // self.spatial_compression_ratio
+        )
+        tile_latent_min_width = (
+            self.tile_sample_min_width // self.spatial_compression_ratio
+        )
+        tile_latent_stride_height = (
+            self.tile_sample_stride_height // self.spatial_compression_ratio
+        )
+        tile_latent_stride_width = (
+            self.tile_sample_stride_width // self.spatial_compression_ratio
+        )
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        # Split z into overlapping tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, height, tile_latent_stride_height):
+            row = []
+            for j in range(0, width, tile_latent_stride_width):
+                self.clear_cache()
+                time = []
+                for k in range(num_frames):
+                    self._conv_idx = [0]
+                    tile = z[
+                        :,
+                        :,
+                        k : k + 1,
+                        i : i + tile_latent_min_height,
+                        j : j + tile_latent_min_width,
+                    ]
+                    tile = self.conv2(tile)
+                    decoded = self.decoder(
+                        tile, feat_cache=self._feat_map, feat_idx=self._conv_idx
+                    )
+                    time.append(decoded)
+                row.append(torch.cat(time, dim=2))
+            rows.append(row)
+        self.clear_cache()
+
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_width)
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        :,
+                        : self.tile_sample_stride_height,
+                        : self.tile_sample_stride_width,
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
+
+        return dec
+
+    def encode(self, x, scale, return_mu=False):
+        self.clear_cache()
+        ## cache
+        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]
+
+        self.clear_cache()
+        if return_mu:
+            return mu, log_var
+        else:
+            return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+
+        # z: [b,c,t,h,w]
+        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,
+                )
+            else:
+                out_ = self.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+
+        self.clear_cache()
+        return out
+
+    def decode_stream(self, z, scale):
+        self.clear_cache()
+
+        # z: [b,c,t,h,w]
+        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]
+            out = self.decoder(
+                x[:, :, i : i + 1, :, :],
+                feat_cache=self._feat_map,
+                feat_idx=self._conv_idx,
+            )
+            yield out
+
+    def cached_decode(self, z, scale):
+        # z: [b,c,t,h,w]
+        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,
+                )
+            else:
+                out_ = self.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        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, scale=[0, 1]):
+        mu, log_var = self.encode(imgs, scale, return_mu=True)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std), mu, log_var
+
+    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 encode_video(self, x, scale=[0, 1]):
+        assert x.ndim == 5  # NTCHW
+        assert x.shape[2] % 3 == 0
+        x = x.transpose(1, 2)
+        y = x.mul(2).sub_(1)
+        y, mu, log_var = self.sample(y, scale=scale)
+        return y.transpose(1, 2).to(x), mu, log_var
+
+    def decode_video(self, x, scale=[0, 1]):
+        assert x.ndim == 5  # NTCHW
+        assert x.shape[2] % self.z_dim == 0
+        x = x.transpose(1, 2)
+        # B, C, T, H, W
+        y = x
+        y = self.decode(y, scale).clamp_(-1, 1)
+        y = y.mul_(0.5).add_(0.5).clamp_(0, 1)  # NCTHW
+        return y.transpose(1, 2).to(x)
+
+
+def _video_vae(
+    pretrained_path=None,
+    z_dim=None,
+    device="cpu",
+    dtype=torch.float,
+    **kwargs,
+):
+    """
+    Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.
+    """
+    # params
+    cfg = dict(
+        dim=96,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[False, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    # init model
+    with torch.device("meta"):
+        model = WanVAE_(**cfg)
+
+    # load checkpoint
+    model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+
+class WanVAE:
+    def __init__(
+        self,
+        z_dim=16,
+        vae_path="cache/vae_step_411000.pth",
+        dtype=torch.float,
+        device="cuda",
+        parallel=False,
+        use_tiling=False,
+        use_2d_split=True,
+    ):
+        self.dtype = dtype
+        self.device = device
+        self.parallel = parallel
+        self.use_tiling = use_tiling
+        self.use_2d_split = use_2d_split
+
+        mean = [
+            -0.7571,
+            -0.7089,
+            -0.9113,
+            0.1075,
+            -0.1745,
+            0.9653,
+            -0.1517,
+            1.5508,
+            0.4134,
+            -0.0715,
+            0.5517,
+            -0.3632,
+            -0.1922,
+            -0.9497,
+            0.2503,
+            -0.2921,
+        ]
+        std = [
+            2.8184,
+            1.4541,
+            2.3275,
+            2.6558,
+            1.2196,
+            1.7708,
+            2.6052,
+            2.0743,
+            3.2687,
+            2.1526,
+            2.8652,
+            1.5579,
+            1.6382,
+            1.1253,
+            2.8251,
+            1.9160,
+        ]
+        self.mean = torch.tensor(mean, dtype=dtype, device=device)
+        self.inv_std = 1.0 / torch.tensor(std, dtype=dtype, device=device)
+        self.scale = [self.mean, self.inv_std]
+
+        # (height, width, world_size) -> (world_size_h, world_size_w)
+        self.grid_table = {
+            # world_size = 2
+            (60, 104, 2): (1, 2),
+            (68, 120, 2): (1, 2),
+            (90, 160, 2): (1, 2),
+            (60, 60, 2): (1, 2),
+            (72, 72, 2): (1, 2),
+            (88, 88, 2): (1, 2),
+            (120, 120, 2): (1, 2),
+            (104, 60, 2): (2, 1),
+            (120, 68, 2): (2, 1),
+            (160, 90, 2): (2, 1),
+            # world_size = 4
+            (60, 104, 4): (2, 2),
+            (68, 120, 4): (2, 2),
+            (90, 160, 4): (2, 2),
+            (60, 60, 4): (2, 2),
+            (72, 72, 4): (2, 2),
+            (88, 88, 4): (2, 2),
+            (120, 120, 4): (2, 2),
+            (104, 60, 4): (2, 2),
+            (120, 68, 4): (2, 2),
+            (160, 90, 4): (2, 2),
+            # world_size = 8
+            (60, 104, 8): (2, 4),
+            (68, 120, 8): (2, 4),
+            (90, 160, 8): (2, 4),
+            (60, 60, 8): (2, 4),
+            (72, 72, 8): (2, 4),
+            (88, 88, 8): (2, 4),
+            (120, 120, 8): (2, 4),
+            (104, 60, 8): (4, 2),
+            (120, 68, 8): (4, 2),
+            (160, 90, 8): (4, 2),
+        }
+
+        # init model
+        self.model = (
+            _video_vae(
+                pretrained_path=vae_path,
+                z_dim=z_dim,
+                dtype=dtype,
+            )
+            .eval()
+            .requires_grad_(False)
+            .to(device)
+            .to(dtype)
+        )
+
+    def _calculate_2d_grid(self, latent_height, latent_width, world_size):
+        if (latent_height, latent_width, world_size) in self.grid_table:
+            best_h, best_w = self.grid_table[(latent_height, latent_width, world_size)]
+            # logger.info(f"Vae using cached 2D grid: {best_h}x{best_w} grid for {latent_height}x{latent_width} latent")
+            return best_h, best_w
+
+        best_h, best_w = 1, world_size
+        min_aspect_diff = float("inf")
+
+        for h in range(1, world_size + 1):
+            if world_size % h == 0:
+                w = world_size // h
+                if latent_height % h == 0 and latent_width % w == 0:
+                    # Calculate how close this grid is to square
+                    aspect_diff = abs((latent_height / h) - (latent_width / w))
+                    if aspect_diff < min_aspect_diff:
+                        min_aspect_diff = aspect_diff
+                        best_h, best_w = h, w
+        # logger.info(f"Vae using 2D grid & Update cache: {best_h}x{best_w} grid for {latent_height}x{latent_width} latent")
+        self.grid_table[(latent_height, latent_width, world_size)] = (best_h, best_w)
+        return best_h, best_w
+
+    def current_device(self):
+        return next(self.model.parameters()).device
+
+    def encode_dist(self, video, world_size, cur_rank, split_dim):
+        spatial_ratio = 8
+
+        if split_dim == 3:
+            total_latent_len = video.shape[3] // spatial_ratio
+        elif split_dim == 4:
+            total_latent_len = video.shape[4] // spatial_ratio
+        else:
+            raise ValueError(f"Unsupported split_dim: {split_dim}")
+
+        splited_chunk_len = total_latent_len // world_size
+        padding_size = 1
+
+        video_chunk_len = splited_chunk_len * spatial_ratio
+        video_padding_len = padding_size * spatial_ratio
+
+        if cur_rank == 0:
+            if split_dim == 3:
+                video_chunk = video[
+                    :, :, :, : video_chunk_len + 2 * video_padding_len, :
+                ].contiguous()
+            elif split_dim == 4:
+                video_chunk = video[
+                    :, :, :, :, : video_chunk_len + 2 * video_padding_len
+                ].contiguous()
+        elif cur_rank == world_size - 1:
+            if split_dim == 3:
+                video_chunk = video[
+                    :, :, :, -(video_chunk_len + 2 * video_padding_len) :, :
+                ].contiguous()
+            elif split_dim == 4:
+                video_chunk = video[
+                    :, :, :, :, -(video_chunk_len + 2 * video_padding_len) :
+                ].contiguous()
+        else:
+            start_idx = cur_rank * video_chunk_len - video_padding_len
+            end_idx = (cur_rank + 1) * video_chunk_len + video_padding_len
+            if split_dim == 3:
+                video_chunk = video[:, :, :, start_idx:end_idx, :].contiguous()
+            elif split_dim == 4:
+                video_chunk = video[:, :, :, :, start_idx:end_idx].contiguous()
+
+        if self.use_tiling:
+            encoded_chunk = self.model.tiled_encode(video_chunk, self.scale)
+        else:
+            encoded_chunk = self.model.encode(video_chunk, self.scale)
+
+        if cur_rank == 0:
+            if split_dim == 3:
+                encoded_chunk = encoded_chunk[
+                    :, :, :, :splited_chunk_len, :
+                ].contiguous()
+            elif split_dim == 4:
+                encoded_chunk = encoded_chunk[
+                    :, :, :, :, :splited_chunk_len
+                ].contiguous()
+        elif cur_rank == world_size - 1:
+            if split_dim == 3:
+                encoded_chunk = encoded_chunk[
+                    :, :, :, -splited_chunk_len:, :
+                ].contiguous()
+            elif split_dim == 4:
+                encoded_chunk = encoded_chunk[
+                    :, :, :, :, -splited_chunk_len:
+                ].contiguous()
+        else:
+            if split_dim == 3:
+                encoded_chunk = encoded_chunk[
+                    :, :, :, padding_size:-padding_size, :
+                ].contiguous()
+            elif split_dim == 4:
+                encoded_chunk = encoded_chunk[
+                    :, :, :, :, padding_size:-padding_size
+                ].contiguous()
+
+        full_encoded = [torch.empty_like(encoded_chunk) for _ in range(world_size)]
+        dist.all_gather(full_encoded, encoded_chunk)
+
+        torch.cuda.synchronize()
+
+        encoded = torch.cat(full_encoded, dim=split_dim)
+
+        return encoded.squeeze(0)
+
+    def encode_dist_2d(self, video, world_size_h, world_size_w, cur_rank_h, cur_rank_w):
+        spatial_ratio = 8
+
+        # Calculate chunk sizes for both dimensions
+        total_latent_h = video.shape[3] // spatial_ratio
+        total_latent_w = video.shape[4] // spatial_ratio
+
+        chunk_h = total_latent_h // world_size_h
+        chunk_w = total_latent_w // world_size_w
+
+        padding_size = 1
+        video_chunk_h = chunk_h * spatial_ratio
+        video_chunk_w = chunk_w * spatial_ratio
+        video_padding_h = padding_size * spatial_ratio
+        video_padding_w = padding_size * spatial_ratio
+
+        # Calculate H dimension slice
+        if cur_rank_h == 0:
+            h_start = 0
+            h_end = video_chunk_h + 2 * video_padding_h
+        elif cur_rank_h == world_size_h - 1:
+            h_start = video.shape[3] - (video_chunk_h + 2 * video_padding_h)
+            h_end = video.shape[3]
+        else:
+            h_start = cur_rank_h * video_chunk_h - video_padding_h
+            h_end = (cur_rank_h + 1) * video_chunk_h + video_padding_h
+
+        # Calculate W dimension slice
+        if cur_rank_w == 0:
+            w_start = 0
+            w_end = video_chunk_w + 2 * video_padding_w
+        elif cur_rank_w == world_size_w - 1:
+            w_start = video.shape[4] - (video_chunk_w + 2 * video_padding_w)
+            w_end = video.shape[4]
+        else:
+            w_start = cur_rank_w * video_chunk_w - video_padding_w
+            w_end = (cur_rank_w + 1) * video_chunk_w + video_padding_w
+
+        # Extract the video chunk for this process
+        video_chunk = video[:, :, :, h_start:h_end, w_start:w_end].contiguous()
+
+        # Encode the chunk
+        if self.use_tiling:
+            encoded_chunk = self.model.tiled_encode(video_chunk, self.scale)
+        else:
+            encoded_chunk = self.model.encode(video_chunk, self.scale)
+
+        # Remove padding from encoded chunk
+        if cur_rank_h == 0:
+            encoded_h_start = 0
+            encoded_h_end = chunk_h
+        elif cur_rank_h == world_size_h - 1:
+            encoded_h_start = encoded_chunk.shape[3] - chunk_h
+            encoded_h_end = encoded_chunk.shape[3]
+        else:
+            encoded_h_start = padding_size
+            encoded_h_end = encoded_chunk.shape[3] - padding_size
+
+        if cur_rank_w == 0:
+            encoded_w_start = 0
+            encoded_w_end = chunk_w
+        elif cur_rank_w == world_size_w - 1:
+            encoded_w_start = encoded_chunk.shape[4] - chunk_w
+            encoded_w_end = encoded_chunk.shape[4]
+        else:
+            encoded_w_start = padding_size
+            encoded_w_end = encoded_chunk.shape[4] - padding_size
+
+        encoded_chunk = encoded_chunk[
+            :, :, :, encoded_h_start:encoded_h_end, encoded_w_start:encoded_w_end
+        ].contiguous()
+
+        # Gather all chunks
+        total_processes = world_size_h * world_size_w
+        full_encoded = [torch.empty_like(encoded_chunk) for _ in range(total_processes)]
+
+        dist.all_gather(full_encoded, encoded_chunk)
+
+        torch.cuda.synchronize()
+
+        # Reconstruct the full encoded tensor
+        encoded_rows = []
+        for h_idx in range(world_size_h):
+            encoded_cols = []
+            for w_idx in range(world_size_w):
+                process_idx = h_idx * world_size_w + w_idx
+                encoded_cols.append(full_encoded[process_idx])
+            encoded_rows.append(torch.cat(encoded_cols, dim=4))
+
+        encoded = torch.cat(encoded_rows, dim=3)
+
+        return encoded.squeeze(0)
+
+    def encode(self, video, world_size_h=None, world_size_w=None):
+        """
+        video: one video  with shape [1, C, T, H, W].
+        """
+        if self.parallel:
+            world_size = dist.get_world_size()
+            cur_rank = dist.get_rank()
+            height, width = video.shape[3], video.shape[4]
+
+            if self.use_2d_split:
+                if world_size_h is None or world_size_w is None:
+                    world_size_h, world_size_w = self._calculate_2d_grid(
+                        height // 8, width // 8, world_size
+                    )
+                cur_rank_h = cur_rank // world_size_w
+                cur_rank_w = cur_rank % world_size_w
+                out = self.encode_dist_2d(
+                    video, world_size_h, world_size_w, cur_rank_h, cur_rank_w
+                )
+            else:
+                # Original 1D splitting logic
+                if width % world_size == 0:
+                    out = self.encode_dist(video, world_size, cur_rank, split_dim=4)
+                elif height % world_size == 0:
+                    out = self.encode_dist(video, world_size, cur_rank, split_dim=3)
+                else:
+                    logger.info("Fall back to naive encode mode")
+                    if self.use_tiling:
+                        out = self.model.tiled_encode(video, self.scale).squeeze(0)
+                    else:
+                        out = self.model.encode(video, self.scale).squeeze(0)
+        else:
+            if self.use_tiling:
+                out = self.model.tiled_encode(video, self.scale).squeeze(0)
+            else:
+                out = self.model.encode(video, self.scale).squeeze(0)
+
+        return out
+
+    def decode_dist(self, zs, world_size, cur_rank, split_dim):
+        splited_total_len = zs.shape[split_dim]
+        splited_chunk_len = splited_total_len // world_size
+        padding_size = 1
+
+        if cur_rank == 0:
+            if split_dim == 2:
+                zs = zs[:, :, : splited_chunk_len + 2 * padding_size, :].contiguous()
+            elif split_dim == 3:
+                zs = zs[:, :, :, : splited_chunk_len + 2 * padding_size].contiguous()
+        elif cur_rank == world_size - 1:
+            if split_dim == 2:
+                zs = zs[:, :, -(splited_chunk_len + 2 * padding_size) :, :].contiguous()
+            elif split_dim == 3:
+                zs = zs[:, :, :, -(splited_chunk_len + 2 * padding_size) :].contiguous()
+        else:
+            if split_dim == 2:
+                zs = zs[
+                    :,
+                    :,
+                    cur_rank * splited_chunk_len - padding_size : (cur_rank + 1)
+                    * splited_chunk_len
+                    + padding_size,
+                    :,
+                ].contiguous()
+            elif split_dim == 3:
+                zs = zs[
+                    :,
+                    :,
+                    :,
+                    cur_rank * splited_chunk_len - padding_size : (cur_rank + 1)
+                    * splited_chunk_len
+                    + padding_size,
+                ].contiguous()
+
+        decode_func = self.model.tiled_decode if self.use_tiling else self.model.decode
+        images = decode_func(zs.unsqueeze(0), self.scale).clamp_(-1, 1)
+
+        if cur_rank == 0:
+            if split_dim == 2:
+                images = images[:, :, :, : splited_chunk_len * 8, :].contiguous()
+            elif split_dim == 3:
+                images = images[:, :, :, :, : splited_chunk_len * 8].contiguous()
+        elif cur_rank == world_size - 1:
+            if split_dim == 2:
+                images = images[:, :, :, -splited_chunk_len * 8 :, :].contiguous()
+            elif split_dim == 3:
+                images = images[:, :, :, :, -splited_chunk_len * 8 :].contiguous()
+        else:
+            if split_dim == 2:
+                images = images[
+                    :, :, :, 8 * padding_size : -8 * padding_size, :
+                ].contiguous()
+            elif split_dim == 3:
+                images = images[
+                    :, :, :, :, 8 * padding_size : -8 * padding_size
+                ].contiguous()
+
+        full_images = [torch.empty_like(images) for _ in range(world_size)]
+        dist.all_gather(full_images, images)
+
+        torch.cuda.synchronize()
+
+        images = torch.cat(full_images, dim=split_dim + 1)
+
+        return images
+
+    def decode_dist_2d(self, zs, world_size_h, world_size_w, cur_rank_h, cur_rank_w):
+        total_h = zs.shape[2]
+        total_w = zs.shape[3]
+
+        chunk_h = total_h // world_size_h
+        chunk_w = total_w // world_size_w
+
+        padding_size = 2
+
+        # Calculate H dimension slice
+        if cur_rank_h == 0:
+            h_start = 0
+            h_end = chunk_h + 2 * padding_size
+        elif cur_rank_h == world_size_h - 1:
+            h_start = total_h - (chunk_h + 2 * padding_size)
+            h_end = total_h
+        else:
+            h_start = cur_rank_h * chunk_h - padding_size
+            h_end = (cur_rank_h + 1) * chunk_h + padding_size
+
+        # Calculate W dimension slice
+        if cur_rank_w == 0:
+            w_start = 0
+            w_end = chunk_w + 2 * padding_size
+        elif cur_rank_w == world_size_w - 1:
+            w_start = total_w - (chunk_w + 2 * padding_size)
+            w_end = total_w
+        else:
+            w_start = cur_rank_w * chunk_w - padding_size
+            w_end = (cur_rank_w + 1) * chunk_w + padding_size
+
+        # Extract the latent chunk for this process
+        zs_chunk = zs[:, :, h_start:h_end, w_start:w_end].contiguous()
+
+        # Decode the chunk
+        decode_func = self.model.tiled_decode if self.use_tiling else self.model.decode
+        images_chunk = decode_func(zs_chunk.unsqueeze(0), self.scale).clamp_(-1, 1)
+
+        # Remove padding from decoded chunk
+        spatial_ratio = 8
+        if cur_rank_h == 0:
+            decoded_h_start = 0
+            decoded_h_end = chunk_h * spatial_ratio
+        elif cur_rank_h == world_size_h - 1:
+            decoded_h_start = images_chunk.shape[3] - chunk_h * spatial_ratio
+            decoded_h_end = images_chunk.shape[3]
+        else:
+            decoded_h_start = padding_size * spatial_ratio
+            decoded_h_end = images_chunk.shape[3] - padding_size * spatial_ratio
+
+        if cur_rank_w == 0:
+            decoded_w_start = 0
+            decoded_w_end = chunk_w * spatial_ratio
+        elif cur_rank_w == world_size_w - 1:
+            decoded_w_start = images_chunk.shape[4] - chunk_w * spatial_ratio
+            decoded_w_end = images_chunk.shape[4]
+        else:
+            decoded_w_start = padding_size * spatial_ratio
+            decoded_w_end = images_chunk.shape[4] - padding_size * spatial_ratio
+
+        images_chunk = images_chunk[
+            :, :, :, decoded_h_start:decoded_h_end, decoded_w_start:decoded_w_end
+        ].contiguous()
+
+        # Gather all chunks
+        total_processes = world_size_h * world_size_w
+        full_images = [torch.empty_like(images_chunk) for _ in range(total_processes)]
+
+        dist.all_gather(full_images, images_chunk)
+
+        torch.cuda.synchronize()
+
+        # Reconstruct the full image tensor
+        image_rows = []
+        for h_idx in range(world_size_h):
+            image_cols = []
+            for w_idx in range(world_size_w):
+                process_idx = h_idx * world_size_w + w_idx
+                image_cols.append(full_images[process_idx])
+            image_rows.append(torch.cat(image_cols, dim=4))
+
+        images = torch.cat(image_rows, dim=3)
+
+        return images
+
+    def decode_dist_2d_stream(
+        self, zs, world_size_h, world_size_w, cur_rank_h, cur_rank_w
+    ):
+        total_h = zs.shape[2]
+        total_w = zs.shape[3]
+
+        chunk_h = total_h // world_size_h
+        chunk_w = total_w // world_size_w
+
+        padding_size = 2
+
+        # Calculate H dimension slice
+        if cur_rank_h == 0:
+            h_start = 0
+            h_end = chunk_h + 2 * padding_size
+        elif cur_rank_h == world_size_h - 1:
+            h_start = total_h - (chunk_h + 2 * padding_size)
+            h_end = total_h
+        else:
+            h_start = cur_rank_h * chunk_h - padding_size
+            h_end = (cur_rank_h + 1) * chunk_h + padding_size
+
+        # Calculate W dimension slice
+        if cur_rank_w == 0:
+            w_start = 0
+            w_end = chunk_w + 2 * padding_size
+        elif cur_rank_w == world_size_w - 1:
+            w_start = total_w - (chunk_w + 2 * padding_size)
+            w_end = total_w
+        else:
+            w_start = cur_rank_w * chunk_w - padding_size
+            w_end = (cur_rank_w + 1) * chunk_w + padding_size
+
+        # Extract the latent chunk for this process
+        zs_chunk = zs[:, :, h_start:h_end, w_start:w_end].contiguous()
+
+        for image in self.model.decode_stream(zs_chunk.unsqueeze(0), self.scale):
+            images_chunk = image.clamp_(-1, 1)
+            # Remove padding from decoded chunk
+            spatial_ratio = 8
+            if cur_rank_h == 0:
+                decoded_h_start = 0
+                decoded_h_end = chunk_h * spatial_ratio
+            elif cur_rank_h == world_size_h - 1:
+                decoded_h_start = images_chunk.shape[3] - chunk_h * spatial_ratio
+                decoded_h_end = images_chunk.shape[3]
+            else:
+                decoded_h_start = padding_size * spatial_ratio
+                decoded_h_end = images_chunk.shape[3] - padding_size * spatial_ratio
+
+            if cur_rank_w == 0:
+                decoded_w_start = 0
+                decoded_w_end = chunk_w * spatial_ratio
+            elif cur_rank_w == world_size_w - 1:
+                decoded_w_start = images_chunk.shape[4] - chunk_w * spatial_ratio
+                decoded_w_end = images_chunk.shape[4]
+            else:
+                decoded_w_start = padding_size * spatial_ratio
+                decoded_w_end = images_chunk.shape[4] - padding_size * spatial_ratio
+
+            images_chunk = images_chunk[
+                :, :, :, decoded_h_start:decoded_h_end, decoded_w_start:decoded_w_end
+            ].contiguous()
+
+            # Gather all chunks
+            total_processes = world_size_h * world_size_w
+            full_images = [
+                torch.empty_like(images_chunk) for _ in range(total_processes)
+            ]
+
+            dist.all_gather(full_images, images_chunk)
+
+            torch.cuda.synchronize()
+
+            # Reconstruct the full image tensor
+            image_rows = []
+            for h_idx in range(world_size_h):
+                image_cols = []
+                for w_idx in range(world_size_w):
+                    process_idx = h_idx * world_size_w + w_idx
+                    image_cols.append(full_images[process_idx])
+                image_rows.append(torch.cat(image_cols, dim=4))
+
+            images = torch.cat(image_rows, dim=3)
+
+            yield images
+
+    def decode(self, zs):
+        if self.parallel:
+            world_size = dist.get_world_size()
+            cur_rank = dist.get_rank()
+            latent_height, latent_width = zs.shape[2], zs.shape[3]
+
+            if self.use_2d_split:
+                world_size_h, world_size_w = self._calculate_2d_grid(
+                    latent_height, latent_width, world_size
+                )
+                cur_rank_h = cur_rank // world_size_w
+                cur_rank_w = cur_rank % world_size_w
+                images = self.decode_dist_2d(
+                    zs, world_size_h, world_size_w, cur_rank_h, cur_rank_w
+                )
+            else:
+                # Original 1D splitting logic
+                if latent_width % world_size == 0:
+                    images = self.decode_dist(zs, world_size, cur_rank, split_dim=3)
+                elif latent_height % world_size == 0:
+                    images = self.decode_dist(zs, world_size, cur_rank, split_dim=2)
+                else:
+                    logger.info("Fall back to naive decode mode")
+                    images = self.model.decode(zs.unsqueeze(0), self.scale).clamp_(
+                        -1, 1
+                    )
+        else:
+            decode_func = (
+                self.model.tiled_decode if self.use_tiling else self.model.decode
+            )
+            images = decode_func(zs.unsqueeze(0), self.scale).clamp_(-1, 1)
+
+        return images
+
+    def decode_stream(self, zs):
+        if self.parallel:
+            world_size = dist.get_world_size()
+            cur_rank = dist.get_rank()
+            latent_height, latent_width = zs.shape[2], zs.shape[3]
+
+            world_size_h, world_size_w = self._calculate_2d_grid(
+                latent_height, latent_width, world_size
+            )
+            cur_rank_h = cur_rank // world_size_w
+            cur_rank_w = cur_rank % world_size_w
+            for images in self.decode_dist_2d_stream(
+                zs, world_size_h, world_size_w, cur_rank_h, cur_rank_w
+            ):
+                yield images
+        else:
+            for image in self.model.decode_stream(zs.unsqueeze(0), self.scale):
+                yield image.clamp_(-1, 1)
+
+    def encode_video(self, vid):
+        return self.model.encode_video(vid)
+
+    def decode_video(self, vid_enc):
+        return self.model.decode_video(vid_enc)

leaptalk_inference.py

@@ -0,0 +1,1156 @@
+
+
+
+
+import argparse
+import contextlib
+import glob
+import os
+import subprocess
+import sys
+import time
+from dataclasses import dataclass
+
+import imageio
+import librosa
+import numpy as np
+import torch
+import torch.distributed as dist
+from loguru import logger
+from peft import PeftModel
+
+
+# `flash_head` and `vibt` are vendored next to this file inside the Space repo.
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+VIBT_ROOT = THIS_DIR
+if THIS_DIR not in sys.path:
+    sys.path.insert(0, THIS_DIR)
+
+from flash_head.src.pipeline.flash_head_pipeline import FlashHeadPipeline, timestep_transform  # noqa: E402
+from flash_head.utils.utils import match_and_blend_colors_torch  # noqa: E402
+from vibt.scheduler import ViBTScheduler  # noqa: E402
+
+
+def _build_infer_timesteps(
+    *,
+    step_list: list[int] | None,
+    num_inference_steps: int,
+    shift_gamma: float,
+    device: str,
+    num_timesteps: int = 1000,
+) -> torch.Tensor:
+    """Return 1D float tensor of timesteps for ViBTScheduler.
+
+    Mirrors training-style warping with `timestep_transform(shift=shift_gamma)`.
+    Does not include a terminal 0 step to avoid an extra Euler update at t=0.
+    """
+
+    if step_list is not None and len(step_list) > 0:
+        base_steps = [int(x) for x in step_list if int(x) > 0]
+        if not base_steps:
+            raise ValueError("--step_list must contain at least one positive timestep")
+        is_strict_desc = all(base_steps[i] > base_steps[i + 1] for i in range(len(base_steps) - 1))
+        if not is_strict_desc:
+            base_steps = sorted(set(base_steps), reverse=True)
+            logger.warning(f"step_list is not strictly descending; using sorted unique list: {base_steps}")
+    else:
+        n = int(num_inference_steps)
+        if n <= 0:
+            raise ValueError("--num_inference_steps must be > 0")
+        if n == 2:
+            base_steps = [1000, 500]
+        elif n == 4:
+            base_steps = [1000, 750, 500, 250]
+        else:
+            base_steps = list(np.linspace(num_timesteps, 1, n, dtype=np.float32))
+            base_steps = [int(round(float(x))) for x in base_steps]
+            base_steps = [x for x in base_steps if x > 0]
+            base_steps = sorted(set(base_steps), reverse=True)
+            if not base_steps:
+                raise ValueError("Derived empty timestep list; check --num_inference_steps")
+    logger.info(f"Using timesteps: {base_steps} (derived from num_inference_steps={num_inference_steps})")
+    ts = [torch.tensor([float(t)], device=device) for t in base_steps]
+    ts = [timestep_transform(t, shift=float(shift_gamma), num_timesteps=num_timesteps) for t in ts]
+    values = [t.reshape(-1)[0].to(device=device, dtype=torch.float32) for t in ts]
+    return torch.stack(values, dim=0)
+
+
+def _pick_latest_audio_proj(path_or_dir: str | None) -> str | None:
+    if not path_or_dir:
+        return None
+    if os.path.isfile(path_or_dir):
+        return path_or_dir
+    candidates = glob.glob(os.path.join(path_or_dir, "audio_proj_step_*.pt"))
+    if not candidates:
+        return None
+
+    def _step_num(p: str) -> int:
+        base = os.path.basename(p)
+        try:
+            return int(base.split("audio_proj_step_")[-1].split(".pt")[0])
+        except Exception:
+            return -1
+
+    candidates.sort(key=_step_num)
+    return candidates[-1]
+
+
+def _pick_latest_lora_dir(path_or_dir: str | None) -> str | None:
+    if not path_or_dir:
+        return None
+    if os.path.isfile(os.path.join(path_or_dir, "adapter_config.json")):
+        return path_or_dir
+
+    release_lora_dir = os.path.join(path_or_dir, "lora")
+    if os.path.isfile(os.path.join(release_lora_dir, "adapter_config.json")):
+        return release_lora_dir
+
+    candidates = glob.glob(os.path.join(path_or_dir, "lora_step_*"))
+    candidates = [p for p in candidates if os.path.isfile(os.path.join(p, "adapter_config.json"))]
+    if not candidates:
+        return None
+
+    def _step_num(p: str) -> int:
+        base = os.path.basename(os.path.normpath(p))
+        try:
+            return int(base.split("lora_step_")[-1])
+        except Exception:
+            return -1
+
+    candidates.sort(key=_step_num)
+    return candidates[-1]
+
+
+def _resolve_lite_tae_path(
+    *,
+    lora_dir: str,
+    resolved_lora_dir: str,
+    ckpt_dir: str,
+    explicit_tae_path: str | None,
+) -> str:
+    if explicit_tae_path:
+        if os.path.isfile(explicit_tae_path):
+            return explicit_tae_path
+        raise SystemExit(f"TAE checkpoint not found: {explicit_tae_path}")
+
+    search_roots = [
+        lora_dir,
+        resolved_lora_dir,
+        os.path.dirname(os.path.abspath(resolved_lora_dir)),
+        os.path.join(ckpt_dir, "VAE_Wan"),
+    ]
+    seen = set()
+    for root in search_roots:
+        if not root:
+            continue
+        root = os.path.abspath(root)
+        if root in seen:
+            continue
+        seen.add(root)
+        for name in ("taew2_1.pth", "taew2_1.safetensors", "taew2_2.pth", "taew2_2.safetensors"):
+            candidate = os.path.join(root, name)
+            if os.path.isfile(candidate):
+                return candidate
+
+    raise SystemExit(
+        "Lite mode requires a TAE checkpoint. Expected taew2_1.pth under --lora_dir "
+        "or its parent directory after downloading https://huggingface.co/z-rx/leaptalk."
+    )
+
+
+def _get_inner_flashhead_model(model):
+    # After PEFT wrapping, training code uses model.base_model.model
+    if hasattr(model, "base_model") and hasattr(model.base_model, "model"):
+        return model.base_model.model
+    return model
+
+
+def _read_num_heads_from_infer_params(vibt_root: str) -> int:
+    """Best-effort parse of flash_head/configs/infer_params.yaml without requiring PyYAML."""
+    cfg_path = os.path.join(vibt_root, "flash_head", "configs", "infer_params.yaml")
+    try:
+        with open(cfg_path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith("#"):
+                    continue
+                if line.startswith("num_heads"):
+                    # e.g. "num_heads: 12"
+                    _, value = line.split(":", 1)
+                    return int(value.strip())
+    except Exception:
+        pass
+    return 12
+
+
+def _lora_checkpoint_needs_compiled_base(lora_dir: str) -> bool:
+    """Detect whether adapter weights were saved from a torch.compile()'d base model."""
+    weights_path = os.path.join(lora_dir, "adapter_model.safetensors")
+    if not os.path.exists(weights_path):
+        return False
+
+    try:
+        from safetensors.torch import safe_open
+
+        with safe_open(weights_path, framework="pt", device="cpu") as f:
+            for k in f.keys():
+                if "._orig_mod." in k:
+                    return True
+        return False
+    except Exception:
+        # Best-effort: if we cannot inspect, do not force compile.
+        return False
+
+
+@dataclass
+class StreamParams:
+    frame_num: int = 33
+    motion_frames_latent_num: int = 2
+    tgt_fps: int = 25
+    sample_rate: int = 16000
+    cached_audio_duration: int = 8
+
+    def init_with_stride(self, vae_stride_t: int) -> "InitializedStreamParams":
+        # motion_frames_num is in *pixel frames*.
+        motion_frames_num = (self.motion_frames_latent_num - 1) * vae_stride_t + 1
+        return InitializedStreamParams(
+            frame_num=self.frame_num,
+            motion_frames_latent_num=self.motion_frames_latent_num,
+            tgt_fps=self.tgt_fps,
+            sample_rate=self.sample_rate,
+            cached_audio_duration=self.cached_audio_duration,
+            motion_frames_num=motion_frames_num,
+        )
+
+
+@dataclass
+class InitializedStreamParams:
+    frame_num: int
+    motion_frames_latent_num: int
+    tgt_fps: int
+    sample_rate: int
+    cached_audio_duration: int
+    motion_frames_num: int
+
+    @property
+    def slice_len(self) -> int:
+        return self.frame_num - self.motion_frames_num
+
+
+def _open_video_writer(tmp_mp4: str, fps: int):
+    os.makedirs(os.path.dirname(os.path.abspath(tmp_mp4)) or ".", exist_ok=True)
+    return imageio.get_writer(
+        tmp_mp4,
+        format="mp4",
+        mode="I",
+        fps=fps,
+        codec="h264",
+        ffmpeg_params=["-bf", "0"],
+    )
+
+
+@torch.no_grad()
+def _preprocess_audio_embeddings(
+    pipeline: FlashHeadPipeline,
+    audio_array: np.ndarray,
+    sr: int,
+    fps: int,
+    device: str,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Return per-frame wav2vec2 embeddings.
+
+    Output: [T, 12, 768] where T ~= len(audio_array) * fps / sr.
+    """
+    video_len = int(len(audio_array) * fps / sr)
+    if video_len <= 0:
+        raise RuntimeError("audio too short: computed video_len <= 0")
+
+    inputs = pipeline.wav2vec_feature_extractor(audio_array, sampling_rate=sr, return_tensors="pt", padding=True)
+    audio_values = inputs.input_values.to(device)
+    audio_out = pipeline.audio_encoder(audio_values, seq_len=video_len, output_hidden_states=True)
+
+    if not getattr(audio_out, "hidden_states", None):
+        raise RuntimeError("wav2vec2 returned no hidden_states")
+
+    hs = audio_out.hidden_states[-12:]
+    audio_emb = torch.stack(hs, dim=2).squeeze(0)  # [T, 12, 768]
+    return audio_emb.to(device=device, dtype=dtype)
+
+
+@torch.no_grad()
+def _audio_context_from_embeddings_range(
+    audio_emb: torch.Tensor,
+    start_idx: int,
+    end_idx: int,
+    device: str,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Build [1, F, 5, 12, 768] from audio_emb [T, 12, 768] for [start_idx, end_idx).
+
+    This mirrors the official implementation in SoulX-FlashHead (vectorized gather)
+    and avoids Python-side loops.
+    """
+
+    if audio_emb.dim() != 3:
+        raise RuntimeError(f"Expected audio_emb [T,12,768], got shape: {tuple(audio_emb.shape)}")
+
+    T = int(audio_emb.shape[0])
+    start_idx = int(start_idx)
+    end_idx = int(end_idx)
+    if end_idx <= start_idx:
+        raise RuntimeError(f"Invalid audio index range: start={start_idx}, end={end_idx}")
+
+    # Clamp to valid range.
+    start_idx = max(0, start_idx)
+    end_idx = min(T, end_idx)
+    F = end_idx - start_idx
+    if F <= 0:
+        raise RuntimeError(f"Invalid clamped range: start={start_idx}, end={end_idx}, T={T}")
+
+    # Official: indices = [-2,-1,0,1,2]
+    dev = audio_emb.device
+    base = torch.arange(start_idx, end_idx, device=dev, dtype=torch.long)  # [F]
+    offsets = torch.arange(-2, 3, device=dev, dtype=torch.long)  # [5]
+    center = base.unsqueeze(1) + offsets.unsqueeze(0)  # [F,5]
+
+    # Match official clamping behavior (max to end_idx-1).
+    center = torch.clamp(center, min=0, max=end_idx - 1)
+    ctx = audio_emb[center]  # [F,5,12,768]
+    ctx = ctx.unsqueeze(0).contiguous()  # [1,F,5,12,768]
+    return ctx.to(device=device, dtype=dtype)
+
+
+@torch.no_grad()
+def _decode_to_cthw(pipeline: FlashHeadPipeline, latents_cthw: torch.Tensor) -> torch.Tensor:
+    decoded = pipeline.vae.decode(latents_cthw)
+    if isinstance(decoded, (tuple, list)):
+        decoded = decoded[0]
+    if decoded.dim() == 5:
+        decoded = decoded[0]
+    if decoded.dim() != 4:
+        raise RuntimeError(f"Unexpected decoded video shape: {tuple(decoded.shape)}")
+
+    # Expect [C,T,H,W]
+    if decoded.shape[0] not in (1, 3, 4):
+        raise RuntimeError(f"Cannot interpret decoded layout as [C,T,H,W]: {tuple(decoded.shape)}")
+    return decoded
+
+
+@torch.no_grad()
+def _maybe_apply_color_correction(
+    pipeline: FlashHeadPipeline,
+    decoded_video_cthw: torch.Tensor,
+) -> torch.Tensor:
+    strength = float(getattr(pipeline, "color_correction_strength", 0.0))
+    if strength <= 0.0:
+        return decoded_video_cthw
+
+    reference = getattr(pipeline, "original_color_reference", None)
+    if reference is None:
+        return decoded_video_cthw
+
+    corrected = match_and_blend_colors_torch(
+        decoded_video_cthw.unsqueeze(0),
+        reference,
+        strength,
+    )
+    return corrected[0]
+
+
+@torch.no_grad()
+def _encode_motion_prefix_from_decoded(
+    pipeline: FlashHeadPipeline,
+    decoded_video_cthw: torch.Tensor,
+    motion_frames_num: int,
+    device: str,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    # decoded_video_cthw: [C,T,H,W] in [-1,1]
+    cond = decoded_video_cthw[:, -motion_frames_num:]
+    cond = cond.unsqueeze(0).to(device=device, dtype=dtype)  # [1,C,T,H,W]
+    return pipeline.vae.encode(cond).squeeze(0)
+
+
+@torch.no_grad()
+def _motion_prefix_from_latent_tail(latent_chunk_cthw: torch.Tensor, motion_frames_latent_num: int) -> torch.Tensor:
+    if motion_frames_latent_num <= 0:
+        raise RuntimeError(f"motion_frames_latent_num must be positive, got {motion_frames_latent_num}")
+    total_latent_len = int(latent_chunk_cthw.shape[1])
+    prefix_len = min(int(motion_frames_latent_num), total_latent_len)
+    return latent_chunk_cthw[:, -prefix_len:].contiguous()
+
+
+@torch.no_grad()
+def _source_suffix_from_ref_latent(ref_latent: torch.Tensor, prefix_len: int, suffix_len: int) -> torch.Tensor:
+    if suffix_len <= 0:
+        return ref_latent[:, 0:0, :, :].contiguous()
+    return ref_latent[:, prefix_len:prefix_len + suffix_len, :, :].contiguous()
+
+
+@torch.no_grad()
+def _bridge_sample_one_chunk(
+    pipeline: FlashHeadPipeline,
+    scheduler: ViBTScheduler,
+    ref_latent: torch.Tensor,
+    audio_context: torch.Tensor,
+    guidance_scale: float,
+    latent_motion_frames: torch.Tensor | None,
+    clamp_latent_len: int,
+    device: str,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Run one chunk with the same source construction used in training.
+
+    Training defines the source suffix from the fixed initial reference latent, while
+    keeping the history prefix hard-clamped for autoregressive continuity.
+    """
+    guidance_scale = float(guidance_scale)
+    y = ref_latent.unsqueeze(0).to(device=device, dtype=dtype)
+
+    if latent_motion_frames is None or clamp_latent_len <= 0:
+        x_t = ref_latent.unsqueeze(0).to(device=device, dtype=dtype).clone()
+    else:
+        history = latent_motion_frames[:, :, :clamp_latent_len].to(device=device, dtype=dtype)
+        total_latent_len = int(ref_latent.shape[1])
+        suffix_len = total_latent_len - clamp_latent_len
+        if suffix_len < 0:
+            raise RuntimeError(f"Invalid latent lengths: total={total_latent_len}, clamp={clamp_latent_len}")
+
+        if suffix_len > 0:
+            source_suffix = _source_suffix_from_ref_latent(ref_latent, clamp_latent_len, suffix_len).unsqueeze(0).to(device=device, dtype=dtype)
+            x_t = torch.cat([history, source_suffix], dim=2)
+        else:
+            x_t = history.clone()
+
+    for t in scheduler.timesteps:
+        if latent_motion_frames is not None and clamp_latent_len > 0:
+            x_t[:, :, :clamp_latent_len] = latent_motion_frames[:, :, :clamp_latent_len]
+        timestep = t.unsqueeze(0).to(device=device, dtype=x_t.dtype)
+        v_pred_cond = pipeline.model(x=x_t, timestep=timestep, context=audio_context, y=y)
+        if isinstance(v_pred_cond, (tuple, list)):
+            v_pred_cond = v_pred_cond[0]
+
+        if guidance_scale == 1.0:
+            v_pred = v_pred_cond
+        else:
+            # Audio-driven classifier-free guidance (CFG)
+            # v_final = v_uncond + s * (v_cond - v_uncond)
+            uncond_audio_context = torch.zeros_like(audio_context)
+            v_pred_uncond = pipeline.model(x=x_t, timestep=timestep, context=uncond_audio_context, y=y)
+            if isinstance(v_pred_uncond, (tuple, list)):
+                v_pred_uncond = v_pred_uncond[0]
+            v_pred = v_pred_uncond + guidance_scale * (v_pred_cond - v_pred_uncond)
+        x_t = scheduler.step(v_pred, t, x_t)[0]
+
+        if latent_motion_frames is not None and clamp_latent_len > 0:
+            x_t[:, :, :clamp_latent_len] = latent_motion_frames[:, :, :clamp_latent_len]
+
+    return x_t[0]
+
+
+def _mux_audio(tmp_mp4: str, audio_path: str, out_path: str):
+    cmd = [
+        "ffmpeg",
+        "-i",
+        tmp_mp4,
+        "-i",
+        audio_path,
+        "-c:v",
+        "copy",
+        "-c:a",
+        "mp3",
+        "-shortest",
+        out_path,
+        "-y",
+    ]
+    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Streaming inference for ViBT-finetuned FlashHead (chunked AR).")
+    parser.add_argument("--ckpt_dir", type=str, default="")
+    parser.add_argument("--wav2vec_dir", type=str, default="")
+    parser.add_argument("--lora_dir", type=str, default="")
+    parser.add_argument(
+        "--audio_proj",
+        type=str,
+        default="",
+        help="Path to audio_proj_step_*.pt (or directory containing them). If omitted, picks latest in the LeapTalk model directory.",
+    )
+
+    parser.add_argument("--cond_image", type=str, default="")
+    parser.add_argument("--audio_path", type=str, default="")
+    parser.add_argument("--out", type=str, default="outputs/leaptalk.mp4")
+
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16", "fp32"])
+    parser.add_argument(
+        "--compile",
+        type=str,
+        default="off",
+        choices=["auto", "on", "off"],
+        help="torch.compile base model: auto uses LoRA key inspection.",
+    )
+
+    parser.add_argument(
+        "--usp",
+        type=str,
+        default="on",
+        choices=["auto", "on", "off"],
+        help="Enable USP multi-GPU (same forward split across ranks). auto enables when WORLD_SIZE>1.",
+    )
+
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--fps", type=int, default=25)
+    parser.add_argument("--sample_rate", type=int, default=16000)
+
+    # Stream params (match SoulX defaults)
+    parser.add_argument("--frame_num", type=int, default=33)
+    parser.add_argument("--motion_frames_latent_num", type=int, default=2)
+    parser.add_argument("--cached_audio_duration", type=int, default=8)
+    parser.add_argument("--audio_encode_mode", type=str, default="stream", choices=["stream", "once"])
+    parser.add_argument("--max_chunks", type=int, default=0, help="0 means run all chunks.")
+    parser.add_argument(
+        "--history_update_mode",
+        type=str,
+        default="roundtrip",
+        choices=["roundtrip", "latent"],
+        help="How to update next-chunk history: SoulX-style VAE round-trip or direct latent tail reuse.",
+    )
+
+    # Bridge sampler params
+    parser.add_argument("--num_inference_steps", type=int, default=1)
+    parser.add_argument(
+        "--step_list",
+        type=int,
+        nargs="+",
+        default=None,
+        help="Explicit discrete timesteps (e.g. --step_list 1000 750 500 250). Overrides --num_inference_steps.",
+    )
+    parser.add_argument("--noise_scale", type=float, default=1.0)
+    parser.add_argument("--shift_gamma", type=float, default=5.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=1.0,
+        help="Audio CFG guidance scale; 1.0 disables guidance.",
+    )
+
+    parser.add_argument("--use_face_crop", action="store_true")
+    parser.add_argument(
+        "--color_correction_strength",
+        type=float,
+        default=1.0,
+        help="Blend decoded chunk colors toward the reference image; 0 disables correction and 1 applies full correction.",
+    )
+    parser.add_argument(
+        "--model_type",
+        type=str,
+        default="pro"
+    )
+
+    # Lite mode uses the TAE backend for fast VAE encode/decode.
+    parser.add_argument(
+        "--lite",
+        dest="lite",
+        action="store_true",
+        default=True,
+        help="Use the TAE checkpoint shipped with the LeapTalk model directory for VAE encode/decode.",
+    )
+    parser.add_argument(
+        "--no_lite",
+        dest="lite",
+        action="store_false",
+        help="Use WanVAE instead of the Lite TAE backend.",
+    )
+
+    # Backward-compatible aliases for older evaluation scripts.
+    parser.add_argument(
+        "--use_tae",
+        dest="lite",
+        action="store_true",
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--no_tae",
+        dest="lite",
+        action="store_false",
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--tae_path",
+        type=str,
+        default=None,
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--tae_model_type",
+        type=str,
+        default="wan21",
+        choices=["wan21", "wan22", "hy15"],
+        help="TAEHV model_type (affects latent_channels/patch_size/clamp behavior).",
+    )
+    args = parser.parse_args()
+
+    missing_paths = [
+        name
+        for name in ("ckpt_dir", "wav2vec_dir", "lora_dir", "cond_image", "audio_path")
+        if not getattr(args, name)
+    ]
+    if missing_paths:
+        raise SystemExit("Missing required path argument(s): " + ", ".join(f"--{name}" for name in missing_paths))
+
+    resolved_lora_dir = _pick_latest_lora_dir(args.lora_dir)
+    if not resolved_lora_dir:
+        raise SystemExit(
+            "LoRA weights not found; pass --lora_dir /path/to/leaptalk, "
+            "/path/to/leaptalk/lora, or a directory containing lora_step_*."
+        )
+
+    tae_path = _resolve_lite_tae_path(
+        lora_dir=args.lora_dir,
+        resolved_lora_dir=resolved_lora_dir,
+        ckpt_dir=args.ckpt_dir,
+        explicit_tae_path=args.tae_path,
+    ) if args.lite else None
+
+    # USP init must happen before pipeline construction.
+    world_size_env = int(os.environ.get("WORLD_SIZE", "1"))
+    use_usp = (args.usp == "on") or (args.usp == "auto" and world_size_env > 1)
+    did_init_dist = False
+    if use_usp and world_size_env > 1:
+        if not torch.cuda.is_available():
+            raise SystemExit("USP requested but CUDA is not available")
+
+        if dist.is_initialized():
+            logger.info("torch.distributed already initialized; reusing existing process group.")
+        else:
+            num_heads = _read_num_heads_from_infer_params(VIBT_ROOT)
+            from flash_head.src.distributed.usp_device import get_device, get_parallel_degree
+
+            ulysses_degree, ring_degree = get_parallel_degree(world_size_env, num_heads)
+            logger.info(
+                f"Initializing USP: WORLD_SIZE={world_size_env}, num_heads={num_heads}, "
+                f"ulysses_degree={ulysses_degree}, ring_degree={ring_degree}"
+            )
+            args.device = str(get_device(ulysses_degree, ring_degree))
+            did_init_dist = True
+    else:
+        use_usp = False
+
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    is_rank0 = rank == 0
+
+    # Under torchrun/USP, every rank executes the same script.
+    # To avoid duplicated progress logs and stdout prints, keep rank0 verbose
+    # and silence non-rank0 stdout while still surfacing errors on stderr.
+    if not is_rank0:
+        try:
+            sys.stdout = open(os.devnull, "w")
+        except Exception:
+            pass
+        try:
+            logger.remove()
+            logger.add(sys.stderr, level="ERROR")
+        except Exception:
+            pass
+
+    if use_usp and args.history_update_mode != "roundtrip":
+        logger.warning("USP enabled: forcing --history_update_mode=roundtrip to match official behavior.")
+        args.history_update_mode = "roundtrip"
+
+    if args.dtype == "bf16":
+        amp_dtype = torch.bfloat16
+    elif args.dtype == "fp16":
+        amp_dtype = torch.float16
+    else:
+        amp_dtype = torch.float32
+
+    # Decide compile mode BEFORE creating pipeline (pipeline module reads globals)
+    needs_compiled = _lora_checkpoint_needs_compiled_base(resolved_lora_dir)
+    if args.compile == "auto":
+        compile_model = bool(needs_compiled)
+    elif args.compile == "on":
+        compile_model = True
+    else:
+        compile_model = False
+
+    # If checkpoint was saved from a compiled base (keys contain '._orig_mod.'), we must compile
+    # BEFORE loading LoRA so module names match. Otherwise, compiling before loading may
+    # introduce a prefix mismatch; compile after loading is safer.
+    compile_before_lora = bool(needs_compiled)
+    compile_after_lora = bool(compile_model and not needs_compiled)
+
+    import flash_head.src.pipeline.flash_head_pipeline as fh_pipe_mod
+
+    fh_pipe_mod.COMPILE_MODEL = bool(compile_before_lora)
+    # TAEHV uses Python-level block loops; torch.compile often provides little benefit and can break.
+    fh_pipe_mod.COMPILE_VAE = bool(not args.lite)
+
+    logger.info(
+        f"Init pipeline: device={args.device}, dtype={args.dtype}, compile_model={fh_pipe_mod.COMPILE_MODEL}, "
+        f"usp={use_usp}, lite={args.lite}, rank={rank}"
+    )
+    pipeline = FlashHeadPipeline(
+        checkpoint_dir=args.ckpt_dir,
+        model_type=args.model_type,
+        wav2vec_dir=args.wav2vec_dir,
+        device=args.device,
+        param_dtype=amp_dtype,
+        use_usp=use_usp,
+        use_tae=bool(args.lite),
+        tae_path=tae_path,
+        tae_model_type=args.tae_model_type,
+    )
+    logger.info(f"Pipeline initialized. Model dtype: {args.model_type}")
+    # Load LoRA
+    logger.info(f"Loading LoRA: {resolved_lora_dir}")
+    pipeline.model = PeftModel.from_pretrained(pipeline.model, resolved_lora_dir, is_trainable=False)
+    # Merge LoRA weights into base model to remove PEFT overhead during inference
+    pipeline.model = pipeline.model.merge_and_unload()
+    pipeline.model.eval().requires_grad_(False)
+
+    if compile_after_lora:
+        logger.info("Compiling model after loading LoRA (--compile on; LoRA keys do not require compiled base).")
+        pipeline.model = torch.compile(pipeline.model)
+
+    # Load audio_proj weights
+    audio_proj_path = (
+        _pick_latest_audio_proj(args.audio_proj)
+        or _pick_latest_audio_proj(args.lora_dir)
+        or _pick_latest_audio_proj(resolved_lora_dir)
+        or _pick_latest_audio_proj(os.path.dirname(resolved_lora_dir))
+    )
+    if not audio_proj_path:
+        raise SystemExit("audio_proj weights not found; pass --audio_proj /path/to/audio_proj_step_XXXX.pt")
+
+    logger.info(f"Loading audio_proj: {audio_proj_path}")
+    audio_proj_state = torch.load(audio_proj_path, map_location="cpu")
+    inner = _get_inner_flashhead_model(pipeline.model)
+    inner.audio_proj.load_state_dict(audio_proj_state, strict=True)
+
+    # Stream params derived from VAE stride
+    vae_stride_t = int(pipeline.config.vae_stride[0])
+    sp = StreamParams(
+        frame_num=int(args.frame_num),
+        motion_frames_latent_num=int(args.motion_frames_latent_num),
+        tgt_fps=int(args.fps),
+        sample_rate=int(args.sample_rate),
+        cached_audio_duration=int(args.cached_audio_duration),
+    ).init_with_stride(vae_stride_t)
+
+    frame_num = sp.frame_num
+    motion_frames_num = sp.motion_frames_num
+    slice_len = sp.slice_len
+
+    if slice_len <= 0:
+        raise SystemExit(f"Invalid streaming params: frame_num={frame_num}, motion_frames_num={motion_frames_num}")
+
+    logger.info(f"Streaming: frame_num={frame_num}, motion_frames_num={motion_frames_num}, slice_len={slice_len}")
+    logger.info(f"History update mode: {args.history_update_mode}")
+
+    # Build X0 using pipeline preprocessing (resizes to H/W and repeats to frame_num)
+    pipeline.prepare_params(
+        cond_image_path_or_dir=args.cond_image,
+        target_size=(args.height, args.width),
+        frame_num=frame_num,
+        motion_frames_num=0,
+        sampling_steps=4,
+        seed=args.seed,
+        shift=args.shift_gamma,
+        color_correction_strength=args.color_correction_strength,
+        use_face_crop=args.use_face_crop,
+    )
+    X0 = pipeline.ref_img_latent.to(device=args.device, dtype=amp_dtype)  # [C_lat,T_lat,h,w]
+    logger.info(f"Initialized static anchor latent with shape: {tuple(X0.shape)}")
+
+    # Scheduler
+    scheduler = ViBTScheduler(num_train_timesteps=1000)
+    scheduler.timesteps = _build_infer_timesteps(
+        step_list=args.step_list,
+        num_inference_steps=args.num_inference_steps,
+        shift_gamma=args.shift_gamma,
+        device=args.device,
+        num_timesteps=1000,
+    )
+    scheduler.num_inference_steps = int(scheduler.timesteps.numel())
+    scheduler.set_parameters(noise_scale=args.noise_scale, shift_gamma=args.shift_gamma, seed=args.seed)
+
+    # Init motion prefix (match pipeline.reset_person_name): only 1 latent frame at start
+    latent_motion_frames = X0[:, :1].unsqueeze(0).clone()  # [1,C_lat,1,h,w]
+    clamp_latent_len = int(latent_motion_frames.shape[2])
+    # Audio preparation
+    audio_all, _ = librosa.load(args.audio_path, sr=args.sample_rate, mono=True)
+    human_speech_array_slice_len = slice_len * args.sample_rate // args.fps
+    human_speech_array_frame_num = frame_num * args.sample_rate // args.fps
+
+    # pad audio with silence to avoid re
+    if args.audio_encode_mode == "once":
+        remainder = (len(audio_all) - human_speech_array_frame_num) % human_speech_array_slice_len
+    else:
+        remainder = len(audio_all) % human_speech_array_slice_len
+    if remainder > 0:
+        pad_len = human_speech_array_slice_len - remainder
+        audio_all = np.concatenate([audio_all, np.zeros(pad_len, dtype=audio_all.dtype)])
+
+    # ensure audio is long enough for at least one window
+    if len(audio_all) < human_speech_array_frame_num:
+        audio_all = np.concatenate([audio_all, np.zeros(human_speech_array_frame_num - len(audio_all), dtype=audio_all.dtype)])
+
+    tmp_mp4 = args.out.replace(".mp4", "_tmp.mp4")
+    # Match official generate_video.py: cache per-chunk video tensors on CPU,
+    # then write tmp mp4 once and mux audio at the end.
+    generated_list: list[torch.Tensor] = []
+
+    total_frames = 0
+    total_time = 0.0
+    total_chunks = 0
+
+    try:
+        if args.audio_encode_mode == "once":
+            # Match official path: use pipeline.preprocess_audio + vectorized gather windowing.
+            audio_emb_all = pipeline.preprocess_audio(audio_all, sr=args.sample_rate, fps=args.fps)
+            if audio_emb_all is None:
+                raise RuntimeError("Failed to extract audio embeddings")
+            audio_emb_all = audio_emb_all.to(device=args.device, dtype=amp_dtype)
+            total_frames = int(audio_emb_all.shape[0])
+            audio_ctx_all = _audio_context_from_embeddings_range(
+                audio_emb_all,
+                start_idx=0,
+                end_idx=total_frames,
+                device=args.device,
+                dtype=amp_dtype,
+            )
+
+            # Always run at least one chunk.
+            num_chunks = 1 + max(0, (total_frames - frame_num) // slice_len)
+            if args.max_chunks and args.max_chunks > 0:
+                num_chunks = min(num_chunks, int(args.max_chunks))
+
+            for chunk_idx in range(num_chunks):
+                use_cuda_timing = args.device.startswith("cuda") and torch.cuda.is_available()
+                if use_cuda_timing:
+                    torch.cuda.synchronize()
+                chunk_start = time.perf_counter()
+                start = chunk_idx * slice_len
+                end = start + frame_num
+                if end > total_frames:
+                    # For the last partial chunk, clamp; audio windows will pad at edges.
+                    end = total_frames
+                audio_ctx = audio_ctx_all[:, start:end].contiguous()
+
+                if use_cuda_timing:
+                    denoise_evt_s = torch.cuda.Event(enable_timing=True)
+                    denoise_evt_e = torch.cuda.Event(enable_timing=True)
+                    denoise_evt_s.record()
+                x_final = _bridge_sample_one_chunk(
+                    pipeline,
+                    scheduler=scheduler,
+                    ref_latent=X0,
+                    audio_context=audio_ctx,
+                    guidance_scale=args.guidance_scale,
+                    latent_motion_frames=latent_motion_frames,
+                    clamp_latent_len=clamp_latent_len,
+                    device=args.device,
+                    dtype=amp_dtype,
+                )
+                if use_cuda_timing:
+                    denoise_evt_e.record()
+
+                if use_cuda_timing:
+                    decode_evt_s = torch.cuda.Event(enable_timing=True)
+                    decode_evt_e = torch.cuda.Event(enable_timing=True)
+                    decode_evt_s.record()
+                decoded_cthw = _decode_to_cthw(pipeline, x_final)
+                if use_cuda_timing:
+                    decode_evt_e.record()
+
+                if use_cuda_timing:
+                    color_evt_s = torch.cuda.Event(enable_timing=True)
+                    color_evt_e = torch.cuda.Event(enable_timing=True)
+                    color_evt_s.record()
+                decoded_cthw = _maybe_apply_color_correction(pipeline, decoded_cthw)
+                if use_cuda_timing:
+                    color_evt_e.record()
+
+                if args.history_update_mode == "roundtrip":
+                    if use_cuda_timing:
+                        hist_evt_s = torch.cuda.Event(enable_timing=True)
+                        hist_evt_e = torch.cuda.Event(enable_timing=True)
+                        hist_evt_s.record()
+                    latent_motion_frames = _encode_motion_prefix_from_decoded(
+                        pipeline,
+                        decoded_video_cthw=decoded_cthw,
+                        motion_frames_num=motion_frames_num,
+                        device=args.device,
+                        dtype=amp_dtype,
+                    ).unsqueeze(0)
+                    if use_cuda_timing:
+                        hist_evt_e.record()
+                else:
+                    if use_cuda_timing:
+                        hist_evt_s = torch.cuda.Event(enable_timing=True)
+                        hist_evt_e = torch.cuda.Event(enable_timing=True)
+                        hist_evt_s.record()
+                    latent_motion_frames = _motion_prefix_from_latent_tail(
+                        x_final,
+                        motion_frames_latent_num=sp.motion_frames_latent_num,
+                    ).unsqueeze(0)
+                    if use_cuda_timing:
+                        hist_evt_e.record()
+                clamp_latent_len = int(latent_motion_frames.shape[2])
+
+                # Write frames: drop overlap on later chunks
+                if chunk_idx != 0:
+                    decoded_cthw = decoded_cthw[:, motion_frames_num:]
+
+                # Match official run_pipeline(): build THWC frames in [0,255] on GPU.
+                video_thwc = (
+                    ((decoded_cthw + 1.0) / 2.0)
+                    .permute(1, 2, 3, 0)
+                    .clamp(0.0, 1.0)
+                    .mul(255.0)
+                    .contiguous()
+                )
+
+                # End chunk_time at GPU boundary (match official timing semantics).
+                if use_cuda_timing:
+                    torch.cuda.synchronize()
+                chunk_time = time.perf_counter() - chunk_start
+
+                # Post-processing / IO time (excluded from chunk_time)
+                io_start = time.perf_counter()
+                video_cpu = video_thwc.detach().cpu()
+                if is_rank0:
+                    generated_list.append(video_cpu)
+                io_time = time.perf_counter() - io_start
+
+                if chunk_idx >= 2:
+                    total_frames += int(video_thwc.shape[0])
+                    total_time += float(chunk_time)
+                    total_chunks += 1
+
+                if use_cuda_timing:
+                    denoise_time = float(denoise_evt_s.elapsed_time(denoise_evt_e)) / 1000.0
+                    decode_time = float(decode_evt_s.elapsed_time(decode_evt_e)) / 1000.0
+                    color_time = float(color_evt_s.elapsed_time(color_evt_e)) / 1000.0
+                    hist_time = float(hist_evt_s.elapsed_time(hist_evt_e)) / 1000.0
+                else:
+                    denoise_time = float("nan")
+                    decode_time = float("nan")
+                    color_time = float("nan")
+                    hist_time = float("nan")
+
+                if is_rank0:
+                    logger.info(
+                        f"Chunk {chunk_idx + 1}/{num_chunks} timings | "
+                        f"denoise={denoise_time:.3f}s | decode={decode_time:.3f}s | color={color_time:.3f}s | "
+                        f"hist={hist_time:.3f}s | io={io_time:.3f}s"
+                    )
+                logger.info(f"Chunk {chunk_idx + 1}/{num_chunks} done")
+
+        else:
+            # stream: cached ring buffer, recompute embeddings for cached audio each chunk
+            from collections import deque
+
+            cached_len = args.sample_rate * sp.cached_audio_duration
+            audio_end_idx = sp.cached_audio_duration * args.fps
+            audio_start_idx = audio_end_idx - frame_num
+
+            if human_speech_array_slice_len <= 0:
+                raise SystemExit("Invalid audio slice length")
+
+            # Ensure divisible for reshape
+            remainder = len(audio_all) % human_speech_array_slice_len
+            if remainder != 0:
+                audio_all = np.concatenate([audio_all, np.zeros(human_speech_array_slice_len - remainder, dtype=audio_all.dtype)])
+
+            slices = audio_all.reshape(-1, human_speech_array_slice_len)
+            if slices.shape[0] == 0:
+                raise SystemExit("audio too short after padding")
+
+            audio_dq = deque([0.0] * cached_len, maxlen=cached_len)
+
+            num_chunks = int(slices.shape[0])
+            if args.max_chunks and args.max_chunks > 0:
+                num_chunks = min(num_chunks, int(args.max_chunks))
+
+            for chunk_idx in range(num_chunks):
+                use_cuda_timing = args.device.startswith("cuda") and torch.cuda.is_available()
+                if use_cuda_timing:
+                    torch.cuda.synchronize()
+                chunk_start = time.perf_counter()
+                audio_dq.extend(slices[chunk_idx].tolist())
+                audio_cache = np.array(audio_dq, dtype=np.float32)
+                audio_emb_cache = pipeline.preprocess_audio(audio_cache, sr=args.sample_rate, fps=args.fps)
+                if audio_emb_cache is None:
+                    raise RuntimeError("Failed to extract audio embeddings")
+                audio_emb_cache = audio_emb_cache.to(device=args.device, dtype=amp_dtype)
+                audio_ctx = _audio_context_from_embeddings_range(
+                    audio_emb_cache,
+                    start_idx=audio_start_idx,
+                    end_idx=audio_end_idx,
+                    device=args.device,
+                    dtype=amp_dtype,
+                )
+
+                if use_cuda_timing:
+                    denoise_evt_s = torch.cuda.Event(enable_timing=True)
+                    denoise_evt_e = torch.cuda.Event(enable_timing=True)
+                    denoise_evt_s.record()
+                x_final = _bridge_sample_one_chunk(
+                    pipeline,
+                    scheduler=scheduler,
+                    ref_latent=X0,
+                    audio_context=audio_ctx,
+                    guidance_scale=args.guidance_scale,
+                    latent_motion_frames=latent_motion_frames,
+                    clamp_latent_len=clamp_latent_len,
+                    device=args.device,
+                    dtype=amp_dtype,
+                )
+                if use_cuda_timing:
+                    denoise_evt_e.record()
+
+                if use_cuda_timing:
+                    decode_evt_s = torch.cuda.Event(enable_timing=True)
+                    decode_evt_e = torch.cuda.Event(enable_timing=True)
+                    decode_evt_s.record()
+                decoded_cthw = _decode_to_cthw(pipeline, x_final)
+                if use_cuda_timing:
+                    decode_evt_e.record()
+
+                if use_cuda_timing:
+                    color_evt_s = torch.cuda.Event(enable_timing=True)
+                    color_evt_e = torch.cuda.Event(enable_timing=True)
+                    color_evt_s.record()
+                decoded_cthw = _maybe_apply_color_correction(pipeline, decoded_cthw)
+                if use_cuda_timing:
+                    color_evt_e.record()
+
+                if args.history_update_mode == "roundtrip":
+                    if use_cuda_timing:
+                        hist_evt_s = torch.cuda.Event(enable_timing=True)
+                        hist_evt_e = torch.cuda.Event(enable_timing=True)
+                        hist_evt_s.record()
+                    latent_motion_frames = _encode_motion_prefix_from_decoded(
+                        pipeline,
+                        decoded_video_cthw=decoded_cthw,
+                        motion_frames_num=motion_frames_num,
+                        device=args.device,
+                        dtype=amp_dtype,
+                    ).unsqueeze(0)
+                    if use_cuda_timing:
+                        hist_evt_e.record()
+                else:
+                    if use_cuda_timing:
+                        hist_evt_s = torch.cuda.Event(enable_timing=True)
+                        hist_evt_e = torch.cuda.Event(enable_timing=True)
+                        hist_evt_s.record()
+                    latent_motion_frames = _motion_prefix_from_latent_tail(
+                        x_final,
+                        motion_frames_latent_num=sp.motion_frames_latent_num,
+                    ).unsqueeze(0)
+                    if use_cuda_timing:
+                        hist_evt_e.record()
+                clamp_latent_len = int(latent_motion_frames.shape[2])
+
+                # In stream mode, match generate_video.py: always drop overlap
+                decoded_cthw = decoded_cthw[:, motion_frames_num:]
+
+                # Match official run_pipeline(): build THWC frames in [0,255] on GPU.
+                video_thwc = (
+                    ((decoded_cthw + 1.0) / 2.0)
+                    .permute(1, 2, 3, 0)
+                    .clamp(0.0, 1.0)
+                    .mul(255.0)
+                    .contiguous()
+                )
+
+                # End chunk_time at GPU boundary (match official timing semantics).
+                if use_cuda_timing:
+                    torch.cuda.synchronize()
+                chunk_time = time.perf_counter() - chunk_start
+
+                # Post-processing / IO time (excluded from chunk_time)
+                io_start = time.perf_counter()
+                video_cpu = video_thwc.detach().cpu()
+                if is_rank0:
+                    generated_list.append(video_cpu)
+                io_time = time.perf_counter() - io_start
+
+                chunk_frames = int(video_thwc.shape[0])
+                chunk_fps = chunk_frames / chunk_time
+
+                if chunk_idx >= 2:
+                    total_frames += chunk_frames
+                    total_time += chunk_time
+                    total_chunks += 1
+
+                if use_cuda_timing:
+                    denoise_time = float(denoise_evt_s.elapsed_time(denoise_evt_e)) / 1000.0
+                    decode_time = float(decode_evt_s.elapsed_time(decode_evt_e)) / 1000.0
+                    color_time = float(color_evt_s.elapsed_time(color_evt_e)) / 1000.0
+                    hist_time = float(hist_evt_s.elapsed_time(hist_evt_e)) / 1000.0
+                else:
+                    denoise_time = float("nan")
+                    decode_time = float("nan")
+                    color_time = float("nan")
+                    hist_time = float("nan")
+
+                logger.info(
+                    f"Chunk {chunk_idx+1}/{num_chunks} | "
+                    f"time={chunk_time:.3f}s | "
+                    f"frames={chunk_frames} | "
+                    f"FPS={chunk_fps:.2f} | "
+                    f"denoise={denoise_time:.3f}s | decode={decode_time:.3f}s | color={color_time:.3f}s | "
+                    f"hist={hist_time:.3f}s | io={io_time:.3f}s"
+                )
+                logger.info(f"Chunk {chunk_idx + 1}/{num_chunks} done")
+
+    finally:
+        pass
+
+    if dist.is_initialized():
+        dist.barrier()
+
+    if is_rank0:
+        os.makedirs(os.path.dirname(os.path.abspath(args.out)) or ".", exist_ok=True)
+
+        # Save video once (official-style) then mux audio.
+        with _open_video_writer(tmp_mp4, fps=args.fps) as writer:
+            for video_cpu in generated_list:
+                # Match official save_video(): numpy().astype(np.uint8) at write time.
+                # NumPy does not support bfloat16; cast at write time (outside chunk_time).
+                if video_cpu.dtype == torch.bfloat16:
+                    video_cpu = video_cpu.to(torch.float16)
+                frames_np = video_cpu.numpy().astype(np.uint8)
+                for fr in frames_np:
+                    writer.append_data(fr)
+        _mux_audio(tmp_mp4, args.audio_path, args.out)
+        try:
+            os.remove(tmp_mp4)
+        except OSError:
+            pass
+
+        if total_time > 0:
+            avg_fps = total_frames / total_time
+            avg_chunk_time = total_time / total_chunks if total_chunks > 0 else float("nan")
+            logger.info(
+                f"Average generation FPS (excl first 2 chunks): {avg_fps:.2f} "
+                f"(total_frames={total_frames}, total_time={total_time:.2f}s)"
+            )
+            logger.info(f"Average chunk time (excl first 2 chunks): {avg_chunk_time:.3f}s")
+        logger.info(f"Saved: {args.out}")
+
+    if dist.is_initialized():
+        dist.barrier()
+    if did_init_dist:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,18 @@
+torchvision
+diffusers==0.38.0
+transformers==4.57.3
+tokenizers
+peft==0.19.1
+accelerate
+safetensors
+einops
+numpy
+pillow
+librosa
+soundfile
+imageio
+imageio-ffmpeg
+loguru
+pyloudnorm
+mediapipe==0.10.35
+tqdm

vibt/scheduler.py

@@ -0,0 +1,59 @@
+from diffusers.schedulers import UniPCMultistepScheduler
+import torch
+import time
+
+
+class ViBTScheduler(UniPCMultistepScheduler):
+    def __init__(self, **kwargs):
+        super().__init__(**{**kwargs, "use_flow_sigmas": True})
+        self.set_parameters()
+
+    def set_parameters(self, noise_scale=1.0, shift_gamma=5.0, seed=None):
+        self.noise_scale = noise_scale
+        self.config.flow_shift = shift_gamma
+        self.generator = (
+            None if seed is None else torch.Generator("cuda").manual_seed(seed)
+        )
+
+    def step(self, model_output, timestep, sample, ode=False, **kwargs):
+        t_val = float(timestep)
+        
+        if not hasattr(self, "_t_list") or len(self._t_list) != len(self.timesteps):
+            self._t_list = self.timesteps.cpu().tolist()
+            
+        idx = -1
+        for i, val in enumerate(self._t_list):
+            if abs(val - t_val) < 1e-4:
+                idx = i
+                break
+                
+        if idx != -1 and idx + 1 < len(self._t_list):
+            next_t = self._t_list[idx + 1]
+        else:
+            next_t = -1.0
+
+        delta_t = (next_t - t_val) / 1000.0
+        current_t = (t_val + 1.0) / 1000.0
+        
+        eta_radicand = -delta_t * (current_t + delta_t) / current_t
+        eta = (max(0.0, eta_radicand)) ** 0.5
+
+        noise = torch.randn(
+            sample.shape,
+            generator=self.generator,
+            device=sample.device,
+            dtype=sample.dtype,
+        )
+
+        latents = sample + delta_t * model_output + eta * self.noise_scale * noise
+
+        return (latents,)
+
+    @classmethod
+    def from_scheduler(
+        cls, scheduler: UniPCMultistepScheduler, noise_scale=1.0, shift_gamma=5.0
+    ):
+        obj = cls.__new__(cls)
+        obj.__dict__ = scheduler.__dict__.copy()
+        obj.set_parameters(noise_scale, shift_gamma)
+        return obj