Upload 54 files

.gitattributes

@@ -33,3 +33,14 @@ 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
+assets/teaser.png filter=lfs diff=lfs merge=lfs -text
+examples/amber.png filter=lfs diff=lfs merge=lfs -text
+examples/armour.png filter=lfs diff=lfs merge=lfs -text
+examples/art.wav filter=lfs diff=lfs merge=lfs -text
+examples/chris.png filter=lfs diff=lfs merge=lfs -text
+examples/dream.mp3 filter=lfs diff=lfs merge=lfs -text
+examples/fictional.wav filter=lfs diff=lfs merge=lfs -text
+examples/fight.wav filter=lfs diff=lfs merge=lfs -text
+examples/jacket.png filter=lfs diff=lfs merge=lfs -text
+examples/naomi.png filter=lfs diff=lfs merge=lfs -text
+examples/vangogh.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 2025 Bytedance
+
+   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,12 +1,13 @@
----
-title: HuMo Local
-emoji: 🌍
-colorFrom: green
-colorTo: red
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: HuMo [Local]
+emoji: 👩‍🦱
+colorFrom: purple
+colorTo: gray
+sdk: gradio
+sdk_version: 5.47.2
+app_file: app.py
+pinned: false
+short_description: Reference based video generation
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,435 @@
+import spaces
+import gradio as gr
+import sys
+import os
+import subprocess
+import uuid
+import shutil
+
+
+
+from huggingface_hub import snapshot_download, list_repo_files, hf_hub_download
+import importlib, site
+
+
+# Re-discover all .pth/.egg-link files
+for sitedir in site.getsitepackages():
+    site.addsitedir(sitedir)
+
+# Clear caches so importlib will pick up new modules
+importlib.invalidate_caches()
+
+def sh(cmd): subprocess.check_call(cmd, shell=True)
+
+flash_attention_installed = False
+
+try:
+    flash_attention_wheel = hf_hub_download(
+            repo_id="alexnasa/flash-attn-3",
+            repo_type="model",
+            filename="128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl",
+        )
+
+    sh(f"pip install {flash_attention_wheel}")
+    print("Attempting to download and install FlashAttention wheel...")
+    # sh("pip install flash-attn")
+    sh("pip install --no-build-isolation transformer_engine-2.5.0+f05f12c9-cp310-cp310-linux_x86_64.whl")
+
+    # tell Python to re-scan site-packages now that the egg-link exists
+    import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()
+
+    flash_attention_installed = True
+
+except Exception as e:
+    print(f"⚠️ Could not install FlashAttention: {e}")
+    print("Continuing without FlashAttention...")
+
+try:
+    te_wheel = hf_hub_download(
+            repo_id="alexnasa/transformer_engine_wheels",
+            repo_type="model",
+            filename="transformer_engine-2.5.0+f05f12c9-cp310-cp310-linux_x86_64.whl",
+        )
+
+    sh(f"pip install {te_wheel}")
+    print("Attempting to download and install Transformer Engine wheel...")
+
+    # tell Python to re-scan site-packages now that the egg-link exists
+    import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()
+
+except Exception as e:
+    print(f"⚠️ Could not install Transformer Engine : {e}")
+    print("Continuing without Transformer Engine ...")
+
+import torch
+print(f"Torch version: {torch.__version__}")
+print(f"FlashAttention available: {flash_attention_installed}")
+
+import tempfile
+from pathlib import Path
+from torch._inductor.runtime.runtime_utils import cache_dir as _inductor_cache_dir
+from huggingface_hub import HfApi
+
+
+snapshot_download(repo_id="bytedance-research/HuMo", local_dir="./weights/HuMo")
+snapshot_download(repo_id="Wan-AI/Wan2.1-T2V-1.3B", local_dir="./weights/Wan2.1-T2V-1.3B")
+snapshot_download(repo_id="openai/whisper-large-v3", local_dir="./weights/whisper-large-v3")
+
+os.environ["PROCESSED_RESULTS"] = f"{os.getcwd()}/proprocess_results"
+
+path_to_insert = "humo"
+if path_to_insert not in sys.path:
+    sys.path.insert(0, path_to_insert)
+
+from common.config import load_config, create_object
+
+config = load_config(
+    "./humo/configs/inference/generate.yaml",
+    [
+        "dit.sp_size=1",
+        "generation.frames=97",
+        "generation.scale_t=5.5",
+        "generation.scale_a=5.0",
+        "generation.mode=TIA",
+        "generation.height=480",
+        "generation.width=832",
+    ],
+)
+runner = create_object(config)
+
+
+os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", f"{os.getcwd()}/torchinductor_space")  # or another writable path
+
+def restore_inductor_cache_from_hub(repo_id: str, filename: str = "torch_compile_cache.zip",
+                                    path_in_repo: str = "inductor_cache", repo_type: str = "model",
+                                    hf_token: str | None = None):
+    cache_root = Path(_inductor_cache_dir()).resolve()
+    cache_root.mkdir(parents=True, exist_ok=True)
+    zip_path = hf_hub_download(repo_id=repo_id, filename=f"{path_in_repo}/{filename}",
+                               repo_type=repo_type, token=hf_token)
+    shutil.unpack_archive(zip_path, extract_dir=str(cache_root))
+    print(f"✓ Restored cache into {cache_root}")
+
+
+# restore_inductor_cache_from_hub("alexnasa/humo-compiled")
+
+
+def get_duration(prompt_text, steps, image_file, audio_file_path, tea_cache_l1_thresh, max_duration, session_id):
+
+    return calculate_required_time(steps, max_duration)
+
+def calculate_required_time(steps, max_duration):
+    
+    warmup_s = 60
+
+    max_duration_duration_mapping = {
+        1: 8,
+        2: 8,
+        3: 11,
+        4: 20,
+        5: 30,
+    }
+    each_step_s = max_duration_duration_mapping[max_duration]
+    duration_s = (each_step_s * steps) + warmup_s
+
+    print(f'estimated duration:{duration_s}')
+
+    return int(duration_s)
+
+def get_required_time_string(steps, max_duration):
+
+    duration_s = calculate_required_time(steps, max_duration)
+    duration_m = duration_s / 60
+
+    return f"<center>⌚ Zero GPU Required: ~{duration_s}.0s ({duration_m:.1f} mins)</center>"
+
+def update_required_time(steps, max_duration):
+
+    return get_required_time_string(steps, max_duration)
+
+
+def generate_scene(prompt_text, steps, image_paths, audio_file_path, tea_cache_l1_thresh, max_duration = 2, session_id = None):
+
+    print(image_paths)
+    prompt_text_check = (prompt_text or "").strip()
+    if not prompt_text_check:
+        raise gr.Error("Please enter a prompt.")
+    
+    if not audio_file_path and not image_paths:
+        raise gr.Error("Please provide a reference image or a lipsync audio.")
+    
+    return run_pipeline(prompt_text, steps, image_paths, audio_file_path, tea_cache_l1_thresh, max_duration, session_id)
+
+
+
+def upload_inductor_cache_to_hub(
+    repo_id: str,
+    path_in_repo: str = "inductor_cache",
+    repo_type: str = "model",   # or "dataset" if you prefer
+    hf_token: str | None = None,
+):
+    """
+    Zips the current TorchInductor cache and uploads it to the given repo path.
+    Assumes the model was already run once with torch.compile() so the cache exists.
+    """
+
+    cache_dir = Path(_inductor_cache_dir()).resolve()
+    if not cache_dir.exists():
+        raise FileNotFoundError(f"TorchInductor cache not found at {cache_dir}. "
+                                "Run a compiled model once to populate it.")
+
+    # Create a zip archive of the entire cache directory
+    with tempfile.TemporaryDirectory() as tmpdir:
+        archive_base = Path(tmpdir) / "torch_compile_cache"
+        archive_path = shutil.make_archive(str(archive_base), "zip", root_dir=str(cache_dir))
+        archive_path = Path(archive_path)
+
+        # Upload to Hub
+        api = HfApi(token=hf_token)
+        api.create_repo(repo_id=repo_id, repo_type=repo_type, exist_ok=True)
+        # Put each artifact under path_in_repo, including a tiny metadata stamp for traceability
+        # Upload the zip
+        dest_path = f"{path_in_repo}/{archive_path.name}"
+        api.upload_file(
+            path_or_fileobj=str(archive_path),
+            path_in_repo=dest_path,
+            repo_id=repo_id,
+            repo_type=repo_type,
+        )
+        # Upload a small metadata file (optional but handy)
+        meta_txt = (
+            f"pytorch={torch.__version__}\n"
+            f"inductor_cache_dir={cache_dir}\n"
+            f"cuda_available={torch.cuda.is_available()}\n"
+            f"cuda_device={torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'cpu'}\n"
+        )
+        api.upload_file(
+            path_or_fileobj=meta_txt.encode(),
+            path_in_repo=f"{path_in_repo}/INDUCTOR_CACHE_METADATA.txt",
+            repo_id=repo_id,
+            repo_type=repo_type,
+        )
+
+    print("✔ Uploaded TorchInductor cache to the Hub.")
+
+
+@spaces.GPU(duration=get_duration)
+def run_pipeline(prompt_text, steps, image_paths, audio_file_path, tea_cache_l1_thresh = 0.0, max_duration = 2, session_id = None):
+
+    if session_id is None:
+        session_id = uuid.uuid4().hex
+
+    inference_mode = "TIA"
+
+    # Validate inputs
+    prompt_text = (prompt_text or "").strip()
+    if not prompt_text:
+        raise gr.Error("Please enter a prompt.")
+    
+    if not audio_file_path and not image_paths:
+        raise gr.Error("Please provide a reference image or a lipsync audio.")
+    
+    if not audio_file_path:
+        inference_mode = "TI"
+        audio_path = None
+    else:
+        audio_path = audio_file_path if isinstance(audio_file_path, str) else getattr(audio_file_path, "name", str(audio_file_path))
+
+    if not image_paths:
+        inference_mode = "TA"
+        img_paths = None
+    else:
+        img_paths = [image_data[0] for image_data in image_paths]
+
+
+    # Prepare output
+    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Random filename
+    filename = f"gen_{uuid.uuid4().hex[:10]}"
+    width, height = 832, 480
+
+    duration_frame_mapping = {
+        1:25,
+        2:45,
+        3:70, 
+        4:97,
+        5:129
+    }
+
+    # Run inference
+    runner.inference_loop(
+        prompt_text,
+        img_paths,
+        audio_path,
+        output_dir,
+        filename,
+        inference_mode,
+        width,
+        height,
+        steps,
+        frames = int(duration_frame_mapping[max_duration]),
+        tea_cache_l1_thresh = tea_cache_l1_thresh,
+    )
+
+    # Return resulting video path
+    video_path = os.path.join(output_dir, f"{filename}.mp4")
+    if os.path.exists(video_path):
+
+        # upload_inductor_cache_to_hub("alexnasa/humo-compiled")
+
+        return video_path
+    else:
+        candidates = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if f.endswith(".mp4")]
+        if candidates:
+            return max(candidates, key=lambda p: os.path.getmtime(p))
+        return None
+
+css = """
+    #col-container {
+        margin: 0 auto;
+        width: 100%;
+        max-width: 720px;
+    }
+    """
+
+def cleanup(request: gr.Request):
+
+    sid = request.session_hash
+    if sid:
+        d1 = os.path.join(os.environ["PROCESSED_RESULTS"], sid)
+        shutil.rmtree(d1, ignore_errors=True)
+        
+def start_session(request: gr.Request):
+
+    return request.session_hash
+
+with gr.Blocks(css=css) as demo:
+
+    session_state = gr.State()
+    demo.load(start_session, outputs=[session_state])
+
+    with gr.Sidebar(width=400):
+
+
+        gr.HTML(
+            """
+            <div style="text-align: center;">
+                <p style="font-size:16px; display: inline; margin: 0;">
+                    <strong>HuMo</strong> – Human-Centric Video Generation via Collaborative Multi-Modal Conditioning
+                </p>
+                <a href="https://github.com/Phantom-video/HuMo" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">
+                    [Github]
+                </a>
+            </div>
+            """
+        )
+
+        gr.Markdown("**REFERENCE IMAGES**")
+
+        img_input = gr.Gallery(
+            show_label=False,
+            label="",
+            interactive=True,
+            rows=1, columns=3, object_fit="contain", height="280",
+            file_types=['image']
+        )
+
+        gr.Markdown("**LIPSYNC AUDIO**")
+
+        audio_input = gr.Audio(
+            sources=["upload"],
+            show_label=False,
+            type="filepath",
+        )
+
+        gr.Markdown("**SETTINGS**")
+
+        default_steps = 10
+        default_max_duration = 2
+
+        max_duration = gr.Slider(minimum=2, maximum=5, value=default_max_duration, step=1, label="Max Duration")
+        steps_input = gr.Slider(minimum=5, maximum=50, value=default_steps, step=5, label="Diffusion Steps")
+        tea_cache_l1_thresh = gr.Slider(minimum=0.0, maximum=1.0, value=0.0, step=0.01, label="Cache", visible=False)
+        
+
+
+    with gr.Column(elem_id="col-container"):
+
+        gr.HTML(
+            """
+            <div style="text-align: center;">
+                <strong>HF Space by:</strong>
+                <a href="https://twitter.com/alexandernasa/" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">
+                    <img src="https://img.shields.io/twitter/url/https/twitter.com/cloudposse.svg?style=social&label=Follow Me" alt="GitHub Repo">
+                </a>
+            </div>
+            """
+        )
+
+        video_output = gr.Video(show_label=False)
+
+        gr.Markdown("<center><h2>PROMPT</h2></center>")
+
+        prompt_tb = gr.Textbox(
+            show_label=False,
+            lines=5,
+            placeholder="Describe the scene and the person talking....",
+        )
+
+        gr.Markdown("")
+        time_required = gr.Markdown(get_required_time_string(default_steps, default_max_duration))
+        run_btn = gr.Button("🎬 Action", variant="primary")
+        
+        gr.Examples(
+            examples=[
+
+                [
+                    "A handheld tracking shot follows a female warrior walking through a cave. Her determined eyes are locked straight ahead. She speaks with intensity.",
+                    5,
+                    ["./examples/naomi.png"], 
+                    "./examples/dream.mp3",                  
+                ],
+
+                [
+                    "A reddish-brown haired and bearded man sits pensively against swirling blue-and-white brushstrokes, dressed in a blue coat and dark waistcoat. The artistic backdrop and his thoughtful pose evoke a Post-Impressionist style in a studio-like setting.",
+                    10,
+                    ["./examples/vangogh.jpg"], 
+                    "./examples/art.wav",               
+                ],
+
+                [
+                    "A handheld tracking shot follows a female through a science lab. Her determined eyes are locked straight ahead. The clip is in black and white and patchy as she is explaining something to someone standing opposite her",
+                    10,
+                    ["./examples/naomi.png"], 
+                    "./examples/science.wav",               
+                ],
+
+                [
+                    "A woman with long, wavy dark hair looking at a person sitting opposite her whilst holding a book, wearing a leather jacket, long-sleeved jacket with a semi purple color one seen on a photo. Warm, window-like light bathes her figure, highlighting the outfit's elegant design and her graceful movements.",
+                    50,
+                    ["./examples/amber.png", "./examples/jacket.png"],
+                    "./examples/fictional.mp3",                  
+                ],
+
+            ],
+            inputs=[prompt_tb, steps_input, img_input, audio_input],
+            outputs=[video_output],
+            fn=run_pipeline,
+            cache_examples=True,
+        )
+        max_duration.change(update_required_time, [steps_input, max_duration], time_required)
+        steps_input.change(update_required_time, [steps_input, max_duration], time_required)
+
+        run_btn.click(
+            fn=generate_scene,
+            inputs=[prompt_tb, steps_input, img_input, audio_input, tea_cache_l1_thresh, max_duration, session_state],
+            outputs=[video_output],
+        )
+
+
+if __name__ == "__main__":
+    demo.unload(cleanup)
+    demo.queue()
+    demo.launch(ssr_mode=False)

examples/art.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72c75df8e93a107e262ea9b002a66e72d3c1cd2084bce1474a31d8afffd0b651
+size 114254

examples/dream.mp3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27248fd9e8f29bd60ccb1163b8df3c6f2630734f358aa3362ffe67e8148e0eb1
+size 108275

examples/fictional.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31b550e6433ea44a0642dee90c326664ff4f568fec184170001f834597b3ad23
+size 167084

examples/fight.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8dbee86c85e992ac6d17820a3730bf753fc9bf5bac6b8a470f84b7e98a64221a
+size 264782

humo/common/config.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+# Codes adapted from [SeedVR]
+# https://github.com/ByteDance-Seed/SeedVR/blob/main/common/config.py
+
+"""
+Configuration utility functions
+"""
+
+import importlib
+from typing import Any, Callable, List, Union
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+OmegaConf.register_new_resolver("eval", eval)
+
+
+def load_config(path: str, argv: List[str] = None) -> Union[DictConfig, ListConfig]:
+    """
+    Load a configuration. Will resolve inheritance.
+    """
+    config = OmegaConf.load(path)
+    if argv is not None:
+        config_argv = OmegaConf.from_dotlist(argv)
+        config = OmegaConf.merge(config, config_argv)
+    config = resolve_recursive(config, resolve_inheritance)
+    return config
+
+
+def resolve_recursive(
+    config: Any,
+    resolver: Callable[[Union[DictConfig, ListConfig]], Union[DictConfig, ListConfig]],
+) -> Any:
+    config = resolver(config)
+    if isinstance(config, DictConfig):
+        for k in config.keys():
+            v = config.get(k)
+            if isinstance(v, (DictConfig, ListConfig)):
+                config[k] = resolve_recursive(v, resolver)
+    if isinstance(config, ListConfig):
+        for i in range(len(config)):
+            v = config.get(i)
+            if isinstance(v, (DictConfig, ListConfig)):
+                config[i] = resolve_recursive(v, resolver)
+    return config
+
+
+def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any:
+    """
+    Recursively resolve inheritance if the config contains:
+    __inherit__: path/to/parent.yaml.
+    """
+    if isinstance(config, DictConfig):
+        inherit = config.pop("__inherit__", None)
+        if inherit:
+            assert isinstance(inherit, str)
+            inherit = load_config(inherit)
+            if len(config.keys()) > 0:
+                config = OmegaConf.merge(inherit, config)
+            else:
+                config = inherit
+    return config
+
+
+def import_item(path: str, name: str) -> Any:
+    """
+    Import a python item. Example: import_item("path.to.file", "MyClass") -> MyClass
+    """
+    return getattr(importlib.import_module(path), name)
+
+
+def create_object(config: DictConfig) -> Any:
+    """
+    Create an object from config.
+    The config is expected to contains the following:
+    __object__:
+      path: path.to.module
+      name: MyClass
+      args: as_config | as_params (default to as_config)
+    """
+    item = import_item(
+        path=config.__object__.path,
+        name=config.__object__.name,
+    )
+    args = config.__object__.get("args", "as_config")
+    if args == "as_config":
+        return item(config)
+    if args == "as_params":
+        config = OmegaConf.to_object(config)
+        config.pop("__object__")
+        return item(**config)
+    raise NotImplementedError(f"Unknown args type: {args}")
+
+
+def create_dataset(path: str, *args, **kwargs) -> Any:
+    """
+    Create a dataset. Requires the file to contain a "create_dataset" function.
+    """
+    return import_item(path, "create_dataset")(*args, **kwargs)

humo/common/distributed/__init__.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+# Codes adapted from [SeedVR]
+# https://github.com/ByteDance-Seed/SeedVR/tree/main/common/distributed
+
+"""
+Distributed package.
+"""
+
+from .basic import (
+    barrier_if_distributed,
+    convert_to_ddp,
+    get_device,
+    get_global_rank,
+    get_local_rank,
+    get_world_size,
+    init_torch,
+    meta_param_init_fn,
+    meta_non_persistent_buffer_init_fn
+)
+
+__all__ = [
+    "barrier_if_distributed",
+    "convert_to_ddp",
+    "get_device",
+    "get_global_rank",
+    "get_local_rank",
+    "get_world_size",
+    "init_torch",
+    "meta_param_init_fn",
+    "meta_non_persistent_buffer_init_fn",
+]

humo/common/distributed/advanced.py

@@ -0,0 +1,484 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+# Codes adapted from [SeedVR]
+# https://github.com/ByteDance-Seed/SeedVR/tree/main/common/distributed
+
+"""
+Advanced distributed functions for sequence parallel.
+"""
+
+import torch
+from typing import Any, List, Optional, Tuple, Union
+import torch.distributed as dist
+from torch import Tensor
+
+from .basic import get_global_rank, get_world_size
+
+
+_DATA_PARALLEL_GROUP = None
+_SEQUENCE_PARALLEL_GROUP = None
+_SEQUENCE_PARALLEL_CPU_GROUP = None
+
+
+_CFG_PARALLEL_GROUP = None
+_CFG_PARALLEL_CPU_GROUP = None
+
+def get_data_parallel_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get data parallel process group.
+    """
+    return _DATA_PARALLEL_GROUP
+
+
+def get_sequence_parallel_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get sequence parallel process group.
+    """
+    return _SEQUENCE_PARALLEL_GROUP
+
+
+def get_sequence_parallel_cpu_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get sequence parallel CPU process group.
+    """
+    return _SEQUENCE_PARALLEL_CPU_GROUP
+
+
+def get_data_parallel_rank() -> int:
+    """
+    Get data parallel rank.
+    """
+    group = get_data_parallel_group()
+    return dist.get_rank(group) if group else get_global_rank()
+
+
+def get_data_parallel_world_size() -> int:
+    """
+    Get data parallel world size.
+    """
+    group = get_data_parallel_group()
+    return dist.get_world_size(group) if group else get_world_size()
+
+
+def get_sequence_parallel_rank() -> int:
+    """
+    Get sequence parallel rank.
+    """
+    group = get_sequence_parallel_group()
+    return dist.get_rank(group) if group else 0
+
+
+def get_sequence_parallel_world_size() -> int:
+    """
+    Get sequence parallel world size.
+    """
+    group = get_sequence_parallel_group()
+    return dist.get_world_size(group) if group else 1
+
+
+def init_unified_parallel(unified_parallel_size):
+    global _SEQUENCE_PARALLEL_GROUP
+    global _SEQUENCE_PARALLEL_CPU_GROUP
+
+    if unified_parallel_size == 1:
+        return
+
+    assert dist.is_initialized()
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    assert world_size % unified_parallel_size == 0
+    data_parallel_size = world_size // unified_parallel_size
+
+    for i in range(data_parallel_size):
+        # build unified parallel group
+        start_rank = i * unified_parallel_size
+        end_rank = start_rank + unified_parallel_size
+        unified_parallel_ranks = range(start_rank, end_rank)
+        unified_parallel_group = dist.new_group(unified_parallel_ranks)
+        unified_parallel_cpu_group = dist.new_group(unified_parallel_ranks, backend="gloo")
+        if rank in unified_parallel_ranks:
+            _SEQUENCE_PARALLEL_GROUP = unified_parallel_group
+            _SEQUENCE_PARALLEL_CPU_GROUP = unified_parallel_cpu_group
+
+
+def get_unified_parallel_group():
+    global _SEQUENCE_PARALLEL_GROUP
+    return _SEQUENCE_PARALLEL_GROUP
+
+
+def get_unified_parallel_cpu_group():
+    global _SEQUENCE_PARALLEL_CPU_GROUP
+    return _SEQUENCE_PARALLEL_CPU_GROUP
+
+
+def get_unified_parallel_rank():
+    group = get_unified_parallel_group()
+    return dist.get_rank(group) if group else 0
+
+
+def get_unified_parallel_world_size():
+    group = get_unified_parallel_group()
+    return dist.get_world_size(group) if group else 1
+
+
+def is_unified_parallel_initialized():
+    group = get_unified_parallel_group()
+    return group is not None
+
+
+def pad_tensor(x: Tensor, dim: int, padding_size: int):
+    shape = list(x.shape)
+    shape[dim] = padding_size
+    pad = torch.zeros(shape, dtype=x.dtype, device=x.device)
+    return torch.cat([x, pad], dim=dim)
+
+
+class Slice(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx: Any, group: dist.ProcessGroup, local_input: Tensor, dim: int, scale_grad: bool) -> Tensor:
+        ctx.group = group
+        ctx.rank = dist.get_rank(group)
+        seq_world_size = dist.get_world_size(group)
+        ctx.seq_world_size = seq_world_size
+        ctx.dim = dim
+        ctx.scale_grad = scale_grad
+        dim_size = local_input.shape[dim]
+        if not ctx.group:
+            return local_input
+        return local_input.split(dim_size // seq_world_size, dim=dim)[ctx.rank].contiguous()
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor, None]:
+        if not ctx.group:
+            return None, grad_output, None, None
+        dim_size = list(grad_output.size())
+        split_size = dim_size[0]
+        dim_size[0] = dim_size[0] * ctx.seq_world_size
+        output = torch.empty(dim_size, dtype=grad_output.dtype, device=torch.cuda.current_device())
+        dist.all_gather_into_tensor(output, grad_output, group=ctx.group)
+        if ctx.scale_grad:
+            output = output / ctx.seq_world_size
+        return (None, torch.cat(output.split(split_size), dim=ctx.dim), None, None)
+
+
+def gather_outputs(
+    x: Tensor,
+    gather_dim: int,
+    padding_dim: Optional[int] = None,
+    unpad_dim_size: Optional[int] = None,
+    scale_grad=True,
+):
+    """
+    A func to gather the outputs for the model result in sequence parallel
+    """
+    group = get_unified_parallel_group()
+    if not group:
+        return x
+    x = Gather.apply(group, x, gather_dim, scale_grad)
+    if padding_dim is not None:
+        x = unpadding_tensor_for_seqeunce_parallel(x, padding_dim, unpad_dim_size)
+    return x
+
+
+def unpadding_tensor_for_seqeunce_parallel(x: Tensor, dim: int, unpadded_dim_size: int):
+    """
+    A func to remove the padding part of the tensor based on its original shape
+    """
+    group = get_unified_parallel_group()
+    if group is None:
+        return x
+    sp_world = get_unified_parallel_world_size()
+    if unpadded_dim_size % sp_world == 0:
+        return x
+    padding_size = sp_world - (unpadded_dim_size % sp_world)
+    assert (padding_size + unpadded_dim_size) % sp_world == 0
+    return unpad_tensor(x, dim=dim, padding_size=padding_size)
+
+
+def gather_seq_scatter_heads_qkv(
+    qkv_tensor: Tensor,
+    seq_dim: int,
+    unpadded_dim_size: Optional[int] = None,
+    restore_shape: bool = True,
+    async_op: bool = False,
+):
+    """
+    A func to sync splited qkv tensor
+    qkv_tensor: the tensor we want to do alltoall with. The last dim must
+        be the projection_idx, which we will split into 3 part. After
+        spliting, the gather idx will be projecttion_idx + 1
+    seq_dim: gather_dim for all2all comm
+    restore_shape: if True, output will has the same shape length as input
+    """
+    group = get_unified_parallel_group()
+    if not group:
+        return qkv_tensor
+    world = get_unified_parallel_world_size()
+    orig_shape = qkv_tensor.shape
+    scatter_dim = qkv_tensor.dim()
+    bef_all2all_shape = list(orig_shape)
+    qkv_proj_dim = bef_all2all_shape[-1]
+    bef_all2all_shape = bef_all2all_shape[:-1] + [3, qkv_proj_dim // 3]
+    qkv_tensor = qkv_tensor.view(bef_all2all_shape)
+    if async_op:
+        return SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, async_op)
+    else:
+        qkv_tensor = SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, async_op)
+
+        if restore_shape:
+            out_shape = list(orig_shape)
+            out_shape[seq_dim] *= world
+            out_shape[-1] = qkv_proj_dim // world
+            qkv_tensor = qkv_tensor.view(out_shape)
+
+        # remove padding
+        if unpadded_dim_size and unpadded_dim_size % world != 0:
+            padding_size = qkv_tensor.size(seq_dim) - unpadded_dim_size
+            qkv_tensor = unpad_tensor(qkv_tensor, seq_dim, padding_size)
+
+        return qkv_tensor
+
+
+def gather_seq_scatter_double_head(
+    qkv_tensor: Tensor,
+    seq_dim: int,
+    unpadded_dim_size: Optional[int] = None,
+    restore_shape: bool = True,
+    async_op: bool = False,
+):
+    """
+    A func to sync splited qkv tensor
+    qkv_tensor: the tensor we want to do alltoall with. The last dim must
+        be the projection_idx, which we will split into 3 part. After
+        spliting, the gather idx will be projecttion_idx + 1
+    seq_dim: gather_dim for all2all comm
+    restore_shape: if True, output will has the same shape length as input
+    """
+    qkv1_shape = qkv_tensor.shape
+    group = get_unified_parallel_group()
+    if not group:
+        return qkv_tensor
+    world = get_unified_parallel_world_size()
+    orig_shape = qkv_tensor.shape
+    scatter_dim = qkv_tensor.dim()
+    bef_all2all_shape = list(orig_shape)
+    qkv_proj_dim = bef_all2all_shape[-1]
+    bef_all2all_shape = bef_all2all_shape[:-1] + [2, qkv_proj_dim // 2]
+    qkv_tensor = qkv_tensor.view(bef_all2all_shape)
+    qkv2_shape = qkv_tensor.shape
+    if async_op:
+        return SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, async_op)
+    else:
+        qkv_tensor = SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, async_op)
+        qkv3_shape = qkv_tensor.shape
+
+        if restore_shape:
+            out_shape = list(orig_shape)
+            out_shape[seq_dim] *= world
+            out_shape[-1] = qkv_proj_dim // world
+            qkv_tensor = qkv_tensor.view(out_shape)
+            qkv4_shape = qkv_tensor.shape
+
+        # remove padding
+        if unpadded_dim_size and unpadded_dim_size % world != 0:
+            padding_size = qkv_tensor.size(seq_dim) - unpadded_dim_size
+            qkv_tensor = unpad_tensor(qkv_tensor, seq_dim, padding_size)
+            qkv5_shape = qkv_tensor.shape
+        
+        return qkv_tensor
+
+
+class SeqAllToAll(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        local_input: Tensor,
+        scatter_dim: int,
+        gather_dim: int,
+        async_op: bool,
+    ) -> Tensor:
+        ctx.group = group
+        ctx.scatter_dim = scatter_dim
+        ctx.gather_dim = gather_dim
+        ctx.async_op = async_op
+        return all_to_all_tensor(local_input, scatter_dim, gather_dim, group, async_op)
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
+        if ctx.async_op:
+            input_t = torch.cat(grad_output[1:], dim=ctx.gather_dim).contiguous()
+        else:
+            input_t = grad_output[0]
+        return (
+            None,
+            all_to_all_tensor(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False),
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+def all_to_all_tensor(
+    x: Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    group: dist.ProcessGroup,
+    async_op: bool = False,
+):
+    if scatter_dim <= 1 and gather_dim <= 1:
+        return _all_to_all_single(x, scatter_dim, gather_dim, group, async_op)
+    else:
+        return _all_to_all(x, scatter_dim, gather_dim, group, async_op)  # 走这里
+
+
+def _all_to_all(
+    local_input: Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    group: dist.ProcessGroup,
+    async_op: bool = False,
+):
+    seq_world_size = dist.get_world_size(group)
+    input_list = [t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]
+    comm = dist.all_to_all(output_list, input_list, group=group, async_op=async_op)
+    if async_op:
+
+        def wait():
+            comm.wait()
+            return torch.cat(output_list, dim=gather_dim).contiguous()
+
+        return wait
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+
+def _all_to_all_single(x: Tensor, scatter_dim: int, gather_dim: int, group: dist.ProcessGroup, async_op: bool = False):
+    """
+    A function to do all-to-all on the first two dim
+    """
+    sp_world_size = dist.get_world_size(group)
+    assert scatter_dim <= 1, "scatter_dim must be 0 or 1 when using all_to_all_single!"
+    assert gather_dim <= 1, "gather_dim must be 0 or 1 when using all_to_all_single!"
+    if scatter_dim != 0:
+        gather_dim_bef = x.shape[gather_dim]
+        scatter_dim_bef = x.shape[scatter_dim]
+        x = (
+            x.reshape([gather_dim_bef, sp_world_size, scatter_dim_bef // sp_world_size] + list(x.shape[2:]))
+            .transpose(0, 1)
+            .reshape([gather_dim_bef * sp_world_size, scatter_dim_bef // sp_world_size] + list(x.shape[2:]))
+            .contiguous()
+        )
+
+    output = torch.empty_like(x)
+    comm = dist.all_to_all_single(output, x.contiguous(), group=group, async_op=async_op)
+
+    if async_op:
+
+        def wait():
+            comm.wait()
+            if scatter_dim == 0:
+                return torch.cat(output.split(x.size(0) // sp_world_size), dim=gather_dim)
+            else:
+                return output
+
+        return wait
+
+    if scatter_dim == 0:
+        output = torch.cat(output.split(x.size(0) // sp_world_size), dim=gather_dim)
+    return output
+
+
+def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int) -> Tensor:
+    """
+    A func to sync attention result with alltoall in sequence parallel
+    """
+    group = get_unified_parallel_group()
+    if not group:
+        return x
+    dim_size = x.size(seq_dim)
+    sp_world = get_unified_parallel_world_size()
+    if dim_size % sp_world != 0:
+        padding_size = sp_world - (dim_size % sp_world)
+        x = pad_tensor(x, seq_dim, padding_size)
+    return SeqAllToAll.apply(group, x, seq_dim, head_dim, False)
+
+
+def unpad_tensor(x: Tensor, dim: int, padding_size: int):
+    slc = [slice(None)] * len(x.shape)
+    slc[dim] = slice(0, -padding_size)
+    return x[slc]
+
+
+class Gather(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        local_input: Tensor,
+        dim: int,
+        grad_scale: Optional[bool] = False,
+    ) -> Tensor:
+        ctx.group = group
+        ctx.rank = dist.get_rank(group)
+        ctx.dim = dim
+        ctx.grad_scale = grad_scale
+        seq_world_size = dist.get_world_size(group)
+        ctx.seq_world_size = seq_world_size
+        dim_size = list(local_input.size())
+        split_size = dim_size[0]
+        ctx.part_size = dim_size[dim]
+        dim_size[0] = dim_size[0] * seq_world_size
+        output = torch.empty(dim_size, dtype=local_input.dtype, device=torch.cuda.current_device())
+        dist.all_gather_into_tensor(output, local_input.contiguous(), group=ctx.group)
+        return torch.cat(output.split(split_size), dim=dim)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor]:
+        if ctx.grad_scale:
+            grad_output = grad_output * ctx.seq_world_size
+        return (
+            None,
+            grad_output.split(ctx.part_size, dim=ctx.dim)[ctx.rank].contiguous(),
+            None,
+            None,
+        )
+
+
+def slice_tensor(tensor, dim, start, end):
+    indices = slice(start, end)
+    return tensor[(slice(None),) * dim + (indices,)]
+
+
+def init_model_shard_cpu_group(sharding_strategy: str, device_mesh: Optional[Tuple] = None):
+    """
+    Initialize CPU process group of model sharding.
+    """
+    global _MODEL_SHARD_CPU_GROUP
+    assert dist.is_initialized()
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    if device_mesh is not None:
+        num_shards_per_group = device_mesh[1]
+    elif "HYBRID" in sharding_strategy:
+        num_shards_per_group = min(8, world_size)
+    else:
+        num_shards_per_group = world_size
+    num_groups = world_size // num_shards_per_group
+    for i in range(num_groups):
+        start_rank = i * num_shards_per_group
+        end_rank = (i + 1) * num_shards_per_group
+        ranks = range(start_rank, end_rank)
+        cpu_group = dist.new_group(ranks, backend="gloo")
+        if rank in ranks:
+            _MODEL_SHARD_CPU_GROUP = cpu_group

humo/common/distributed/basic.py

@@ -0,0 +1,143 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+# Codes adapted from [SeedVR]
+# https://github.com/ByteDance-Seed/SeedVR/tree/main/common/distributed
+
+"""
+Distributed basic functions.
+"""
+
+import os
+import torch
+from torch import nn
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torch.distributed.fsdp._common_utils import _is_fsdp_flattened
+
+
+def get_global_rank() -> int:
+    """
+    Get the global rank, the global index of the GPU.
+    """
+    return int(os.environ.get("RANK", "0"))
+
+
+def get_local_rank() -> int:
+    """
+    Get the local rank, the local index of the GPU.
+    """
+    return int(os.environ.get("LOCAL_RANK", "0"))
+
+
+def get_world_size() -> int:
+    """
+    Get the world size, the total amount of GPUs.
+    """
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def get_device() -> torch.device:
+    """
+    Get current rank device.
+    """
+    return torch.device("cuda", get_local_rank())
+
+
+def barrier_if_distributed(*args, **kwargs):
+    """
+    Synchronizes all processes if under distributed context.
+    """
+    if dist.is_initialized():
+        return dist.barrier(*args, **kwargs)
+
+
+def init_torch(cudnn_benchmark=True):
+    """
+    Common PyTorch initialization configuration.
+    """
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = cudnn_benchmark
+    torch.cuda.set_device(get_local_rank())
+    dist.init_process_group(
+        backend="nccl",
+        rank=get_global_rank(),
+        world_size=get_world_size(),
+    )
+
+
+def convert_to_ddp(module: torch.nn.Module, **kwargs) -> DistributedDataParallel:
+    return DistributedDataParallel(
+        module=module,
+        device_ids=[get_local_rank()],
+        output_device=get_local_rank(),
+        **kwargs,
+    )
+
+
+def meta_param_init_fn(module: nn.Module) -> None:
+    """
+    Used for model inited onto meta device.
+    Init meta param/buffer with empty tensor.
+    We don't care numerical correctness in this func.
+    FSDP will sync param/buffer state from rank0 to the other ranks.
+    """
+
+    with torch.no_grad():
+        for submodule in module.modules():
+            for param_name, param in submodule.named_parameters(recurse=False):
+                if not _is_fsdp_flattened(param) and param.is_meta:
+                    materialized_param = nn.Parameter(torch.empty_like(param, device="cpu"))
+                    setattr(submodule, param_name, materialized_param)
+            for buffer_name, buffer in submodule.named_buffers(recurse=False):
+                if not _is_fsdp_flattened(buffer) and buffer.is_meta:
+                    materialized_param = torch.empty_like(buffer, device="cpu")
+                    setattr(submodule, buffer_name, materialized_param)
+            torch.cuda.empty_cache()
+
+
+def meta_non_persistent_buffer_init_fn(module: nn.Module) -> nn.Module:
+    """
+    Materialize meta device buffers that are not persistent in state_dict.
+    Handles special cases like RotaryEmbedding.freqs.
+    """
+    with torch.no_grad():
+        for submodule in module.modules():
+            if hasattr(submodule, "freqs"):
+                freqs = getattr(submodule, "freqs")
+                if isinstance(freqs, torch.Tensor) and freqs.is_meta:
+                    dim = submodule.dim
+                    def rope_params(max_seq_len, dim, theta=10000):
+                        assert dim % 2 == 0
+                        freqs = torch.outer(
+                            torch.arange(max_seq_len),
+                            1.0 / torch.pow(theta,
+                                torch.arange(0, dim, 2).to(torch.float32).div(dim)))
+                        freqs = torch.polar(torch.ones_like(freqs), freqs)
+                        return freqs
+                    
+                    dim = 5120  # 1536 
+                    num_heads = 40  # 12
+                    # dim = 1536 
+                    # num_heads = 12
+                    d = dim // num_heads
+                    freqs_tensor = torch.cat([
+                        rope_params(1024, d - 4 * (d // 6)),
+                        rope_params(1024, 2 * (d // 6)),
+                        rope_params(1024, 2 * (d // 6))
+                    ], dim=1).to(dtype=torch.cfloat, device="cpu")
+                    
+                    setattr(submodule, "freqs", freqs_tensor)
+                    print(f"Successfully materialized freqs for {submodule.__class__.__name__}")
+
+    assert not any(b.is_meta for n, b in module.named_buffers())
+    return module

humo/common/logger.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+# Codes adapted from [SeedVR]
+# https://github.com/ByteDance-Seed/SeedVR/blob/main/common/logger.py
+
+"""
+Logging utility functions.
+"""
+
+import logging
+import sys
+from typing import Optional
+
+from common.distributed import get_global_rank, get_local_rank, get_world_size
+
+_default_handler = logging.StreamHandler(sys.stdout)
+_default_handler.setFormatter(
+    logging.Formatter(
+        "%(asctime)s "
+        + (f"[Rank:{get_global_rank()}]" if get_world_size() > 1 else "")
+        + (f"[LocalRank:{get_local_rank()}]" if get_world_size() > 1 else "")
+        + "[%(threadName).12s][%(name)s][%(levelname).5s] "
+        + "%(message)s"
+    )
+)
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+    """
+    Get a logger.
+    """
+    logger = logging.getLogger(name)
+    logger.addHandler(_default_handler)
+    logger.setLevel(logging.INFO)
+    return logger

humo/configs/inference/generate.yaml

@@ -0,0 +1,78 @@
+__object__:
+  path: humo.generate
+  name: Generator
+
+dit:
+  model:
+    __inherit__: humo/configs/models/Wan_14B_I2V.yaml
+    __object__:
+      path: humo.models.wan_modules.model_humo
+      name: WanModel
+    insert_audio: True
+  zero_vae_path: ./weights/HuMo/zero_vae_129frame.pt
+  zero_vae_720p_path: ./weights/HuMo/zero_vae_720p_161frame.pt
+  checkpoint_dir: ./weights/HuMo/HuMo-17B
+  compile: False
+  init_with_meta_device: True
+  gradient_checkpoint: True
+  fsdp:
+    sharding_strategy: _HYBRID_SHARD_ZERO2
+  sp_size: 1
+  dtype: bfloat16
+
+vae:
+  checkpoint: ./weights/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth
+  vae_stride: [ 4, 8, 8 ]
+  scaling_factor: 0.9152
+  compile: False
+  grouping: True
+  use_sample: False
+  dtype: bfloat16
+
+text:
+  t5_checkpoint: ./weights/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth
+  t5_tokenizer: ./weights/Wan2.1-T2V-1.3B/google/umt5-xxl
+  dropout: 0.1
+  dtype: bfloat16
+  fsdp:
+    enabled: True
+    sharding_strategy: HYBRID_SHARD
+
+diffusion:
+  schedule:
+    type: lerp
+    T: 1000.0
+  sampler:
+    type: euler
+    prediction_type: v_lerp
+  timesteps:
+    training:
+      type: logitnormal
+      loc: 0.0
+      scale: 1.0
+    sampling:
+      type: uniform_trailing
+      steps: 50
+      shift: 5.0
+
+audio:
+  vocal_separator: ./weights/HuMo/audio_separator/Kim_Vocal_2.onnx
+  wav2vec_model: ./weights/whisper-large-v3
+
+generation:
+  mode: "TIA"  # TA, TIA
+  extract_audio_feat: True
+  seed: 666666
+  frames: 97
+  fps: 25
+  height: 480 # 720 480
+  width: 832 # 1280 832
+  batch_size: 1
+  sequence_parallel: 8
+  output:
+    dir: ./output
+  # positive_prompt: ./examples/test_case.json
+  sample_neg_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+  scale_a: 5.5
+  scale_t: 5.0
+  step_change: 980

humo/configs/inference/generate_1_7B.yaml

@@ -0,0 +1,76 @@
+__object__:
+  path: humo.generate_1_7B
+  name: Generator
+
+dit:
+  model:
+    __inherit__: humo/configs/models/Wan_1.3B.yaml
+    __object__:
+      path: humo.models.wan_modules.model_humo
+      name: WanModel
+    insert_audio: True
+  zero_vae_path: ./weights/HuMo/zero_vae_129frame.pt
+  zero_vae_720p_path: ./weights/HuMo/zero_vae_720p_161frame.pt
+  checkpoint_dir: ./weights/HuMo/HuMo-1.7B/ema.pth #./weights/HuMo/HuMo-17B
+  compile: False
+  init_with_meta_device: True
+  gradient_checkpoint: True
+  fsdp:
+    sharding_strategy: _HYBRID_SHARD_ZERO2
+  sp_size: 1
+
+vae:
+  checkpoint: ./weights/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth
+  vae_stride: [ 4, 8, 8 ]
+  scaling_factor: 0.9152
+  compile: False
+  grouping: True
+  use_sample: False
+  dtype: bfloat16
+
+text:
+  t5_checkpoint: ./weights/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth
+  t5_tokenizer: ./weights/Wan2.1-T2V-1.3B/google/umt5-xxl
+  dropout: 0.1
+  dtype: bfloat16
+  fsdp:
+    enabled: True
+    sharding_strategy: HYBRID_SHARD
+
+diffusion:
+  schedule:
+    type: lerp
+    T: 1000.0
+  sampler:
+    type: euler
+    prediction_type: v_lerp
+  timesteps:
+    training:
+      type: logitnormal
+      loc: 0.0
+      scale: 1.0
+    sampling:
+      type: uniform_trailing
+      steps: 50
+      shift: 5.0
+
+audio:
+  vocal_separator: ./weights/audio_separator/Kim_Vocal_2.onnx
+  wav2vec_model: ./weights/whisper-large-v3
+
+generation:
+  mode: "TIA"  # TA, TIA
+  extract_audio_feat: True
+  seed: 666666
+  frames: 97
+  fps: 25
+  height: 720 # 480
+  width: 1280 # 832
+  batch_size: 1
+  output:
+    dir: ./output
+  sample_neg_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+  scale_t: 7.5
+  scale_i: 4.0
+  scale_a: 7.5
+  # step_change: 980

humo/configs/models/Wan_1.3B.yaml

@@ -0,0 +1,17 @@
+__object__:
+  path: ???
+  name: ???
+  args: as_params
+
+text_len: 512
+patch_size: [ 1, 2, 2 ]
+dim: 1536
+ffn_dim: 8960
+freq_dim: 256
+model_type: "t2v"
+num_heads: 12
+num_layers: 30
+window_size: [ -1, -1 ]
+qk_norm: True
+cross_attn_norm: True
+eps: 1e-6

humo/configs/models/Wan_1.3B_I2V.yaml

@@ -0,0 +1,18 @@
+__object__:
+  path: ???
+  name: ???
+  args: as_params
+
+text_len: 512
+patch_size: [ 1, 2, 2 ]
+dim: 1536
+ffn_dim: 8960
+freq_dim: 256
+in_dim: 36
+model_type: "i2v"
+num_heads: 12
+num_layers: 30
+window_size: [ -1, -1 ]
+qk_norm: True
+cross_attn_norm: True
+eps: 1e-6

humo/configs/models/Wan_14B.yaml

@@ -0,0 +1,17 @@
+__object__:
+  path: ???
+  name: ???
+  args: as_params
+
+text_len: 512
+patch_size: [ 1, 2, 2 ]
+dim: 5120
+ffn_dim: 13824
+freq_dim: 256
+model_type: "t2v"
+num_heads: 40
+num_layers: 40
+window_size: [ -1, -1 ]
+qk_norm: True
+cross_attn_norm: True
+eps: 1e-6

humo/configs/models/Wan_14B_I2V.yaml

@@ -0,0 +1,18 @@
+__object__:
+  path: ???
+  name: ???
+  args: as_params
+
+text_len: 512
+patch_size: [ 1, 2, 2 ]
+dim: 5120
+ffn_dim: 13824
+freq_dim: 256
+in_dim: 36
+model_type: "i2v"
+num_heads: 40
+num_layers: 40
+window_size: [ -1, -1 ]
+qk_norm: True
+cross_attn_norm: True
+eps: 1e-6

humo/generate.py

@@ -0,0 +1,984 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+# Inference codes adapted from [SeedVR]
+# https://github.com/ByteDance-Seed/SeedVR/blob/main/projects/inference_seedvr2_7b.py
+
+import math
+import os
+import gc
+import random
+import sys
+import mediapy
+import numpy as np
+import torch
+import torch.distributed as dist
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from einops import rearrange
+from omegaconf import OmegaConf
+from PIL import Image, ImageOps
+from torchvision.transforms import ToTensor
+from tqdm import tqdm
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import (
+    BackwardPrefetch,
+    FullyShardedDataParallel,
+    MixedPrecision,
+    ShardingStrategy,
+)
+from common.distributed import (
+    get_device,
+    get_global_rank,
+    get_local_rank,
+    meta_param_init_fn,
+    meta_non_persistent_buffer_init_fn,
+    init_torch,
+)
+from common.distributed.advanced import (
+    init_unified_parallel,
+    get_unified_parallel_world_size,
+    get_sequence_parallel_rank,
+    init_model_shard_cpu_group,
+)
+from common.logger import get_logger
+from common.config import create_object
+from common.distributed import get_device, get_global_rank
+from torchvision.transforms import Compose, Normalize, ToTensor
+from humo.models.wan_modules.t5 import T5EncoderModel
+from humo.models.wan_modules.vae import WanVAE
+from humo.models.utils.utils import tensor_to_video, prepare_json_dataset
+from contextlib import contextmanager
+import torch.cuda.amp as amp
+from humo.models.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from humo.utils.audio_processor_whisper import AudioProcessor
+from humo.utils.wav2vec import linear_interpolation_fps
+from torchao.quantization import quantize_
+
+import torch._dynamo as dynamo
+dynamo.config.capture_scalar_outputs = True
+torch.set_float32_matmul_precision("high")
+
+import torch
+import torch.nn as nn
+import transformer_engine.pytorch as te
+
+image_transform = Compose([
+    ToTensor(),
+    Normalize(mean=0.5, std=0.5),
+])
+
+SIZE_CONFIGS = {
+    '720*1280': (720, 1280),
+    '1280*720': (1280, 720),
+    '480*832': (480, 832),
+    '832*480': (832, 480),
+    '1024*1024': (1024, 1024),
+}
+
+def clever_format(nums, format="%.2f"):
+    from typing import Iterable
+    if not isinstance(nums, Iterable):
+        nums = [nums]
+    clever_nums = []
+    for num in nums:
+        if num > 1e12:
+            clever_nums.append(format % (num / 1e12) + "T")
+        elif num > 1e9:
+            clever_nums.append(format % (num / 1e9) + "G")
+        elif num > 1e6:
+            clever_nums.append(format % (num / 1e6) + "M")
+        elif num > 1e3:
+            clever_nums.append(format % (num / 1e3) + "K")
+        else:
+            clever_nums.append(format % num + "B")
+
+    clever_nums = clever_nums[0] if len(clever_nums) == 1 else (*clever_nums,)
+
+    return clever_nums
+
+
+    
+# --- put near your imports ---
+import torch
+import torch.nn as nn
+import contextlib
+import transformer_engine.pytorch as te
+
+# FP8 autocast compatibility for different TE versions
+try:
+    # Preferred modern API
+    from transformer_engine.pytorch import fp8_autocast
+    try:
+        # Newer TE: use recipe-based API
+        from transformer_engine.common.recipe import DelayedScaling, Format
+        def make_fp8_ctx(enabled: bool = True):
+            if not enabled:
+                return contextlib.nullcontext()
+            fp8_recipe = DelayedScaling(fp8_format=Format.E4M3)  # E4M3 format
+            return fp8_autocast(enabled=True, fp8_recipe=fp8_recipe)
+    except Exception:
+        # Very old variant that might still accept fp8_format directly
+        def make_fp8_ctx(enabled: bool = True):
+            # If TE doesn't have FP8Format, just no-op
+            if not hasattr(te, "FP8Format"):
+                return contextlib.nullcontext()
+            return te.fp8_autocast(enabled=enabled, fp8_format=te.FP8Format.E4M3)
+except Exception:
+    # TE not present or totally incompatible — no-op
+    def make_fp8_ctx(enabled: bool = True):
+        return contextlib.nullcontext()
+
+
+# TE sometimes exposes Linear at different paths; this normalizes it.
+try:
+    TELinear = te.Linear
+except AttributeError:  # very old layouts
+    from transformer_engine.pytorch.modules.linear import Linear as TELinear  # type: ignore
+
+# --- near imports ---
+import torch
+import torch.nn as nn
+import transformer_engine.pytorch as te
+
+try:
+    TELinear = te.Linear
+except AttributeError:
+    from transformer_engine.pytorch.modules.linear import Linear as TELinear  # type: ignore
+
+import torch
+import torch.nn as nn
+import transformer_engine.pytorch as te
+
+try:
+    TELinear = te.Linear
+except AttributeError:
+    from transformer_engine.pytorch.modules.linear import Linear as TELinear  # type: ignore
+
+def _default_te_allow(fullname: str, lin: nn.Linear) -> bool:
+    """
+    Allow TE only where it's shape-safe & beneficial.
+    Skip small/special layers (time/timestep/pos embeds, heads).
+    Enforce multiples of 16 for in/out features (FP8 kernel friendly).
+    Also skip very small projections likely to see M=1.
+    """
+    blocked_keywords = (
+        "time_embedding", "timestep", "time_embed",
+        "time_projection", "pos_embedding", "pos_embed",
+        "to_logits", "logits", "final_proj", "proj_out", "output_projection",
+    )
+    if any(k in fullname for k in blocked_keywords):
+        return False
+
+    # TE FP8 kernels like K, N divisible by 16
+    if lin.in_features % 16 != 0 or lin.out_features % 16 != 0:
+        return False
+
+    # Heuristic: avoid tiny layers; keeps attention/MLP, skips small MLPs
+    if lin.in_features < 512 or lin.out_features < 512:
+        return False
+
+    # Whitelist: only convert inside transformer blocks if you know their prefix
+    # This further reduces risk of catching special heads elsewhere.
+    allowed_context = ("blocks", "layers", "transformer", "attn", "mlp", "ffn")
+    if not any(tok in fullname for tok in allowed_context):
+        return False
+
+    return True
+
+@torch.no_grad()
+def convert_linears_to_te_fp8(module: nn.Module, allow_pred=_default_te_allow, _prefix=""):
+    for name, child in list(module.named_children()):
+        full = f"{_prefix}.{name}" if _prefix else name
+        convert_linears_to_te_fp8(child, allow_pred, full)
+
+        if isinstance(child, nn.Linear):
+            if allow_pred is not None and not allow_pred(full, child):
+                continue
+
+            te_lin = TELinear(
+                in_features=child.in_features,
+                out_features=child.out_features,
+                bias=(child.bias is not None),
+                params_dtype=torch.bfloat16,
+            ).to(child.weight.device)
+
+            te_lin.weight.copy_(child.weight.to(te_lin.weight.dtype))
+            if child.bias is not None:
+                te_lin.bias.copy_(child.bias.to(te_lin.bias.dtype))
+
+            setattr(module, name, te_lin)
+    return module
+
+class Generator():
+    def __init__(self, config: DictConfig):
+        self.config = config.copy()
+        OmegaConf.set_readonly(self.config, True)
+        self.logger = get_logger(self.__class__.__name__)
+        
+        # init_torch(cudnn_benchmark=False)
+        self.configure_models()
+
+    def entrypoint(self):
+        
+        self.inference_loop()
+    
+    def get_fsdp_sharding_config(self, sharding_strategy, device_mesh_config):
+        device_mesh = None
+        fsdp_strategy = ShardingStrategy[sharding_strategy]
+        if (
+            fsdp_strategy in [ShardingStrategy._HYBRID_SHARD_ZERO2, ShardingStrategy.HYBRID_SHARD]
+            and device_mesh_config is not None
+        ):
+            device_mesh = init_device_mesh("cuda", tuple(device_mesh_config))
+        return device_mesh, fsdp_strategy
+
+        
+    def configure_models(self):
+        self.configure_dit_model(device="cuda")
+
+        self.dit.eval().to("cuda")
+        convert_linears_to_te_fp8(self.dit)
+
+        self.dit = torch.compile(self.dit, )
+
+
+        self.configure_vae_model(device="cuda")
+        if self.config.generation.get('extract_audio_feat', False):
+            self.configure_wav2vec(device="cpu")
+        self.configure_text_model(device="cuda")
+
+        # # Initialize fsdp.
+        # self.configure_dit_fsdp_model()
+        # self.configure_text_fsdp_model()
+
+        # quantize_(self.text_encoder, Int8WeightOnlyConfig())
+        # quantize_(self.dit, Float8DynamicActivationFloat8WeightConfig())
+
+    
+    def configure_dit_model(self, device=get_device()):
+
+        init_unified_parallel(self.config.dit.sp_size)
+        self.sp_size = get_unified_parallel_world_size()
+
+        # Create DiT model on meta, then mark dtype as bfloat16 (no real allocation yet).
+        init_device = "meta"
+        with torch.device(init_device):
+            self.dit = create_object(self.config.dit.model)
+            self.dit = self.dit.to(dtype=torch.bfloat16)  # or: self.dit.bfloat16()
+        self.logger.info(f"Load DiT model on {init_device}.")
+        self.dit.eval().requires_grad_(False)
+
+        # Load dit checkpoint.
+        path = self.config.dit.checkpoint_dir
+
+        def _cast_state_dict_to_bf16(state):
+            for k, v in state.items():
+                if isinstance(v, torch.Tensor) and v.is_floating_point():
+                    state[k] = v.to(dtype=torch.bfloat16, copy=False)
+            return state
+
+        if path.endswith(".pth"):
+            # Load to CPU first; we’ll move the model later.
+            state = torch.load(path, map_location="cpu", mmap=True)
+            state = _cast_state_dict_to_bf16(state)
+            missing_keys, unexpected_keys = self.dit.load_state_dict(state, strict=False, assign=True)
+            self.logger.info(
+                f"dit loaded from {path}. Missing keys: {len(missing_keys)}, Unexpected keys: {len(unexpected_keys)}"
+            )
+        else:
+            from safetensors.torch import load_file
+            import json
+            def load_custom_sharded_weights(model_dir, base_name):
+                index_path = f"{model_dir}/{base_name}.safetensors.index.json"
+                with open(index_path, "r") as f:
+                    index = json.load(f)
+                weight_map = index["weight_map"]
+                shard_files = set(weight_map.values())
+                state_dict = {}
+                for shard_file in shard_files:
+                    shard_path = f"{model_dir}/{shard_file}"
+                    # Load on CPU, then cast to bf16; we’ll move the whole module later.
+                    shard_state = load_file(shard_path, device="cpu")
+                    shard_state = {k: (v.to(dtype=torch.bfloat16, copy=False) if v.is_floating_point() else v)
+                                for k, v in shard_state.items()}
+                    state_dict.update(shard_state)
+                return state_dict
+
+            state = load_custom_sharded_weights(path, 'humo')
+            self.dit.load_state_dict(state, strict=False, assign=True)
+
+        self.dit = meta_non_persistent_buffer_init_fn(self.dit)
+
+        target_device = get_device() if device in [get_device(), "cuda"] else device
+        self.dit.to(target_device)  # dtype already bf16
+
+        # Print model size.
+        params = sum(p.numel() for p in self.dit.parameters())
+        self.logger.info(
+            f"[RANK:{get_global_rank()}] DiT Parameters: {clever_format(params, '%.3f')}"
+        )
+
+        
+    def configure_vae_model(self, device=get_device()):
+        self.vae_stride = self.config.vae.vae_stride
+        self.vae = WanVAE(
+            vae_pth=self.config.vae.checkpoint,
+            device=device)
+        
+        if self.config.generation.height == 480:
+            self.zero_vae = torch.load(self.config.dit.zero_vae_path)
+        elif self.config.generation.height == 720:
+            self.zero_vae = torch.load(self.config.dit.zero_vae_720p_path)
+        else:
+            raise ValueError(f"Unsupported height {self.config.generation.height} for zero-vae.")
+    
+    def configure_wav2vec(self, device=get_device()):
+        audio_separator_model_file = self.config.audio.vocal_separator
+        wav2vec_model_path = self.config.audio.wav2vec_model
+
+        self.audio_processor = AudioProcessor(
+            16000,
+            25,
+            wav2vec_model_path,
+            "all",
+            audio_separator_model_file,
+            None,  # not seperate
+            os.path.join(self.config.generation.output.dir, "vocals"),
+            device=device,
+        )
+
+    def configure_text_model(self, device=get_device()):
+        self.text_encoder = T5EncoderModel(
+            text_len=self.config.dit.model.text_len,
+            dtype=torch.bfloat16,
+            device=device,
+            checkpoint_path=self.config.text.t5_checkpoint,
+            tokenizer_path=self.config.text.t5_tokenizer,
+            )
+
+    
+    def configure_dit_fsdp_model(self):
+        from humo.models.wan_modules.model_humo import WanAttentionBlock
+
+        dit_blocks = (WanAttentionBlock,)
+
+        # Init model_shard_cpu_group for saving checkpoint with sharded state_dict.
+        init_model_shard_cpu_group(
+            self.config.dit.fsdp.sharding_strategy,
+            self.config.dit.fsdp.get("device_mesh", None),
+        )
+
+        # Assert that dit has wrappable blocks.
+        assert any(isinstance(m, dit_blocks) for m in self.dit.modules())
+
+        # Define wrap policy on all dit blocks.
+        def custom_auto_wrap_policy(module, recurse, *args, **kwargs):
+            return recurse or isinstance(module, dit_blocks)
+
+        # Configure FSDP settings.
+        device_mesh, fsdp_strategy = self.get_fsdp_sharding_config(
+            self.config.dit.fsdp.sharding_strategy,
+            self.config.dit.fsdp.get("device_mesh", None),
+        )
+        settings = dict(
+            auto_wrap_policy=custom_auto_wrap_policy,
+            sharding_strategy=fsdp_strategy,
+            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+            device_id=get_local_rank(),
+            use_orig_params=False,
+            sync_module_states=True,
+            forward_prefetch=True,
+            limit_all_gathers=False,  # False for ZERO2.
+            mixed_precision=MixedPrecision(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.float32,
+                buffer_dtype=torch.float32,
+            ),
+            device_mesh=device_mesh,
+            param_init_fn=meta_param_init_fn,
+        )
+
+        # Apply FSDP.
+        self.dit = FullyShardedDataParallel(self.dit, **settings)
+        # self.dit.to(get_device())
+
+
+    def configure_text_fsdp_model(self):
+        # If FSDP is not enabled, put text_encoder to GPU and return.
+        if not self.config.text.fsdp.enabled:
+            self.text_encoder.to(get_device())
+            return
+
+        # from transformers.models.t5.modeling_t5 import T5Block
+        from humo.models.wan_modules.t5 import T5SelfAttention
+
+        text_blocks = (torch.nn.Embedding, T5SelfAttention)
+        # text_blocks_names = ("QWenBlock", "QWenModel")  # QWen cannot be imported. Use str.
+
+        def custom_auto_wrap_policy(module, recurse, *args, **kwargs):
+            return (
+                recurse
+                or isinstance(module, text_blocks)
+            )
+
+        # Apply FSDP.
+        text_encoder_dtype = getattr(torch, self.config.text.dtype)
+        device_mesh, fsdp_strategy = self.get_fsdp_sharding_config(
+            self.config.text.fsdp.sharding_strategy,
+            self.config.text.fsdp.get("device_mesh", None),
+        )
+        self.text_encoder = FullyShardedDataParallel(
+            module=self.text_encoder,
+            auto_wrap_policy=custom_auto_wrap_policy,
+            sharding_strategy=fsdp_strategy,
+            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+            device_id=get_local_rank(),
+            use_orig_params=False,
+            sync_module_states=False,
+            forward_prefetch=True,
+            limit_all_gathers=True,
+            mixed_precision=MixedPrecision(
+                param_dtype=text_encoder_dtype,
+                reduce_dtype=text_encoder_dtype,
+                buffer_dtype=text_encoder_dtype,
+            ),
+            device_mesh=device_mesh,
+        )
+        self.text_encoder.to(get_device()).requires_grad_(False)
+
+
+    def load_image_latent_ref_id(self, path: str, size, device):
+        # Load size.
+        h, w = size[1], size[0]
+
+        # Load image.
+        if len(path) > 1 and not isinstance(path, str):
+            ref_vae_latents = []
+            for image_path in path:
+                with Image.open(image_path) as img:
+                    img = img.convert("RGB")
+
+                    # Calculate the required size to keep aspect ratio and fill the rest with padding.
+                    img_ratio = img.width / img.height
+                    target_ratio = w / h
+                    
+                    if img_ratio > target_ratio:  # Image is wider than target
+                        new_width = w
+                        new_height = int(new_width / img_ratio)
+                    else:  # Image is taller than target
+                        new_height = h
+                        new_width = int(new_height * img_ratio)
+                    
+                    # img = img.resize((new_width, new_height), Image.ANTIALIAS)
+                    img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+                    # Create a new image with the target size and place the resized image in the center
+                    delta_w = w - img.size[0]
+                    delta_h = h - img.size[1]
+                    padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+                    new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))
+
+                    # Transform to tensor and normalize.
+                    transform = Compose(
+                        [
+                            ToTensor(),
+                            Normalize(0.5, 0.5),
+                        ]
+                    )
+                    new_img = transform(new_img)
+                    # img_vae_latent = self.vae_encode([new_img.unsqueeze(1)])[0]
+                    img_vae_latent = self.vae.encode([new_img.unsqueeze(1)], device)
+                    ref_vae_latents.append(img_vae_latent[0])
+
+            return [torch.cat(ref_vae_latents, dim=1)]
+        else:
+            if not isinstance(path, str):
+                path = path[0]
+            with Image.open(path) as img:
+                img = img.convert("RGB")
+
+                # Calculate the required size to keep aspect ratio and fill the rest with padding.
+                img_ratio = img.width / img.height
+                target_ratio = w / h
+                
+                if img_ratio > target_ratio:  # Image is wider than target
+                    new_width = w
+                    new_height = int(new_width / img_ratio)
+                else:  # Image is taller than target
+                    new_height = h
+                    new_width = int(new_height * img_ratio)
+                
+                # img = img.resize((new_width, new_height), Image.ANTIALIAS)
+                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+                # Create a new image with the target size and place the resized image in the center
+                delta_w = w - img.size[0]
+                delta_h = h - img.size[1]
+                padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+                new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))
+
+                # Transform to tensor and normalize.
+                transform = Compose(
+                    [
+                        ToTensor(),
+                        Normalize(0.5, 0.5),
+                    ]
+                )
+                new_img = transform(new_img)
+                img_vae_latent = self.vae.encode([new_img.unsqueeze(1)], device)
+
+            # Vae encode.
+            return img_vae_latent
+    
+    def get_audio_emb_window(self, audio_emb, frame_num, frame0_idx, audio_shift=2):
+        zero_audio_embed = torch.zeros((audio_emb.shape[1], audio_emb.shape[2]), dtype=audio_emb.dtype, device=audio_emb.device)
+        zero_audio_embed_3 = torch.zeros((3, audio_emb.shape[1], audio_emb.shape[2]), dtype=audio_emb.dtype, device=audio_emb.device)  # device=audio_emb.device
+        iter_ = 1 + (frame_num - 1) // 4
+        audio_emb_wind = []
+        for lt_i in range(iter_):
+            if lt_i == 0:
+                st = frame0_idx + lt_i - 2
+                ed = frame0_idx + lt_i + 3
+                wind_feat = torch.stack([
+                    audio_emb[i] if (0 <= i < audio_emb.shape[0]) else zero_audio_embed
+                    for i in range(st, ed)
+                ], dim=0)
+                wind_feat = torch.cat((zero_audio_embed_3, wind_feat), dim=0)
+            else:
+                st = frame0_idx + 1 + 4 * (lt_i - 1) - audio_shift
+                ed = frame0_idx + 1 + 4 * lt_i + audio_shift
+                wind_feat = torch.stack([
+                    audio_emb[i] if (0 <= i < audio_emb.shape[0]) else zero_audio_embed
+                    for i in range(st, ed)
+                ], dim=0)
+            audio_emb_wind.append(wind_feat)
+        audio_emb_wind = torch.stack(audio_emb_wind, dim=0)
+
+        return audio_emb_wind, ed - audio_shift
+    
+    def audio_emb_enc(self, audio_emb, wav_enc_type="whisper"):
+        if wav_enc_type == "wav2vec":
+            feat_merge = audio_emb
+        elif wav_enc_type == "whisper":
+            feat0 = linear_interpolation_fps(audio_emb[:, :, 0: 8].mean(dim=2), 50, 25)
+            feat1 = linear_interpolation_fps(audio_emb[:, :, 8: 16].mean(dim=2), 50, 25)
+            feat2 = linear_interpolation_fps(audio_emb[:, :, 16: 24].mean(dim=2), 50, 25)
+            feat3 = linear_interpolation_fps(audio_emb[:, :, 24: 32].mean(dim=2), 50, 25)
+            feat4 = linear_interpolation_fps(audio_emb[:, :, 32], 50, 25)
+            feat_merge = torch.stack([feat0, feat1, feat2, feat3, feat4], dim=2)[0]
+        else:
+            raise ValueError(f"Unsupported wav_enc_type: {wav_enc_type}")
+        
+        return feat_merge
+    
+    def parse_output(self, output):
+        latent = output[0]
+        mask = None
+        return latent, mask
+    
+    def forward_tia(self, latents, timestep, t, step_change, arg_tia, arg_ti, arg_i, arg_null):
+        pos_tia, _ = self.parse_output(self.dit(
+            latents, t=timestep, **arg_tia
+            ))
+        torch.cuda.empty_cache()
+
+        pos_ti, _ = self.parse_output(self.dit(
+            latents, t=timestep, **arg_ti
+            ))
+        torch.cuda.empty_cache()
+
+        if t > step_change:
+            neg, _ = self.parse_output(self.dit(
+                latents, t=timestep, **arg_i
+                ))  # img included in null, same with official Wan-2.1
+            torch.cuda.empty_cache()
+
+            noise_pred = self.config.generation.scale_a * (pos_tia - pos_ti) + \
+                    self.config.generation.scale_t * (pos_ti - neg) + \
+                    neg
+        else:
+            neg, _ = self.parse_output(self.dit(
+                latents, t=timestep, **arg_null
+                ))  # img not included in null
+            torch.cuda.empty_cache()
+
+            noise_pred = self.config.generation.scale_a * (pos_tia - pos_ti) + \
+                    (self.config.generation.scale_t - 2.0) * (pos_ti - neg) + \
+                    neg
+        return noise_pred
+    
+    def forward_ti(self, latents, timestep, t, step_change, arg_ti, arg_t, arg_i, arg_null):
+        # Positive with text+image (no audio)
+        pos_ti, _ = self.parse_output(self.dit(
+            latents, t=timestep, **arg_ti
+        ))
+        torch.cuda.empty_cache()
+
+        # Positive with text only (no image, no audio)
+        pos_t, _ = self.parse_output(self.dit(
+            latents, t=timestep, **arg_t
+        ))
+        torch.cuda.empty_cache()
+
+        # Negative branch: before step_change, don't include image in null; after, include image (like Wan-2.1)
+        if t > step_change:
+            neg, _ = self.parse_output(self.dit(
+                latents, t=timestep, **arg_i
+            ))  # img included in null
+        else:
+            neg, _ = self.parse_output(self.dit(
+                latents, t=timestep, **arg_null
+            ))  # img NOT included in null
+        torch.cuda.empty_cache()
+
+        # Guidance blend: replace "scale_a" below with "scale_i" if you add a separate image scale in config
+        noise_pred = self.config.generation.scale_a * (pos_ti - pos_t) + \
+                    self.config.generation.scale_t * (pos_t - neg) + \
+                    neg
+        return noise_pred
+
+    def forward_ta(self, latents, timestep, arg_ta, arg_t, arg_null):
+        pos_ta, _ = self.parse_output(self.dit(
+            latents, t=timestep, **arg_ta
+            ))
+        torch.cuda.empty_cache()
+
+        pos_t, _ = self.parse_output(self.dit(
+            latents, t=timestep, **arg_t
+            ))
+        torch.cuda.empty_cache()
+
+        neg, _ = self.parse_output(self.dit(
+                latents, t=timestep, **arg_null
+                ))
+        torch.cuda.empty_cache()
+            
+        noise_pred = self.config.generation.scale_a * (pos_ta - pos_t) + \
+                self.config.generation.scale_t * (pos_t - neg) + \
+                neg
+        return noise_pred
+                    
+    @torch.no_grad()
+    def inference(self,
+                 input_prompt,
+                 img_path,
+                 audio_path,
+                 size=(1280, 720),
+                 frame_num=81,
+                 shift=5.0,
+                 sample_solver='unipc',
+                 inference_mode='TIA',
+                 sampling_steps=50,
+                 n_prompt="",
+                 seed=-1,
+                 tea_cache_l1_thresh = 0.0,
+                 device = get_device(),
+        ):
+
+        print("inference started")
+
+        # self.vae.model.to(device=device)
+        if img_path is not None:
+            latents_ref = self.load_image_latent_ref_id(img_path, size, device)
+        else:
+            latents_ref = [torch.zeros(16, 1, size[1]//8, size[0]//8).to(device)]
+            
+        # self.vae.model.to(device="cpu")
+
+        print("vae finished")
+
+        latents_ref_neg = [torch.zeros_like(latent_ref) for latent_ref in latents_ref]
+        
+        # audio
+        if audio_path is not None:
+            if self.config.generation.extract_audio_feat:
+                self.audio_processor.whisper.to(device=device)
+                audio_emb, audio_length = self.audio_processor.preprocess(audio_path)
+                self.audio_processor.whisper.to(device='cpu')
+            else:
+                audio_emb_path = audio_path.replace(".wav", ".pt")
+                audio_emb = torch.load(audio_emb_path).to(device=device)
+                audio_emb = self.audio_emb_enc(audio_emb, wav_enc_type="whisper")
+                self.logger.info("使用预先提取好的音频特征: %s", audio_emb_path)
+        else:
+            audio_emb = torch.zeros(frame_num, 5, 1280).to(device)
+            
+        frame_num = frame_num if frame_num != -1 else audio_length
+        frame_num = 4 * ((frame_num - 1) // 4) + 1
+        audio_emb, _ = self.get_audio_emb_window(audio_emb, frame_num, frame0_idx=0)
+        zero_audio_pad = torch.zeros(latents_ref[0].shape[1], *audio_emb.shape[1:]).to(audio_emb.device)
+        audio_emb = torch.cat([audio_emb, zero_audio_pad], dim=0)
+        audio_emb = [audio_emb.to(device)]
+        audio_emb_neg = [torch.zeros_like(audio_emb[0])]
+        
+        # preprocess
+        self.patch_size = self.config.dit.model.patch_size
+        F = frame_num
+        target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1 + latents_ref[0].shape[1],
+                        size[1] // self.vae_stride[1],
+                        size[0] // self.vae_stride[2])
+
+        seq_len = math.ceil((target_shape[2] * target_shape[3]) /
+                            (self.patch_size[1] * self.patch_size[2]) *
+                            target_shape[1] / self.sp_size) * self.sp_size
+
+        if n_prompt == "":
+            n_prompt = self.config.generation.sample_neg_prompt
+        seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
+        seed_g = torch.Generator(device=device)
+        seed_g.manual_seed(seed)
+
+        # self.text_encoder.model.to(device)
+        context = self.text_encoder([input_prompt], device)
+        context_null = self.text_encoder([n_prompt], device)
+        # self.text_encoder.model.cpu()
+
+        print("text encoder finished")
+
+        noise = [
+            torch.randn(
+                target_shape[0],
+                target_shape[1], # - latents_ref[0].shape[1],
+                target_shape[2],
+                target_shape[3],
+                dtype=torch.float32,
+                device=device,
+                generator=seed_g)
+        ]
+
+        @contextmanager
+        def noop_no_sync():
+            yield
+
+        no_sync = getattr(self.dit, 'no_sync', noop_no_sync)
+        step_change = self.config.generation.step_change # 980
+
+        # evaluation mode
+        with make_fp8_ctx(True), torch.autocast('cuda', dtype=torch.bfloat16), torch.no_grad(), no_sync():
+
+            if sample_solver == 'unipc':
+                sample_scheduler = FlowUniPCMultistepScheduler(
+                    num_train_timesteps=1000,
+                    shift=1,
+                    use_dynamic_shifting=False)
+                sample_scheduler.set_timesteps(
+                    sampling_steps, device=device, shift=shift)
+                timesteps = sample_scheduler.timesteps
+
+            # sample videos
+            latents = noise
+
+            msk = torch.ones(4, target_shape[1], target_shape[2], target_shape[3], device=get_device())
+            msk[:,:-latents_ref[0].shape[1]] = 0
+
+            zero_vae = self.zero_vae[:, :(target_shape[1]-latents_ref[0].shape[1])].to(
+                device=get_device(), dtype=latents_ref[0].dtype)
+            y_c = torch.cat([
+                zero_vae,
+                latents_ref[0]
+                ], dim=1)
+            y_c = [torch.concat([msk, y_c])]
+
+            y_null = self.zero_vae[:, :target_shape[1]].to(
+                device=get_device(), dtype=latents_ref[0].dtype)
+            y_null = [torch.concat([msk, y_null])]
+
+            tea_cache_l1_thresh = tea_cache_l1_thresh
+            tea_cache_model_id = "Wan2.1-T2V-14B"
+
+            arg_null = {'seq_len': seq_len, 'audio': audio_emb_neg, 'y': y_null, 'context': context_null, "tea_cache": TeaCache(sampling_steps, rel_l1_thresh=tea_cache_l1_thresh, model_id=tea_cache_model_id) if tea_cache_l1_thresh is not None and tea_cache_l1_thresh > 0 else None}
+            arg_t = {'seq_len': seq_len, 'audio': audio_emb_neg, 'y': y_null, 'context': context, "tea_cache": TeaCache(sampling_steps, rel_l1_thresh=tea_cache_l1_thresh, model_id=tea_cache_model_id) if tea_cache_l1_thresh is not None and tea_cache_l1_thresh > 0 else None}
+            arg_i = {'seq_len': seq_len, 'audio': audio_emb_neg, 'y': y_c, 'context': context_null, "tea_cache": TeaCache(sampling_steps, rel_l1_thresh=tea_cache_l1_thresh, model_id=tea_cache_model_id) if tea_cache_l1_thresh is not None and tea_cache_l1_thresh > 0 else None}
+            arg_ti = {'seq_len': seq_len, 'audio': audio_emb_neg, 'y': y_c, 'context': context, "tea_cache": TeaCache(sampling_steps, rel_l1_thresh=tea_cache_l1_thresh, model_id=tea_cache_model_id) if tea_cache_l1_thresh is not None and tea_cache_l1_thresh > 0 else None}
+            arg_ta = {'seq_len': seq_len, 'audio': audio_emb, 'y': y_null, 'context': context, "tea_cache": TeaCache(sampling_steps, rel_l1_thresh=tea_cache_l1_thresh, model_id=tea_cache_model_id) if tea_cache_l1_thresh is not None and tea_cache_l1_thresh > 0 else None}
+            arg_tia = {'seq_len': seq_len, 'audio': audio_emb, 'y': y_c, 'context': context, "tea_cache": TeaCache(sampling_steps, rel_l1_thresh=tea_cache_l1_thresh, model_id=tea_cache_model_id) if tea_cache_l1_thresh is not None and tea_cache_l1_thresh > 0 else None}
+            
+            torch.cuda.empty_cache()
+            # self.dit.to(device=get_device())
+            for _, t in enumerate(tqdm(timesteps)):
+                timestep = [t]
+                timestep = torch.stack(timestep)
+
+                if inference_mode == "TIA":
+                    noise_pred = self.forward_tia(latents, timestep, t, step_change, 
+                                                  arg_tia, arg_ti, arg_i, arg_null)
+                elif inference_mode == "TA":
+                    noise_pred = self.forward_ta(latents, timestep, arg_ta, arg_t, arg_null)
+                elif inference_mode == "TI":
+                    noise_pred = self.forward_ti(latents, timestep, t, step_change,
+                                                arg_ti, arg_t, arg_i, arg_null)
+                else:
+                    raise ValueError(f"Unsupported generation mode: {self.config.generation.mode}")
+
+                temp_x0 = sample_scheduler.step(
+                    noise_pred.unsqueeze(0),
+                    t,
+                    latents[0].unsqueeze(0),
+                    return_dict=False,
+                    generator=seed_g)[0]
+                latents = [temp_x0.squeeze(0)]
+
+                del timestep
+                torch.cuda.empty_cache()
+
+            x0 = latents
+            x0 = [x0_[:,:-latents_ref[0].shape[1]] for x0_ in x0]
+
+            # if offload_model:
+            # self.dit.cpu()
+
+            print("dit finished")
+
+            torch.cuda.empty_cache()
+            # if get_local_rank() == 0:
+            # self.vae.model.to(device=device)
+            videos = self.vae.decode(x0)
+            # self.vae.model.to(device="cpu")
+
+            print("vae 2 finished")
+
+        del noise, latents, noise_pred
+        del audio_emb, audio_emb_neg, latents_ref, latents_ref_neg, context, context_null
+        del x0, temp_x0
+        del sample_scheduler
+        torch.cuda.empty_cache()
+        gc.collect()
+        torch.cuda.synchronize()
+        if dist.is_initialized():
+            dist.barrier()
+
+        return videos[0] # if get_local_rank() == 0 else None
+
+
+    def inference_loop(self, prompt, ref_img_path, audio_path, output_dir, filename, inference_mode = "TIA", width = 832, height = 480, steps=50, frames = 97, tea_cache_l1_thresh = 0.0, seed = 0):
+
+        video = self.inference(
+            prompt,
+            ref_img_path,
+            audio_path,
+            size=SIZE_CONFIGS[f"{width}*{height}"],
+            frame_num=frames,
+            shift=self.config.diffusion.timesteps.sampling.shift,
+            sample_solver='unipc',
+            sampling_steps=steps,
+            inference_mode = inference_mode,
+            tea_cache_l1_thresh = tea_cache_l1_thresh,
+            seed=seed
+        )
+
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        # Save samples.
+        if get_sequence_parallel_rank() == 0:
+            pathname = self.save_sample(
+                sample=video,
+                audio_path=audio_path,
+                output_dir = output_dir,
+                filename=filename,
+            )
+            self.logger.info(f"Finished {filename}, saved to {pathname}.")
+        
+        del video, prompt
+        torch.cuda.empty_cache()
+        gc.collect()
+            
+
+    def save_sample(self, *, sample: torch.Tensor, audio_path: str, output_dir: str, filename: str):
+        gen_config = self.config.generation
+        # Prepare file path.
+        extension = ".mp4" if sample.ndim == 4 else ".png"
+        filename += extension
+        pathname = os.path.join(output_dir, filename)
+        # Convert sample.
+        sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).to("cpu", torch.uint8)
+        sample = rearrange(sample, "c t h w -> t h w c")
+        # Save file.
+        if sample.ndim == 4:
+            if audio_path is not None:
+                tensor_to_video(
+                    sample.numpy(),
+                    pathname,
+                    audio_path,
+                    fps=gen_config.fps)
+            else:
+                mediapy.write_video(
+                path=pathname,
+                images=sample.numpy(),
+                fps=gen_config.fps,
+            )
+        else:
+            raise ValueError
+        return pathname
+    
+
+    def prepare_positive_prompts(self):
+        pos_prompts = self.config.generation.positive_prompt
+        if pos_prompts.endswith(".json"):
+            pos_prompts = prepare_json_dataset(pos_prompts)
+        else:
+            raise NotImplementedError
+        assert isinstance(pos_prompts, ListConfig)
+
+        return pos_prompts
+    
+class TeaCache:
+    def __init__(self, num_inference_steps, rel_l1_thresh, model_id):
+        self.num_inference_steps = num_inference_steps
+        self.step = 0
+        self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = None
+        self.rel_l1_thresh = rel_l1_thresh
+        self.previous_residual = None
+        self.previous_hidden_states = None
+        
+        self.coefficients_dict = {
+            "Wan2.1-T2V-1.3B": [-5.21862437e+04, 9.23041404e+03, -5.28275948e+02, 1.36987616e+01, -4.99875664e-02],
+            "Wan2.1-T2V-14B": [-3.03318725e+05, 4.90537029e+04, -2.65530556e+03, 5.87365115e+01, -3.15583525e-01],
+            "Wan2.1-I2V-14B-480P": [2.57151496e+05, -3.54229917e+04,  1.40286849e+03, -1.35890334e+01, 1.32517977e-01],
+            "Wan2.1-I2V-14B-720P": [ 8.10705460e+03,  2.13393892e+03, -3.72934672e+02,  1.66203073e+01, -4.17769401e-02],
+        }
+        if model_id not in self.coefficients_dict:
+            supported_model_ids = ", ".join([i for i in self.coefficients_dict])
+            raise ValueError(f"{model_id} is not a supported TeaCache model id. Please choose a valid model id in ({supported_model_ids}).")
+        self.coefficients = self.coefficients_dict[model_id]
+
+    def check(self, dit, x, t_mod):
+        modulated_inp = t_mod.clone()
+        if self.step == 0 or self.step == self.num_inference_steps - 1:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0
+        else:
+            coefficients = self.coefficients
+            rescale_func = np.poly1d(coefficients)
+            self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+            if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = modulated_inp
+        self.step += 1
+        if self.step == self.num_inference_steps:
+            self.step = 0
+        if should_calc:
+            self.previous_hidden_states = x.clone()
+        return not should_calc
+
+    def store(self, hidden_states):
+        if self.previous_hidden_states is None:
+            return
+        self.previous_residual = hidden_states - self.previous_hidden_states
+        self.previous_hidden_states = None
+
+    def update(self, hidden_states):
+        hidden_states = hidden_states + self.previous_residual
+        return hidden_states

humo/generate_1_7B.py

@@ -0,0 +1,622 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+# Inference codes adapted from [SeedVR]
+# https://github.com/ByteDance-Seed/SeedVR/blob/main/projects/inference_seedvr2_7b.py
+
+import math
+import os
+import gc
+import random
+import sys
+import mediapy
+import torch
+import torch.distributed as dist
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from einops import rearrange
+from omegaconf import OmegaConf
+from PIL import Image, ImageOps
+from torchvision.transforms import ToTensor
+from tqdm import tqdm
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import (
+    BackwardPrefetch,
+    FullyShardedDataParallel,
+    MixedPrecision,
+    ShardingStrategy,
+)
+from common.distributed import (
+    get_device,
+    get_global_rank,
+    get_local_rank,
+    meta_param_init_fn,
+    meta_non_persistent_buffer_init_fn,
+    init_torch,
+)
+from common.distributed.advanced import (
+    init_unified_parallel,
+    get_unified_parallel_world_size,
+    get_sequence_parallel_rank,
+    init_model_shard_cpu_group,
+)
+from common.logger import get_logger
+from common.config import create_object
+from common.distributed import get_device, get_global_rank
+from torchvision.transforms import Compose, Normalize, ToTensor
+from humo.models.wan_modules.t5 import T5EncoderModel
+from humo.models.wan_modules.vae import WanVAE
+from humo.models.utils.utils import tensor_to_video, prepare_json_dataset
+from contextlib import contextmanager
+import torch.cuda.amp as amp
+from humo.models.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from humo.utils.audio_processor_whisper import AudioProcessor
+from humo.utils.wav2vec import linear_interpolation_fps
+
+
+image_transform = Compose([
+    ToTensor(),
+    Normalize(mean=0.5, std=0.5),
+])
+
+SIZE_CONFIGS = {
+    '720*1280': (720, 1280),
+    '1280*720': (1280, 720),
+    '480*832': (480, 832),
+    '832*480': (832, 480),
+    '1024*1024': (1024, 1024),
+}
+
+def clever_format(nums, format="%.2f"):
+    from typing import Iterable
+    if not isinstance(nums, Iterable):
+        nums = [nums]
+    clever_nums = []
+    for num in nums:
+        if num > 1e12:
+            clever_nums.append(format % (num / 1e12) + "T")
+        elif num > 1e9:
+            clever_nums.append(format % (num / 1e9) + "G")
+        elif num > 1e6:
+            clever_nums.append(format % (num / 1e6) + "M")
+        elif num > 1e3:
+            clever_nums.append(format % (num / 1e3) + "K")
+        else:
+            clever_nums.append(format % num + "B")
+
+    clever_nums = clever_nums[0] if len(clever_nums) == 1 else (*clever_nums,)
+
+    return clever_nums
+
+
+class Generator():
+    def __init__(self, config: DictConfig):
+        self.config = config.copy()
+        OmegaConf.set_readonly(self.config, True)
+        self.logger = get_logger(self.__class__.__name__)
+        self.configure_models()
+        
+        # init_torch(cudnn_benchmark=False)
+    
+    def get_fsdp_sharding_config(self, sharding_strategy, device_mesh_config):
+        device_mesh = None
+        fsdp_strategy = ShardingStrategy[sharding_strategy]
+        if (
+            fsdp_strategy in [ShardingStrategy._HYBRID_SHARD_ZERO2, ShardingStrategy.HYBRID_SHARD]
+            and device_mesh_config is not None
+        ):
+            device_mesh = init_device_mesh("cuda", tuple(device_mesh_config))
+        return device_mesh, fsdp_strategy
+
+    def configure_models(self):
+        self.configure_dit_model(device="cpu")
+        self.configure_vae_model()
+        if self.config.generation.get('extract_audio_feat', False):
+            self.configure_wav2vec(device="cpu")
+        self.configure_text_model(device="cpu")
+
+        # Initialize fsdp.
+        self.configure_dit_fsdp_model()
+        self.configure_text_fsdp_model()
+    
+    def configure_dit_model(self, device=get_device()):
+
+        init_unified_parallel(self.config.dit.sp_size)
+        self.sp_size = get_unified_parallel_world_size()
+        
+        # Create dit model.
+        init_device = "meta"
+        with torch.device(init_device):
+            self.dit = create_object(self.config.dit.model)
+        self.logger.info(f"Load DiT model on {init_device}.")
+        self.dit.eval().requires_grad_(False)
+
+        # Load dit checkpoint.
+        path = self.config.dit.checkpoint_dir
+        if path.endswith(".pth"):
+            state = torch.load(path, map_location=device, mmap=True)
+            missing_keys, unexpected_keys = self.dit.load_state_dict(state, strict=False, assign=True)
+            self.logger.info(
+                f"dit loaded from {path}. "
+                f"Missing keys: {len(missing_keys)}, "
+                f"Unexpected keys: {len(unexpected_keys)}"
+            )
+        else:
+            from safetensors.torch import load_file
+            import json
+            def load_custom_sharded_weights(model_dir, base_name, device=device):
+                index_path = f"{model_dir}/{base_name}.safetensors.index.json"
+                with open(index_path, "r") as f:
+                    index = json.load(f)
+                weight_map = index["weight_map"]
+                shard_files = set(weight_map.values())
+                state_dict = {}
+                for shard_file in shard_files:
+                    shard_path = f"{model_dir}/{shard_file}"
+                    shard_state = load_file(shard_path)
+                    shard_state = {k: v.to(device) for k, v in shard_state.items()}
+                    state_dict.update(shard_state)
+                return state_dict
+            state = load_custom_sharded_weights(path, 'humo', device)
+            self.dit.load_state_dict(state, strict=False, assign=True)
+        
+        self.dit = meta_non_persistent_buffer_init_fn(self.dit)
+        if device in [get_device(), "cuda"]:
+            self.dit.to(get_device())
+
+        # Print model size.
+        params = sum(p.numel() for p in self.dit.parameters())
+        self.logger.info(
+            f"[RANK:{get_global_rank()}] DiT Parameters: {clever_format(params, '%.3f')}"
+        )
+    
+    def configure_vae_model(self, device=get_device()):
+        self.vae_stride = self.config.vae.vae_stride
+        self.vae = WanVAE(
+            vae_pth=self.config.vae.checkpoint,
+            device=device)
+        
+        if self.config.generation.height == 480:
+            self.zero_vae = torch.load(self.config.dit.zero_vae_path)
+        elif self.config.generation.height == 720:
+            self.zero_vae = torch.load(self.config.dit.zero_vae_720p_path)
+        else:
+            raise ValueError(f"Unsupported height {self.config.generation.height} for zero-vae.")
+    
+    def configure_wav2vec(self, device=get_device()):
+        audio_separator_model_file = self.config.audio.vocal_separator
+        wav2vec_model_path = self.config.audio.wav2vec_model
+
+        self.audio_processor = AudioProcessor(
+            16000,
+            25,
+            wav2vec_model_path,
+            "all",
+            audio_separator_model_file,
+            None,  # not seperate
+            os.path.join(self.config.generation.output.dir, "vocals"),
+            device=device,
+        )
+
+    def configure_text_model(self, device=get_device()):
+        self.text_encoder = T5EncoderModel(
+            text_len=self.config.dit.model.text_len,
+            dtype=torch.bfloat16,
+            device=device,
+            checkpoint_path=self.config.text.t5_checkpoint,
+            tokenizer_path=self.config.text.t5_tokenizer,
+            )
+
+    
+    def configure_dit_fsdp_model(self):
+        self.dit.to(get_device())
+        
+        return   
+
+
+    def configure_text_fsdp_model(self):
+        self.text_encoder.to(get_device())
+        
+        return
+
+
+    def load_image_latent_ref_id(self, path: str, size, device):
+        # Load size.
+        h, w = size[1], size[0]
+
+        # Load image.
+        if len(path) > 1 and not isinstance(path, str):
+            ref_vae_latents = []
+            for image_path in path:
+                with Image.open(image_path) as img:
+                    img = img.convert("RGB")
+
+                    # Calculate the required size to keep aspect ratio and fill the rest with padding.
+                    img_ratio = img.width / img.height
+                    target_ratio = w / h
+                    
+                    if img_ratio > target_ratio:  # Image is wider than target
+                        new_width = w
+                        new_height = int(new_width / img_ratio)
+                    else:  # Image is taller than target
+                        new_height = h
+                        new_width = int(new_height * img_ratio)
+                    
+                    # img = img.resize((new_width, new_height), Image.ANTIALIAS)
+                    img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+                    # Create a new image with the target size and place the resized image in the center
+                    delta_w = w - img.size[0]
+                    delta_h = h - img.size[1]
+                    padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+                    new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))
+
+                    # Transform to tensor and normalize.
+                    transform = Compose(
+                        [
+                            ToTensor(),
+                            Normalize(0.5, 0.5),
+                        ]
+                    )
+                    new_img = transform(new_img)
+                    # img_vae_latent = self.vae_encode([new_img.unsqueeze(1)])[0]
+                    img_vae_latent = self.vae.encode([new_img.unsqueeze(1)], device)
+                    ref_vae_latents.append(img_vae_latent[0])
+
+            return [torch.cat(ref_vae_latents, dim=1)]
+        else:
+            if not isinstance(path, str):
+                path = path[0]
+            with Image.open(path) as img:
+                img = img.convert("RGB")
+
+                # Calculate the required size to keep aspect ratio and fill the rest with padding.
+                img_ratio = img.width / img.height
+                target_ratio = w / h
+                
+                if img_ratio > target_ratio:  # Image is wider than target
+                    new_width = w
+                    new_height = int(new_width / img_ratio)
+                else:  # Image is taller than target
+                    new_height = h
+                    new_width = int(new_height * img_ratio)
+                
+                # img = img.resize((new_width, new_height), Image.ANTIALIAS)
+                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+                # Create a new image with the target size and place the resized image in the center
+                delta_w = w - img.size[0]
+                delta_h = h - img.size[1]
+                padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+                new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))
+
+                # Transform to tensor and normalize.
+                transform = Compose(
+                    [
+                        ToTensor(),
+                        Normalize(0.5, 0.5),
+                    ]
+                )
+                new_img = transform(new_img)
+                img_vae_latent = self.vae.encode([new_img.unsqueeze(1)], device)
+
+            # Vae encode.
+            return img_vae_latent
+    
+    def get_audio_emb_window(self, audio_emb, frame_num, frame0_idx, audio_shift=2):
+        zero_audio_embed = torch.zeros((audio_emb.shape[1], audio_emb.shape[2]), dtype=audio_emb.dtype, device=audio_emb.device)
+        zero_audio_embed_3 = torch.zeros((3, audio_emb.shape[1], audio_emb.shape[2]), dtype=audio_emb.dtype, device=audio_emb.device)  # device=audio_emb.device
+        iter_ = 1 + (frame_num - 1) // 4
+        audio_emb_wind = []
+        for lt_i in range(iter_):
+            if lt_i == 0:
+                st = frame0_idx + lt_i - 2
+                ed = frame0_idx + lt_i + 3
+                wind_feat = torch.stack([
+                    audio_emb[i] if (0 <= i < audio_emb.shape[0]) else zero_audio_embed
+                    for i in range(st, ed)
+                ], dim=0)
+                wind_feat = torch.cat((zero_audio_embed_3, wind_feat), dim=0)
+            else:
+                st = frame0_idx + 1 + 4 * (lt_i - 1) - audio_shift
+                ed = frame0_idx + 1 + 4 * lt_i + audio_shift
+                wind_feat = torch.stack([
+                    audio_emb[i] if (0 <= i < audio_emb.shape[0]) else zero_audio_embed
+                    for i in range(st, ed)
+                ], dim=0)
+            audio_emb_wind.append(wind_feat)
+        audio_emb_wind = torch.stack(audio_emb_wind, dim=0)
+
+        return audio_emb_wind, ed - audio_shift
+    
+    def audio_emb_enc(self, audio_emb, wav_enc_type="whisper"):
+        if wav_enc_type == "wav2vec":
+            feat_merge = audio_emb
+        elif wav_enc_type == "whisper":
+            feat0 = linear_interpolation_fps(audio_emb[:, :, 0: 8].mean(dim=2), 50, 25)
+            feat1 = linear_interpolation_fps(audio_emb[:, :, 8: 16].mean(dim=2), 50, 25)
+            feat2 = linear_interpolation_fps(audio_emb[:, :, 16: 24].mean(dim=2), 50, 25)
+            feat3 = linear_interpolation_fps(audio_emb[:, :, 24: 32].mean(dim=2), 50, 25)
+            feat4 = linear_interpolation_fps(audio_emb[:, :, 32], 50, 25)
+            feat_merge = torch.stack([feat0, feat1, feat2, feat3, feat4], dim=2)[0]
+        else:
+            raise ValueError(f"Unsupported wav_enc_type: {wav_enc_type}")
+        
+        return feat_merge
+    
+    def forward_tia(self, latents, latents_ref, latents_ref_neg, timestep, arg_t, arg_ta, arg_null):
+        neg = self.dit(
+            [torch.cat([latent[:,:-latent_ref_neg.shape[1]], latent_ref_neg], dim=1) for latent, latent_ref_neg in zip(latents, latents_ref_neg)], t=timestep, **arg_null
+            )[0]
+        
+        pos_t = self.dit(
+            [torch.cat([latent[:,:-latent_ref_neg.shape[1]], latent_ref_neg], dim=1) for latent, latent_ref_neg in zip(latents, latents_ref_neg)], t=timestep, **arg_t
+            )[0]
+        pos_ta = self.dit(
+            [torch.cat([latent[:,:-latent_ref_neg.shape[1]], latent_ref_neg], dim=1) for latent, latent_ref_neg in zip(latents, latents_ref_neg)], t=timestep, **arg_ta
+            )[0]
+        pos_tia = self.dit(
+            [torch.cat([latent[:,:-latent_ref.shape[1]], latent_ref], dim=1) for latent, latent_ref in zip(latents, latents_ref)], t=timestep, **arg_ta
+            )[0]
+        
+        noise_pred = self.config.generation.scale_i * (pos_tia - pos_ta) + \
+                    self.config.generation.scale_a * (pos_ta - pos_t) + \
+                    self.config.generation.scale_t * (pos_t - neg) + \
+                    neg
+        
+        return noise_pred
+    
+    def forward_ta(self, latents, latents_ref_neg, timestep, arg_t, arg_ta, arg_null):
+        neg = self.dit(
+            [torch.cat([latent[:,:-latent_ref_neg.shape[1]], latent_ref_neg], dim=1) for latent, latent_ref_neg in zip(latents, latents_ref_neg)], t=timestep, **arg_null
+            )[0]
+        
+        pos_t = self.dit(
+            [torch.cat([latent[:,:-latent_ref_neg.shape[1]], latent_ref_neg], dim=1) for latent, latent_ref_neg in zip(latents, latents_ref_neg)], t=timestep, **arg_t
+            )[0]
+        pos_ta = self.dit(
+            [torch.cat([latent[:,:-latent_ref_neg.shape[1]], latent_ref_neg], dim=1) for latent, latent_ref_neg in zip(latents, latents_ref_neg)], t=timestep, **arg_ta
+            )[0]
+        
+        noise_pred = self.config.generation.scale_a * (pos_ta - pos_t) + \
+                    self.config.generation.scale_t * (pos_t - neg) + \
+                    neg
+        
+        return noise_pred
+            
+                    
+    @torch.no_grad()
+    def inference(self,
+                 input_prompt,
+                 img_path,
+                 audio_path,
+                 size=(1280, 720),
+                 frame_num=81,
+                 shift=5.0,
+                 sample_solver='unipc',
+                 sampling_steps=50,
+                 n_prompt="",
+                 seed=-1,
+                 offload_model=True,
+                 device = get_device(),
+        ):
+
+        self.vae.model.to(device=device)
+        if img_path is not None:
+            latents_ref = self.load_image_latent_ref_id(img_path, size, device)
+        else:
+            latents_ref = [torch.zeros(16, 1, size[1]//8, size[0]//8).to(device)]
+            
+        self.vae.model.to(device="cpu")
+        latents_ref_neg = [torch.zeros_like(latent_ref) for latent_ref in latents_ref]
+        
+        # audio
+        if audio_path is not None:
+            if self.config.generation.extract_audio_feat:
+                self.audio_processor.whisper.to(device=device)
+                audio_emb, audio_length = self.audio_processor.preprocess(audio_path)
+                self.audio_processor.whisper.to(device='cpu')
+            else:
+                audio_emb_path = audio_path.replace(".wav", ".pt")
+                audio_emb = torch.load(audio_emb_path).to(device=device)
+                audio_emb = self.audio_emb_enc(audio_emb, wav_enc_type="whisper")
+                self.logger.info("使用预先提取好的音频特征: %s", audio_emb_path)
+        else:
+            audio_emb = torch.zeros(frame_num, 5, 1280).to(device)
+            
+        frame_num = frame_num if frame_num != -1 else audio_length
+        frame_num = 4 * ((frame_num - 1) // 4) + 1
+        audio_emb, _ = self.get_audio_emb_window(audio_emb, frame_num, frame0_idx=0)
+        zero_audio_pad = torch.zeros(latents_ref[0].shape[1], *audio_emb.shape[1:]).to(audio_emb.device)
+        audio_emb = torch.cat([audio_emb, zero_audio_pad], dim=0)
+        audio_emb = [audio_emb.to(device)]
+        audio_emb_neg = [torch.zeros_like(audio_emb[0])]
+        
+        # preprocess
+        self.patch_size = self.config.dit.model.patch_size
+        F = frame_num
+        target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1 + latents_ref[0].shape[1],
+                        size[1] // self.vae_stride[1],
+                        size[0] // self.vae_stride[2])
+
+        seq_len = math.ceil((target_shape[2] * target_shape[3]) /
+                            (self.patch_size[1] * self.patch_size[2]) *
+                            target_shape[1] / self.sp_size) * self.sp_size
+
+        if n_prompt == "":
+            n_prompt = self.config.generation.sample_neg_prompt
+        seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
+        seed_g = torch.Generator(device=device)
+        seed_g.manual_seed(seed)
+
+        self.text_encoder.model.to(device)
+        context = self.text_encoder([input_prompt], device)
+        context_null = self.text_encoder([n_prompt], device)
+        self.text_encoder.model.cpu()
+
+        noise = [
+            torch.randn(
+                target_shape[0],
+                target_shape[1], # - latents_ref[0].shape[1],
+                target_shape[2],
+                target_shape[3],
+                dtype=torch.float32,
+                device=device,
+                generator=seed_g)
+        ]
+
+        @contextmanager
+        def noop_no_sync():
+            yield
+
+        no_sync = getattr(self.dit, 'no_sync', noop_no_sync)
+        # step_change = self.config.generation.step_change # 980
+
+        # evaluation mode
+        with amp.autocast(dtype=torch.bfloat16), torch.no_grad(), no_sync():
+
+            if sample_solver == 'unipc':
+                sample_scheduler = FlowUniPCMultistepScheduler(
+                    num_train_timesteps=1000,
+                    shift=1,
+                    use_dynamic_shifting=False)
+                sample_scheduler.set_timesteps(
+                    sampling_steps, device=device, shift=shift)
+                timesteps = sample_scheduler.timesteps
+
+            # sample videos
+            latents = noise
+
+            # referene image在下面的输入中手动指定, 不在arg中指定
+            arg_ta = {'context': context, 'seq_len': seq_len, 'audio': audio_emb}
+            arg_t = {'context': context, 'seq_len': seq_len, 'audio': audio_emb_neg}
+            arg_null = {'context': context_null, 'seq_len': seq_len, 'audio': audio_emb_neg}
+            
+            torch.cuda.empty_cache()
+            self.dit.to(device=get_device())
+            for _, t in enumerate(tqdm(timesteps)):
+                timestep = [t]
+                timestep = torch.stack(timestep)
+
+                if self.config.generation.mode == "TIA":
+                    noise_pred = self.forward_tia(latents, latents_ref, latents_ref_neg, timestep, arg_t, arg_ta, arg_null)
+                elif self.config.generation.mode == "TA":
+                    noise_pred = self.forward_ta(latents, latents_ref_neg, timestep, arg_t, arg_ta, arg_null)
+                else:
+                    raise ValueError(f"Unsupported generation mode: {self.config.generation.mode}")
+
+                temp_x0 = sample_scheduler.step(
+                    noise_pred.unsqueeze(0),
+                    t,
+                    latents[0].unsqueeze(0),
+                    return_dict=False,
+                    generator=seed_g)[0]
+                latents = [temp_x0.squeeze(0)]
+
+                del timestep
+                torch.cuda.empty_cache()
+
+            x0 = latents
+            x0 = [x0_[:,:-latents_ref[0].shape[1]] for x0_ in x0]
+
+            # if offload_model:
+            self.dit.cpu()
+            torch.cuda.empty_cache()
+            # if get_local_rank() == 0:
+            self.vae.model.to(device=device)
+            videos = self.vae.decode(x0)
+            self.vae.model.to(device="cpu")
+
+        del noise, latents, noise_pred
+        del audio_emb, audio_emb_neg, latents_ref, latents_ref_neg, context, context_null
+        del x0, temp_x0
+        del sample_scheduler
+        torch.cuda.empty_cache()
+        gc.collect()
+        torch.cuda.synchronize()
+        if dist.is_initialized():
+            dist.barrier()
+
+        return videos[0] # if get_local_rank() == 0 else None
+
+
+    def inference_loop(self, prompt, ref_img_path, audio_path, output_dir, filename, width = 832, height = 480, steps=50, frames = 97, seed = 0):
+        print(f'ref_img_path:{ref_img_path}')
+
+        video = self.inference(
+            prompt,
+            ref_img_path,
+            audio_path,
+            size=SIZE_CONFIGS[f"{width}*{height}"],
+            frame_num=frames,
+            shift=self.config.diffusion.timesteps.sampling.shift,
+            sample_solver='unipc',
+            sampling_steps=steps,
+            seed=seed,
+            offload_model=False,
+        )
+
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+
+        # Save samples.
+        if get_sequence_parallel_rank() == 0:
+            pathname = self.save_sample(
+                sample=video,
+                audio_path=audio_path,
+                output_dir = output_dir,
+                filename=filename,
+            )
+            self.logger.info(f"Finished {filename}, saved to {pathname}.")
+        
+        del video, prompt
+        torch.cuda.empty_cache()
+        gc.collect()
+
+            
+
+    def save_sample(self, *, sample: torch.Tensor, audio_path: str, output_dir: str, filename: str):
+        gen_config = self.config.generation
+        # Prepare file path.
+        extension = ".mp4" if sample.ndim == 4 else ".png"
+        filename += extension
+        pathname = os.path.join(output_dir, filename)
+        # Convert sample.
+        sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).to("cpu", torch.uint8)
+        sample = rearrange(sample, "c t h w -> t h w c")
+        # Save file.
+        if sample.ndim == 4:
+            if audio_path is not None:
+                tensor_to_video(
+                    sample.numpy(),
+                    pathname,
+                    audio_path,
+                    fps=gen_config.fps)
+            else:
+                mediapy.write_video(
+                path=pathname,
+                images=sample.numpy(),
+                fps=gen_config.fps,
+            )
+        else:
+            raise ValueError
+        return pathname
+    
+
+    def prepare_positive_prompts(self):
+        pos_prompts = self.config.generation.positive_prompt
+        if pos_prompts.endswith(".json"):
+            pos_prompts = prepare_json_dataset(pos_prompts)
+        else:
+            raise NotImplementedError
+        assert isinstance(pos_prompts, ListConfig)
+
+        return pos_prompts

humo/models/audio/audio_proj.py

@@ -0,0 +1,87 @@
+import torch
+from einops import rearrange
+from torch import nn
+from einops import rearrange
+
+class WanRMSNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class DummyAdapterLayer(nn.Module):
+    def __init__(self, layer):
+        super().__init__()
+        self.layer = layer
+
+    def forward(self, *args, **kwargs):
+        return self.layer(*args, **kwargs)
+    
+
+class AudioProjModel(nn.Module):
+    def __init__(
+        self,
+        seq_len=5,
+        blocks=13,  # add a new parameter blocks
+        channels=768,  # add a new parameter channels
+        intermediate_dim=512,
+        output_dim=1536,
+        context_tokens=16,
+    ):
+        super().__init__()
+
+        self.seq_len = seq_len
+        self.blocks = blocks
+        self.channels = channels
+        self.input_dim = seq_len * blocks * channels  # update input_dim to be the product of blocks and channels.
+        self.intermediate_dim = intermediate_dim
+        self.context_tokens = context_tokens
+        self.output_dim = output_dim
+
+        # define multiple linear layers
+        self.audio_proj_glob_1 = DummyAdapterLayer(nn.Linear(self.input_dim, intermediate_dim))
+        self.audio_proj_glob_2 = DummyAdapterLayer(nn.Linear(intermediate_dim, intermediate_dim))
+        self.audio_proj_glob_3 = DummyAdapterLayer(nn.Linear(intermediate_dim, context_tokens * output_dim))
+
+        self.audio_proj_glob_norm = DummyAdapterLayer(nn.LayerNorm(output_dim))
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+    def forward(self, audio_embeds):
+        video_length = audio_embeds.shape[1]
+        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)
+
+        audio_embeds = torch.relu(self.audio_proj_glob_1(audio_embeds))
+        audio_embeds = torch.relu(self.audio_proj_glob_2(audio_embeds))
+
+        context_tokens = self.audio_proj_glob_3(audio_embeds).reshape(batch_size, self.context_tokens, self.output_dim)
+
+        context_tokens = self.audio_proj_glob_norm(context_tokens)
+        context_tokens = rearrange(context_tokens, "(bz f) m c -> bz f m c", f=video_length)
+
+        return context_tokens

humo/models/distributed/dit_ulysses_sequence_parallel.py

@@ -0,0 +1,270 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+import torch
+from einops import rearrange
+from common.distributed import get_device
+
+from common.distributed.advanced import (
+    get_unified_parallel_world_size,
+    get_unified_parallel_group,
+    pad_tensor,
+    Slice,
+    gather_outputs,
+    gather_seq_scatter_heads_qkv,
+    gather_seq_scatter_double_head,
+    gather_heads_scatter_seq,
+    unpad_tensor
+)
+from humo.models.wan_modules.attention import flash_attention
+from humo.models.wan_modules.model_humo import rope_apply, sinusoidal_embedding_1d
+
+
+def ulysses_dit_forward(
+    self,
+    x,
+    t,
+    context,
+    seq_len,
+    audio=None,
+    y=None
+):
+    """
+    x:              A list of videos each with shape [C, T, H, W].
+    t:              [B].
+    context:        A list of text embeddings each with shape [L, C].
+    """
+    if self.model_type == 'i2v':
+        # assert clip_fea is not None and y is not None
+        assert y is not None
+    # params
+    device = self.patch_embedding.weight.device
+    if self.freqs.device != device:
+        self.freqs = self.freqs.to(device)
+
+    if y is not None:
+        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+    # embeddings
+    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+    grid_sizes = torch.stack(
+        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+    x = [u.flatten(2).transpose(1, 2) for u in x]
+    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long, device=device)
+
+    assert seq_lens.max() <= seq_len
+    x = torch.cat([
+        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
+        for u in x
+    ])
+
+    # time embeddings
+    with torch.amp.autocast('cuda', dtype=torch.float32):
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t).float()).float()
+        e0 = self.time_projection(e).unflatten(1, (6, self.dim)).float()
+        assert e.dtype == torch.float32 and e0.dtype == torch.float32
+        
+    # context
+    context_lens = None
+    context = self.text_embedding(
+        torch.stack([
+            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+            for u in context
+        ]))
+
+    if self.insert_audio:
+        audio = [self.audio_proj(au.unsqueeze(0)).permute(0, 3, 1, 2) for au in audio]
+
+        audio_seq_len = torch.tensor(max([au.shape[2] for au in audio]) * audio[0].shape[3], device=get_device())
+        audio = [au.flatten(2).transpose(1, 2) for au in audio] # [1, t*32, 1536]
+        audio_seq_lens = torch.tensor([au.size(1) for au in audio], dtype=torch.long, device=device)
+        audio = torch.cat([
+            torch.cat([au, au.new_zeros(1, audio_seq_len - au.size(1), au.size(2))],
+                        dim=1) for au in audio
+        ])
+    else:
+        audio = None
+        audio_seq_len = None
+        audio_seq_lens = None
+
+    # ulysses support
+    sp_world = get_unified_parallel_world_size()
+    group = get_unified_parallel_group()
+    if seq_len % sp_world:
+        padding_size = sp_world - (seq_len % sp_world)
+        x = pad_tensor(x, dim=1, padding_size=padding_size)
+
+        if self.insert_audio:
+            audio_padding_size = sp_world - (audio_seq_len % sp_world)
+            audio = pad_tensor(audio, dim=1, padding_size=audio_padding_size)
+        
+    x = Slice.apply(group, x, 1, True)
+
+    if self.insert_audio:
+        audio = Slice.apply(group, audio, 1, True)
+
+    # arguments
+    kwargs = dict(
+        e=e0,
+        seq_lens=seq_lens,
+        grid_sizes=grid_sizes,
+        freqs=self.freqs,
+        context=context,
+        context_lens=context_lens,
+        audio=audio,
+        audio_seq_len=audio_seq_len)
+    
+    for block in self.blocks:
+        x = block(x, **kwargs)
+
+    # head
+    x = self.head(x, e)
+
+    # ulysses support
+    x = gather_outputs(x, gather_dim=1, padding_dim=1, unpad_dim_size=seq_len, scale_grad=True)
+    
+    # unpatchify
+    x = self.unpatchify(x, grid_sizes)
+    return [u.float() for u in x]
+
+
+def ulysses_attn_forward(
+    self,
+    x,
+    seq_lens,
+    grid_sizes,
+    freqs,
+    dtype=torch.bfloat16
+):
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    seq_len = seq_lens.max()
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(x))
+        v = self.v(x)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+
+    # ulysses support
+    sp_size = get_unified_parallel_world_size()
+    if n % sp_size:
+        pad_size = sp_size - (n % sp_size)
+        pad_size = pad_size * d
+        pad_inner_dim = n * d + pad_size
+        q = pad_tensor(q, dim=2, padding_size=pad_size)
+        k = pad_tensor(k, dim=2, padding_size=pad_size)
+        v = pad_tensor(v, dim=2, padding_size=pad_size)
+    else:
+        pad_inner_dim = n * d
+
+    qkv = torch.cat([q, k, v], dim=2)
+    qkv = gather_seq_scatter_heads_qkv(qkv, seq_dim=1, unpadded_dim_size=seq_len)
+    q, k, v = qkv.split(pad_inner_dim // sp_size, dim=2)
+
+    pad_n = pad_inner_dim // d
+    pad_split_n = pad_n // sp_size
+    q = q.view(b, seq_len, pad_split_n, d)
+    k = k.view(b, seq_len, pad_split_n, d)
+    v = v.view(b, seq_len, pad_split_n, d)
+
+    q = rope_apply(q, grid_sizes, freqs)
+    k = rope_apply(k, grid_sizes, freqs)
+
+    x = flash_attention(
+        q=half(q),
+        k=half(k),
+        v=half(v),
+        k_lens=seq_lens,
+        window_size=self.window_size
+    )
+
+    # ulysses support
+    x = x.flatten(2)
+    x = gather_heads_scatter_seq(x, head_dim=2, seq_dim=1)
+    if n % sp_size:
+        x = unpad_tensor(x, dim=2, unpad_dim_size=seq_len)
+
+    x = self.o(x)
+    return x
+
+
+def ulysses_audio_cross_attn_forward(
+    self,
+    x,
+    audio,
+    seq_lens,
+    grid_sizes,
+    freqs,
+    audio_seq_len,
+    dtype=torch.bfloat16
+):
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    seq_len = seq_lens.max()
+
+    q = self.norm_q(self.q(x))
+    k = self.norm_k(self.k(audio))
+    v = self.v(audio)
+
+    # ulysses support
+    sp_size = get_unified_parallel_world_size()
+    if n % sp_size:
+        pad_size = sp_size - (n % sp_size)
+        pad_size = pad_size * d
+        pad_inner_dim = n * d + pad_size
+        q = pad_tensor(q, dim=2, padding_size=pad_size)
+        k = pad_tensor(k, dim=2, padding_size=pad_size)
+        v = pad_tensor(v, dim=2, padding_size=pad_size)
+    else:
+        pad_inner_dim = n * d
+
+    qq = torch.cat([q, q], dim=2)
+    kv = torch.cat([k, v], dim=2)
+    qq = gather_seq_scatter_double_head(qq, seq_dim=1, unpadded_dim_size=seq_len)
+    kv = gather_seq_scatter_double_head(kv, seq_dim=1, unpadded_dim_size=audio_seq_len)
+    q, _ = qq.split(pad_inner_dim // sp_size, dim=2)
+    k, v = kv.split(pad_inner_dim // sp_size, dim=2)
+
+    pad_n = pad_inner_dim // d
+    pad_split_n = pad_n // sp_size
+    q = q.view(b, seq_len, pad_split_n, d)
+    k = k.view(b, audio_seq_len, pad_split_n, d)
+    v = v.view(b, audio_seq_len, pad_split_n, d)
+
+    hlen_wlen = int(grid_sizes[0][1] * grid_sizes[0][2])
+    assert hlen_wlen == 1560 or hlen_wlen == 3600
+    q = q.reshape(-1, hlen_wlen, pad_split_n, d)
+    k = k.reshape(-1, 16, pad_split_n, d)
+    v = v.reshape(-1, 16, pad_split_n, d)
+
+    x = flash_attention(
+        q=q,
+        k=k,
+        v=v,
+        k_lens=None,
+    )
+    x = x.view(b, -1, pad_split_n, d)
+
+    # ulysses support
+    x = x.flatten(2)
+    x = gather_heads_scatter_seq(x, head_dim=2, seq_dim=1)
+    if n % sp_size:
+        x = unpad_tensor(x, dim=2, unpad_dim_size=seq_len)
+
+    x = self.o(x)
+    return x

humo/models/distributed/fsdp.py

@@ -0,0 +1,42 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+from functools import partial
+
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
+from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
+
+
+def shard_model(
+    model,
+    device_id,
+    param_dtype=torch.bfloat16,
+    reduce_dtype=torch.float32,
+    buffer_dtype=torch.float32,
+    process_group=None,
+    sharding_strategy=ShardingStrategy.FULL_SHARD,
+    sync_module_states=True,
+):
+    model = FSDP(
+        module=model,
+        process_group=process_group,
+        sharding_strategy=sharding_strategy,
+        auto_wrap_policy=partial(
+            lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks),
+        mixed_precision=MixedPrecision(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype,
+            buffer_dtype=buffer_dtype),
+        device_id=device_id,
+        sync_module_states=sync_module_states)
+    return model

humo/models/text/encoder.py

@@ -0,0 +1,173 @@
+import os
+from dataclasses import dataclass
+from typing import List, Optional, Union
+import torch
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    CLIPTextModel,
+    CLIPTokenizerFast,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+from transformers.tokenization_utils_base import BatchEncoding
+
+from common.fs import download_and_extract
+from common.logger import get_logger
+
+logger = get_logger(__name__)
+
+MODEL_TYPES = {
+    "clip": (CLIPTokenizerFast, CLIPTextModel),
+    "t5": (T5TokenizerFast, T5EncoderModel),
+    "llm14b": (AutoTokenizer, AutoModelForCausalLM),
+}
+
+
+@dataclass
+class TextEncoderOutput:
+    embeddings: Union[torch.FloatTensor, List[torch.FloatTensor]]
+    masks: Union[torch.BoolTensor, List[torch.BoolTensor]]
+    pooled: Optional[Union[torch.FloatTensor, List[torch.FloatTensor]]]
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, config: DictConfig):
+        super().__init__()
+        self.config = config
+        self.tokenizers = []
+        self.models = nn.ModuleList([])
+
+        # Disable tokenizer parallelism since we already use distributed training.
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+        for model in config.models:
+            tokenizer_cls, model_cls = MODEL_TYPES[model.type]
+            path = download_and_extract(model.path)
+            max_length = model.max_length
+
+            if model.type == "llm14b":
+                tokenizer = tokenizer_cls.from_pretrained(
+                    path,
+                    model_max_length=max_length,
+                    use_fast=False,
+                    trust_remote_code=True,
+                    padding_side="right",
+                    truncation_side="right",
+                    add_eod_token=True,
+                )
+                tokenizer.add_special_tokens({"pad_token": "<|endoftext|>"})
+                model = model_cls.from_pretrained(path, trust_remote_code=True, bf16=True)
+            else:
+                tokenizer = tokenizer_cls.from_pretrained(path, model_max_length=max_length)
+                model = model_cls.from_pretrained(path, torch_dtype=torch.bfloat16)
+            self.tokenizers.append(tokenizer)
+            self.models.append(model)
+
+    def forward(self, text: Union[str, List[str]]) -> TextEncoderOutput:
+        embeddings, masks, pooled = [], [], []
+
+        for encoder_config, tokenizer, model in zip(
+            self.config.models, self.tokenizers, self.models
+        ):
+            if encoder_config.type == "llm14b":
+                use_mask = encoder_config.get("mask", True)
+                tokens = tokenizer(
+                    text,
+                    return_tensors="pt",
+                    padding="max_length",
+                    truncation=True,
+                ).to(model.device)
+                token_ids = tokens["input_ids"]
+                attention_mask = tokens["attention_mask"]
+                num_tokens = attention_mask.sum(dim=1)
+                range_ids = torch.arange(len(token_ids), device=token_ids.device, dtype=torch.long)
+                token_ids[range_ids, num_tokens.clamp(max=token_ids.size(1) - 1)] = (
+                    tokenizer.pad_token_id
+                )
+                attention_mask[range_ids, num_tokens.clamp(max=token_ids.size(1) - 1)] = 1
+                tokens = BatchEncoding({"input_ids": token_ids, "attention_mask": attention_mask})
+                output = model.transformer(
+                    input_ids=tokens.input_ids,
+                    attention_mask=attention_mask if use_mask else None,
+                    output_hidden_states=False,
+                    use_cache=False,
+                )
+                emb = output.last_hidden_state  # batch_size, num_tokens, feat_dim
+                # emb *= tokens.attention_mask.unsqueeze(-1)
+
+                embeddings.append(emb)
+                masks.append(
+                    tokens.attention_mask.bool() if use_mask else tokens.attention_mask > -1
+                )
+
+            else:
+                # Tokenizer
+                tokens = tokenizer(
+                    text=text,
+                    truncation=True,
+                    padding="max_length",
+                    return_tensors="pt",
+                )
+
+                # Encoder
+                use_mask = encoder_config.get("mask", True)
+                input_ids = tokens.input_ids.to(model.device)
+                attention_mask = tokens.attention_mask.to(model.device)
+                output = model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask if use_mask else None,
+                    output_hidden_states=True,
+                )
+
+                # Save embeddings from the defined layer.
+                layer = encoder_config.get("layer", "last")
+                if layer == "last":
+                    embeddings.append(output.last_hidden_state)
+                elif layer == "penultimate":
+                    embeddings.append(model.text_model.final_layer_norm(output.hidden_states[-2]))
+                elif layer == "penultimate_nonorm":
+                    embeddings.append(output.hidden_states[-2])
+                else:
+                    raise NotImplementedError(f"Unknown layer type: {layer}.")
+
+                # Save masks
+                masks.append(attention_mask.bool() if use_mask else attention_mask > -1)
+
+                # Save pooled output if available.
+                if hasattr(output, "pooler_output"):
+                    pooled.append(output.pooler_output)
+
+            output_config = self.config.get("output") or OmegaConf.create()
+            embedding_output_type = output_config.get("embedding_and_mask", "undefined")
+            pooled_output_type = output_config.get("pooled", "undefined")
+
+            # Select or merge embeddings and mask if needed.
+            if embedding_output_type == "undefined" and len(self.models) == 1:
+                embeddings = embeddings[0]
+                masks = masks[0]
+            elif embedding_output_type == "channel_concat":
+                embeddings = torch.cat(embeddings, dim=-1)
+                masks = sum(masks).bool()
+            elif embedding_output_type == "last":
+                embeddings = embeddings[-1]
+                masks = masks[-1]
+            else:
+                raise NotImplementedError(f"output.embedding_and_mask: {embedding_output_type}")
+
+            # Select or merge pooled output if needed.
+            if pooled_output_type == "undefined":
+                pooled = None
+            elif pooled_output_type == "channel_concat":
+                pooled = torch.cat(pooled, dim=-1)
+            elif pooled_output_type == "first":
+                pooled = pooled[0]
+            elif pooled_output_type == "last":
+                pooled = pooled[-1]
+            else:
+                raise NotImplementedError(f"output.pooled: {pooled_output_type}")
+
+        # Return final results.
+        return TextEncoderOutput(embeddings, masks, pooled)

humo/models/utils/fm_solvers.py

@@ -0,0 +1,857 @@
+# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
+# Convert dpm solver for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import inspect
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+                                                   SchedulerMixin,
+                                                   SchedulerOutput)
+from diffusers.utils import deprecate, is_scipy_available
+from diffusers.utils.torch_utils import randn_tensor
+
+if is_scipy_available():
+    pass
+
+
+def get_sampling_sigmas(sampling_steps, shift):
+    sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps]
+    sigma = (shift * sigma / (1 + (shift - 1) * sigma))
+
+    return sigma
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps=None,
+    device=None,
+    timesteps=None,
+    sigmas=None,
+    **kwargs,
+):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError(
+            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+        )
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class FlowDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `FlowDPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs.
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model. This determines the resolution of the diffusion process.
+        solver_order (`int`, defaults to 2):
+            The DPMSolver order which can be `1`, `2`, or `3`. It is recommended to use `solver_order=2` for guided
+            sampling, and `solver_order=3` for unconditional sampling. This affects the number of model outputs stored
+            and used in multistep updates.
+        prediction_type (`str`, defaults to "flow_prediction"):
+            Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+            the flow of the diffusion process.
+        shift (`float`, *optional*, defaults to 1.0):
+            A factor used to adjust the sigmas in the noise schedule. It modifies the step sizes during the sampling
+            process.
+        use_dynamic_shifting (`bool`, defaults to `False`):
+            Whether to apply dynamic shifting to the timesteps based on image resolution. If `True`, the shifting is
+            applied on the fly.
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This method adjusts the predicted sample to prevent
+            saturation and improve photorealism.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and
+            `algorithm_type="dpmsolver++"`.
+        algorithm_type (`str`, defaults to `dpmsolver++`):
+            Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The
+            `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927)
+            paper, and the `dpmsolver++` type implements the algorithms in the
+            [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or
+            `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion.
+        solver_type (`str`, defaults to `midpoint`):
+            Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the
+            sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers.
+        lower_order_final (`bool`, defaults to `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        euler_at_final (`bool`, defaults to `False`):
+            Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail
+            richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference
+            steps, but sometimes may result in blurring.
+        final_sigmas_type (`str`, *optional*, defaults to "zero"):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+        lambda_min_clipped (`float`, defaults to `-inf`):
+            Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the
+            cosine (`squaredcos_cap_v2`) noise schedule.
+        variance_type (`str`, *optional*):
+            Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output
+            contains the predicted Gaussian variance.
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        solver_order: int = 2,
+        prediction_type: str = "flow_prediction",
+        shift: Optional[float] = 1.0,
+        use_dynamic_shifting=False,
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
+        lower_order_final: bool = True,
+        euler_at_final: bool = False,
+        final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+        lambda_min_clipped: float = -float("inf"),
+        variance_type: Optional[str] = None,
+        invert_sigmas: bool = False,
+    ):
+        if algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"
+            deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0",
+                      deprecation_message)
+
+        # settings for DPM-Solver
+        if algorithm_type not in [
+                "dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"
+        ]:
+            if algorithm_type == "deis":
+                self.register_to_config(algorithm_type="dpmsolver++")
+            else:
+                raise NotImplementedError(
+                    f"{algorithm_type} is not implemented for {self.__class__}")
+
+        if solver_type not in ["midpoint", "heun"]:
+            if solver_type in ["logrho", "bh1", "bh2"]:
+                self.register_to_config(solver_type="midpoint")
+            else:
+                raise NotImplementedError(
+                    f"{solver_type} is not implemented for {self.__class__}")
+
+        if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"
+                                 ] and final_sigmas_type == "zero":
+            raise ValueError(
+                f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead."
+            )
+
+        # setable values
+        self.num_inference_steps = None
+        alphas = np.linspace(1, 1 / num_train_timesteps,
+                             num_train_timesteps)[::-1].copy()
+        sigmas = 1.0 - alphas
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 +
+                                       (shift - 1) * sigmas)  # pyright: ignore
+
+        self.sigmas = sigmas
+        self.timesteps = sigmas * num_train_timesteps
+
+        self.model_outputs = [None] * solver_order
+        self.lower_order_nums = 0
+        self._step_index = None
+        self._begin_index = None
+
+        # self.sigmas = self.sigmas.to(
+        #     "cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+    def set_timesteps(
+        self,
+        num_inference_steps: Union[int, None] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[Union[float, None]] = None,
+        shift: Optional[Union[float, None]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                Total number of the spacing of the time steps.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is None:
+            sigmas = np.linspace(self.sigma_max, self.sigma_min,
+                                 num_inference_steps +
+                                 1).copy()[:-1]  # pyright: ignore
+
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
+        else:
+            if shift is None:
+                shift = self.config.shift
+            sigmas = shift * sigmas / (1 +
+                                       (shift - 1) * sigmas)  # pyright: ignore
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = ((1 - self.alphas_cumprod[0]) /
+                          self.alphas_cumprod[0])**0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        timesteps = sigmas * self.config.num_train_timesteps
+        sigmas = np.concatenate([sigmas, [sigma_last]
+                                ]).astype(np.float32)  # pyright: ignore
+
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(
+            device=device, dtype=torch.int64)
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+
+        self._step_index = None
+        self._begin_index = None
+        # self.sigmas = self.sigmas.to(
+        #     "cpu")  # to avoid too much CPU/GPU communication
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float(
+            )  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(
+            abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(
+            1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(
+            sample, -s, s
+        ) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        return 1 - sigma, sigma
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.convert_model_output
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is
+        designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an
+        integral of the data prediction model.
+        <Tip>
+        The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise
+        prediction and data prediction models.
+        </Tip>
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(
+                    "missing `sample` as a required keyward argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        # DPM-Solver++ needs to solve an integral of the data prediction model.
+        if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction`, or `flow_prediction` for the FlowDPMSolverMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+
+        # DPM-Solver needs to solve an integral of the noise prediction model.
+        elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                epsilon = sample - (1 - sigma_t) * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the FlowDPMSolverMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = model_output + x0_pred
+
+            return epsilon
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.dpm_solver_first_order_update
+    def dpm_solver_first_order_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the first-order DPMSolver (equivalent to DDIM).
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop(
+            "prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(
+                    " missing `sample` as a required keyward argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[
+            self.step_index]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s)
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s = torch.log(alpha_s) - torch.log(sigma_s)
+
+        h = lambda_t - lambda_s
+        if self.config.algorithm_type == "dpmsolver++":
+            x_t = (sigma_t /
+                   sigma_s) * sample - (alpha_t *
+                                        (torch.exp(-h) - 1.0)) * model_output
+        elif self.config.algorithm_type == "dpmsolver":
+            x_t = (alpha_t /
+                   alpha_s) * sample - (sigma_t *
+                                        (torch.exp(h) - 1.0)) * model_output
+        elif self.config.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            x_t = ((sigma_t / sigma_s * torch.exp(-h)) * sample +
+                   (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output +
+                   sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise)
+        elif self.config.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            x_t = ((alpha_t / alpha_s) * sample - 2.0 *
+                   (sigma_t * (torch.exp(h) - 1.0)) * model_output +
+                   sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise)
+        return x_t  # pyright: ignore
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_second_order_update
+    def multistep_dpm_solver_second_order_update(
+        self,
+        model_output_list: List[torch.Tensor],
+        *args,
+        sample: torch.Tensor = None,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the second-order multistep DPMSolver.
+        Args:
+            model_output_list (`List[torch.Tensor]`):
+                The direct outputs from learned diffusion model at current and latter timesteps.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        timestep_list = args[0] if len(args) > 0 else kwargs.pop(
+            "timestep_list", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop(
+            "prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(
+                    " missing `sample` as a required keyward argument")
+        if timestep_list is not None:
+            deprecate(
+                "timestep_list",
+                "1.0.0",
+                "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s0, sigma_s1 = (
+            self.sigmas[self.step_index + 1],  # pyright: ignore
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],  # pyright: ignore
+        )
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+        alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+
+        m0, m1 = model_output_list[-1], model_output_list[-2]
+
+        h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
+        r0 = h_0 / h
+        D0, D1 = m0, (1.0 / r0) * (m0 - m1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2211.01095 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                x_t = ((sigma_t / sigma_s0) * sample -
+                       (alpha_t * (torch.exp(-h) - 1.0)) * D0 - 0.5 *
+                       (alpha_t * (torch.exp(-h) - 1.0)) * D1)
+            elif self.config.solver_type == "heun":
+                x_t = ((sigma_t / sigma_s0) * sample -
+                       (alpha_t * (torch.exp(-h) - 1.0)) * D0 +
+                       (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1)
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                x_t = ((alpha_t / alpha_s0) * sample -
+                       (sigma_t * (torch.exp(h) - 1.0)) * D0 - 0.5 *
+                       (sigma_t * (torch.exp(h) - 1.0)) * D1)
+            elif self.config.solver_type == "heun":
+                x_t = ((alpha_t / alpha_s0) * sample -
+                       (sigma_t * (torch.exp(h) - 1.0)) * D0 -
+                       (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1)
+        elif self.config.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            if self.config.solver_type == "midpoint":
+                x_t = ((sigma_t / sigma_s0 * torch.exp(-h)) * sample +
+                       (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 + 0.5 *
+                       (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 +
+                       sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise)
+            elif self.config.solver_type == "heun":
+                x_t = ((sigma_t / sigma_s0 * torch.exp(-h)) * sample +
+                       (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 +
+                       (alpha_t * ((1.0 - torch.exp(-2.0 * h)) /
+                                   (-2.0 * h) + 1.0)) * D1 +
+                       sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise)
+        elif self.config.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            if self.config.solver_type == "midpoint":
+                x_t = ((alpha_t / alpha_s0) * sample - 2.0 *
+                       (sigma_t * (torch.exp(h) - 1.0)) * D0 -
+                       (sigma_t * (torch.exp(h) - 1.0)) * D1 +
+                       sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise)
+            elif self.config.solver_type == "heun":
+                x_t = ((alpha_t / alpha_s0) * sample - 2.0 *
+                       (sigma_t * (torch.exp(h) - 1.0)) * D0 - 2.0 *
+                       (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 +
+                       sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise)
+        return x_t  # pyright: ignore
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_third_order_update
+    def multistep_dpm_solver_third_order_update(
+        self,
+        model_output_list: List[torch.Tensor],
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the third-order multistep DPMSolver.
+        Args:
+            model_output_list (`List[torch.Tensor]`):
+                The direct outputs from learned diffusion model at current and latter timesteps.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by diffusion process.
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+
+        timestep_list = args[0] if len(args) > 0 else kwargs.pop(
+            "timestep_list", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop(
+            "prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(
+                    " missing`sample` as a required keyward argument")
+        if timestep_list is not None:
+            deprecate(
+                "timestep_list",
+                "1.0.0",
+                "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s0, sigma_s1, sigma_s2 = (
+            self.sigmas[self.step_index + 1],  # pyright: ignore
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],  # pyright: ignore
+            self.sigmas[self.step_index - 2],  # pyright: ignore
+        )
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+        alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+        alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+        lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2)
+
+        m0, m1, m2 = model_output_list[-1], model_output_list[
+            -2], model_output_list[-3]
+
+        h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2
+        r0, r1 = h_0 / h, h_1 / h
+        D0 = m0
+        D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)
+        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+        D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = ((sigma_t / sigma_s0) * sample -
+                   (alpha_t * (torch.exp(-h) - 1.0)) * D0 +
+                   (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 -
+                   (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2)
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = ((alpha_t / alpha_s0) * sample - (sigma_t *
+                                                    (torch.exp(h) - 1.0)) * D0 -
+                   (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 -
+                   (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2)
+        return x_t  # pyright: ignore
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    # Modified from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.step
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[int, torch.Tensor],
+        sample: torch.Tensor,
+        generator=None,
+        variance_noise: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[SchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep DPMSolver.
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            variance_noise (`torch.Tensor`):
+                Alternative to generating noise with `generator` by directly providing the noise for the variance
+                itself. Useful for methods such as [`LEdits++`].
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Improve numerical stability for small number of steps
+        lower_order_final = (self.step_index == len(self.timesteps) - 1) and (
+            self.config.euler_at_final or
+            (self.config.lower_order_final and len(self.timesteps) < 15) or
+            self.config.final_sigmas_type == "zero")
+        lower_order_second = ((self.step_index == len(self.timesteps) - 2) and
+                              self.config.lower_order_final and
+                              len(self.timesteps) < 15)
+
+        model_output = self.convert_model_output(model_output, sample=sample)
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+        self.model_outputs[-1] = model_output
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"
+                                         ] and variance_noise is None:
+            noise = randn_tensor(
+                model_output.shape,
+                generator=generator,
+                device=model_output.device,
+                dtype=torch.float32)
+        elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
+            noise = variance_noise.to(
+                device=model_output.device,
+                dtype=torch.float32)  # pyright: ignore
+        else:
+            noise = None
+
+        if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
+            prev_sample = self.dpm_solver_first_order_update(
+                model_output, sample=sample, noise=noise)
+        elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:
+            prev_sample = self.multistep_dpm_solver_second_order_update(
+                self.model_outputs, sample=sample, noise=noise)
+        else:
+            prev_sample = self.multistep_dpm_solver_third_order_update(
+                self.model_outputs, sample=sample)
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # Cast sample back to expected dtype
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        # upon completion increase step index by one
+        self._step_index += 1  # pyright: ignore
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input
+    def scale_model_input(self, sample: torch.Tensor, *args,
+                          **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(
+            device=original_samples.device, dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(
+                timesteps):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(
+                original_samples.device, dtype=torch.float32)
+            timesteps = timesteps.to(
+                original_samples.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(original_samples.device)
+            timesteps = timesteps.to(original_samples.device)
+
+        # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps)
+                for t in timesteps
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timesteps.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timesteps.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps

humo/models/utils/fm_solvers_unipc.py

@@ -0,0 +1,800 @@
+# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py
+# Convert unipc for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+                                                   SchedulerMixin,
+                                                   SchedulerOutput)
+from diffusers.utils import deprecate, is_scipy_available
+
+if is_scipy_available():
+    import scipy.stats
+
+
+class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        solver_order (`int`, default `2`):
+            The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
+            due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
+            unconditional sampling.
+        prediction_type (`str`, defaults to "flow_prediction"):
+            Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+            the flow of the diffusion process.
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
+        predict_x0 (`bool`, defaults to `True`):
+            Whether to use the updating algorithm on the predicted x0.
+        solver_type (`str`, default `bh2`):
+            Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
+            otherwise.
+        lower_order_final (`bool`, default `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        disable_corrector (`list`, default `[]`):
+            Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
+            and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
+            usually disabled during the first few steps.
+        solver_p (`SchedulerMixin`, default `None`):
+            Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
+        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+            the sigmas are determined according to a sequence of noise levels {σi}.
+        use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        final_sigmas_type (`str`, defaults to `"zero"`):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+            self,
+            num_train_timesteps: int = 1000,
+            solver_order: int = 2,
+            prediction_type: str = "flow_prediction",
+            shift: Optional[float] = 1.0,
+            use_dynamic_shifting=False,
+            thresholding: bool = False,
+            dynamic_thresholding_ratio: float = 0.995,
+            sample_max_value: float = 1.0,
+            predict_x0: bool = True,
+            solver_type: str = "bh2",
+            lower_order_final: bool = True,
+            disable_corrector: List[int] = [],
+            solver_p: SchedulerMixin = None,
+            timestep_spacing: str = "linspace",
+            steps_offset: int = 0,
+            final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+    ):
+
+        if solver_type not in ["bh1", "bh2"]:
+            if solver_type in ["midpoint", "heun", "logrho"]:
+                self.register_to_config(solver_type="bh2")
+            else:
+                raise NotImplementedError(
+                    f"{solver_type} is not implemented for {self.__class__}")
+
+        self.predict_x0 = predict_x0
+        # setable values
+        self.num_inference_steps = None
+        alphas = np.linspace(1, 1 / num_train_timesteps,
+                             num_train_timesteps)[::-1].copy()
+        sigmas = 1.0 - alphas
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 +
+                                       (shift - 1) * sigmas)  # pyright: ignore
+
+        self.sigmas = sigmas
+        self.timesteps = sigmas * num_train_timesteps
+
+        self.model_outputs = [None] * solver_order
+        self.timestep_list = [None] * solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = disable_corrector
+        self.solver_p = solver_p
+        self.last_sample = None
+        self._step_index = None
+        self._begin_index = None
+
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+    def set_timesteps(
+        self,
+        num_inference_steps: Union[int, None] = None,
+        device: Union[str, torch.device] = None,
+        sigmas: Optional[List[float]] = None,
+        mu: Optional[Union[float, None]] = None,
+        shift: Optional[Union[float, None]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+        Args:
+            num_inference_steps (`int`):
+                Total number of the spacing of the time steps.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        """
+
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError(
+                " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+            )
+
+        if sigmas is None:
+            sigmas = np.linspace(self.sigma_max, self.sigma_min,
+                                 num_inference_steps +
+                                 1).copy()[:-1]  # pyright: ignore
+
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)  # pyright: ignore
+        else:
+            if shift is None:
+                shift = self.config.shift
+            sigmas = shift * sigmas / (1 +
+                                       (shift - 1) * sigmas)  # pyright: ignore
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = ((1 - self.alphas_cumprod[0]) /
+                          self.alphas_cumprod[0])**0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        timesteps = sigmas * self.config.num_train_timesteps
+        sigmas = np.concatenate([sigmas, [sigma_last]
+                                ]).astype(np.float32)  # pyright: ignore
+
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(
+            device=device, dtype=torch.int64)
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+        self.last_sample = None
+        if self.solver_p:
+            self.solver_p.set_timesteps(self.num_inference_steps, device=device)
+
+        # add an index counter for schedulers that allow duplicated timesteps
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to(
+            "cpu")  # to avoid too much CPU/GPU communication
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float(
+            )  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(
+            abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(
+            1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(
+            sample, -s, s
+        ) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma):
+        return sigma * self.config.num_train_timesteps
+
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        return 1 - sigma, sigma
+
+    # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Convert the model output to the corresponding type the UniPC algorithm needs.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(
+                    "missing `sample` as a required keyward argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma = self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
+        if self.predict_x0:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+        else:
+            if self.config.prediction_type == "flow_prediction":
+                sigma_t = self.sigmas[self.step_index]
+                epsilon = sample - (1 - sigma_t) * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+                    " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma_t = self.sigmas[self.step_index]
+                x0_pred = sample - sigma_t * model_output
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = model_output + x0_pred
+
+            return epsilon
+
+    def multistep_uni_p_bh_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        order: int = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model at the current timestep.
+            prev_timestep (`int`):
+                The previous discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            order (`int`):
+                The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        prev_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "prev_timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError(
+                    " missing `sample` as a required keyward argument")
+        if order is None:
+            if len(args) > 2:
+                order = args[2]
+            else:
+                raise ValueError(
+                    " missing `order` as a required keyward argument")
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+        model_output_list = self.model_outputs
+
+        s0 = self.timestep_list[-1]
+        m0 = model_output_list[-1]
+        x = sample
+
+        if self.solver_p:
+            x_t = self.solver_p.step(model_output, s0, x).prev_sample
+            return x_t
+
+        sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[
+            self.step_index]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - i  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)  # (B, K)
+            # for order 2, we use a simplified version
+            if order == 2:
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+            else:
+                rhos_p = torch.linalg.solve(R[:-1, :-1],
+                                            b[:-1]).to(device).to(x.dtype)
+        else:
+            D1s = None
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+                                        D1s)  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - alpha_t * B_h * pred_res
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+                                        D1s)  # pyright: ignore
+            else:
+                pred_res = 0
+            x_t = x_t_ - sigma_t * B_h * pred_res
+
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def multistep_uni_c_bh_update(
+        self,
+        this_model_output: torch.Tensor,
+        *args,
+        last_sample: torch.Tensor = None,
+        this_sample: torch.Tensor = None,
+        order: int = None,  # pyright: ignore
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the UniC (B(h) version).
+
+        Args:
+            this_model_output (`torch.Tensor`):
+                The model outputs at `x_t`.
+            this_timestep (`int`):
+                The current timestep `t`.
+            last_sample (`torch.Tensor`):
+                The generated sample before the last predictor `x_{t-1}`.
+            this_sample (`torch.Tensor`):
+                The generated sample after the last predictor `x_{t}`.
+            order (`int`):
+                The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+
+        Returns:
+            `torch.Tensor`:
+                The corrected sample tensor at the current timestep.
+        """
+        this_timestep = args[0] if len(args) > 0 else kwargs.pop(
+            "this_timestep", None)
+        if last_sample is None:
+            if len(args) > 1:
+                last_sample = args[1]
+            else:
+                raise ValueError(
+                    " missing`last_sample` as a required keyward argument")
+        if this_sample is None:
+            if len(args) > 2:
+                this_sample = args[2]
+            else:
+                raise ValueError(
+                    " missing`this_sample` as a required keyward argument")
+        if order is None:
+            if len(args) > 3:
+                order = args[3]
+            else:
+                raise ValueError(
+                    " missing`order` as a required keyward argument")
+        if this_timestep is not None:
+            deprecate(
+                "this_timestep",
+                "1.0.0",
+                "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        model_output_list = self.model_outputs
+
+        m0 = model_output_list[-1]
+        x = last_sample
+        x_t = this_sample
+        model_t = this_model_output
+
+        sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[
+            self.step_index - 1]  # pyright: ignore
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+        h = lambda_t - lambda_s0
+        device = this_sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - (i + 1)  # pyright: ignore
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)  # pyright: ignore
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)  # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+        else:
+            D1s = None
+
+        # for order 1, we use a simplified version
+        if order == 1:
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+        else:
+            rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            if D1s is not None:
+                corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+            else:
+                corr_res = 0
+            D1_t = model_t - m0
+            x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        x_t = x_t.to(x.dtype)
+        return x_t
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
+    def _init_step_index(self, timestep):
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(self,
+             model_output: torch.Tensor,
+             timestep: Union[int, torch.Tensor],
+             sample: torch.Tensor,
+             return_dict: bool = True,
+             generator=None) -> Union[SchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep UniPC.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
+
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        use_corrector = (
+            self.step_index > 0 and
+            self.step_index - 1 not in self.disable_corrector and
+            self.last_sample is not None  # pyright: ignore
+        )
+
+        model_output_convert = self.convert_model_output(
+            model_output, sample=sample)
+        if use_corrector:
+            sample = self.multistep_uni_c_bh_update(
+                this_model_output=model_output_convert,
+                last_sample=self.last_sample,
+                this_sample=sample,
+                order=self.this_order,
+            )
+
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+            self.timestep_list[i] = self.timestep_list[i + 1]
+
+        self.model_outputs[-1] = model_output_convert
+        self.timestep_list[-1] = timestep  # pyright: ignore
+
+        if self.config.lower_order_final:
+            this_order = min(self.config.solver_order,
+                             len(self.timesteps) -
+                             self.step_index)  # pyright: ignore
+        else:
+            this_order = self.config.solver_order
+
+        self.this_order = min(this_order,
+                              self.lower_order_nums + 1)  # warmup for multistep
+        assert self.this_order > 0
+
+        self.last_sample = sample
+        prev_sample = self.multistep_uni_p_bh_update(
+            model_output=model_output,  # pass the original non-converted model output, in case solver-p is used
+            sample=sample,
+            order=self.this_order,
+        )
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # upon completion increase step index by one
+        self._step_index += 1  # pyright: ignore
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    def scale_model_input(self, sample: torch.Tensor, *args,
+                          **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(
+            device=original_samples.device, dtype=original_samples.dtype)
+        if original_samples.device.type == "mps" and torch.is_floating_point(
+                timesteps):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(
+                original_samples.device, dtype=torch.float32)
+            timesteps = timesteps.to(
+                original_samples.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(original_samples.device)
+            timesteps = timesteps.to(original_samples.device)
+
+        # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [
+                self.index_for_timestep(t, schedule_timesteps)
+                for t in timesteps
+            ]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timesteps.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timesteps.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(original_samples.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+
+    def __len__(self):
+        return self.config.num_train_timesteps

humo/models/utils/utils.py

@@ -0,0 +1,58 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import argparse
+import binascii
+import os
+import os.path as osp
+import json
+from omegaconf import OmegaConf
+
+import imageio
+import torch
+import torchvision
+from moviepy.editor import AudioFileClip, VideoClip
+
+__all__ = ['tensor_to_video', 'prepare_json_dataset']
+    
+
+def tensor_to_video(tensor, output_video_path, input_audio_path, fps=25):
+    """
+    Converts a Tensor with shape [c, f, h, w] into a video and adds an audio track from the specified audio file.
+
+    Args:
+        tensor (numpy): The Tensor to be converted, shaped [f, h, w, c].
+        output_video_path (str): The file path where the output video will be saved.
+        input_audio_path (str): The path to the audio file (WAV file) that contains the audio track to be added.
+        fps (int): The frame rate of the output video. Default is 30 fps.
+    """
+    def make_frame(t):
+        frame_index = min(int(t * fps), tensor.shape[0] - 1)
+        return tensor[frame_index]
+
+    video_duration = tensor.shape[0] / fps
+    audio_clip = AudioFileClip(input_audio_path)
+    audio_duration = audio_clip.duration
+    final_duration = min(video_duration, audio_duration)
+    audio_clip = audio_clip.subclip(0, final_duration)
+    new_video_clip = VideoClip(make_frame, duration=final_duration)
+    new_video_clip = new_video_clip.set_audio(audio_clip)
+    new_video_clip.write_videofile(output_video_path, fps=fps, audio_codec="aac")
+
+
+def prepare_json_dataset(json_path):
+    samples = []
+    with open(json_path, "rb") as f:
+        data = json.load(f)
+    for itemname, row in data.items():
+        text = row['prompt'].strip().replace("_", " ").strip('"')
+        audio_path = row['audio_path']
+        ref_img_path = [x for x in row['img_paths']]
+
+        samples.append({
+            "text": text,
+            "ref_img": ref_img_path,
+            "audio": audio_path,
+            "itemname": itemname
+        })
+    samples = OmegaConf.create(samples)
+    
+    return samples

humo/models/wan_modules/__init__.py

@@ -0,0 +1,16 @@
+from .attention import flash_attention
+from .model import WanModel
+from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
+from .tokenizers import HuggingfaceTokenizer
+from .vae import WanVAE
+
+__all__ = [
+    'WanVAE',
+    'WanModel',
+    'T5Model',
+    'T5Encoder',
+    'T5Decoder',
+    'T5EncoderModel',
+    'HuggingfaceTokenizer',
+    'flash_attention',
+]

humo/models/wan_modules/attention.py

@@ -0,0 +1,256 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import warnings
+import torch
+from typing import Optional, Tuple
+
+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
+
+
+__all__ = [
+    'flash_attention',
+    'attention',
+]
+
+
+# ---------------------------
+# Custom op + fake kernel
+# ---------------------------
+from typing import Optional, Sequence  # <- add Sequence
+
+# ... imports unchanged ...
+from typing import Optional, Sequence
+
+@torch.library.custom_op("wan::flash_attention", mutates_args=())
+def _wan_flash_attention_op(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    q_lens: Optional[torch.Tensor] = None,
+    k_lens: Optional[torch.Tensor] = None,
+    dropout_p: float = 0.0,
+    softmax_scale: Optional[float] = None,
+    q_scale: Optional[float] = None,
+    causal: bool = False,
+    # IMPORTANT: schema-friendly default (None), not a tuple
+    window_size: Optional[Sequence[int]] = None,
+    deterministic: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
+    version: Optional[int] = None,
+) -> torch.Tensor:
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.size(-1) <= 256
+
+    # normalize window_size to a 2-tuple for FA2 API
+    if window_size is None:
+        ws = (-1, -1)
+    else:
+        ws = tuple(window_size)
+        if len(ws) != 2:
+            raise ValueError(f"window_size must have length 2; got {window_size!r}")
+
+    b, lq, nheads = q.shape[0], q.shape[1], q.shape[2]
+    lk = k.shape[1]
+    out_dtype = q.dtype
+
+    def half(x: torch.Tensor) -> torch.Tensor:
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # --- preprocess (unchanged) ---
+    if q_lens is None:
+        q_flat = half(q.flatten(0, 1))
+        q_lens = torch.tensor([lq] * b, dtype=torch.int32)
+    else:
+        q_flat = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+    if k_lens is None:
+        k_flat = half(k.flatten(0, 1))
+        v_flat = half(v.flatten(0, 1))
+        k_lens = torch.tensor([lk] * b, dtype=torch.int32)
+    else:
+        k_flat = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+        v_flat = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+    q_flat = q_flat.to(v_flat.dtype); k_flat = k_flat.to(v_flat.dtype)
+    if q_scale is not None:
+        q_flat = q_flat * q_scale
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn('Flash attention 3 is not available, use flash attention 2 instead.')
+
+    if FLASH_ATTN_3_AVAILABLE:
+        ret = flash_attn_interface.flash_attn_varlen_func(
+            q=q_flat,
+            k=k_flat,
+            v=v_flat,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(0, dtype=torch.int32).to(q_flat.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(0, dtype=torch.int32).to(k_flat.device, non_blocking=True),
+            seqused_q=None,
+            seqused_k=None,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic,
+        )
+        out0 = ret[0] if isinstance(ret, (tuple, list)) else ret
+        total_q = b * lq
+        if out0.dim() != 3:
+            raise RuntimeError(f"Unexpected FA3 output rank {out0.dim()} shape={tuple(out0.shape)}")
+        if out0.shape[0] == total_q:
+            out_flat = out0
+        elif out0.shape[0] == nheads and out0.shape[1] == total_q:
+            out_flat = out0.transpose(0, 1).contiguous()
+        else:
+            raise RuntimeError(f"Unexpected FA3 output shape {tuple(out0.shape)}")
+        out = out_flat.unflatten(0, (b, lq))
+
+    elif FLASH_ATTN_2_AVAILABLE:
+        out = flash_attn.flash_attn_varlen_func(
+            q=q_flat,
+            k=k_flat,
+            v=v_flat,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(0, dtype=torch.int32).to(q_flat.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(0, dtype=torch.int32).to(q_flat.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=ws,                 # <- pass 2-tuple
+            deterministic=deterministic,
+        ).unflatten(0, (b, lq))
+    else:
+        q_s = q.transpose(1, 2).to(dtype)
+        k_s = k.transpose(1, 2).to(dtype)
+        v_s = v.transpose(1, 2).to(dtype)
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q_s, k_s, v_s, attn_mask=None, is_causal=causal, dropout_p=dropout_p
+        ).transpose(1, 2).contiguous()
+
+    return out.to(out_dtype)
+
+@_wan_flash_attention_op.register_fake
+def _wan_flash_attention_op_fake(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p: float = 0.0,
+    softmax_scale=None,
+    q_scale=None,
+    causal: bool = False,
+    window_size: Optional[Sequence[int]] = None,
+    deterministic: bool = False,
+    dtype: torch.dtype = torch.bfloat16,
+    version: Optional[int] = None,
+):
+    # Match output shape: (B, Lq, Nq, Dh_v) and keep the SAME fake device as `q`
+    B, Lq, Nq, _ = q.shape
+    Dh_v = v.shape[-1]
+    return q.new_empty((B, Lq, Nq, Dh_v), dtype=q.dtype)
+
+
+
+# ---------------------------
+# Public API (unchanged signature)
+# ---------------------------
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    # Simply delegate to the custom op so Dynamo/AOT treats it as a single node;
+    # our eager kernel inside _wan_flash_attention_op keeps the original behavior.
+    return _wan_flash_attention_op(
+        q, k, v,
+        q_lens=q_lens,
+        k_lens=k_lens,
+        dropout_p=dropout_p,
+        softmax_scale=softmax_scale,
+        q_scale=q_scale,
+        causal=causal,
+        window_size=window_size,
+        deterministic=deterministic,
+        dtype=dtype,
+        version=version,
+    )
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=False,
+    dtype=torch.bfloat16,
+    fa_version=None,
+):
+    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+    else:
+        if q_lens is not None or k_lens is not None:
+            warnings.warn(
+                'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+            )
+        q_ = q.transpose(1, 2).to(dtype)
+        k_ = k.transpose(1, 2).to(dtype)
+        v_ = v.transpose(1, 2).to(dtype)
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q_, k_, v_, attn_mask=None, is_causal=causal, dropout_p=dropout_p
+        )
+        return out.transpose(1, 2).contiguous()

humo/models/wan_modules/clip.py

@@ -0,0 +1,542 @@
+# Modified from ``https://github.com/openai/CLIP'' and ``https://github.com/mlfoundations/open_clip''
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+
+from .attention import flash_attention
+from .tokenizers import HuggingfaceTokenizer
+from .xlm_roberta import XLMRoberta
+
+__all__ = [
+    'XLMRobertaCLIP',
+    'clip_xlm_roberta_vit_h_14',
+    'CLIPModel',
+]
+
+
+def pos_interpolate(pos, seq_len):
+    if pos.size(1) == seq_len:
+        return pos
+    else:
+        src_grid = int(math.sqrt(pos.size(1)))
+        tar_grid = int(math.sqrt(seq_len))
+        n = pos.size(1) - src_grid * src_grid
+        return torch.cat([
+            pos[:, :n],
+            F.interpolate(
+                pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute(
+                    0, 3, 1, 2),
+                size=(tar_grid, tar_grid),
+                mode='bicubic',
+                align_corners=False).flatten(2).transpose(1, 2)
+        ],
+                         dim=1)
+
+
+class QuickGELU(nn.Module):
+
+    def forward(self, x):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class LayerNorm(nn.LayerNorm):
+
+    def forward(self, x):
+        return super().forward(x.float()).type_as(x)
+
+
+class SelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 causal=False,
+                 attn_dropout=0.0,
+                 proj_dropout=0.0):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.causal = causal
+        self.attn_dropout = attn_dropout
+        self.proj_dropout = proj_dropout
+
+        # layers
+        self.to_qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        """
+        x:   [B, L, C].
+        """
+        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q, k, v = self.to_qkv(x).view(b, s, 3, n, d).unbind(2)
+
+        # compute attention
+        p = self.attn_dropout if self.training else 0.0
+        x = flash_attention(q, k, v, dropout_p=p, causal=self.causal, version=2)
+        x = x.reshape(b, s, c)
+
+        # output
+        x = self.proj(x)
+        x = F.dropout(x, self.proj_dropout, self.training)
+        return x
+
+
+class SwiGLU(nn.Module):
+
+    def __init__(self, dim, mid_dim):
+        super().__init__()
+        self.dim = dim
+        self.mid_dim = mid_dim
+
+        # layers
+        self.fc1 = nn.Linear(dim, mid_dim)
+        self.fc2 = nn.Linear(dim, mid_dim)
+        self.fc3 = nn.Linear(mid_dim, dim)
+
+    def forward(self, x):
+        x = F.silu(self.fc1(x)) * self.fc2(x)
+        x = self.fc3(x)
+        return x
+
+
+class AttentionBlock(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 mlp_ratio,
+                 num_heads,
+                 post_norm=False,
+                 causal=False,
+                 activation='quick_gelu',
+                 attn_dropout=0.0,
+                 proj_dropout=0.0,
+                 norm_eps=1e-5):
+        assert activation in ['quick_gelu', 'gelu', 'swi_glu']
+        super().__init__()
+        self.dim = dim
+        self.mlp_ratio = mlp_ratio
+        self.num_heads = num_heads
+        self.post_norm = post_norm
+        self.causal = causal
+        self.norm_eps = norm_eps
+
+        # layers
+        self.norm1 = LayerNorm(dim, eps=norm_eps)
+        self.attn = SelfAttention(dim, num_heads, causal, attn_dropout,
+                                  proj_dropout)
+        self.norm2 = LayerNorm(dim, eps=norm_eps)
+        if activation == 'swi_glu':
+            self.mlp = SwiGLU(dim, int(dim * mlp_ratio))
+        else:
+            self.mlp = nn.Sequential(
+                nn.Linear(dim, int(dim * mlp_ratio)),
+                QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
+                nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
+
+    def forward(self, x):
+        if self.post_norm:
+            x = x + self.norm1(self.attn(x))
+            x = x + self.norm2(self.mlp(x))
+        else:
+            x = x + self.attn(self.norm1(x))
+            x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class AttentionPool(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 mlp_ratio,
+                 num_heads,
+                 activation='gelu',
+                 proj_dropout=0.0,
+                 norm_eps=1e-5):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.mlp_ratio = mlp_ratio
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.proj_dropout = proj_dropout
+        self.norm_eps = norm_eps
+
+        # layers
+        gain = 1.0 / math.sqrt(dim)
+        self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
+        self.to_q = nn.Linear(dim, dim)
+        self.to_kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+        self.norm = LayerNorm(dim, eps=norm_eps)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, int(dim * mlp_ratio)),
+            QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
+            nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
+
+    def forward(self, x):
+        """
+        x:  [B, L, C].
+        """
+        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.to_q(self.cls_embedding).view(1, 1, n, d).expand(b, -1, -1, -1)
+        k, v = self.to_kv(x).view(b, s, 2, n, d).unbind(2)
+
+        # compute attention
+        x = flash_attention(q, k, v, version=2)
+        x = x.reshape(b, 1, c)
+
+        # output
+        x = self.proj(x)
+        x = F.dropout(x, self.proj_dropout, self.training)
+
+        # mlp
+        x = x + self.mlp(self.norm(x))
+        return x[:, 0]
+
+
+class VisionTransformer(nn.Module):
+
+    def __init__(self,
+                 image_size=224,
+                 patch_size=16,
+                 dim=768,
+                 mlp_ratio=4,
+                 out_dim=512,
+                 num_heads=12,
+                 num_layers=12,
+                 pool_type='token',
+                 pre_norm=True,
+                 post_norm=False,
+                 activation='quick_gelu',
+                 attn_dropout=0.0,
+                 proj_dropout=0.0,
+                 embedding_dropout=0.0,
+                 norm_eps=1e-5):
+        if image_size % patch_size != 0:
+            print(
+                '[WARNING] image_size is not divisible by patch_size',
+                flush=True)
+        assert pool_type in ('token', 'token_fc', 'attn_pool')
+        out_dim = out_dim or dim
+        super().__init__()
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_patches = (image_size // patch_size)**2
+        self.dim = dim
+        self.mlp_ratio = mlp_ratio
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.pool_type = pool_type
+        self.post_norm = post_norm
+        self.norm_eps = norm_eps
+
+        # embeddings
+        gain = 1.0 / math.sqrt(dim)
+        self.patch_embedding = nn.Conv2d(
+            3,
+            dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=not pre_norm)
+        if pool_type in ('token', 'token_fc'):
+            self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
+        self.pos_embedding = nn.Parameter(gain * torch.randn(
+            1, self.num_patches +
+            (1 if pool_type in ('token', 'token_fc') else 0), dim))
+        self.dropout = nn.Dropout(embedding_dropout)
+
+        # transformer
+        self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None
+        self.transformer = nn.Sequential(*[
+            AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False,
+                           activation, attn_dropout, proj_dropout, norm_eps)
+            for _ in range(num_layers)
+        ])
+        self.post_norm = LayerNorm(dim, eps=norm_eps)
+
+        # head
+        if pool_type == 'token':
+            self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
+        elif pool_type == 'token_fc':
+            self.head = nn.Linear(dim, out_dim)
+        elif pool_type == 'attn_pool':
+            self.head = AttentionPool(dim, mlp_ratio, num_heads, activation,
+                                      proj_dropout, norm_eps)
+
+    def forward(self, x, interpolation=False, use_31_block=False):
+        b = x.size(0)
+
+        # embeddings
+        x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
+        if self.pool_type in ('token', 'token_fc'):
+            x = torch.cat([self.cls_embedding.expand(b, -1, -1), x], dim=1)
+        if interpolation:
+            e = pos_interpolate(self.pos_embedding, x.size(1))
+        else:
+            e = self.pos_embedding
+        x = self.dropout(x + e)
+        if self.pre_norm is not None:
+            x = self.pre_norm(x)
+
+        # transformer
+        if use_31_block:
+            x = self.transformer[:-1](x)
+            return x
+        else:
+            x = self.transformer(x)
+            return x
+
+
+class XLMRobertaWithHead(XLMRoberta):
+
+    def __init__(self, **kwargs):
+        self.out_dim = kwargs.pop('out_dim')
+        super().__init__(**kwargs)
+
+        # head
+        mid_dim = (self.dim + self.out_dim) // 2
+        self.head = nn.Sequential(
+            nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(),
+            nn.Linear(mid_dim, self.out_dim, bias=False))
+
+    def forward(self, ids):
+        # xlm-roberta
+        x = super().forward(ids)
+
+        # average pooling
+        mask = ids.ne(self.pad_id).unsqueeze(-1).to(x)
+        x = (x * mask).sum(dim=1) / mask.sum(dim=1)
+
+        # head
+        x = self.head(x)
+        return x
+
+
+class XLMRobertaCLIP(nn.Module):
+
+    def __init__(self,
+                 embed_dim=1024,
+                 image_size=224,
+                 patch_size=14,
+                 vision_dim=1280,
+                 vision_mlp_ratio=4,
+                 vision_heads=16,
+                 vision_layers=32,
+                 vision_pool='token',
+                 vision_pre_norm=True,
+                 vision_post_norm=False,
+                 activation='gelu',
+                 vocab_size=250002,
+                 max_text_len=514,
+                 type_size=1,
+                 pad_id=1,
+                 text_dim=1024,
+                 text_heads=16,
+                 text_layers=24,
+                 text_post_norm=True,
+                 text_dropout=0.1,
+                 attn_dropout=0.0,
+                 proj_dropout=0.0,
+                 embedding_dropout=0.0,
+                 norm_eps=1e-5):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.vision_dim = vision_dim
+        self.vision_mlp_ratio = vision_mlp_ratio
+        self.vision_heads = vision_heads
+        self.vision_layers = vision_layers
+        self.vision_pre_norm = vision_pre_norm
+        self.vision_post_norm = vision_post_norm
+        self.activation = activation
+        self.vocab_size = vocab_size
+        self.max_text_len = max_text_len
+        self.type_size = type_size
+        self.pad_id = pad_id
+        self.text_dim = text_dim
+        self.text_heads = text_heads
+        self.text_layers = text_layers
+        self.text_post_norm = text_post_norm
+        self.norm_eps = norm_eps
+
+        # models
+        self.visual = VisionTransformer(
+            image_size=image_size,
+            patch_size=patch_size,
+            dim=vision_dim,
+            mlp_ratio=vision_mlp_ratio,
+            out_dim=embed_dim,
+            num_heads=vision_heads,
+            num_layers=vision_layers,
+            pool_type=vision_pool,
+            pre_norm=vision_pre_norm,
+            post_norm=vision_post_norm,
+            activation=activation,
+            attn_dropout=attn_dropout,
+            proj_dropout=proj_dropout,
+            embedding_dropout=embedding_dropout,
+            norm_eps=norm_eps)
+        self.textual = XLMRobertaWithHead(
+            vocab_size=vocab_size,
+            max_seq_len=max_text_len,
+            type_size=type_size,
+            pad_id=pad_id,
+            dim=text_dim,
+            out_dim=embed_dim,
+            num_heads=text_heads,
+            num_layers=text_layers,
+            post_norm=text_post_norm,
+            dropout=text_dropout)
+        self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([]))
+
+    def forward(self, imgs, txt_ids):
+        """
+        imgs:       [B, 3, H, W] of torch.float32.
+        - mean:     [0.48145466, 0.4578275, 0.40821073]
+        - std:      [0.26862954, 0.26130258, 0.27577711]
+        txt_ids:    [B, L] of torch.long.
+                    Encoded by data.CLIPTokenizer.
+        """
+        xi = self.visual(imgs)
+        xt = self.textual(txt_ids)
+        return xi, xt
+
+    def param_groups(self):
+        groups = [{
+            'params': [
+                p for n, p in self.named_parameters()
+                if 'norm' in n or n.endswith('bias')
+            ],
+            'weight_decay': 0.0
+        }, {
+            'params': [
+                p for n, p in self.named_parameters()
+                if not ('norm' in n or n.endswith('bias'))
+            ]
+        }]
+        return groups
+
+
+def _clip(pretrained=False,
+          pretrained_name=None,
+          model_cls=XLMRobertaCLIP,
+          return_transforms=False,
+          return_tokenizer=False,
+          tokenizer_padding='eos',
+          dtype=torch.float32,
+          device='cpu',
+          **kwargs):
+    # init a model on device
+    with torch.device(device):
+        model = model_cls(**kwargs)
+
+    # set device
+    model = model.to(dtype=dtype, device=device)
+    output = (model,)
+
+    # init transforms
+    if return_transforms:
+        # mean and std
+        if 'siglip' in pretrained_name.lower():
+            mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
+        else:
+            mean = [0.48145466, 0.4578275, 0.40821073]
+            std = [0.26862954, 0.26130258, 0.27577711]
+
+        # transforms
+        transforms = T.Compose([
+            T.Resize((model.image_size, model.image_size),
+                     interpolation=T.InterpolationMode.BICUBIC),
+            T.ToTensor(),
+            T.Normalize(mean=mean, std=std)
+        ])
+        output += (transforms,)
+    return output[0] if len(output) == 1 else output
+
+
+def clip_xlm_roberta_vit_h_14(
+        pretrained=False,
+        pretrained_name='open-clip-xlm-roberta-large-vit-huge-14',
+        **kwargs):
+    cfg = dict(
+        embed_dim=1024,
+        image_size=224,
+        patch_size=14,
+        vision_dim=1280,
+        vision_mlp_ratio=4,
+        vision_heads=16,
+        vision_layers=32,
+        vision_pool='token',
+        activation='gelu',
+        vocab_size=250002,
+        max_text_len=514,
+        type_size=1,
+        pad_id=1,
+        text_dim=1024,
+        text_heads=16,
+        text_layers=24,
+        text_post_norm=True,
+        text_dropout=0.1,
+        attn_dropout=0.0,
+        proj_dropout=0.0,
+        embedding_dropout=0.0)
+    cfg.update(**kwargs)
+    return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg)
+
+
+class CLIPModel:
+
+    def __init__(self, dtype, device, checkpoint_path, tokenizer_path):
+        self.dtype = dtype
+        self.device = device
+        self.checkpoint_path = checkpoint_path
+        self.tokenizer_path = tokenizer_path
+
+        # init model
+        self.model, self.transforms = clip_xlm_roberta_vit_h_14(
+            pretrained=False,
+            return_transforms=True,
+            return_tokenizer=False,
+            dtype=dtype,
+            device=device)
+        self.model = self.model.eval().requires_grad_(False)
+        logging.info(f'loading {checkpoint_path}')
+        self.model.load_state_dict(
+            torch.load(checkpoint_path, map_location='cpu'))
+
+        # init tokenizer
+        self.tokenizer = HuggingfaceTokenizer(
+            name=tokenizer_path,
+            seq_len=self.model.max_text_len - 2,
+            clean='whitespace')
+
+    def visual(self, videos):
+        # preprocess
+        size = (self.model.image_size,) * 2
+        videos = torch.cat([
+            F.interpolate(
+                u.transpose(0, 1),
+                size=size,
+                mode='bicubic',
+                align_corners=False) for u in videos
+        ])
+        videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))
+
+        # forward
+        with torch.amp.autocast('cuda', dtype=self.dtype):
+            out = self.model.visual(videos, use_31_block=True)
+            return out

humo/models/wan_modules/model.py

@@ -0,0 +1,619 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+from .attention import flash_attention
+
+__all__ = ['WanModel']
+
+
+def sinusoidal_embedding_1d(dim, position):
+    # preprocess
+    assert dim % 2 == 0
+    half = dim // 2
+    position = position.type(torch.float64)
+
+    # calculation
+    sinusoid = torch.outer(
+        position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x
+
+
+@torch.amp.autocast("cuda", enabled=False)
+def rope_params(max_seq_len, dim, theta=10000):
+    assert dim % 2 == 0
+    freqs = torch.outer(
+        torch.arange(max_seq_len),
+        1.0 / torch.pow(theta,
+                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
+    freqs = torch.polar(torch.ones_like(freqs), freqs)
+    return freqs
+
+
+@torch.amp.autocast("cuda", enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
+            seq_len, n, -1, 2))
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+                            dim=-1).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[i, seq_len:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).float()
+
+
+class WanRMSNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class WanLayerNorm(nn.LayerNorm):
+
+    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class WanSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, seq_lens, grid_sizes, freqs):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            q = self.norm_q(self.q(x)).view(b, s, n, d)
+            k = self.norm_k(self.k(x)).view(b, s, n, d)
+            v = self.v(x).view(b, s, n, d)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+
+        x = flash_attention(
+            q=rope_apply(q, grid_sizes, freqs),
+            k=rope_apply(k, grid_sizes, freqs),
+            v=v,
+            k_lens=seq_lens,
+            window_size=self.window_size)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+
+        # compute attention
+        x = flash_attention(q, k, v, k_lens=context_lens)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+        self.k_img = nn.Linear(dim, dim)
+        self.v_img = nn.Linear(dim, dim)
+        # self.alpha = nn.Parameter(torch.zeros((1, )))
+        self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        context_img = context[:, :257]
+        context = context[:, 257:]
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+        k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)
+        v_img = self.v_img(context_img).view(b, -1, n, d)
+        img_x = flash_attention(q, k_img, v_img, k_lens=None)
+        # compute attention
+        x = flash_attention(q, k, v, k_lens=context_lens)
+
+        # output
+        x = x.flatten(2)
+        img_x = img_x.flatten(2)
+        x = x + img_x
+        x = self.o(x)
+        return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+    't2v_cross_attn': WanT2VCrossAttention,
+    'i2v_cross_attn': WanI2VCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+
+    def __init__(self,
+                 cross_attn_type,
+                 dim,
+                 ffn_dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=False,
+                 eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # layers
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+                                          eps)
+        self.norm3 = WanLayerNorm(
+            dim, eps,
+            elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,
+                                                                      num_heads,
+                                                                      (-1, -1),
+                                                                      qk_norm,
+                                                                      eps)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+            nn.Linear(ffn_dim, dim))
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        assert e.dtype == torch.float32
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            e = (self.modulation + e).chunk(6, dim=1)
+        assert e[0].dtype == torch.float32
+
+        # self-attention
+        y = self.self_attn(
+            self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, grid_sizes,
+            freqs)
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            x = x + y * e[2]
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e):
+            x = x + self.cross_attn(self.norm3(x), context, context_lens)
+            y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3])
+            with torch.amp.autocast('cuda', dtype=torch.float32):
+                x = x + y * e[5]
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+        return x
+
+
+class Head(nn.Module):
+
+    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+
+        # layers
+        out_dim = math.prod(patch_size) * out_dim
+        self.norm = WanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_dim)
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, e):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            e(Tensor): Shape [B, C]
+        """
+        assert e.dtype == torch.float32
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+            x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+        return x
+
+
+class MLPProj(torch.nn.Module):
+
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+
+        self.proj = torch.nn.Sequential(
+            torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
+            torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
+            torch.nn.LayerNorm(out_dim))
+
+    def forward(self, image_embeds):
+        clip_extra_context_tokens = self.proj(image_embeds)
+        return clip_extra_context_tokens
+
+
+class WanModel(ModelMixin, ConfigMixin):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    ignore_for_config = [
+        'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+    ]
+    _no_split_modules = ['WanAttentionBlock']
+
+    @register_to_config
+    def __init__(self,
+                 model_type='t2v',
+                 patch_size=(1, 2, 2),
+                 text_len=512,
+                 in_dim=16,
+                 dim=5120,
+                 ffn_dim=13824,
+                 freq_dim=256,
+                 text_dim=4096,
+                 out_dim=16,
+                 num_heads=40,
+                 num_layers=40,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=True,
+                 eps=1e-6):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        assert model_type in ['t2v', 'i2v']
+        self.model_type = model_type
+
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # embeddings
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim))
+
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks
+        cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+        self.blocks = nn.ModuleList([
+            WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+                              window_size, qk_norm, cross_attn_norm, eps)
+            for _ in range(num_layers)
+        ])
+
+        # head
+        self.head = Head(dim, out_dim, patch_size, eps)
+
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+        self.freqs = torch.cat([
+            rope_params(1024, d - 4 * (d // 6)),
+            rope_params(1024, 2 * (d // 6)),
+            rope_params(1024, 2 * (d // 6))
+        ],
+                               dim=1)
+
+        if model_type == 'i2v':
+            self.img_emb = MLPProj(1280, dim)
+
+        # initialize weights
+        self.init_weights()
+
+    def forward(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        clip_fea=None,
+        y=None,
+    ):
+        r"""
+        Forward pass through the diffusion model
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            clip_fea (Tensor, *optional*):
+                CLIP image features for image-to-video mode
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+        if self.model_type == 'i2v':
+            assert clip_fea is not None and y is not None
+        # params
+        device = self.patch_embedding.weight.device
+        freqs = self.freqs.to(device)
+
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        grid_sizes = torch.stack(
+            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat([
+            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+                      dim=1) for u in x
+        ])
+
+        # time embeddings
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            e = self.time_embedding(
+                sinusoidal_embedding_1d(self.freq_dim, t).float())
+            e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+            assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([
+                torch.cat(
+                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                for u in context
+            ]))
+
+        if clip_fea is not None:
+            context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
+            context = torch.concat([context_clip, context], dim=1)
+
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=freqs,
+            context=context,
+            context_lens=context_lens)
+
+        for block in self.blocks:
+            x = block(x, **kwargs)
+
+        # head
+        x = self.head(x, e)
+
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+        return [u.float() for u in x]
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum('fhwpqrc->cfphqwr', u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out
+
+    def init_weights(self):
+        r"""
+        Initialize model parameters using Xavier initialization.
+        """
+
+        # basic init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # init embeddings
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.text_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+
+        # init output layer
+        nn.init.zeros_(self.head.head.weight)

humo/models/wan_modules/model_humo.py

@@ -0,0 +1,803 @@
+import torch
+from torch import nn
+
+from common.distributed import get_device
+from models.audio.audio_proj import AudioProjModel
+
+import torch.cuda.amp as amp
+import math
+from humo.models.wan_modules.attention import flash_attention
+from common.distributed.advanced import is_unified_parallel_initialized
+
+import types
+
+def sinusoidal_embedding_1d(dim, position):
+    # preprocess
+    assert dim % 2 == 0
+    half = dim // 2
+    position = position.type(torch.float64)
+
+    # calculation
+    sinusoid = torch.outer(
+        position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x
+
+
+@amp.autocast(enabled=False)
+def rope_params(max_seq_len, dim, theta=10000):
+    assert dim % 2 == 0
+    freqs = torch.outer(
+        torch.arange(max_seq_len),
+        1.0 / torch.pow(theta,
+                        torch.arange(0, dim, 2).to(torch.float32).div(dim)))
+    freqs = torch.polar(torch.ones_like(freqs), freqs)
+    return freqs
+
+
+@amp.autocast(enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+    
+    # loop over samples
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float32).reshape(
+            seq_len, n, -1, 2))
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+                            dim=-1).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i = torch.cat([x_i, x[i, seq_len:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).float()
+
+
+class WanRMSNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return self._norm(x.float()).type_as(x) * self.weight
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class WanLayerNorm(nn.LayerNorm):
+
+    def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+        super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+    def forward(self, x):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class WanSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, seq_lens, grid_sizes, freqs):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads], torch.Size([1, 9360, 5120])
+            seq_lens(Tensor): Shape [B], tensor([9360])
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W), tensor([[ 6, 30, 52]])
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            q = self.norm_q(self.q(x)).view(b, s, n, d)
+            k = self.norm_k(self.k(x)).view(b, s, n, d)
+            v = self.v(x).view(b, s, n, d)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+
+        x = flash_attention(
+            q=rope_apply(q, grid_sizes, freqs),
+            k=rope_apply(k, grid_sizes, freqs),
+            v=v,
+            k_lens=seq_lens,
+            window_size=self.window_size)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanSelfAttentionSepKVDim(nn.Module):
+
+    def __init__(self,
+                 kv_dim,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(kv_dim, dim)
+        self.v = nn.Linear(kv_dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, seq_lens, grid_sizes, freqs):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads], torch.Size([1, 9360, 5120])
+            seq_lens(Tensor): Shape [B], tensor([9360])
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W), tensor([[ 6, 30, 52]])
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            q = self.norm_q(self.q(x)).view(b, s, n, d)
+            k = self.norm_k(self.k(x)).view(b, s, n, d)
+            v = self.v(x).view(b, s, n, d)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+
+        x = flash_attention(
+            q=rope_apply(q, grid_sizes, freqs),
+            k=rope_apply(k, grid_sizes, freqs),
+            v=v,
+            k_lens=seq_lens,
+            window_size=self.window_size)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+
+        # compute attention
+        x = flash_attention(q, k, v, k_lens=context_lens)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanT2VCrossAttentionGather(WanSelfAttentionSepKVDim):
+
+    def forward(self, x, context, context_lens, grid_sizes, freqs, audio_seq_len):
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+
+        # --- NEW: derive sizes from shapes (SymInts), no int(tensor) casts ---
+        Lq = q.shape[1]                 # total video tokens per sample
+        # audio has 16 tokens per frame -> frames = audio_tokens // 16
+        frames = (context.shape[1] // 16)
+        hlen_wlen = Lq // frames        # tokens per frame = H*W
+
+        # Now reshape using SymInt-derived sizes
+        q = q.reshape(-1, hlen_wlen, n, d)
+        k = k.reshape(-1, 16, n, d)
+        v = v.reshape(-1, 16, n, d)
+
+        x = flash_attention(q, k, v, k_lens=None)
+        x = x.view(b, -1, n, d).flatten(2)
+        x = self.o(x)
+        return x
+
+    # def forward(self, x, context, context_lens, grid_sizes, freqs, audio_seq_len):
+    #     r"""
+    #     Args:
+    #         x(Tensor): Shape [B, L1, C] - video tokens
+    #         context(Tensor): Shape [B, L2, C] - audio tokens with shape [B, frames*16, 1536]
+    #         context_lens(Tensor): Shape [B] - actually seq_lens from call (video sequence length)
+    #         grid_sizes(Tensor): Shape [B, 3] - video grid dimensions (F, H, W)
+    #         freqs(Tensor): RoPE frequencies
+    #         audio_seq_len(Tensor): Actual audio sequence length (frames * 16)
+    #     """
+    #     b, n, d = x.size(0), self.num_heads, self.head_dim
+
+    #     q = self.norm_q(self.q(x)).view(b, -1, n, d)
+    #     k = self.norm_k(self.k(context)).view(b, -1, n, d)
+    #     v = self.v(context).view(b, -1, n, d)
+
+    #     # Handle video spatial structure
+    #     hlen_wlen = int(grid_sizes[0][1] * grid_sizes[0][2])
+    #     q = q.reshape(-1, hlen_wlen, n, d)
+        
+    #     # Handle audio temporal structure (16 tokens per frame)
+    #     k = k.reshape(-1, 16, n, d)
+    #     v = v.reshape(-1, 16, n, d)
+
+    #     # Cross-attention
+    #     x = flash_attention(q, k, v, k_lens=None)  # No masking for audio
+        
+    #     x = x.view(b, -1, n, d).flatten(2)
+    #     x = self.o(x)
+    #     return x
+
+
+class AudioCrossAttentionWrapper(nn.Module):
+    def __init__(self, dim, kv_dim, num_heads, qk_norm=True, eps=1e-6,):
+        super().__init__()
+
+        self.audio_cross_attn = WanT2VCrossAttentionGather(
+                kv_dim, dim, num_heads, (-1, -1), qk_norm, eps)
+        self.norm1_audio = WanLayerNorm(dim, eps,
+            elementwise_affine=True)
+
+    def forward(self, x, audio, seq_lens, grid_sizes, freqs, audio_seq_len):
+        x = x + self.audio_cross_attn(
+            self.norm1_audio(x), audio, seq_lens, grid_sizes, freqs, audio_seq_len)
+        return x
+        
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 eps=1e-6):
+        super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+    def forward(self, x, context, context_lens):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            context(Tensor): Shape [B, L2, C]
+            context_lens(Tensor): Shape [B]
+        """
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+        x = flash_attention(q, k, v, k_lens=context_lens)
+        
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+    't2v_cross_attn': WanT2VCrossAttention,
+    'i2v_cross_attn': WanI2VCrossAttention,
+}
+
+class WanAttentionBlock(nn.Module):
+
+    def __init__(self,
+                 cross_attn_type,
+                 dim,
+                 ffn_dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=False,
+                 eps=1e-6,
+                 use_audio=True):
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # layers
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+                                          eps)
+        self.norm3 = WanLayerNorm(
+            dim, eps,
+            elementwise_affine=True) if cross_attn_norm else nn.Identity()
+        self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,
+                                                                      num_heads,
+                                                                      (-1, -1),
+                                                                      qk_norm,
+                                                                      eps)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+            nn.Linear(ffn_dim, dim))
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+        self.use_audio = use_audio
+        if use_audio:
+            self.audio_cross_attn_wrapper = AudioCrossAttentionWrapper(dim, 1536, num_heads, qk_norm, eps)
+
+    def forward(
+        self,
+        x, # torch.Size([1, 9360, 5120])
+        e, # torch.Size([1, 6, 5120])
+        seq_lens, # tensor([9360])
+        grid_sizes, # tensor([[ 6, 30, 52]])
+        freqs, # torch.Size([1024, 64])
+        context, # torch.Size([1, 512, 5120])
+        context_lens, # None
+        audio=None, # None
+        audio_seq_len=None,
+        ref_num_list=None,
+    ):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, L, C]
+            audio(Tensor): Shape [B, L, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+            ref_num_list: 配合seq_lens可以查到reference image在倒数第几个
+        """
+        assert e.dtype == torch.float32
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            e = (self.modulation + e).chunk(6, dim=1)
+        assert e[0].dtype == torch.float32
+
+        # self-attention
+        y = self.self_attn(
+            self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, grid_sizes,
+            freqs)
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            x = x + y * e[2]
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e):
+            x = x + self.cross_attn(self.norm3(x), context, context_lens)
+            
+            if self.use_audio:
+                x = self.audio_cross_attn_wrapper(x, audio, seq_lens, grid_sizes, freqs, audio_seq_len)
+
+            y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3])
+            with torch.amp.autocast('cuda', dtype=torch.float32):
+                x = x + y * e[5]
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+
+        return x
+
+
+class Head(nn.Module):
+
+    def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.patch_size = patch_size
+        self.eps = eps
+
+        # layers
+        out_dim = math.prod(patch_size) * out_dim
+        self.norm = WanLayerNorm(dim, eps)
+        self.head = nn.Linear(dim, out_dim)
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, e):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            e(Tensor): Shape [B, C]
+        """
+        assert e.dtype == torch.float32
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+            x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+        return x
+
+
+class MLPProj(torch.nn.Module):
+
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+
+        self.proj = torch.nn.Sequential(
+            torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
+            torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
+            torch.nn.LayerNorm(out_dim))
+
+    def forward(self, image_embeds):
+        clip_extra_context_tokens = self.proj(image_embeds)
+        return clip_extra_context_tokens
+
+
+class WanModel(nn.Module):
+    r"""
+    Wan diffusion backbone supporting both text-to-video and image-to-video.
+    """
+
+    ignore_for_config = [
+        'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+    ]
+    _no_split_modules = ['WanAttentionBlock']
+
+    gradient_checkpointing = False
+
+    def __init__(self,
+                 model_type='t2v',
+                 patch_size=(1, 2, 2),
+                 text_len=512,
+                 in_dim=16,
+                 dim=2048,
+                 ffn_dim=13824,
+                 freq_dim=256,
+                 text_dim=4096,
+                 out_dim=16,
+                 num_heads=40,
+                 num_layers=40,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=True,
+                 eps=1e-6,
+                 audio_token_num=16,
+                 insert_audio=True):
+        r"""
+        Initialize the diffusion model backbone.
+
+        Args:
+            model_type (`str`, *optional*, defaults to 't2v'):
+                Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+            patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+                3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+            text_len (`int`, *optional*, defaults to 512):
+                Fixed length for text embeddings
+            in_dim (`int`, *optional*, defaults to 16):
+                Input video channels (C_in)
+            dim (`int`, *optional*, defaults to 2048):
+                Hidden dimension of the transformer
+            ffn_dim (`int`, *optional*, defaults to 8192):
+                Intermediate dimension in feed-forward network
+            freq_dim (`int`, *optional*, defaults to 256):
+                Dimension for sinusoidal time embeddings
+            text_dim (`int`, *optional*, defaults to 4096):
+                Input dimension for text embeddings
+            out_dim (`int`, *optional*, defaults to 16):
+                Output video channels (C_out)
+            num_heads (`int`, *optional*, defaults to 16):
+                Number of attention heads
+            num_layers (`int`, *optional*, defaults to 32):
+                Number of transformer blocks
+            window_size (`tuple`, *optional*, defaults to (-1, -1)):
+                Window size for local attention (-1 indicates global attention)
+            qk_norm (`bool`, *optional*, defaults to True):
+                Enable query/key normalization
+            cross_attn_norm (`bool`, *optional*, defaults to False):
+                Enable cross-attention normalization
+            eps (`float`, *optional*, defaults to 1e-6):
+                Epsilon value for normalization layers
+        """
+
+        super().__init__()
+
+        assert model_type in ['t2v', 'i2v']
+        self.model_type = model_type
+
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # embeddings
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim))
+
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks
+        cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+        self.insert_audio = insert_audio
+        self.blocks = nn.ModuleList([
+            WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+                        window_size, qk_norm, cross_attn_norm,
+                        eps, use_audio=self.insert_audio)
+            for _ in range(num_layers)
+        ])
+
+        # head
+        self.head = Head(dim, out_dim, patch_size, eps)
+
+        if self.insert_audio:
+            self.audio_proj = AudioProjModel(seq_len=8, blocks=5, channels=1280, 
+                intermediate_dim=512, output_dim=1536, context_tokens=audio_token_num)
+
+        # RoPE freqs: register as a buffer so it moves with .to() / DDP and is tracked by compile
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+
+        _freqs = torch.cat([
+            rope_params(1024, d - 4 * (d // 6)),
+            rope_params(1024, 2 * (d // 6)),
+            rope_params(1024, 2 * (d // 6))
+        ], dim=1)
+        self.register_buffer("freqs", _freqs, persistent=False)
+
+        # initialize weights
+        self.init_weights()
+
+        # initialize unified parallel
+        if is_unified_parallel_initialized():
+            print(f"Initializing WanModel with unified parallel initialized")
+            from humo.models.distributed.dit_ulysses_sequence_parallel import ulysses_attn_forward, ulysses_dit_forward, ulysses_audio_cross_attn_forward
+            for block in self.blocks:
+                block.self_attn.forward = types.MethodType(ulysses_attn_forward, block.self_attn)
+                if block.use_audio:
+                    block.audio_cross_attn_wrapper.audio_cross_attn.forward = types.MethodType(ulysses_audio_cross_attn_forward, block.audio_cross_attn_wrapper.audio_cross_attn)
+            self.forward = types.MethodType(ulysses_dit_forward, self)
+        
+    def forward(
+        self,
+        x,
+        t,
+        context,
+        seq_len,
+        audio=None,
+        y=None,
+        tea_cache=None,
+    ):
+        r"""
+        Forward pass through the diffusion model
+
+        Args:
+            x (List[Tensor]):
+                List of input video tensors, each with shape [C_in, F, H, W]
+            t (Tensor):
+                Diffusion timesteps tensor of shape [B]
+            context (List[Tensor]):
+                List of text embeddings each with shape [L, C]
+            seq_len (`int`):
+                Maximum sequence length for positional encoding
+            clip_fea (Tensor, *optional*):
+                CLIP image features for image-to-video mode
+            y (List[Tensor], *optional*):
+                Conditional video inputs for image-to-video mode, same shape as x
+
+        Returns:
+            List[Tensor]:
+                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+        """
+        if self.model_type == 'i2v':
+            assert y is not None
+
+        # params
+        freqs = self.freqs
+
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len
+
+        # pad to uniform length and batch
+        x = torch.cat([
+            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
+            for u in x
+        ])  # shape: [B, seq_len, C]
+
+        # time embeddings
+        with torch.amp.autocast('cuda', dtype=torch.float32):
+            e = self.time_embedding(
+                sinusoidal_embedding_1d(self.freq_dim, t).float()
+            ).float()
+            e0 = self.time_projection(e).unflatten(1, (6, self.dim)).float()
+            assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([
+                torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                for u in context
+            ])
+        )
+
+        # audio (unchanged; not cached)
+        if self.insert_audio:
+            audio = [self.audio_proj(au.unsqueeze(0)).permute(0, 3, 1, 2) for au in audio]
+            audio_seq_len = max(au.shape[2] for au in audio) * audio[0].shape[3]
+
+            audio = [au.flatten(2).transpose(1, 2) for au in audio]  # [1, t*32, 1536]
+            audio = torch.cat([
+                torch.cat([au, au.new_zeros(1, int(audio_seq_len) - au.size(1), au.size(2))], dim=1)
+                for au in audio
+            ])
+        else:
+            audio = None
+            audio_seq_len = None
+
+        # ---- tea_cache integration (mirrors your working model) ----
+        if tea_cache is not None:
+            # Use the pre-block tokens 'x' and time-mod 'e0' to decide whether to reuse cache
+            tea_cache_update = tea_cache.check(self, x, e0)
+        else:
+            tea_cache_update = False
+
+        ori_x_len = x.shape[1]  # remember original token length before potential cache extension
+
+        if tea_cache_update:
+            # Let the cache inject/append any needed past states/tokens for reuse
+            x = tea_cache.update(x)
+        else:
+            # arguments for blocks
+            kwargs = dict(
+                e=e0,
+                seq_lens=seq_lens,
+                grid_sizes=grid_sizes,
+                freqs=freqs,
+                context=context,
+                context_lens=context_lens,
+                audio=audio,
+                audio_seq_len=audio_seq_len
+            )
+
+            # transformer blocks
+            for block in self.blocks:
+                x = block(x, **kwargs)
+
+            if tea_cache is not None:
+                x_cache = x[:, :ori_x_len]
+                tea_cache.store(x_cache)
+
+        # head
+        x = self.head(x, e)
+
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+        return [u.float() for u in x]
+
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum('fhwpqrc->cfphqwr', u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out
+
+    def init_weights(self):
+        r"""
+        Initialize model parameters using Xavier initialization.
+        """
+
+        # basic init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # init embeddings
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.text_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+
+        # init output layer
+        nn.init.zeros_(self.head.head.weight)

humo/models/wan_modules/t5.py

@@ -0,0 +1,525 @@
+# Modified from transformers.models.t5.modeling_t5
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .tokenizers import HuggingfaceTokenizer
+
+__all__ = [
+    'T5Model',
+    'T5Encoder',
+    'T5Decoder',
+    'T5EncoderModel',
+]
+
+
+def fp16_clamp(x):
+    if x.dtype == torch.float16 and torch.isinf(x).any():
+        clamp = torch.finfo(x.dtype).max - 1000
+        x = torch.clamp(x, min=-clamp, max=clamp)
+    return x
+
+
+def init_weights(m):
+    if isinstance(m, T5LayerNorm):
+        nn.init.ones_(m.weight)
+    elif isinstance(m, T5Model):
+        nn.init.normal_(m.token_embedding.weight, std=1.0)
+    elif isinstance(m, T5FeedForward):
+        nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+    elif isinstance(m, T5Attention):
+        nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)
+        nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)
+    elif isinstance(m, T5RelativeEmbedding):
+        nn.init.normal_(
+            m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)
+
+
+class GELU(nn.Module):
+
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.tanh(
+            math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-6):
+        super(T5LayerNorm, self).__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +
+                            self.eps)
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            x = x.type_as(self.weight)
+        return self.weight * x
+
+
+class T5Attention(nn.Module):
+
+    def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+        assert dim_attn % num_heads == 0
+        super(T5Attention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.num_heads = num_heads
+        self.head_dim = dim_attn // num_heads
+
+        # layers
+        self.q = nn.Linear(dim, dim_attn, bias=False)
+        self.k = nn.Linear(dim, dim_attn, bias=False)
+        self.v = nn.Linear(dim, dim_attn, bias=False)
+        self.o = nn.Linear(dim_attn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context=None, mask=None, pos_bias=None):
+        """
+        x:          [B, L1, C].
+        context:    [B, L2, C] or None.
+        mask:       [B, L2] or [B, L1, L2] or None.
+        """
+        # check inputs
+        context = x if context is None else context
+        b, n, c = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.q(x).view(b, -1, n, c)
+        k = self.k(context).view(b, -1, n, c)
+        v = self.v(context).view(b, -1, n, c)
+
+        # attention bias
+        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+        if pos_bias is not None:
+            attn_bias += pos_bias
+        if mask is not None:
+            assert mask.ndim in [2, 3]
+            mask = mask.view(b, 1, 1,
+                             -1) if mask.ndim == 2 else mask.unsqueeze(1)
+            attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+        # compute attention (T5 does not use scaling)
+        attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum('bnij,bjnc->binc', attn, v)
+
+        # output
+        x = x.reshape(b, -1, n * c)
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5FeedForward(nn.Module):
+
+    def __init__(self, dim, dim_ffn, dropout=0.1):
+        super(T5FeedForward, self).__init__()
+        self.dim = dim
+        self.dim_ffn = dim_ffn
+
+        # layers
+        self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+        self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+        self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x) * self.gate(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5SelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5SelfAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True)
+
+    def forward(self, x, mask=None, pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.ffn(self.norm2(x)))
+        return x
+
+
+class T5CrossAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5CrossAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm3 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=False)
+
+    def forward(self,
+                x,
+                mask=None,
+                encoder_states=None,
+                encoder_mask=None,
+                pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.cross_attn(
+            self.norm2(x), context=encoder_states, mask=encoder_mask))
+        x = fp16_clamp(x + self.ffn(self.norm3(x)))
+        return x
+
+
+class T5RelativeEmbedding(nn.Module):
+
+    def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+        super(T5RelativeEmbedding, self).__init__()
+        self.num_buckets = num_buckets
+        self.num_heads = num_heads
+        self.bidirectional = bidirectional
+        self.max_dist = max_dist
+
+        # layers
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, lq, lk):
+        device = self.embedding.weight.device
+        # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \
+        #     torch.arange(lq).unsqueeze(1).to(device)
+        rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \
+            torch.arange(lq, device=device).unsqueeze(1)
+        rel_pos = self._relative_position_bucket(rel_pos)
+        rel_pos_embeds = self.embedding(rel_pos)
+        rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(
+            0)  # [1, N, Lq, Lk]
+        return rel_pos_embeds.contiguous()
+
+    def _relative_position_bucket(self, rel_pos):
+        # preprocess
+        if self.bidirectional:
+            num_buckets = self.num_buckets // 2
+            rel_buckets = (rel_pos > 0).long() * num_buckets
+            rel_pos = torch.abs(rel_pos)
+        else:
+            num_buckets = self.num_buckets
+            rel_buckets = 0
+            rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+        # embeddings for small and large positions
+        max_exact = num_buckets // 2
+        rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /
+                                     math.log(self.max_dist / max_exact) *
+                                     (num_buckets - max_exact)).long()
+        rel_pos_large = torch.min(
+            rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+        rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+        return rel_buckets
+
+
+class T5Encoder(nn.Module):
+
+    def __init__(self,
+                 vocab,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Encoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+            else nn.Embedding(vocab, dim)
+        self.pos_embedding = T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True) if shared_pos else None
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([
+            T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+                            shared_pos, dropout) for _ in range(num_layers)
+        ])
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None):
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1),
+                               x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5Decoder(nn.Module):
+
+    def __init__(self,
+                 vocab,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Decoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+            else nn.Embedding(vocab, dim)
+        self.pos_embedding = T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=False) if shared_pos else None
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([
+            T5CrossAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+                             shared_pos, dropout) for _ in range(num_layers)
+        ])
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None):
+        b, s = ids.size()
+
+        # causal mask
+        if mask is None:
+            mask = torch.tril(torch.ones(1, s, s).to(ids.device))
+        elif mask.ndim == 2:
+            mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1))
+
+        # layers
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1),
+                               x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, encoder_states, encoder_mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5Model(nn.Module):
+
+    def __init__(self,
+                 vocab_size,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 encoder_layers,
+                 decoder_layers,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5Model, self).__init__()
+        self.vocab_size = vocab_size
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.encoder_layers = encoder_layers
+        self.decoder_layers = decoder_layers
+        self.num_buckets = num_buckets
+
+        # layers
+        self.token_embedding = nn.Embedding(vocab_size, dim)
+        self.encoder = T5Encoder(self.token_embedding, dim, dim_attn, dim_ffn,
+                                 num_heads, encoder_layers, num_buckets,
+                                 shared_pos, dropout)
+        self.decoder = T5Decoder(self.token_embedding, dim, dim_attn, dim_ffn,
+                                 num_heads, decoder_layers, num_buckets,
+                                 shared_pos, dropout)
+        self.head = nn.Linear(dim, vocab_size, bias=False)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask):
+        x = self.encoder(encoder_ids, encoder_mask)
+        x = self.decoder(decoder_ids, decoder_mask, x, encoder_mask)
+        x = self.head(x)
+        return x
+
+
+def _t5(name,
+        encoder_only=False,
+        decoder_only=False,
+        return_tokenizer=False,
+        tokenizer_kwargs={},
+        dtype=torch.float32,
+        device='cpu',
+        **kwargs):
+    # sanity check
+    assert not (encoder_only and decoder_only)
+
+    # params
+    if encoder_only:
+        model_cls = T5Encoder
+        kwargs['vocab'] = kwargs.pop('vocab_size')
+        kwargs['num_layers'] = kwargs.pop('encoder_layers')
+        _ = kwargs.pop('decoder_layers')
+    elif decoder_only:
+        model_cls = T5Decoder
+        kwargs['vocab'] = kwargs.pop('vocab_size')
+        kwargs['num_layers'] = kwargs.pop('decoder_layers')
+        _ = kwargs.pop('encoder_layers')
+    else:
+        model_cls = T5Model
+
+    # init model
+    with torch.device(device):
+        model = model_cls(**kwargs)
+
+    # set device
+    model = model.to(dtype=dtype, device=device)
+
+    # init tokenizer
+    if return_tokenizer:
+        from .tokenizers import HuggingfaceTokenizer
+        tokenizer = HuggingfaceTokenizer(f'google/{name}', **tokenizer_kwargs)
+        return model, tokenizer
+    else:
+        return model
+
+
+def umt5_xxl(**kwargs):
+    cfg = dict(
+        vocab_size=256384,
+        dim=4096,
+        dim_attn=4096,
+        dim_ffn=10240,
+        num_heads=64,
+        encoder_layers=24,
+        decoder_layers=24,
+        num_buckets=32,
+        shared_pos=False,
+        dropout=0.1)
+    cfg.update(**kwargs)
+    return _t5('umt5-xxl', **cfg)
+
+
+class T5EncoderModel(nn.Module):
+
+    def __init__(
+        self,
+        text_len,
+        dtype=torch.bfloat16,
+        device=torch.cuda.current_device(),
+        checkpoint_path=None,
+        tokenizer_path=None,
+        shard_fn=None,
+    ):
+        super(T5EncoderModel, self).__init__()
+        self.text_len = text_len
+        self.dtype = dtype
+        self.device = device
+        self.checkpoint_path = checkpoint_path
+        self.tokenizer_path = tokenizer_path
+
+        with torch.device(device):
+            self.model = T5Encoder(
+                vocab=256384,
+                dim=4096,
+                dim_attn=4096,
+                dim_ffn=10240,
+                num_heads=64,
+                num_layers=24,
+                num_buckets=32,
+                shared_pos=False,
+                dropout=0.1
+            )
+        # set device
+        self.model = self.model.to(dtype=dtype, device=device).eval().requires_grad_(False)
+
+        logging.info(f'loading {checkpoint_path}')
+        if checkpoint_path is not None:
+            self.model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
+        
+        if shard_fn is not None:
+            self.model = shard_fn(self.model, sync_module_states=False)
+        else:
+            self.model.to(self.device)
+        # init tokenizer
+        self.tokenizer = HuggingfaceTokenizer(
+            name=tokenizer_path, seq_len=text_len, clean='whitespace')
+
+    @torch.no_grad()
+    def __call__(self, texts, device):
+        ids, mask = self.tokenizer(
+            texts, return_mask=True, add_special_tokens=True)
+        ids = ids.to(device)
+        mask = mask.to(device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        context = self.model(ids, mask)
+        return [u[:v] for u, v in zip(context, seq_lens)]

humo/models/wan_modules/tokenizers.py

@@ -0,0 +1,82 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import html
+import string
+
+import ftfy
+import regex as re
+from transformers import AutoTokenizer
+
+__all__ = ['HuggingfaceTokenizer']
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+def canonicalize(text, keep_punctuation_exact_string=None):
+    text = text.replace('_', ' ')
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans('', '', string.punctuation))
+            for part in text.split(keep_punctuation_exact_string))
+    else:
+        text = text.translate(str.maketrans('', '', string.punctuation))
+    text = text.lower()
+    text = re.sub(r'\s+', ' ', text)
+    return text.strip()
+
+
+class HuggingfaceTokenizer:
+
+    def __init__(self, name, seq_len=None, clean=None, **kwargs):
+        assert clean in (None, 'whitespace', 'lower', 'canonicalize')
+        self.name = name
+        self.seq_len = seq_len
+        self.clean = clean
+
+        # init tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)
+        self.vocab_size = self.tokenizer.vocab_size
+
+    def __call__(self, sequence, **kwargs):
+        return_mask = kwargs.pop('return_mask', False)
+
+        # arguments
+        _kwargs = {'return_tensors': 'pt'}
+        if self.seq_len is not None:
+            _kwargs.update({
+                'padding': 'max_length',
+                'truncation': True,
+                'max_length': self.seq_len
+            })
+        _kwargs.update(**kwargs)
+
+        # tokenization
+        if isinstance(sequence, str):
+            sequence = [sequence]
+        if self.clean:
+            sequence = [self._clean(u) for u in sequence]
+        ids = self.tokenizer(sequence, **_kwargs)
+
+        # output
+        if return_mask:
+            return ids.input_ids, ids.attention_mask
+        else:
+            return ids.input_ids
+
+    def _clean(self, text):
+        if self.clean == 'whitespace':
+            text = whitespace_clean(basic_clean(text))
+        elif self.clean == 'lower':
+            text = whitespace_clean(basic_clean(text)).lower()
+        elif self.clean == 'canonicalize':
+            text = canonicalize(basic_clean(text))
+        return text

humo/models/wan_modules/vae.py

@@ -0,0 +1,666 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__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.
+
+    def forward(self, x):
+        return F.normalize(
+            x, dim=(1 if self.channel_first else
+                    -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',
+                        'downsample3d')
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == 'upsample2d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+        elif mode == 'upsample3d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), 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]
+
+        # 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 encode(self, x, scale):
+        self.clear_cache()
+        ## cache
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        ## 对encode输入的x，按时间拆分为1、4、4、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()
+        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 reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        #cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **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
+    logging.info(f'loading {pretrained_path}')
+    if pretrained_path is not None:
+        model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True)
+
+    return model
+
+
+class WanVAE:
+
+    def __init__(self,
+                 z_dim=16,
+                 vae_pth=None,
+                 dtype=torch.float,
+                 device="cuda"):
+        self.dtype = dtype
+        self.device = device
+
+        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.std = torch.tensor(std, dtype=dtype, device=device)
+        self.scale = [self.mean, 1.0 / self.std]
+
+        # init model
+        self.model = _video_vae(
+            pretrained_path=vae_pth,
+            z_dim=z_dim,
+        ).eval().requires_grad_(False).to(device)
+
+    @torch.no_grad()
+    def encode(self, videos, device):
+        """
+        videos: A list of videos each with shape [C, T, H, W].
+        """
+
+        with torch.amp.autocast('cuda', dtype=self.dtype):
+            return [
+                self.model.encode(u.unsqueeze(0).to(device,self.dtype), self.scale).float().squeeze(0)
+                for u in videos
+            ]
+
+    @torch.no_grad()
+    def decode(self, zs):
+        with torch.amp.autocast('cuda', dtype=self.dtype):
+            return [
+                self.model.decode(u.unsqueeze(0),
+                                  self.scale).float().clamp_(-1, 1).squeeze(0)
+                for u in zs
+            ]

humo/models/wan_modules/xlm_roberta.py

@@ -0,0 +1,170 @@
+# Modified from transformers.models.xlm_roberta.modeling_xlm_roberta
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['XLMRoberta', 'xlm_roberta_large']
+
+
+class SelfAttention(nn.Module):
+
+    def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.eps = eps
+
+        # layers
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, mask):
+        """
+        x:   [B, L, C].
+        """
+        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.q(x).reshape(b, s, n, d).permute(0, 2, 1, 3)
+        k = self.k(x).reshape(b, s, n, d).permute(0, 2, 1, 3)
+        v = self.v(x).reshape(b, s, n, d).permute(0, 2, 1, 3)
+
+        # compute attention
+        p = self.dropout.p if self.training else 0.0
+        x = F.scaled_dot_product_attention(q, k, v, mask, p)
+        x = x.permute(0, 2, 1, 3).reshape(b, s, c)
+
+        # output
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+
+class AttentionBlock(nn.Module):
+
+    def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.post_norm = post_norm
+        self.eps = eps
+
+        # layers
+        self.attn = SelfAttention(dim, num_heads, dropout, eps)
+        self.norm1 = nn.LayerNorm(dim, eps=eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim),
+            nn.Dropout(dropout))
+        self.norm2 = nn.LayerNorm(dim, eps=eps)
+
+    def forward(self, x, mask):
+        if self.post_norm:
+            x = self.norm1(x + self.attn(x, mask))
+            x = self.norm2(x + self.ffn(x))
+        else:
+            x = x + self.attn(self.norm1(x), mask)
+            x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class XLMRoberta(nn.Module):
+    """
+    XLMRobertaModel with no pooler and no LM head.
+    """
+
+    def __init__(self,
+                 vocab_size=250002,
+                 max_seq_len=514,
+                 type_size=1,
+                 pad_id=1,
+                 dim=1024,
+                 num_heads=16,
+                 num_layers=24,
+                 post_norm=True,
+                 dropout=0.1,
+                 eps=1e-5):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.type_size = type_size
+        self.pad_id = pad_id
+        self.dim = dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.post_norm = post_norm
+        self.eps = eps
+
+        # embeddings
+        self.token_embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id)
+        self.type_embedding = nn.Embedding(type_size, dim)
+        self.pos_embedding = nn.Embedding(max_seq_len, dim, padding_idx=pad_id)
+        self.dropout = nn.Dropout(dropout)
+
+        # blocks
+        self.blocks = nn.ModuleList([
+            AttentionBlock(dim, num_heads, post_norm, dropout, eps)
+            for _ in range(num_layers)
+        ])
+
+        # norm layer
+        self.norm = nn.LayerNorm(dim, eps=eps)
+
+    def forward(self, ids):
+        """
+        ids: [B, L] of torch.LongTensor.
+        """
+        b, s = ids.shape
+        mask = ids.ne(self.pad_id).long()
+
+        # embeddings
+        x = self.token_embedding(ids) + \
+            self.type_embedding(torch.zeros_like(ids)) + \
+            self.pos_embedding(self.pad_id + torch.cumsum(mask, dim=1) * mask)
+        if self.post_norm:
+            x = self.norm(x)
+        x = self.dropout(x)
+
+        # blocks
+        mask = torch.where(
+            mask.view(b, 1, 1, s).gt(0), 0.0,
+            torch.finfo(x.dtype).min)
+        for block in self.blocks:
+            x = block(x, mask)
+
+        # output
+        if not self.post_norm:
+            x = self.norm(x)
+        return x
+
+
+def xlm_roberta_large(pretrained=False,
+                      return_tokenizer=False,
+                      device='cpu',
+                      **kwargs):
+    """
+    XLMRobertaLarge adapted from Huggingface.
+    """
+    # params
+    cfg = dict(
+        vocab_size=250002,
+        max_seq_len=514,
+        type_size=1,
+        pad_id=1,
+        dim=1024,
+        num_heads=16,
+        num_layers=24,
+        post_norm=True,
+        dropout=0.1,
+        eps=1e-5)
+    cfg.update(**kwargs)
+
+    # init a model on device
+    with torch.device(device):
+        model = XLMRoberta(**cfg)
+    return model

humo/utils/audio_processor_whisper.py

@@ -0,0 +1,173 @@
+# pylint: disable=C0301
+'''
+This module contains the AudioProcessor class and related functions for processing audio data.
+It utilizes various libraries and models to perform tasks such as preprocessing, feature extraction,
+and audio separation. The class is initialized with configuration parameters and can process
+audio files using the provided models.
+'''
+import os
+import subprocess
+
+import librosa
+import numpy as np
+import torch
+from audio_separator.separator import Separator
+from transformers import WhisperModel, AutoFeatureExtractor
+import torch.nn.functional as F
+
+
+def linear_interpolation_fps(features, input_fps, output_fps, output_len=None):
+    features = features.transpose(1, 2)  # [1, C, T]
+    seq_len = features.shape[2] / float(input_fps)
+    if output_len is None:
+        output_len = int(seq_len * output_fps)
+    output_features = F.interpolate(features, size=output_len, align_corners=True, mode='linear')
+    return output_features.transpose(1, 2)
+
+
+def resample_audio(input_audio_file: str, output_audio_file: str, sample_rate: int):
+    p = subprocess.Popen([
+        "ffmpeg", "-y", "-v", "error", "-i", input_audio_file, "-ar", str(sample_rate), output_audio_file
+    ])
+    ret = p.wait()
+    assert ret == 0, "Resample audio failed!"
+    return output_audio_file
+
+class AudioProcessor:
+    """
+    AudioProcessor is a class that handles the processing of audio files.
+    It takes care of preprocessing the audio files, extracting features
+    using wav2vec models, and separating audio signals if needed.
+
+    :param sample_rate: Sampling rate of the audio file
+    :param fps: Frames per second for the extracted features
+    :param wav2vec_model_path: Path to the wav2vec model
+    :param only_last_features: Whether to only use the last features
+    :param audio_separator_model_path: Path to the audio separator model
+    :param audio_separator_model_name: Name of the audio separator model
+    :param cache_dir: Directory to cache the intermediate results
+    :param device: Device to run the processing on
+    """
+    def __init__(
+        self,
+        sample_rate,
+        fps,
+        wav2vec_model_path,
+        wav2vec_feature_type,
+        audio_separator_model_path:str=None,
+        audio_separator_model_name:str=None,
+        cache_dir:str='',
+        device="cuda:0",
+    ) -> None:
+        self.sample_rate = sample_rate
+        self.fps = fps
+        self.device = device
+
+        self.whisper = WhisperModel.from_pretrained(wav2vec_model_path).to(device).eval()
+        self.whisper.requires_grad_(False)
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(wav2vec_model_path)
+
+        if audio_separator_model_name is not None:
+            try:
+                os.makedirs(cache_dir, exist_ok=True)
+            except OSError as _:
+                print("Fail to create the output cache dir.")
+            self.audio_separator = Separator(
+                output_dir=cache_dir,
+                output_single_stem="vocals",
+                model_file_dir=audio_separator_model_path,
+            )
+            self.audio_separator.load_model(audio_separator_model_name)
+            assert self.audio_separator.model_instance is not None, "Fail to load audio separate model."
+        else:
+            self.audio_separator=None
+            print("Use audio directly without vocals seperator.")        
+
+
+    def get_audio_feature(self, audio_path):
+        audio_input, sampling_rate = librosa.load(audio_path, sr=16000)
+        assert sampling_rate == 16000
+
+        audio_features = []
+        window = 750*640
+        for i in range(0, len(audio_input), window):
+            audio_feature = self.feature_extractor(audio_input[i:i+window], 
+                                            sampling_rate=sampling_rate, 
+                                            return_tensors="pt", 
+                                            ).input_features
+            audio_features.append(audio_feature)
+        audio_features = torch.cat(audio_features, dim=-1)
+        return audio_features, len(audio_input) // 640
+
+
+    def preprocess(self, audio_path: str):
+        audio_input, audio_len = self.get_audio_feature(audio_path)
+        audio_feature = audio_input.to(self.whisper.device).float()
+        window = 3000
+        audio_prompts = []
+        for i in range(0, audio_feature.shape[-1], window):
+            audio_prompt = self.whisper.encoder(audio_feature[:,:,i:i+window], output_hidden_states=True).hidden_states
+            audio_prompt = torch.stack(audio_prompt, dim=2)
+            audio_prompts.append(audio_prompt)
+
+        audio_prompts = torch.cat(audio_prompts, dim=1)
+        audio_prompts = audio_prompts[:,:audio_len*2]
+
+        audio_emb = self.audio_emb_enc(audio_prompts, wav_enc_type="whisper")
+
+        return audio_emb, audio_emb.shape[0]
+    
+    def audio_emb_enc(self, audio_emb, wav_enc_type="whisper"):
+        if wav_enc_type == "wav2vec":
+            feat_merge = audio_emb
+        elif wav_enc_type == "whisper":
+            # [1, T, 33, 1280]
+            feat0 = linear_interpolation_fps(audio_emb[:, :, 0: 8].mean(dim=2), 50, 25)
+            feat1 = linear_interpolation_fps(audio_emb[:, :, 8: 16].mean(dim=2), 50, 25)
+            feat2 = linear_interpolation_fps(audio_emb[:, :, 16: 24].mean(dim=2), 50, 25)
+            feat3 = linear_interpolation_fps(audio_emb[:, :, 24: 32].mean(dim=2), 50, 25)
+            feat4 = linear_interpolation_fps(audio_emb[:, :, 32], 50, 25)
+            feat_merge = torch.stack([feat0, feat1, feat2, feat3, feat4], dim=2)[0]  # [T, 5, 1280]
+        else:
+            raise ValueError(f"Unsupported wav_enc_type: {wav_enc_type}")
+        
+        return feat_merge
+    
+    def get_audio_emb_window(self, audio_emb, frame_num, frame0_idx, audio_shift=2):
+        zero_audio_embed = torch.zeros((audio_emb.shape[1], audio_emb.shape[2]), dtype=audio_emb.dtype, device=audio_emb.device)
+        zero_audio_embed_3 = torch.zeros((3, audio_emb.shape[1], audio_emb.shape[2]), dtype=audio_emb.dtype, device=audio_emb.device)  # device=audio_emb.device
+        iter_ = 1 + (frame_num - 1) // 4
+        audio_emb_wind = []
+        for lt_i in range(iter_):
+            if lt_i == 0:  # latent_i
+                # 提取第一帧VAElatent，audio左侧补0，标识出
+                st = frame0_idx + lt_i - 2
+                ed = frame0_idx + lt_i + 3
+                wind_feat = torch.stack([
+                    audio_emb[i] if (0 <= i < audio_emb.shape[0]) else zero_audio_embed
+                    for i in range(st, ed)
+                ], dim=0)  # [5, 13, 768]
+                wind_feat = torch.cat((zero_audio_embed_3, wind_feat), dim=0)  # [8, 13, 768]
+            else:
+                st = frame0_idx + 1 + 4 * (lt_i - 1) - audio_shift
+                ed = frame0_idx + 1 + 4 * lt_i + audio_shift
+                wind_feat = torch.stack([
+                    audio_emb[i] if (0 <= i < audio_emb.shape[0]) else zero_audio_embed
+                    for i in range(st, ed)
+                ], dim=0)  # [8, 13, 768]
+            audio_emb_wind.append(wind_feat)
+        audio_emb_wind = torch.stack(audio_emb_wind, dim=0)  # [iter_, 8, 13, 768]
+
+        return audio_emb_wind, ed - audio_shift
+
+    def close(self):
+        """
+        TODO: to be implemented
+        """
+        return self
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, _exc_type, _exc_val, _exc_tb):
+        self.close()

humo/utils/wav2vec.py

@@ -0,0 +1,218 @@
+# pylint: disable=R0901
+# src/models/wav2vec.py
+
+"""
+This module defines the Wav2Vec model, which is a pre-trained model for speech recognition and understanding.
+It inherits from the Wav2Vec2Model class in the transformers library and provides additional functionalities
+such as feature extraction and encoding.
+
+Classes:
+    Wav2VecModel: Inherits from Wav2Vec2Model and adds additional methods for feature extraction and encoding.
+
+Functions:
+    linear_interpolation: Interpolates the features based on the sequence length.
+"""
+
+import torch.nn.functional as F
+from transformers import Wav2Vec2Model
+from transformers.modeling_outputs import BaseModelOutput
+
+
+class Wav2VecModel(Wav2Vec2Model):
+    """
+    Wav2VecModel is a custom model class that extends the Wav2Vec2Model class from the transformers library. 
+    It inherits all the functionality of the Wav2Vec2Model and adds additional methods for feature extraction and encoding.
+    ...
+
+    Attributes:
+        base_model (Wav2Vec2Model): The base Wav2Vec2Model object.
+
+    Methods:
+        forward(input_values, seq_len, attention_mask=None, mask_time_indices=None
+        , output_attentions=None, output_hidden_states=None, return_dict=None):
+            Forward pass of the Wav2VecModel. 
+            It takes input_values, seq_len, and other optional parameters as input and returns the output of the base model.
+
+        feature_extract(input_values, seq_len):
+            Extracts features from the input_values using the base model.
+
+        encode(extract_features, attention_mask=None, mask_time_indices=None, output_attentions=None, output_hidden_states=None, return_dict=None):
+            Encodes the extracted features using the base model and returns the encoded features.
+    """
+    def forward(
+        self,
+        input_values,
+        seq_len,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        """
+        Forward pass of the Wav2Vec model.
+
+        Args:
+            self: The instance of the model.
+            input_values: The input values (waveform) to the model.
+            seq_len: The sequence length of the input values.
+            attention_mask: Attention mask to be used for the model.
+            mask_time_indices: Mask indices to be used for the model.
+            output_attentions: If set to True, returns attentions.
+            output_hidden_states: If set to True, returns hidden states.
+            return_dict: If set to True, returns a BaseModelOutput instead of a tuple.
+
+        Returns:
+            The output of the Wav2Vec model.
+        """
+        self.config.output_attentions = True
+
+        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,
+    ):
+        """
+        Extracts features from the input values and returns the extracted features.
+
+        Parameters:
+        input_values (torch.Tensor): The input values to be processed.
+        seq_len (torch.Tensor): The sequence lengths of the input values.
+
+        Returns:
+        extracted_features (torch.Tensor): The extracted features from the input values.
+        """
+        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,
+    ):
+        """
+        Encodes the input features into the output space.
+
+        Args:
+            extract_features (torch.Tensor): The extracted features from the audio signal.
+            attention_mask (torch.Tensor, optional): Attention mask to be used for padding.
+            mask_time_indices (torch.Tensor, optional): Masked indices for the time dimension.
+            output_attentions (bool, optional): If set to True, returns the attention weights.
+            output_hidden_states (bool, optional): If set to True, returns all hidden states.
+            return_dict (bool, optional): If set to True, returns a BaseModelOutput instead of the tuple.
+
+        Returns:
+            The encoded output features.
+        """
+        self.config.output_attentions = True
+
+        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,
+        )
+
+
+def linear_interpolation(features, seq_len):
+    """
+    Transpose the features to interpolate linearly.
+
+    Args:
+        features (torch.Tensor): The extracted features to be interpolated.
+        seq_len (torch.Tensor): The sequence lengths of the features.
+
+    Returns:
+        torch.Tensor: The interpolated features.
+    """
+    features = features.transpose(1, 2)
+    output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear')
+    return output_features.transpose(1, 2)
+
+
+def linear_interpolation_fps(features, input_fps, output_fps, output_len=None):
+    features = features.transpose(1, 2)  # [1, C, T]
+    seq_len = features.shape[2] / float(input_fps)
+    if output_len is None:
+        output_len = int(seq_len * output_fps)
+    output_features = F.interpolate(features, size=output_len, align_corners=True, mode='linear')
+    return output_features.transpose(1, 2)

main.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# 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.
+
+# Inference codes adapted from [SeedVR]
+# https://github.com/ByteDance-Seed/SeedVR/blob/main/projects/inference_seedvr2_7b.py
+
+from sys import argv
+import sys
+
+path_to_insert = "humo"
+if path_to_insert not in sys.path:
+    sys.path.insert(0, path_to_insert)
+
+from common.config import load_config, create_object
+
+# Load config.
+config = load_config(argv[1], argv[2:])
+
+runner = create_object(config)
+runner.entrypoint()