add code

.gitignore

@@ -0,0 +1,10 @@
+__pycache__
+*.egg-info
+
+output*
+
+*logs/
+
+.history
+
+saved_labels*/

LICENSE

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

README copy.md

@@ -0,0 +1,95 @@
+<p align="center">
+<h1 align="center">DragStream</h1>
+<h3 align="center">Streaming Drag-Oriented Interactive Video Manipulation: Drag Anything, Anytime!</h3>
+</p>
+<p align="center">
+  <p align="center">
+    <a>Junbao Zhou</a><sup>1</sup>
+    ·
+    <a>Yuan Zhou</a><sup>1</sup>
+    ·
+    <a>Kesen Zhao</a><sup>1</sup>
+    ·
+    <a>Qingshan Xu</a><sup>1</sup>
+    ·
+    <a>Beier Zhu</a><sup>1</sup>
+    ·
+    <a>Richang Hong</a><sup>2</sup>
+    ·
+    <a>Hanwang Zhang</a><sup>1</sup><br>
+    <sup>1</sup>Nanyang Technological University <sup>2</sup>Hefei University of Technology
+  </p>
+  <h3 align="center">
+    <a href="https://arxiv.org/abs/2510.03550"><img src="https://img.shields.io/badge/Arxiv-2510.03550-B31B1B.svg?logo=arXiv"></a>
+    <a href="https://junbao-zhou.github.io/DragStream.github.io/"><img src="https://img.shields.io/badge/Project_Page-Lots_of_Demos-Green"></a>
+    <br>
+  </h3>
+</p>
+
+---
+
+Achieving streaming, fine-grained control over the outputs of autoregressive video diffusion models remains challenging, making it difficult to ensure that they consistently align with user expectations. To bridge this gap, we propose **stReaming drag-oriEnted interactiVe vidEo manipuLation (REVEL)**, a new task that enables users to modify generated videos *anytime* on *anything* via fine-grained, interactive drag. Beyond DragVideo and SG-I2V, REVEL unifies drag-style video manipulation as editing and animating video frames with both supporting user-specified translation, deformation, and rotation effects, making drag operations versatile. In resolving REVEL, we observe: *i*) drag-induced perturbations accumulate in latent space, causing severe latent distribution drift that halts the drag process; *ii*) streaming drag is easily disturbed by context frames, thereby yielding visually unnatural outcomes. We thus propose a training-free approach, **DragStream**, comprising: *i*) an adaptive distribution self-rectification strategy that leverages neighboring frames' statistics to effectively constrain the drift of latent embeddings; *ii*) a spatial-frequency selective optimization mechanism, allowing the model to fully exploit contextual information while mitigating its interference via selectively propagating visual cues along generation. Our method can be seamlessly integrated into existing autoregressive video diffusion models, and extensive experiments firmly demonstrate the effectiveness of our DragStream
+
+---
+
+![alt text](asset/figure1.png)
+
+## Requirements
+We tested this repo on the following setup:
+* Nvidia GPU with at least 40 GB memory.
+* Linux operating system.
+* 64 GB RAM.
+
+Other hardware setup could also work but hasn't been tested.
+
+## Installation
+
+### 1. Follow Self-Forcing to Install Dependencies
+
+Create a conda environment and install dependencies:
+```
+conda create -n drag-stream python=3.10 -y
+conda activate drag-stream
+pip install -r requirements.txt
+pip install flash-attn --no-build-isolation
+python setup.py develop
+```
+
+### 2. Follow Self-Forcing to Download Checkpoints
+```
+huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B --local-dir wan_models/Wan2.1-T2V-1.3B
+huggingface-cli download gdhe17/Self-Forcing checkpoints --local-dir ./checkpoints
+```
+
+### 3. Follow Segment-Anything to Install SAM
+
+```
+git clone git@github.com:facebookresearch/segment-anything.git
+cd segment-anything; pip install -e .
+```
+
+### 4. Follow Segment-Anything to Download SAM Checkpoint
+
+- [ViT-H SAM model.](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth)
+
+## Drag/Animate Video with GUI
+
+```
+python app.py
+```
+
+## CLI Inference with Saved Trajectories
+
+```
+python offline_run.py
+```
+
+## Reproducibility
+
+To ensure every Drag/Animation is performed on the same generated video given the same input conditions, we set the random seed before the initialization of random noise and before the generation process.
+
+Please refer to the `set_seed(seed)` in `inference.py`, `stream_inference.py`, `click_gui_video.py`, and `offline_run.py` for details.
+
+
+## Acknowledgements
+This codebase is built on top of the open-source implementation of [Self-Forcing](https://self-forcing.github.io/).

app.py

@@ -0,0 +1,1464 @@
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+import copy
+import torch
+from torchvision.io import write_video
+from torch.utils.data import Dataset
+import numpy as np
+from pathlib import Path
+from hydra import initialize, compose
+from hydra.core.global_hydra import GlobalHydra
+
+from b_spline import build_clamped_bspline, equidistant_points_on_spline
+
+torch.set_grad_enabled(False)
+
+from palette import _palette
+
+import gradio as gr
+import numpy as np
+from scipy import ndimage
+from PIL import Image
+import os
+from pathlib import Path
+import cv2
+
+# from sam_segment import predict_masks_with_sam
+from segment_anything import SamPredictor, sam_model_registry
+
+from tensor_utils import (
+    image_to_pil,
+    image_to_np,
+    bbox_from_mask,
+    draw_bbox_on_image,
+    draw_mask_on_image,
+    draw_points_on_image,
+    draw_lines_on_image,
+    trajectory_interpolate,
+    dilate_mask,
+    dilate_masks,
+)
+
+from optimize_utils import (
+    MultiTrajectory,
+    Trajectory,
+)
+
+import sys
+
+from utils.misc import set_seed
+
+from stream_inference_wrapper import StreamInferenceWrapper
+from stream_drag_inference_wrapper import StreamDragInferenceWrapper
+from utils.dataset import TextDataset
+
+from video_operations import generate_video, optimize_video
+
+# from compute_objmc import visualize_ground_truth_from_trajectory_file
+
+
+def extract_layer_as_mask(image_editor, layer_index=0):
+    if len(image_editor["layers"]) > layer_index:
+        layer = image_editor["layers"][layer_index]
+        return image_to_np(layer.convert("L")) > 0
+    return None
+
+
+def apply_mask_to_image(
+    mask: np.ndarray | None,
+    image: np.ndarray | Image.Image,
+    mask_color: list[int],
+    alpha: float,
+) -> None | Image.Image:
+    if image is None:
+        return None
+    if mask is None:
+        return image_to_pil(image)
+    mask = np.array(mask)
+    new_image = draw_mask_on_image(
+        image,
+        mask,
+        mask_color=mask_color,
+        alpha=alpha,
+    )
+    return new_image
+
+
+def apply_movable_mask_to_image(
+    mask: np.ndarray | None,
+    image: np.ndarray | Image.Image,
+):
+    return apply_mask_to_image(
+        mask=mask,
+        image=image,
+        mask_color=(255, 255, 255),
+        alpha=0.35,
+    )
+
+
+def apply_target_mask_to_image(
+    mask: np.ndarray | None,
+    image: np.ndarray | Image.Image,
+):
+    return apply_mask_to_image(
+        mask=mask,
+        image=image,
+        mask_color=(255, 64, 64),
+        alpha=0.5,
+    )
+
+
+def get_video_last_frame(
+    # video: Optional[torch.Tensor],  # None or (t, h, w, c)
+    video_path: str,
+):
+    """
+    Loads the last frame from a video.
+
+    Returns:
+        Image: The last frame as a PIL Image.
+    """
+    print(f"Getting last frame from video: {video_path = }")
+    if video_path is None:
+        return None
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        print(f"Failed to open video: {video_path}")
+        return None
+    try:
+        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        if frame_count <= 0:
+            print(f"Video has non-positive frame count: {frame_count}")
+            cap.release()
+            return None
+
+        # Try direct seek to last frame
+        target_index = frame_count - 1
+        cap.set(cv2.CAP_PROP_POS_FRAMES, target_index)
+        ret, frame = cap.read()
+
+        # Fallback: iterate to last frame if random access failed
+        if (not ret) or frame is None:
+            print("Direct seek failed, iterating through frames...")
+            cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+            last_valid = None
+            while True:
+                ret_i, frame_i = cap.read()
+                if not ret_i:
+                    break
+                last_valid = frame_i
+            frame = last_valid
+
+        if frame is None:
+            print("Could not retrieve last frame.")
+            return None
+
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        last_frame_image = Image.fromarray(frame)
+        return last_frame_image
+    except Exception as e:
+        print(f"Error extracting last frame: {e}")
+        return None
+    finally:
+        cap.release()
+
+
+def sam_predict_segmentation(
+    sam_predictor: SamPredictor,
+    origin_image: Image.Image | np.ndarray,
+    restriction_mask: np.ndarray,  # (h, w), bool
+    click_points: list[tuple[int, int]],
+    previous_sam_logits: np.ndarray | None,  # (3, 256, 256)
+):
+    # print(f"{restriction_mask.shape = }")
+
+    origin_image_np = image_to_np(origin_image)
+    # print(f"{origin_image_np.shape = }")
+    sam_predictor.set_image(origin_image_np)
+
+    if previous_sam_logits is not None:
+        print(f"{previous_sam_logits.shape = }")
+    else:
+        print(f"{previous_sam_logits = }")
+    masks, scores, logits = sam_predictor.predict(
+        point_coords=np.array(click_points),
+        point_labels=np.ones((len(click_points),)),
+        mask_input=(previous_sam_logits[0:1] if previous_sam_logits is not None else None),
+        multimask_output=True,
+    )
+    # mask: np.ndarray
+    # scores: np.ndarray
+    # logits: np.ndarray
+    # print(f"{masks.shape = }") # (3, 480, 832)
+    # print(f"{logits.shape = }") # (3, 256, 256)
+
+    mask = masks[0]
+    mask *= restriction_mask
+
+    logits *= cv2.resize(
+        restriction_mask.astype(np.uint8),
+        dsize=(256, 256),
+        interpolation=cv2.INTER_LINEAR,
+    )
+
+    return mask, logits
+
+
+def sam_predict_segmentation_wrapper(
+    sam_predictor: SamPredictor,
+    original_image: Image.Image | np.ndarray,
+    restriction_mask: np.ndarray | None,
+    previous_click_points: list[tuple[int, int]],
+    previous_sam_logits: np.ndarray | None,
+    bypass_sam_model: bool,
+    evt: gr.SelectData,
+):
+    # print(f"{restriction_mask = }")
+    original_image = image_to_pil(original_image).convert("RGB")
+
+    if restriction_mask is None:
+        labeled_restriction_mask = np.zeros(
+            (original_image.height, original_image.width), dtype=np.int32
+        )
+    else:
+        labeled_restriction_mask, _ = ndimage.label(restriction_mask, structure=np.ones((3, 3)))
+    # print(f"{labeled_restriction_mask = }")
+    current_click_label = labeled_restriction_mask[evt.index[1], evt.index[0]]
+    # print(f"{current_click_label = }")
+
+    if current_click_label == 0:
+        selected_component_mask = np.zeros_like(labeled_restriction_mask, dtype=bool)
+    else:
+        selected_component_mask = labeled_restriction_mask == current_click_label
+    # print(f"{selected_component_mask = }")
+
+    if bypass_sam_model:
+        click_points = [evt.index]
+        mask = selected_component_mask
+        logits = None
+    else:
+        click_points = previous_click_points + [evt.index]
+        mask, logits = sam_predict_segmentation(
+            sam_predictor=sam_predictor,
+            origin_image=original_image,
+            restriction_mask=selected_component_mask,
+            click_points=click_points,
+            previous_sam_logits=previous_sam_logits,
+        )
+
+    return mask, click_points, logits
+
+
+def draw_all_sam_masks(image: Image.Image | None, mask_list: list[np.ndarray]):
+    if image is None:
+        return None
+    if len(mask_list) == 0:
+        pass
+    else:
+        for mask_idx, mask in enumerate(mask_list):
+            image = draw_mask_on_image(
+                image,
+                mask,
+                mask_color=tuple(_palette[mask_idx + 1]),
+                alpha=0.65,
+            )
+    return image
+
+
+def draw_sam_mask_wrapper(
+    original_image,
+    movable_mask,
+    current_mask: np.ndarray | None,
+    previous_masks: list[np.ndarray],
+    click_points: list[tuple[int, int]],
+):
+    image = apply_movable_mask_to_image(
+        image=original_image,
+        mask=movable_mask,
+    )
+    if image is None:
+        return None
+    image = draw_all_sam_masks(
+        image,
+        previous_masks + ([current_mask] if current_mask is not None else []),
+    )
+    image = draw_points_on_image(
+        image,
+        click_points,
+        color=[(0, 255, 0, 255) for l in click_points],
+        radius=5,
+    )
+    return image
+
+
+def save_sam_masks(
+    current_mask: np.ndarray | None,
+    previous_masks: list[np.ndarray],
+):
+    new_masks = previous_masks + ([current_mask] if current_mask is not None else [])
+    return None, new_masks, [], None
+
+
+def select_target_sam_mask(
+    masks_list: list[np.ndarray],
+    evt: gr.SelectData,
+):
+    is_match_mask = False
+    for mask_index, sam_mask in enumerate(masks_list):
+        # check if evt point in sam_mask
+        if sam_mask[evt.index[1], evt.index[0]]:
+            is_match_mask = True
+            break
+
+    if not is_match_mask:
+        print(f"Mask not found for {evt.index = }")
+        mask_index = -1
+    return mask_index
+
+
+def draw_rotation_trajectory(
+    image,
+    points,
+):
+    image = draw_points_on_image(
+        image,
+        [points[0]],
+        color="green",
+        radius=15,
+    )
+    if len(points) > 1:
+        image = draw_points_on_image(
+            image,
+            points[1:],
+            color=[
+                (
+                    255 - int(float(i) / len(points[1:]) * 255.0),
+                    64,
+                    int(float(i) / len(points[1:]) * 255.0),
+                    255,
+                )
+                for i in range(len(points[1:]))
+            ],
+            radius=5,
+        )
+        for point in points[1:]:
+            image = draw_lines_on_image(
+                image,
+                [points[0], point],
+                color="green",
+                width=3,
+            )
+
+    return image
+
+
+def draw_translation_trajectory(
+    image,
+    points,
+    control_points: list[tuple[int, int]] = [],
+    is_draw_control_points: bool = True,
+):
+    if len(points) == 1:
+        image = draw_points_on_image(
+            image,
+            points,
+            color=[(255, 64, 0, 255)],
+            radius=6,
+        )
+        return image
+    if is_draw_control_points and (len(control_points) >= 2):
+        image = draw_points_on_image(
+            image,
+            control_points,
+            color=[(0, 255, 0, 255) for _ in control_points],
+            radius=3,
+        )
+        image = draw_lines_on_image(
+            image,
+            control_points,
+            color=[(0, 255, 0, 255) for _ in control_points],
+            width=2,
+        )
+    image = draw_lines_on_image(
+        image,
+        points,
+        color=[
+            (
+                255 - int(float(i) / len(points[1:]) * 255.0),
+                64,
+                int(float(i) / len(points[1:]) * 255.0),
+                255,
+            )
+            for i in range(len(points))
+        ],
+        width=4,
+    )
+    image = draw_points_on_image(
+        image,
+        points,
+        color=[
+            (
+                255 - int(float(i) / len(points[1:]) * 255.0),
+                64,
+                int(float(i) / len(points[1:]) * 255.0),
+                255,
+            )
+            for i in range(len(points))
+        ],
+        radius=6,
+    )
+
+    return image
+
+
+def draw_all_trajectories(
+    image,
+    trajectory: MultiTrajectory,
+    is_draw_control_points: bool = True,
+):
+    print(
+        f"""
+draw_all_trajectories:
+"""
+    )
+    if trajectory.trajectories is None:
+        return image
+    for traj in trajectory.trajectories:
+        if traj.original_trajectory is None:
+            continue
+        original_traj = traj.original_trajectory
+        if original_traj["is_rotation"]:
+            image = draw_rotation_trajectory(image, original_traj["points"])
+        else:
+            image = draw_translation_trajectory(
+                image,
+                original_traj["points"],
+                original_traj.get("control_points", []),
+                is_draw_control_points=is_draw_control_points,
+            )
+
+    return image
+
+
+def draw_trajectory_image(
+    original_image,
+    movable_mask,
+    mask_index,
+    masks_list: list[np.ndarray],
+    trajectory: MultiTrajectory,
+    is_draw_bbox: bool = True,
+    is_draw_control_points: bool = True,
+):
+    print(
+        f"""
+draw_trajectory_image:
+    {mask_index = }
+"""
+    )
+    image = apply_movable_mask_to_image(
+        mask=movable_mask,
+        image=original_image,
+    )
+    image = draw_all_sam_masks(image, masks_list)
+    if (
+        (mask_index is not None)
+        and (mask_index >= 0)
+        and (mask_index < len(masks_list))
+        and is_draw_bbox
+    ):
+        image = draw_bbox_on_image(image, bbox_from_mask(masks_list[mask_index]))
+    image = draw_all_trajectories(
+        image,
+        trajectory,
+        is_draw_control_points=is_draw_control_points,
+    )
+    return image
+
+
+def update_trajectory(
+    trajectory: MultiTrajectory,
+    mask_index: int,
+    drag_animation_select: str,
+    translate_rotate_select: str,
+    evt: gr.SelectData,
+):
+    print(f"update_trajectory")
+
+    # Work on a deep copy so Gradio sees a new object
+    trajectory = copy.deepcopy(trajectory)
+
+    if mask_index < 0:
+        print(f"Invalid mask_index: {mask_index}")
+        return trajectory
+
+    # print(f"{evt.index = }")
+    x_center, y_center = evt.index  # evt.value is (x, y)
+
+    clicked_point = (x_center, y_center)
+    print(f"{clicked_point = }")
+
+    # Ensure trajectories list is large enough
+    while len(trajectory.trajectories) <= mask_index:
+        trajectory.trajectories.append(Trajectory())
+
+    existing_traj_obj = trajectory.trajectories[mask_index]
+    if existing_traj_obj.original_trajectory is not None:
+        current_trajectory = dict(existing_traj_obj.original_trajectory)
+    else:
+        current_trajectory = {}
+
+    if translate_rotate_select == "Translation":
+        current_trajectory["is_rotation"] = False
+
+        # Append clicked control point
+        control_points = current_trajectory.get("control_points", [])
+        control_points = control_points + [clicked_point]
+
+        # Drag vs Animation behavior
+        if drag_animation_select == "Drag":
+            # Restrict to last two control points, sample exactly 2 points
+            if len(control_points) > 2:
+                control_points = [clicked_point]
+            num_traj_points = 2
+        elif drag_animation_select == "Animation":
+            # No restriction on control points, sample N = 1 + 3 * block_number
+            num_traj_points = 1 + 3 * int(trajectory.block_number)
+        else:
+            raise ValueError(f"Invalid drag_animation_select: {drag_animation_select}")
+
+        current_trajectory["control_points"] = control_points
+
+        # Compute trajectory points along BSpline (or pad if not enough controls)
+        if len(control_points) < 2:
+            sampled_pts = [control_points[0]] * num_traj_points
+        else:
+            spline = build_clamped_bspline(control_points, degree=3)
+            pts = equidistant_points_on_spline(spline, num_points=num_traj_points, grid=8000)
+            sampled_pts = [(int(round(px)), int(round(py))) for px, py in pts]
+
+        current_trajectory["points"] = sampled_pts
+
+    elif translate_rotate_select == "Rotation":
+        current_trajectory["is_rotation"] = True
+
+        # Initialize if missing, else apply 3-point logic
+        if "points" not in current_trajectory or current_trajectory["points"] is None:
+            current_trajectory["points"] = [clicked_point]
+        else:
+            pts = current_trajectory["points"] + [clicked_point]
+
+            # If about to exceed 3, reset to the new point
+            if len(pts) > 3:
+                current_trajectory["points"] = [clicked_point]
+            # If less than 3, just append
+            elif len(pts) < 3:
+                current_trajectory["points"] = pts
+            else:
+                # len(pts) == 3: pts[0] is rotation center
+                if drag_animation_select == "Animation":
+                    first = trajectory_interpolate(pts[1:], scale=int(trajectory.block_number))
+                    second = trajectory_interpolate(first, scale=3)
+                    current_trajectory["points"] = pts[0:1] + second
+                else:
+                    # Drag: do not interpolate
+                    current_trajectory["points"] = pts
+    else:
+        raise ValueError("Invalid translation/rotation selection")
+
+    # Update the Trajectory object in-place (recomputes block_trajectories)
+    existing_traj_obj.set_original_trajectory(current_trajectory)
+    # print(f"{trajectory = }")
+
+    return trajectory
+
+
+def save_trajectory(
+    save_dir: Path,
+    saved_trajectory: MultiTrajectory,
+    original_image: Image.Image,
+    current_block_index: int,
+    masks: list[np.ndarray],
+):
+    print(f"save_trajectory")
+    print(f"{save_dir = }")
+    print(f"{saved_trajectory = }")
+    save_dir = Path(save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    drag_animation_select = saved_trajectory.drag_or_animation_select or "Drag"
+    save_prefix = f"block_{current_block_index}_{drag_animation_select}"
+
+    # Use MultiTrajectory's save method
+    saved_trajectory.save(
+        save_dir=save_dir,
+        prefix=save_prefix,
+    )
+
+    # Save the trajectory image
+    trajectory_image = draw_trajectory_image(
+        original_image=original_image,
+        movable_mask=saved_trajectory.movable_mask,
+        mask_index=None,
+        masks_list=masks,
+        trajectory=saved_trajectory,
+        is_draw_bbox=False,
+        is_draw_control_points=False,
+    )
+    trajectory_image.save(save_dir / f"{save_prefix}_trajectory.png")
+
+
+def clear_current_trajectory(
+    idx: int,
+    trajectory: MultiTrajectory,
+):
+    trajectory = copy.deepcopy(trajectory)
+    """Clear the trajectory at the given mask index."""
+    try:
+        idx_int = int(idx)
+    except Exception:
+        return trajectory
+
+    if not trajectory.trajectories:
+        return trajectory
+
+    if idx_int < 0 or idx_int >= len(trajectory.trajectories):
+        return trajectory
+
+    # Reset this trajectory (keep the mask)
+    mask = trajectory.trajectories[idx_int].mask
+    trajectory.trajectories[idx_int] = Trajectory(mask=mask)
+
+    return trajectory
+
+
+def clear_all_trajectories(
+    trajectory: MultiTrajectory,
+):
+    trajectory = copy.deepcopy(trajectory)
+    """Clear all trajectories but keep the masks."""
+    if trajectory.trajectories is not None:
+        for i in range(len(trajectory.trajectories)):
+            mask = trajectory.trajectories[i].mask
+            trajectory.trajectories[i] = Trajectory(mask=mask)
+    return trajectory
+
+
+def sync_trajectory_masks(saved_trajectory: MultiTrajectory, dilated_masks: list[np.ndarray]):
+    """Resize saved_trajectory.trajectories to match the number of dilated masks,
+    and update each Trajectory.mask with the corresponding dilated mask."""
+    saved_trajectory = copy.deepcopy(saved_trajectory)
+
+    current_len = len(saved_trajectory.trajectories)
+    target_len = len(dilated_masks) if dilated_masks else 0
+
+    if target_len > current_len:
+        # Expand: append new empty Trajectory objects
+        for _ in range(target_len - current_len):
+            saved_trajectory.trajectories.append(Trajectory())
+    elif target_len < current_len:
+        # Shrink: truncate
+        saved_trajectory.trajectories = saved_trajectory.trajectories[:target_len]
+
+    # Update each Trajectory.mask
+    for i, mask in enumerate(dilated_masks):
+        saved_trajectory.trajectories[i].mask = mask
+
+    return saved_trajectory
+
+
+def add_listeners_to_trajectory(
+    saved_trajectory: MultiTrajectory,
+    prompt_box: gr.Textbox,
+    trajectory_block_number_slider: gr.Slider,
+    drag_animation_select: gr.Dropdown,
+    movable_area_mask: gr.State,
+    dilated_saved_sam_predicted_masks: gr.State,
+):
+    # Sync prompt into saved_trajectory when prompt_box changes
+
+    def sync_trajectory_prompt(saved_trajectory: MultiTrajectory, prompt: str):
+        saved_trajectory.prompt = prompt
+        return saved_trajectory
+
+    prompt_box.change(
+        fn=sync_trajectory_prompt,
+        inputs=[saved_trajectory, prompt_box],
+        outputs=saved_trajectory,
+        trigger_mode="always_last",
+    )
+
+    # Sync block_number into saved_trajectory when trajectory_block_number_slider changes
+    def sync_trajectory_block_number(saved_trajectory: MultiTrajectory, block_number: int):
+        saved_trajectory.block_number = block_number
+        return saved_trajectory
+
+    trajectory_block_number_slider.change(
+        fn=sync_trajectory_block_number,
+        inputs=[saved_trajectory, trajectory_block_number_slider],
+        outputs=saved_trajectory,
+        trigger_mode="always_last",
+    )
+
+    # Sync drag_or_animation_select into saved_trajectory when drag_animation_select changes
+    def sync_trajectory_drag_animation(
+        saved_trajectory: MultiTrajectory, drag_animation_select: str
+    ):
+        saved_trajectory.drag_or_animation_select = drag_animation_select
+        return saved_trajectory
+
+    drag_animation_select.change(
+        fn=sync_trajectory_drag_animation,
+        inputs=[saved_trajectory, drag_animation_select],
+        outputs=saved_trajectory,
+        trigger_mode="always_last",
+    )
+
+    # Sync movable_area_mask into saved_trajectory when it changes
+    def sync_trajectory_movable_mask(saved_trajectory: MultiTrajectory, movable_mask):
+        saved_trajectory.movable_mask = movable_mask
+        return saved_trajectory
+
+    movable_area_mask.change(
+        fn=sync_trajectory_movable_mask,
+        inputs=[saved_trajectory, movable_area_mask],
+        outputs=saved_trajectory,
+        trigger_mode="always_last",
+    )
+
+    # Sync dilated_saved_sam_predicted_masks into saved_trajectory when it changes
+    dilated_saved_sam_predicted_masks.change(
+        fn=sync_trajectory_masks,
+        inputs=[saved_trajectory, dilated_saved_sam_predicted_masks],
+        outputs=saved_trajectory,
+        trigger_mode="always_last",
+    )
+
+
+def create_generate_video_ui(
+    label_root: str | Path,
+    text_dataset: Dataset,
+    video_path: gr.State,
+    stream_drag_inference: StreamDragInferenceWrapper,
+    output_dir: str | Path,
+    original_image: gr.State,
+):
+    with gr.Row():
+        prompt_index_number = gr.Number(
+            label="Step 1: Select Prompt Index Here",
+            interactive=True,
+            scale=1,
+        )
+        prompt_box = gr.Textbox(
+            label="Prompt",
+            interactive=True,
+            scale=3,
+        )
+        save_dir_text_box = gr.Textbox(
+            label="Save Directory",
+            interactive=False,
+            scale=1,
+        )
+    prompt_index_number.change(
+        fn=lambda prompt_index_number: text_dataset[prompt_index_number]["prompts"],
+        inputs=prompt_index_number,
+        outputs=[
+            prompt_box,
+        ],
+    )
+    gr.on(
+        triggers=[
+            prompt_box.change,
+        ],
+        fn=lambda prompt_index_number, prompt: str(
+            label_root / f"{prompt_index_number:04d}-{prompt[:50].replace(' ', '_')}"
+        ),
+        inputs=[prompt_index_number, prompt_box],
+        outputs=save_dir_text_box,
+        trigger_mode="always_last",
+    )
+    with gr.Row():
+        current_block_index_slider = gr.Slider(
+            label="Current Start Block Index",
+            minimum=0,
+            maximum=50,
+            value=0,
+            step=1,
+        )
+        generate_block_number_slider = gr.Slider(
+            label="Step 2: Select Number of Blocks to Generate",
+            minimum=1,
+            maximum=50,
+            value=2,
+            step=1,
+        )
+    with gr.Row():
+        begin_generate_button = gr.Button(
+            value="Step 3: Click Here to Begin Generation",
+        )
+        refresh_video_display_button = gr.Button(value="Refresh Video Display")
+
+    with gr.Row():
+        video_display = gr.Video()
+
+    begin_generate_button.click(
+        fn=lambda pi, p, sbi, bn: generate_video(
+            stream_inference_model=stream_drag_inference,
+            prompt_index=pi,
+            prompt=p,
+            start_block_index=sbi,
+            block_number=bn,
+            output_dir=output_dir,
+        ),
+        inputs=[
+            prompt_index_number,
+            prompt_box,
+            current_block_index_slider,
+            generate_block_number_slider,
+        ],
+        outputs=[video_path, current_block_index_slider],
+    )
+    gr.on(
+        triggers=[
+            refresh_video_display_button.click,
+            video_path.change,
+        ],
+        fn=lambda video_path: video_path,
+        inputs=video_path,
+        outputs=video_display,
+        trigger_mode="always_last",
+    )
+
+    with gr.Row():
+        get_last_frame_button = gr.Button(
+            value="Get Last Frame (Normally No Need to Click This, In Case the Last Frame Fails to Update due to Gradio Bug)",
+        )
+    gr.on(
+        triggers=[
+            video_path.change,
+            get_last_frame_button.click,
+        ],
+        fn=get_video_last_frame,
+        inputs=video_path,
+        outputs=original_image,
+    )
+
+    return (
+        prompt_index_number,
+        save_dir_text_box,
+        prompt_box,
+        current_block_index_slider,
+        generate_block_number_slider,
+    )
+
+
+def create_movable_area_ui(
+    movable_area_mask: gr.State,
+    original_image: gr.State,
+):
+
+    with gr.Row():
+        movable_area_image_editor = gr.ImageEditor(
+            label="Step 4: This is Last Frame of Video, Draw Editable Area Here. (Normally This Should Be Large and Cover all Possible Area Where the Object You Want to Move/Animate to)",
+            type="pil",
+            interactive=True,
+            brush=gr.Brush(
+                default_size=100,
+                colors=[
+                    "rgba(0, 0, 255, 0.5)",
+                ],
+                default_color="auto",
+                color_mode="defaults",
+            ),
+        )
+    movable_area_image_editor.change(
+        fn=extract_layer_as_mask,
+        inputs=movable_area_image_editor,
+        outputs=movable_area_mask,
+        trigger_mode="always_last",
+    )
+    original_image.change(
+        fn=lambda image: image,
+        inputs=original_image,
+        outputs=movable_area_image_editor,
+        trigger_mode="always_last",
+    )
+    with gr.Row():
+        refresh_movable_area_button = gr.Button(
+            value="Refresh Movable Area (Normally No Need to Click This, In Case the Mask Fails to Update due to Gradio Bug)"
+        )
+    refresh_movable_area_button.click(
+        fn=extract_layer_as_mask,
+        inputs=movable_area_image_editor,
+        outputs=movable_area_mask,
+        trigger_mode="always_last",
+    )
+
+
+def create_target_area_ui(
+    target_area_mask: gr.State,
+    original_image: gr.State,
+    movable_area_mask: gr.State,
+):
+    with gr.Row():
+        target_area_image_editor = gr.ImageEditor(
+            label="Step 5: Draw Target Area on the Object You Want to Move/Animate (Normally This Should Be a Subset of Editable Area) (Normally This Mask should be Bigger than the Desired Object)",
+            type="pil",
+            interactive=True,
+            brush=gr.Brush(
+                default_size=50,
+                colors=[
+                    "rgba(255, 0, 0, 0.5)",
+                ],
+                default_color="auto",
+                color_mode="defaults",
+            ),
+        )
+    target_area_image_editor.change(
+        fn=extract_layer_as_mask,
+        inputs=target_area_image_editor,
+        outputs=target_area_mask,
+        trigger_mode="always_last",
+    )
+    gr.on(
+        triggers=[
+            original_image.change,
+            movable_area_mask.change,
+        ],
+        fn=apply_movable_mask_to_image,
+        inputs=[
+            movable_area_mask,
+            original_image,
+        ],
+        outputs=target_area_image_editor,
+        trigger_mode="always_last",
+    )
+
+    with gr.Row():
+        refresh_target_area_button = gr.Button(
+            value="Refresh Target Area (Normally No Need to Click This, In Case the Mask Fails to Update due to Gradio Bug)"
+        )
+    refresh_target_area_button.click(
+        fn=extract_layer_as_mask,
+        inputs=target_area_image_editor,
+        outputs=target_area_mask,
+        trigger_mode="always_last",
+    )
+
+
+def create_sam_segmentation_ui(
+    original_image: gr.State,
+    movable_area_mask: gr.State,
+    target_area_mask: gr.State,
+    sam_predictor: SamPredictor,
+    sam_click_points: gr.State,
+    sam_saved_logits: gr.State,
+    current_sam_predicted_mask: gr.State,
+    saved_sam_predicted_masks: gr.State,
+    dilated_current_sam_predicted_mask: gr.State,
+    dilated_saved_sam_predicted_masks: gr.State,
+):
+    with gr.Row():
+        refresh_sam_segment_click_image_button = gr.Button(
+            value="Refresh Target Area Mask Display (Normally No Need to Click This, In Case the Mask Fails to Update due to Gradio Bug)"
+        )
+    with gr.Row():
+        sam_segment_click_image = gr.Image(
+            label="Step 6: Click to Perform SAM Segment on Target Area, Segment the Object You Want to Move/Animate. The SAM Mask is Restricted within the Target Area Mask",
+            type="pil",
+            interactive=True,
+        )
+    gr.on(
+        triggers=[
+            original_image.change,
+            movable_area_mask.change,
+            target_area_mask.change,
+            refresh_sam_segment_click_image_button.click,
+        ],
+        fn=lambda movable_mask, target_mask, image: apply_target_mask_to_image(
+            target_mask,
+            apply_movable_mask_to_image(
+                movable_mask,
+                image,
+            ),
+        ),
+        inputs=[
+            movable_area_mask,
+            target_area_mask,
+            original_image,
+        ],
+        outputs=sam_segment_click_image,
+        trigger_mode="always_last",
+    )
+
+    with gr.Row():
+        dilate_mask_slider = gr.Slider(
+            label="Dilate Mask Pixel",
+            minimum=0,
+            maximum=50,
+            value=15,
+            step=1,
+        )
+        bypass_sam_model_check_box = gr.Checkbox(
+            label="Bypass SAM Model",
+            value=False,
+        )
+
+    def sam_predict_segmentation_wrapper_wrapper(
+        oi,
+        rm,
+        pcp,
+        psl,
+        bs,
+        evt: gr.SelectData,
+    ):
+        return sam_predict_segmentation_wrapper(
+            sam_predictor=sam_predictor,
+            original_image=oi,
+            restriction_mask=rm,
+            previous_click_points=pcp,
+            previous_sam_logits=psl,
+            bypass_sam_model=bs,
+            evt=evt,
+        )
+
+    sam_segment_click_image.select(
+        fn=sam_predict_segmentation_wrapper_wrapper,
+        inputs=[
+            original_image,
+            target_area_mask,
+            sam_click_points,
+            sam_saved_logits,
+            bypass_sam_model_check_box,
+        ],
+        outputs=[
+            current_sam_predicted_mask,
+            sam_click_points,
+            sam_saved_logits,
+        ],
+        trigger_mode="always_last",
+    )
+    gr.on(
+        triggers=[
+            current_sam_predicted_mask.change,
+            dilate_mask_slider.change,
+        ],
+        fn=dilate_mask,
+        inputs=[
+            current_sam_predicted_mask,
+            dilate_mask_slider,
+        ],
+        outputs=dilated_current_sam_predicted_mask,
+        trigger_mode="always_last",
+    )
+    gr.on(
+        triggers=[
+            saved_sam_predicted_masks.change,
+            dilate_mask_slider.change,
+        ],
+        fn=dilate_masks,
+        inputs=[
+            saved_sam_predicted_masks,
+            dilate_mask_slider,
+        ],
+        outputs=dilated_saved_sam_predicted_masks,
+        trigger_mode="always_last",
+    )
+
+
+def create_sam_mask_management_ui(
+    original_image: gr.State,
+    movable_area_mask: gr.State,
+    dilated_current_sam_predicted_mask: gr.State,
+    dilated_saved_sam_predicted_masks: gr.State,
+    sam_click_points: gr.State,
+    current_sam_predicted_mask: gr.State,
+    saved_sam_predicted_masks: gr.State,
+    sam_saved_logits: gr.State,
+):
+    with gr.Row():
+        save_sam_masks_button = gr.Button(
+            value="Step 7: Save the Current SAM Mask",
+        )
+        cancel_sam_mask_button = gr.Button(value="Cancel Current SAM Mask")
+        delete_sam_mask_button = gr.Button(value="Delete All SAM Masks")
+    save_sam_masks_button.click(
+        fn=save_sam_masks,
+        inputs=[
+            current_sam_predicted_mask,
+            saved_sam_predicted_masks,
+        ],
+        outputs=[
+            current_sam_predicted_mask,
+            saved_sam_predicted_masks,
+            sam_click_points,
+            sam_saved_logits,
+        ],
+        trigger_mode="always_last",
+    )
+    with gr.Row():
+        sam_segment_display_image = gr.Image(
+            label="Step 8: Display the SAM Segmentation, Click to Select Target Object to Create Trajectory",
+            type="pil",
+            interactive=True,
+        )
+    gr.on(
+        triggers=[
+            original_image.change,
+            movable_area_mask.change,
+            dilated_current_sam_predicted_mask.change,
+            dilated_saved_sam_predicted_masks.change,
+            sam_click_points.change,
+        ],
+        fn=draw_sam_mask_wrapper,
+        inputs=[
+            original_image,
+            movable_area_mask,
+            dilated_current_sam_predicted_mask,
+            dilated_saved_sam_predicted_masks,
+            sam_click_points,
+        ],
+        outputs=sam_segment_display_image,
+        trigger_mode="always_last",
+    )
+    cancel_sam_mask_button.click(
+        fn=lambda: (None, [], None),
+        outputs=[
+            current_sam_predicted_mask,
+            sam_click_points,
+            sam_saved_logits,
+        ],
+        trigger_mode="always_last",
+    )
+    gr.on(
+        triggers=[
+            # target_area_mask.change,
+            delete_sam_mask_button.click,
+        ],
+        fn=lambda: (None, [], [], None),
+        outputs=[
+            current_sam_predicted_mask,
+            saved_sam_predicted_masks,
+            sam_click_points,
+            sam_saved_logits,
+        ],
+        trigger_mode="always_last",
+    )
+    with gr.Row():
+        current_selected_mask_index_number = gr.Number(
+            label="Current Selected Mask Index",
+            interactive=False,
+        )
+
+    sam_segment_display_image.select(
+        fn=select_target_sam_mask,
+        inputs=[
+            saved_sam_predicted_masks,
+        ],
+        outputs=[
+            current_selected_mask_index_number,
+        ],
+        trigger_mode="always_last",
+    )
+
+    return current_selected_mask_index_number
+
+
+def create_trajectory_display_ui(
+    original_image: gr.State,
+    movable_area_mask: gr.State,
+    dilated_saved_sam_predicted_masks: gr.State,
+    saved_trajectory: gr.State,
+    current_selected_mask_index_number: gr.State,
+):
+    with gr.Row():
+        trajectory_block_number_slider = gr.Slider(
+            label="Step 9: Select Number of Trajectory Blocks (For Animation Only, More Blocks Means Longer Animation, For Drag, This Should be 1)",
+            minimum=1,
+            maximum=10,
+            value=1,
+            step=1,
+        )
+    with gr.Row():
+        drag_animation_select = gr.Dropdown(
+            choices=["Drag", "Animation"],
+            label="Step 10: Select Drag or Animation",
+        )
+        translate_rotate_select = gr.Dropdown(
+            choices=["Translation", "Rotation"],
+            label="Step 11: Select Translation or Rotation",
+        )
+
+    with gr.Row():
+        trajectory_display_image = gr.Image(
+            label="Step 12: Click on the Object in the Image to Create Trajectory. The Translation Trajectory is Controlled by Bspline Interpolation. The Rotation Trajectory is Controlled by 3 Points",
+            type="pil",
+            interactive=False,
+        )
+    gr.on(
+        triggers=[
+            original_image.change,
+            movable_area_mask.change,
+            current_selected_mask_index_number.change,
+            dilated_saved_sam_predicted_masks.change,
+            saved_trajectory.change,
+        ],
+        fn=draw_trajectory_image,
+        inputs=[
+            original_image,
+            movable_area_mask,
+            current_selected_mask_index_number,
+            dilated_saved_sam_predicted_masks,
+            saved_trajectory,
+        ],
+        outputs=trajectory_display_image,
+        trigger_mode="always_last",
+    )
+
+    trajectory_display_image.select(
+        fn=update_trajectory,
+        inputs=[
+            saved_trajectory,
+            current_selected_mask_index_number,
+            drag_animation_select,
+            translate_rotate_select,
+        ],
+        outputs=saved_trajectory,
+    )
+
+    return drag_animation_select, trajectory_block_number_slider
+
+
+def create_trajectory_management_ui(
+    save_dir_text_box: gr.Textbox,
+    original_image: gr.State,
+    current_block_index_slider: gr.Slider,
+    saved_trajectory: gr.State,
+    dilated_saved_sam_predicted_masks: gr.State,
+    current_selected_mask_index_number: gr.Number,
+):
+    with gr.Row():
+        save_trajectory_button = gr.Button(
+            value="Step 13: Save Trajectory",
+        )
+        delete_current_trajectory_button = gr.Button(value="Delete Current Trajectory")
+        delete_all_trajectory_button = gr.Button(value="Delete All Trajectories")
+    save_trajectory_button.click(
+        fn=save_trajectory,
+        inputs=[
+            save_dir_text_box,
+            saved_trajectory,
+            original_image,
+            current_block_index_slider,
+            dilated_saved_sam_predicted_masks,
+        ],
+    )
+    delete_current_trajectory_button.click(
+        fn=clear_current_trajectory,
+        inputs=[current_selected_mask_index_number, saved_trajectory],
+        outputs=[saved_trajectory],
+    )
+    delete_all_trajectory_button.click(
+        fn=clear_all_trajectories,
+        inputs=[saved_trajectory],
+        outputs=[saved_trajectory],
+    )
+
+
+def create_ui(
+    text_dataset: Dataset,
+    label_root: str | Path,
+    output_dir: str | Path,
+    sam_predictor: SamPredictor,
+    stream_drag_inference: StreamDragInferenceWrapper,
+):
+    with gr.Blocks() as demo:
+        video_path = gr.State(value=None)
+        original_image = gr.State(value=None)
+        movable_area_mask = gr.State(value=None)
+        target_area_mask = gr.State(value=None)
+
+        sam_click_points = gr.State(value=[])
+        sam_saved_logits = gr.State(value=None)
+        saved_sam_predicted_masks = gr.State(value=[])
+        current_sam_predicted_mask = gr.State(value=None)
+
+        dilated_current_sam_predicted_mask = gr.State(value=None)
+        dilated_saved_sam_predicted_masks = gr.State(value=[])
+
+        saved_trajectory = gr.State(value=MultiTrajectory())
+
+        (
+            prompt_index_number,
+            save_dir_text_box,
+            prompt_box,
+            current_block_index_slider,
+            generate_block_number_slider,
+        ) = create_generate_video_ui(
+            label_root=label_root,
+            text_dataset=text_dataset,
+            video_path=video_path,
+            stream_drag_inference=stream_drag_inference,
+            output_dir=output_dir,
+            original_image=original_image,
+        )
+
+        create_movable_area_ui(movable_area_mask, original_image)
+        create_target_area_ui(target_area_mask, original_image, movable_area_mask)
+        create_sam_segmentation_ui(
+            original_image=original_image,
+            movable_area_mask=movable_area_mask,
+            target_area_mask=target_area_mask,
+            sam_predictor=sam_predictor,
+            sam_click_points=sam_click_points,
+            sam_saved_logits=sam_saved_logits,
+            current_sam_predicted_mask=current_sam_predicted_mask,
+            saved_sam_predicted_masks=saved_sam_predicted_masks,
+            dilated_current_sam_predicted_mask=dilated_current_sam_predicted_mask,
+            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
+        )
+
+        current_selected_mask_index_number = create_sam_mask_management_ui(
+            original_image=original_image,
+            movable_area_mask=movable_area_mask,
+            dilated_current_sam_predicted_mask=dilated_current_sam_predicted_mask,
+            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
+            sam_click_points=sam_click_points,
+            current_sam_predicted_mask=current_sam_predicted_mask,
+            saved_sam_predicted_masks=saved_sam_predicted_masks,
+            sam_saved_logits=sam_saved_logits,
+        )
+
+        drag_animation_select, trajectory_block_number_slider = create_trajectory_display_ui(
+            original_image=original_image,
+            movable_area_mask=movable_area_mask,
+            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
+            saved_trajectory=saved_trajectory,
+            current_selected_mask_index_number=current_selected_mask_index_number,
+        )
+        create_trajectory_management_ui(
+            save_dir_text_box=save_dir_text_box,
+            original_image=original_image,
+            current_block_index_slider=current_block_index_slider,
+            saved_trajectory=saved_trajectory,
+            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
+            current_selected_mask_index_number=current_selected_mask_index_number,
+        )
+
+        add_listeners_to_trajectory(
+            saved_trajectory=saved_trajectory,
+            prompt_box=prompt_box,
+            trajectory_block_number_slider=trajectory_block_number_slider,
+            drag_animation_select=drag_animation_select,
+            movable_area_mask=movable_area_mask,
+            dilated_saved_sam_predicted_masks=dilated_saved_sam_predicted_masks,
+        )
+
+        with gr.Row():
+            begin_optimize_button = gr.Button(
+                value="Step 14: Click Here to Begin Optimize, Wait for a Moment and the Dragged/Animated Video will be Displayed Above",
+            )
+        begin_optimize_button.click(
+            fn=lambda pi, sbi, st: optimize_video(
+                stream_drag_inference_model=stream_drag_inference,
+                output_dir=output_dir,
+                prompt_index=pi,
+                start_block_index=sbi,
+                multi_trajectory=st,
+            ),
+            inputs=[
+                prompt_index_number,
+                current_block_index_slider,
+                saved_trajectory,
+            ],
+            outputs=[
+                video_path,
+                current_block_index_slider,
+            ],
+        )
+        with gr.Row():
+            clear_all_button = gr.Button(
+                value="Step 15: Remember to Click Here to Clear All Before Generation/Editing on Next Video, Otherwise the Previous KV Cache will Affect the Generation/Editing of Next Video",
+            )
+
+        def clear_all():
+            stream_drag_inference.reset()
+
+            return (
+                0,
+                None,
+                None,
+                None,
+                None,
+                [],
+                None,
+                [],
+                None,
+                MultiTrajectory(),
+            )
+
+        clear_all_button.click(
+            fn=clear_all,
+            outputs=[
+                current_block_index_slider,
+                video_path,
+                original_image,
+                movable_area_mask,
+                target_area_mask,
+                sam_click_points,
+                sam_saved_logits,
+                saved_sam_predicted_masks,
+                current_sam_predicted_mask,
+                saved_trajectory,
+            ],
+        )
+
+    return demo
+
+
+def main():
+    sam_model = sam_model_registry["vit_h"](checkpoint="../segment-anything/sam_vit_h_4b8939.pth")
+    sam_model.to(device="cuda")
+    sam_predictor = SamPredictor(sam_model)
+
+    SEED = 42
+
+    text_dataset = TextDataset(prompt_path="prompts/MovieGenVideoBench_extended.txt")
+
+    if GlobalHydra.instance().is_initialized():
+        GlobalHydra.instance().clear()
+
+    config_dir = "configs"
+    stream_config_name = "self_forcing_dmd_vsink_stream_drag"
+    with initialize(version_base=None, config_path=config_dir):
+        stream_config = compose(config_name=stream_config_name)
+    print(f"{stream_config = }")
+
+    stream_drag_inference = StreamDragInferenceWrapper(
+        stream_model_config=stream_config,
+        checkpoint_path="./checkpoints/self_forcing_dmd.pt",
+        total_generate_block_number=36,
+        use_ema=True,
+        seed=SEED,
+    )
+    label_save_dir = Path("./saved_labels")
+    label_save_dir = label_save_dir / f"{stream_config_name}-seed{SEED}"
+    label_save_dir.mkdir(parents=True, exist_ok=True)
+
+    output_save_dir = Path("outputs-editing")
+    output_save_dir = output_save_dir / f"{stream_config_name}-seed{SEED}"
+    output_save_dir.mkdir(parents=True, exist_ok=True)
+
+    demo = create_ui(
+        text_dataset=text_dataset,
+        label_root=label_save_dir,
+        output_dir=output_save_dir,
+        sam_predictor=sam_predictor,
+        stream_drag_inference=stream_drag_inference,
+    )
+    demo.launch(server_name="0.0.0.0")
+
+
+if __name__ == "__main__":
+    main()

b_spline.py

@@ -0,0 +1,160 @@
+# Click to place control points, then see a clamped B-spline plus
+# equally spaced points along the curve.
+#
+# Controls:
+# - Left-click: add a control point
+# - Right-click: remove the last control point
+# - '+' / '-': increase/decrease the number of equally spaced points
+# - 'c': clear all points
+#
+# Requirements: numpy, matplotlib, scipy
+
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.interpolate import BSpline
+from scipy.integrate import cumulative_trapezoid
+
+# ------------- Spline utilities -------------
+
+
+def _open_uniform_knots(n, k):
+    # n = number of control points, k = degree
+    m = n + k + 1
+    t = np.zeros(m, dtype=float)
+    t[-(k + 1) :] = 1.0
+    num_interior = n - k - 1
+    if num_interior > 0:
+        t[k + 1 : n] = np.linspace(1 / (n - k), (n - k - 1) / (n - k), num_interior)
+    return t
+
+
+def build_clamped_bspline(control_points, degree=3):
+    """
+    Build an open-uniform (clamped) B-spline that passes through the
+    first and last control points.
+    """
+    P = np.asarray(control_points, dtype=float)
+    n = len(P)
+    if n < 2:
+        raise ValueError("Need at least 2 control points")
+    k = min(degree, n - 1)  # degree cannot exceed n-1
+    t = _open_uniform_knots(n, k)
+    return BSpline(t, P, k, axis=0)
+
+
+def equidistant_points_on_spline(spline, num_points, grid=6000):
+    """
+    Return `num_points` points equally spaced in arc length along `spline`.
+    """
+    if num_points < 2:
+        raise ValueError("num_points must be >= 2")
+
+    u = np.linspace(0.0, 1.0, grid)
+    dCdu = spline.derivative()(u)
+    speed = np.linalg.norm(dCdu, axis=1)
+
+    s = cumulative_trapezoid(speed, u, initial=0.0)
+    total_len = s[-1]
+
+    if total_len <= 1e-12:
+        P0 = spline(0.0)
+        return np.repeat(P0[None, :], num_points, axis=0)
+
+    s_targets = np.linspace(0.0, total_len, num_points)
+    u_targets = np.interp(s_targets, s, u)
+    return spline(u_targets)
+
+
+# ------------- Interactive demo -------------
+
+if __name__ == "__main__":
+    # Optional: draw over an image. Uncomment and set path if needed.
+    # img = plt.imread("your_image.png")
+    # H, W = img.shape[:2]
+
+    points = []
+    sample_count = [25]  # use a list so we can modify inside callbacks
+    sample_count = [2]  # use a list so we can modify inside callbacks
+
+    fig, ax = plt.subplots()
+    ax.set_aspect("equal", adjustable="box")
+
+    # If drawing over an image, uncomment:
+    # ax.imshow(img, extent=[0, W, H, 0], origin='upper')
+    # ax.set_xlim(0, W)
+    # ax.set_ylim(H, 0)
+    # Otherwise, use a unit square canvas:
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+
+    title_template = "Left-click: add, Right-click: undo | +/-: change N | c: clear | N = {}"
+    ax.set_title(title_template.format(sample_count[0]))
+
+    # Artists
+    scatter_ctrl = ax.scatter([], [], c="r", s=25, zorder=3, label="control points")
+    (ctrl_line,) = ax.plot([], [], "r--", lw=1, alpha=0.6, zorder=2, label="control polygon")
+    (curve_line,) = ax.plot([], [], "g-", lw=2, zorder=1, label="B-spline")
+    eq_scatter = ax.scatter([], [], c="b", s=20, zorder=4, label="equally spaced points")
+
+    ax.legend(loc="upper right")
+
+    def update_plot():
+        if points:
+            P = np.array(points, dtype=float)
+            scatter_ctrl.set_offsets(P)
+            ctrl_line.set_data(P[:, 0], P[:, 1])
+        else:
+            scatter_ctrl.set_offsets(np.empty((0, 2)))
+            ctrl_line.set_data([], [])
+
+        if len(points) >= 2:
+            try:
+                spl = build_clamped_bspline(points, degree=3)
+                # For a smooth preview of the curve
+                C = spl(np.linspace(0, 1, 1000))
+                curve_line.set_data(C[:, 0], C[:, 1])
+
+                # Equally spaced points along the curve
+                N = max(2, sample_count[0])
+                eq_pts = equidistant_points_on_spline(spl, num_points=N, grid=8000)
+                eq_scatter.set_offsets(eq_pts)
+            except Exception as e:
+                print("Error building spline:", e)
+                curve_line.set_data([], [])
+                eq_scatter.set_offsets(np.empty((0, 2)))
+        else:
+            curve_line.set_data([], [])
+            eq_scatter.set_offsets(np.empty((0, 2)))
+
+        ax.set_title(title_template.format(sample_count[0]))
+        fig.canvas.draw_idle()
+
+    def onclick(event):
+        if event.inaxes != ax:
+            return
+        if event.button == 1:  # left: add
+            if event.xdata is None or event.ydata is None:
+                return
+            points.append([event.xdata, event.ydata])
+            update_plot()
+        elif event.button == 3:  # right: undo last
+            if points:
+                points.pop()
+                update_plot()
+
+    def onkey(event):
+        if event.key in ["+", "="]:
+            sample_count[0] += 1
+            update_plot()
+        elif event.key == "-":
+            if sample_count[0] > 2:
+                sample_count[0] -= 1
+                update_plot()
+        elif event.key in ["c", "C"]:
+            points.clear()
+            update_plot()
+
+    fig.canvas.mpl_connect("button_press_event", onclick)
+    fig.canvas.mpl_connect("key_press_event", onkey)
+
+    plt.show()

configs/default_config.yaml

@@ -0,0 +1,31 @@
+independent_first_frame: false
+warp_denoising_step: false
+weight_decay: 0.01
+same_step_across_blocks: true
+discriminator_lr_multiplier: 1.0
+last_step_only: false
+i2v: false
+num_training_frames: 21
+gc_interval: 100
+context_noise: 0
+causal: true
+
+ckpt_step: 0
+prompt_name: MovieGenVideoBench
+prompt_path: prompts/MovieGenVideoBench.txt
+eval_first_n: 64
+num_samples: 1
+height: 480
+width: 832
+num_frames: 81
+
+no_save: false
+no_visualize: false
+logdir: logs
+wandb_save_dir: ""
+disable_wandb: false
+
+# Recommended: prevent Hydra from chdir into run dir (optional but common)
+hydra:
+  job:
+    chdir: false

configs/self_forcing_dmd.yaml

@@ -0,0 +1,56 @@
+defaults:
+  - default_config
+  - _self_
+
+generator_ckpt: checkpoints/ode_init.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-14B
+# real_name: Wan2.1-T2V-1.3B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: WANDB_HOST
+wandb_key: WANDB_KEY
+wandb_entity: WANDB_ENTITY
+wandb_project: WANDB_PROJECT
+sharding_strategy: hybrid_full
+lr: 2.0e-06
+lr_critic: 4.0e-07
+beta1: 0.0
+beta2: 0.999
+beta1_critic: 0.0
+beta2_critic: 0.999
+data_path: prompts/vidprom_filtered_extended.txt
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 64
+log_iters: 500
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 3
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0

configs/self_forcing_dmd_vsink.yaml

@@ -0,0 +1,58 @@
+defaults:
+  - default_config
+  - _self_
+
+generator_ckpt: checkpoints/ode_init.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+# real_name: Wan2.1-T2V-1.3B
+real_name: Wan2.1-T2V-1.3B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: WANDB_HOST
+wandb_key: WANDB_KEY
+wandb_entity: WANDB_ENTITY
+wandb_project: WANDB_PROJECT
+sharding_strategy: hybrid_full
+lr: 2.0e-06
+lr_critic: 4.0e-07
+beta1: 0.0
+beta2: 0.999
+beta1_critic: 0.0
+beta2_critic: 0.999
+data_path: prompts/vidprom_filtered_extended.txt
+batch_size: 1
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 64
+log_iters: 400
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 3
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0
+  local_attn_size: 21
+  sink_size: 3

configs/self_forcing_dmd_vsink_stream.yaml

@@ -0,0 +1,5 @@
+defaults:
+  - self_forcing_dmd_vsink
+  - _self_
+
+vae_offload_cpu: false

configs/self_forcing_dmd_vsink_stream_drag.yaml

@@ -0,0 +1,33 @@
+defaults:
+  - self_forcing_dmd_vsink_stream
+  - _self_
+
+drag_optim_config:
+  record_feature_block_indexes:
+    - 12
+    - 13
+    - 14
+    - 15
+
+  optimize_denoising_steps_indexes:
+    - 1
+
+  optimize_iter: 5
+  optimize_lr: 0.03
+
+  translation_step: 16.0
+  rotation_step: 8.0
+
+  normalize_latent_after_drag_optimize: True
+  normalize_latent_after_post_merge: True
+  dynamic_chunk_normalization_block_number: 2
+
+  feature_scaling_factor: 2.0
+
+  gradient_gaussian_padding: 1.5
+  gradient_gaussian_sigma: 1.5
+
+  feature_fft_cutoff:
+    - 0.2
+    - 0.4
+    - 0.6

configs/self_forcing_sid.yaml

@@ -0,0 +1,57 @@
+defaults:
+  - default_config
+  - _self_
+
+generator_ckpt: checkpoints/ode_init.pt
+generator_fsdp_wrap_strategy: size
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+real_name: Wan2.1-T2V-1.3B
+text_encoder_fsdp_wrap_strategy: size
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true # need to remove - 0 in denoising_step_list if warp_denoising_step is true
+ts_schedule: false
+num_train_timestep: 1000
+timestep_shift: 5.0
+guidance_scale: 3.0
+denoising_loss_type: flow
+mixed_precision: true
+seed: 0
+wandb_host: WANDB_HOST
+wandb_key: WANDB_KEY
+wandb_entity: WANDB_ENTITY
+wandb_project: WANDB_PROJECT
+sharding_strategy: hybrid_full
+lr: 2.0e-06
+lr_critic: 2.0e-06
+beta1: 0.0
+beta2: 0.999
+beta1_critic: 0.0
+beta2_critic: 0.999
+weight_decay: 0.0
+data_path: prompts/vidprom_filtered_extended.txt
+batch_size: 1
+sid_alpha: 1.0
+ema_weight: 0.99
+ema_start_step: 200
+total_batch_size: 64
+log_iters: 50
+negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+dfake_gen_update_ratio: 5
+image_or_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+distribution_loss: dmd
+trainer: score_distillation
+gradient_checkpointing: true
+num_frame_per_block: 3
+load_raw_video: false
+model_kwargs:
+  timestep_shift: 5.0

demo.py

@@ -0,0 +1,715 @@
+"""
+Demo for Self-Forcing.
+"""
+
+import os
+import re
+import random
+import time
+import base64
+import argparse
+import hashlib
+import subprocess
+import urllib.request
+from io import BytesIO
+from PIL import Image
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from flask import Flask, render_template, jsonify
+from flask_socketio import SocketIO, emit
+import queue
+from threading import Thread, Event
+
+from pipeline import CausalInferencePipeline
+from demo_utils.constant import ZERO_VAE_CACHE
+from demo_utils.vae_block3 import VAEDecoderWrapper
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder
+from demo_utils.utils import generate_timestamp
+from demo_utils.memory import (
+    gpu,
+    get_cuda_free_memory_gb,
+    DynamicSwapInstaller,
+    move_model_to_device_with_memory_preservation,
+)
+
+# Parse arguments
+parser = argparse.ArgumentParser()
+parser.add_argument("--port", type=int, default=5001)
+parser.add_argument("--host", type=str, default="0.0.0.0")
+parser.add_argument("--checkpoint_path", type=str, default="./checkpoints/self_forcing_dmd.pt")
+parser.add_argument("--config_path", type=str, default="./configs/self_forcing_dmd.yaml")
+parser.add_argument("--trt", action="store_true")
+args = parser.parse_args()
+
+print(f"Free VRAM {get_cuda_free_memory_gb(gpu)} GB")
+low_memory = get_cuda_free_memory_gb(gpu) < 40
+
+# Load models
+config = OmegaConf.load(args.config_path)
+default_config = OmegaConf.load("configs/default_config.yaml")
+config = OmegaConf.merge(default_config, config)
+
+text_encoder = WanTextEncoder()
+
+# Global variables for dynamic model switching
+current_vae_decoder = None
+current_use_taehv = False
+fp8_applied = False
+torch_compile_applied = False
+global frame_number
+frame_number = 0
+anim_name = ""
+frame_rate = 6
+
+
+def initialize_vae_decoder(use_taehv=False, use_trt=False):
+    """Initialize VAE decoder based on the selected option"""
+    global current_vae_decoder, current_use_taehv
+
+    if use_trt:
+        from demo_utils.vae import VAETRTWrapper
+
+        current_vae_decoder = VAETRTWrapper()
+        return current_vae_decoder
+
+    if use_taehv:
+        from demo_utils.taehv import TAEHV
+
+        # Check if taew2_1.pth exists in checkpoints folder, download if missing
+        taehv_checkpoint_path = "checkpoints/taew2_1.pth"
+        if not os.path.exists(taehv_checkpoint_path):
+            print(
+                f"taew2_1.pth not found in checkpoints folder {taehv_checkpoint_path}. Downloading..."
+            )
+            os.makedirs("checkpoints", exist_ok=True)
+            download_url = "https://github.com/madebyollin/taehv/raw/main/taew2_1.pth"
+            try:
+                urllib.request.urlretrieve(download_url, taehv_checkpoint_path)
+                print(f"Successfully downloaded taew2_1.pth to {taehv_checkpoint_path}")
+            except Exception as e:
+                print(f"Failed to download taew2_1.pth: {e}")
+                raise
+
+        class DotDict(dict):
+            __getattr__ = dict.__getitem__
+            __setattr__ = dict.__setitem__
+
+        class TAEHVDiffusersWrapper(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.dtype = torch.float16
+                self.taehv = TAEHV(checkpoint_path=taehv_checkpoint_path).to(self.dtype)
+                self.config = DotDict(scaling_factor=1.0)
+
+            def decode(self, latents, return_dict=None):
+                # n, c, t, h, w = latents.shape
+                # low-memory, set parallel=True for faster + higher memory
+                return self.taehv.decode_video(latents, parallel=False).mul_(2).sub_(1)
+
+        current_vae_decoder = TAEHVDiffusersWrapper()
+    else:
+        current_vae_decoder = VAEDecoderWrapper()
+        vae_state_dict = torch.load("wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth", map_location="cpu")
+        decoder_state_dict = {}
+        for key, value in vae_state_dict.items():
+            if "decoder." in key or "conv2" in key:
+                decoder_state_dict[key] = value
+        current_vae_decoder.load_state_dict(decoder_state_dict)
+
+    current_vae_decoder.eval()
+    current_vae_decoder.to(dtype=torch.float16)
+    current_vae_decoder.requires_grad_(False)
+    current_vae_decoder.to(gpu)
+    current_use_taehv = use_taehv
+
+    print(f"✅ VAE decoder initialized with {'TAEHV' if use_taehv else 'default VAE'}")
+    return current_vae_decoder
+
+
+# Initialize with default VAE
+vae_decoder = initialize_vae_decoder(use_taehv=False, use_trt=args.trt)
+
+transformer = WanDiffusionWrapper(is_causal=True)
+state_dict = torch.load(args.checkpoint_path, map_location="cpu")
+transformer.load_state_dict(state_dict["generator_ema"])
+
+text_encoder.eval()
+transformer.eval()
+
+transformer.to(dtype=torch.float16)
+text_encoder.to(dtype=torch.bfloat16)
+
+text_encoder.requires_grad_(False)
+transformer.requires_grad_(False)
+
+pipeline = CausalInferencePipeline(
+    config,
+    device=gpu,
+    generator=transformer,
+    text_encoder=text_encoder,
+    vae=vae_decoder,
+)
+
+if low_memory:
+    DynamicSwapInstaller.install_model(text_encoder, device=gpu)
+else:
+    text_encoder.to(gpu)
+transformer.to(gpu)
+
+# Flask and SocketIO setup
+app = Flask(__name__)
+app.config["SECRET_KEY"] = "frontend_buffered_demo"
+socketio = SocketIO(app, cors_allowed_origins="*")
+
+generation_active = False
+stop_event = Event()
+frame_send_queue = queue.Queue()
+sender_thread = None
+models_compiled = False
+
+
+def tensor_to_base64_frame(frame_tensor):
+    """Convert a single frame tensor to base64 image string."""
+    global frame_number, anim_name
+    # Clamp and normalize to 0-255
+    frame = torch.clamp(frame_tensor.float(), -1.0, 1.0) * 127.5 + 127.5
+    frame = frame.to(torch.uint8).cpu().numpy()
+
+    # CHW -> HWC
+    if len(frame.shape) == 3:
+        frame = np.transpose(frame, (1, 2, 0))
+
+    # Convert to PIL Image
+    if frame.shape[2] == 3:  # RGB
+        image = Image.fromarray(frame, "RGB")
+    else:  # Handle other formats
+        image = Image.fromarray(frame)
+
+    # Convert to base64
+    buffer = BytesIO()
+    image.save(buffer, format="JPEG", quality=100)
+    if not os.path.exists("./images/%s" % anim_name):
+        os.makedirs("./images/%s" % anim_name)
+    frame_number += 1
+    image.save("./images/%s/%s_%03d.jpg" % (anim_name, anim_name, frame_number))
+    img_str = base64.b64encode(buffer.getvalue()).decode()
+    return f"data:image/jpeg;base64,{img_str}"
+
+
+def frame_sender_worker():
+    """Background thread that processes frame send queue non-blocking."""
+    global frame_send_queue, generation_active, stop_event
+
+    print("📡 Frame sender thread started")
+
+    while True:
+        frame_data = None
+        try:
+            # Get frame data from queue
+            frame_data = frame_send_queue.get(timeout=1.0)
+
+            if frame_data is None:  # Shutdown signal
+                frame_send_queue.task_done()  # Mark shutdown signal as done
+                break
+
+            frame_tensor, frame_index, block_index, job_id = frame_data
+
+            # Convert tensor to base64
+            base64_frame = tensor_to_base64_frame(frame_tensor)
+
+            # Send via SocketIO
+            try:
+                socketio.emit(
+                    "frame_ready",
+                    {
+                        "data": base64_frame,
+                        "frame_index": frame_index,
+                        "block_index": block_index,
+                        "job_id": job_id,
+                    },
+                )
+            except Exception as e:
+                print(f"⚠️ Failed to send frame {frame_index}: {e}")
+
+            frame_send_queue.task_done()
+
+        except queue.Empty:
+            # Check if we should continue running
+            if not generation_active and frame_send_queue.empty():
+                break
+        except Exception as e:
+            print(f"❌ Frame sender error: {e}")
+            # Make sure to mark task as done even if there's an error
+            if frame_data is not None:
+                try:
+                    frame_send_queue.task_done()
+                except Exception as e:
+                    print(f"❌ Failed to mark frame task as done: {e}")
+            break
+
+    print("📡 Frame sender thread stopped")
+
+
+@torch.no_grad()
+def generate_video_stream(
+    prompt, seed, enable_torch_compile=False, enable_fp8=False, use_taehv=False
+):
+    """Generate video and push frames immediately to frontend."""
+    global generation_active, stop_event, frame_send_queue, sender_thread, models_compiled, torch_compile_applied, fp8_applied, current_vae_decoder, current_use_taehv, frame_rate, anim_name
+
+    try:
+        generation_active = True
+        stop_event.clear()
+        job_id = generate_timestamp()
+
+        # Start frame sender thread if not already running
+        if sender_thread is None or not sender_thread.is_alive():
+            sender_thread = Thread(target=frame_sender_worker, daemon=True)
+            sender_thread.start()
+
+        # Emit progress updates
+        def emit_progress(message, progress):
+            try:
+                socketio.emit(
+                    "progress",
+                    {
+                        "message": message,
+                        "progress": progress,
+                        "job_id": job_id,
+                    },
+                )
+            except Exception as e:
+                print(f"❌ Failed to emit progress: {e}")
+
+        emit_progress("Starting generation...", 0)
+
+        # Handle VAE decoder switching
+        if use_taehv != current_use_taehv:
+            emit_progress("Switching VAE decoder...", 2)
+            print(f"🔄 Switching VAE decoder to {'TAEHV' if use_taehv else 'default VAE'}")
+            current_vae_decoder = initialize_vae_decoder(use_taehv=use_taehv)
+            # Update pipeline with new VAE decoder
+            pipeline.vae = current_vae_decoder
+
+        # Handle FP8 quantization
+        if enable_fp8 and not fp8_applied:
+            emit_progress("Applying FP8 quantization...", 3)
+            print("🔧 Applying FP8 quantization to transformer")
+            from torchao.quantization.quant_api import (
+                quantize_,
+                Float8DynamicActivationFloat8WeightConfig,
+                PerTensor,
+            )
+
+            quantize_(
+                transformer,
+                Float8DynamicActivationFloat8WeightConfig(granularity=PerTensor()),
+            )
+            fp8_applied = True
+
+        # Text encoding
+        emit_progress("Encoding text prompt...", 8)
+        conditional_dict = text_encoder(text_prompts=[prompt])
+        for key, value in conditional_dict.items():
+            conditional_dict[key] = value.to(dtype=torch.float16)
+        if low_memory:
+            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
+            move_model_to_device_with_memory_preservation(
+                text_encoder,
+                target_device=gpu,
+                preserved_memory_gb=gpu_memory_preservation,
+            )
+
+        # Handle torch.compile if enabled
+        torch_compile_applied = enable_torch_compile
+        if enable_torch_compile and not models_compiled:
+            # Compile transformer and decoder
+            transformer.compile(mode="max-autotune-no-cudagraphs")
+            if not current_use_taehv and not low_memory and not args.trt:
+                current_vae_decoder.compile(mode="max-autotune-no-cudagraphs")
+
+        # Initialize generation
+        emit_progress("Initializing generation...", 12)
+
+        rnd = torch.Generator(gpu).manual_seed(seed)
+        # all_latents = torch.zeros([1, 21, 16, 60, 104], device=gpu, dtype=torch.bfloat16)
+
+        pipeline._initialize_kv_cache(batch_size=1, dtype=torch.float16, device=gpu)
+        pipeline._initialize_crossattn_cache(batch_size=1, dtype=torch.float16, device=gpu)
+
+        noise = torch.randn([1, 21, 16, 60, 104], device=gpu, dtype=torch.float16, generator=rnd)
+
+        # Generation parameters
+        num_blocks = 7
+        current_start_frame = 0
+        num_input_frames = 0
+        all_num_frames = [pipeline.num_frame_per_block] * num_blocks
+        if current_use_taehv:
+            vae_cache = None
+        else:
+            vae_cache = ZERO_VAE_CACHE
+            for i in range(len(vae_cache)):
+                vae_cache[i] = vae_cache[i].to(device=gpu, dtype=torch.float16)
+
+        total_frames_sent = 0
+        generation_start_time = time.time()
+
+        emit_progress("Generating frames... (frontend handles timing)", 15)
+
+        for idx, current_num_frames in enumerate(all_num_frames):
+            if not generation_active or stop_event.is_set():
+                break
+
+            progress = int(((idx + 1) / len(all_num_frames)) * 80) + 15
+
+            # Special message for first block with torch.compile
+            if idx == 0 and torch_compile_applied and not models_compiled:
+                emit_progress(
+                    f"Processing block 1/{len(all_num_frames)} - Compiling models (may take 5-10 minutes)...",
+                    progress,
+                )
+                print(f"🔥 Processing block {idx+1}/{len(all_num_frames)}")
+                models_compiled = True
+            else:
+                emit_progress(
+                    f"Processing block {idx+1}/{len(all_num_frames)}...",
+                    progress,
+                )
+                print(f"🔄 Processing block {idx+1}/{len(all_num_frames)}")
+
+            block_start_time = time.time()
+
+            noisy_input = noise[
+                :,
+                current_start_frame
+                - num_input_frames : current_start_frame
+                + current_num_frames
+                - num_input_frames,
+            ]
+
+            # Denoising loop
+            denoising_start = time.time()
+            for index, current_timestep in enumerate(pipeline.denoising_step_list):
+                if not generation_active or stop_event.is_set():
+                    break
+
+                timestep = (
+                    torch.ones(
+                        [1, current_num_frames],
+                        device=noise.device,
+                        dtype=torch.int64,
+                    )
+                    * current_timestep
+                )
+
+                if index < len(pipeline.denoising_step_list) - 1:
+                    _, denoised_pred = transformer(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=pipeline.kv_cache1,
+                        crossattn_cache=pipeline.crossattn_cache,
+                        current_start=current_start_frame * pipeline.frame_seq_length,
+                    )
+                    next_timestep = pipeline.denoising_step_list[index + 1]
+                    noisy_input = pipeline.scheduler.add_noise(
+                        denoised_pred.flatten(0, 1),
+                        torch.randn_like(denoised_pred.flatten(0, 1)),
+                        next_timestep
+                        * torch.ones(
+                            [1 * current_num_frames],
+                            device=noise.device,
+                            dtype=torch.long,
+                        ),
+                    ).unflatten(0, denoised_pred.shape[:2])
+                else:
+                    _, denoised_pred = transformer(
+                        noisy_image_or_video=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=pipeline.kv_cache1,
+                        crossattn_cache=pipeline.crossattn_cache,
+                        current_start=current_start_frame * pipeline.frame_seq_length,
+                    )
+
+            if not generation_active or stop_event.is_set():
+                break
+
+            denoising_time = time.time() - denoising_start
+            print(f"⚡ Block {idx+1} denoising completed in {denoising_time:.2f}s")
+
+            # Record output
+            # all_latents[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
+
+            # Update KV cache for next block
+            if idx != len(all_num_frames) - 1:
+                transformer(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=conditional_dict,
+                    timestep=torch.zeros_like(timestep),
+                    kv_cache=pipeline.kv_cache1,
+                    crossattn_cache=pipeline.crossattn_cache,
+                    current_start=current_start_frame * pipeline.frame_seq_length,
+                )
+
+            # Decode to pixels and send frames immediately
+            print(f"🎨 Decoding block {idx+1} to pixels...")
+            decode_start = time.time()
+            if args.trt:
+                all_current_pixels = []
+                for i in range(denoised_pred.shape[1]):
+                    is_first_frame = (
+                        torch.tensor(1.0).cuda().half()
+                        if idx == 0 and i == 0
+                        else torch.tensor(0.0).cuda().half()
+                    )
+                    outputs = vae_decoder.forward(
+                        denoised_pred[:, i : i + 1, :, :, :].half(),
+                        is_first_frame,
+                        *vae_cache,
+                    )
+                    # outputs = vae_decoder.forward(denoised_pred.float(), *vae_cache)
+                    current_pixels, vae_cache = outputs[0], outputs[1:]
+                    print(current_pixels.max(), current_pixels.min())
+                    all_current_pixels.append(current_pixels.clone())
+                pixels = torch.cat(all_current_pixels, dim=1)
+                if idx == 0:
+                    pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+            else:
+                if current_use_taehv:
+                    if vae_cache is None:
+                        vae_cache = denoised_pred
+                    else:
+                        denoised_pred = torch.cat([vae_cache, denoised_pred], dim=1)
+                        vae_cache = denoised_pred[:, -3:, :, :, :]
+                    pixels = current_vae_decoder.decode(denoised_pred)
+                    print(f"denoised_pred shape: {denoised_pred.shape}")
+                    print(f"pixels shape: {pixels.shape}")
+                    if idx == 0:
+                        pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+                    else:
+                        pixels = pixels[:, 12:, :, :, :]
+
+                else:
+                    pixels, vae_cache = current_vae_decoder(denoised_pred.half(), *vae_cache)
+                    if idx == 0:
+                        pixels = pixels[:, 3:, :, :, :]  # Skip first 3 frames of first block
+
+            decode_time = time.time() - decode_start
+            print(f"🎨 Block {idx+1} VAE decoding completed in {decode_time:.2f}s")
+
+            # Queue frames for non-blocking sending
+            block_frames = pixels.shape[1]
+            print(f"📡 Queueing {block_frames} frames from block {idx+1} for sending...")
+            queue_start = time.time()
+
+            for frame_idx in range(block_frames):
+                if not generation_active or stop_event.is_set():
+                    break
+
+                frame_tensor = pixels[0, frame_idx].cpu()
+
+                # Queue frame data in non-blocking way
+                frame_send_queue.put((frame_tensor, total_frames_sent, idx, job_id))
+                total_frames_sent += 1
+
+            queue_time = time.time() - queue_start
+            block_time = time.time() - block_start_time
+            print(
+                f"✅ Block {idx+1} completed in {block_time:.2f}s ({block_frames} frames queued in {queue_time:.3f}s)"
+            )
+
+            current_start_frame += current_num_frames
+
+        generation_time = time.time() - generation_start_time
+        print(
+            f"🎉 Generation completed in {generation_time:.2f}s! {total_frames_sent} frames queued for sending"
+        )
+
+        # Wait for all frames to be sent before completing
+        emit_progress("Waiting for all frames to be sent...", 97)
+        print("⏳ Waiting for all frames to be sent...")
+        frame_send_queue.join()  # Wait for all queued frames to be processed
+        print("✅ All frames sent successfully!")
+
+        generate_mp4_from_images("./images", "./videos/" + anim_name + ".mp4", frame_rate)
+        # Final progress update
+        emit_progress("Generation complete!", 100)
+
+        try:
+            socketio.emit(
+                "generation_complete",
+                {
+                    "message": "Video generation completed!",
+                    "total_frames": total_frames_sent,
+                    "generation_time": f"{generation_time:.2f}s",
+                    "job_id": job_id,
+                },
+            )
+        except Exception as e:
+            print(f"❌ Failed to emit generation complete: {e}")
+
+    except Exception as e:
+        print(f"❌ Generation failed: {e}")
+        try:
+            socketio.emit(
+                "error",
+                {"message": f"Generation failed: {str(e)}", "job_id": job_id},
+            )
+        except Exception as e:
+            print(f"❌ Failed to emit error: {e}")
+    finally:
+        generation_active = False
+        stop_event.set()
+
+        # Clean up sender thread
+        try:
+            frame_send_queue.put(None)
+        except Exception as e:
+            print(f"❌ Failed to put None in frame_send_queue: {e}")
+
+
+def generate_mp4_from_images(image_directory, output_video_path, fps=24):
+    """
+    Generate an MP4 video from a directory of images ordered alphabetically.
+
+    :param image_directory: Path to the directory containing images.
+    :param output_video_path: Path where the output MP4 will be saved.
+    :param fps: Frames per second for the output video.
+    """
+    global anim_name
+    # Construct the ffmpeg command
+    cmd = [
+        "ffmpeg",
+        "-framerate",
+        str(fps),
+        "-i",
+        os.path.join(
+            image_directory, anim_name + "/" + anim_name + "_%03d.jpg"
+        ),  # Adjust the pattern if necessary
+        "-c:v",
+        "libx264",
+        "-pix_fmt",
+        "yuv420p",
+        output_video_path,
+    ]
+    try:
+        subprocess.run(cmd, check=True)
+        print(f"Video saved to {output_video_path}")
+    except subprocess.CalledProcessError as e:
+        print(f"An error occurred: {e}")
+
+
+def calculate_sha256(data):
+    # Convert data to bytes if it's not already
+    if isinstance(data, str):
+        data = data.encode()
+    # Calculate SHA-256 hash
+    sha256_hash = hashlib.sha256(data).hexdigest()
+    return sha256_hash
+
+
+# Socket.IO event handlers
+@socketio.on("connect")
+def handle_connect():
+    print("Client connected")
+    emit("status", {"message": "Connected to frontend-buffered demo server"})
+
+
+@socketio.on("disconnect")
+def handle_disconnect():
+    print("Client disconnected")
+
+
+@socketio.on("start_generation")
+def handle_start_generation(data):
+    global generation_active, frame_number, anim_name, frame_rate
+
+    frame_number = 0
+    if generation_active:
+        emit("error", {"message": "Generation already in progress"})
+        return
+
+    prompt = data.get("prompt", "")
+
+    seed = data.get("seed", -1)
+    if seed == -1:
+        seed = random.randint(0, 2**32)
+
+    # Extract words up to the first punctuation or newline
+    words_up_to_punctuation = re.split(r"[^\w\s]", prompt)[0].strip() if prompt else ""
+    if not words_up_to_punctuation:
+        words_up_to_punctuation = re.split(r"[\n\r]", prompt)[0].strip()
+
+    # Calculate SHA-256 hash of the entire prompt
+    sha256_hash = calculate_sha256(prompt)
+
+    # Create anim_name with the extracted words and first 10 characters of the hash
+    anim_name = f"{words_up_to_punctuation[:20]}_{str(seed)}_{sha256_hash[:10]}"
+
+    generation_active = True
+    generation_start_time = time.time()
+    enable_torch_compile = data.get("enable_torch_compile", False)
+    enable_fp8 = data.get("enable_fp8", False)
+    use_taehv = data.get("use_taehv", False)
+    frame_rate = data.get("fps", 6)
+
+    if not prompt:
+        emit("error", {"message": "Prompt is required"})
+        return
+
+    # Start generation in background thread
+    socketio.start_background_task(
+        generate_video_stream,
+        prompt,
+        seed,
+        enable_torch_compile,
+        enable_fp8,
+        use_taehv,
+    )
+    emit(
+        "status",
+        {"message": "Generation started - frames will be sent immediately"},
+    )
+
+
+@socketio.on("stop_generation")
+def handle_stop_generation():
+    global generation_active, stop_event, frame_send_queue
+    generation_active = False
+    stop_event.set()
+
+    # Signal sender thread to stop (will be processed after current frames)
+    try:
+        frame_send_queue.put(None)
+    except Exception as e:
+        print(f"❌ Failed to put None in frame_send_queue: {e}")
+
+    emit("status", {"message": "Generation stopped"})
+
+
+# Web routes
+
+
+@app.route("/")
+def index():
+    return render_template("demo.html")
+
+
+@app.route("/api/status")
+def api_status():
+    return jsonify(
+        {
+            "generation_active": generation_active,
+            "free_vram_gb": get_cuda_free_memory_gb(gpu),
+            "fp8_applied": fp8_applied,
+            "torch_compile_applied": torch_compile_applied,
+            "current_use_taehv": current_use_taehv,
+        }
+    )
+
+
+if __name__ == "__main__":
+    print(f"🚀 Starting demo on http://{args.host}:{args.port}")
+    socketio.run(app, host=args.host, port=args.port, debug=False)

demo_utils/constant.py

@@ -0,0 +1,39 @@
+import torch
+
+ZERO_VAE_CACHE = [
+    torch.zeros(1, 16, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 384, 2, 60, 104),
+    torch.zeros(1, 192, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 384, 2, 120, 208),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 192, 2, 240, 416),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+    torch.zeros(1, 96, 2, 480, 832),
+]
+
+feat_names = [f"vae_cache_{i}" for i in range(len(ZERO_VAE_CACHE))]
+ALL_INPUTS_NAMES = ["z", "use_cache"] + feat_names

demo_utils/memory.py

@@ -0,0 +1,174 @@
+# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils
+# Apache-2.0 License
+# By lllyasviel
+
+import torch
+
+cpu = torch.device("cpu")
+gpu = torch.device(f"cuda:{torch.cuda.current_device()}")
+gpu_complete_modules = []
+
+
+class DynamicSwapInstaller:
+    @staticmethod
+    def _install_module(
+        module: torch.nn.Module,
+        **kwargs,
+    ):
+        original_class = module.__class__
+        module.__dict__["forge_backup_original_class"] = original_class
+
+        def hacked_get_attr(
+            self,
+            name: str,
+        ):
+            if "_parameters" in self.__dict__:
+                _parameters = self.__dict__["_parameters"]
+                if name in _parameters:
+                    p = _parameters[name]
+                    if p is None:
+                        return None
+                    if p.__class__ == torch.nn.Parameter:
+                        return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
+                    else:
+                        return p.to(**kwargs)
+            if "_buffers" in self.__dict__:
+                _buffers = self.__dict__["_buffers"]
+                if name in _buffers:
+                    return _buffers[name].to(**kwargs)
+            return super(original_class, self).__getattr__(name)
+
+        module.__class__ = type(
+            "DynamicSwap_" + original_class.__name__,
+            (original_class,),
+            {
+                "__getattr__": hacked_get_attr,
+            },
+        )
+
+        return
+
+    @staticmethod
+    def _uninstall_module(
+        module: torch.nn.Module,
+    ):
+        if "forge_backup_original_class" in module.__dict__:
+            module.__class__ = module.__dict__.pop("forge_backup_original_class")
+        return
+
+    @staticmethod
+    def install_model(
+        model: torch.nn.Module,
+        **kwargs,
+    ):
+        for m in model.modules():
+            DynamicSwapInstaller._install_module(m, **kwargs)
+        return
+
+    @staticmethod
+    def uninstall_model(
+        model: torch.nn.Module,
+    ):
+        for m in model.modules():
+            DynamicSwapInstaller._uninstall_module(m)
+        return
+
+
+def fake_diffusers_current_device(
+    model: torch.nn.Module,
+    target_device: torch.device,
+):
+    if hasattr(model, "scale_shift_table"):
+        model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
+        return
+
+    for k, p in model.named_modules():
+        if hasattr(p, "weight"):
+            p.to(target_device)
+            return
+
+
+def get_cuda_free_memory_gb(
+    device=None,
+):
+    if device is None:
+        device = gpu
+
+    memory_stats = torch.cuda.memory_stats(device)
+    bytes_active = memory_stats["active_bytes.all.current"]
+    bytes_reserved = memory_stats["reserved_bytes.all.current"]
+    bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
+    bytes_inactive_reserved = bytes_reserved - bytes_active
+    bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
+    return bytes_total_available / (1024**3)
+
+
+def move_model_to_device_with_memory_preservation(
+    model,
+    target_device,
+    preserved_memory_gb=0,
+):
+    print(
+        f"Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB"
+    )
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, "weight"):
+            m.to(device=target_device)
+
+    model.to(device=target_device)
+    torch.cuda.empty_cache()
+    return
+
+
+def offload_model_from_device_for_memory_preservation(
+    model,
+    target_device,
+    preserved_memory_gb=0,
+):
+    print(
+        f"Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB"
+    )
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, "weight"):
+            m.to(device=cpu)
+
+    model.to(device=cpu)
+    torch.cuda.empty_cache()
+    return
+
+
+def unload_complete_models(
+    *args,
+):
+    for m in gpu_complete_modules + list(args):
+        m.to(device=cpu)
+        print(f"Unloaded {m.__class__.__name__} as complete.")
+
+    gpu_complete_modules.clear()
+    torch.cuda.empty_cache()
+    return
+
+
+def load_model_as_complete(
+    model,
+    target_device,
+    unload=True,
+):
+    if unload:
+        unload_complete_models()
+
+    model.to(device=target_device)
+    print(f"Loaded {model.__class__.__name__} to {target_device} as complete.")
+
+    gpu_complete_modules.append(model)
+    return

demo_utils/taehv.py

@@ -0,0 +1,477 @@
+#!/usr/bin/env python3
+"""
+Tiny AutoEncoder for Hunyuan Video
+(DNN for encoding / decoding videos to Hunyuan Video's latent space)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm.auto import tqdm
+from collections import namedtuple
+
+DecoderResult = namedtuple("DecoderResult", ("frame", "memory"))
+TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
+
+
+def conv(
+    n_in,
+    n_out,
+    **kwargs,
+):
+    return nn.Conv2d(
+        n_in,
+        n_out,
+        3,
+        padding=1,
+        **kwargs,
+    )
+
+
+class Clamp(nn.Module):
+    def forward(
+        self,
+        x,
+    ):
+        return torch.tanh(x / 3) * 3
+
+
+class MemBlock(nn.Module):
+    def __init__(
+        self,
+        n_in,
+        n_out,
+    ):
+        super().__init__()
+        self.conv = nn.Sequential(
+            conv(n_in * 2, n_out),
+            nn.ReLU(inplace=True),
+            conv(n_out, n_out),
+            nn.ReLU(inplace=True),
+            conv(n_out, n_out),
+        )
+        self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(
+        self,
+        x,
+        past,
+    ):
+        return self.act(
+            self.conv(
+                torch.cat([x, past], 1),
+            )
+            + self.skip(x),
+        )
+
+
+class TPool(nn.Module):
+    def __init__(
+        self,
+        n_f,
+        stride,
+    ):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(
+            n_f * stride,
+            n_f,
+            1,
+            bias=False,
+        )
+
+    def forward(
+        self,
+        x,
+    ):
+        _NT, C, H, W = x.shape
+        return self.conv(
+            x.reshape(-1, self.stride * C, H, W),
+        )
+
+
+class TGrow(nn.Module):
+    def __init__(
+        self,
+        n_f,
+        stride,
+    ):
+        super().__init__()
+        self.stride = stride
+        self.conv = nn.Conv2d(
+            n_f,
+            n_f * stride,
+            1,
+            bias=False,
+        )
+
+    def forward(
+        self,
+        x,
+    ):
+        _NT, C, H, W = x.shape
+        x = self.conv(x)
+        return x.reshape(-1, C, H, W)
+
+
+def apply_model_with_memblocks(
+    model,
+    x,
+    parallel,
+    show_progress_bar,
+):
+    """
+    Apply a sequential model with memblocks to the given input.
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - parallel: if True, parallelize over timesteps (fast but uses O(T) memory)
+        if False, each timestep will be processed sequentially (slow but uses O(1) memory)
+    - show_progress_bar: if True, enables tqdm progressbar display
+
+    Returns NTCHW tensor of output data.
+    """
+    assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
+    N, T, C, H, W = x.shape
+    if parallel:
+        x = x.reshape(N * T, C, H, W)
+        # parallel over input timesteps, iterate over blocks
+        for b in tqdm(model, disable=not show_progress_bar):
+            if isinstance(b, MemBlock):
+                NT, C, H, W = x.shape
+                T = NT // N
+                _x = x.reshape(N, T, C, H, W)
+                mem = F.pad(_x, (0, 0, 0, 0, 0, 0, 1, 0), value=0)[:, :T].reshape(x.shape)
+                x = b(x, mem)
+            else:
+                x = b(x)
+        NT, C, H, W = x.shape
+        T = NT // N
+        x = x.view(N, T, C, H, W)
+    else:
+        # TODO(oboerbohan): at least on macos this still gradually uses more memory during decode...
+        # need to fix :(
+        out = []
+        # iterate over input timesteps and also iterate over blocks.
+        # because of the cursed TPool/TGrow blocks, this is not a nested loop,
+        # it's actually a ***graph traversal*** problem! so let's make a queue
+        work_queue = [
+            TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))
+        ]
+        # in addition to manually managing our queue, we also need to manually manage our progressbar.
+        # we'll update it for every source node that we consume.
+        progress_bar = tqdm(range(T), disable=not show_progress_bar)
+        # we'll also need a separate addressable memory per node as well
+        mem = [None] * len(model)
+        while work_queue:
+            xt, i = work_queue.pop(0)
+            if i == 0:
+                # new source node consumed
+                progress_bar.update(1)
+            if i == len(model):
+                # reached end of the graph, append result to output list
+                out.append(xt)
+            else:
+                # fetch the block to process
+                b = model[i]
+                if isinstance(b, MemBlock):
+                    # mem blocks are simple since we're visiting the graph in causal order
+                    if mem[i] is None:
+                        xt_new = b(xt, xt * 0)
+                        mem[i] = xt
+                    else:
+                        xt_new = b(xt, mem[i])
+                        mem[i].copy_(
+                            xt
+                        )  # inplace might reduce mysterious pytorch memory allocations? doesn't help though
+                    # add successor to work queue
+                    work_queue.insert(0, TWorkItem(xt_new, i + 1))
+                elif isinstance(b, TPool):
+                    # pool blocks are miserable
+                    if mem[i] is None:
+                        mem[i] = []  # pool memory is itself a queue of inputs to pool
+                    mem[i].append(xt)
+                    if len(mem[i]) > b.stride:
+                        # pool mem is in invalid state, we should have pooled before this
+                        raise ValueError("???")
+                    elif len(mem[i]) < b.stride:
+                        # pool mem is not yet full, go back to processing the work queue
+                        pass
+                    else:
+                        # pool mem is ready, run the pool block
+                        N, C, H, W = xt.shape
+                        xt = b(torch.cat(mem[i], 1).view(N * b.stride, C, H, W))
+                        # reset the pool mem
+                        mem[i] = []
+                        # add successor to work queue
+                        work_queue.insert(0, TWorkItem(xt, i + 1))
+                elif isinstance(b, TGrow):
+                    xt = b(xt)
+                    NT, C, H, W = xt.shape
+                    # each tgrow has multiple successor nodes
+                    for xt_next in reversed(xt.view(N, b.stride * C, H, W).chunk(b.stride, 1)):
+                        # add successor to work queue
+                        work_queue.insert(0, TWorkItem(xt_next, i + 1))
+                else:
+                    # normal block with no funny business
+                    xt = b(xt)
+                    # add successor to work queue
+                    work_queue.insert(0, TWorkItem(xt, i + 1))
+        progress_bar.close()
+        x = torch.stack(out, 1)
+    return x
+
+
+class TAEHV(nn.Module):
+    latent_channels = 16
+    image_channels = 3
+
+    def __init__(
+        self,
+        checkpoint_path="taehv.pth",
+        decoder_time_upscale=(True, True),
+        decoder_space_upscale=(True, True, True),
+    ):
+        """Initialize pretrained TAEHV from the given checkpoint.
+
+        Arg:
+            checkpoint_path: path to weight file to load. taehv.pth for Hunyuan, taew2_1.pth for Wan 2.1.
+            decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+            decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+        """
+        super().__init__()
+        self.encoder = nn.Sequential(
+            conv(TAEHV.image_channels, 64),
+            nn.ReLU(inplace=True),
+            TPool(64, 2),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            TPool(64, 2),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            TPool(64, 1),
+            conv(64, 64, stride=2, bias=False),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            MemBlock(64, 64),
+            conv(64, TAEHV.latent_channels),
+        )
+        n_f = [256, 128, 64, 64]
+        self.frames_to_trim = 2 ** sum(decoder_time_upscale) - 1
+        self.decoder = nn.Sequential(
+            Clamp(),
+            conv(TAEHV.latent_channels, n_f[0]),
+            nn.ReLU(inplace=True),
+            MemBlock(n_f[0], n_f[0]),
+            MemBlock(n_f[0], n_f[0]),
+            MemBlock(n_f[0], n_f[0]),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1),
+            TGrow(n_f[0], 1),
+            conv(n_f[0], n_f[1], bias=False),
+            MemBlock(n_f[1], n_f[1]),
+            MemBlock(n_f[1], n_f[1]),
+            MemBlock(n_f[1], n_f[1]),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1),
+            TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1),
+            conv(n_f[1], n_f[2], bias=False),
+            MemBlock(n_f[2], n_f[2]),
+            MemBlock(n_f[2], n_f[2]),
+            MemBlock(n_f[2], n_f[2]),
+            nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1),
+            TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1),
+            conv(n_f[2], n_f[3], bias=False),
+            nn.ReLU(inplace=True),
+            conv(n_f[3], TAEHV.image_channels),
+        )
+        if checkpoint_path is not None:
+            self.load_state_dict(
+                self.patch_tgrow_layers(
+                    torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+                )
+            )
+
+    def patch_tgrow_layers(
+        self,
+        sd,
+    ):
+        """Patch TGrow layers to use a smaller kernel if needed.
+
+        Args:
+            sd: state dict to patch
+        """
+        new_sd = self.state_dict()
+        for i, layer in enumerate(self.decoder):
+            if isinstance(layer, TGrow):
+                key = f"decoder.{i}.conv.weight"
+                if sd[key].shape[0] > new_sd[key].shape[0]:
+                    # take the last-timestep output channels
+                    sd[key] = sd[key][-new_sd[key].shape[0] :]
+        return sd
+
+    def encode_video(
+        self,
+        x,
+        parallel=True,
+        show_progress_bar=True,
+    ):
+        """Encode a sequence of frames.
+
+        Args:
+            x: input NTCHW RGB (C=3) tensor with values in [0, 1].
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW latent tensor with ~Gaussian values.
+        """
+        return apply_model_with_memblocks(
+            self.encoder,
+            x,
+            parallel,
+            show_progress_bar,
+        )
+
+    def decode_video(
+        self,
+        x,
+        parallel=True,
+        show_progress_bar=False,
+    ):
+        """Decode a sequence of frames.
+
+        Args:
+            x: input NTCHW latent (C=12) tensor with ~Gaussian values.
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW RGB tensor with ~[0, 1] values.
+        """
+        x = apply_model_with_memblocks(
+            self.decoder,
+            x,
+            parallel,
+            show_progress_bar,
+        )
+        # return x[:, self.frames_to_trim:]
+        return x
+
+    def forward(
+        self,
+        x,
+    ):
+        return self.c(x)
+
+
+@torch.no_grad()
+def main():
+    """Run TAEHV roundtrip reconstruction on the given video paths."""
+    import os
+    import sys
+    import cv2  # no highly esteemed deed is commemorated here
+
+    class VideoTensorReader:
+        def __init__(
+            self,
+            video_file_path,
+        ):
+            self.cap = cv2.VideoCapture(video_file_path)
+            assert self.cap.isOpened(), f"Could not load {video_file_path}"
+            self.fps = self.cap.get(cv2.CAP_PROP_FPS)
+
+        def __iter__(
+            self,
+        ):
+            return self
+
+        def __next__(
+            self,
+        ):
+            ret, frame = self.cap.read()
+            if not ret:
+                self.cap.release()
+                raise StopIteration  # End of video or error
+            return torch.from_numpy(
+                cv2.cvtColor(frame, cv2.COLOR_BGR2RGB),
+            ).permute(
+                2, 0, 1
+            )  # BGR HWC -> RGB CHW
+
+    class VideoTensorWriter:
+        def __init__(
+            self,
+            video_file_path,
+            width_height,
+            fps=30,
+        ):
+            self.writer = cv2.VideoWriter(
+                video_file_path,
+                cv2.VideoWriter_fourcc(*"mp4v"),
+                fps,
+                width_height,
+            )
+            assert self.writer.isOpened(), f"Could not create writer for {video_file_path}"
+
+        def write(
+            self,
+            frame_tensor,
+        ):
+            assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??"
+            self.writer.write(
+                cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)
+            )  # RGB CHW -> BGR HWC
+
+        def __del__(
+            self,
+        ):
+            if hasattr(self, "writer"):
+                self.writer.release()
+
+    dev = torch.device(
+        "cuda"
+        if torch.cuda.is_available()
+        else "mps" if torch.backends.mps.is_available() else "cpu"
+    )
+    dtype = torch.float16
+    checkpoint_path = os.getenv("TAEHV_CHECKPOINT_PATH", "taehv.pth")
+    checkpoint_name = os.path.splitext(os.path.basename(checkpoint_path))[0]
+    print(
+        f"Using device \033[31m{dev}\033[0m, dtype \033[32m{dtype}\033[0m, checkpoint \033[34m{checkpoint_name}\033[0m ({checkpoint_path})"
+    )
+    taehv = TAEHV(checkpoint_path=checkpoint_path).to(dev, dtype)
+    for video_path in sys.argv[1:]:
+        print(f"Processing {video_path}...")
+        video_in = VideoTensorReader(video_path)
+        video = torch.stack(list(video_in), 0)[None]
+        vid_dev = video.to(dev, dtype).div_(255.0)
+        # convert to device tensor
+        if video.numel() < 100_000_000:
+            print(f"  {video_path} seems small enough, will process all frames in parallel")
+            # convert to device tensor
+            vid_enc = taehv.encode_video(vid_dev)
+            print(f"  Encoded {video_path} -> {vid_enc.shape}. Decoding...")
+            vid_dec = taehv.decode_video(vid_enc)
+            print(f"  Decoded {video_path} -> {vid_dec.shape}")
+        else:
+            print(f"  {video_path} seems large, will process each frame sequentially")
+            # convert to device tensor
+            vid_enc = taehv.encode_video(vid_dev, parallel=False)
+            print(f"  Encoded {video_path} -> {vid_enc.shape}. Decoding...")
+            vid_dec = taehv.decode_video(vid_enc, parallel=False)
+            print(f"  Decoded {video_path} -> {vid_dec.shape}")
+        video_out_path = video_path + f".reconstructed_by_{checkpoint_name}.mp4"
+        video_out = VideoTensorWriter(
+            video_out_path,
+            (vid_dec.shape[-1], vid_dec.shape[-2]),
+            fps=int(round(video_in.fps)),
+        )
+        for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]:
+            video_out.write(frame)
+        print(f"  Saved to {video_out_path}")
+
+
+if __name__ == "__main__":
+    main()

demo_utils/utils.py

@@ -0,0 +1,809 @@
+# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils
+# Apache-2.0 License
+# By lllyasviel
+
+import os
+import cv2
+import json
+import random
+import glob
+import torch
+import einops
+import numpy as np
+import datetime
+import torchvision
+
+from PIL import Image
+
+
+def min_resize(
+    x,
+    m,
+):
+    if x.shape[0] < x.shape[1]:
+        s0 = m
+        s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
+    else:
+        s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
+        s1 = m
+    new_max = max(s1, s0)
+    raw_max = max(x.shape[0], x.shape[1])
+    if new_max < raw_max:
+        interpolation = cv2.INTER_AREA
+    else:
+        interpolation = cv2.INTER_LANCZOS4
+    y = cv2.resize(x, (s1, s0), interpolation=interpolation)
+    return y
+
+
+def d_resize(
+    x,
+    y,
+):
+    H, W, C = y.shape
+    new_min = min(H, W)
+    raw_min = min(x.shape[0], x.shape[1])
+    if new_min < raw_min:
+        interpolation = cv2.INTER_AREA
+    else:
+        interpolation = cv2.INTER_LANCZOS4
+    y = cv2.resize(x, (W, H), interpolation=interpolation)
+    return y
+
+
+def resize_and_center_crop(
+    image,
+    target_width,
+    target_height,
+):
+    if target_height == image.shape[0] and target_width == image.shape[1]:
+        return image
+
+    pil_image = Image.fromarray(image)
+    original_width, original_height = pil_image.size
+    scale_factor = max(target_width / original_width, target_height / original_height)
+    resized_width = int(round(original_width * scale_factor))
+    resized_height = int(round(original_height * scale_factor))
+    resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
+    left = (resized_width - target_width) / 2
+    top = (resized_height - target_height) / 2
+    right = (resized_width + target_width) / 2
+    bottom = (resized_height + target_height) / 2
+    cropped_image = resized_image.crop((left, top, right, bottom))
+    return np.array(cropped_image)
+
+
+def resize_and_center_crop_pytorch(
+    image,
+    target_width,
+    target_height,
+):
+    B, C, H, W = image.shape
+
+    if H == target_height and W == target_width:
+        return image
+
+    scale_factor = max(target_width / W, target_height / H)
+    resized_width = int(round(W * scale_factor))
+    resized_height = int(round(H * scale_factor))
+
+    resized = torch.nn.functional.interpolate(
+        image,
+        size=(resized_height, resized_width),
+        mode="bilinear",
+        align_corners=False,
+    )
+
+    top = (resized_height - target_height) // 2
+    left = (resized_width - target_width) // 2
+    cropped = resized[:, :, top : top + target_height, left : left + target_width]
+
+    return cropped
+
+
+def resize_without_crop(
+    image,
+    target_width,
+    target_height,
+):
+    if target_height == image.shape[0] and target_width == image.shape[1]:
+        return image
+
+    pil_image = Image.fromarray(image)
+    resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
+    return np.array(resized_image)
+
+
+def just_crop(
+    image,
+    w,
+    h,
+):
+    if h == image.shape[0] and w == image.shape[1]:
+        return image
+
+    original_height, original_width = image.shape[:2]
+    k = min(original_height / h, original_width / w)
+    new_width = int(round(w * k))
+    new_height = int(round(h * k))
+    x_start = (original_width - new_width) // 2
+    y_start = (original_height - new_height) // 2
+    cropped_image = image[y_start : y_start + new_height, x_start : x_start + new_width]
+    return cropped_image
+
+
+def write_to_json(
+    data,
+    file_path,
+):
+    temp_file_path = file_path + ".tmp"
+    with open(temp_file_path, "wt", encoding="utf-8") as temp_file:
+        json.dump(data, temp_file, indent=4)
+    os.replace(temp_file_path, file_path)
+    return
+
+
+def read_from_json(
+    file_path,
+):
+    with open(file_path, "rt", encoding="utf-8") as file:
+        data = json.load(file)
+    return data
+
+
+def get_active_parameters(
+    m,
+):
+    return {k: v for k, v in m.named_parameters() if v.requires_grad}
+
+
+def cast_training_params(
+    m,
+    dtype=torch.float32,
+):
+    result = {}
+    for n, param in m.named_parameters():
+        if param.requires_grad:
+            param.data = param.to(dtype)
+            result[n] = param
+    return result
+
+
+def separate_lora_AB(
+    parameters,
+    B_patterns=None,
+):
+    parameters_normal = {}
+    parameters_B = {}
+
+    if B_patterns is None:
+        B_patterns = [".lora_B.", "__zero__"]
+
+    for k, v in parameters.items():
+        if any(B_pattern in k for B_pattern in B_patterns):
+            parameters_B[k] = v
+        else:
+            parameters_normal[k] = v
+
+    return parameters_normal, parameters_B
+
+
+def set_attr_recursive(
+    obj,
+    attr,
+    value,
+):
+    attrs = attr.split(".")
+    for name in attrs[:-1]:
+        obj = getattr(obj, name)
+    setattr(obj, attrs[-1], value)
+    return
+
+
+def print_tensor_list_size(
+    tensors,
+):
+    total_size = 0
+    total_elements = 0
+
+    if isinstance(tensors, dict):
+        tensors = tensors.values()
+
+    for tensor in tensors:
+        total_size += tensor.nelement() * tensor.element_size()
+        total_elements += tensor.nelement()
+
+    total_size_MB = total_size / (1024**2)
+    total_elements_B = total_elements / 1e9
+
+    print(f"Total number of tensors: {len(tensors)}")
+    print(f"Total size of tensors: {total_size_MB:.2f} MB")
+    print(f"Total number of parameters: {total_elements_B:.3f} billion")
+    return
+
+
+@torch.no_grad()
+def batch_mixture(
+    a,
+    b=None,
+    probability_a=0.5,
+    mask_a=None,
+):
+    batch_size = a.size(0)
+
+    if b is None:
+        b = torch.zeros_like(a)
+
+    if mask_a is None:
+        mask_a = torch.rand(batch_size) < probability_a
+
+    mask_a = mask_a.to(a.device)
+    mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1))
+    result = torch.where(mask_a, a, b)
+    return result
+
+
+@torch.no_grad()
+def zero_module(
+    module,
+):
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+@torch.no_grad()
+def supress_lower_channels(
+    m,
+    k,
+    alpha=0.01,
+):
+    data = m.weight.data.clone()
+
+    assert int(data.shape[1]) >= k
+
+    data[:, :k] = data[:, :k] * alpha
+    m.weight.data = data.contiguous().clone()
+    return m
+
+
+def freeze_module(
+    m,
+):
+    if not hasattr(m, "_forward_inside_frozen_module"):
+        m._forward_inside_frozen_module = m.forward
+    m.requires_grad_(False)
+    m.forward = torch.no_grad()(m.forward)
+    return m
+
+
+def get_latest_safetensors(
+    folder_path,
+):
+    safetensors_files = glob.glob(os.path.join(folder_path, "*.safetensors"))
+
+    if not safetensors_files:
+        raise ValueError("No file to resume!")
+
+    latest_file = max(safetensors_files, key=os.path.getmtime)
+    latest_file = os.path.abspath(os.path.realpath(latest_file))
+    return latest_file
+
+
+def generate_random_prompt_from_tags(
+    tags_str,
+    min_length=3,
+    max_length=32,
+):
+    tags = tags_str.split(", ")
+    tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags)))
+    prompt = ", ".join(tags)
+    return prompt
+
+
+def interpolate_numbers(
+    a,
+    b,
+    n,
+    round_to_int=False,
+    gamma=1.0,
+):
+    numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma)
+    if round_to_int:
+        numbers = np.round(numbers).astype(int)
+    return numbers.tolist()
+
+
+def uniform_random_by_intervals(
+    inclusive,
+    exclusive,
+    n,
+    round_to_int=False,
+):
+    edges = np.linspace(0, 1, n + 1)
+    points = np.random.uniform(edges[:-1], edges[1:])
+    numbers = inclusive + (exclusive - inclusive) * points
+    if round_to_int:
+        numbers = np.round(numbers).astype(int)
+    return numbers.tolist()
+
+
+def soft_append_bcthw(
+    history,
+    current,
+    overlap=0,
+):
+    if overlap <= 0:
+        return torch.cat([history, current], dim=2)
+
+    assert (
+        history.shape[2] >= overlap
+    ), f"History length ({history.shape[2]}) must be >= overlap ({overlap})"
+    assert (
+        current.shape[2] >= overlap
+    ), f"Current length ({current.shape[2]}) must be >= overlap ({overlap})"
+
+    weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(
+        1, 1, -1, 1, 1
+    )
+    blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap]
+    output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2)
+
+    return output.to(history)
+
+
+def save_bcthw_as_mp4(
+    x,
+    output_filename,
+    fps=10,
+    crf=0,
+):
+    b, c, t, h, w = x.shape
+
+    per_row = b
+    for p in [6, 5, 4, 3, 2]:
+        if b % p == 0:
+            per_row = p
+            break
+
+    os.makedirs(
+        os.path.dirname(os.path.abspath(os.path.realpath(output_filename))),
+        exist_ok=True,
+    )
+    x = torch.clamp(x.float(), -1.0, 1.0) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, "(m n) c t h w -> t (m h) (n w) c", n=per_row)
+    torchvision.io.write_video(
+        output_filename,
+        x,
+        fps=fps,
+        video_codec="libx264",
+        options={"crf": str(int(crf))},
+    )
+    return x
+
+
+def save_bcthw_as_png(
+    x,
+    output_filename,
+):
+    os.makedirs(
+        os.path.dirname(os.path.abspath(os.path.realpath(output_filename))),
+        exist_ok=True,
+    )
+    x = torch.clamp(x.float(), -1.0, 1.0) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, "b c t h w -> c (b h) (t w)")
+    torchvision.io.write_png(x, output_filename)
+    return output_filename
+
+
+def save_bchw_as_png(
+    x,
+    output_filename,
+):
+    os.makedirs(
+        os.path.dirname(os.path.abspath(os.path.realpath(output_filename))),
+        exist_ok=True,
+    )
+    x = torch.clamp(x.float(), -1.0, 1.0) * 127.5 + 127.5
+    x = x.detach().cpu().to(torch.uint8)
+    x = einops.rearrange(x, "b c h w -> c h (b w)")
+    torchvision.io.write_png(x, output_filename)
+    return output_filename
+
+
+def add_tensors_with_padding(
+    tensor1,
+    tensor2,
+):
+    if tensor1.shape == tensor2.shape:
+        return tensor1 + tensor2
+
+    shape1 = tensor1.shape
+    shape2 = tensor2.shape
+
+    new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2))
+
+    padded_tensor1 = torch.zeros(new_shape)
+    padded_tensor2 = torch.zeros(new_shape)
+
+    padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1
+    padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2
+
+    result = padded_tensor1 + padded_tensor2
+    return result
+
+
+def print_free_mem():
+    torch.cuda.empty_cache()
+    free_mem, total_mem = torch.cuda.mem_get_info(0)
+    free_mem_mb = free_mem / (1024**2)
+    total_mem_mb = total_mem / (1024**2)
+    print(f"Free memory: {free_mem_mb:.2f} MB")
+    print(f"Total memory: {total_mem_mb:.2f} MB")
+    return
+
+
+def print_gpu_parameters(
+    device,
+    state_dict,
+    log_count=1,
+):
+    summary = {"device": device, "keys_count": len(state_dict)}
+
+    logged_params = {}
+    for i, (key, tensor) in enumerate(state_dict.items()):
+        if i >= log_count:
+            break
+        logged_params[key] = tensor.flatten()[:3].tolist()
+
+    summary["params"] = logged_params
+
+    print(str(summary))
+    return
+
+
+def visualize_txt_as_img(
+    width,
+    height,
+    text,
+    font_path="font/DejaVuSans.ttf",
+    size=18,
+):
+    from PIL import Image, ImageDraw, ImageFont
+
+    txt = Image.new("RGB", (width, height), color="white")
+    draw = ImageDraw.Draw(txt)
+    font = ImageFont.truetype(font_path, size=size)
+
+    if text == "":
+        return np.array(txt)
+
+    # Split text into lines that fit within the image width
+    lines = []
+    words = text.split()
+    current_line = words[0]
+
+    for word in words[1:]:
+        line_with_word = f"{current_line} {word}"
+        if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width:
+            current_line = line_with_word
+        else:
+            lines.append(current_line)
+            current_line = word
+
+    lines.append(current_line)
+
+    # Draw the text line by line
+    y = 0
+    line_height = draw.textbbox((0, 0), "A", font=font)[3]
+
+    for line in lines:
+        if y + line_height > height:
+            break  # stop drawing if the next line will be outside the image
+        draw.text((0, y), line, fill="black", font=font)
+        y += line_height
+
+    return np.array(txt)
+
+
+def blue_mark(
+    x,
+):
+    x = x.copy()
+    c = x[:, :, 2]
+    b = cv2.blur(c, (9, 9))
+    x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1)
+    return x
+
+
+def green_mark(
+    x,
+):
+    x = x.copy()
+    x[:, :, 2] = -1
+    x[:, :, 0] = -1
+    return x
+
+
+def frame_mark(
+    x,
+):
+    x = x.copy()
+    x[:64] = -1
+    x[-64:] = -1
+    x[:, :8] = 1
+    x[:, -8:] = 1
+    return x
+
+
+@torch.inference_mode()
+def pytorch2numpy(
+    imgs,
+):
+    results = []
+    for x in imgs:
+        y = x.movedim(0, -1)
+        y = y * 127.5 + 127.5
+        y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
+        results.append(y)
+    return results
+
+
+@torch.inference_mode()
+def numpy2pytorch(
+    imgs,
+):
+    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0
+    h = h.movedim(-1, 1)
+    return h
+
+
+@torch.no_grad()
+def duplicate_prefix_to_suffix(
+    x,
+    count,
+    zero_out=False,
+):
+    if zero_out:
+        return torch.cat([x, torch.zeros_like(x[:count])], dim=0)
+    else:
+        return torch.cat([x, x[:count]], dim=0)
+
+
+def weighted_mse(a, b, weight):
+    return torch.mean(weight.float() * (a.float() - b.float()) ** 2)
+
+
+def clamped_linear_interpolation(
+    x,
+    x_min,
+    y_min,
+    x_max,
+    y_max,
+    sigma=1.0,
+):
+    x = (x - x_min) / (x_max - x_min)
+    x = max(0.0, min(x, 1.0))
+    x = x**sigma
+    return y_min + x * (y_max - y_min)
+
+
+def expand_to_dims(x, target_dims):
+    return x.view(*x.shape, *([1] * max(0, target_dims - x.dim())))
+
+
+def repeat_to_batch_size(
+    tensor: torch.Tensor,
+    batch_size: int,
+):
+    if tensor is None:
+        return None
+
+    first_dim = tensor.shape[0]
+
+    if first_dim == batch_size:
+        return tensor
+
+    if batch_size % first_dim != 0:
+        raise ValueError(
+            f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}."
+        )
+
+    repeat_times = batch_size // first_dim
+
+    return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1))
+
+
+def dim5(
+    x,
+):
+    return expand_to_dims(x, 5)
+
+
+def dim4(
+    x,
+):
+    return expand_to_dims(x, 4)
+
+
+def dim3(
+    x,
+):
+    return expand_to_dims(x, 3)
+
+
+def crop_or_pad_yield_mask(
+    x,
+    length,
+):
+    B, F, C = x.shape
+    device = x.device
+    dtype = x.dtype
+
+    if F < length:
+        y = torch.zeros((B, length, C), dtype=dtype, device=device)
+        mask = torch.zeros((B, length), dtype=torch.bool, device=device)
+        y[:, :F, :] = x
+        mask[:, :F] = True
+        return y, mask
+
+    return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device)
+
+
+def extend_dim(
+    x,
+    dim,
+    minimal_length,
+    zero_pad=False,
+):
+    original_length = int(x.shape[dim])
+
+    if original_length >= minimal_length:
+        return x
+
+    if zero_pad:
+        padding_shape = list(x.shape)
+        padding_shape[dim] = minimal_length - original_length
+        padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device)
+    else:
+        idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1)
+        last_element = x[idx]
+        padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim)
+
+    return torch.cat([x, padding], dim=dim)
+
+
+def lazy_positional_encoding(
+    t,
+    repeats=None,
+):
+    if not isinstance(t, list):
+        t = [t]
+
+    from diffusers.models.embeddings import get_timestep_embedding
+
+    te = torch.tensor(t)
+    te = get_timestep_embedding(
+        timesteps=te,
+        embedding_dim=256,
+        flip_sin_to_cos=True,
+        downscale_freq_shift=0.0,
+        scale=1.0,
+    )
+
+    if repeats is None:
+        return te
+
+    te = te[:, None, :].expand(-1, repeats, -1)
+
+    return te
+
+
+def state_dict_offset_merge(
+    A,
+    B,
+    C=None,
+):
+    result = {}
+    keys = A.keys()
+
+    for key in keys:
+        A_value = A[key]
+        B_value = B[key].to(A_value)
+
+        if C is None:
+            result[key] = A_value + B_value
+        else:
+            C_value = C[key].to(A_value)
+            result[key] = A_value + B_value - C_value
+
+    return result
+
+
+def state_dict_weighted_merge(
+    state_dicts,
+    weights,
+):
+    if len(state_dicts) != len(weights):
+        raise ValueError("Number of state dictionaries must match number of weights")
+
+    if not state_dicts:
+        return {}
+
+    total_weight = sum(weights)
+
+    if total_weight == 0:
+        raise ValueError("Sum of weights cannot be zero")
+
+    normalized_weights = [w / total_weight for w in weights]
+
+    keys = state_dicts[0].keys()
+    result = {}
+
+    for key in keys:
+        result[key] = state_dicts[0][key] * normalized_weights[0]
+
+        for i in range(1, len(state_dicts)):
+            state_dict_value = state_dicts[i][key].to(result[key])
+            result[key] += state_dict_value * normalized_weights[i]
+
+    return result
+
+
+def group_files_by_folder(
+    all_files,
+):
+    grouped_files = {}
+
+    for file in all_files:
+        folder_name = os.path.basename(os.path.dirname(file))
+        if folder_name not in grouped_files:
+            grouped_files[folder_name] = []
+        grouped_files[folder_name].append(file)
+
+    list_of_lists = list(grouped_files.values())
+    return list_of_lists
+
+
+def generate_timestamp():
+    now = datetime.datetime.now()
+    timestamp = now.strftime("%y%m%d_%H%M%S")
+    milliseconds = f"{int(now.microsecond / 1000):03d}"
+    random_number = random.randint(0, 9999)
+    return f"{timestamp}_{milliseconds}_{random_number}"
+
+
+def write_PIL_image_with_png_info(
+    image,
+    metadata,
+    path,
+):
+    from PIL.PngImagePlugin import PngInfo
+
+    png_info = PngInfo()
+    for key, value in metadata.items():
+        png_info.add_text(key, value)
+
+    image.save(path, "PNG", pnginfo=png_info)
+    return image
+
+
+def torch_safe_save(
+    content,
+    path,
+):
+    torch.save(content, path + "_tmp")
+    os.replace(path + "_tmp", path)
+    return path
+
+
+def move_optimizer_to_device(
+    optimizer,
+    device,
+):
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if isinstance(v, torch.Tensor):
+                state[k] = v.to(device)

demo_utils/vae.py

@@ -0,0 +1,477 @@
+from typing import List
+from einops import rearrange
+import tensorrt as trt
+import torch
+import torch.nn as nn
+
+from demo_utils.constant import ALL_INPUTS_NAMES, ZERO_VAE_CACHE
+from wan.modules.vae import AttentionBlock, CausalConv3d, RMS_norm, Upsample
+
+CACHE_T = 2
+
+
+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_1,
+        feat_cache_2,
+    ):
+        h = self.shortcut(x)
+        feat_cache = feat_cache_1
+        out_feat_cache = []
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d):
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat(
+                        [
+                            feat_cache[:, :, -1, :, :].unsqueeze(2).to(cache_x.device),
+                            cache_x,
+                        ],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache)
+                out_feat_cache.append(cache_x)
+                feat_cache = feat_cache_2
+            else:
+                x = layer(x)
+        return x + h, *out_feat_cache
+
+
+class Resample(nn.Module):
+
+    def __init__(
+        self,
+        dim,
+        mode,
+    ):
+        assert mode in ("none", "upsample2d", "upsample3d")
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+            self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(
+        self,
+        x,
+        is_first_frame,
+        feat_cache,
+    ):
+        if self.mode == "upsample3d":
+            b, c, t, h, w = x.size()
+            # x, out_feat_cache = torch.cond(
+            #     is_first_frame,
+            #     lambda: (torch.cat([torch.zeros_like(x), x], dim=2), feat_cache.clone()),
+            #     lambda: self.temporal_conv(x, feat_cache),
+            # )
+            # x, out_feat_cache = torch.cond(
+            #     is_first_frame,
+            #     lambda: (torch.cat([torch.zeros_like(x), x], dim=2), feat_cache.clone()),
+            #     lambda: self.temporal_conv(x, feat_cache),
+            # )
+            x, out_feat_cache = self.temporal_conv(x, is_first_frame, feat_cache)
+            out_feat_cache = torch.cond(
+                is_first_frame,
+                lambda: feat_cache.clone().contiguous(),
+                lambda: out_feat_cache.clone().contiguous(),
+            )
+            # if is_first_frame:
+            #     x = torch.cat([torch.zeros_like(x), x], dim=2)
+            #     out_feat_cache = feat_cache.clone()
+            # else:
+            #     x, out_feat_cache = self.temporal_conv(x, feat_cache)
+        else:
+            out_feat_cache = None
+        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)
+        return x, out_feat_cache
+
+    def temporal_conv(
+        self,
+        x,
+        is_first_frame,
+        feat_cache,
+    ):
+        b, c, t, h, w = x.size()
+        cache_x = x[:, :, -CACHE_T:, :, :].clone()
+        if cache_x.shape[2] < 2 and feat_cache is not None:
+            cache_x = torch.cat([torch.zeros_like(cache_x), cache_x], dim=2)
+        x = torch.cond(
+            is_first_frame,
+            lambda: torch.cat([torch.zeros_like(x), x], dim=1).contiguous(),
+            lambda: self.time_conv(x, feat_cache).contiguous(),
+        )
+        # x = self.time_conv(x, feat_cache)
+        out_feat_cache = cache_x
+
+        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)
+        return x.contiguous(), out_feat_cache.contiguous()
+
+    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 VAEDecoderWrapperSingle(nn.Module):
+    def __init__(
+        self,
+    ):
+        super().__init__()
+        self.decoder = VAEDecoder3d()
+        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=torch.float32)
+        self.std = torch.tensor(std, dtype=torch.float32)
+        self.z_dim = 16
+        self.conv2 = CausalConv3d(self.z_dim, self.z_dim, 1)
+
+    def forward(
+        self,
+        z: torch.Tensor,
+        is_first_frame: torch.Tensor,
+        *feat_cache: List[torch.Tensor],
+    ):
+        # from [batch_size, num_frames, num_channels, height, width]
+        # to [batch_size, num_channels, num_frames, height, width]
+        z = z.permute(0, 2, 1, 3, 4)
+        assert z.shape[2] == 1
+        feat_cache = list(feat_cache)
+        is_first_frame = is_first_frame.bool()
+
+        device, dtype = z.device, z.dtype
+        scale = [
+            self.mean.to(device=device, dtype=dtype),
+            1.0 / self.std.to(device=device, dtype=dtype),
+        ]
+
+        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]
+        x = self.conv2(z)
+        out, feat_cache = self.decoder(x, is_first_frame, feat_cache=feat_cache)
+        out = out.clamp_(-1, 1)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        out = out.permute(0, 2, 1, 3, 4)
+        return out, feat_cache
+
+
+class VAEDecoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=96,
+        z_dim=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[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_upsample = temperal_upsample
+        self.cache_t = 2
+        self.decoder_conv_num = 32
+
+        # 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: torch.Tensor,
+        is_first_frame: torch.Tensor,
+        feat_cache: List[torch.Tensor],
+    ):
+        idx = 0
+        out_feat_cache = []
+
+        # conv1
+        cache_x = x[:, :, -self.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])
+        out_feat_cache.append(cache_x)
+        idx += 1
+
+        # middle
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x, out_feat_cache_1, out_feat_cache_2 = layer(
+                    x, feat_cache[idx], feat_cache[idx + 1]
+                )
+                idx += 2
+                out_feat_cache.append(out_feat_cache_1)
+                out_feat_cache.append(out_feat_cache_2)
+            else:
+                x = layer(x)
+
+        # upsamples
+        for layer in self.upsamples:
+            if isinstance(layer, Resample):
+                x, cache_x = layer(x, is_first_frame, feat_cache[idx])
+                if cache_x is not None:
+                    out_feat_cache.append(cache_x)
+                    idx += 1
+            else:
+                x, out_feat_cache_1, out_feat_cache_2 = layer(
+                    x, feat_cache[idx], feat_cache[idx + 1]
+                )
+                idx += 2
+                out_feat_cache.append(out_feat_cache_1)
+                out_feat_cache.append(out_feat_cache_2)
+
+        # head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                cache_x = x[:, :, -self.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])
+                out_feat_cache.append(cache_x)
+                idx += 1
+            else:
+                x = layer(x)
+        return x, out_feat_cache
+
+
+class VAETRTWrapper:
+    def __init__(
+        self,
+    ):
+        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+        with open("checkpoints/vae_decoder_int8.trt", "rb") as f, trt.Runtime(TRT_LOGGER) as rt:
+            self.engine: trt.ICudaEngine = rt.deserialize_cuda_engine(f.read())
+
+        self.context: trt.IExecutionContext = self.engine.create_execution_context()
+        self.stream = torch.cuda.current_stream().cuda_stream
+
+        # ──────────────────────────────
+        # 2️⃣  Feed the engine with tensors
+        #     (name-based API in TRT ≥10)
+        # ──────────────────────────────
+        self.dtype_map = {
+            trt.float32: torch.float32,
+            trt.float16: torch.float16,
+            trt.int8: torch.int8,
+            trt.int32: torch.int32,
+        }
+        test_input = torch.zeros(1, 16, 1, 60, 104).cuda().half()
+        is_first_frame = torch.tensor(1.0).cuda().half()
+        test_cache_inputs = [c.cuda().half() for c in ZERO_VAE_CACHE]
+        test_inputs = [test_input, is_first_frame] + test_cache_inputs
+
+        # keep references so buffers stay alive
+        self.device_buffers, self.outputs = {}, []
+
+        # ---- inputs ----
+        for i, name in enumerate(ALL_INPUTS_NAMES):
+            tensor, scale = test_inputs[i], 1 / 127
+            tensor = self.quantize_if_needed(tensor, self.engine.get_tensor_dtype(name), scale)
+
+            # dynamic shapes
+            if -1 in self.engine.get_tensor_shape(name):
+                # new API :contentReference[oaicite:0]{index=0}
+                self.context.set_input_shape(name, tuple(tensor.shape))
+
+            # replaces bindings[] :contentReference[oaicite:1]{index=1}
+            self.context.set_tensor_address(name, int(tensor.data_ptr()))
+            self.device_buffers[name] = tensor  # keep pointer alive
+
+        # ---- (after all input shapes are known) infer output shapes ----
+        # propagates shapes :contentReference[oaicite:2]{index=2}
+        self.context.infer_shapes()
+
+        for i in range(self.engine.num_io_tensors):
+            name = self.engine.get_tensor_name(i)
+            # replaces binding_is_input :contentReference[oaicite:3]{index=3}
+            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:
+                shape = tuple(self.context.get_tensor_shape(name))
+                dtype = self.dtype_map[self.engine.get_tensor_dtype(name)]
+                out = torch.empty(shape, dtype=dtype, device="cuda").contiguous()
+
+                self.context.set_tensor_address(name, int(out.data_ptr()))
+                self.outputs.append(out)
+                self.device_buffers[name] = out
+
+    # helper to quant-convert on the fly
+    def quantize_if_needed(
+        self,
+        t,
+        expected_dtype,
+        scale,
+    ):
+        if expected_dtype == trt.int8 and t.dtype != torch.int8:
+            t = torch.clamp((t / scale).round(), -128, 127).to(torch.int8).contiguous()
+        return t  # keep pointer alive
+
+    def forward(
+        self,
+        *test_inputs,
+    ):
+        for i, name in enumerate(ALL_INPUTS_NAMES):
+            tensor, scale = test_inputs[i], 1 / 127
+            tensor = self.quantize_if_needed(tensor, self.engine.get_tensor_dtype(name), scale)
+            self.context.set_tensor_address(name, int(tensor.data_ptr()))
+            self.device_buffers[name] = tensor
+
+        self.context.execute_async_v3(stream_handle=self.stream)
+        torch.cuda.current_stream().synchronize()
+        return self.outputs

demo_utils/vae_block3.py

@@ -0,0 +1,364 @@
+from typing import List
+from einops import rearrange
+import torch
+import torch.nn as nn
+
+from wan.modules.vae import (
+    AttentionBlock,
+    CausalConv3d,
+    RMS_norm,
+    ResidualBlock,
+    Upsample,
+)
+
+
+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
+        self.cache_t = 2
+
+        # layers
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest"),
+                nn.Conv2d(dim, dim // 2, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest"),
+                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[:, :, -self.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 VAEDecoderWrapper(nn.Module):
+    def __init__(
+        self,
+    ):
+        super().__init__()
+        self.decoder = VAEDecoder3d()
+        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=torch.float32)
+        self.std = torch.tensor(std, dtype=torch.float32)
+        self.z_dim = 16
+        self.conv2 = CausalConv3d(self.z_dim, self.z_dim, 1)
+
+    def forward(
+        self,
+        z: torch.Tensor,
+        *feat_cache: List[torch.Tensor],
+    ):
+        # from [batch_size, num_frames, num_channels, height, width]
+        # to [batch_size, num_channels, num_frames, height, width]
+        z = z.permute(0, 2, 1, 3, 4)
+        feat_cache = list(feat_cache)
+        print("Length of feat_cache: ", len(feat_cache))
+
+        device, dtype = z.device, z.dtype
+        scale = [
+            self.mean.to(device=device, dtype=dtype),
+            1.0 / self.std.to(device=device, dtype=dtype),
+        ]
+
+        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_):
+            if i == 0:
+                out, feat_cache = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=feat_cache)
+            else:
+                out_, feat_cache = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=feat_cache)
+                out = torch.cat([out, out_], 2)
+
+        out = out.float().clamp_(-1, 1)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        out = out.permute(0, 2, 1, 3, 4)
+        return out, feat_cache
+
+
+class VAEDecoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=96,
+        z_dim=16,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_upsample=[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_upsample = temperal_upsample
+        self.cache_t = 2
+        self.decoder_conv_num = 32
+
+        # 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: torch.Tensor,
+        feat_cache: List[torch.Tensor],
+    ):
+        feat_idx = [0]
+
+        # conv1
+        idx = feat_idx[0]
+        cache_x = x[:, :, -self.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
+
+        # 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:
+            x = layer(x, feat_cache, feat_idx)
+
+        # head
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -self.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, feat_cache

demo_utils/vae_torch2trt.py

@@ -0,0 +1,335 @@
+# ---- INT8 (optional) ----
+from demo_utils.vae import (
+    VAEDecoderWrapperSingle,  # main nn.Module
+    ZERO_VAE_CACHE,  # helper constants shipped with your code base
+)
+import pycuda.driver as cuda  # ← add
+import pycuda.autoinit  # noqa
+
+import sys
+from pathlib import Path
+
+import torch
+import tensorrt as trt
+
+from utils.dataset import ShardingLMDBDataset
+
+data_path = "/mnt/localssd/wanx_14B_shift-3.0_cfg-5.0_lmdb_oneshard"
+dataset = ShardingLMDBDataset(data_path, max_pair=int(1e8))
+dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=0)
+
+# ─────────────────────────────────────────────────────────
+# 1️⃣  Bring the PyTorch model into scope
+#     (all code you pasted lives in `vae_decoder.py`)
+# ─────────────────────────────────────────────────────────
+
+# --- dummy tensors (exact shapes you posted) ---
+dummy_input = torch.randn(1, 1, 16, 60, 104).half().cuda()
+is_first_frame = torch.tensor([1.0], device="cuda", dtype=torch.float16)
+dummy_cache_input = [
+    torch.randn(*s.shape).half().cuda() if isinstance(s, torch.Tensor) else s
+    for s in ZERO_VAE_CACHE  # keep exactly the same ordering
+]
+inputs = [dummy_input, is_first_frame, *dummy_cache_input]
+
+# ─────────────────────────────────────────────────────────
+# 2️⃣  Export → ONNX
+# ─────────────────────────────────────────────────────────
+model = VAEDecoderWrapperSingle().half().cuda().eval()
+
+vae_state_dict = torch.load("wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth", map_location="cpu")
+decoder_state_dict = {}
+for key, value in vae_state_dict.items():
+    if "decoder." in key or "conv2" in key:
+        decoder_state_dict[key] = value
+model.load_state_dict(decoder_state_dict)
+model = model.half().cuda().eval()  # only batch dim dynamic
+
+onnx_path = Path("vae_decoder.onnx")
+feat_names = [f"vae_cache_{i}" for i in range(len(dummy_cache_input))]
+all_inputs_names = ["z", "use_cache"] + feat_names
+
+with torch.inference_mode():
+    torch.onnx.export(
+        model,
+        tuple(inputs),  # must be a tuple
+        onnx_path.as_posix(),
+        input_names=all_inputs_names,
+        output_names=["rgb_out", "cache_out"],
+        opset_version=17,
+        do_constant_folding=True,
+        dynamo=True,
+    )
+print(f"✅  ONNX graph saved to {onnx_path.resolve()}")
+
+# (Optional) quick sanity-check with ONNX-Runtime
+try:
+    import onnxruntime as ort
+
+    sess = ort.InferenceSession(onnx_path.as_posix(), providers=["CUDAExecutionProvider"])
+    ort_inputs = {n: t.cpu().numpy() for n, t in zip(all_inputs_names, inputs)}
+    _ = sess.run(None, ort_inputs)
+    print("✅  ONNX graph is executable")
+except Exception as e:
+    print("⚠️  ONNX check failed:", e)
+
+# ─────────────────────────────────────────────────────────
+# 3️⃣  Build the TensorRT engine
+# ─────────────────────────────────────────────────────────
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+builder = trt.Builder(TRT_LOGGER)
+network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
+parser = trt.OnnxParser(network, TRT_LOGGER)
+
+with open(onnx_path, "rb") as f:
+    if not parser.parse(f.read()):
+        for i in range(parser.num_errors):
+            print(parser.get_error(i))
+        sys.exit("❌  ONNX → TRT parsing failed")
+
+config = builder.create_builder_config()
+
+
+def set_workspace(
+    config,
+    bytes_,
+):
+    """Version-agnostic workspace limit."""
+    if hasattr(config, "max_workspace_size"):  # TRT 8 / 9
+        config.max_workspace_size = bytes_
+    else:  # TRT 10+
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, bytes_)
+
+
+# …
+config = builder.create_builder_config()
+set_workspace(config, 4 << 30)  # 4 GB
+# 4 GB
+
+if builder.platform_has_fast_fp16:
+    config.set_flag(trt.BuilderFlag.FP16)
+
+# ---- INT8 (optional) ----
+# provide a calibrator if you need an INT8 engine; comment this
+# block if you only care about FP16.
+# ─────────────────────────────────────────────────────────
+# helper: version-agnostic workspace limit
+# ─────────────────────────────────────────────────────────
+
+
+def set_workspace(
+    config: trt.IBuilderConfig,
+    bytes_: int = 4 << 30,
+):
+    """
+    TRT < 10.x  →  config.max_workspace_size
+    TRT ≥ 10.x  →  config.set_memory_pool_limit(...)
+    """
+    if hasattr(config, "max_workspace_size"):  # TRT 8 / 9
+        config.max_workspace_size = bytes_
+    else:  # TRT 10+
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, bytes_)
+
+
+# ─────────────────────────────────────────────────────────
+# (optional) INT-8 calibrator
+# ─────────────────────────────────────────────────────────
+# ‼ Only keep this block if you really need INT-8 ‼                      # gracefully skip if PyCUDA not present
+
+
+class VAECalibrator(trt.IInt8EntropyCalibrator2):
+    def __init__(
+        self,
+        loader,
+        cache="calibration.cache",
+        max_batches=10,
+    ):
+        super().__init__()
+        self.loader = iter(loader)
+        self.batch_size = loader.batch_size or 1
+        self.max_batches = max_batches
+        self.count = 0
+        self.cache_file = cache
+        self.stream = cuda.Stream()
+        self.dev_ptrs = {}
+
+    # --- TRT 10 needs BOTH spellings ---
+    def get_batch_size(
+        self,
+    ):
+        return self.batch_size
+
+    def getBatchSize(
+        self,
+    ):
+        return self.batch_size
+
+    def get_batch(
+        self,
+        names,
+    ):
+        if self.count >= self.max_batches:
+            return None
+
+        # Randomly sample a number from 1 to 10
+        import random
+
+        vae_idx = random.randint(0, 10)
+        data = next(self.loader)
+
+        latent = data["ode_latent"][0][:, :1]
+        is_first_frame = torch.tensor([1.0], device="cuda", dtype=torch.float16)
+        feat_cache = ZERO_VAE_CACHE
+        for i in range(vae_idx):
+            inputs = [latent, is_first_frame, *feat_cache]
+            with torch.inference_mode():
+                outputs = model(*inputs)
+            latent = data["ode_latent"][0][:, i + 1 : i + 2]
+            is_first_frame = torch.tensor([0.0], device="cuda", dtype=torch.float16)
+            feat_cache = outputs[1:]
+
+        # -------- ensure context is current --------
+        z_np = latent.cpu().numpy().astype("float32")
+
+        ptrs = []  # list[int] – one entry per name
+        for name in names:  # <-- match TRT's binding order
+            if name == "z":
+                arr = z_np
+            elif name == "use_cache":
+                arr = is_first_frame.cpu().numpy().astype("float32")
+            else:
+                idx = int(name.split("_")[-1])  # "vae_cache_17" -> 17
+                arr = feat_cache[idx].cpu().numpy().astype("float32")
+
+            if name not in self.dev_ptrs:
+                self.dev_ptrs[name] = cuda.mem_alloc(arr.nbytes)
+
+            cuda.memcpy_htod_async(self.dev_ptrs[name], arr, self.stream)
+            ptrs.append(int(self.dev_ptrs[name]))  # ***int() is required***
+
+        self.stream.synchronize()
+        self.count += 1
+        print(f"Calibration batch {self.count}/{self.max_batches}")
+        return ptrs
+
+    # --- calibration-cache helpers (both spellings) ---
+    def read_calibration_cache(
+        self,
+    ):
+        try:
+            with open(self.cache_file, "rb") as f:
+                return f.read()
+        except FileNotFoundError:
+            return None
+
+    def readCalibrationCache(
+        self,
+    ):
+        return self.read_calibration_cache()
+
+    def write_calibration_cache(
+        self,
+        cache,
+    ):
+        with open(self.cache_file, "wb") as f:
+            f.write(cache)
+
+    def writeCalibrationCache(
+        self,
+        cache,
+    ):
+        self.write_calibration_cache(cache)
+
+
+# ─────────────────────────────────────────────────────────
+# Builder-config + optimisation profile
+# ─────────────────────────────────────────────────────────
+config = builder.create_builder_config()
+set_workspace(config, 4 << 30)  # 4 GB
+
+# ► enable FP16 if possible
+if builder.platform_has_fast_fp16:
+    config.set_flag(trt.BuilderFlag.FP16)
+
+# ► enable INT-8  (delete this block if you don’t need it)
+if cuda is not None:
+    config.set_flag(trt.BuilderFlag.INT8)
+    # supply any representative batch you like – here we reuse the latent z
+    calib = VAECalibrator(dataloader)
+    # TRT-10 renamed the setter:
+    if hasattr(config, "set_int8_calibrator"):  # TRT 10+
+        config.set_int8_calibrator(calib)
+    else:  # TRT ≤ 9
+        config.int8_calibrator = calib
+
+# ---- optimisation profile ----
+profile = builder.create_optimization_profile()
+profile.set_shape(
+    all_inputs_names[0],  # latent z
+    min=(1, 1, 16, 60, 104),
+    opt=(1, 1, 16, 60, 104),
+    max=(1, 1, 16, 60, 104),
+)
+profile.set_shape("use_cache", min=(1,), opt=(1,), max=(1,))  # scalar flag
+for name, tensor in zip(all_inputs_names[2:], dummy_cache_input):
+    profile.set_shape(name, tensor.shape, tensor.shape, tensor.shape)
+
+config.add_optimization_profile(profile)
+
+# ─────────────────────────────────────────────────────────
+# Build the engine  (API changed in TRT-10)
+# ─────────────────────────────────────────────────────────
+print("⚙️  Building engine … (can take a minute)")
+
+if hasattr(builder, "build_serialized_network"):  # TRT 10+
+    serialized_engine = builder.build_serialized_network(network, config)
+    assert serialized_engine is not None, "build_serialized_network() failed"
+    plan_path = Path("checkpoints/vae_decoder_int8.trt")
+    plan_path.write_bytes(serialized_engine)
+    engine_bytes = serialized_engine  # keep for smoke-test
+else:  # TRT ≤ 9
+    engine = builder.build_engine(network, config)
+    assert engine is not None, "build_engine() returned None"
+    plan_path = Path("checkpoints/vae_decoder_int8.trt")
+    plan_path.write_bytes(engine.serialize())
+    engine_bytes = engine.serialize()
+
+print(f"✅  TensorRT engine written to {plan_path.resolve()}")
+
+# ─────────────────────────────────────────────────────────
+# 4️⃣  Quick smoke-test with the brand-new engine
+# ─────────────────────────────────────────────────────────
+with trt.Runtime(TRT_LOGGER) as rt:
+    engine = rt.deserialize_cuda_engine(engine_bytes)
+    context = engine.create_execution_context()
+    stream = torch.cuda.current_stream().cuda_stream
+
+    # pre-allocate device buffers once
+    device_buffers, outputs = {}, []
+    dtype_map = {
+        trt.float32: torch.float32,
+        trt.float16: torch.float16,
+        trt.int8: torch.int8,
+        trt.int32: torch.int32,
+    }
+
+    for name, tensor in zip(all_inputs_names, inputs):
+        if -1 in engine.get_tensor_shape(name):  # dynamic input
+            context.set_input_shape(name, tensor.shape)
+        context.set_tensor_address(name, int(tensor.data_ptr()))
+        device_buffers[name] = tensor
+
+    context.infer_shapes()  # propagate ⇢ outputs
+    for i in range(engine.num_io_tensors):
+        name = engine.get_tensor_name(i)
+        if engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:
+            shape = tuple(context.get_tensor_shape(name))
+            dtype = dtype_map[engine.get_tensor_dtype(name)]
+            out = torch.empty(shape, dtype=dtype, device="cuda")
+            context.set_tensor_address(name, int(out.data_ptr()))
+            outputs.append(out)
+            print(f"output {name} shape: {shape}")
+
+    context.execute_async_v3(stream_handle=stream)
+    torch.cuda.current_stream().synchronize()
+    print("✅  TRT execution OK – first output shape:", outputs[0].shape)

frequency_utils.py

@@ -0,0 +1,1020 @@
+import math
+from typing import List, Sequence, Tuple, Union
+import torch
+from torch import Tensor
+
+from PIL import Image
+import numpy as np
+import os
+from pathlib import Path
+
+torch.set_printoptions(
+    linewidth=10000,
+)
+
+
+def _get_center_distance(size: Tuple[int], device: str = "cpu") -> Tensor:
+    """Compute the distance of each matrix element to the center.
+
+    Args:
+        size (Tuple[int]): [m, n].
+        device (str, optional): cpu/cuda. Defaults to 'cpu'.
+
+    Returns:
+        Tensor: [m, n].
+    """
+    m, n = size
+    i_ind = torch.tile(
+        torch.tensor([[[i]] for i in range(m)], device=device), dims=[1, n, 1]
+    ).float()  # [m, n, 1]
+    j_ind = torch.tile(
+        torch.tensor([[[i] for i in range(n)]], device=device), dims=[m, 1, 1]
+    ).float()  # [m, n, 1]
+    ij_ind = torch.cat([i_ind, j_ind], dim=-1)  # [m, n, 2]
+    ij_ind = ij_ind.reshape([m * n, 1, 2])  # [m * n, 1, 2]
+    center_ij = torch.tensor(((m - 1) / 2, (n - 1) / 2), device=device).reshape(1, 2)
+    center_ij = torch.tile(center_ij, dims=[m * n, 1, 1])
+    dist = torch.cdist(ij_ind, center_ij, p=2).reshape([m, n])
+    return dist
+
+
+def _get_ideal_weights(
+    size: Tuple[int], D0: int, lowpass: bool = True, device: str = "cpu"
+) -> Tensor:
+    """Get H(u, v) of ideal bandpass filter.
+
+    Args:
+        size (Tuple[int]): [H, W].
+        D0 (int): The cutoff frequency.
+        lowpass (bool): True for low-pass filter, otherwise for high-pass filter. Defaults to True.
+        device (str, optional): cpu/cuda. Defaults to 'cpu'.
+
+    Returns:
+        Tensor: [H, W].
+    """
+    center_distance = _get_center_distance(size, device)
+    center_distance[center_distance > D0] = -1
+    center_distance[center_distance != -1] = 1
+    if lowpass is True:
+        center_distance[center_distance == -1] = 0
+    else:
+        center_distance[center_distance == 1] = 0
+        center_distance[center_distance == -1] = 1
+    return center_distance
+
+
+def _to_freq(image: Tensor) -> Tensor:
+    """Convert from spatial domain to frequency domain.
+
+    Args:
+        image (Tensor): [B, C, H, W].
+
+    Returns:
+        Tensor: [B, C, H, W]
+    """
+    img_fft = torch.fft.fft2(image)
+    img_fft_shift = torch.fft.fftshift(img_fft)
+    return img_fft_shift
+
+
+def _to_space(image_fft: Tensor) -> Tensor:
+    """Convert from frequency domain to spatial domain.
+
+    Args:
+        image_fft (Tensor): [B, C, H, W].
+
+    Returns:
+        Tensor: [B, C, H, W].
+    """
+    img_ifft_shift = torch.fft.ifftshift(image_fft)
+    img_ifft = torch.fft.ifft2(img_ifft_shift)
+    img = img_ifft.real.clamp(0, 1)
+    return img
+
+
+def ideal_bandpass(image: Tensor, D0: int, lowpass: bool = True) -> Tensor:
+    """Low-pass filter for images.
+
+    Args:
+        image (Tensor): [B, C, H, W].
+        D0 (int): Cutoff frequency.
+        lowpass (bool): True for low-pass filter, otherwise for high-pass filter. Defaults to True.
+
+    Returns:
+        Tensor: [B, C, H, W].
+    """
+    img_fft = _to_freq(image)
+    weights = _get_ideal_weights(img_fft.shape[-2:], D0=D0, lowpass=lowpass, device=image.device)
+    img_fft = img_fft * weights
+    img = _to_space(img_fft)
+    return img
+
+
+# Butterworth
+
+
+def _get_butterworth_weights(size: Tuple[int], D0: int, n: int, device: str = "cpu") -> Tensor:
+    """Get H(u, v) of Butterworth filter.
+
+    Args:
+        size (Tuple[int]): [H, W].
+        D0 (int): The cutoff frequency.
+        n (int): Order of Butterworth filters.
+        device (str, optional): cpu/cuda. Defaults to 'cpu'.
+
+    Returns:
+        Tensor: [H, W].
+    """
+    center_distance = _get_center_distance(size=size, device=device)
+    weights = 1 / (1 + torch.pow(center_distance / D0, 2 * n))
+    return weights
+
+
+def butterworth(image: Tensor, D0: int, n: int) -> Tensor:
+    """Butterworth low-pass filter for images.
+
+    Args:
+        image (Tensor): [B, C, H, W].
+        D0 (int): Cutoff frequency.
+        n (int): Order of the Butterworth low-pass filter.
+
+    Returns:
+        Tensor: [B, C, H, W].
+    """
+    img_fft = _to_freq(image)
+    weights = _get_butterworth_weights(image.shape[-2:], D0, n, device=image.device)
+    img_fft = weights * img_fft
+    img = _to_space(img_fft)
+    return img
+
+
+# def my_butterworth_low_pass_filter(
+#     shape,
+#     stop_freqs: List[float],
+#     n=4,
+# ):
+#     assert len(shape) == len(stop_freqs)
+
+#     grid = torch.meshgrid(
+#         *[torch.arange(s, dtype=torch.float32) for s in shape],
+#         indexing='ij',
+#     )
+#     # ( [shape[0], shape[1], ..., shape[N]] ) * len(shape)
+#     indices = torch.stack(grid, dim=-1).float()
+#     # print(f"{indices.shape = }")
+#     # [shape[0], shape[1], ..., shape[N], len(shape)]
+
+#     max_len = torch.tensor(shape).float()
+#     max_len -= 1.0
+#     max_len /= 2.0
+#     # print(f"{max_len = }")
+#     # print(f"{max_len.shape = }")
+#     # [len(shape)]
+#     max_len = max_len.view(*([1]*len(shape)), -1)
+#     # print(f"{max_len.shape = }")
+#     # [1, 1, ..., 1, len(shape)]
+
+#     normalized_indices = indices / max_len
+#     # [shape[0], shape[1], ..., shape[N], len(shape)]
+
+#     normalized_indices_offset = normalized_indices - 1
+#     # print(f"{normalized_indices_offset.shape = }")
+#     # [shape[0], shape[1], ..., shape[N], len(shape)]
+
+
+#     stop_freqs_torch = torch.tensor(stop_freqs).float().view(*([1]*len(shape)), -1)
+#     # print(f"{stop_freqs_torch.shape = }")
+#     # [1, 1, ..., 1, len(shape)]
+
+#     scaled_normalized_indices_offset = normalized_indices_offset / stop_freqs_torch
+#     # print(f"{scaled_normalized_indices_offset.shape = }")
+#     # [shape[0], shape[1], ..., shape[N], len(shape)]
+
+#     filter_ = 1.0 / (1.0 + torch.pow(scaled_normalized_indices_offset.norm(p=2, dim=-1), 2 * n))
+#     return filter_
+
+
+# def my_butterworth_low_pass_filter_non_center(
+#     shape,
+#     stop_freqs: List[float],
+#     n=4,
+# ):
+#     new_shape = [
+#         2*i-1
+#         for i in shape
+#     ]
+#     filter_ = my_butterworth_low_pass_filter(
+#         new_shape,
+#         n=n,
+#         stop_freqs=stop_freqs,
+#     )
+
+#     if len(shape) == 1:
+#         crop_filter = filter_[-shape[0]:]
+#     elif len(shape) == 2:
+#         crop_filter = filter_[-shape[0]:, -shape[1]:]
+#     elif len(shape) == 3:
+#         crop_filter = filter_[-shape[0]:, -shape[1]:, -shape[2]:]
+#     else:
+#         raise ValueError("Shape must be 1D, 2D, or 3D.")
+#     return crop_filter
+
+
+# def my_butterworth_high_pass_filter(
+#     shape,
+#     stop_freqs: List[float],
+#     n=4,
+# ):
+#     assert len(shape) == len(stop_freqs)
+
+#     grid = torch.meshgrid(
+#         *[torch.arange(s, dtype=torch.float32) for s in shape],
+#         indexing='ij',
+#     )
+#     # ( [shape[0], shape[1], ..., shape[N]] ) * len(shape)
+#     indices = torch.stack(grid, dim=-1).float()
+#     # print(f"{indices.shape = }")
+#     # [shape[0], shape[1], ..., shape[N], len(shape)]
+
+#     max_len = torch.tensor(shape).float()
+#     max_len -= 1.0
+#     max_len /= 2.0
+#     # print(f"{max_len = }")
+#     # print(f"{max_len.shape = }")
+#     # [len(shape)]
+#     max_len = max_len.view(*([1]*len(shape)), -1)
+#     # print(f"{max_len.shape = }")
+#     # [1, 1, ..., 1, len(shape)]
+
+#     normalized_indices = indices / max_len
+#     # [shape[0], shape[1], ..., shape[N], len(shape)]
+
+#     normalized_indices_offset = normalized_indices - 1
+#     # print(f"{normalized_indices_offset.shape = }")
+#     # [shape[0], shape[1], ..., shape[N], len(shape)]
+
+
+#     stop_freqs_torch = torch.tensor(stop_freqs).float().view(*([1]*len(shape)), -1)
+#     # print(f"{stop_freqs_torch.shape = }")
+#     # [1, 1, ..., 1, len(shape)]
+
+#     scaled_normalized_indices_offset = stop_freqs_torch / normalized_indices_offset
+#     # print(f"{scaled_normalized_indices_offset.shape = }")
+#     # [shape[0], shape[1], ..., shape[N], len(shape)]
+
+#     filter_ = 1.0 / (1.0 + torch.pow(scaled_normalized_indices_offset.norm(p=2, dim=-1), 2 * n))
+#     return filter_
+
+
+# def my_butterworth_high_pass_filter_non_center(
+#     shape,
+#     stop_freqs: List[float],
+#     n=4,
+# ):
+#     new_shape = [
+#         2*i-1
+#         for i in shape
+#     ]
+#     filter_ = my_butterworth_high_pass_filter(
+#         new_shape,
+#         n=n,
+#         stop_freqs=stop_freqs,
+#     )
+
+#     if len(shape) == 1:
+#         crop_filter = filter_[-shape[0]:]
+#     elif len(shape) == 2:
+#         crop_filter = filter_[-shape[0]:, -shape[1]:]
+#     elif len(shape) == 3:
+#         crop_filter = filter_[-shape[0]:, -shape[1]:, -shape[2]:]
+#     else:
+#         raise ValueError("Shape must be 1D, 2D, or 3D.")
+#     return crop_filter
+
+
+# ------------------------ Image loading ------------------------
+def load_grayscale_image():
+    # Try common sample images; fall back to skimage if available; else ask user to put an image in cwd
+    candidates = ["onion.png", "cameraman.tif", "peppers.png", "lena.png", "camera.png"]
+    for name in candidates:
+        if os.path.exists(name):
+            # img = Image.open(name).convert('L')
+            img = Image.open(name).convert("RGB")
+            image_np = np.asarray(img, dtype=np.float64)
+            # print(f"{image_np = }")
+            image_np = image_np / 255.0
+            # print(f"{image_np = }")
+            return image_np
+
+    raise FileNotFoundError(
+        "Could not find a local image. Place an image (e.g., cameraman.tif/peppers.png) in the working directory."
+    )
+
+
+# ------------------------ DCT implementations (orthonormal) ------------------------
+
+
+def dct2_matrix_ortho(N, device="cpu", dtype=torch.float32):
+    # T2[k, n] = sqrt(2/N) * beta(k) * cos(pi/N * (n + 0.5) * k), beta(0)=1/sqrt(2)
+    n = torch.arange(N, device=device, dtype=dtype)
+    k = torch.arange(N, device=device, dtype=dtype).unsqueeze(1)
+    W = torch.cos(math.pi / N * (n + 0.5) * k)  # [N, N]
+    beta = torch.ones(N, device=device, dtype=dtype)
+    beta[0] = 1 / math.sqrt(2.0)
+    T = (math.sqrt(2.0 / N) * beta).unsqueeze(1) * W
+    return T  # orthonormal; inverse is T.T
+
+
+def dct1_matrix_ortho(N, device="cpu", dtype=torch.float32):
+    # T1[k, n] = sqrt(2/(N-1)) * alpha(k) * alpha(n) * cos(pi/(N-1) * n*k)
+    # alpha(0)=alpha(N-1)=1/sqrt(2), else 1. Self-inverse (orthonormal and symmetric).
+    if N < 2:
+        # N=1 trivial case
+        return torch.ones((1, 1), device=device, dtype=dtype)
+    n = torch.arange(N, device=device, dtype=dtype)
+    k = torch.arange(N, device=device, dtype=dtype).unsqueeze(1)
+    C = torch.cos(math.pi / (N - 1) * (n * k))  # [N, N]
+    alpha = torch.ones(N, device=device, dtype=dtype)
+    alpha[0] = 1 / math.sqrt(2.0)
+    alpha[-1] = 1 / math.sqrt(2.0)
+    T = math.sqrt(2.0 / (N - 1)) * (alpha.unsqueeze(1) * C * alpha.unsqueeze(0))
+    return T  # orthonormal, symmetric, self-inverse
+
+
+def dct2_ortho(x, T2=None):
+    # x: [N] float tensor. Returns DCT-II (orthonormal) [N].
+    x = x.reshape(-1)
+    N = x.numel()
+    if T2 is None:
+        T2 = dct2_matrix_ortho(N, device=x.device, dtype=x.dtype)
+    return T2 @ x
+
+
+def idct2_ortho(X, T2=None):
+    # Inverse of DCT-II (orthonormal) is transpose
+    X = X.reshape(-1)
+    N = X.numel()
+    if T2 is None:
+        T2 = dct2_matrix_ortho(N, device=X.device, dtype=X.dtype)
+    return T2.t() @ X
+
+
+def dct1_ortho(x, T1=None):
+    # x: [N] float tensor. Returns DCT-I (orthonormal) [N].
+    x = x.reshape(-1)
+    N = x.numel()
+    if T1 is None:
+        T1 = dct1_matrix_ortho(N, device=x.device, dtype=x.dtype)
+    return T1 @ x
+
+
+def idct1_ortho(X, T1=None):
+    # DCT-I orthonormal is self-inverse
+    X = X.reshape(-1)
+    N = X.numel()
+    if T1 is None:
+        T1 = dct1_matrix_ortho(N, device=X.device, dtype=X.dtype)
+    return T1 @ X
+
+
+def _complex_dtype_from_real(real_dtype):
+    if real_dtype == torch.float32:
+        return torch.complex64
+    if real_dtype == torch.float64:
+        return torch.complex128
+    raise TypeError("Only float32/float64 supported.")
+
+
+def dct2_fft(x, dim=-1, norm="ortho"):
+    """
+    DCT-II via even-symmetric 2N extension and torch.fft.rfft.
+    x: real tensor (..., N)
+    Returns: real tensor (..., N)
+    norm: 'ortho' (orthonormal, like scipy.fft.dct(..., type=2, norm='ortho')) or None (unnormalized).
+    """
+    if not torch.is_floating_point(x):
+        raise TypeError("x must be float tensor")
+    N = x.shape[dim]
+    if N < 1:
+        return x.clone()
+
+    # Even extension [x, flip(x)]
+    x_flip = torch.flip(x, dims=(dim,))
+    s = torch.cat([x, x_flip], dim=dim)  # (..., 2N)
+
+    # RFFT over length 2N
+    S = torch.fft.rfft(s, n=2 * N, dim=dim)  # (..., N+1)
+
+    # k = 0..N-1
+    k = torch.arange(N, device=x.device, dtype=x.dtype)
+    # exp(-j*pi*k/(2N))
+    ctype = _complex_dtype_from_real(x.dtype)
+    twiddle = torch.exp(-1j * math.pi * k / (2.0 * N)).to(dtype=ctype, device=x.device)
+    for _ in range(dim, S.dim() - 1):
+        twiddle = twiddle.unsqueeze(-1)
+
+    # Take real part; factor 1/2 (see derivation)
+    C = (S.narrow(dim, 0, N) * twiddle).real * 0.5  # (..., N)
+
+    if norm == "ortho":
+        # Orthonormal scaling: sqrt(2/N) * beta(k), beta(0)=1/sqrt(2)
+        C = C * math.sqrt(2.0 / N)
+        index0 = [slice(None)] * C.dim()
+        index0[dim] = 0
+        C[tuple(index0)] /= math.sqrt(2.0)
+    elif norm is None:
+        pass
+    else:
+        raise ValueError("norm must be 'ortho' or None")
+    return C
+
+
+def idct2_fft(C, dim=-1, norm="ortho"):
+    """
+    Inverse of dct2_fft (i.e., DCT-III) using torch.fft.irfft.
+    C: real tensor (..., N) with same norm used in dct2_fft.
+    Returns real tensor (..., N).
+    """
+    if not torch.is_floating_point(C):
+        raise TypeError("C must be float tensor")
+    N = C.shape[dim]
+    if N < 1:
+        return C.clone()
+
+    # Undo orthonormal scaling to get "unnormalized" DCT-II coefficients
+    Cun = C
+    if norm == "ortho":
+        Cun = C / math.sqrt(2.0 / N)
+        index0 = [slice(None)] * Cun.dim()
+        index0[dim] = 0
+        Cun = Cun.clone()
+        Cun[tuple(index0)] *= math.sqrt(2.0)
+    elif norm is None:
+        Cun = C
+    else:
+        raise ValueError("norm must be 'ortho' or None")
+
+    # Build unique half-spectrum (length N+1) for the 2N-length irfft
+    # S[k] = 2*Cun[k] * exp(+j*pi*k/(2N)), for k=0..N-1
+    k = torch.arange(N, device=C.device, dtype=C.dtype)
+    ctype = _complex_dtype_from_real(C.dtype)
+    twiddle = torch.exp(+1j * math.pi * k / (2.0 * N)).to(dtype=ctype, device=C.device)
+    for _ in range(dim, C.dim() - 1):
+        twiddle = twiddle.unsqueeze(-1)
+
+    # Allocate (..., N+1)
+    new_shape = list(Cun.shape)
+    new_shape[dim] = N + 1
+    S_half = torch.zeros(*new_shape, dtype=ctype, device=C.device)
+
+    # Fill 0..N-1
+    # real times complex -> cast below
+    S_part = (2.0 * Cun) * twiddle.real - 0j
+    S_part = (2.0 * Cun).to(ctype) * twiddle
+    S_half.narrow(dim, 0, N).copy_(S_part)
+
+    # Nyquist (k=N) is zero for the chosen even-symmetric extension
+    indexN = [slice(None)] * S_half.dim()
+    indexN[dim] = N
+    S_half[tuple(indexN)] = 0
+
+    # irfft to length 2N, take first N samples
+    s = torch.fft.irfft(S_half, n=2 * N, dim=dim)  # (..., 2N)
+
+    # Slice first N along dim
+    x = s.narrow(dim, 0, N)
+    return x
+
+
+# --------- N-D (multi-axis) DCT-II / IDCT-II built from the 1D versions ---------
+def _normalize_dims(dims, ndim):
+    if isinstance(dims, int):
+        dims = (dims,)
+    dims = tuple(d if d >= 0 else d + ndim for d in dims)
+    if any(d < 0 or d >= ndim for d in dims):
+        raise ValueError("dims out of range for input tensor.")
+    # You can enforce uniqueness if desired:
+    if len(set(dims)) != len(dims):
+        raise ValueError("dims must be unique.")
+    return dims
+
+
+def dct2_nd_fft(x, dims, norm="ortho"):
+    """
+    N-D DCT-II applied along the specified dimensions.
+    x: real tensor
+    dims: tuple of axes (e.g., (-2,-1) for 2D, (-3,-2,-1) for 3D)
+    norm: 'ortho' or None
+    """
+    dims = _normalize_dims(dims, x.ndim)
+    y = x
+    for d in dims:
+        y = dct2_fft(y, dim=d, norm=norm)
+    return y
+
+
+def idct2_nd_fft(X, dims, norm="ortho"):
+    """
+    N-D inverse of DCT-II (DCT-III) along the specified dimensions.
+    """
+    dims = _normalize_dims(dims, X.ndim)
+    y = X
+    for d in dims:
+        y = idct2_fft(y, dim=d, norm=norm)
+    return y
+
+
+def _to_device_dtype(x, device, dtype):
+    if device is None:
+        device = x.device if isinstance(x, torch.Tensor) else "cpu"
+    if dtype is None:
+        dtype = torch.float64  # match MATLAB double
+    return device, dtype
+
+
+def _omega_grid_1d(N, shifted, device, dtype):
+    # Digital radian frequency samples on FFT bins.
+    # unshifted: ω_k = 2π k / N, k=0..N-1 (DC at index 0)
+    # shifted: fftshift layout (DC at center), monotonically increasing from negative to positive
+    k = torch.arange(N, device=device, dtype=dtype)
+    w = 2.0 * math.pi * k / N
+    # [0, 2π)
+    if shifted:
+        w = torch.fft.fftshift(w)  # center DC
+    return w
+
+
+def _tan_half_abs(w, eps=1e-12):
+    # Safe |tan(w/2)| to avoid overflow at w=π.
+    half = 0.5 * w
+    c = torch.cos(half)
+    s = torch.sin(half)
+    # Where cos is near zero, use a very large value (approach infinity)
+    # large but not inf to avoid NaNs downstream
+    large = torch.finfo(w.dtype).max ** 0.5
+    t = torch.where(c.abs() < eps, torch.sign(s) * large, s / c)
+    return t.abs()
+
+
+def butterworth_mask_1d(
+    N,
+    fc,
+    order,
+    btype="low",
+    shifted=False,
+    device=None,
+    dtype=None,
+):
+    """
+    1D Butterworth frequency mask equivalent to MATLAB butter+freqz magnitude.
+    - N: number of FFT bins
+    - fc: normalized cutoff(s) in cycles/sample (relative to 1 sample) with 0 < fc < 0.5
+      low/high: scalar; bandpass/stop: [f1, f2] with 0 < f1 < f2 < 0.5
+        * fc is equivalent to Wn / 2 in MATLAB's butter function. e.g. butter(4, 0.25) is equivalent to fc=0.125 here.
+    - order: integer >= 1
+    - btype: 'low', 'high', 'bandpass', 'stop'
+    - shifted: if True, return mask in fftshift layout (DC at center)
+    """
+    assert isinstance(N, int) and N >= 2
+    assert isinstance(order, int) and order >= 1
+    btype = btype.lower()
+    if btype in ("low", "high"):
+        fc = float(fc)
+        assert 0.0 < fc < 0.5
+    else:
+        assert len(fc) == 2
+        f1, f2 = float(fc[0]), float(fc[1])
+        assert 0.0 < f1 < f2 < 0.5
+        fc = (f1, f2)
+
+    device, dtype = _to_device_dtype(torch.empty(0), device, dtype)
+    w = _omega_grid_1d(N, shifted=shifted, device=device, dtype=dtype)  # 0..2π (or centered)
+    # Bilinear mapping (prewarped): Ω = 2 * tan(ω/2)
+    Om = 2.0 * _tan_half_abs(w)  # analog rad/sec (normalized T=1)
+
+    if btype == "low":
+        # Prewarp analog cutoff: Ωc = 2*tan(π*fc)
+        Oc = 2.0 * math.tan(math.pi * fc)
+        ratio = (Om / Oc).clamp_min(0)
+        mag = 1.0 / torch.sqrt(1.0 + ratio.pow(2 * order))
+    elif btype == "high":
+        Oc = 2.0 * math.tan(math.pi * fc)
+        # Handle Om=0 => magnitude=0
+        ratio = torch.where(Om > 0, (Oc / Om), torch.full_like(Om, float("inf")))
+        mag = 1.0 / torch.sqrt(1.0 + ratio.pow(2 * order))
+    elif btype == "bandpass":
+        f1, f2 = fc
+        O1 = 2.0 * math.tan(math.pi * f1)
+        O2 = 2.0 * math.tan(math.pi * f2)
+        B = O2 - O1
+        O0 = math.sqrt(O1 * O2)
+        # D(Ω) = (Ω^2 - Ω0^2)/(B*Ω)
+        denom = B * Om
+        # denom=0 at Om=0 -> D=inf, magnitude=0
+        D = torch.where(denom != 0, (Om.pow(2) - O0**2) / denom, torch.full_like(Om, float("inf")))
+        mag = 1.0 / torch.sqrt(1.0 + D.abs().pow(2 * order))
+    elif btype in ("stop", "bandstop", "bandreject"):
+        f1, f2 = fc
+        O1 = 2.0 * math.tan(math.pi * f1)
+        O2 = 2.0 * math.tan(math.pi * f2)
+        B = O2 - O1
+        O0 = math.sqrt(O1 * O2)
+        # D(Ω) = (B*Ω)/(Ω^2 - Ω0^2)
+        denom = Om.pow(2) - O0**2
+        # denom=0 at Om=O0 -> D=inf, magnitude=0
+        D = torch.where(denom != 0, (B * Om) / denom, torch.full_like(Om, float("inf")))
+        mag = 1.0 / torch.sqrt(1.0 + D.abs().pow(2 * order))
+    else:
+        raise ValueError("btype must be 'low', 'high', 'bandpass', or 'stop'.")
+
+    return mag.to(dtype=dtype, device=device)
+
+
+def butterworth_mask_2d_separable(
+    shape,
+    fc,
+    order,
+    btype="low",
+    shifted=False,
+    device=None,
+    dtype=None,
+):
+    """
+    2D separable Butterworth mask (rows × cols), equivalent to applying 1D Butterworth along rows and columns (zero-phase). Not an isotropic circular Butterworth.
+    - shape: (M, N)
+    - fc: scalar or 2-tuple for low/high; for band types, pass 2-tuples for each axis: ([f1y,f2y], [f1x,f2x]) You can also pass scalar or 2-tuple to apply same cutoffs on both axes.
+    - order: integer or 2-tuple for (order_y, order_x)
+    - btype: 'low', 'high', 'bandpass', 'stop'
+    - shifted: if True, both axes are centered (fftshift layout)
+    """
+    M, N = int(shape[0]), int(shape[1])
+    assert M >= 2 and N >= 2
+    device, dtype = _to_device_dtype(torch.empty(0), device, dtype)
+
+    # Normalize fc/order to per-axis tuples
+    if btype in ("low", "high"):
+        if not isinstance(fc, (list, tuple)):
+            fcy = fcx = fc
+        else:
+            assert len(fc) == 2
+            fcy, fcx = fc
+    else:
+        # band types
+        if isinstance(fc[0], (list, tuple)) and isinstance(fc[1], (list, tuple)):
+            fcy, fcx = fc
+        else:
+            # same band on both axes
+            fcy = fcx = fc
+
+    if isinstance(order, (list, tuple)):
+        oy, ox = int(order[0]), int(order[1])
+    else:
+        oy = ox = int(order)
+
+    Hy = butterworth_mask_1d(M, fcy, oy, btype=btype, shifted=shifted, device=device, dtype=dtype)
+    Hx = butterworth_mask_1d(N, fcx, ox, btype=btype, shifted=shifted, device=device, dtype=dtype)
+
+    # Outer product to build separable 2D mask
+    H2 = Hy.reshape(M, 1) * Hx.reshape(1, N)
+    return H2
+
+
+def _freqvec_norm(
+    N: int,
+    shifted: bool,
+    device=None,
+    dtype=None,
+):
+    """
+    Normalized frequency vector in [-0.5, 0.5), length N.
+    - shifted=False: DC at index 0 (unshifted FFT layout)
+    - shifted=True: DC at center (fftshift layout)
+    """
+    if device is None:
+        device = "cpu"
+    if dtype is None:
+        dtype = torch.float64
+    k = torch.arange(N, device=device, dtype=dtype)
+    if shifted:
+        f = (k - torch.floor(torch.tensor(N / 2, dtype=dtype, device=device))) / N
+    else:
+        f = k / N
+    f = torch.where(f >= 0.5, f - 1.0, f)  # wrap into [-0.5, 0.5)
+    return f  # [N]
+
+
+def _radial_frequency_nd(
+    shape: Sequence[int],
+    shifted: bool,
+    device=None,
+    dtype=None,
+):
+    """
+    Radial normalized frequency R in [-0.5,0.5) computed over all axes.
+    Returns R with shape 'shape'.
+    """
+    if device is None:
+        device = "cpu"
+    if dtype is None:
+        dtype = torch.float64
+    grids = [_freqvec_norm(N, shifted=shifted, device=device, dtype=dtype) for N in shape]
+    # list of tensors, each shape = shape
+    meshes = torch.meshgrid(*grids, indexing="ij")
+    R2 = torch.zeros(shape, dtype=dtype, device=device)
+    for g in meshes:
+        R2 = R2 + g**2
+    R = torch.sqrt(R2)
+    return R
+
+
+def butterworth_nd(
+    shape: Sequence[int],
+    cutoff: Union[float, Tuple[float, float]],
+    order: int,
+    btype: str = "low",
+    shifted: bool = False,
+    device=None,
+    dtype=None,
+):
+    """Isotropic N-D Butterworth mask (low/high/bandpass/bandstop).
+    Args:
+    shape: iterable of ints, e.g., (H, W) or (D, H, W) ...
+    cutoff:
+        - 'low'/'high': scalar D0 in (0, 0.5]
+        - 'bandpass'/'bandstop': tuple (D1, D2) with 0 < D1 < D2 <= 0.5
+    order: integer >= 1
+    btype: 'low' | 'high' | 'bandpass' | 'bandstop' (alias 'stop')
+    shifted: if True, mask is centered (fftshift layout); else unshifted
+    device, dtype: optional torch device/dtype (defaults: CPU, float64)
+
+    Returns:
+    H: tensor with shape 'shape', values in [0, 1].
+    """
+    assert len(shape) >= 1 and all(int(s) >= 1 for s in shape), "Invalid shape."
+    order = int(order)
+    assert order >= 1, "order must be >= 1"
+    btype = btype.lower()
+    if btype in ("low", "high"):
+        D0 = float(cutoff)
+        # assert 0.0 < D0 <= 0.5, "cutoff must be in (0, 0.5]"
+    else:
+        D1, D2 = float(cutoff[0]), float(cutoff[1])
+        # assert 0.0 < D1 < D2 <= 0.5, "for band types: 0 < D1 < D2 <= 0.5"
+        B = D2 - D1
+        D0 = math.sqrt(D1 * D2)
+
+    if device is None:
+        device = "cpu"
+    if dtype is None:
+        dtype = torch.float64
+
+    R = _radial_frequency_nd(
+        tuple(int(s) for s in shape), shifted=shifted, device=device, dtype=dtype
+    )
+    eps = torch.finfo(dtype).eps
+    # print(f"{R = }")
+
+    if btype == "low":
+        # H = 1 / (1 + (R/D0)^(2n))
+        ratio = (R / D0).clamp_min(0)
+        H = 1.0 / (1.0 + ratio.pow(2 * order))
+
+    elif btype == "high":
+        # H = 1 / (1 + (D0/R)^(2n)), H(DC)=0
+        # avoid divide-by-zero at R=0
+        safe_R = torch.where(R > 0, R, torch.tensor(1.0, device=device, dtype=dtype))  # dummy
+        ratio = D0 / safe_R
+        H = 1.0 / (1.0 + ratio.pow(2 * order))
+        # enforce DC = 0
+        H = torch.where(R > 0, H, torch.zeros_like(H))
+
+    elif btype == "bandpass":
+        # D = (R^2 - D0^2) / (B*R); H = 1 / (1 + |D|^(2n))
+        # Handle R=0 -> D=inf -> H=0
+        denom = B * R
+        D = torch.where(denom != 0, (R.pow(2) - D0**2) / denom, torch.full_like(R, float("inf")))
+        H = 1.0 / (1.0 + D.abs().pow(2 * order))
+
+    elif btype in ("bandstop", "stop", "bandreject"):
+        # D = (B*R) / (R^2 - D0^2); H = 1 / (1 + |D|^(2n))
+        # Handle R^2 - D0^2 = 0 -> D=inf -> H=0 (deep notch at R=D0)
+        denom = R.pow(2) - D0**2
+        D = torch.where(denom != 0, (B * R) / denom, torch.full_like(R, float("inf")))
+        H = 1.0 / (1.0 + D.abs().pow(2 * order))
+
+    else:
+        raise ValueError("btype must be 'low', 'high', 'bandpass', or 'bandstop'.")
+
+    return H
+
+
+def butterworth_low_pass_filter(
+    tensor: torch.Tensor,
+    dims: Sequence[int],
+    cutoff: float,
+    order: int,
+    shifted: bool = False,
+    device=None,
+    dtype=None,
+):
+    """
+    Applies a Butterworth low-pass filter to the input tensor.
+
+    the dims specify which dim should be perform filtering
+
+    return filtered tensor
+    """
+    if not isinstance(dims, (list, tuple)):
+        dims = (dims,)
+    ndims_total = tensor.ndim
+    # Normalize dims (handle negatives)
+    norm_dims = _normalize_dims(dims, ndim=ndims_total)
+
+    original_dtype = tensor.dtype
+    work_dtype = dtype or (tensor.dtype if torch.is_floating_point(tensor) else torch.float32)
+    if work_dtype == torch.bfloat16 or work_dtype == torch.float16:
+        work_dtype = torch.float32
+    device = device or tensor.device
+
+    # Prepare frequency-domain representation
+    x = tensor.to(device=device, dtype=work_dtype)
+    X = torch.fft.fftn(x, dim=norm_dims)
+    if shifted:
+        X = torch.fft.fftshift(X, dim=norm_dims)
+
+    # Build isotropic Butterworth mask over the selected dims
+    shape_subset = [x.shape[d] for d in norm_dims]
+    H_small = butterworth_nd(
+        shape=shape_subset,
+        cutoff=cutoff,
+        order=order,
+        btype="low",
+        shifted=shifted,
+        device=device,
+        dtype=work_dtype,
+    )
+
+    # Broadcast mask into full tensor shape
+    mask_shape = [1] * ndims_total
+    for i, d in enumerate(norm_dims):
+        mask_shape[d] = shape_subset[i]
+    H = H_small.view(*mask_shape)
+
+    # Apply mask
+    X_filtered = X * H
+
+    # Inverse FFT
+    if shifted:
+        X_filtered = torch.fft.ifftshift(X_filtered, dim=norm_dims)
+    x_filtered = torch.fft.ifftn(X_filtered, dim=norm_dims).real
+
+    return x_filtered.to(dtype=original_dtype)
+
+
+# def fft_denoise(tensor, dim, fft_ratio):
+#     assert len(dim) == 2
+#     original_dtype = tensor.dtype
+#     tensor = tensor.to(torch.float32)
+#     # Create low pass filter
+#     LPF = butterworth_low_pass_filter(
+#         (tensor.shape[dim[0]], tensor.shape[dim[1]]),
+#         n=4,
+#         d_s=fft_ratio,
+#     )
+#     LPF = LPF.to(dtype=tensor.dtype, device=tensor.device)
+#     # print(f"{LPF = }")
+#     # print(f"{LPF.shape = }")
+#     for _ in range(dim[0]):
+#         LPF = LPF.unsqueeze(0)
+#     for _ in range(dim[1] + 1, len(tensor.shape)):
+#         LPF = LPF.unsqueeze(-1)
+#     # print(f"{LPF.shape = }")
+#     # FFT
+#     latents_freq_k = torch.fft.fftn(tensor, dim=dim)
+#     # print(f"{latents_freq_k.shape = }")
+#     latents_freq_k = torch.fft.fftshift(latents_freq_k, dim=dim)
+#     # print(f"{latents_freq_k.shape = }")
+
+#     new_freq_k = latents_freq_k * LPF
+
+#     # IFFT
+#     new_freq_k = torch.fft.ifftshift(new_freq_k, dim=dim)
+#     denoised_k = torch.fft.ifftn(new_freq_k, dim=dim).real
+#     denoised_k = denoised_k.to(original_dtype)
+#     return denoised_k
+
+
+if __name__ == "__main__":
+    # x = torch.linspace(0, 2 * np.pi, 8)
+    # y = torch.linspace(0, 2 * np.pi, 8)
+    # X, Y = torch.meshgrid(x, y, indexing='ij')
+    # latents = (
+    #     torch.sin(2 * X + Y) +
+    #     torch.sin(X + 3 * Y) +
+    #     torch.sin(3 * X - 2 * Y)
+    # ) + 1
+    # latents += 0.01 * torch.randn_like(latents)  # Add Gaussian noise
+    # # latents = torch.randn([8, 8])
+    # print(f"latents = \n{latents}")
+
+    # latents_freq = torch.fft.fftn(latents, dim=(-2, -1))
+    # print(f"latents_freq = \n{torch.abs(latents_freq)}")
+
+    # latents_freq_shift = torch.fft.fftshift(latents_freq, dim=(-2, -1))
+    # print(f"latents_freq_shift = \n{torch.abs(latents_freq_shift)}")
+
+    # latents_freq_dct = dct_2d(latents)
+    # print(f"latents_freq_dct = \n{latents_freq_dct}")
+
+    # LPF_1 = butterworth_low_pass_filter(latents=latents, d_s=-1.0)
+
+    # print(f"LPF_1 = \n{LPF_1}")
+
+    # LPF_2 = my_butterworth_low_pass_filter_non_center(
+    #     shape=latents.shape,
+    #     stop_freqs=[0.25, 0.25],
+    #     n=4,
+    # )
+    # print(f"LPF_2 = \n{LPF_2}")
+
+    # LPF_3 = my_butterworth_low_pass_filter(
+    #     shape=latents.shape,
+    #     stop_freqs=[0.25, 0.25],
+    #     n=4,
+    # )
+    # print(f"LPF_3 = \n{LPF_3}")
+
+    # img = load_grayscale_image()
+    # # Extract middle column as 1-D signal
+    # col = img.shape[1] // 2 - 1
+    # print(f"{col = }")
+    # x_np = img[:, col].astype(np.float32)  # [H]
+    # # print(f"{x_np = }")
+    # N = x_np.shape[0]
+    # print(f"{N = }")
+
+    # device = 'cpu'
+    # dtype = torch.float64
+
+    # x = torch.from_numpy(img).to(device=device, dtype=dtype)
+    # print(f"{x = }")
+
+    # # Transforms
+    # Xf = torch.fft.fftn(x, dim=(-3, -2, -1), norm=None)            # complex64
+    # print(f"{Xf = }")
+    # x_reconstructed = torch.fft.ifftn(Xf, dim=(-3, -2, -1), norm=None)
+    # print(f"{x_reconstructed = }")
+    # print(f"{(x - x_reconstructed).abs().max() = }")
+
+    # Xd2 = dct2_nd_fft(x, dims=(-3, -2, -1), norm="ortho")       # float
+    # print(f"{Xd2 = }")
+    # x_reconstructed = idct2_nd_fft(Xd2, dims=(-1, -2, -3), norm="ortho")
+    # print(f"{x_reconstructed = }")
+    # print(f"{(x - x_reconstructed).abs().max() = }")
+
+    # H1 = butterworth_mask_1d(16, 0.125, 4, btype='low', shifted=True)
+    # print(f"{H1 = }")
+
+    H2 = butterworth_nd([30, 52], 1.0, 4, btype="low", shifted=True)
+    print(f"{H2 = }")
+
+    # ---- Planar wave demo with Butterworth low-pass filtering ----
+    def demo_planar_wave():
+        # Generate 2D planar wave: low-frequency + added high-frequency component
+        H, W = 128, 128
+        device = "cpu"
+        y = torch.arange(H, device=device).view(H, 1)
+        x = torch.arange(W, device=device).view(1, W)
+
+        # Low-frequency component
+        kx_low, ky_low = 2, 3
+        low = torch.sin(2 * math.pi * (kx_low * x / W + ky_low * y / H))
+
+        # High-frequency component
+        kx_high, ky_high = 20, 24
+        high = 0.5 * torch.sin(2 * math.pi * (kx_high * x / W + ky_high * y / H))
+
+        signal = low + high
+
+        # Apply Butterworth low-pass (cutoff chosen to keep low freq, attenuate high freq)
+        cutoff = 0.12  # normalized radial cutoff (<=0.5)
+        order = 4
+        filtered = butterworth_low_pass_filter(
+            signal, dims=(-2, -1), cutoff=cutoff, order=order, shifted=True
+        )
+
+        # Metrics
+        mse_before = (signal - low).pow(2).mean()
+        mse_after = (filtered - low).pow(2).mean()
+        residual_energy_ratio = (filtered - low).pow(2).sum() / (signal - low).pow(2).sum()
+
+        print("Planar wave demo:")
+        print(f"mse_before={mse_before.item():.6e}")
+        print(f"mse_after ={mse_after.item():.6e}")
+        print(f"residual_energy_ratio={residual_energy_ratio.item():.4%}")
+        # Quick sanity: high frequency suppression (should be << 1)
+        assert (
+            mse_after < mse_before
+        ), "Filtering did not reduce error to low-frequency ground truth."
+
+    demo_planar_wave()

inference.py

@@ -0,0 +1,263 @@
+import argparse
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+import torch
+from omegaconf import OmegaConf
+from tqdm import tqdm
+from torchvision import transforms
+from torchvision.io import write_video
+from einops import rearrange
+import torch.distributed as dist
+from torch.utils.data import DataLoader, SequentialSampler
+from torch.utils.data.distributed import DistributedSampler
+
+from pipeline import (
+    CausalDiffusionInferencePipeline,
+    CausalInferencePipeline,
+)
+from utils.dataset import TextDataset, TextImagePairDataset
+from utils.misc import set_seed
+from hydra import initialize, compose
+from hydra.core.global_hydra import GlobalHydra
+
+from demo_utils.memory import gpu, get_cuda_free_memory_gb, DynamicSwapInstaller
+
+from pathlib import Path
+
+config_name = "self_forcing_dmd_vsink"
+output_chunk_number = 21
+output_latent_frame_number = 21
+# output_latent_frame_number = 81
+seed = 42
+import sys
+
+sys.argv.extend(
+    [
+        "--output_folder",
+        f"outputs/{output_latent_frame_number}-{config_name}-seed{seed}",
+        # f"outputs-test/{output_latent_frame_number}-{config_name}-seed{seed}",
+        "--config_dir",
+        "configs",
+        "--config_name",
+        config_name,
+        "--num_output_frames",
+        f"{output_latent_frame_number}",
+        "--data_path",
+        "prompts/MovieGenVideoBench_extended.txt",
+        "--checkpoint_path",
+        "./checkpoints/self_forcing_dmd.pt",
+        "--use_ema",
+        "--seed",
+        f"{seed}",
+    ]
+)
+print(f"{sys.argv = }")
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config_dir", type=str, help="Directory to the config file")
+parser.add_argument("--config_name", type=str, help="Name to the config file")
+parser.add_argument("--checkpoint_path", type=str, help="Path to the checkpoint folder")
+parser.add_argument("--data_path", type=str, help="Path to the dataset")
+parser.add_argument("--extended_prompt_path", type=str, help="Path to the extended prompt")
+parser.add_argument("--output_folder", type=str, help="Output folder")
+parser.add_argument(
+    "--num_output_frames",
+    type=int,
+    default=21,
+    help="Number of overlap frames between sliding windows",
+)
+parser.add_argument(
+    "--i2v",
+    action="store_true",
+    help="Whether to perform I2V (or T2V by default)",
+)
+parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA parameters")
+parser.add_argument("--seed", type=int, default=0, help="Random seed")
+parser.add_argument(
+    "--num_samples",
+    type=int,
+    default=1,
+    help="Number of samples to generate per prompt",
+)
+args = parser.parse_args()
+
+# Initialize distributed inference
+if "LOCAL_RANK" in os.environ:
+    dist.init_process_group(backend="nccl")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    world_size = dist.get_world_size()
+    set_seed(args.seed + local_rank)
+else:
+    device = torch.device("cuda")
+    local_rank = 0
+    world_size = 1
+    set_seed(args.seed)
+
+print(f"Free VRAM {get_cuda_free_memory_gb(gpu)} GB")
+low_memory = get_cuda_free_memory_gb(gpu) < 40
+
+torch.set_grad_enabled(False)
+
+if GlobalHydra.instance().is_initialized():
+    GlobalHydra.instance().clear()
+
+with initialize(version_base=None, config_path=args.config_dir):
+    config = compose(config_name=args.config_name)
+print(f"{config = }")
+
+# Initialize pipeline
+if hasattr(config, "denoising_step_list"):
+    # Few-step inference
+    pipeline = CausalInferencePipeline(config, device=device)
+else:
+    # Multi-step diffusion inference
+    pipeline = CausalDiffusionInferencePipeline(config, device=device)
+
+if args.checkpoint_path:
+    state_dict = torch.load(args.checkpoint_path, map_location="cpu")
+    pipeline.generator.load_state_dict(
+        state_dict["generator" if not args.use_ema else "generator_ema"]
+    )
+
+pipeline = pipeline.to(dtype=torch.bfloat16)
+if low_memory:
+    DynamicSwapInstaller.install_model(pipeline.text_encoder, device=gpu)
+else:
+    pipeline.text_encoder.to(device=gpu)
+pipeline.generator.to(device=gpu)
+pipeline.vae.to(device=gpu)
+
+
+# Create dataset
+if args.i2v:
+    assert not dist.is_initialized(), "I2V does not support distributed inference yet"
+    transform = transforms.Compose(
+        [
+            transforms.Resize((480, 832)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),
+        ]
+    )
+    dataset = TextImagePairDataset(args.data_path, transform=transform)
+else:
+    dataset = TextDataset(
+        prompt_path=args.data_path,
+        extended_prompt_path=args.extended_prompt_path,
+    )
+num_prompts = len(dataset)
+print(f"Number of prompts: {num_prompts}")
+
+if dist.is_initialized():
+    sampler = DistributedSampler(dataset, shuffle=False, drop_last=True)
+else:
+    sampler = SequentialSampler(dataset)
+dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
+
+# Create output directory (only on main process to avoid race conditions)
+if local_rank == 0:
+    os.makedirs(args.output_folder, exist_ok=True)
+
+if dist.is_initialized():
+    dist.barrier()
+
+
+def encode(self, videos: torch.Tensor) -> torch.Tensor:
+    device, dtype = videos[0].device, videos[0].dtype
+    scale = [
+        self.mean.to(device=device, dtype=dtype),
+        1.0 / self.std.to(device=device, dtype=dtype),
+    ]
+    output = [self.model.encode(u.unsqueeze(0), scale).float().squeeze(0) for u in videos]
+
+    output = torch.stack(output, dim=0)
+    return output
+
+
+for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)):
+    idx = batch_data["idx"].item()
+
+    # For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container
+    # Unpack the batch data for convenience
+    if isinstance(batch_data, dict):
+        batch = batch_data
+    elif isinstance(batch_data, list):
+        batch = batch_data[0]  # First (and only) item in the batch
+
+    all_video = []
+    num_generated_frames = 0  # Number of generated (latent) frames
+
+    set_seed(args.seed)
+    if args.i2v:
+        # For image-to-video, batch contains image and caption
+        prompt = batch["prompts"][0]  # Get caption from batch
+        prompts = [prompt] * args.num_samples
+
+        # Process the image
+        image = (
+            batch["image"]
+            .squeeze(0)
+            .unsqueeze(0)
+            .unsqueeze(2)
+            .to(device=device, dtype=torch.bfloat16)
+        )
+
+        # Encode the input image as the first latent
+        initial_latent = pipeline.vae.encode_to_latent(image).to(
+            device=device, dtype=torch.bfloat16
+        )
+        initial_latent = initial_latent.repeat(args.num_samples, 1, 1, 1, 1)
+
+        sampled_noise = torch.randn(
+            [args.num_samples, args.num_output_frames - 1, 16, 60, 104],
+            device=device,
+            dtype=torch.bfloat16,
+        )
+    else:
+        # For text-to-video, batch is just the text prompt
+        prompt = batch["prompts"][0]
+        extended_prompt = batch["extended_prompts"][0] if "extended_prompts" in batch else None
+        if extended_prompt is not None:
+            prompts = [extended_prompt] * args.num_samples
+        else:
+            prompts = [prompt] * args.num_samples
+        initial_latent = None
+
+        sampled_noise = torch.randn(
+            [args.num_samples, args.num_output_frames, 16, 60, 104],
+            device=device,
+            dtype=torch.bfloat16,
+        )
+
+    set_seed(args.seed)
+    # Generate 81 frames
+    video, latents = pipeline.inference(
+        noise=sampled_noise,
+        text_prompts=prompts,
+        return_latents=True,
+        initial_latent=initial_latent,
+        low_memory=low_memory,
+    )
+    current_video = rearrange(video, "b t c h w -> b t h w c").cpu()
+    all_video.append(current_video)
+    num_generated_frames += latents.shape[1]
+
+    # Final output video
+    video = 255.0 * torch.cat(all_video, dim=1)
+
+    # Clear VAE cache
+    pipeline.vae.model.clear_cache()
+
+    # Save the video if the current prompt is not a dummy prompt
+    if idx < num_prompts:
+        model = "regular" if not args.use_ema else "ema"
+        for seed_idx in range(args.num_samples):
+            # All processes save their videos
+            output_path = os.path.join(
+                args.output_folder,
+                f"{idx}-{prompt[:50].replace(' ', '_')}-{seed_idx}_{model}.mp4",
+            )
+            write_video(output_path, video[seed_idx], fps=16)

model/__init__.py

@@ -0,0 +1,15 @@
+from .diffusion import CausalDiffusion
+from .causvid import CausVid
+from .dmd import DMD
+from .gan import GAN
+from .sid import SiD
+from .ode_regression import ODERegression
+
+__all__ = [
+    "CausalDiffusion",
+    "CausVid",
+    "DMD",
+    "GAN",
+    "SiD",
+    "ODERegression",
+]

model/base.py

@@ -0,0 +1,256 @@
+from typing import Tuple
+from einops import rearrange
+from torch import nn
+import torch.distributed as dist
+import torch
+
+from pipeline import SelfForcingTrainingPipeline
+from utils.loss import get_denoising_loss
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class BaseModel(nn.Module):
+    def __init__(
+        self,
+        args,
+        device,
+    ):
+        super().__init__()
+        self._initialize_models(args, device)
+
+        self.device = device
+        self.args = args
+        self.dtype = torch.bfloat16 if args.mixed_precision else torch.float32
+        if hasattr(args, "denoising_step_list"):
+            self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long)
+            if args.warp_denoising_step:
+                timesteps = torch.cat(
+                    (
+                        self.scheduler.timesteps.cpu(),
+                        torch.tensor([0], dtype=torch.float32),
+                    )
+                )
+                self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+    def _initialize_models(self, args, device):
+        self.real_model_name = getattr(args, "real_name", "Wan2.1-T2V-1.3B")
+        self.fake_model_name = getattr(args, "fake_name", "Wan2.1-T2V-1.3B")
+
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        self.real_score = WanDiffusionWrapper(model_name=self.real_model_name, is_causal=False)
+        self.real_score.model.requires_grad_(False)
+
+        self.fake_score = WanDiffusionWrapper(model_name=self.fake_model_name, is_causal=False)
+        self.fake_score.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+        self.scheduler = self.generator.get_scheduler()
+        self.scheduler.timesteps = self.scheduler.timesteps.to(device)
+
+    def _get_timestep(
+        self,
+        min_timestep: int,
+        max_timestep: int,
+        batch_size: int,
+        num_frame: int,
+        num_frame_per_block: int,
+        uniform_timestep: bool = False,
+    ) -> torch.Tensor:
+        """
+        Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep
+        from the range [min_timestep, max_timestep], and returns a tensor of shape [batch_size, num_frame].
+        - If uniform_timestep, it will use the same timestep for all frames.
+        - If not uniform_timestep, it will use a different timestep for each block.
+        """
+        if uniform_timestep:
+            timestep = torch.randint(
+                min_timestep,
+                max_timestep,
+                [batch_size, 1],
+                device=self.device,
+                dtype=torch.long,
+            ).repeat(1, num_frame)
+            return timestep
+        else:
+            timestep = torch.randint(
+                min_timestep,
+                max_timestep,
+                [batch_size, num_frame],
+                device=self.device,
+                dtype=torch.long,
+            )
+            # make the noise level the same within every block
+            if self.independent_first_frame:
+                # the first frame is always kept the same
+                timestep_from_second = timestep[:, 1:]
+                timestep_from_second = timestep_from_second.reshape(
+                    timestep_from_second.shape[0], -1, num_frame_per_block
+                )
+                timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1]
+                timestep_from_second = timestep_from_second.reshape(
+                    timestep_from_second.shape[0], -1
+                )
+                timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1)
+            else:
+                timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block)
+                timestep[:, :, 1:] = timestep[:, :, 0:1]
+                timestep = timestep.reshape(timestep.shape[0], -1)
+            return timestep
+
+
+class SelfForcingModel(BaseModel):
+    def __init__(
+        self,
+        args,
+        device,
+    ):
+        super().__init__(args, device)
+        self.denoising_loss_func = get_denoising_loss(args.denoising_loss_type)()
+
+    def _run_generator(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        initial_latent: torch.tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Optionally simulate the generator's input from noise using backward simulation
+        and then run the generator for one-step.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+            - initial_latent: a tensor containing the initial latents [B, F, C, H, W].
+        Output:
+            - pred_image: a tensor with shape [B, F, C, H, W].
+            - denoised_timestep: an integer
+        """
+        # Step 1: Sample noise and backward simulate the generator's input
+        assert getattr(
+            self.args, "backward_simulation", True
+        ), "Backward simulation needs to be enabled"
+        if initial_latent is not None:
+            conditional_dict["initial_latent"] = initial_latent
+        if self.args.i2v:
+            noise_shape = [
+                image_or_video_shape[0],
+                image_or_video_shape[1] - 1,
+                *image_or_video_shape[2:],
+            ]
+        else:
+            noise_shape = image_or_video_shape.copy()
+
+        # During training, the number of generated frames should be uniformly sampled from
+        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block
+        min_num_frames = 20 if self.args.independent_first_frame else 21
+        max_num_frames = (
+            self.num_training_frames - 1
+            if self.args.independent_first_frame
+            else self.num_training_frames
+        )
+        assert max_num_frames % self.num_frame_per_block == 0
+        assert min_num_frames % self.num_frame_per_block == 0
+        max_num_blocks = max_num_frames // self.num_frame_per_block
+        min_num_blocks = min_num_frames // self.num_frame_per_block
+        num_generated_blocks = torch.randint(
+            min_num_blocks, max_num_blocks + 1, (1,), device=self.device
+        )
+        dist.broadcast(num_generated_blocks, src=0)
+        num_generated_blocks = num_generated_blocks.item()
+        num_generated_frames = num_generated_blocks * self.num_frame_per_block
+        if self.args.independent_first_frame and initial_latent is None:
+            num_generated_frames += 1
+            min_num_frames += 1
+        # Sync num_generated_frames across all processes
+        noise_shape[1] = num_generated_frames
+
+        pred_image_or_video, denoised_timestep_from, denoised_timestep_to = (
+            self._consistency_backward_simulation(
+                noise=torch.randn(noise_shape, device=self.device, dtype=self.dtype),
+                **conditional_dict,
+            )
+        )
+        # Slice last 21 frames
+        if pred_image_or_video.shape[1] > 21:
+            with torch.no_grad():
+                # Reencode to get image latent
+                latent_to_decode = pred_image_or_video[:, :-20, ...]
+                # Deccode to video
+                pixels = self.vae.decode_to_pixel(latent_to_decode)
+                frame = pixels[:, -1:, ...].to(self.dtype)
+                frame = rearrange(frame, "b t c h w -> b c t h w")
+                # Encode frame to get image latent
+                image_latent = self.vae.encode_to_latent(frame).to(self.dtype)
+            pred_image_or_video_last_21 = torch.cat(
+                [image_latent, pred_image_or_video[:, -20:, ...]], dim=1
+            )
+        else:
+            pred_image_or_video_last_21 = pred_image_or_video
+
+        if num_generated_frames != min_num_frames:
+            # Currently, we do not use gradient for the first chunk, since it contains image latents
+            gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool)
+            if self.args.independent_first_frame:
+                gradient_mask[:, :1] = False
+            else:
+                gradient_mask[:, : self.num_frame_per_block] = False
+        else:
+            gradient_mask = None
+
+        pred_image_or_video_last_21 = pred_image_or_video_last_21.to(self.dtype)
+        return (
+            pred_image_or_video_last_21,
+            gradient_mask,
+            denoised_timestep_from,
+            denoised_timestep_to,
+        )
+
+    def _consistency_backward_simulation(
+        self,
+        noise: torch.Tensor,
+        **conditional_dict: dict,
+    ) -> torch.Tensor:
+        """
+        Simulate the generator's input from noise to avoid training/inference mismatch.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Here we use the consistency sampler (https://arxiv.org/abs/2303.01469)
+        Input:
+            - noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+        Output:
+            - output: a tensor with shape [B, T, F, C, H, W].
+            T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0
+            represents the x0 prediction at each timestep.
+        """
+        if self.inference_pipeline is None:
+            self._initialize_inference_pipeline()
+
+        return self.inference_pipeline.inference_with_trajectory(noise=noise, **conditional_dict)
+
+    def _initialize_inference_pipeline(
+        self,
+    ):
+        """
+        Lazy initialize the inference pipeline during the first backward simulation run.
+        Here we encapsulate the inference code with a model-dependent outside function.
+        We pass our FSDP-wrapped modules into the pipeline to save memory.
+        """
+        self.inference_pipeline = SelfForcingTrainingPipeline(
+            denoising_step_list=self.denoising_step_list,
+            scheduler=self.scheduler,
+            generator=self.generator,
+            num_frame_per_block=self.num_frame_per_block,
+            independent_first_frame=self.args.independent_first_frame,
+            same_step_across_blocks=self.args.same_step_across_blocks,
+            last_step_only=self.args.last_step_only,
+            num_max_frames=self.num_training_frames,
+            context_noise=self.args.context_noise,
+        )

model/causvid.py

@@ -0,0 +1,417 @@
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+
+
+class CausVid(BaseModel):
+    def __init__(
+        self,
+        args,
+        device,
+    ):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.num_training_frames = getattr(args, "num_training_frames", 21)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.teacher_forcing = getattr(args, "teacher_forcing", False)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _compute_kl_grad(
+        self,
+        noisy_image_or_video: torch.Tensor,
+        estimated_clean_image_or_video: torch.Tensor,
+        timestep: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        normalization: bool = True,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video.
+            - timestep: a tensor with shape [B, F] containing the randomly generated timestep.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - normalization: a boolean indicating whether to normalize the gradient.
+        Output:
+            - kl_grad: a tensor representing the KL grad.
+            - kl_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Compute the fake score
+        _, pred_fake_image_cond = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+        )
+
+        if self.fake_guidance_scale != 0.0:
+            _, pred_fake_image_uncond = self.fake_score(
+                noisy_image_or_video=noisy_image_or_video,
+                conditional_dict=unconditional_dict,
+                timestep=timestep,
+            )
+            pred_fake_image = (
+                pred_fake_image_cond
+                + (pred_fake_image_cond - pred_fake_image_uncond) * self.fake_guidance_scale
+            )
+        else:
+            pred_fake_image = pred_fake_image_cond
+
+        # Step 2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep,
+        )
+
+        pred_real_image = (
+            pred_real_image_cond
+            + (pred_real_image_cond - pred_real_image_uncond) * self.real_guidance_scale
+        )
+
+        # Step 3: Compute the DMD gradient (DMD paper eq. 7).
+        grad = pred_fake_image - pred_real_image
+
+        # TODO: Change the normalizer for causal teacher
+        if normalization:
+            # Step 4: Gradient normalization (DMD paper eq. 8).
+            p_real = estimated_clean_image_or_video - pred_real_image
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+            grad = grad / normalizer
+        grad = torch.nan_to_num(grad)
+
+        return grad, {
+            "dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(),
+            "timestep": timestep.detach(),
+        }
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        with torch.no_grad():
+            # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+            timestep = self._get_timestep(
+                0,
+                self.num_train_timestep,
+                batch_size,
+                num_frame,
+                self.num_frame_per_block,
+                uniform_timestep=True,
+            )
+
+            if self.timestep_shift > 1:
+                timestep = (
+                    self.timestep_shift
+                    * (timestep / 1000)
+                    / (1 + (self.timestep_shift - 1) * (timestep / 1000))
+                    * 1000
+                )
+            timestep = timestep.clamp(self.min_step, self.max_step)
+
+            noise = torch.randn_like(image_or_video)
+            noisy_latent = (
+                self.scheduler.add_noise(
+                    image_or_video.flatten(0, 1),
+                    noise.flatten(0, 1),
+                    timestep.flatten(0, 1),
+                )
+                .detach()
+                .unflatten(0, (batch_size, num_frame))
+            )
+
+            # Step 2: Compute the KL grad
+            grad, dmd_log_dict = self._compute_kl_grad(
+                noisy_image_or_video=noisy_latent,
+                estimated_clean_image_or_video=original_latent,
+                timestep=timestep,
+                conditional_dict=conditional_dict,
+                unconditional_dict=unconditional_dict,
+            )
+
+        if gradient_mask is not None:
+            dmd_loss = 0.5 * F.mse_loss(
+                original_latent.double()[gradient_mask],
+                (original_latent.double() - grad.double()).detach()[gradient_mask],
+                reduction="mean",
+            )
+        else:
+            dmd_loss = 0.5 * F.mse_loss(
+                original_latent.double(),
+                (original_latent.double() - grad.double()).detach(),
+                reduction="mean",
+            )
+        return dmd_loss, dmd_log_dict
+
+    def _run_generator(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        clean_latent: torch.tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Optionally simulate the generator's input from noise using backward simulation
+        and then run the generator for one-step.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+            - initial_latent: a tensor containing the initial latents [B, F, C, H, W].
+        Output:
+            - pred_image: a tensor with shape [B, F, C, H, W].
+        """
+        simulated_noisy_input = []
+        for timestep in self.denoising_step_list:
+            noise = torch.randn(image_or_video_shape, device=self.device, dtype=self.dtype)
+
+            noisy_timestep = timestep * torch.ones(
+                image_or_video_shape[:2], device=self.device, dtype=torch.long
+            )
+
+            if timestep != 0:
+                noisy_image = self.scheduler.add_noise(
+                    clean_latent.flatten(0, 1),
+                    noise.flatten(0, 1),
+                    noisy_timestep.flatten(0, 1),
+                ).unflatten(0, image_or_video_shape[:2])
+            else:
+                noisy_image = clean_latent
+
+            simulated_noisy_input.append(noisy_image)
+
+        simulated_noisy_input = torch.stack(simulated_noisy_input, dim=1)
+
+        # Step 2: Randomly sample a timestep and pick the corresponding input
+        index = self._get_timestep(
+            0,
+            len(self.denoising_step_list),
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=False,
+        )
+
+        # select the corresponding timestep's noisy input from the stacked tensor [B, T, F, C, H, W]
+        noisy_input = torch.gather(
+            simulated_noisy_input,
+            dim=1,
+            index=index.reshape(index.shape[0], 1, index.shape[1], 1, 1, 1)
+            .expand(-1, -1, -1, *image_or_video_shape[2:])
+            .to(self.device),
+        ).squeeze(1)
+
+        timestep = self.denoising_step_list[index].to(self.device)
+
+        _, pred_image_or_video = self.generator(
+            noisy_image_or_video=noisy_input,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            clean_x=clean_latent if self.teacher_forcing else None,
+        )
+
+        gradient_mask = None  # timestep != 0
+
+        pred_image_or_video = pred_image_or_video.type_as(noisy_input)
+
+        return pred_image_or_video, gradient_mask
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Run generator on backward simulated noisy input
+        pred_image, gradient_mask = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            clean_latent=clean_latent,
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask,
+        )
+
+        # Step 3: TODO: Implement the GAN loss
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _ = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                clean_latent=clean_latent,
+            )
+
+        # Step 2: Compute the fake prediction
+        critic_timestep = self._get_timestep(
+            0,
+            self.num_train_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True,
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = (
+                self.timestep_shift
+                * (critic_timestep / 1000)
+                / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000))
+                * 1000
+            )
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1),
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep,
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1),
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1),
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred,
+        )
+
+        # Step 4: TODO: Compute the GAN loss
+
+        # Step 5: Debugging Log
+        critic_log_dict = {"critic_timestep": critic_timestep.detach()}
+
+        return denoising_loss, critic_log_dict

model/diffusion.py

@@ -0,0 +1,131 @@
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class CausalDiffusion(BaseModel):
+    def __init__(
+        self,
+        args,
+        device,
+    ):
+        """
+        Initialize the Diffusion loss module.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        self.guidance_scale = args.guidance_scale
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.teacher_forcing = getattr(args, "teacher_forcing", False)
+        # Noise augmentation in teacher forcing, we add small noise to clean context latents
+        self.noise_augmentation_max_timestep = getattr(args, "noise_augmentation_max_timestep", 0)
+
+    def _initialize_models(self, args):
+        self.generator = WanDiffusionWrapper(
+            **getattr(args, "model_kwargs", {}),
+            is_causal=True,
+        )
+        self.generator.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        noise = torch.randn_like(clean_latent)
+        batch_size, num_frame = image_or_video_shape[:2]
+
+        # Step 2: Randomly sample a timestep and add noise to denoiser inputs
+        index = self._get_timestep(
+            0,
+            self.scheduler.num_train_timesteps,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=False,
+        )
+        timestep = self.scheduler.timesteps[index].to(dtype=self.dtype, device=self.device)
+        noisy_latents = self.scheduler.add_noise(
+            clean_latent.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep.flatten(0, 1),
+        ).unflatten(0, (batch_size, num_frame))
+        training_target = self.scheduler.training_target(clean_latent, noise, timestep)
+
+        # Step 3: Noise augmentation, also add small noise to clean context latents
+        if self.noise_augmentation_max_timestep > 0:
+            index_clean_aug = self._get_timestep(
+                0,
+                self.noise_augmentation_max_timestep,
+                image_or_video_shape[0],
+                image_or_video_shape[1],
+                self.num_frame_per_block,
+                uniform_timestep=False,
+            )
+            timestep_clean_aug = self.scheduler.timesteps[index_clean_aug].to(
+                dtype=self.dtype, device=self.device
+            )
+            clean_latent_aug = self.scheduler.add_noise(
+                clean_latent.flatten(0, 1),
+                noise.flatten(0, 1),
+                timestep_clean_aug.flatten(0, 1),
+            ).unflatten(0, (batch_size, num_frame))
+        else:
+            clean_latent_aug = clean_latent
+            timestep_clean_aug = None
+
+        # Compute loss
+        flow_pred, x0_pred = self.generator(
+            noisy_image_or_video=noisy_latents,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            clean_x=clean_latent_aug if self.teacher_forcing else None,
+            aug_t=timestep_clean_aug if self.teacher_forcing else None,
+        )
+        # loss = torch.nn.functional.mse_loss(flow_pred.float(), training_target.float())
+        loss = torch.nn.functional.mse_loss(
+            flow_pred.float(), training_target.float(), reduction="none"
+        ).mean(dim=(2, 3, 4))
+        loss = loss * self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame))
+        loss = loss.mean()
+
+        log_dict = {"x0": clean_latent.detach(), "x0_pred": x0_pred.detach()}
+        return loss, log_dict

model/dmd.py

@@ -0,0 +1,377 @@
+from pipeline import SelfForcingTrainingPipeline
+import torch.nn.functional as F
+from typing import Optional, Tuple
+import torch
+
+from model.base import SelfForcingModel
+
+
+class DMD(SelfForcingModel):
+    def __init__(
+        self,
+        args,
+        device,
+    ):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True)
+        self.num_training_frames = getattr(args, "num_training_frames", 21)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: SelfForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+        self.min_score_timestep = getattr(args, "min_score_timestep", 0)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _compute_kl_grad(
+        self,
+        noisy_image_or_video: torch.Tensor,
+        estimated_clean_image_or_video: torch.Tensor,
+        timestep: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        normalization: bool = True,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video.
+            - timestep: a tensor with shape [B, F] containing the randomly generated timestep.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - normalization: a boolean indicating whether to normalize the gradient.
+        Output:
+            - kl_grad: a tensor representing the KL grad.
+            - kl_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Compute the fake score
+        _, pred_fake_image_cond = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+        )
+
+        if self.fake_guidance_scale != 0.0:
+            _, pred_fake_image_uncond = self.fake_score(
+                noisy_image_or_video=noisy_image_or_video,
+                conditional_dict=unconditional_dict,
+                timestep=timestep,
+            )
+            pred_fake_image = (
+                pred_fake_image_cond
+                + (pred_fake_image_cond - pred_fake_image_uncond) * self.fake_guidance_scale
+            )
+        else:
+            pred_fake_image = pred_fake_image_cond
+
+        # Step 2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep,
+        )
+
+        pred_real_image = (
+            pred_real_image_cond
+            + (pred_real_image_cond - pred_real_image_uncond) * self.real_guidance_scale
+        )
+
+        # Step 3: Compute the DMD gradient (DMD paper eq. 7).
+        grad = pred_fake_image - pred_real_image
+
+        # TODO: Change the normalizer for causal teacher
+        if normalization:
+            # Step 4: Gradient normalization (DMD paper eq. 8).
+            p_real = estimated_clean_image_or_video - pred_real_image
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+            grad = grad / normalizer
+        grad = torch.nan_to_num(grad)
+
+        return grad, {
+            "dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(),
+            "timestep": timestep.detach(),
+        }
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: Optional[torch.Tensor] = None,
+        denoised_timestep_from: int = 0,
+        denoised_timestep_to: int = 0,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        with torch.no_grad():
+            # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+            min_timestep = (
+                denoised_timestep_to
+                if self.ts_schedule and denoised_timestep_to is not None
+                else self.min_score_timestep
+            )
+            max_timestep = (
+                denoised_timestep_from
+                if self.ts_schedule_max and denoised_timestep_from is not None
+                else self.num_train_timestep
+            )
+            timestep = self._get_timestep(
+                min_timestep,
+                max_timestep,
+                batch_size,
+                num_frame,
+                self.num_frame_per_block,
+                uniform_timestep=True,
+            )
+
+            # TODO:should we change it to `timestep = self.scheduler.timesteps[timestep]`?
+            if self.timestep_shift > 1:
+                timestep = (
+                    self.timestep_shift
+                    * (timestep / 1000)
+                    / (1 + (self.timestep_shift - 1) * (timestep / 1000))
+                    * 1000
+                )
+            timestep = timestep.clamp(self.min_step, self.max_step)
+
+            noise = torch.randn_like(image_or_video)
+            noisy_latent = (
+                self.scheduler.add_noise(
+                    image_or_video.flatten(0, 1),
+                    noise.flatten(0, 1),
+                    timestep.flatten(0, 1),
+                )
+                .detach()
+                .unflatten(0, (batch_size, num_frame))
+            )
+
+            # Step 2: Compute the KL grad
+            grad, dmd_log_dict = self._compute_kl_grad(
+                noisy_image_or_video=noisy_latent,
+                estimated_clean_image_or_video=original_latent,
+                timestep=timestep,
+                conditional_dict=conditional_dict,
+                unconditional_dict=unconditional_dict,
+            )
+
+        if gradient_mask is not None:
+            dmd_loss = 0.5 * F.mse_loss(
+                original_latent.double()[gradient_mask],
+                (original_latent.double() - grad.double()).detach()[gradient_mask],
+                reduction="mean",
+            )
+        else:
+            dmd_loss = 0.5 * F.mse_loss(
+                original_latent.double(),
+                (original_latent.double() - grad.double()).detach(),
+                reduction="mean",
+            )
+        return dmd_loss, dmd_log_dict
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        (
+            pred_image,
+            gradient_mask,
+            denoised_timestep_from,
+            denoised_timestep_to,
+        ) = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent,
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask,
+            denoised_timestep_from=denoised_timestep_from,
+            denoised_timestep_to=denoised_timestep_to,
+        )
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent,
+            )
+
+        # Step 2: Compute the fake prediction
+        min_timestep = (
+            denoised_timestep_to
+            if self.ts_schedule and denoised_timestep_to is not None
+            else self.min_score_timestep
+        )
+        max_timestep = (
+            denoised_timestep_from
+            if self.ts_schedule_max and denoised_timestep_from is not None
+            else self.num_train_timestep
+        )
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True,
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = (
+                self.timestep_shift
+                * (critic_timestep / 1000)
+                / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000))
+                * 1000
+            )
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1),
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep,
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1),
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1),
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred,
+        )
+
+        # Step 5: Debugging Log
+        critic_log_dict = {"critic_timestep": critic_timestep.detach()}
+
+        return denoising_loss, critic_log_dict

model/gan.py

@@ -0,0 +1,352 @@
+import copy
+from pipeline import SelfForcingTrainingPipeline
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import SelfForcingModel
+
+
+class GAN(SelfForcingModel):
+    def __init__(
+        self,
+        args,
+        device,
+    ):
+        """
+        Initialize the GAN module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True)
+        self.concat_time_embeddings = getattr(args, "concat_time_embeddings", False)
+        self.num_class = args.num_class
+        self.relativistic_discriminator = getattr(args, "relativistic_discriminator", False)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.fake_score.adding_cls_branch(
+            atten_dim=1536,
+            num_class=args.num_class,
+            time_embed_dim=1536 if self.concat_time_embeddings else 0,
+        )
+        self.fake_score.model.requires_grad_(True)
+
+        self.independent_first_frame = getattr(args, "independent_first_frame", False)
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: SelfForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+            self.fake_guidance_scale = args.fake_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+            self.fake_guidance_scale = 0.0
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.critic_timestep_shift = getattr(args, "critic_timestep_shift", self.timestep_shift)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+        self.min_score_timestep = getattr(args, "min_score_timestep", 0)
+
+        self.gan_g_weight = getattr(args, "gan_g_weight", 1e-2)
+        self.gan_d_weight = getattr(args, "gan_d_weight", 1e-2)
+        self.r1_weight = getattr(args, "r1_weight", 0.0)
+        self.r2_weight = getattr(args, "r2_weight", 0.0)
+        self.r1_sigma = getattr(args, "r1_sigma", 0.01)
+        self.r2_sigma = getattr(args, "r2_sigma", 0.01)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def _run_cls_pred_branch(
+        self,
+        noisy_image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Run the classifier prediction branch on the generated image or video.
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W].
+        Output:
+            - cls_pred: a tensor with shape [B, 1, 1, 1, 1] representing the feature map for classification.
+        """
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings,
+        )
+
+        return noisy_logit
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        (
+            pred_image,
+            gradient_mask,
+            denoised_timestep_from,
+            denoised_timestep_to,
+        ) = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent,
+        )
+
+        # Step 2: Get timestep and add noise to generated/real latents
+        min_timestep = (
+            denoised_timestep_to
+            if self.ts_schedule and denoised_timestep_to is not None
+            else self.min_score_timestep
+        )
+        max_timestep = (
+            denoised_timestep_from
+            if self.ts_schedule_max and denoised_timestep_from is not None
+            else self.num_train_timestep
+        )
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True,
+        )
+
+        if self.critic_timestep_shift > 1:
+            critic_timestep = (
+                self.critic_timestep_shift
+                * (critic_timestep / 1000)
+                / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000))
+                * 1000
+            )
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(pred_image)
+        noisy_fake_latent = self.scheduler.add_noise(
+            pred_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1),
+        ).unflatten(0, image_or_video_shape[:2])
+
+        # Step 4: Compute the real GAN discriminator loss
+        real_image_or_video = clean_latent.clone()
+        critic_noise = torch.randn_like(real_image_or_video)
+        noisy_real_latent = self.scheduler.add_noise(
+            real_image_or_video.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1),
+        ).unflatten(0, image_or_video_shape[:2])
+
+        conditional_dict["prompt_embeds"] = torch.concatenate(
+            (
+                conditional_dict["prompt_embeds"],
+                conditional_dict["prompt_embeds"],
+            ),
+            dim=0,
+        )
+        critic_timestep = torch.concatenate((critic_timestep, critic_timestep), dim=0)
+        noisy_latent = torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0)
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=noisy_latent,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep,
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings,
+        )
+        noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)
+
+        if not self.relativistic_discriminator:
+            gan_G_loss = F.softplus(-noisy_fake_logit.float()).mean() * self.gan_g_weight
+        else:
+            relative_fake_logit = noisy_fake_logit - noisy_real_logit
+            gan_G_loss = F.softplus(-relative_fake_logit.float()).mean() * self.gan_g_weight
+
+        return gan_G_loss
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        real_image_or_video: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            (
+                generated_image,
+                _,
+                denoised_timestep_from,
+                denoised_timestep_to,
+                num_sim_steps,
+            ) = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent,
+            )
+
+        # Step 2: Get timestep and add noise to generated/real latents
+        min_timestep = (
+            denoised_timestep_to
+            if self.ts_schedule and denoised_timestep_to is not None
+            else self.min_score_timestep
+        )
+        max_timestep = (
+            denoised_timestep_from
+            if self.ts_schedule_max and denoised_timestep_from is not None
+            else self.num_train_timestep
+        )
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True,
+        )
+
+        if self.critic_timestep_shift > 1:
+            critic_timestep = (
+                self.critic_timestep_shift
+                * (critic_timestep / 1000)
+                / (1 + (self.critic_timestep_shift - 1) * (critic_timestep / 1000))
+                * 1000
+            )
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_fake_latent = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1),
+        ).unflatten(0, image_or_video_shape[:2])
+
+        # Step 4: Compute the real GAN discriminator loss
+        noisy_real_latent = self.scheduler.add_noise(
+            real_image_or_video.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1),
+        ).unflatten(0, image_or_video_shape[:2])
+
+        conditional_dict_cloned = copy.deepcopy(conditional_dict)
+        conditional_dict_cloned["prompt_embeds"] = torch.concatenate(
+            (
+                conditional_dict_cloned["prompt_embeds"],
+                conditional_dict_cloned["prompt_embeds"],
+            ),
+            dim=0,
+        )
+        _, _, noisy_logit = self.fake_score(
+            noisy_image_or_video=torch.concatenate((noisy_fake_latent, noisy_real_latent), dim=0),
+            conditional_dict=conditional_dict_cloned,
+            timestep=torch.concatenate((critic_timestep, critic_timestep), dim=0),
+            classify_mode=True,
+            concat_time_embeddings=self.concat_time_embeddings,
+        )
+        noisy_fake_logit, noisy_real_logit = noisy_logit.chunk(2, dim=0)
+
+        if not self.relativistic_discriminator:
+            gan_D_loss = (
+                F.softplus(-noisy_real_logit.float()).mean()
+                + F.softplus(noisy_fake_logit.float()).mean()
+            )
+        else:
+            relative_real_logit = noisy_real_logit - noisy_fake_logit
+            gan_D_loss = F.softplus(-relative_real_logit.float()).mean()
+        gan_D_loss = gan_D_loss * self.gan_d_weight
+
+        # R1 regularization
+        if self.r1_weight > 0.0:
+            noisy_real_latent_perturbed = noisy_real_latent.clone()
+            epison_real = self.r1_sigma * torch.randn_like(noisy_real_latent_perturbed)
+            noisy_real_latent_perturbed = noisy_real_latent_perturbed + epison_real
+            noisy_real_logit_perturbed = self._run_cls_pred_branch(
+                noisy_image_or_video=noisy_real_latent_perturbed,
+                conditional_dict=conditional_dict,
+                timestep=critic_timestep,
+            )
+
+            r1_grad = (noisy_real_logit_perturbed - noisy_real_logit) / self.r1_sigma
+            r1_loss = self.r1_weight * torch.mean((r1_grad) ** 2)
+        else:
+            r1_loss = torch.zeros_like(gan_D_loss)
+
+        # R2 regularization
+        if self.r2_weight > 0.0:
+            noisy_fake_latent_perturbed = noisy_fake_latent.clone()
+            epison_generated = self.r2_sigma * torch.randn_like(noisy_fake_latent_perturbed)
+            noisy_fake_latent_perturbed = noisy_fake_latent_perturbed + epison_generated
+            noisy_fake_logit_perturbed = self._run_cls_pred_branch(
+                noisy_image_or_video=noisy_fake_latent_perturbed,
+                conditional_dict=conditional_dict,
+                timestep=critic_timestep,
+            )
+
+            r2_grad = (noisy_fake_logit_perturbed - noisy_fake_logit) / self.r2_sigma
+            r2_loss = self.r2_weight * torch.mean((r2_grad) ** 2)
+        else:
+            r2_loss = torch.zeros_like(r2_loss)
+
+        critic_log_dict = {
+            "critic_timestep": critic_timestep.detach(),
+            "noisy_real_logit": noisy_real_logit.detach(),
+            "noisy_fake_logit": noisy_fake_logit.detach(),
+        }
+
+        return (gan_D_loss, r1_loss, r2_loss), critic_log_dict

model/ode_regression.py

@@ -0,0 +1,162 @@
+import torch.nn.functional as F
+from typing import Tuple
+import torch
+
+from model.base import BaseModel
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class ODERegression(BaseModel):
+    def __init__(
+        self,
+        args,
+        device,
+    ):
+        """
+        Initialize the ODERegression module.
+        This class is self-contained and compute generator losses
+        in the forward pass given precomputed ode solution pairs.
+        This class supports the ode regression loss for both causal and bidirectional models.
+        See Sec 4.3 of CausVid https://arxiv.org/abs/2412.07772 for details
+        """
+        super().__init__(args, device)
+
+        # Step 1: Initialize all models
+
+        self.generator = WanDiffusionWrapper(
+            **getattr(args, "model_kwargs", {}),
+            is_causal=True,
+        )
+        self.generator.model.requires_grad_(True)
+        if getattr(args, "generator_ckpt", False):
+            print(f"Loading pretrained generator from {args.generator_ckpt}")
+            state_dict = torch.load(args.generator_ckpt, map_location="cpu")["generator"]
+            self.generator.load_state_dict(state_dict, strict=True)
+
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        self.independent_first_frame = getattr(
+            args,
+            "independent_first_frame",
+            False,
+        )
+        if self.independent_first_frame:
+            self.generator.model.independent_first_frame = True
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+
+        # Step 2: Initialize all hyperparameters
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+
+    def _initialize_models(self, args):
+        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
+        self.generator.model.requires_grad_(True)
+
+        self.text_encoder = WanTextEncoder()
+        self.text_encoder.requires_grad_(False)
+
+        self.vae = WanVAEWrapper()
+        self.vae.requires_grad_(False)
+
+    @torch.no_grad()
+    def _prepare_generator_input(
+        self,
+        ode_latent: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Given a tensor containing the whole ODE sampling trajectories,
+        randomly choose an intermediate timestep and return the latent as well as the corresponding timestep.
+        Input:
+            - ode_latent: a tensor containing the whole ODE sampling trajectories [batch_size, num_denoising_steps, num_frames, num_channels, height, width].
+        Output:
+            - noisy_input: a tensor containing the selected latent [batch_size, num_frames, num_channels, height, width].
+            - timestep: a tensor containing the corresponding timestep [batch_size].
+        """
+        (
+            batch_size,
+            num_denoising_steps,
+            num_frames,
+            num_channels,
+            height,
+            width,
+        ) = ode_latent.shape
+
+        # Step 1: Randomly choose a timestep for each frame
+        index = self._get_timestep(
+            0,
+            len(self.denoising_step_list),
+            batch_size,
+            num_frames,
+            self.num_frame_per_block,
+            uniform_timestep=False,
+        )
+        if self.args.i2v:
+            index[:, 0] = len(self.denoising_step_list) - 1
+
+        noisy_input = torch.gather(
+            ode_latent,
+            dim=1,
+            index=index.reshape(batch_size, 1, num_frames, 1, 1, 1)
+            .expand(-1, -1, -1, num_channels, height, width)
+            .to(self.device),
+        ).squeeze(1)
+
+        timestep = self.denoising_step_list[index].to(self.device)
+
+        # if self.extra_noise_step > 0:
+        #     random_timestep = torch.randint(0, self.extra_noise_step, [
+        #                                     batch_size, num_frames], device=self.device, dtype=torch.long)
+        #     perturbed_noisy_input = self.scheduler.add_noise(
+        #         noisy_input.flatten(0, 1),
+        #         torch.randn_like(noisy_input.flatten(0, 1)),
+        #         random_timestep.flatten(0, 1)
+        #     ).detach().unflatten(0, (batch_size, num_frames)).type_as(noisy_input)
+
+        #     noisy_input[timestep == 0] = perturbed_noisy_input[timestep == 0]
+
+        return noisy_input, timestep
+
+    def generator_loss(
+        self,
+        ode_latent: torch.Tensor,
+        conditional_dict: dict,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noisy latents and compute the ODE regression loss.
+        Input:
+            - ode_latent: a tensor containing the ODE latents [batch_size, num_denoising_steps, num_frames, num_channels, height, width].
+            They are ordered from most noisy to clean latents.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - log_dict: a dictionary containing additional information for loss timestep breakdown.
+        """
+        # Step 1: Run generator on noisy latents
+        target_latent = ode_latent[:, -1]
+
+        noisy_input, timestep = self._prepare_generator_input(ode_latent=ode_latent)
+
+        _, pred_image_or_video = self.generator(
+            noisy_image_or_video=noisy_input,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+        )
+
+        # Step 2: Compute the regression loss
+        mask = timestep != 0
+
+        loss = F.mse_loss(pred_image_or_video[mask], target_latent[mask], reduction="mean")
+
+        log_dict = {
+            "unnormalized_loss": F.mse_loss(pred_image_or_video, target_latent, reduction="none")
+            .mean(dim=[1, 2, 3, 4])
+            .detach(),
+            "timestep": timestep.float().mean(dim=1).detach(),
+            "input": noisy_input.detach(),
+            "output": pred_image_or_video.detach(),
+        }
+
+        return loss, log_dict

model/sid.py

@@ -0,0 +1,318 @@
+from pipeline import SelfForcingTrainingPipeline
+from typing import Optional, Tuple
+import torch
+
+from model.base import SelfForcingModel
+
+
+class SiD(SelfForcingModel):
+    def __init__(
+        self,
+        args,
+        device,
+    ):
+        """
+        Initialize the DMD (Distribution Matching Distillation) module.
+        This class is self-contained and compute generator and fake score losses
+        in the forward pass.
+        """
+        super().__init__(args, device)
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+        if args.gradient_checkpointing:
+            self.generator.enable_gradient_checkpointing()
+            self.fake_score.enable_gradient_checkpointing()
+            self.real_score.enable_gradient_checkpointing()
+
+        # this will be init later with fsdp-wrapped modules
+        self.inference_pipeline: SelfForcingTrainingPipeline = None
+
+        # Step 2: Initialize all dmd hyperparameters
+        self.num_train_timestep = args.num_train_timestep
+        self.min_step = int(0.02 * self.num_train_timestep)
+        self.max_step = int(0.98 * self.num_train_timestep)
+        if hasattr(args, "real_guidance_scale"):
+            self.real_guidance_scale = args.real_guidance_scale
+        else:
+            self.real_guidance_scale = args.guidance_scale
+        self.timestep_shift = getattr(args, "timestep_shift", 1.0)
+        self.sid_alpha = getattr(args, "sid_alpha", 1.0)
+        self.ts_schedule = getattr(args, "ts_schedule", True)
+        self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
+
+        if getattr(self.scheduler, "alphas_cumprod", None) is not None:
+            self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
+        else:
+            self.scheduler.alphas_cumprod = None
+
+    def compute_distribution_matching_loss(
+        self,
+        image_or_video: torch.Tensor,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        gradient_mask: Optional[torch.Tensor] = None,
+        denoised_timestep_from: int = 0,
+        denoised_timestep_to: int = 0,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
+        Input:
+            - image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
+        Output:
+            - dmd_loss: a scalar tensor representing the DMD loss.
+            - dmd_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        original_latent = image_or_video
+
+        batch_size, num_frame = image_or_video.shape[:2]
+
+        # Step 1: Randomly sample timestep based on the given schedule and corresponding noise
+        min_timestep = (
+            denoised_timestep_to
+            if self.ts_schedule and denoised_timestep_to is not None
+            else self.min_score_timestep
+        )
+        max_timestep = (
+            denoised_timestep_from
+            if self.ts_schedule_max and denoised_timestep_from is not None
+            else self.num_train_timestep
+        )
+        timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            batch_size,
+            num_frame,
+            self.num_frame_per_block,
+            uniform_timestep=True,
+        )
+
+        if self.timestep_shift > 1:
+            timestep = (
+                self.timestep_shift
+                * (timestep / 1000)
+                / (1 + (self.timestep_shift - 1) * (timestep / 1000))
+                * 1000
+            )
+        timestep = timestep.clamp(self.min_step, self.max_step)
+
+        noise = torch.randn_like(image_or_video)
+        noisy_latent = self.scheduler.add_noise(
+            image_or_video.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep.flatten(0, 1),
+        ).unflatten(0, (batch_size, num_frame))
+
+        # Step 2: SiD (May be wrap it?)
+        noisy_image_or_video = noisy_latent
+        # Step 2.1: Compute the fake score
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+        )
+        # Step 2.2: Compute the real score
+        # We compute the conditional and unconditional prediction
+        # and add them together to achieve cfg (https://arxiv.org/abs/2207.12598)
+        # NOTE: This step may cause OOM issue, which can be addressed by the CFG-free technique
+
+        _, pred_real_image_cond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+        )
+
+        _, pred_real_image_uncond = self.real_score(
+            noisy_image_or_video=noisy_image_or_video,
+            conditional_dict=unconditional_dict,
+            timestep=timestep,
+        )
+
+        pred_real_image = (
+            pred_real_image_cond
+            + (pred_real_image_cond - pred_real_image_uncond) * self.real_guidance_scale
+        )
+
+        # Step 2.3: SiD Loss
+        # TODO: Add alpha
+        # TODO: Double?
+        sid_loss = (pred_real_image.double() - pred_fake_image.double()) * (
+            (pred_real_image.double() - original_latent.double())
+            - self.sid_alpha * (pred_real_image.double() - pred_fake_image.double())
+        )
+
+        # Step 2.4: Loss normalizer
+        with torch.no_grad():
+            p_real = original_latent - pred_real_image
+            normalizer = torch.abs(p_real).mean(dim=[1, 2, 3, 4], keepdim=True)
+        sid_loss = sid_loss / normalizer
+
+        sid_loss = torch.nan_to_num(sid_loss)
+        num_frame = sid_loss.shape[1]
+        sid_loss = sid_loss.mean()
+
+        sid_log_dict = {
+            "dmdtrain_gradient_norm": torch.zeros_like(sid_loss),
+            "timestep": timestep.detach(),
+        }
+
+        return sid_loss, sid_log_dict
+
+    def generator_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and compute the DMD loss.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - generator_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+        # Step 1: Unroll generator to obtain fake videos
+        (
+            pred_image,
+            gradient_mask,
+            denoised_timestep_from,
+            denoised_timestep_to,
+        ) = self._run_generator(
+            image_or_video_shape=image_or_video_shape,
+            conditional_dict=conditional_dict,
+            initial_latent=initial_latent,
+        )
+
+        # Step 2: Compute the DMD loss
+        dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
+            image_or_video=pred_image,
+            conditional_dict=conditional_dict,
+            unconditional_dict=unconditional_dict,
+            gradient_mask=gradient_mask,
+            denoised_timestep_from=denoised_timestep_from,
+            denoised_timestep_to=denoised_timestep_to,
+        )
+
+        return dmd_loss, dmd_log_dict
+
+    def critic_loss(
+        self,
+        image_or_video_shape,
+        conditional_dict: dict,
+        unconditional_dict: dict,
+        clean_latent: torch.Tensor,
+        initial_latent: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Generate image/videos from noise and train the critic with generated samples.
+        The noisy input to the generator is backward simulated.
+        This removes the need of any datasets during distillation.
+        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
+        Input:
+            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
+            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
+            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
+            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
+        Output:
+            - loss: a scalar tensor representing the generator loss.
+            - critic_log_dict: a dictionary containing the intermediate tensors for logging.
+        """
+
+        # Step 1: Run generator on backward simulated noisy input
+        with torch.no_grad():
+            generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
+                image_or_video_shape=image_or_video_shape,
+                conditional_dict=conditional_dict,
+                initial_latent=initial_latent,
+            )
+
+        # Step 2: Compute the fake prediction
+        min_timestep = (
+            denoised_timestep_to
+            if self.ts_schedule and denoised_timestep_to is not None
+            else self.min_score_timestep
+        )
+        max_timestep = (
+            denoised_timestep_from
+            if self.ts_schedule_max and denoised_timestep_from is not None
+            else self.num_train_timestep
+        )
+        critic_timestep = self._get_timestep(
+            min_timestep,
+            max_timestep,
+            image_or_video_shape[0],
+            image_or_video_shape[1],
+            self.num_frame_per_block,
+            uniform_timestep=True,
+        )
+
+        if self.timestep_shift > 1:
+            critic_timestep = (
+                self.timestep_shift
+                * (critic_timestep / 1000)
+                / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000))
+                * 1000
+            )
+
+        critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
+
+        critic_noise = torch.randn_like(generated_image)
+        noisy_generated_image = self.scheduler.add_noise(
+            generated_image.flatten(0, 1),
+            critic_noise.flatten(0, 1),
+            critic_timestep.flatten(0, 1),
+        ).unflatten(0, image_or_video_shape[:2])
+
+        _, pred_fake_image = self.fake_score(
+            noisy_image_or_video=noisy_generated_image,
+            conditional_dict=conditional_dict,
+            timestep=critic_timestep,
+        )
+
+        # Step 3: Compute the denoising loss for the fake critic
+        if self.args.denoising_loss_type == "flow":
+            from utils.wan_wrapper import WanDiffusionWrapper
+
+            flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
+                scheduler=self.scheduler,
+                x0_pred=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1),
+            )
+            pred_fake_noise = None
+        else:
+            flow_pred = None
+            pred_fake_noise = self.scheduler.convert_x0_to_noise(
+                x0=pred_fake_image.flatten(0, 1),
+                xt=noisy_generated_image.flatten(0, 1),
+                timestep=critic_timestep.flatten(0, 1),
+            ).unflatten(0, image_or_video_shape[:2])
+
+        denoising_loss = self.denoising_loss_func(
+            x=generated_image.flatten(0, 1),
+            x_pred=pred_fake_image.flatten(0, 1),
+            noise=critic_noise.flatten(0, 1),
+            noise_pred=pred_fake_noise,
+            alphas_cumprod=self.scheduler.alphas_cumprod,
+            timestep=critic_timestep.flatten(0, 1),
+            flow_pred=flow_pred,
+        )
+
+        # Step 5: Debugging Log
+        critic_log_dict = {"critic_timestep": critic_timestep.detach()}
+
+        return denoising_loss, critic_log_dict

offline_run.py

@@ -0,0 +1,97 @@
+import argparse
+from pathlib import Path
+
+import torch
+from hydra import compose, initialize
+from hydra.core.global_hydra import GlobalHydra
+from torchvision.io import write_video
+
+from optimize_utils import MultiTrajectory
+from stream_drag_inference_wrapper import StreamDragInferenceWrapper
+from utils.misc import set_seed
+from video_operations import run_optimization, save_videos
+
+
+def build_stream_drag_inference(
+    config_dir: str,
+    config_name: str,
+    checkpoint_path: str,
+    total_generate_block_number: int,
+    use_ema: bool,
+    seed: int,
+) -> StreamDragInferenceWrapper:
+    if GlobalHydra.instance().is_initialized():
+        GlobalHydra.instance().clear()
+
+    with initialize(version_base=None, config_path=config_dir):
+        stream_config = compose(config_name=config_name)
+
+    return StreamDragInferenceWrapper(
+        stream_model_config=stream_config,
+        checkpoint_path=checkpoint_path,
+        total_generate_block_number=total_generate_block_number,
+        use_ema=use_ema,
+        seed=seed,
+    )
+
+
+def main() -> None:
+    prompt_index = 4
+    trajectory_dir = "./saved_labels/self_forcing_dmd_vsink_stream_drag-seed42/0004-A_close-up_3D_animated_scene_of_a_short,_fluffy_mo"
+    start_block_index = 3
+    trajectory_prefix = "block_3_Animation"
+    config_dir = "configs"
+    config_name = "self_forcing_dmd_vsink_stream_drag"
+    checkpoint_path = "./checkpoints/self_forcing_dmd.pt"
+    total_generate_block_number = 36
+    seed = 42
+    fps = 8
+    output_dir = "outputs-editing/self_forcing_dmd_vsink_stream_drag-seed42"
+    use_ema = True
+
+    torch.set_grad_enabled(False)
+
+    trajectory = MultiTrajectory.load(
+        save_dir=trajectory_dir,
+        prefix=trajectory_prefix,
+    )
+
+    model = build_stream_drag_inference(
+        config_dir=config_dir,
+        config_name=config_name,
+        checkpoint_path=checkpoint_path,
+        total_generate_block_number=total_generate_block_number,
+        use_ema=use_ema,
+        seed=seed,
+    )
+
+    set_seed(seed)
+    model.reset()
+    model.inference(
+        start_block_index=0,
+        end_block_index=start_block_index,
+        prompt=trajectory.prompt,
+    )
+
+    all_video, current_video, end_block_index = run_optimization(
+        model=model,
+        trajectory=trajectory,
+        start_block_index=start_block_index,
+    )
+
+    saved_path = save_videos(
+        all_video=all_video,
+        current_video=current_video,
+        output_dir=Path(output_dir),
+        prompt_index=prompt_index,
+        prompt=trajectory.prompt,
+        start_block_index=start_block_index,
+        end_block_index=end_block_index,
+        mode=trajectory.drag_or_animation_select,
+        fps=fps,
+    )
+    print(str(saved_path))
+
+
+if __name__ == "__main__":
+    main()

optimize_utils.py

@@ -0,0 +1,275 @@
+import json
+from pathlib import Path
+from attr import dataclass
+import numpy as np
+import torch
+
+from tensor_utils import (
+    calculate_angle_from_points,
+    read_mask_from_file,
+    save_mask_to_file,
+)
+
+
+@dataclass
+class Trajectory:
+    original_trajectory: dict[str, bool | list[torch.Tensor]] = None
+    """
+    trajectory is dict, keys include 'is_rotation', 'points', if translation also has 'control_points'
+    """
+    block_trajectories: list[dict[str, bool | list[torch.Tensor]]] = []
+    """block_num x trajectory
+    trajectory has keys 'is_rotation' 'deltas' 'start_point'
+    if is_rotation: trajectory also has 'rotation_center'
+    """
+    mask: np.ndarray = None
+    """
+    target mask for the trajectory
+    """
+
+    def __init__(
+        self,
+        original_trajectory: dict[str, bool | list[torch.Tensor]] = None,
+        mask: np.ndarray = None,
+    ):
+        self.original_trajectory = original_trajectory
+        self.mask = mask
+        if original_trajectory is not None:
+            self.block_trajectories = self.original_to_block_trajectories(original_trajectory)
+        else:
+            self.block_trajectories = []
+
+    @staticmethod
+    def original_to_block_trajectories(
+        original_trajectory: dict[str, bool | list[torch.Tensor]],
+        block_length: int = 3,
+    ) -> list[dict[str, bool | list[torch.Tensor]]]:
+        """Convert an original trajectory (with 'points') into per-block trajectories (with 'deltas').
+
+        For translation:
+            deltas[i] = points[i+1] - points[0]   (displacement from start)
+            Each block gets `block_length` consecutive deltas.
+
+        For rotation:
+            points[0] is the rotation center.
+            deltas[i] = angle(center, points[1], points[i+2])
+            Each block gets `block_length` consecutive deltas,
+            plus 'rotation_center' and 'start_point'.
+        """
+        is_rotation = original_trajectory.get("is_rotation", False)
+        points = original_trajectory.get("points", [])
+
+        if is_rotation:
+            # points[0] = rotation center, points[1] = start arm, points[2:] = subsequent arms
+            if len(points) < 2:
+                return []
+            rotation_center = points[0]
+            start_point = points[1]
+            deltas = [
+                calculate_angle_from_points(
+                    rotation_center,
+                    start_point,
+                    point,
+                )
+                for point in points[2:]
+            ]
+        else:
+            # Translation: points[0] = start, points[1:] = subsequent positions
+            if len(points) < 1:
+                return []
+            start_point = points[0]
+            deltas = [torch.Tensor(point) - torch.Tensor(start_point) for point in points[1:]]
+
+        block_trajectories = []
+        for i in range(0, len(deltas), block_length):
+            block_traj = {
+                "is_rotation": is_rotation,
+                "deltas": deltas[i : i + block_length],
+                "start_point": start_point,
+            }
+            if is_rotation:
+                block_traj["rotation_center"] = rotation_center
+            block_trajectories.append(block_traj)
+        return block_trajectories
+
+    def set_original_trajectory(
+        self,
+        original_trajectory: dict[str, bool | list[torch.Tensor]] = None,
+    ):
+        self.original_trajectory = original_trajectory
+        if original_trajectory is not None:
+            self.block_trajectories = self.original_to_block_trajectories(original_trajectory)
+        else:
+            self.block_trajectories = []
+
+    @staticmethod
+    def _serialize_value(
+        v,
+    ):
+        """Recursively serialize a value to JSON-compatible types."""
+        if isinstance(v, torch.Tensor):
+            return v.tolist()
+        elif isinstance(v, np.ndarray):
+            return v.tolist()
+        elif isinstance(v, dict):
+            return {k: Trajectory._serialize_value(val) for k, val in v.items()}
+        elif isinstance(v, list):
+            return [Trajectory._serialize_value(item) for item in v]
+        else:
+            return v
+
+    def to_dict(
+        self,
+        mask_filename: str = None,
+    ) -> dict:
+        """Convert the Trajectory to a JSON-serializable dictionary.
+
+        Args:
+            mask_filename: If provided, store this filename instead of the mask array.
+        """
+        result = {}
+
+        if self.original_trajectory is not None:
+            result["original_trajectory"] = self._serialize_value(self.original_trajectory)
+        else:
+            result["original_trajectory"] = None
+
+        result["block_trajectories"] = self._serialize_value(self.block_trajectories)
+
+        if mask_filename is not None:
+            result["mask_file"] = mask_filename
+
+        return result
+
+    def save_mask(
+        self,
+        save_path: Path,
+    ) -> None:
+        """Save the mask as a PNG image."""
+        if self.mask is not None:
+            save_mask_to_file(self.mask, save_path)
+
+    @staticmethod
+    def load(
+        data: dict,
+        save_dir: Path,
+    ) -> "Trajectory":
+        """Load a Trajectory from a dictionary and directory."""
+        traj = Trajectory()
+        traj.original_trajectory = data.get("original_trajectory", None)
+        traj.block_trajectories = data.get("block_trajectories", [])
+        mask_file = data.get("mask_file", None)
+        if mask_file is not None:
+            traj.mask = read_mask_from_file(save_dir / mask_file)
+        return traj
+
+
+@dataclass
+class MultiTrajectory:
+    block_number: int = 1
+    prompt: str = ""
+    drag_or_animation_select: str = "Drag"
+    trajectories: list[Trajectory] = []
+    """
+    multiple trajectories for a single prompt, each trajectory has its own mask
+    """
+    movable_mask: np.ndarray = None
+    """
+    the movable area mask for the whole image
+    """
+
+    def save(
+        self,
+        save_dir: str | Path,
+        prefix: str = "multi_traj",
+    ) -> Path:
+        """Save the MultiTrajectory to a directory.
+
+        Masks are saved as PNG images, and metadata is saved as a JSON file.
+
+        Args:
+            save_dir: Directory to save files into.
+            prefix: Filename prefix for all saved files.
+
+        Returns:
+            Path to the saved JSON file.
+        """
+        save_dir = Path(save_dir)
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+        result = {
+            "block_number": self.block_number,
+            "prompt": self.prompt,
+            "drag_or_animation_select": self.drag_or_animation_select,
+        }
+
+        # Save movable_mask
+        if self.movable_mask is not None:
+            movable_mask_filename = f"{prefix}_movable_mask.png"
+            save_mask_to_file(self.movable_mask, save_dir / movable_mask_filename)
+            result["movable_area_mask_file"] = movable_mask_filename
+        else:
+            result["movable_area_mask_file"] = None
+
+        # Save each trajectory and its mask
+        traj_dicts = []
+        if self.trajectories is not None:
+            for i, traj in enumerate(self.trajectories):
+                mask_filename = None
+                if traj.mask is not None:
+                    mask_filename = f"{prefix}_traj_{i}_mask.png"
+                    traj.save_mask(save_dir / mask_filename)
+                traj_dicts.append(traj.to_dict(mask_filename=mask_filename))
+        result["trajectories"] = traj_dicts
+
+        # Write JSON
+        json_path = save_dir / f"{prefix}_trajectory.json"
+        with open(json_path, "w") as f:
+            json.dump(result, f, indent=2)
+
+        return json_path
+
+    @staticmethod
+    def load(
+        save_dir: str | Path,
+        prefix: str = "multi_traj",
+    ) -> "MultiTrajectory":
+        """Load a MultiTrajectory from a directory.
+
+        Args:
+            save_dir: Directory containing the saved files.
+            prefix: Filename prefix used when saving.
+
+        Returns:
+            The loaded MultiTrajectory instance.
+        """
+        save_dir = Path(save_dir)
+        json_path = save_dir / f"{prefix}_trajectory.json"
+
+        with open(json_path, "r") as f:
+            data = json.load(f)
+
+        mt = MultiTrajectory()
+        mt.block_number = data.get("block_number", 1)
+        mt.prompt = data.get("prompt", "")
+        mt.drag_or_animation_select = data.get("drag_or_animation_select", "Drag")
+        # Load movable_mask
+        movable_file = data.get("movable_area_mask_file", None)
+        if movable_file is not None:
+            mt.movable_mask = read_mask_from_file(save_dir / movable_file)
+
+        # Load trajectories
+        mt.trajectories = []
+        for traj_data in data.get("trajectories", []):
+            mt.trajectories.append(Trajectory.load(traj_data, save_dir))
+
+        return mt
+
+
+def transpose_dict_2d(d):
+    """Transpose a 2D dict: dict[key1][key2] -> dict[key2][key1]."""
+    result = {}
+    for key1, inner in d.items():
+        for key2, item in inner.items():
+            result.setdefault(key2, {})[key1] = item
+    return result

palette.py

@@ -0,0 +1,774 @@
+import numpy as np
+
+_palette = [
+    0,
+    0,
+    0,
+    128,
+    0,
+    0,
+    0,
+    128,
+    0,
+    128,
+    128,
+    0,
+    0,
+    0,
+    128,
+    128,
+    0,
+    128,
+    0,
+    128,
+    128,
+    128,
+    128,
+    128,
+    64,
+    0,
+    0,
+    191,
+    0,
+    0,
+    64,
+    128,
+    0,
+    191,
+    128,
+    0,
+    64,
+    0,
+    128,
+    191,
+    0,
+    128,
+    64,
+    128,
+    128,
+    191,
+    128,
+    128,
+    0,
+    64,
+    0,
+    128,
+    64,
+    0,
+    0,
+    191,
+    0,
+    128,
+    191,
+    0,
+    0,
+    64,
+    128,
+    128,
+    64,
+    128,
+    22,
+    22,
+    22,
+    23,
+    23,
+    23,
+    24,
+    24,
+    24,
+    25,
+    25,
+    25,
+    26,
+    26,
+    26,
+    27,
+    27,
+    27,
+    28,
+    28,
+    28,
+    29,
+    29,
+    29,
+    30,
+    30,
+    30,
+    31,
+    31,
+    31,
+    32,
+    32,
+    32,
+    33,
+    33,
+    33,
+    34,
+    34,
+    34,
+    35,
+    35,
+    35,
+    36,
+    36,
+    36,
+    37,
+    37,
+    37,
+    38,
+    38,
+    38,
+    39,
+    39,
+    39,
+    40,
+    40,
+    40,
+    41,
+    41,
+    41,
+    42,
+    42,
+    42,
+    43,
+    43,
+    43,
+    44,
+    44,
+    44,
+    45,
+    45,
+    45,
+    46,
+    46,
+    46,
+    47,
+    47,
+    47,
+    48,
+    48,
+    48,
+    49,
+    49,
+    49,
+    50,
+    50,
+    50,
+    51,
+    51,
+    51,
+    52,
+    52,
+    52,
+    53,
+    53,
+    53,
+    54,
+    54,
+    54,
+    55,
+    55,
+    55,
+    56,
+    56,
+    56,
+    57,
+    57,
+    57,
+    58,
+    58,
+    58,
+    59,
+    59,
+    59,
+    60,
+    60,
+    60,
+    61,
+    61,
+    61,
+    62,
+    62,
+    62,
+    63,
+    63,
+    63,
+    64,
+    64,
+    64,
+    65,
+    65,
+    65,
+    66,
+    66,
+    66,
+    67,
+    67,
+    67,
+    68,
+    68,
+    68,
+    69,
+    69,
+    69,
+    70,
+    70,
+    70,
+    71,
+    71,
+    71,
+    72,
+    72,
+    72,
+    73,
+    73,
+    73,
+    74,
+    74,
+    74,
+    75,
+    75,
+    75,
+    76,
+    76,
+    76,
+    77,
+    77,
+    77,
+    78,
+    78,
+    78,
+    79,
+    79,
+    79,
+    80,
+    80,
+    80,
+    81,
+    81,
+    81,
+    82,
+    82,
+    82,
+    83,
+    83,
+    83,
+    84,
+    84,
+    84,
+    85,
+    85,
+    85,
+    86,
+    86,
+    86,
+    87,
+    87,
+    87,
+    88,
+    88,
+    88,
+    89,
+    89,
+    89,
+    90,
+    90,
+    90,
+    91,
+    91,
+    91,
+    92,
+    92,
+    92,
+    93,
+    93,
+    93,
+    94,
+    94,
+    94,
+    95,
+    95,
+    95,
+    96,
+    96,
+    96,
+    97,
+    97,
+    97,
+    98,
+    98,
+    98,
+    99,
+    99,
+    99,
+    100,
+    100,
+    100,
+    101,
+    101,
+    101,
+    102,
+    102,
+    102,
+    103,
+    103,
+    103,
+    104,
+    104,
+    104,
+    105,
+    105,
+    105,
+    106,
+    106,
+    106,
+    107,
+    107,
+    107,
+    108,
+    108,
+    108,
+    109,
+    109,
+    109,
+    110,
+    110,
+    110,
+    111,
+    111,
+    111,
+    112,
+    112,
+    112,
+    113,
+    113,
+    113,
+    114,
+    114,
+    114,
+    115,
+    115,
+    115,
+    116,
+    116,
+    116,
+    117,
+    117,
+    117,
+    118,
+    118,
+    118,
+    119,
+    119,
+    119,
+    120,
+    120,
+    120,
+    121,
+    121,
+    121,
+    122,
+    122,
+    122,
+    123,
+    123,
+    123,
+    124,
+    124,
+    124,
+    125,
+    125,
+    125,
+    126,
+    126,
+    126,
+    127,
+    127,
+    127,
+    128,
+    128,
+    128,
+    129,
+    129,
+    129,
+    130,
+    130,
+    130,
+    131,
+    131,
+    131,
+    132,
+    132,
+    132,
+    133,
+    133,
+    133,
+    134,
+    134,
+    134,
+    135,
+    135,
+    135,
+    136,
+    136,
+    136,
+    137,
+    137,
+    137,
+    138,
+    138,
+    138,
+    139,
+    139,
+    139,
+    140,
+    140,
+    140,
+    141,
+    141,
+    141,
+    142,
+    142,
+    142,
+    143,
+    143,
+    143,
+    144,
+    144,
+    144,
+    145,
+    145,
+    145,
+    146,
+    146,
+    146,
+    147,
+    147,
+    147,
+    148,
+    148,
+    148,
+    149,
+    149,
+    149,
+    150,
+    150,
+    150,
+    151,
+    151,
+    151,
+    152,
+    152,
+    152,
+    153,
+    153,
+    153,
+    154,
+    154,
+    154,
+    155,
+    155,
+    155,
+    156,
+    156,
+    156,
+    157,
+    157,
+    157,
+    158,
+    158,
+    158,
+    159,
+    159,
+    159,
+    160,
+    160,
+    160,
+    161,
+    161,
+    161,
+    162,
+    162,
+    162,
+    163,
+    163,
+    163,
+    164,
+    164,
+    164,
+    165,
+    165,
+    165,
+    166,
+    166,
+    166,
+    167,
+    167,
+    167,
+    168,
+    168,
+    168,
+    169,
+    169,
+    169,
+    170,
+    170,
+    170,
+    171,
+    171,
+    171,
+    172,
+    172,
+    172,
+    173,
+    173,
+    173,
+    174,
+    174,
+    174,
+    175,
+    175,
+    175,
+    176,
+    176,
+    176,
+    177,
+    177,
+    177,
+    178,
+    178,
+    178,
+    179,
+    179,
+    179,
+    180,
+    180,
+    180,
+    181,
+    181,
+    181,
+    182,
+    182,
+    182,
+    183,
+    183,
+    183,
+    184,
+    184,
+    184,
+    185,
+    185,
+    185,
+    186,
+    186,
+    186,
+    187,
+    187,
+    187,
+    188,
+    188,
+    188,
+    189,
+    189,
+    189,
+    190,
+    190,
+    190,
+    191,
+    191,
+    191,
+    192,
+    192,
+    192,
+    193,
+    193,
+    193,
+    194,
+    194,
+    194,
+    195,
+    195,
+    195,
+    196,
+    196,
+    196,
+    197,
+    197,
+    197,
+    198,
+    198,
+    198,
+    199,
+    199,
+    199,
+    200,
+    200,
+    200,
+    201,
+    201,
+    201,
+    202,
+    202,
+    202,
+    203,
+    203,
+    203,
+    204,
+    204,
+    204,
+    205,
+    205,
+    205,
+    206,
+    206,
+    206,
+    207,
+    207,
+    207,
+    208,
+    208,
+    208,
+    209,
+    209,
+    209,
+    210,
+    210,
+    210,
+    211,
+    211,
+    211,
+    212,
+    212,
+    212,
+    213,
+    213,
+    213,
+    214,
+    214,
+    214,
+    215,
+    215,
+    215,
+    216,
+    216,
+    216,
+    217,
+    217,
+    217,
+    218,
+    218,
+    218,
+    219,
+    219,
+    219,
+    220,
+    220,
+    220,
+    221,
+    221,
+    221,
+    222,
+    222,
+    222,
+    223,
+    223,
+    223,
+    224,
+    224,
+    224,
+    225,
+    225,
+    225,
+    226,
+    226,
+    226,
+    227,
+    227,
+    227,
+    228,
+    228,
+    228,
+    229,
+    229,
+    229,
+    230,
+    230,
+    230,
+    231,
+    231,
+    231,
+    232,
+    232,
+    232,
+    233,
+    233,
+    233,
+    234,
+    234,
+    234,
+    235,
+    235,
+    235,
+    236,
+    236,
+    236,
+    237,
+    237,
+    237,
+    238,
+    238,
+    238,
+    239,
+    239,
+    239,
+    240,
+    240,
+    240,
+    241,
+    241,
+    241,
+    242,
+    242,
+    242,
+    243,
+    243,
+    243,
+    244,
+    244,
+    244,
+    245,
+    245,
+    245,
+    246,
+    246,
+    246,
+    247,
+    247,
+    247,
+    248,
+    248,
+    248,
+    249,
+    249,
+    249,
+    250,
+    250,
+    250,
+    251,
+    251,
+    251,
+    252,
+    252,
+    252,
+    253,
+    253,
+    253,
+    254,
+    254,
+    254,
+    255,
+    255,
+    255,
+]
+_palette = np.array(_palette).reshape((-1, 3))
+_palette = _palette.tolist()

pipeline/__init__.py

@@ -0,0 +1,15 @@
+from .bidirectional_diffusion_inference import (
+    BidirectionalDiffusionInferencePipeline,
+)
+from .bidirectional_inference import BidirectionalInferencePipeline
+from .causal_diffusion_inference import CausalDiffusionInferencePipeline
+from .causal_inference import CausalInferencePipeline
+from .self_forcing_training import SelfForcingTrainingPipeline
+
+__all__ = [
+    "BidirectionalDiffusionInferencePipeline",
+    "BidirectionalInferencePipeline",
+    "CausalDiffusionInferencePipeline",
+    "CausalInferencePipeline",
+    "SelfForcingTrainingPipeline",
+]

pipeline/bidirectional_diffusion_inference.py

@@ -0,0 +1,131 @@
+from tqdm import tqdm
+from typing import List
+import torch
+
+from wan.utils.fm_solvers import (
+    FlowDPMSolverMultistepScheduler,
+    get_sampling_sigmas,
+    retrieve_timesteps,
+)
+from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from utils.wan_wrapper import (
+    WanDiffusionWrapper,
+    WanTextEncoder,
+    WanVAEWrapper,
+)
+
+
+class BidirectionalDiffusionInferencePipeline(torch.nn.Module):
+    def __init__(
+        self,
+        args,
+        device,
+        generator=None,
+        text_encoder=None,
+        vae=None,
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        self.generator = (
+            WanDiffusionWrapper(
+                **getattr(args, "model_kwargs", {}),
+                is_causal=False,
+            )
+            if generator is None
+            else generator
+        )
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize scheduler
+        self.num_train_timesteps = args.num_train_timestep
+        self.sampling_steps = 50
+        self.sample_solver = "unipc"
+        self.shift = 8.0
+
+        self.args = args
+
+    def inference(
+        self,
+        noise: torch.Tensor,
+        text_prompts: List[str],
+        return_latents=False,
+    ) -> torch.Tensor:
+        """
+        Perform inference on the given noise and text prompts.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].
+        """
+
+        conditional_dict = self.text_encoder(text_prompts=text_prompts)
+        unconditional_dict = self.text_encoder(
+            text_prompts=[self.args.negative_prompt] * len(text_prompts)
+        )
+
+        latents = noise
+
+        sample_scheduler = self._initialize_sample_scheduler(noise)
+        for _, t in enumerate(tqdm(sample_scheduler.timesteps)):
+            latent_model_input = latents
+            timestep = t * torch.ones(
+                [latents.shape[0], 21], device=noise.device, dtype=torch.float32
+            )
+
+            flow_pred_cond, _ = self.generator(latent_model_input, conditional_dict, timestep)
+            flow_pred_uncond, _ = self.generator(latent_model_input, unconditional_dict, timestep)
+
+            flow_pred = flow_pred_uncond + self.args.guidance_scale * (
+                flow_pred_cond - flow_pred_uncond
+            )
+
+            temp_x0 = sample_scheduler.step(
+                flow_pred.unsqueeze(0),
+                t,
+                latents.unsqueeze(0),
+                return_dict=False,
+            )[0]
+            latents = temp_x0.squeeze(0)
+
+        x0 = latents
+        video = self.vae.decode_to_pixel(x0)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+
+        del sample_scheduler
+
+        if return_latents:
+            return video, latents
+        else:
+            return video
+
+    def _initialize_sample_scheduler(
+        self,
+        noise,
+    ):
+        if self.sample_solver == "unipc":
+            sample_scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False,
+            )
+            sample_scheduler.set_timesteps(
+                self.sampling_steps, device=noise.device, shift=self.shift
+            )
+            self.timesteps = sample_scheduler.timesteps
+        elif self.sample_solver == "dpm++":
+            sample_scheduler = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False,
+            )
+            sampling_sigmas = get_sampling_sigmas(self.sampling_steps, self.shift)
+            self.timesteps, _ = retrieve_timesteps(
+                sample_scheduler, device=noise.device, sigmas=sampling_sigmas
+            )
+        else:
+            raise NotImplementedError("Unsupported solver.")
+        return sample_scheduler

pipeline/bidirectional_inference.py

@@ -0,0 +1,85 @@
+from typing import List
+import torch
+
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
+
+
+class BidirectionalInferencePipeline(torch.nn.Module):
+    def __init__(
+        self,
+        args,
+        device,
+        generator=None,
+        text_encoder=None,
+        vae=None,
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        self.generator = (
+            WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=False)
+            if generator is None
+            else generator
+        )
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize all bidirectional wan hyperparmeters
+        self.scheduler = self.generator.get_scheduler()
+        self.denoising_step_list = torch.tensor(
+            args.denoising_step_list, dtype=torch.long, device=device
+        )
+        if self.denoising_step_list[-1] == 0:
+            self.denoising_step_list = self.denoising_step_list[
+                :-1
+            ]  # remove the zero timestep for inference
+        if args.warp_denoising_step:
+            timesteps = torch.cat(
+                (
+                    self.scheduler.timesteps.cpu(),
+                    torch.tensor([0], dtype=torch.float32),
+                )
+            )
+            self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+    def inference(
+        self,
+        noise: torch.Tensor,
+        text_prompts: List[str],
+    ) -> torch.Tensor:
+        """
+        Perform inference on the given noise and text prompts.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].
+        """
+        conditional_dict = self.text_encoder(text_prompts=text_prompts)
+
+        # initial point
+        noisy_image_or_video = noise
+
+        # use the last n-1 timesteps to simulate the generator's input
+        for index, current_timestep in enumerate(self.denoising_step_list[:-1]):
+            _, pred_image_or_video = self.generator(
+                noisy_image_or_video=noisy_image_or_video,
+                conditional_dict=conditional_dict,
+                timestep=torch.ones(noise.shape[:2], dtype=torch.long, device=noise.device)
+                * current_timestep,
+            )  # [B, F, C, H, W]
+
+            next_timestep = self.denoising_step_list[index + 1] * torch.ones(
+                noise.shape[:2], dtype=torch.long, device=noise.device
+            )
+
+            noisy_image_or_video = self.scheduler.add_noise(
+                pred_image_or_video.flatten(0, 1),
+                torch.randn_like(pred_image_or_video.flatten(0, 1)),
+                next_timestep.flatten(0, 1),
+            ).unflatten(0, noise.shape[:2])
+
+        video = self.vae.decode_to_pixel(pred_image_or_video)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+        return video

pipeline/causal_diffusion_inference.py

@@ -0,0 +1,402 @@
+from tqdm import tqdm
+from typing import List, Optional
+import torch
+
+from wan.utils.fm_solvers import (
+    FlowDPMSolverMultistepScheduler,
+    get_sampling_sigmas,
+    retrieve_timesteps,
+)
+from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from utils.wan_wrapper import (
+    WanDiffusionWrapper,
+    WanTextEncoder,
+    WanVAEWrapper,
+)
+
+
+class CausalDiffusionInferencePipeline(torch.nn.Module):
+    def __init__(
+        self,
+        args,
+        device,
+        generator=None,
+        text_encoder=None,
+        vae=None,
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        self.generator = (
+            WanDiffusionWrapper(
+                **getattr(args, "model_kwargs", {}),
+                is_causal=True,
+            )
+            if generator is None
+            else generator
+        )
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize scheduler
+        self.num_train_timesteps = args.num_train_timestep
+        self.sampling_steps = 50
+        self.sample_solver = "unipc"
+        self.shift = args.timestep_shift
+
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+
+        self.kv_cache_pos = None
+        self.kv_cache_neg = None
+        self.crossattn_cache_pos = None
+        self.crossattn_cache_neg = None
+        self.args = args
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.independent_first_frame = args.independent_first_frame
+        self.local_attn_size = self.generator.model.local_attn_size
+
+        print(f"KV inference with {self.num_frame_per_block} frames per block")
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+    def inference(
+        self,
+        noise: torch.Tensor,
+        text_prompts: List[str],
+        initial_latent: Optional[torch.Tensor] = None,
+        return_latents: bool = False,
+        start_frame_index: Optional[int] = 0,
+    ) -> torch.Tensor:
+        """
+        Perform inference on the given noise and text prompts.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+            initial_latent (torch.Tensor): The initial latent tensor of shape
+                (batch_size, num_input_frames, num_channels, height, width).
+                If num_input_frames is 1, perform image to video.
+                If num_input_frames is greater than 1, perform video extension.
+            return_latents (bool): Whether to return the latents.
+            start_frame_index (int): In long video generation, where does the current window start?
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].
+        """
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (
+            self.independent_first_frame and initial_latent is not None
+        ):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        elif self.independent_first_frame and initial_latent is None:
+            # Using a [1, 4, 4, 4, 4, 4] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        conditional_dict = self.text_encoder(text_prompts=text_prompts)
+        unconditional_dict = self.text_encoder(
+            text_prompts=[self.args.negative_prompt] * len(text_prompts)
+        )
+
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype,
+        )
+
+        # Step 1: Initialize KV cache to all zeros
+        if self.kv_cache_pos is None:
+            self._initialize_kv_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device)
+            self._initialize_crossattn_cache(
+                batch_size=batch_size, dtype=noise.dtype, device=noise.device
+            )
+        else:
+            # reset cross attn cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache_pos[block_index]["is_init"] = False
+                self.crossattn_cache_neg[block_index]["is_init"] = False
+            # reset kv cache
+            for block_index in range(len(self.kv_cache_pos)):
+                self.kv_cache_pos[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device
+                )
+                self.kv_cache_pos[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device
+                )
+                self.kv_cache_neg[block_index]["global_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device
+                )
+                self.kv_cache_neg[block_index]["local_end_index"] = torch.tensor(
+                    [0], dtype=torch.long, device=noise.device
+                )
+
+        # Step 2: Cache context feature
+        current_start_frame = start_frame_index
+        cache_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            if self.independent_first_frame:
+                # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+                assert (num_input_frames - 1) % self.num_frame_per_block == 0
+                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
+                output[:, :1] = initial_latent[:, :1]
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                self.generator(
+                    noisy_image_or_video=initial_latent[:, :1],
+                    conditional_dict=unconditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_neg,
+                    crossattn_cache=self.crossattn_cache_neg,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += 1
+                cache_start_frame += 1
+            else:
+                # Assume num_input_frames is self.num_frame_per_block * num_input_blocks
+                assert num_input_frames % self.num_frame_per_block == 0
+                num_input_blocks = num_input_frames // self.num_frame_per_block
+
+            for block_index in range(num_input_blocks):
+                current_ref_latents = initial_latent[
+                    :,
+                    cache_start_frame : cache_start_frame + self.num_frame_per_block,
+                ]
+                output[
+                    :,
+                    cache_start_frame : cache_start_frame + self.num_frame_per_block,
+                ] = current_ref_latents
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                self.generator(
+                    noisy_image_or_video=current_ref_latents,
+                    conditional_dict=unconditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache_neg,
+                    crossattn_cache=self.crossattn_cache_neg,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                current_start_frame += self.num_frame_per_block
+                cache_start_frame += self.num_frame_per_block
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        for current_num_frames in all_num_frames:
+            noisy_input = noise[
+                :,
+                cache_start_frame
+                - num_input_frames : cache_start_frame
+                + current_num_frames
+                - num_input_frames,
+            ]
+            latents = noisy_input
+
+            # Step 3.1: Spatial denoising loop
+            sample_scheduler = self._initialize_sample_scheduler(noise)
+            for _, t in enumerate(tqdm(sample_scheduler.timesteps)):
+                latent_model_input = latents
+                timestep = t * torch.ones(
+                    [batch_size, current_num_frames],
+                    device=noise.device,
+                    dtype=torch.float32,
+                )
+
+                flow_pred_cond, _ = self.generator(
+                    noisy_image_or_video=latent_model_input,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache_pos,
+                    crossattn_cache=self.crossattn_cache_pos,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+                flow_pred_uncond, _ = self.generator(
+                    noisy_image_or_video=latent_model_input,
+                    conditional_dict=unconditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache_neg,
+                    crossattn_cache=self.crossattn_cache_neg,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    cache_start=cache_start_frame * self.frame_seq_length,
+                )
+
+                flow_pred = flow_pred_uncond + self.args.guidance_scale * (
+                    flow_pred_cond - flow_pred_uncond
+                )
+
+                temp_x0 = sample_scheduler.step(flow_pred, t, latents, return_dict=False)[0]
+                latents = temp_x0
+                print(f"kv_cache['local_end_index']: {self.kv_cache_pos[0]['local_end_index']}")
+                print(f"kv_cache['global_end_index']: {self.kv_cache_pos[0]['global_end_index']}")
+
+            # Step 3.2: record the model's output
+            output[:, cache_start_frame : cache_start_frame + current_num_frames] = latents
+
+            # Step 3.3: rerun with timestep zero to update KV cache using clean context
+            self.generator(
+                noisy_image_or_video=latents,
+                conditional_dict=conditional_dict,
+                timestep=timestep * 0,
+                kv_cache=self.kv_cache_pos,
+                crossattn_cache=self.crossattn_cache_pos,
+                current_start=current_start_frame * self.frame_seq_length,
+                cache_start=cache_start_frame * self.frame_seq_length,
+            )
+            self.generator(
+                noisy_image_or_video=latents,
+                conditional_dict=unconditional_dict,
+                timestep=timestep * 0,
+                kv_cache=self.kv_cache_neg,
+                crossattn_cache=self.crossattn_cache_neg,
+                current_start=current_start_frame * self.frame_seq_length,
+                cache_start=cache_start_frame * self.frame_seq_length,
+            )
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+            cache_start_frame += current_num_frames
+
+        # Step 4: Decode the output
+        video = self.vae.decode_to_pixel(output)
+        video = (video * 0.5 + 0.5).clamp(0, 1)
+
+        if return_latents:
+            return video, output
+        else:
+            return video
+
+    def _initialize_kv_cache(
+        self,
+        batch_size,
+        dtype,
+        device,
+    ):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache_pos = []
+        kv_cache_neg = []
+        if self.local_attn_size != -1:
+            # Use the local attention size to compute the KV cache size
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            # Use the default KV cache size
+            kv_cache_size = 32760
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache_pos.append(
+                {
+                    "k": torch.zeros(
+                        [batch_size, kv_cache_size, 12, 128],
+                        dtype=dtype,
+                        device=device,
+                    ),
+                    "v": torch.zeros(
+                        [batch_size, kv_cache_size, 12, 128],
+                        dtype=dtype,
+                        device=device,
+                    ),
+                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                }
+            )
+            kv_cache_neg.append(
+                {
+                    "k": torch.zeros(
+                        [batch_size, kv_cache_size, 12, 128],
+                        dtype=dtype,
+                        device=device,
+                    ),
+                    "v": torch.zeros(
+                        [batch_size, kv_cache_size, 12, 128],
+                        dtype=dtype,
+                        device=device,
+                    ),
+                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                }
+            )
+
+        self.kv_cache_pos = kv_cache_pos  # always store the clean cache
+        self.kv_cache_neg = kv_cache_neg  # always store the clean cache
+
+    def _initialize_crossattn_cache(
+        self,
+        batch_size,
+        dtype,
+        device,
+    ):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache_pos = []
+        crossattn_cache_neg = []
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache_pos.append(
+                {
+                    "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                    "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                    "is_init": False,
+                }
+            )
+            crossattn_cache_neg.append(
+                {
+                    "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                    "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                    "is_init": False,
+                }
+            )
+
+        self.crossattn_cache_pos = crossattn_cache_pos  # always store the clean cache
+        self.crossattn_cache_neg = crossattn_cache_neg  # always store the clean cache
+
+    def _initialize_sample_scheduler(
+        self,
+        noise,
+    ):
+        if self.sample_solver == "unipc":
+            sample_scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False,
+            )
+            sample_scheduler.set_timesteps(
+                self.sampling_steps, device=noise.device, shift=self.shift
+            )
+            self.timesteps = sample_scheduler.timesteps
+        elif self.sample_solver == "dpm++":
+            sample_scheduler = FlowDPMSolverMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False,
+            )
+            sampling_sigmas = get_sampling_sigmas(self.sampling_steps, self.shift)
+            self.timesteps, _ = retrieve_timesteps(
+                sample_scheduler, device=noise.device, sigmas=sampling_sigmas
+            )
+        else:
+            raise NotImplementedError("Unsupported solver.")
+        return sample_scheduler

pipeline/causal_inference.py

@@ -0,0 +1,1193 @@
+import gc
+import random
+import time
+from typing import List, Optional
+import numpy as np
+from omegaconf import DictConfig, OmegaConf
+import torch
+
+import torch.nn.functional as F
+from torch.nn.utils import clip_grad_norm_
+from tqdm import tqdm
+
+from optimize_utils import transpose_dict_2d
+from tensor_utils import (
+    build_anisotropic_gaussian_from_mask,
+    combine_gaussian_maps,
+    combine_masks_or,
+    normalize_tensor_to_match_tensor,
+    resize_tensor,
+    warp_tensor,
+    warp_tensor_sequence,
+)
+from utils.wan_wrapper import (
+    WanDiffusionWrapper,
+    WanTextEncoder,
+    WanVAEWrapper,
+)
+
+from demo_utils.memory import (
+    gpu,
+    get_cuda_free_memory_gb,
+    DynamicSwapInstaller,
+    move_model_to_device_with_memory_preservation,
+)
+
+IMAGE_HEIGHT = 480.0
+
+
+def split_trajectories_segments(
+    trajectories: list[dict[str, bool | list[torch.Tensor]]],
+    translation_step: float,
+    rotation_step: float,
+) -> List[list[dict[str, bool | list[torch.Tensor]]]]:
+    """
+    Split drag trajectories into evenly spaced intermediate segments for
+    progressive (coarse-to-fine) optimization.
+
+    Given N trajectories (each with per-frame deltas), this function:
+      1. Determines the maximum number of segments needed across all
+         trajectories based on the magnitude of their deltas and the
+         provided step sizes.
+      2. Divides every trajectory's deltas uniformly into that many segments.
+      3. Produces a list of cumulative intermediate trajectory snapshots,
+         where segment k contains deltas scaled by (k / max_segments).
+
+    :param trajectories:
+        N x trajectory dicts.
+        Each dict has keys:
+          - 'is_rotation' (bool): Whether this trajectory is a rotation.
+          - 'deltas' (list): Per-frame displacement values.
+              For translation: each delta is a 2D vector (dx, dy).
+              For rotation: each delta is a scalar angle.
+          - 'start_point': The starting pixel coordinate of the drag.
+          - 'rotation_center' (only if is_rotation): The center of rotation.
+
+    :param translation_step:
+        The pixel distance that defines one segment for translation
+        trajectories. Larger values produce fewer segments.
+
+    :param rotation_step:
+        The angle (in the same units as deltas) that defines one segment
+        for rotation trajectories. Larger values produce fewer segments.
+
+    :returns:
+        segment_num x N x trajectory dicts.
+        A list of length `max_segment_number`, where each element is a list
+        of N trajectory dicts. The k-th element (1-indexed) contains
+        trajectories whose deltas are scaled to (k / max_segment_number)
+        of the original deltas — i.e., cumulative intermediate waypoints.
+    """
+    # -------------------------------------------------------------------------
+    # Phase 1: Convert raw deltas to torch tensors (ensure uniform type)
+    # -------------------------------------------------------------------------
+    for trajectory in trajectories:
+        trajectory["deltas"] = [torch.tensor(delta, device="cpu") for delta in trajectory["deltas"]]
+
+    # -------------------------------------------------------------------------
+    # Phase 2: Determine the maximum number of segments across all trajectories.
+    #   - For rotations: segment count = |angle_delta| // rotation_step
+    #   - For translations: segment count = ||displacement_delta||₂ // translation_step
+    #   - We take the global maximum so every trajectory is split into the
+    #     same number of segments (ensuring synchronized progressive steps).
+    # -------------------------------------------------------------------------
+    max_segment_number = 1  # at least one segment
+    for trajectory in trajectories:
+        print(f"{trajectory['is_rotation'] = }")
+        for delta in trajectory["deltas"]:
+            if trajectory["is_rotation"]:
+                magnitude = abs(delta)
+                step = rotation_step
+            else:
+                magnitude = abs(torch.norm(delta))
+                step = translation_step
+            segment_number = int(magnitude // step)
+            print(f"{delta = } {magnitude = } {segment_number = }")
+            max_segment_number = max(max_segment_number, segment_number)
+    print(f"{max_segment_number = }")
+
+    # -------------------------------------------------------------------------
+    # Phase 3: Compute per-segment step sizes for each trajectory.
+    #   Each trajectory's deltas are divided by max_segment_number to get
+    #   the uniform per-segment increment.
+    # -------------------------------------------------------------------------
+    split_trajectory_steps = []
+    for trajectory in trajectories:
+        print(f"{trajectory['is_rotation'] = }")
+        # Divide each frame's delta by the total number of segments
+        trajectory_steps = [delta / float(max_segment_number) for delta in trajectory["deltas"]]
+        print(f"{trajectory_steps = }")
+        # Build the per-trajectory step metadata
+        split_trajectory_step = {
+            "is_rotation": trajectory["is_rotation"],
+            "steps": trajectory_steps,  # per-segment increment per frame
+            "start_point": trajectory["start_point"],
+        }
+        if trajectory["is_rotation"]:
+            split_trajectory_step["rotation_center"] = trajectory["rotation_center"]
+        split_trajectory_steps.append(split_trajectory_step)
+
+    # -------------------------------------------------------------------------
+    # Phase 4: Build cumulative intermediate trajectory lists.
+    #   For segment_index k (1-indexed from 1 to max_segment_number):
+    #     delta_k = step * k
+    #   This produces progressively larger displacements, enabling the
+    #   optimizer to move features gradually toward the final target.
+    # -------------------------------------------------------------------------
+    new_trajectories_list = []
+    for segment_index in range(max_segment_number):
+        segment_index += 1  # 1-indexed: cumulative scale factor
+        new_trajectories = []
+        for trajectory_step in split_trajectory_steps:
+            new_trajectory = {
+                "is_rotation": trajectory_step["is_rotation"],
+                "deltas": [step * segment_index for step in trajectory_step["steps"]],
+                "start_point": trajectory_step["start_point"],
+            }
+            if trajectory_step["is_rotation"]:
+                new_trajectory["rotation_center"] = trajectory_step["rotation_center"]
+            new_trajectories.append(new_trajectory)
+        print(f"{new_trajectories = }")
+        new_trajectories_list.append(new_trajectories)
+
+    # Return: list of length max_segment_number, each containing N trajectory dicts
+    # with cumulatively scaled deltas (segment 1 = smallest, last = full original delta)
+    return new_trajectories_list
+
+
+class CausalInferencePipeline(torch.nn.Module):
+    def __init__(
+        self,
+        args,
+        device,
+        generator=None,
+        text_encoder=None,
+        vae=None,
+    ):
+        super().__init__()
+        # Step 1: Initialize all models
+        self.generator = (
+            WanDiffusionWrapper(
+                **getattr(args, "model_kwargs", {}),
+                is_causal=True,
+            )
+            if generator is None
+            else generator
+        )
+        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
+        self.vae = WanVAEWrapper() if vae is None else vae
+
+        # Step 2: Initialize all causal hyperparmeters
+        self.scheduler = self.generator.get_scheduler()
+        self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long)
+        if args.warp_denoising_step:
+            timesteps = torch.cat(
+                (
+                    self.scheduler.timesteps.cpu(),
+                    torch.tensor([0], dtype=torch.float32),
+                )
+            )
+            self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
+
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+
+        self.kv_cache1 = None
+        self.args = args
+        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
+        self.independent_first_frame = args.independent_first_frame
+        self.local_attn_size = self.generator.model.local_attn_size
+
+        print(f"KV inference with {self.num_frame_per_block} frames per block")
+
+        if self.num_frame_per_block > 1:
+            self.generator.model.num_frame_per_block = self.num_frame_per_block
+
+    def stack_features(
+        self,
+        record_features: dict[int, torch.Tensor],
+        # dict[block_index] -> Tensor[1, 3, 30, 52, 1536]
+    ):
+        """
+        Reorder and stack per-block attention features into one tensor.
+
+        Input:
+            record_features: Mapping `block_index -> Tensor` with shape
+                `[1, 3, 30, 52, 1536]`.
+
+        Returns:
+            Tensor with shape `[1536 * L, 3, 30, 52]`,
+            where `L = len(record_features)`.
+        """
+        attention_values = {
+            k: v.permute(0, 4, 1, 2, 3).contiguous() for k, v in record_features.items()
+        }
+        # [ [1, 1536, 3, 30, 52] ]
+
+        features = []
+        for fet in attention_values.values():
+            # [1536, 3, 30, 52]
+            features.append(fet.squeeze(0).contiguous())
+        del fet
+
+        # Concatenate all features
+        feature = torch.cat(features, dim=0)
+        # [1536 * L, 3, 30, 52], L is the number of blocks
+        return feature
+
+    def generate_features(
+        self,
+        latents,
+        conditional_dict,
+        timestep,
+        kv_cache,
+        crossattn_cache,
+        current_start,
+        # below are for drag optimization
+        model_config: DictConfig = None,
+    ):
+        """
+        Run one generator forward pass and return prediction + stacked features.
+
+        Notes:
+        - KV cache is deep-cloned/detached before forward to avoid mutating
+          the caller's cache during optimization.
+        - Forward runs under CUDA bfloat16 autocast.
+        - Returned `record_features` are converted via `stack_features(...)`.
+
+        Returns:
+            denoised_pred:
+                Model denoised prediction tensor.
+            record_features:
+                Dict `variant_key -> Tensor[1536 * L, 3, 30, 52]`. L is the number of blocks.
+        """
+        temp_kv_cache = [
+            {
+                "k": kv_cache[block_index]["k"].clone().detach(),
+                "v": kv_cache[block_index]["v"].clone().detach(),
+                "global_end_index": kv_cache[block_index]["global_end_index"].clone().detach(),
+                "local_end_index": kv_cache[block_index]["local_end_index"].clone().detach(),
+            }
+            for block_index in range(self.num_transformer_blocks)
+        ]
+        # print(f"{temp_kv_cache[0]['k'].shape = }")
+        # Forward pass through the transformer with user-specified autocast dtype
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            # Extract features during forward pass
+            _, denoised_pred = self.generator(
+                noisy_image_or_video=latents,
+                conditional_dict=conditional_dict,
+                timestep=timestep,
+                kv_cache=temp_kv_cache,
+                crossattn_cache=crossattn_cache,
+                current_start=current_start,
+                # below are for drag optimization
+                model_config=model_config,
+            )
+            denoised_pred, record_features = denoised_pred
+        # record_features: dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536]
+
+        record_features = transpose_dict_2d(record_features)
+        # record_features: Dict[key] -> Dict[block_index -> Tensor] [1, 3, 30, 52, 1536]
+        record_features = {k: self.stack_features(v) for k, v in record_features.items()}
+        # record_features: Dict[key] -> Tensor [1536 * L, 3, 30, 52]
+        return denoised_pred, record_features
+
+    @staticmethod
+    def precompute_conditions(
+        trajectories: list[dict[str, bool | list[torch.Tensor]]],
+        # N x trajectory
+        masks: list[np.ndarray],  # N x mask
+        dtype: torch.dtype,
+        device: torch.device,
+        model_config: DictConfig,
+        downsample_movable_mask: torch.Tensor,
+        previous_record_features: dict[int, dict[str, torch.Tensor]],
+        latent_spatial_size: tuple[int, int],  # (H_lat, W_lat)
+    ):
+        with torch.no_grad():
+            _any_prev_feat = next(iter(previous_record_features.values()))
+            feat_spatial_size = tuple(_any_prev_feat.shape[-2:])  # (Hf, Wf)
+
+            # 1) Warped masks in image space
+            warped_masks = [
+                warp_tensor_sequence(
+                    tensor=torch.tensor(mask, device=device).float(),
+                    is_rotation=trajectory["is_rotation"],
+                    deltas=trajectory["deltas"],
+                    rotation_center=trajectory.get("rotation_center", None),
+                    original_height=IMAGE_HEIGHT,
+                    mode="nearest",
+                    cumulative=False,
+                )
+                for trajectory, mask in zip(trajectories, masks)
+            ]  # N x frame x [H_img, W_img]
+            # print(f"{len(warped_masks) = }, {len(warped_masks[0]) = }, {warped_masks[0][0].shape = }")
+
+            # 2) Downsampled warped masks in feature space
+            down_warp_masks = [
+                [
+                    resize_tensor(
+                        warped_mask.detach(),
+                        size=feat_spatial_size,
+                        mode="nearest",
+                    ).detach()
+                    for warped_mask in traj_warped_masks
+                ]
+                for traj_warped_masks in warped_masks
+            ]  # N x frame x [Hf, Wf]
+
+            # 3) Gaussian heatmaps per trajectory
+            gaussian_heatmaps_per_traj = [
+                [
+                    build_anisotropic_gaussian_from_mask(
+                        warped_mask,
+                        padding_scale=model_config.drag_optim_config.gradient_gaussian_padding,
+                        sigma_scale=model_config.drag_optim_config.gradient_gaussian_sigma,
+                    ).detach()
+                    for warped_mask in traj_warped_masks
+                ]
+                for traj_warped_masks in warped_masks
+            ]  # N x frame x  [H_img, W_img]
+
+            # 4) Combined downsampled movable mask (OR of all warped + original)
+            all_down_warp_masks = [
+                dwm for traj_down_masks in down_warp_masks for dwm in traj_down_masks
+            ]
+            all_down_warp_masks.append(downsample_movable_mask.clone())
+            combined_downsample_movable_mask = combine_masks_or(all_down_warp_masks)
+
+            # 5) Precompute warped attention values per variant
+            warped_for_prev: dict[str | int, List[List[torch.Tensor]]] = {
+                key: [
+                    warp_tensor_sequence(
+                        tensor=prev_feat.to(dtype=dtype, device=device),
+                        is_rotation=trajectory["is_rotation"],
+                        deltas=[
+                            d.to(
+                                dtype=dtype,
+                                device=device,
+                            )
+                            for d in trajectory["deltas"]
+                        ],
+                        rotation_center=trajectory.get("rotation_center", None),
+                        original_height=IMAGE_HEIGHT,
+                        mode="nearest",
+                        cumulative=False,
+                    )
+                    for trajectory in trajectories
+                ]
+                for key, prev_feat in previous_record_features.items()
+            }
+            # warped_for_prev: dict[key] -> list[traj_index] -> list[frame_index] -> Tensor [1536 * L, 30 * scaling, 52 * scaling]
+
+            # 6) Combined Gaussian heatmaps in latent space
+            combined_gaussian_heatmaps = None
+            num_frames = len(gaussian_heatmaps_per_traj[0])
+            num_trajs = len(gaussian_heatmaps_per_traj)
+            combined_gaussian_heatmaps = torch.stack(
+                [
+                    combine_gaussian_maps(
+                        [
+                            gaussian_heatmaps_per_traj[traj_idx][frame_idx]
+                            for traj_idx in range(num_trajs)
+                        ]
+                    )  # [H_img, W_img]
+                    for frame_idx in range(num_frames)
+                ],
+                dim=0,
+            ).to(
+                device=device
+            )  # [F, H_img, W_img]
+
+            combined_gaussian_heatmaps = resize_tensor(
+                combined_gaussian_heatmaps,
+                size=latent_spatial_size,
+                mode="bilinear",
+            ).detach()  # [F, H_lat, W_lat]
+            combined_gaussian_heatmaps = combined_gaussian_heatmaps.to(dtype=dtype)
+        return (
+            warped_masks,
+            down_warp_masks,
+            gaussian_heatmaps_per_traj,
+            combined_downsample_movable_mask,
+            warped_for_prev,
+            combined_gaussian_heatmaps,
+        )
+
+    def optimize_latent(
+        self,
+        latents,
+        conditional_dict,
+        timestep,
+        kv_cache,
+        crossattn_cache,
+        current_start,
+        # below are for drag optimization
+        trajectories: list[dict[str, bool | list[torch.Tensor]]],
+        # N x trajectory,
+        # trajectory has keys 'is_rotation' 'deltas' 'start_point'
+        # if is_rotation: trajectory also has 'rotation_center'
+        masks: list[np.ndarray],  # N x mask
+        movable_mask: np.ndarray,
+        clean_previous_record_feature: dict[int, dict[str, torch.Tensor]],
+        # dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536]
+        noisy_previous_record_feature: dict[int, dict[str, torch.Tensor]],
+        # dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536]
+        model_config: DictConfig,
+        optimize_target_latent_index: int = -1,
+    ):
+        """
+        :param trajectories:
+            N x trajectory,
+            trajectory has keys 'is_rotation' 'deltas' 'start_point'
+            if is_rotation: trajectory also has 'rotation_center'
+        :param masks:
+            N x mask
+        """
+        assert isinstance(model_config.drag_optim_config.optimize_iter, int)
+        assert isinstance(model_config.drag_optim_config.optimize_lr, float)
+        assert len(model_config.drag_optim_config.record_feature_block_indexes) > 0
+        assert len(trajectories) == len(masks)
+        if len(trajectories) == 0:
+            return latents
+
+        print(f"{trajectories = }")
+        print(f"{len(masks) = }")
+
+        original_latents = latents.clone().detach()
+
+        original_denoised_pred = self.generate_features(
+            latents=latents,
+            conditional_dict=conditional_dict,
+            timestep=timestep,
+            kv_cache=kv_cache,
+            crossattn_cache=crossattn_cache,
+            current_start=current_start,
+            # below are for drag optimization
+            model_config=model_config,
+        )
+        original_denoised_pred, _ = original_denoised_pred
+
+        # Check if optimization is enabled
+        if model_config.drag_optim_config.optimize_iter <= 0:
+            return latents
+        # Cast latents to configured dtype for optimization
+        latent_original_dtype = latents.dtype
+        latents = latents.to(dtype=torch.bfloat16)
+        timestep_original_dtype = timestep.dtype
+        timestep = timestep.to(dtype=latents.dtype)
+
+        for param in self.generator.parameters():
+            generator_original_dtype = param.dtype
+            break
+        self.generator = self.generator.to(dtype=latents.dtype)
+        # self.generator.train(True)
+        for param in self.generator.parameters():
+            param.requires_grad = False
+
+        split_trajectories_list = split_trajectories_segments(
+            trajectories=trajectories,
+            translation_step=model_config.drag_optim_config.translation_step,
+            rotation_step=model_config.drag_optim_config.rotation_step,
+        )
+        # split_trajectories_list: list[segment_index] -> list[trajectory_index] -> trajectory dict
+
+        def _select_variant(feat):
+            if isinstance(feat, dict):
+                keys = list(feat.keys())
+                if not keys:
+                    raise ValueError("Empty feature dict provided.")
+                non_orig = [k for k in keys if str(k) != "original"]
+                key = random.choice(keys)
+                print(f"Selected feature variant {key = } from {keys = }")
+                # key = "original" if "original" in feat else keys[0]
+                return feat[key]
+            return feat
+
+        def get_previous_last(
+            prev: dict[int, dict[str, torch.Tensor]],
+            # dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536]
+        ) -> dict[str, torch.Tensor]:
+            out = transpose_dict_2d(prev)
+            # out: dict[key] -> dict[block_index -> Tensor] [1, 3, 30, 52, 1536]
+            out = {k: self.stack_features(v) for k, v in out.items()}
+            # out: dict[key] -> Tensor [1536 * L, 3, 30, 52]
+            out = {k: v[:, -1, ...].detach() for k, v in out.items()}
+            # out: dict[key] -> Tensor [1536 * L, 30, 52]
+            out = {
+                k: resize_tensor(
+                    v.detach(),
+                    scale_factor=model_config.drag_optim_config.feature_scaling_factor,
+                    mode="bilinear",
+                ).detach()
+                for k, v in out.items()
+            }
+            # out: dict[key] -> Tensor [1536 * L, 30 * scaling, 52 * scaling]
+            return out
+
+        previous_record_features: dict[str, torch.Tensor] = get_previous_last(
+            noisy_previous_record_feature
+        )
+        # previous_record_features: dict[key] -> Tensor [1536 * L, 30 * scaling, 52 * scaling]
+
+        movable_mask_torch = torch.tensor(movable_mask, device=latents.device).float()
+        downsample_movable_mask = resize_tensor(
+            movable_mask_torch.detach(),
+            size=tuple(original_denoised_pred.shape[-2:]),
+            mode="nearest",
+        ).detach()
+        # print(f"{downsample_movable_mask.shape = }") # [60, 104]
+
+        with torch.enable_grad():
+
+            latents.requires_grad_(True)
+            optimizer = torch.optim.AdamW(
+                [latents],
+                lr=model_config.drag_optim_config.optimize_lr,
+            )
+
+            for split_traj_idx, split_trajectories in enumerate(split_trajectories_list):
+                # split_trajectories: N x trajectory, list[trajectory_index] -> trajectory dict
+                (
+                    warped_masks,
+                    down_warp_masks,
+                    gaussian_heatmaps_per_traj,
+                    combined_downsample_movable_mask,
+                    warped_previous_record_features,
+                    combined_gaussian_heatmaps,
+                ) = CausalInferencePipeline.precompute_conditions(
+                    trajectories=split_trajectories,
+                    masks=masks,
+                    dtype=latents.dtype,
+                    device=latents.device,
+                    model_config=model_config,
+                    downsample_movable_mask=downsample_movable_mask,
+                    previous_record_features=previous_record_features,
+                    latent_spatial_size=tuple(latents.shape[-2:]),
+                )
+                # -------------------------
+                # Optimization iterations (reuse precomputed items)
+                # -------------------------
+                for optimize_iter_idx in range(model_config.drag_optim_config.optimize_iter):
+                    print(f"{optimize_iter_idx = }")
+                    print(f"{latents.mean((0, 2, 3, 4)) = }")
+                    print(f"{latents.std((0, 2, 3, 4)) = }")
+
+                    denoised_pred, record_features = self.generate_features(
+                        latents=latents,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=kv_cache,
+                        crossattn_cache=crossattn_cache,
+                        current_start=current_start,
+                        model_config=model_config,
+                    )
+                    # denoised_pred.shape [1, 3, 16, 60, 104]
+                    # record_features: dict[key] -> Tensor [1536 * L, 3, 30, 52]
+
+                    print(f"selecting features for optimization")
+                    record_features_selected = _select_variant(
+                        record_features,
+                    )
+                    # print(f"{record_features_selected.shape = }")
+                    # record_features_selected: Tensor [1536 * L, 3, 30, 52]
+                    del record_features
+
+                    if optimize_target_latent_index >= 0:
+                        record_features_selected = record_features_selected[
+                            :,
+                            optimize_target_latent_index : optimize_target_latent_index + 1,
+                        ]
+
+                    record_features_selected = resize_tensor(
+                        record_features_selected,
+                        scale_factor=model_config.drag_optim_config.feature_scaling_factor,
+                        mode="bilinear",
+                    )
+                    # print(
+                    #     f"{record_features_selected.shape = }"
+                    # )  # [1536 * L, 3, 30 * scaling, 52 * scaling]
+
+                    print(f"selecting warped previous features")
+                    warped_previous_feature_selected = _select_variant(
+                        warped_previous_record_features,
+                    )  # list[traj_index] -> list[frame_index] -> Tensor [1536 * L, 30 * scaling, 52 * scaling]
+
+                    loss = 0
+                    loss_cnt = 0
+
+                    # Iterate over each trajectory point
+                    for trajectory_index, trajectory in enumerate(split_trajectories):
+                        assert record_features_selected.shape[1] == len(trajectory["deltas"])
+
+                        for frame_index in range(len(trajectory["deltas"])):
+                            warped_attention_values = warped_previous_feature_selected[
+                                trajectory_index
+                            ][frame_index]
+
+                            downsample_warped_mask = down_warp_masks[trajectory_index][frame_index]
+
+                            pixel_wise_loss = F.mse_loss(
+                                warped_attention_values * downsample_warped_mask,
+                                record_features_selected[:, frame_index] * downsample_warped_mask,
+                                reduction="none",
+                            ).mean(dim=0)
+                            # print(f"{pixel_wise_loss.shape = }")  # [60, 104]
+
+                            # Add weighted loss
+                            loss = loss + (downsample_warped_mask * pixel_wise_loss).sum()
+                            loss_cnt += downsample_warped_mask.sum()
+
+                    print(f"{loss = }  /  {loss_cnt = }")
+                    loss = loss / max(1e-8, loss_cnt)
+                    print(f"{loss = }")
+
+                    unchanged_mask = 1.0 - combined_downsample_movable_mask
+                    unchanged_loss = F.mse_loss(
+                        denoised_pred * unchanged_mask.detach(),
+                        original_denoised_pred.detach() * unchanged_mask.detach(),
+                    )
+                    print(f"{unchanged_loss = }")
+                    loss = loss + unchanged_loss * 1.0
+
+                    # Update latents
+                    self.generator.zero_grad()
+                    optimizer.zero_grad()
+                    if loss_cnt > 0:
+                        loss.backward()
+                        assert (
+                            combined_gaussian_heatmaps.shape[0] == 1
+                            or combined_gaussian_heatmaps.shape[0] == latents.shape[-4]
+                        )
+                        assert combined_gaussian_heatmaps.shape[-2:] == latents.shape[-2:]
+                        latents.grad.mul_(combined_gaussian_heatmaps[:, None, :, :])
+                        # Clip gradients
+                        clip_grad_norm_(
+                            [latents],
+                            max_norm=1.0,
+                            norm_type=2,
+                        )
+                        optimizer.step()
+                    if model_config.drag_optim_config.normalize_latent_after_drag_optimize:
+                        print(f"Normalizing latents after optimize iteration")
+                        latents = (
+                            normalize_tensor_to_match_tensor(
+                                latents.detach().clone(),
+                                dim=(0, 3, 4),
+                                reference_tensor=original_latents.to(dtype=latents.dtype),
+                            )
+                            .detach()
+                            .clone()
+                        )
+                        # latents = latents.clamp(
+                        #     min=latents_min,
+                        #     max=latents_max,
+                        # ).detach().clone()
+                        latents.requires_grad_(True)
+                        optimizer = torch.optim.AdamW(
+                            [latents],
+                            lr=model_config.drag_optim_config.optimize_lr,
+                        )
+                    # Clean up to save memory
+                    gc.collect()
+                    torch.cuda.empty_cache()
+
+        latents = latents.detach().requires_grad_(False)
+
+        if model_config.drag_optim_config.normalize_latent_after_post_merge:
+            latents = (
+                normalize_tensor_to_match_tensor(
+                    latents,
+                    dim=None,
+                    reference_tensor=original_latents,
+                )
+                .detach()
+                .clone()
+            )
+
+        # Convert back to original dtype
+        self.generator = self.generator.to(dtype=generator_original_dtype)
+        self.generator.train(False)
+        latents = latents.to(dtype=latent_original_dtype)
+        timestep = timestep.to(dtype=timestep_original_dtype)
+
+        # Detach latents and remove gradient
+        latents = latents.detach().requires_grad_(False)
+        return latents
+
+    def inference(
+        self,
+        noise: torch.Tensor,
+        text_prompts: List[str],
+        initial_latent: Optional[torch.Tensor] = None,
+        return_latents: bool = False,
+        profile: bool = False,
+        low_memory: bool = False,
+        do_not_decode_video: bool = False,
+        do_not_recompute_initial_latents: bool = False,
+        # below are for drag optimization
+        model_config: DictConfig = None,
+        previous_record_feature_list: dict[int, dict[int, dict[str, torch.Tensor]]] = None,
+        # dict[denoising_step] -> dict[block_index] -> dict[key -> Tensor] [1, 3, 30, 52, 1536]
+        is_drag_optimize: bool = False,
+        block_trajectories: list[list[dict[str, bool | list[torch.Tensor]]]] = [],
+        masks: list[np.ndarray] = [],
+        movable_mask: np.ndarray = None,
+        drag_optimize_target_latent_index: int = -1,
+    ) -> torch.Tensor:
+        """
+        Perform inference on the given noise and text prompts.
+        Inputs:
+            noise (torch.Tensor): The input noise tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+            text_prompts (List[str]): The list of text prompts.
+            initial_latent (torch.Tensor): The initial latent tensor of shape
+                (batch_size, num_input_frames, num_channels, height, width).
+                If num_input_frames is 1, perform image to video.
+                If num_input_frames is greater than 1, perform video extension.
+            return_latents (bool): Whether to return the latents.
+            :param block_trajectories:
+                block_num x N x trajectory,
+                trajectory has keys 'is_rotation' 'deltas' 'start_point'
+                if is_rotation: trajectory also has 'rotation_center'
+            :param masks:
+                N x mask
+        Outputs:
+            video (torch.Tensor): The generated video tensor of shape
+                (batch_size, num_output_frames, num_channels, height, width).
+                It is normalized to be in the range [0, 1].
+        """
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (
+            self.independent_first_frame and initial_latent is not None
+        ):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        conditional_dict = self.text_encoder(text_prompts=text_prompts)
+
+        if low_memory:
+            gpu_memory_preservation = get_cuda_free_memory_gb(gpu) + 5
+            move_model_to_device_with_memory_preservation(
+                self.text_encoder,
+                target_device=gpu,
+                preserved_memory_gb=gpu_memory_preservation,
+            )
+
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype,
+        )
+
+        # Set up profiling if requested
+        if profile:
+            init_start = torch.cuda.Event(enable_timing=True)
+            init_end = torch.cuda.Event(enable_timing=True)
+            diffusion_start = torch.cuda.Event(enable_timing=True)
+            diffusion_end = torch.cuda.Event(enable_timing=True)
+            vae_start = torch.cuda.Event(enable_timing=True)
+            vae_end = torch.cuda.Event(enable_timing=True)
+            block_times = []
+            block_start = torch.cuda.Event(enable_timing=True)
+            block_end = torch.cuda.Event(enable_timing=True)
+            init_start.record()
+
+        # Step 1: Initialize KV cache to all zeros
+        if self.kv_cache1 is None:
+            self._initialize_kv_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device)
+            self._initialize_crossattn_cache(
+                batch_size=batch_size, dtype=noise.dtype, device=noise.device
+            )
+        else:
+            if do_not_recompute_initial_latents:
+                pass
+            else:
+                print(f"Resetting caches")
+                self._reset_crossattn_cache()
+                self._reset_kv_cache()
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            if self.independent_first_frame:
+                # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+                assert (num_input_frames - 1) % self.num_frame_per_block == 0
+                num_input_blocks = (num_input_frames - 1) // self.num_frame_per_block
+                output[:, :1] = initial_latent[:, :1]
+                if do_not_recompute_initial_latents:
+                    pass
+                else:
+                    print(f"Recompute KV cache based on Initial Latents")
+                    self.generator(
+                        noisy_image_or_video=initial_latent[:, :1],
+                        conditional_dict=conditional_dict,
+                        timestep=timestep * 0,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length,
+                    )
+                current_start_frame += 1
+            else:
+                # Assume num_input_frames is self.num_frame_per_block * num_input_blocks
+                assert num_input_frames % self.num_frame_per_block == 0
+                num_input_blocks = num_input_frames // self.num_frame_per_block
+
+            for _ in range(num_input_blocks):
+                current_ref_latents = initial_latent[
+                    :,
+                    current_start_frame : current_start_frame + self.num_frame_per_block,
+                ]
+                output[
+                    :,
+                    current_start_frame : current_start_frame + self.num_frame_per_block,
+                ] = current_ref_latents
+                if do_not_recompute_initial_latents:
+                    pass
+                else:
+                    print(f"Recompute KV cache based on Initial Latents")
+                    self.generator(
+                        noisy_image_or_video=current_ref_latents,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep * 0,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length,
+                    )
+                current_start_frame += self.num_frame_per_block
+
+        if profile:
+            init_end.record()
+            torch.cuda.synchronize()
+            diffusion_start.record()
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        for current_chunk_index, current_num_frames in enumerate(
+            tqdm(all_num_frames),
+            start=num_input_blocks if initial_latent is not None else 0,
+        ):
+            print(f"\n{current_chunk_index = } ; {current_start_frame = }")
+            if profile:
+                block_start.record()
+
+            noisy_input = noise[
+                :,
+                current_start_frame
+                - num_input_frames : current_start_frame
+                + current_num_frames
+                - num_input_frames,
+            ]
+
+            if model_config is not None and OmegaConf.select(
+                model_config, "drag_optim_config.record_feature_block_indexes"
+            ):
+                record_attention_values_list = {}
+            # Step 3.1: Spatial denoising loop
+            for time_step_index, current_timestep in enumerate(self.denoising_step_list):
+                print(f"{time_step_index = } ; {current_timestep = }")
+                # set current timestep
+                timestep = (
+                    torch.ones(
+                        [batch_size, current_num_frames],
+                        device=noise.device,
+                        dtype=torch.int64,
+                    )
+                    * current_timestep
+                )
+
+                if (
+                    is_drag_optimize
+                    and time_step_index
+                    in model_config.drag_optim_config.optimize_denoising_steps_indexes
+                ):
+
+                    noisy_input = self.optimize_latent(
+                        latents=noisy_input,
+                        conditional_dict=conditional_dict,
+                        timestep=timestep,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame * self.frame_seq_length,
+                        # below are for drag optimization
+                        trajectories=block_trajectories[current_chunk_index - num_input_blocks],
+                        masks=masks,
+                        movable_mask=movable_mask,
+                        clean_previous_record_feature=previous_record_feature_list[-1],
+                        noisy_previous_record_feature=previous_record_feature_list[time_step_index],
+                        model_config=model_config,
+                        optimize_target_latent_index=drag_optimize_target_latent_index,
+                    )
+                    print(f"{noisy_input.mean() = }")
+                    print(f"{noisy_input.std() = }")
+
+                _, denoised_pred = self.generator(
+                    noisy_image_or_video=noisy_input,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                    model_config=model_config,
+                )
+                if model_config is not None and OmegaConf.select(
+                    model_config,
+                    "drag_optim_config.record_feature_block_indexes",
+                ):
+                    denoised_pred, record_features = denoised_pred
+                    if (
+                        time_step_index
+                        in model_config.drag_optim_config.optimize_denoising_steps_indexes
+                    ):
+                        record_attention_values_list[time_step_index] = record_features
+
+                if (
+                    model_config is not None
+                    and OmegaConf.select(
+                        model_config,
+                        "drag_optim_config.dynamic_chunk_normalization_block_number",
+                        default=0,
+                    )
+                    > 0
+                ):
+                    num_norm_blocks = (
+                        model_config.drag_optim_config.dynamic_chunk_normalization_block_number
+                    )
+                    # Exclude the first chunk (independent first frame latent) by starting no earlier than num_frame_per_block
+                    dynamic_normalize_start_frame_index = max(
+                        self.num_frame_per_block,
+                        (current_chunk_index - num_norm_blocks) * self.num_frame_per_block,
+                    )
+                    # print(f"{dynamic_normalize_start_frame_index = }")
+                    if dynamic_normalize_start_frame_index < current_start_frame:
+                        reference_tensor = torch.cat(
+                            [
+                                output[
+                                    :,
+                                    dynamic_normalize_start_frame_index:current_start_frame,
+                                ],
+                                denoised_pred,
+                            ],
+                            dim=1,
+                        )
+                        denoised_pred = normalize_tensor_to_match_tensor(
+                            denoised_pred,
+                            dim=None,
+                            reference_tensor=reference_tensor,
+                        )
+                        # print(f"{denoised_pred.mean() = }")
+                        # print(f"{denoised_pred.std() = }")
+
+                if time_step_index < len(self.denoising_step_list) - 1:
+                    next_timestep = self.denoising_step_list[time_step_index + 1]
+                    noisy_input = self.scheduler.add_noise(
+                        denoised_pred.flatten(0, 1),
+                        torch.randn_like(denoised_pred.flatten(0, 1)),
+                        next_timestep
+                        * torch.ones(
+                            [batch_size * current_num_frames],
+                            device=noise.device,
+                            dtype=torch.long,
+                        ),
+                    ).unflatten(0, denoised_pred.shape[:2])
+
+            # Step 3.2: record the model's output
+            output[
+                :,
+                current_start_frame : current_start_frame + current_num_frames,
+            ] = denoised_pred
+
+            # Step 3.3: rerun with timestep zero to update KV cache using clean context
+            context_timestep = torch.ones_like(timestep) * self.args.context_noise
+            _, denoised_pred = self.generator(
+                noisy_image_or_video=denoised_pred,
+                conditional_dict=conditional_dict,
+                timestep=context_timestep,
+                kv_cache=self.kv_cache1,
+                crossattn_cache=self.crossattn_cache,
+                current_start=current_start_frame * self.frame_seq_length,
+                model_config=model_config,
+            )
+            if model_config is not None and OmegaConf.select(
+                model_config, "drag_optim_config.record_feature_block_indexes"
+            ):
+                denoised_pred, record_features = denoised_pred
+                record_attention_values_list[-1] = record_features
+
+            if profile:
+                block_end.record()
+                torch.cuda.synchronize()
+                block_time = block_start.elapsed_time(block_end)
+                block_times.append(block_time)
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        if profile:
+            # End diffusion timing and synchronize CUDA
+            diffusion_end.record()
+            torch.cuda.synchronize()
+            diffusion_time = diffusion_start.elapsed_time(diffusion_end)
+            init_time = init_start.elapsed_time(init_end)
+            vae_start.record()
+
+        # Step 4: Decode the output
+        if not do_not_decode_video:
+            start_decode_time = time.time()
+            video = self.vae.decode_to_pixel(output, use_cache=False)
+            video = (video * 0.5 + 0.5).clamp(0, 1)
+            print(
+                f"{self.__class__.__name__}.inference() VAE decode time: {time.time() - start_decode_time:.2f} seconds"
+            )
+
+        if profile:
+            # End VAE timing and synchronize CUDA
+            vae_end.record()
+            torch.cuda.synchronize()
+            vae_time = vae_start.elapsed_time(vae_end)
+            total_time = init_time + diffusion_time + vae_time
+
+            print("Profiling results:")
+            print(
+                f"  - Initialization/caching time: {init_time:.2f} ms ({100 * init_time / total_time:.2f}%)"
+            )
+            print(
+                f"  - Diffusion generation time: {diffusion_time:.2f} ms ({100 * diffusion_time / total_time:.2f}%)"
+            )
+            for i, block_time in enumerate(block_times):
+                print(
+                    f"    - Block {i} generation time: {block_time:.2f} ms ({100 * block_time / diffusion_time:.2f}% of diffusion)"
+                )
+            print(f"  - VAE decoding time: {vae_time:.2f} ms ({100 * vae_time / total_time:.2f}%)")
+            print(f"  - Total time: {total_time:.2f} ms")
+
+        return_values = []
+        if not do_not_decode_video:
+            return_values.append(video)
+        if return_latents:
+            return_values.append(output)
+        if model_config is not None and OmegaConf.select(
+            model_config, "drag_optim_config.record_feature_block_indexes"
+        ):
+            return_values.append(record_attention_values_list)
+
+        if len(return_values) == 0:
+            return
+        elif len(return_values) == 1:
+            return return_values[0]
+        else:
+            return tuple(return_values)
+
+    def _initialize_kv_cache(
+        self,
+        batch_size,
+        dtype,
+        device,
+    ):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        print(
+            f"""
+{type(self).__name__}._initialize_kv_cache
+    {batch_size = }
+    {dtype = }
+    {device = }
+"""
+        )
+        kv_cache1 = []
+        if self.local_attn_size != -1:
+            print(f"use {self.local_attn_size = }")
+            # Use the local attention size to compute the KV cache size
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            # Use the default KV cache size
+            kv_cache_size = 32760
+        print(f"{kv_cache_size = }")
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.append(
+                {
+                    "k": torch.zeros(
+                        [batch_size, kv_cache_size, 12, 128],
+                        dtype=dtype,
+                        device=device,
+                    ),
+                    "v": torch.zeros(
+                        [batch_size, kv_cache_size, 12, 128],
+                        dtype=dtype,
+                        device=device,
+                    ),
+                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                }
+            )
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+
+    def _initialize_crossattn_cache(
+        self,
+        batch_size,
+        dtype,
+        device,
+    ):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        print(
+            f"""
+{type(self).__name__}._initialize_crossattn_cache
+    {batch_size = }
+    {dtype = }
+    {device = }
+"""
+        )
+        crossattn_cache = []
+
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append(
+                {
+                    "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                    "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                    "is_init": False,
+                }
+            )
+        self.crossattn_cache = crossattn_cache
+
+    def _reset_crossattn_cache(self):
+        # reset cross attn cache
+        print(f"{type(self).__name__}._reset_crossattn_cache")
+        for block_index in range(self.num_transformer_blocks):
+            self.crossattn_cache[block_index]["is_init"] = False
+
+    def _reset_kv_cache(self):
+        # reset kv cache
+        print(f"{type(self).__name__}._reset_kv_cache")
+        for block_index in range(len(self.kv_cache1)):
+            self.kv_cache1[block_index]["global_end_index"] = torch.tensor(
+                [0],
+                dtype=torch.long,
+                device=self.kv_cache1[block_index]["global_end_index"].device,
+            )
+            self.kv_cache1[block_index]["local_end_index"] = torch.tensor(
+                [0],
+                dtype=torch.long,
+                device=self.kv_cache1[block_index]["local_end_index"].device,
+            )
+
+    def is_kv_cache_initialized(self):
+        return hasattr(self, "kv_cache1") and self.kv_cache1 is not None
+
+    def is_crossattn_cache_initialized(self):
+        return hasattr(self, "crossattn_cache") and self.crossattn_cache is not None

pipeline/self_forcing_training.py

@@ -0,0 +1,351 @@
+from utils.wan_wrapper import WanDiffusionWrapper
+from utils.scheduler import SchedulerInterface
+from typing import List, Optional
+import torch
+import torch.distributed as dist
+
+
+class SelfForcingTrainingPipeline:
+    def __init__(
+        self,
+        denoising_step_list: List[int],
+        scheduler: SchedulerInterface,
+        generator: WanDiffusionWrapper,
+        num_frame_per_block=3,
+        independent_first_frame: bool = False,
+        same_step_across_blocks: bool = False,
+        last_step_only: bool = False,
+        num_max_frames: int = 21,
+        context_noise: int = 0,
+        **kwargs
+    ):
+        super().__init__()
+        self.scheduler = scheduler
+        self.generator = generator
+        self.denoising_step_list = denoising_step_list
+        if self.denoising_step_list[-1] == 0:
+            self.denoising_step_list = self.denoising_step_list[
+                :-1
+            ]  # remove the zero timestep for inference
+
+        # Wan specific hyperparameters
+        self.num_transformer_blocks = 30
+        self.frame_seq_length = 1560
+        self.num_frame_per_block = num_frame_per_block
+        self.context_noise = context_noise
+        self.i2v = False
+
+        self.kv_cache1 = None
+        self.kv_cache2 = None
+        self.independent_first_frame = independent_first_frame
+        self.same_step_across_blocks = same_step_across_blocks
+        self.last_step_only = last_step_only
+        self.kv_cache_size = num_max_frames * self.frame_seq_length
+
+    def generate_and_sync_list(
+        self,
+        num_blocks,
+        num_denoising_steps,
+        device,
+    ):
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        if rank == 0:
+            # Generate random indices
+            indices = torch.randint(
+                low=0,
+                high=num_denoising_steps,
+                size=(num_blocks,),
+                device=device,
+            )
+            if self.last_step_only:
+                indices = torch.ones_like(indices) * (num_denoising_steps - 1)
+        else:
+            indices = torch.empty(num_blocks, dtype=torch.long, device=device)
+
+        dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
+        return indices.tolist()
+
+    def inference_with_trajectory(
+        self,
+        noise: torch.Tensor,
+        initial_latent: Optional[torch.Tensor] = None,
+        return_sim_step: bool = False,
+        **conditional_dict
+    ) -> torch.Tensor:
+        batch_size, num_frames, num_channels, height, width = noise.shape
+        if not self.independent_first_frame or (
+            self.independent_first_frame and initial_latent is not None
+        ):
+            # If the first frame is independent and the first frame is provided, then the number of frames in the
+            # noise should still be a multiple of num_frame_per_block
+            assert num_frames % self.num_frame_per_block == 0
+            num_blocks = num_frames // self.num_frame_per_block
+        else:
+            # Using a [1, 4, 4, 4, 4, 4, ...] model to generate a video without image conditioning
+            assert (num_frames - 1) % self.num_frame_per_block == 0
+            num_blocks = (num_frames - 1) // self.num_frame_per_block
+        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
+        num_output_frames = num_frames + num_input_frames  # add the initial latent frames
+        output = torch.zeros(
+            [batch_size, num_output_frames, num_channels, height, width],
+            device=noise.device,
+            dtype=noise.dtype,
+        )
+
+        # Step 1: Initialize KV cache to all zeros
+        self._initialize_kv_cache(batch_size=batch_size, dtype=noise.dtype, device=noise.device)
+        self._initialize_crossattn_cache(
+            batch_size=batch_size, dtype=noise.dtype, device=noise.device
+        )
+        # if self.kv_cache1 is None:
+        #     self._initialize_kv_cache(
+        #         batch_size=batch_size,
+        #         dtype=noise.dtype,
+        #         device=noise.device,
+        #     )
+        #     self._initialize_crossattn_cache(
+        #         batch_size=batch_size,
+        #         dtype=noise.dtype,
+        #         device=noise.device
+        #     )
+        # else:
+        #     # reset cross attn cache
+        #     for block_index in range(self.num_transformer_blocks):
+        #         self.crossattn_cache[block_index]["is_init"] = False
+        #     # reset kv cache
+        #     for block_index in range(len(self.kv_cache1)):
+        #         self.kv_cache1[block_index]["global_end_index"] = torch.tensor(
+        #             [0], dtype=torch.long, device=noise.device)
+        #         self.kv_cache1[block_index]["local_end_index"] = torch.tensor(
+        #             [0], dtype=torch.long, device=noise.device)
+
+        # Step 2: Cache context feature
+        current_start_frame = 0
+        if initial_latent is not None:
+            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
+            # Assume num_input_frames is 1 + self.num_frame_per_block * num_input_blocks
+            output[:, :1] = initial_latent
+            with torch.no_grad():
+                self.generator(
+                    noisy_image_or_video=initial_latent,
+                    conditional_dict=conditional_dict,
+                    timestep=timestep * 0,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                )
+            current_start_frame += 1
+
+        # Step 3: Temporal denoising loop
+        all_num_frames = [self.num_frame_per_block] * num_blocks
+        if self.independent_first_frame and initial_latent is None:
+            all_num_frames = [1] + all_num_frames
+        num_denoising_steps = len(self.denoising_step_list)
+        exit_flags = self.generate_and_sync_list(
+            len(all_num_frames), num_denoising_steps, device=noise.device
+        )
+        start_gradient_frame_index = num_output_frames - 21
+
+        # for block_index in range(num_blocks):
+        for block_index, current_num_frames in enumerate(all_num_frames):
+            noisy_input = noise[
+                :,
+                current_start_frame
+                - num_input_frames : current_start_frame
+                + current_num_frames
+                - num_input_frames,
+            ]
+
+            # Step 3.1: Spatial denoising loop
+            for index, current_timestep in enumerate(self.denoising_step_list):
+                if self.same_step_across_blocks:
+                    exit_flag = index == exit_flags[0]
+                else:
+                    exit_flag = (
+                        index == exit_flags[block_index]
+                    )  # Only backprop at the randomly selected timestep (consistent across all ranks)
+                timestep = (
+                    torch.ones(
+                        [batch_size, current_num_frames],
+                        device=noise.device,
+                        dtype=torch.int64,
+                    )
+                    * current_timestep
+                )
+
+                if not exit_flag:
+                    with torch.no_grad():
+                        _, denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=current_start_frame * self.frame_seq_length,
+                        )
+                        next_timestep = self.denoising_step_list[index + 1]
+                        noisy_input = self.scheduler.add_noise(
+                            denoised_pred.flatten(0, 1),
+                            torch.randn_like(denoised_pred.flatten(0, 1)),
+                            next_timestep
+                            * torch.ones(
+                                [batch_size * current_num_frames],
+                                device=noise.device,
+                                dtype=torch.long,
+                            ),
+                        ).unflatten(0, denoised_pred.shape[:2])
+                else:
+                    # for getting real output
+                    # with torch.set_grad_enabled(current_start_frame >= start_gradient_frame_index):
+                    if current_start_frame < start_gradient_frame_index:
+                        with torch.no_grad():
+                            _, denoised_pred = self.generator(
+                                noisy_image_or_video=noisy_input,
+                                conditional_dict=conditional_dict,
+                                timestep=timestep,
+                                kv_cache=self.kv_cache1,
+                                crossattn_cache=self.crossattn_cache,
+                                current_start=current_start_frame * self.frame_seq_length,
+                            )
+                    else:
+                        _, denoised_pred = self.generator(
+                            noisy_image_or_video=noisy_input,
+                            conditional_dict=conditional_dict,
+                            timestep=timestep,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=current_start_frame * self.frame_seq_length,
+                        )
+                    break
+
+            # Step 3.2: record the model's output
+            output[
+                :,
+                current_start_frame : current_start_frame + current_num_frames,
+            ] = denoised_pred
+
+            # Step 3.3: rerun with timestep zero to update the cache
+            context_timestep = torch.ones_like(timestep) * self.context_noise
+            # add context noise
+            denoised_pred = self.scheduler.add_noise(
+                denoised_pred.flatten(0, 1),
+                torch.randn_like(denoised_pred.flatten(0, 1)),
+                context_timestep
+                * torch.ones(
+                    [batch_size * current_num_frames],
+                    device=noise.device,
+                    dtype=torch.long,
+                ),
+            ).unflatten(0, denoised_pred.shape[:2])
+            with torch.no_grad():
+                self.generator(
+                    noisy_image_or_video=denoised_pred,
+                    conditional_dict=conditional_dict,
+                    timestep=context_timestep,
+                    kv_cache=self.kv_cache1,
+                    crossattn_cache=self.crossattn_cache,
+                    current_start=current_start_frame * self.frame_seq_length,
+                )
+
+            # Step 3.4: update the start and end frame indices
+            current_start_frame += current_num_frames
+
+        # Step 3.5: Return the denoised timestep
+        if not self.same_step_across_blocks:
+            denoised_timestep_from, denoised_timestep_to = None, None
+        elif exit_flags[0] == len(self.denoising_step_list) - 1:
+            denoised_timestep_to = 0
+            denoised_timestep_from = (
+                1000
+                - torch.argmin(
+                    (
+                        self.scheduler.timesteps.cuda()
+                        - self.denoising_step_list[exit_flags[0]].cuda()
+                    ).abs(),
+                    dim=0,
+                ).item()
+            )
+        else:
+            denoised_timestep_to = (
+                1000
+                - torch.argmin(
+                    (
+                        self.scheduler.timesteps.cuda()
+                        - self.denoising_step_list[exit_flags[0] + 1].cuda()
+                    ).abs(),
+                    dim=0,
+                ).item()
+            )
+            denoised_timestep_from = (
+                1000
+                - torch.argmin(
+                    (
+                        self.scheduler.timesteps.cuda()
+                        - self.denoising_step_list[exit_flags[0]].cuda()
+                    ).abs(),
+                    dim=0,
+                ).item()
+            )
+
+        if return_sim_step:
+            return (
+                output,
+                denoised_timestep_from,
+                denoised_timestep_to,
+                exit_flags[0] + 1,
+            )
+
+        return output, denoised_timestep_from, denoised_timestep_to
+
+    def _initialize_kv_cache(
+        self,
+        batch_size,
+        dtype,
+        device,
+    ):
+        """
+        Initialize a Per-GPU KV cache for the Wan model.
+        """
+        kv_cache1 = []
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.append(
+                {
+                    "k": torch.zeros(
+                        [batch_size, self.kv_cache_size, 12, 128],
+                        dtype=dtype,
+                        device=device,
+                    ),
+                    "v": torch.zeros(
+                        [batch_size, self.kv_cache_size, 12, 128],
+                        dtype=dtype,
+                        device=device,
+                    ),
+                    "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                    "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+                }
+            )
+
+        self.kv_cache1 = kv_cache1  # always store the clean cache
+
+    def _initialize_crossattn_cache(
+        self,
+        batch_size,
+        dtype,
+        device,
+    ):
+        """
+        Initialize a Per-GPU cross-attention cache for the Wan model.
+        """
+        crossattn_cache = []
+
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append(
+                {
+                    "k": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                    "v": torch.zeros([batch_size, 512, 12, 128], dtype=dtype, device=device),
+                    "is_init": False,
+                }
+            )
+        self.crossattn_cache = crossattn_cache

prompts/vbench/all_dimension.txt

@@ -0,0 +1,946 @@
+In a still frame, a stop sign
+a toilet, frozen in time
+a laptop, frozen in time
+A tranquil tableau of alley
+A tranquil tableau of bar
+A tranquil tableau of barn
+A tranquil tableau of bathroom
+A tranquil tableau of bedroom
+A tranquil tableau of cliff
+In a still frame, courtyard
+In a still frame, gas station
+A tranquil tableau of house
+indoor gymnasium, frozen in time
+A tranquil tableau of indoor library
+A tranquil tableau of kitchen
+A tranquil tableau of palace
+In a still frame, parking lot
+In a still frame, phone booth
+A tranquil tableau of restaurant
+A tranquil tableau of tower
+A tranquil tableau of a bowl
+A tranquil tableau of an apple
+A tranquil tableau of a bench
+A tranquil tableau of a bed
+A tranquil tableau of a chair
+A tranquil tableau of a cup
+A tranquil tableau of a dining table
+In a still frame, a pear
+A tranquil tableau of a bunch of grapes
+A tranquil tableau of a bowl on the kitchen counter
+A tranquil tableau of a beautiful, handcrafted ceramic bowl
+A tranquil tableau of an antique bowl
+A tranquil tableau of an exquisite mahogany dining table
+A tranquil tableau of a wooden bench in the park
+A tranquil tableau of a beautiful wrought-iron bench surrounded by blooming flowers
+In a still frame, a park bench with a view of the lake
+A tranquil tableau of a vintage rocking chair was placed on the porch
+A tranquil tableau of the jail cell was small and dimly lit, with cold, steel bars
+A tranquil tableau of the phone booth was tucked away in a quiet alley
+a dilapidated phone booth stood as a relic of a bygone era on the sidewalk, frozen in time
+A tranquil tableau of the old red barn stood weathered and iconic against the backdrop of the countryside
+A tranquil tableau of a picturesque barn was painted a warm shade of red and nestled in a picturesque meadow
+In a still frame, within the desolate desert, an oasis unfolded, characterized by the stoic presence of palm trees and a motionless, glassy pool of water
+In a still frame, the Parthenon's majestic Doric columns stand in serene solitude atop the Acropolis, framed by the tranquil Athenian landscape
+In a still frame, the Temple of Hephaestus, with its timeless Doric grace, stands stoically against the backdrop of a quiet Athens
+In a still frame, the ornate Victorian streetlamp stands solemnly, adorned with intricate ironwork and stained glass panels
+A tranquil tableau of the Stonehenge presented itself as an enigmatic puzzle, each colossal stone meticulously placed against the backdrop of tranquility
+In a still frame, in the vast desert, an oasis nestled among dunes, featuring tall palm trees and an air of serenity
+static view on a desert scene with an oasis, palm trees, and a clear, calm pool of water
+A tranquil tableau of an ornate Victorian streetlamp standing on a cobblestone street corner, illuminating the empty night
+A tranquil tableau of a tranquil lakeside cabin nestled among tall pines, its reflection mirrored perfectly in the calm water
+In a still frame, a vintage gas lantern, adorned with intricate details, gracing a historic cobblestone square
+In a still frame, a tranquil Japanese tea ceremony room, with tatami mats, a delicate tea set, and a bonsai tree in the corner
+A tranquil tableau of the Parthenon stands resolute in its classical elegance, a timeless symbol of Athens' cultural legacy
+A tranquil tableau of in the heart of Plaka, the neoclassical architecture of the old city harmonizes with the ancient ruins
+A tranquil tableau of in the desolate beauty of the American Southwest, Chaco Canyon's ancient ruins whispered tales of an enigmatic civilization that once thrived amidst the arid landscapes
+A tranquil tableau of at the edge of the Arabian Desert, the ancient city of Petra beckoned with its enigmatic rock-carved façades
+In a still frame, amidst the cobblestone streets, an Art Nouveau lamppost stood tall
+A tranquil tableau of in the quaint village square, a traditional wrought-iron streetlamp featured delicate filigree patterns and amber-hued glass panels
+A tranquil tableau of the lampposts were adorned with Art Deco motifs, their geometric shapes and frosted glass creating a sense of vintage glamour
+In a still frame, in the picturesque square, a Gothic-style lamppost adorned with intricate stone carvings added a touch of medieval charm to the setting
+In a still frame, in the heart of the old city, a row of ornate lantern-style streetlamps bathed the narrow alleyway in a warm, welcoming light
+A tranquil tableau of in the heart of the Utah desert, a massive sandstone arch spanned the horizon
+A tranquil tableau of in the Arizona desert, a massive stone bridge arched across a rugged canyon
+A tranquil tableau of in the corner of the minimalist tea room, a bonsai tree added a touch of nature's beauty to the otherwise simple and elegant space
+In a still frame, amidst the hushed ambiance of the traditional tea room, a meticulously arranged tea set awaited, with porcelain cups, a bamboo whisk
+In a still frame, nestled in the Zen garden, a rustic teahouse featured tatami seating and a traditional charcoal brazier
+A tranquil tableau of a country estate's library featured elegant wooden shelves
+A tranquil tableau of beneath the shade of a solitary oak tree, an old wooden park bench sat patiently
+A tranquil tableau of beside a tranquil pond, a weeping willow tree draped its branches gracefully over the water's surface, creating a serene tableau of reflection and calm
+A tranquil tableau of in the Zen garden, a perfectly raked gravel path led to a serene rock garden
+In a still frame, a tranquil pond was fringed by weeping cherry trees, their blossoms drifting lazily onto the glassy surface
+In a still frame, within the historic library's reading room, rows of antique leather chairs and mahogany tables offered a serene haven for literary contemplation
+A tranquil tableau of a peaceful orchid garden showcased a variety of delicate blooms
+A tranquil tableau of in the serene courtyard, a centuries-old stone well stood as a symbol of a bygone era, its mossy stones bearing witness to the passage of time
+a bird and a cat
+a cat and a dog
+a dog and a horse
+a horse and a sheep
+a sheep and a cow
+a cow and an elephant
+an elephant and a bear
+a bear and a zebra
+a zebra and a giraffe
+a giraffe and a bird
+a chair and a couch
+a couch and a potted plant
+a potted plant and a tv
+a tv and a laptop
+a laptop and a remote
+a remote and a keyboard
+a keyboard and a cell phone
+a cell phone and a book
+a book and a clock
+a clock and a backpack
+a backpack and an umbrella
+an umbrella and a handbag
+a handbag and a tie
+a tie and a suitcase
+a suitcase and a vase
+a vase and scissors
+scissors and a teddy bear
+a teddy bear and a frisbee
+a frisbee and skis
+skis and a snowboard
+a snowboard and a sports ball
+a sports ball and a kite
+a kite and a baseball bat
+a baseball bat and a baseball glove
+a baseball glove and a skateboard
+a skateboard and a surfboard
+a surfboard and a tennis racket
+a tennis racket and a bottle
+a bottle and a chair
+an airplane and a train
+a train and a boat
+a boat and an airplane
+a bicycle and a car
+a car and a motorcycle
+a motorcycle and a bus
+a bus and a traffic light
+a traffic light and a fire hydrant
+a fire hydrant and a stop sign
+a stop sign and a parking meter
+a parking meter and a truck
+a truck and a bicycle
+a toilet and a hair drier
+a hair drier and a toothbrush
+a toothbrush and a sink
+a sink and a toilet
+a wine glass and a chair
+a cup and a couch
+a fork and a potted plant
+a knife and a tv
+a spoon and a laptop
+a bowl and a remote
+a banana and a keyboard
+an apple and a cell phone
+a sandwich and a book
+an orange and a clock
+broccoli and a backpack
+a carrot and an umbrella
+a hot dog and a handbag
+a pizza and a tie
+a donut and a suitcase
+a cake and a vase
+an oven and scissors
+a toaster and a teddy bear
+a microwave and a frisbee
+a refrigerator and skis
+a bicycle and an airplane
+a car and a train
+a motorcycle and a boat
+a person and a toilet
+a person and a hair drier
+a person and a toothbrush
+a person and a sink
+A person is riding a bike
+A person is marching
+A person is roller skating
+A person is tasting beer
+A person is clapping
+A person is drawing
+A person is petting animal (not cat)
+A person is eating watermelon
+A person is playing harp
+A person is wrestling
+A person is riding scooter
+A person is sweeping floor
+A person is skateboarding
+A person is dunking basketball
+A person is playing flute
+A person is stretching leg
+A person is tying tie
+A person is skydiving
+A person is shooting goal (soccer)
+A person is playing piano
+A person is finger snapping
+A person is canoeing or kayaking
+A person is laughing
+A person is digging
+A person is clay pottery making
+A person is shooting basketball
+A person is bending back
+A person is shaking hands
+A person is bandaging
+A person is push up
+A person is catching or throwing frisbee
+A person is playing trumpet
+A person is flying kite
+A person is filling eyebrows
+A person is shuffling cards
+A person is folding clothes
+A person is smoking
+A person is tai chi
+A person is squat
+A person is playing controller
+A person is throwing axe
+A person is giving or receiving award
+A person is air drumming
+A person is taking a shower
+A person is planting trees
+A person is sharpening knives
+A person is robot dancing
+A person is rock climbing
+A person is hula hooping
+A person is writing
+A person is bungee jumping
+A person is pushing cart
+A person is cleaning windows
+A person is cutting watermelon
+A person is cheerleading
+A person is washing hands
+A person is ironing
+A person is cutting nails
+A person is hugging
+A person is trimming or shaving beard
+A person is jogging
+A person is making bed
+A person is washing dishes
+A person is grooming dog
+A person is doing laundry
+A person is knitting
+A person is reading book
+A person is baby waking up
+A person is massaging legs
+A person is brushing teeth
+A person is crawling baby
+A person is motorcycling
+A person is driving car
+A person is sticking tongue out
+A person is shaking head
+A person is sword fighting
+A person is doing aerobics
+A person is strumming guitar
+A person is riding or walking with horse
+A person is archery
+A person is catching or throwing baseball
+A person is playing chess
+A person is rock scissors paper
+A person is using computer
+A person is arranging flowers
+A person is bending metal
+A person is ice skating
+A person is climbing a rope
+A person is crying
+A person is dancing ballet
+A person is getting a haircut
+A person is running on treadmill
+A person is kissing
+A person is counting money
+A person is barbequing
+A person is peeling apples
+A person is milking cow
+A person is shining shoes
+A person is making snowman
+A person is sailing
+a person swimming in ocean
+a person giving a presentation to a room full of colleagues
+a person washing the dishes
+a person eating a burger
+a person walking in the snowstorm
+a person drinking coffee in a cafe
+a person playing guitar
+a bicycle leaning against a tree
+a bicycle gliding through a snowy field
+a bicycle slowing down to stop
+a bicycle accelerating to gain speed
+a car stuck in traffic during rush hour
+a car turning a corner
+a car slowing down to stop
+a car accelerating to gain speed
+a motorcycle cruising along a coastal highway
+a motorcycle turning a corner
+a motorcycle slowing down to stop
+a motorcycle gliding through a snowy field
+a motorcycle accelerating to gain speed
+an airplane soaring through a clear blue sky
+an airplane taking off
+an airplane landing smoothly on a runway
+an airplane accelerating to gain speed
+a bus turning a corner
+a bus stuck in traffic during rush hour
+a bus accelerating to gain speed
+a train speeding down the tracks
+a train crossing over a tall bridge
+a train accelerating to gain speed
+a truck turning a corner
+a truck anchored in a tranquil bay
+a truck stuck in traffic during rush hour
+a truck slowing down to stop
+a truck accelerating to gain speed
+a boat sailing smoothly on a calm lake
+a boat slowing down to stop
+a boat accelerating to gain speed
+a bird soaring gracefully in the sky
+a bird building a nest from twigs and leaves
+a bird flying over a snowy forest
+a cat grooming itself meticulously with its tongue
+a cat playing in park
+a cat drinking water
+a cat running happily
+a dog enjoying a peaceful walk
+a dog playing in park
+a dog drinking water
+a dog running happily
+a horse bending down to drink water from a river
+a horse galloping across an open field
+a horse taking a peaceful walk
+a horse running to join a herd of its kind
+a sheep bending down to drink water from a river
+a sheep taking a peaceful walk
+a sheep running to join a herd of its kind
+a cow bending down to drink water from a river
+a cow chewing cud while resting in a tranquil barn
+a cow running to join a herd of its kind
+an elephant spraying itself with water using its trunk to cool down
+an elephant taking a peaceful walk
+an elephant running to join a herd of its kind
+a bear catching a salmon in its powerful jaws
+a bear sniffing the air for scents of food
+a bear climbing a tree
+a bear hunting for prey
+a zebra bending down to drink water from a river
+a zebra running to join a herd of its kind
+a zebra taking a peaceful walk
+a giraffe bending down to drink water from a river
+a giraffe taking a peaceful walk
+a giraffe running to join a herd of its kind
+a person
+a bicycle
+a car
+a motorcycle
+an airplane
+a bus
+a train
+a truck
+a boat
+a traffic light
+a fire hydrant
+a stop sign
+a parking meter
+a bench
+a bird
+a cat
+a dog
+a horse
+a sheep
+a cow
+an elephant
+a bear
+a zebra
+a giraffe
+a backpack
+an umbrella
+a handbag
+a tie
+a suitcase
+a frisbee
+skis
+a snowboard
+a sports ball
+a kite
+a baseball bat
+a baseball glove
+a skateboard
+a surfboard
+a tennis racket
+a bottle
+a wine glass
+a cup
+a fork
+a knife
+a spoon
+a bowl
+a banana
+an apple
+a sandwich
+an orange
+broccoli
+a carrot
+a hot dog
+a pizza
+a donut
+a cake
+a chair
+a couch
+a potted plant
+a bed
+a dining table
+a toilet
+a tv
+a laptop
+a remote
+a keyboard
+a cell phone
+a microwave
+an oven
+a toaster
+a sink
+a refrigerator
+a book
+a clock
+a vase
+scissors
+a teddy bear
+a hair drier
+a toothbrush
+a red bicycle
+a green bicycle
+a blue bicycle
+a yellow bicycle
+an orange bicycle
+a purple bicycle
+a pink bicycle
+a black bicycle
+a white bicycle
+a red car
+a green car
+a blue car
+a yellow car
+an orange car
+a purple car
+a pink car
+a black car
+a white car
+a red bird
+a green bird
+a blue bird
+a yellow bird
+an orange bird
+a purple bird
+a pink bird
+a black bird
+a white bird
+a black cat
+a white cat
+an orange cat
+a yellow cat
+a red umbrella
+a green umbrella
+a blue umbrella
+a yellow umbrella
+an orange umbrella
+a purple umbrella
+a pink umbrella
+a black umbrella
+a white umbrella
+a red suitcase
+a green suitcase
+a blue suitcase
+a yellow suitcase
+an orange suitcase
+a purple suitcase
+a pink suitcase
+a black suitcase
+a white suitcase
+a red bowl
+a green bowl
+a blue bowl
+a yellow bowl
+an orange bowl
+a purple bowl
+a pink bowl
+a black bowl
+a white bowl
+a red chair
+a green chair
+a blue chair
+a yellow chair
+an orange chair
+a purple chair
+a pink chair
+a black chair
+a white chair
+a red clock
+a green clock
+a blue clock
+a yellow clock
+an orange clock
+a purple clock
+a pink clock
+a black clock
+a white clock
+a red vase
+a green vase
+a blue vase
+a yellow vase
+an orange vase
+a purple vase
+a pink vase
+a black vase
+a white vase
+A beautiful coastal beach in spring, waves lapping on sand, Van Gogh style
+A beautiful coastal beach in spring, waves lapping on sand, oil painting
+A beautiful coastal beach in spring, waves lapping on sand by Hokusai, in the style of Ukiyo
+A beautiful coastal beach in spring, waves lapping on sand, black and white
+A beautiful coastal beach in spring, waves lapping on sand, pixel art
+A beautiful coastal beach in spring, waves lapping on sand, in cyberpunk style
+A beautiful coastal beach in spring, waves lapping on sand, animated style
+A beautiful coastal beach in spring, waves lapping on sand, watercolor painting
+A beautiful coastal beach in spring, waves lapping on sand, surrealism style
+The bund Shanghai, Van Gogh style
+The bund Shanghai, oil painting
+The bund Shanghai by Hokusai, in the style of Ukiyo
+The bund Shanghai, black and white
+The bund Shanghai, pixel art
+The bund Shanghai, in cyberpunk style
+The bund Shanghai, animated style
+The bund Shanghai, watercolor painting
+The bund Shanghai, surrealism style
+a shark is swimming in the ocean, Van Gogh style
+a shark is swimming in the ocean, oil painting
+a shark is swimming in the ocean by Hokusai, in the style of Ukiyo
+a shark is swimming in the ocean, black and white
+a shark is swimming in the ocean, pixel art
+a shark is swimming in the ocean, in cyberpunk style
+a shark is swimming in the ocean, animated style
+a shark is swimming in the ocean, watercolor painting
+a shark is swimming in the ocean, surrealism style
+A panda drinking coffee in a cafe in Paris, Van Gogh style
+A panda drinking coffee in a cafe in Paris, oil painting
+A panda drinking coffee in a cafe in Paris by Hokusai, in the style of Ukiyo
+A panda drinking coffee in a cafe in Paris, black and white
+A panda drinking coffee in a cafe in Paris, pixel art
+A panda drinking coffee in a cafe in Paris, in cyberpunk style
+A panda drinking coffee in a cafe in Paris, animated style
+A panda drinking coffee in a cafe in Paris, watercolor painting
+A panda drinking coffee in a cafe in Paris, surrealism style
+A cute happy Corgi playing in park, sunset, Van Gogh style
+A cute happy Corgi playing in park, sunset, oil painting
+A cute happy Corgi playing in park, sunset by Hokusai, in the style of Ukiyo
+A cute happy Corgi playing in park, sunset, black and white
+A cute happy Corgi playing in park, sunset, pixel art
+A cute happy Corgi playing in park, sunset, in cyberpunk style
+A cute happy Corgi playing in park, sunset, animated style
+A cute happy Corgi playing in park, sunset, watercolor painting
+A cute happy Corgi playing in park, sunset, surrealism style
+Gwen Stacy reading a book, Van Gogh style
+Gwen Stacy reading a book, oil painting
+Gwen Stacy reading a book by Hokusai, in the style of Ukiyo
+Gwen Stacy reading a book, black and white
+Gwen Stacy reading a book, pixel art
+Gwen Stacy reading a book, in cyberpunk style
+Gwen Stacy reading a book, animated style
+Gwen Stacy reading a book, watercolor painting
+Gwen Stacy reading a book, surrealism style
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, Van Gogh style
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, oil painting
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background by Hokusai, in the style of Ukiyo
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, black and white
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, pixel art
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, in cyberpunk style
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, animated style
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, watercolor painting
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, surrealism style
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, Van Gogh style
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, oil painting
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas by Hokusai, in the style of Ukiyo
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, black and white
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, pixel art
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, in cyberpunk style
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, animated style
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, watercolor painting
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, surrealism style
+An astronaut flying in space, Van Gogh style
+An astronaut flying in space, oil painting
+An astronaut flying in space by Hokusai, in the style of Ukiyo
+An astronaut flying in space, black and white
+An astronaut flying in space, pixel art
+An astronaut flying in space, in cyberpunk style
+An astronaut flying in space, animated style
+An astronaut flying in space, watercolor painting
+An astronaut flying in space, surrealism style
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, Van Gogh style
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, oil painting
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks by Hokusai, in the style of Ukiyo
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, black and white
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, pixel art
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, in cyberpunk style
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, animated style
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, watercolor painting
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, surrealism style
+A beautiful coastal beach in spring, waves lapping on sand, in super slow motion
+A beautiful coastal beach in spring, waves lapping on sand, zoom in
+A beautiful coastal beach in spring, waves lapping on sand, zoom out
+A beautiful coastal beach in spring, waves lapping on sand, pan left
+A beautiful coastal beach in spring, waves lapping on sand, pan right
+A beautiful coastal beach in spring, waves lapping on sand, tilt up
+A beautiful coastal beach in spring, waves lapping on sand, tilt down
+A beautiful coastal beach in spring, waves lapping on sand, with an intense shaking effect
+A beautiful coastal beach in spring, waves lapping on sand, featuring a steady and smooth perspective
+A beautiful coastal beach in spring, waves lapping on sand, racking focus
+The bund Shanghai, in super slow motion
+The bund Shanghai, zoom in
+The bund Shanghai, zoom out
+The bund Shanghai, pan left
+The bund Shanghai, pan right
+The bund Shanghai, tilt up
+The bund Shanghai, tilt down
+The bund Shanghai, with an intense shaking effect
+The bund Shanghai, featuring a steady and smooth perspective
+The bund Shanghai, racking focus
+a shark is swimming in the ocean, in super slow motion
+a shark is swimming in the ocean, zoom in
+a shark is swimming in the ocean, zoom out
+a shark is swimming in the ocean, pan left
+a shark is swimming in the ocean, pan right
+a shark is swimming in the ocean, tilt up
+a shark is swimming in the ocean, tilt down
+a shark is swimming in the ocean, with an intense shaking effect
+a shark is swimming in the ocean, featuring a steady and smooth perspective
+a shark is swimming in the ocean, racking focus
+A panda drinking coffee in a cafe in Paris, in super slow motion
+A panda drinking coffee in a cafe in Paris, zoom in
+A panda drinking coffee in a cafe in Paris, zoom out
+A panda drinking coffee in a cafe in Paris, pan left
+A panda drinking coffee in a cafe in Paris, pan right
+A panda drinking coffee in a cafe in Paris, tilt up
+A panda drinking coffee in a cafe in Paris, tilt down
+A panda drinking coffee in a cafe in Paris, with an intense shaking effect
+A panda drinking coffee in a cafe in Paris, featuring a steady and smooth perspective
+A panda drinking coffee in a cafe in Paris, racking focus
+A cute happy Corgi playing in park, sunset, in super slow motion
+A cute happy Corgi playing in park, sunset, zoom in
+A cute happy Corgi playing in park, sunset, zoom out
+A cute happy Corgi playing in park, sunset, pan left
+A cute happy Corgi playing in park, sunset, pan right
+A cute happy Corgi playing in park, sunset, tilt up
+A cute happy Corgi playing in park, sunset, tilt down
+A cute happy Corgi playing in park, sunset, with an intense shaking effect
+A cute happy Corgi playing in park, sunset, featuring a steady and smooth perspective
+A cute happy Corgi playing in park, sunset, racking focus
+Gwen Stacy reading a book, in super slow motion
+Gwen Stacy reading a book, zoom in
+Gwen Stacy reading a book, zoom out
+Gwen Stacy reading a book, pan left
+Gwen Stacy reading a book, pan right
+Gwen Stacy reading a book, tilt up
+Gwen Stacy reading a book, tilt down
+Gwen Stacy reading a book, with an intense shaking effect
+Gwen Stacy reading a book, featuring a steady and smooth perspective
+Gwen Stacy reading a book, racking focus
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, in super slow motion
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, zoom in
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, zoom out
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, pan left
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, pan right
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, tilt up
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, tilt down
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, with an intense shaking effect
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, featuring a steady and smooth perspective
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background, racking focus
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, in super slow motion
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, zoom in
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, zoom out
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, pan left
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, pan right
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, tilt up
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, tilt down
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, with an intense shaking effect
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, featuring a steady and smooth perspective
+A couple in formal evening wear going home get caught in a heavy downpour with umbrellas, racking focus
+An astronaut flying in space, in super slow motion
+An astronaut flying in space, zoom in
+An astronaut flying in space, zoom out
+An astronaut flying in space, pan left
+An astronaut flying in space, pan right
+An astronaut flying in space, tilt up
+An astronaut flying in space, tilt down
+An astronaut flying in space, with an intense shaking effect
+An astronaut flying in space, featuring a steady and smooth perspective
+An astronaut flying in space, racking focus
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, in super slow motion
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, zoom in
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, zoom out
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, pan left
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, pan right
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, tilt up
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, tilt down
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, with an intense shaking effect
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, featuring a steady and smooth perspective
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks, racking focus
+Close up of grapes on a rotating table.
+Turtle swimming in ocean.
+A storm trooper vacuuming the beach.
+A panda standing on a surfboard in the ocean in sunset.
+An astronaut feeding ducks on a sunny afternoon, reflection from the water.
+Two pandas discussing an academic paper.
+Sunset time lapse at the beach with moving clouds and colors in the sky.
+A fat rabbit wearing a purple robe walking through a fantasy landscape.
+A koala bear playing piano in the forest.
+An astronaut flying in space.
+Fireworks.
+An animated painting of fluffy white clouds moving in sky.
+Flying through fantasy landscapes.
+A bigfoot walking in the snowstorm.
+A squirrel eating a burger.
+A cat wearing sunglasses and working as a lifeguard at a pool.
+Snow rocky mountains peaks canyon. snow blanketed rocky mountains surround and shadow deep canyons. the canyons twist and bend through the high elevated mountain peaks.
+Splash of turquoise water in extreme slow motion, alpha channel included.
+an ice cream is melting on the table.
+a drone flying over a snowy forest.
+a shark is swimming in the ocean.
+Aerial panoramic video from a drone of a fantasy land.
+a teddy bear is swimming in the ocean.
+time lapse of sunrise on mars.
+golden fish swimming in the ocean.
+An artist brush painting on a canvas close up.
+A drone view of celebration with Christmas tree and fireworks, starry sky - background.
+happy dog wearing a yellow turtleneck, studio, portrait, facing camera, dark background
+Origami dancers in white paper, 3D render, on white background, studio shot, dancing modern dance.
+Campfire at night in a snowy forest with starry sky in the background.
+a fantasy landscape
+A 3D model of a 1800s victorian house.
+this is how I do makeup in the morning.
+A raccoon that looks like a turtle, digital art.
+Robot dancing in Times Square.
+Busy freeway at night.
+Balloon full of water exploding in extreme slow motion.
+An astronaut is riding a horse in the space in a photorealistic style.
+Macro slo-mo. Slow motion cropped closeup of roasted coffee beans falling into an empty bowl.
+Sewing machine, old sewing machine working.
+Motion colour drop in water, ink swirling in water, colourful ink in water, abstraction fancy dream cloud of ink.
+Few big purple plums rotating on the turntable. water drops appear on the skin during rotation. isolated on the white background. close-up. macro.
+Vampire makeup face of beautiful girl, red contact lenses.
+Ashtray full of butts on table, smoke flowing on black background, close-up
+Pacific coast, carmel by the sea ocean and waves.
+A teddy bear is playing drum kit in NYC Times Square.
+A corgi is playing drum kit.
+An Iron man is playing the electronic guitar, high electronic guitar.
+A raccoon is playing the electronic guitar.
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background by Vincent van Gogh
+A corgi's head depicted as an explosion of a nebula
+A fantasy landscape
+A future where humans have achieved teleportation technology
+A jellyfish floating through the ocean, with bioluminescent tentacles
+A Mars rover moving on Mars
+A panda drinking coffee in a cafe in Paris
+A space shuttle launching into orbit, with flames and smoke billowing out from the engines
+A steam train moving on a mountainside
+A super cool giant robot in Cyberpunk Beijing
+A tropical beach at sunrise, with palm trees and crystal-clear water in the foreground
+Cinematic shot of Van Gogh's selfie, Van Gogh style
+Gwen Stacy reading a book
+Iron Man flying in the sky
+The bund Shanghai, oil painting
+Yoda playing guitar on the stage
+A beautiful coastal beach in spring, waves lapping on sand by Hokusai, in the style of Ukiyo
+A beautiful coastal beach in spring, waves lapping on sand by Vincent van Gogh
+A boat sailing leisurely along the Seine River with the Eiffel Tower in background
+A car moving slowly on an empty street, rainy evening
+A cat eating food out of a bowl
+A cat wearing sunglasses at a pool
+A confused panda in calculus class
+A cute fluffy panda eating Chinese food in a restaurant
+A cute happy Corgi playing in park, sunset
+A cute raccoon playing guitar in a boat on the ocean
+A happy fuzzy panda playing guitar nearby a campfire, snow mountain in the background
+A lightning striking atop of eiffel tower, dark clouds in the sky
+A modern art museum, with colorful paintings
+A panda cooking in the kitchen
+A panda playing on a swing set
+A polar bear is playing guitar
+A raccoon dressed in suit playing the trumpet, stage background
+A robot DJ is playing the turntable, in heavy raining futuristic tokyo rooftop cyberpunk night, sci-fi, fantasy
+A shark swimming in clear Caribbean ocean
+A super robot protecting city
+A teddy bear washing the dishes
+An epic tornado attacking above a glowing city at night, the tornado is made of smoke
+An oil painting of a couple in formal evening wear going home get caught in a heavy downpour with umbrellas
+Clown fish swimming through the coral reef
+Hyper-realistic spaceship landing on Mars
+The bund Shanghai, vibrant color
+Vincent van Gogh is painting in the room
+Yellow flowers swing in the wind
+alley
+amusement park
+aquarium
+arch
+art gallery
+bathroom
+bakery shop
+ballroom
+bar
+barn
+basement
+beach
+bedroom
+bridge
+botanical garden
+cafeteria
+campsite
+campus
+carrousel
+castle
+cemetery
+classroom
+cliff
+crosswalk
+construction site
+corridor
+courtyard
+desert
+downtown
+driveway
+farm
+food court
+football field
+forest road
+fountain
+gas station
+glacier
+golf course
+indoor gymnasium
+harbor
+highway
+hospital
+house
+iceberg
+industrial area
+jail cell
+junkyard
+kitchen
+indoor library
+lighthouse
+laboratory
+mansion
+marsh
+mountain
+indoor movie theater
+indoor museum
+music studio
+nursery
+ocean
+office
+palace
+parking lot
+pharmacy
+phone booth
+raceway
+restaurant
+river
+science museum
+shower
+ski slope
+sky
+skyscraper
+baseball stadium
+staircase
+street
+supermarket
+indoor swimming pool
+tower
+outdoor track
+train railway
+train station platform
+underwater coral reef
+valley
+volcano
+waterfall
+windmill
+a bicycle on the left of a car, front view
+a car on the right of a motorcycle, front view
+a motorcycle on the left of a bus, front view
+a bus on the right of a traffic light, front view
+a traffic light on the left of a fire hydrant, front view
+a fire hydrant on the right of a stop sign, front view
+a stop sign on the left of a parking meter, front view
+a parking meter on the right of a bench, front view
+a bench on the left of a truck, front view
+a truck on the right of a bicycle, front view
+a bird on the left of a cat, front view
+a cat on the right of a dog, front view
+a dog on the left of a horse, front view
+a horse on the right of a sheep, front view
+a sheep on the left of a cow, front view
+a cow on the right of an elephant, front view
+an elephant on the left of a bear, front view
+a bear on the right of a zebra, front view
+a zebra on the left of a giraffe, front view
+a giraffe on the right of a bird, front view
+a bottle on the left of a wine glass, front view
+a wine glass on the right of a cup, front view
+a cup on the left of a fork, front view
+a fork on the right of a knife, front view
+a knife on the left of a spoon, front view
+a spoon on the right of a bowl, front view
+a bowl on the left of a bottle, front view
+a potted plant on the left of a remote, front view
+a remote on the right of a clock, front view
+a clock on the left of a vase, front view
+a vase on the right of scissors, front view
+scissors on the left of a teddy bear, front view
+a teddy bear on the right of a potted plant, front view
+a frisbee on the left of a sports ball, front view
+a sports ball on the right of a baseball bat, front view
+a baseball bat on the left of a baseball glove, front view
+a baseball glove on the right of a tennis racket, front view
+a tennis racket on the left of a frisbee, front view
+a toilet on the left of a hair drier, front view
+a hair drier on the right of a toothbrush, front view
+a toothbrush on the left of a sink, front view
+a sink on the right of a toilet, front view
+a chair on the left of a couch, front view
+a couch on the right of a bed, front view
+a bed on the left of a tv, front view
+a tv on the right of a dining table, front view
+a dining table on the left of a chair, front view
+an airplane on the left of a train, front view
+a train on the right of a boat, front view
+a boat on the left of an airplane, front view
+an oven on the top of a toaster, front view
+an oven on the bottom of a toaster, front view
+a toaster on the top of a microwave, front view
+a toaster on the bottom of a microwave, front view
+a microwave on the top of an oven, front view
+a microwave on the bottom of an oven, front view
+a banana on the top of an apple, front view
+a banana on the bottom of an apple, front view
+an apple on the top of a sandwich, front view
+an apple on the bottom of a sandwich, front view
+a sandwich on the top of an orange, front view
+a sandwich on the bottom of an orange, front view
+an orange on the top of a carrot, front view
+an orange on the bottom of a carrot, front view
+a carrot on the top of a hot dog, front view
+a carrot on the bottom of a hot dog, front view
+a hot dog on the top of a pizza, front view
+a hot dog on the bottom of a pizza, front view
+a pizza on the top of a donut, front view
+a pizza on the bottom of a donut, front view
+a donut on the top of broccoli, front view
+a donut on the bottom of broccoli, front view
+broccoli on the top of a banana, front view
+broccoli on the bottom of a banana, front view
+skis on the top of a snowboard, front view
+skis on the bottom of a snowboard, front view
+a snowboard on the top of a kite, front view
+a snowboard on the bottom of a kite, front view
+a kite on the top of a skateboard, front view
+a kite on the bottom of a skateboard, front view
+a skateboard on the top of a surfboard, front view
+a skateboard on the bottom of a surfboard, front view
+a surfboard on the top of skis, front view
+a surfboard on the bottom of skis, front view

requirements.txt

@@ -0,0 +1,43 @@
+torch>=2.4.0
+torchvision>=0.19.0
+opencv-python>=4.9.0.80
+diffusers==0.31.0
+transformers>=4.49.0
+tokenizers>=0.20.3
+accelerate>=1.1.1
+tqdm
+imageio
+easydict
+ftfy
+dashscope
+imageio-ffmpeg
+numpy==1.24.4
+wandb
+omegaconf
+einops
+av==13.1.0
+opencv-python
+git+https://github.com/openai/CLIP.git
+open_clip_torch
+starlette
+pycocotools
+lmdb
+matplotlib
+sentencepiece
+pydantic==2.10.6
+scikit-image
+huggingface_hub[cli]
+dominate
+nvidia-pyindex
+nvidia-tensorrt
+pycuda
+onnx
+onnxruntime
+onnxscript
+onnxconverter_common
+flask
+flask-socketio
+torchao
+hydra-core
+git+https://github.com/facebookresearch/segment-anything.git
+

scripts/create_lmdb_14b_shards.py

@@ -0,0 +1,103 @@
+"""
+python create_lmdb_14b_shards.py \
+--data_path /mnt/localssd/wanx_14b_data \
+--lmdb_path /mnt/localssd/wanx_14B_shift-3.0_cfg-5.0_lmdb
+"""
+
+from tqdm import tqdm
+import numpy as np
+import argparse
+import torch
+import lmdb
+import glob
+import os
+
+from utils.lmdb import store_arrays_to_lmdb, process_data_dict
+
+
+def main():
+    """
+    Aggregate all ode pairs inside a folder into a lmdb dataset.
+    Each pt file should contain a (key, value) pair representing a
+    video's ODE trajectories.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_path", type=str, required=True, help="path to ode pairs")
+    parser.add_argument("--lmdb_path", type=str, required=True, help="path to lmdb")
+    parser.add_argument("--num_shards", type=int, default=16, help="num_shards")
+
+    args = parser.parse_args()
+
+    all_dirs = sorted(os.listdir(args.data_path))
+
+    # figure out the maximum map size needed
+    map_size = int(1e12)  # adapt to your need, set to 1TB by default
+    os.makedirs(args.lmdb_path, exist_ok=True)
+    # 1) Open one LMDB env per shard
+    envs = []
+    num_shards = args.num_shards
+    for shard_id in range(num_shards):
+        print("shard_id ", shard_id)
+        path = os.path.join(args.lmdb_path, f"shard_{shard_id}")
+        env = lmdb.open(
+            path,
+            map_size=map_size,
+            subdir=True,  # set to True if you want a directory per env
+            readonly=False,
+            metasync=True,
+            sync=True,
+            lock=True,
+            readahead=False,
+            meminit=False,
+        )
+        envs.append(env)
+
+    counters = [0] * num_shards
+    seen_prompts = set()  # for deduplication
+    total_samples = 0
+    all_files = []
+
+    for part_dir in all_dirs:
+        all_files += sorted(glob.glob(os.path.join(args.data_path, part_dir, "*.pt")))
+
+    # 2) Prepare a write transaction for each shard
+    for idx, file in tqdm(enumerate(all_files)):
+        try:
+            data_dict = torch.load(file)
+            data_dict = process_data_dict(data_dict, seen_prompts)
+        except Exception as e:
+            print(f"Error processing {file}: {e}")
+            continue
+
+        if data_dict["latents"].shape != (1, 21, 16, 60, 104):
+            continue
+
+        shard_id = idx % num_shards
+        # write to lmdb file
+        store_arrays_to_lmdb(envs[shard_id], data_dict, start_index=counters[shard_id])
+        counters[shard_id] += len(data_dict["prompts"])
+        data_shape = data_dict["latents"].shape
+
+    total_samples += len(all_files)
+
+    print(len(seen_prompts))
+
+    # save each entry's shape to lmdb
+    for shard_id, env in enumerate(envs):
+        with env.begin(write=True) as txn:
+            for key, val in data_dict.items():
+                assert len(data_shape) == 5
+                array_shape = np.array(data_shape)  # val.shape)
+                array_shape[0] = counters[shard_id]
+                shape_key = f"{key}_shape".encode()
+                print(shape_key, array_shape)
+                shape_str = " ".join(map(str, array_shape))
+                txn.put(shape_key, shape_str.encode())
+
+    print(
+        f"Finished writing {total_samples} examples into {num_shards} shards under {args.lmdb_path}"
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/create_lmdb_iterative.py

@@ -0,0 +1,58 @@
+from tqdm import tqdm
+import numpy as np
+import argparse
+import torch
+import lmdb
+import glob
+import os
+
+from utils.lmdb import store_arrays_to_lmdb, process_data_dict
+
+
+def main():
+    """
+    Aggregate all ode pairs inside a folder into a lmdb dataset.
+    Each pt file should contain a (key, value) pair representing a
+    video's ODE trajectories.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_path", type=str, required=True, help="path to ode pairs")
+    parser.add_argument("--lmdb_path", type=str, required=True, help="path to lmdb")
+
+    args = parser.parse_args()
+
+    all_files = sorted(glob.glob(os.path.join(args.data_path, "*.pt")))
+
+    # figure out the maximum map size needed
+    total_array_size = 5000000000000  # adapt to your need, set to 5TB by default
+
+    env = lmdb.open(args.lmdb_path, map_size=total_array_size * 2)
+
+    counter = 0
+
+    seen_prompts = set()  # for deduplication
+
+    for index, file in tqdm(enumerate(all_files)):
+        # read from disk
+        data_dict = torch.load(file)
+
+        data_dict = process_data_dict(data_dict, seen_prompts)
+
+        # write to lmdb file
+        store_arrays_to_lmdb(env, data_dict, start_index=counter)
+        counter += len(data_dict["prompts"])
+
+    # save each entry's shape to lmdb
+    with env.begin(write=True) as txn:
+        for key, val in data_dict.items():
+            print(key, val)
+            array_shape = np.array(val.shape)
+            array_shape[0] = counter
+
+            shape_key = f"{key}_shape".encode()
+            shape_str = " ".join(map(str, array_shape))
+            txn.put(shape_key, shape_str.encode())
+
+
+if __name__ == "__main__":
+    main()

scripts/generate_ode_pairs.py

@@ -0,0 +1,113 @@
+from utils.distributed import launch_distributed_job
+from utils.scheduler import FlowMatchScheduler
+from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder
+from utils.dataset import TextDataset
+import torch.distributed as dist
+from tqdm import tqdm
+import argparse
+import torch
+import math
+import os
+
+
+def init_model(
+    device,
+):
+    model = WanDiffusionWrapper().to(device).to(torch.float32)
+    encoder = WanTextEncoder().to(device).to(torch.float32)
+    model.model.requires_grad_(False)
+
+    scheduler = FlowMatchScheduler(shift=8.0, sigma_min=0.0, extra_one_step=True)
+    scheduler.set_timesteps(num_inference_steps=48, denoising_strength=1.0)
+    scheduler.sigmas = scheduler.sigmas.to(device)
+
+    sample_neg_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+
+    unconditional_dict = encoder(text_prompts=[sample_neg_prompt])
+
+    return model, encoder, scheduler, unconditional_dict
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local_rank", type=int, default=-1)
+    parser.add_argument("--output_folder", type=str)
+    parser.add_argument("--caption_path", type=str)
+    parser.add_argument("--guidance_scale", type=float, default=6.0)
+
+    args = parser.parse_args()
+
+    # launch_distributed_job()
+    launch_distributed_job()
+
+    device = torch.cuda.current_device()
+
+    torch.set_grad_enabled(False)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    model, encoder, scheduler, unconditional_dict = init_model(device=device)
+
+    dataset = TextDataset(args.caption_path)
+
+    # if global_rank == 0:
+    os.makedirs(args.output_folder, exist_ok=True)
+
+    for index in tqdm(
+        range(int(math.ceil(len(dataset) / dist.get_world_size()))),
+        disable=dist.get_rank() != 0,
+    ):
+        prompt_index = index * dist.get_world_size() + dist.get_rank()
+        if prompt_index >= len(dataset):
+            continue
+        prompt = dataset[prompt_index]
+
+        conditional_dict = encoder(text_prompts=prompt)
+
+        latents = torch.randn([1, 21, 16, 60, 104], dtype=torch.float32, device=device)
+
+        noisy_input = []
+
+        for progress_id, t in enumerate(tqdm(scheduler.timesteps)):
+            timestep = t * torch.ones([1, 21], device=device, dtype=torch.float32)
+
+            noisy_input.append(latents)
+
+            _, x0_pred_cond = model(latents, conditional_dict, timestep)
+
+            _, x0_pred_uncond = model(latents, unconditional_dict, timestep)
+
+            x0_pred = x0_pred_uncond + args.guidance_scale * (x0_pred_cond - x0_pred_uncond)
+
+            flow_pred = model._convert_x0_to_flow_pred(
+                scheduler=scheduler,
+                x0_pred=x0_pred.flatten(0, 1),
+                xt=latents.flatten(0, 1),
+                timestep=timestep.flatten(0, 1),
+            ).unflatten(0, x0_pred.shape[:2])
+
+            latents = scheduler.step(
+                flow_pred.flatten(0, 1),
+                scheduler.timesteps[progress_id]
+                * torch.ones([1, 21], device=device, dtype=torch.long).flatten(0, 1),
+                latents.flatten(0, 1),
+            ).unflatten(dim=0, sizes=flow_pred.shape[:2])
+
+        noisy_input.append(latents)
+
+        noisy_inputs = torch.stack(noisy_input, dim=1)
+
+        noisy_inputs = noisy_inputs[:, [0, 12, 24, 36, -1]]
+
+        stored_data = noisy_inputs
+
+        torch.save(
+            {prompt: stored_data.cpu().detach()},
+            os.path.join(args.output_folder, f"{prompt_index:05d}.pt"),
+        )
+
+    dist.barrier()
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -0,0 +1,7 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="self_forcing",
+    version="0.0.1",
+    packages=find_packages(),
+)

stream_drag_inference_wrapper.py

@@ -0,0 +1,201 @@
+import torch
+import numpy as np
+from omegaconf import DictConfig
+from optimize_utils import MultiTrajectory
+from stream_inference_wrapper import StreamInferenceWrapper
+
+
+def _extract_block_trajectories(
+    multi_traj: MultiTrajectory,
+) -> tuple[
+    list[list[dict[str, bool | list[torch.Tensor]]]],
+    list[np.ndarray],
+    np.ndarray | None,
+]:
+    """Extract block_trajectories from a MultiTrajectory in the format expected by begin_optimize.
+
+    Returns:
+        block_trajectories: block_num x N x trajectory dict
+            Each trajectory dict has keys 'is_rotation', 'deltas', 'start_point',
+            and optionally 'rotation_center'.
+        masks: list of N masks corresponding to each trajectory
+        movable_mask: the movable area mask for the whole image
+    """
+    if multi_traj.trajectories is None or len(multi_traj.trajectories) == 0:
+        return [], [], None
+
+    movable_mask = multi_traj.movable_mask
+
+    # Collect per-trajectory masks
+    masks = [traj.mask for traj in multi_traj.trajectories]
+
+    # Find the maximum number of blocks across all trajectories
+    max_blocks = (
+        max(
+            len(traj.block_trajectories)
+            for traj in multi_traj.trajectories
+            if traj.block_trajectories
+        )
+        if any(traj.block_trajectories for traj in multi_traj.trajectories)
+        else 0
+    )
+
+    if max_blocks == 0:
+        return [], masks, movable_mask
+
+    block_trajectories = []
+    for block_idx in range(max_blocks):
+        block = []
+        for traj in multi_traj.trajectories:
+            if traj.block_trajectories and block_idx < len(traj.block_trajectories):
+                block.append(traj.block_trajectories[block_idx])
+            else:
+                # Provide an empty placeholder
+                block.append(
+                    {
+                        "is_rotation": False,
+                        "deltas": [],
+                        "start_point": (0, 0),
+                    }
+                )
+        block_trajectories.append(block)
+
+    # Assert: the N of each block in block_trajectories should equal the length of masks
+    for block_idx, block in enumerate(block_trajectories):
+        assert len(block) == len(masks), (
+            f"Block {block_idx} has {len(block)} trajectories, " f"but there are {len(masks)} masks"
+        )
+
+    assert ((len(block_trajectories) == 0) and (movable_mask is None)) or (
+        (len(block_trajectories) > 0) and (movable_mask is not None)
+    ), "block_trajectories and movable_mask must both be present or both be absent"
+
+    return block_trajectories, masks, movable_mask
+
+
+class StreamDragInferenceWrapper(StreamInferenceWrapper):
+    def __init__(
+        self,
+        stream_model_config: DictConfig,
+        checkpoint_path: str,
+        total_generate_block_number: int,
+        use_ema: bool = True,
+        seed: int = 0,
+    ):
+        super().__init__(
+            stream_model_config=stream_model_config,
+            checkpoint_path=checkpoint_path,
+            total_generate_block_number=total_generate_block_number,
+            use_ema=use_ema,
+            seed=seed,
+        )
+        self.previous_record_feature_list = None
+
+    def inference(
+        self,
+        start_block_index: int,
+        end_block_index: int,
+        prompt: str,
+        # below are for drag optimization
+        multiple_trajectory: MultiTrajectory = None,
+    ):
+        assert start_block_index >= 0
+        assert end_block_index > start_block_index
+        print(f"""
+{self.__class__.__name__}.inference():
+    {start_block_index = }  |  {end_block_index = }
+""")
+        sampled_noise = self.get_sampled_noise(start_block_index, end_block_index)
+        prompts = [prompt]
+
+        # Extract block_trajectories, masks, and movable_mask from multiple_trajectory
+        drag_optimize_target_latent_index = -1
+        if multiple_trajectory is not None:
+            block_trajectories, masks, movable_mask = _extract_block_trajectories(
+                multiple_trajectory
+            )
+            assert multiple_trajectory.drag_or_animation_select in [
+                "Drag",
+                "Animation",
+            ]
+            if multiple_trajectory.drag_or_animation_select == "Drag":
+                drag_optimize_target_latent_index = 2
+        else:
+            block_trajectories, masks, movable_mask = [], [], None
+
+        if len(block_trajectories) > 0:
+            is_drag_optimize = True
+        else:
+            is_drag_optimize = False
+
+        initial_latents = self.get_initial_latents(
+            start_block_index,
+        )
+        if initial_latents is not None:
+            print(f"{initial_latents.shape = }")
+
+        print(f"{block_trajectories = }")
+        print(f"{len(masks) = }")
+        latents_result = self.pipeline.inference(
+            noise=sampled_noise,
+            text_prompts=prompts,
+            return_latents=True,
+            initial_latent=initial_latents,
+            do_not_decode_video=True,
+            do_not_recompute_initial_latents=True,
+            # below are for drag optimization
+            model_config=self.stream_model_config,
+            previous_record_feature_list=self.previous_record_feature_list,
+            is_drag_optimize=is_drag_optimize,
+            block_trajectories=block_trajectories,
+            masks=masks,
+            movable_mask=movable_mask,
+            drag_optimize_target_latent_index=drag_optimize_target_latent_index,
+        )
+        if self.stream_model_config.drag_optim_config.record_feature_block_indexes:
+            latents, record_attention_values_list = latents_result
+        else:
+            latents = latents_result
+            record_attention_values_list = None
+        if self.recorded_latents is None:
+            self.recorded_latents = latents
+        else:
+            self.recorded_latents = torch.concat(
+                [
+                    self.recorded_latents[:, :0],
+                    latents,
+                ],
+                dim=1,
+            )
+
+        if record_attention_values_list is not None:
+
+            def dict_first_value(d: dict):
+                return next(iter(d.values()))
+
+            print(f"{record_attention_values_list.keys() = }")  # denoising timesteps
+            print(
+                f"{dict_first_value(record_attention_values_list).keys() = }"
+            )  # attention block layers
+            print(
+                f"{dict_first_value(dict_first_value(record_attention_values_list)).keys() = }"
+            )  # attention types name
+            print(
+                f"{dict_first_value(dict_first_value(dict_first_value(record_attention_values_list))).shape = }"
+            )  # [1, 3, 30, 52, 1536]
+        else:
+            print(f"{record_attention_values_list = }")
+        self.previous_record_feature_list = record_attention_values_list
+
+        self.decode_and_update_video(start_block_index, end_block_index)
+
+        return (
+            self.video,
+            self.video[self.block_to_video_index(start_block_index) :],
+        )
+
+    def reset(
+        self,
+    ):
+        super().reset()
+        self.previous_record_feature_list = None

stream_inference.py

@@ -0,0 +1,111 @@
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+from tqdm import tqdm
+from torchvision.io import write_video
+from torch.utils.data import DataLoader, SequentialSampler
+
+from stream_inference_wrapper import StreamInferenceWrapper
+
+# from stream_drag_inference_wrapper import StreamDragInferenceWrapper
+from utils.dataset import TextDataset
+from utils.misc import set_seed
+from hydra import initialize, compose
+from hydra.core.global_hydra import GlobalHydra
+
+from demo_utils.memory import gpu, get_cuda_free_memory_gb
+
+
+def main():
+
+    output_block_number = 27
+
+    config_dir = "configs"
+    stream_config_name = "self_forcing_dmd_vsink_stream"
+    # stream_config_name = "self_forcing_dmd_vsink_stream_drag"
+
+    data_path = "prompts/MovieGenVideoBench_extended.txt"
+
+    seed = 42
+    set_seed(seed)
+
+    output_folder = "outputs-stream"
+    output_folder = f"{output_folder}/blk{output_block_number}-{stream_config_name}-seed{seed}"
+
+    print(f"Free VRAM {get_cuda_free_memory_gb(gpu)} GB")
+    # low_memory = get_cuda_free_memory_gb(gpu) < 40
+
+    # Create dataset
+    dataset = TextDataset(prompt_path=data_path)
+    num_prompts = len(dataset)
+    print(f"Number of prompts: {num_prompts}")
+
+    sampler = SequentialSampler(dataset)
+    dataloader = DataLoader(dataset, batch_size=1, sampler=sampler, num_workers=0, drop_last=False)
+
+    os.makedirs(output_folder, exist_ok=True)
+
+    if GlobalHydra.instance().is_initialized():
+        GlobalHydra.instance().clear()
+
+    with initialize(version_base=None, config_path=config_dir):
+        stream_config = compose(config_name=stream_config_name)
+    print(f"{stream_config = }")
+
+    stream_inference = StreamInferenceWrapper(
+        stream_model_config=stream_config,
+        checkpoint_path="./checkpoints/self_forcing_dmd.pt",
+        total_generate_block_number=output_block_number,
+        use_ema=True,
+        seed=seed,
+    )
+
+    for i, batch_data in tqdm(enumerate(dataloader)):
+        idx = batch_data["idx"].item()
+        print(f"{idx = }")
+
+        # For DataLoader batch_size=1, the batch_data is already a single item, but in a batch container
+        # Unpack the batch data for convenience
+        if isinstance(batch_data, dict):
+            batch = batch_data
+        elif isinstance(batch_data, list):
+            batch = batch_data[0]  # First (and only) item in the batch
+
+        # For text-to-video, batch is just the text prompt
+        prompt = batch["prompts"][0]
+        print(f"{prompt = }")
+        extended_prompt = batch["extended_prompts"][0] if "extended_prompts" in batch else None
+        print(f"{extended_prompt = }")
+
+        set_seed(seed)
+        stream_inference.reset()
+
+        current_block_index = 0
+        block_step = 3
+        while current_block_index < output_block_number:
+            end_block_index = current_block_index + block_step
+            all_video, current_video = stream_inference.inference(
+                start_block_index=current_block_index,
+                end_block_index=end_block_index,
+                prompt=prompt,
+            )
+
+            # Save the video if the current prompt is not a dummy prompt
+            if idx < num_prompts:
+                current_video_output_path = os.path.join(
+                    output_folder,
+                    f"{idx:04d}-{prompt[:50].replace(' ', '_')}-{current_block_index:02d}-{end_block_index:02d}.mp4",
+                )
+                write_video(current_video_output_path, current_video, fps=16)
+                all_video_output_path = os.path.join(
+                    output_folder,
+                    f"{idx:04d}-{prompt[:50].replace(' ', '_')}-{0:02d}-{end_block_index:02d}.mp4",
+                )
+                write_video(all_video_output_path, all_video, fps=16)
+
+            current_block_index = end_block_index
+
+
+if __name__ == "__main__":
+    main()