Upload 31 files

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 2024 Genmo
+
+   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.

pyproject.toml

@@ -0,0 +1,29 @@
+[project]
+name = "genmo"
+version = "0.1.0"
+description = "Genmo models"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "addict>=2.4.0",
+    "click>=8.1.7",
+    "einops>=0.8.0",
+    "gradio>=3.36.1",
+    "omegaconf>=2.3.0",
+    "pillow>=11.0.0",
+    "pyyaml>=6.0.2",
+    "ray>=2.37.0",
+    "sentencepiece>=0.2.0",
+    "setuptools>=75.2.0",
+    "torch>=2.4.1",
+    "transformers>=4.45.2",
+]
+
+[project.optional-dependencies]
+flash = [
+    "flash-attn>=2.6.3",
+]
+
+[tool.ruff]
+# Allow lines to be as long as 120.
+line-length = 120

pyrightconfig.json

@@ -0,0 +1,4 @@
+{
+    "include": ["src/genmo/mochi_preview/pipelines.py"]
+}
+  

requirements.txt

@@ -0,0 +1,12 @@
+addict>=2.4.0
+click>=8.1.7
+einops>=0.8.0
+gradio>=3.36.1
+omegaconf>=2.3.0
+pillow>=11.0.0
+pyyaml>=6.0.2
+ray>=2.37.0
+sentencepiece>=0.2.0
+setuptools>=75.2.0
+torch>=2.4.1
+transformers>=4.45.2

scripts/format.bash

@@ -0,0 +1,4 @@
+#! /bin/bash
+set -euxo pipefail
+ruff format src
+ruff check --fix --select I src

scripts/pytorch_to_safe_tensors.py

@@ -0,0 +1,20 @@
+#! /usr/bin/env python3
+from pathlib import Path
+
+import click
+import torch
+from safetensors.torch import save_file
+
+
+@click.command()
+@click.argument("input_path", type=click.Path(exists=True))
+def convert_to_safetensors(input_path):
+    model = torch.load(input_path)
+    input_path = Path(input_path)
+    output_path = input_path.with_suffix(".safetensors")
+    save_file(model, str(output_path))
+    click.echo(f"Converted {input_path} to {output_path}")
+
+
+if __name__ == "__main__":
+    convert_to_safetensors()

scripts/typecheck.bash

@@ -0,0 +1,2 @@
+#! /bin/bash
+npx pyright

src/genmo.egg-info/PKG-INFO

@@ -0,0 +1,154 @@
+Metadata-Version: 2.1
+Name: genmo
+Version: 0.1.0
+Summary: Genmo models
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: addict>=2.4.0
+Requires-Dist: click>=8.1.7
+Requires-Dist: einops>=0.8.0
+Requires-Dist: gradio>=3.36.1
+Requires-Dist: omegaconf>=2.3.0
+Requires-Dist: pillow>=11.0.0
+Requires-Dist: pyyaml>=6.0.2
+Requires-Dist: ray>=2.37.0
+Requires-Dist: sentencepiece>=0.2.0
+Requires-Dist: setuptools>=75.2.0
+Requires-Dist: torch>=2.4.1
+Requires-Dist: transformers>=4.45.2
+Provides-Extra: flash
+Requires-Dist: flash-attn>=2.6.3; extra == "flash"
+
+# Mochi 1
+[Blog](https://www.genmo.ai/blog) | [Hugging Face](https://huggingface.co/genmo/mochi-1-preview) | [Playground](https://www.genmo.ai/play) | [Careers](https://jobs.ashbyhq.com/genmo)
+
+A state of the art video generation model by [Genmo](https://genmo.ai).
+
+https://github.com/user-attachments/assets/4d268d02-906d-4cb0-87cc-f467f1497108
+
+## Overview
+
+Mochi 1 preview is an open state-of-the-art video generation model with high-fidelity motion and strong prompt adherence in preliminary evaluation. This model dramatically closes the gap between closed and open video generation systems. We’re releasing the model under a permissive Apache 2.0 license. Try this model for free on [our playground](https://genmo.ai/play).
+
+## Installation
+
+Install using [uv](https://github.com/astral-sh/uv):
+
+```bash
+git clone https://github.com/genmoai/models
+cd models 
+pip install uv
+uv venv .venv
+source .venv/bin/activate
+uv pip install -e . --no-build-isolation
+```
+
+If you want to install flash attention, you can use:
+```
+uv pip install -e .[flash] --no-build-isolation
+```
+
+You will also need to install [FFMPEG](https://www.ffmpeg.org/) to turn your outputs into videos.
+
+## Download Weights
+
+Download the weights from [Hugging Face](https://huggingface.co/genmo/mochi-1-preview/tree/main) or via `magnet:?xt=urn:btih:441da1af7a16bcaa4f556964f8028d7113d21cbb&dn=weights&tr=udp://tracker.opentrackr.org:1337/announce` to a folder on your computer.
+
+## Running
+
+Start the gradio UI with
+
+```bash
+python3 ./demos/gradio_ui.py --model_dir "<path_to_downloaded_directory>"
+```
+
+Or generate videos directly from the CLI with
+
+```bash
+python3 ./demos/cli.py --model_dir "<path_to_downloaded_directory>"
+```
+
+Replace `<path_to_downloaded_directory>` with the path to your model directory.
+
+## API
+
+This repository comes with a simple, composable API, so you can programmatically call the model. You can find a full example [here](demos/api_example.py). But, roughly, it looks like this:
+
+```python
+from genmo.mochi_preview.pipelines import (
+    DecoderModelFactory,
+    DitModelFactory,
+    MochiSingleGPUPipeline,
+    T5ModelFactory,
+    linear_quadratic_schedule,
+)
+
+pipeline = MochiSingleGPUPipeline(
+    text_encoder_factory=T5ModelFactory(),
+    dit_factory=DitModelFactory(
+        model_path=f"{MOCHI_DIR}/dit.safetensors", model_dtype="bf16"
+    ),
+    decoder_factory=DecoderModelFactory(
+        model_path=f"{MOCHI_DIR}/vae.safetensors",
+        model_stats_path=f"{MOCHI_DIR}/vae_stats.json",
+    ),
+    cpu_offload=True,
+    decode_type="tiled_full",
+)
+
+video = pipeline(
+    height=480,
+    width=848,
+    num_frames=31,
+    num_inference_steps=64,
+    sigma_schedule=linear_quadratic_schedule(64, 0.025),
+    cfg_schedule=[4.5] * 64,
+    batch_cfg=False,
+    prompt="your favorite prompt here ...",
+    negative_prompt="",
+    seed=12345,
+)
+```
+
+## Model Architecture
+
+Mochi 1 represents a significant advancement in open-source video generation, featuring a 10 billion parameter diffusion model built on our novel Asymmetric Diffusion Transformer (AsymmDiT) architecture. Trained entirely from scratch, it is the largest video generative model ever openly released. And best of all, it’s a simple, hackable architecture. Additionally, we are releasing an inference harness that includes an efficient context parallel implementation. 
+
+Alongside Mochi, we are open-sourcing our video AsymmVAE. We use an asymmetric encoder-decoder structure to build an efficient high quality compression model. Our AsymmVAE causally compresses videos to a 128x smaller size, with an 8x8 spatial and a 6x temporal compression to a 12-channel latent space. 
+
+### AsymmVAE Model Specs
+|Params <br> Count | Enc Base <br>  Channels | Dec Base <br> Channels |Latent <br> Dim | Spatial <br> Compression | Temporal <br> Compression | 
+|:--:|:--:|:--:|:--:|:--:|:--:|
+|362M   | 64  | 128  | 12   | 8x8   | 6x   | 
+
+An AsymmDiT efficiently processes user prompts alongside compressed video tokens by streamlining text processing and focusing neural network capacity on visual reasoning. AsymmDiT jointly attends to text and visual tokens with multi-modal self-attention and learns separate MLP layers for each modality, similar to Stable Diffusion 3. However, our visual stream has nearly 4 times as many parameters as the text stream via a larger hidden dimension. To unify the modalities in self-attention, we use non-square QKV and output projection layers. This asymmetric design reduces inference memory requirements.
+Many modern diffusion models use multiple pretrained language models to represent user prompts. In contrast, Mochi 1 simply encodes prompts with a single T5-XXL language model.
+
+### AsymmDiT Model Specs
+|Params <br> Count | Num <br> Layers | Num <br> Heads | Visual <br> Dim | Text <br> Dim | Visual <br> Tokens | Text <br> Tokens | 
+|:--:|:--:|:--:|:--:|:--:|:--:|:--:|
+|10B   | 48   | 24   | 3072   | 1536   | 44520   |   256   |
+
+## Hardware Requirements
+
+The model requires at least 4 H100 GPUs to run. We welcome contributions from the community to reduce this requirement.
+
+## Safety
+Genmo video models are general text-to-video diffusion models that inherently reflect the biases and preconceptions found in their training data. While steps have been taken to limit NSFW content, organizations should implement additional safety protocols and careful consideration before deploying these model weights in any commercial services or products.
+
+## Limitations
+Under the research preview, Mochi 1 is a living and evolving checkpoint. There are a few known limitations. The initial release generates videos at 480p today. In some edge cases with extreme motion, minor warping and distortions can also occur. Mochi 1 is also optimized for photorealistic styles so does not perform well with animated content. We also anticipate that the community will fine-tune the model to suit various aesthetic preferences.
+
+## Related Work
+- [ComfyUI-MochiWrapper](https://github.com/kijai/ComfyUI-MochiWrapper) adds ComfyUI support for Mochi. The integration of Pytorch's SDPA attention was taken from their repository.
+
+
+## BibTeX
+```
+@misc{genmo2024mochi,
+      title={Mochi},
+      author={Genmo Team},
+      year={2024}
+}
+```

src/genmo.egg-info/SOURCES.txt

@@ -0,0 +1,25 @@
+LICENSE
+README.md
+pyproject.toml
+src/genmo.egg-info/PKG-INFO
+src/genmo.egg-info/SOURCES.txt
+src/genmo.egg-info/dependency_links.txt
+src/genmo.egg-info/requires.txt
+src/genmo.egg-info/top_level.txt
+src/genmo/lib/attn_imports.py
+src/genmo/lib/progress.py
+src/genmo/lib/utils.py
+src/genmo/mochi_preview/__init__.py
+src/genmo/mochi_preview/pipelines.py
+src/genmo/mochi_preview/dit/joint_model/__init__.py
+src/genmo/mochi_preview/dit/joint_model/asymm_models_joint.py
+src/genmo/mochi_preview/dit/joint_model/context_parallel.py
+src/genmo/mochi_preview/dit/joint_model/layers.py
+src/genmo/mochi_preview/dit/joint_model/mod_rmsnorm.py
+src/genmo/mochi_preview/dit/joint_model/residual_tanh_gated_rmsnorm.py
+src/genmo/mochi_preview/dit/joint_model/rope_mixed.py
+src/genmo/mochi_preview/dit/joint_model/temporal_rope.py
+src/genmo/mochi_preview/dit/joint_model/utils.py
+src/genmo/mochi_preview/vae/__init__.py
+src/genmo/mochi_preview/vae/cp_conv.py
+src/genmo/mochi_preview/vae/model.py

src/genmo.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

src/genmo.egg-info/requires.txt

@@ -0,0 +1,15 @@
+addict>=2.4.0
+click>=8.1.7
+einops>=0.8.0
+gradio>=3.36.1
+omegaconf>=2.3.0
+pillow>=11.0.0
+pyyaml>=6.0.2
+ray>=2.37.0
+sentencepiece>=0.2.0
+setuptools>=75.2.0
+torch>=2.4.1
+transformers>=4.45.2
+
+[flash]
+flash-attn>=2.6.3

src/genmo.egg-info/top_level.txt

@@ -0,0 +1 @@
+genmo

src/genmo/lib/attn_imports.py

@@ -0,0 +1,35 @@
+from contextlib import contextmanager
+
+import torch
+
+try:
+    from flash_attn import flash_attn_varlen_qkvpacked_func as flash_varlen_qkvpacked_attn
+except ImportError:
+    flash_varlen_qkvpacked_attn = None
+
+try:
+    from sageattention import sageattn as sage_attn
+except ImportError:
+    sage_attn = None
+
+try:
+    from comfy.ldm.modules.attention import comfy_optimized_attention as comfy_attn
+except ImportError:
+    comfy_attn = None
+
+
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+backends = []
+if torch.cuda.get_device_properties(0).major < 7:
+    backends.append(SDPBackend.MATH)
+if torch.cuda.get_device_properties(0).major >= 9.0:
+    backends.append(SDPBackend.CUDNN_ATTENTION)
+else:
+    backends.append(SDPBackend.EFFICIENT_ATTENTION)
+
+
+@contextmanager
+def sdpa_attn_ctx():
+    with sdpa_kernel(backends):
+        yield

src/genmo/lib/progress.py

@@ -0,0 +1,87 @@
+import contextlib
+from typing import Any, Iterable, Iterator, Optional
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    tqdm = None
+
+try:
+    from ray.experimental.tqdm_ray import tqdm as ray_tqdm
+except:
+    ray_tqdm = None
+
+# Global state
+_current_progress_type = "none"
+_is_progress_bar_active = False
+
+
+class DummyProgressBar:
+    """A no-op progress bar that mimics tqdm interface"""
+
+    def __init__(self, iterable=None, **kwargs):
+        self.iterable = iterable
+
+    def __iter__(self):
+        return iter(self.iterable)
+
+    def update(self, n=1):
+        pass
+
+    def close(self):
+        pass
+
+    def set_description(self, desc):
+        pass
+
+
+def get_new_progress_bar(iterable: Optional[Iterable] = None, **kwargs) -> Any:
+    if not _is_progress_bar_active:
+        return DummyProgressBar(iterable=iterable, **kwargs)
+
+    if _current_progress_type == "tqdm":
+        if tqdm is None:
+            raise ImportError("tqdm is required but not installed. Please install tqdm to use the tqdm progress bar.")
+        return tqdm(iterable=iterable, **kwargs)
+    elif _current_progress_type == "ray_tqdm":
+        if ray_tqdm is None:
+            raise ImportError("ray is required but not installed. Please install ray to use the ray_tqdm progress bar.")
+        return ray_tqdm(iterable=iterable, **kwargs)
+    return DummyProgressBar(iterable=iterable, **kwargs)
+
+
+@contextlib.contextmanager
+def progress_bar(type: str = "none", enabled=True):
+    """
+    Context manager for setting progress bar type and options.
+
+    Args:
+        type: Type of progress bar ("none" or "tqdm")
+        **options: Options to pass to the progress bar (e.g., total, desc)
+
+    Raises:
+        ValueError: If progress bar type is invalid
+        RuntimeError: If progress bars are nested
+
+    Example:
+        with progress_bar(type="tqdm", total=100):
+            for i in get_new_progress_bar(range(100)):
+                process(i)
+    """
+    if type not in ("none", "tqdm", "ray_tqdm"):
+        raise ValueError("Progress bar type must be 'none' or 'tqdm' or 'ray_tqdm'")
+    if not enabled:
+        type = "none"
+    global _current_progress_type, _is_progress_bar_active
+
+    if _is_progress_bar_active:
+        raise RuntimeError("Nested progress bars are not supported")
+
+    _is_progress_bar_active = True
+    _current_progress_type = type
+
+    try:
+        yield
+    finally:
+        _is_progress_bar_active = False
+        _current_progress_type = "none"

src/genmo/lib/utils.py

@@ -0,0 +1,58 @@
+import os
+import subprocess
+import tempfile
+import time
+
+import numpy as np
+from PIL import Image
+
+from genmo.lib.progress import get_new_progress_bar
+
+class Timer:
+    def __init__(self):
+        self.times = {}  # Dictionary to store times per stage
+
+    def __call__(self, name):
+        print(f"Timing {name}")
+        return self.TimerContextManager(self, name)
+
+    def print_stats(self):
+        total_time = sum(self.times.values())
+        # Print table header
+        print("{:<20} {:>10} {:>10}".format("Stage", "Time(s)", "Percent"))
+        for name, t in self.times.items():
+            percent = (t / total_time) * 100 if total_time > 0 else 0
+            print("{:<20} {:>10.2f} {:>9.2f}%".format(name, t, percent))
+
+    class TimerContextManager:
+        def __init__(self, outer, name):
+            self.outer = outer  # Reference to the Timer instance
+            self.name = name
+            self.start_time = None
+
+        def __enter__(self):
+            self.start_time = time.perf_counter()
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            end_time = time.perf_counter()
+            elapsed = end_time - self.start_time
+            self.outer.times[self.name] = self.outer.times.get(self.name, 0) + elapsed
+
+
+def save_video(final_frames, output_path):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        frame_paths = []
+        for i, frame in enumerate(get_new_progress_bar(final_frames)):
+            frame = (frame * 255).astype(np.uint8)
+            frame_img = Image.fromarray(frame)
+            frame_path = os.path.join(tmpdir, f"frame_{i:04d}.png")
+            frame_img.save(frame_path)
+            frame_paths.append(frame_path)
+
+        frame_pattern = os.path.join(tmpdir, "frame_%04d.png")
+        ffmpeg_cmd = f"ffmpeg -y -r 30 -i {frame_pattern} -vcodec libx264 -pix_fmt yuv420p {output_path}"
+        try:
+            subprocess.run(ffmpeg_cmd, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        except subprocess.CalledProcessError as e:
+            print(f"Error occurred while running ffmpeg:\n{e.stderr.decode()}")

src/genmo/mochi_preview/dit/joint_model/asymm_models_joint.py

@@ -0,0 +1,629 @@
+import os
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch.nn.attention import sdpa_kernel
+
+import genmo.mochi_preview.dit.joint_model.context_parallel as cp
+from genmo.mochi_preview.dit.joint_model.layers import (
+    FeedForward,
+    PatchEmbed,
+    RMSNorm,
+    TimestepEmbedder,
+)
+from genmo.mochi_preview.dit.joint_model.mod_rmsnorm import modulated_rmsnorm
+from genmo.mochi_preview.dit.joint_model.residual_tanh_gated_rmsnorm import (
+    residual_tanh_gated_rmsnorm,
+)
+from genmo.mochi_preview.dit.joint_model.rope_mixed import (
+    compute_mixed_rotation,
+    create_position_matrix,
+)
+from genmo.mochi_preview.dit.joint_model.temporal_rope import apply_rotary_emb_qk_real
+from genmo.mochi_preview.dit.joint_model.utils import (
+    AttentionPool,
+    modulate,
+    pad_and_split_xy,
+    unify_streams,
+)
+
+COMPILE_FINAL_LAYER = os.environ.get("COMPILE_DIT") == "1"
+COMPILE_MMDIT_BLOCK = os.environ.get("COMPILE_DIT") == "1"
+
+from genmo.lib.attn_imports import comfy_attn, flash_varlen_qkvpacked_attn, sage_attn, sdpa_attn_ctx
+
+
+class AsymmetricAttention(nn.Module):
+    def __init__(
+        self,
+        dim_x: int,
+        dim_y: int,
+        num_heads: int = 8,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
+        update_y: bool = True,
+        out_bias: bool = True,
+        attention_mode: str = "flash",
+        softmax_scale: Optional[float] = None,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        self.attention_mode = attention_mode
+        self.dim_x = dim_x
+        self.dim_y = dim_y
+        self.num_heads = num_heads
+        self.head_dim = dim_x // num_heads
+        self.update_y = update_y
+        self.softmax_scale = softmax_scale
+        if dim_x % num_heads != 0:
+            raise ValueError(f"dim_x={dim_x} should be divisible by num_heads={num_heads}")
+
+        # Input layers.
+        self.qkv_bias = qkv_bias
+        self.qkv_x = nn.Linear(dim_x, 3 * dim_x, bias=qkv_bias, device=device)
+        # Project text features to match visual features (dim_y -> dim_x)
+        self.qkv_y = nn.Linear(dim_y, 3 * dim_x, bias=qkv_bias, device=device)
+
+        # Query and key normalization for stability.
+        assert qk_norm
+        self.q_norm_x = RMSNorm(self.head_dim, device=device)
+        self.k_norm_x = RMSNorm(self.head_dim, device=device)
+        self.q_norm_y = RMSNorm(self.head_dim, device=device)
+        self.k_norm_y = RMSNorm(self.head_dim, device=device)
+
+        # Output layers. y features go back down from dim_x -> dim_y.
+        self.proj_x = nn.Linear(dim_x, dim_x, bias=out_bias, device=device)
+        self.proj_y = nn.Linear(dim_x, dim_y, bias=out_bias, device=device) if update_y else nn.Identity()
+
+    def run_qkv_y(self, y):
+        cp_rank, cp_size = cp.get_cp_rank_size()
+        local_heads = self.num_heads // cp_size
+
+        if cp.is_cp_active():
+            # Only predict local heads.
+            assert not self.qkv_bias
+            W_qkv_y = self.qkv_y.weight.view(3, self.num_heads, self.head_dim, self.dim_y)
+            W_qkv_y = W_qkv_y.narrow(1, cp_rank * local_heads, local_heads)
+            W_qkv_y = W_qkv_y.reshape(3 * local_heads * self.head_dim, self.dim_y)
+            qkv_y = F.linear(y, W_qkv_y, None)  # (B, L, 3 * local_h * head_dim)
+        else:
+            qkv_y = self.qkv_y(y)  # (B, L, 3 * dim)
+
+        qkv_y = qkv_y.view(qkv_y.size(0), qkv_y.size(1), 3, local_heads, self.head_dim)
+        q_y, k_y, v_y = qkv_y.unbind(2)
+        return q_y, k_y, v_y
+
+    def prepare_qkv(
+        self,
+        x: torch.Tensor,  # (B, N, dim_x)
+        y: torch.Tensor,  # (B, L, dim_y)
+        *,
+        scale_x: torch.Tensor,
+        scale_y: torch.Tensor,
+        rope_cos: torch.Tensor,
+        rope_sin: torch.Tensor,
+        valid_token_indices: torch.Tensor,
+    ):
+        # Pre-norm for visual features
+        x = modulated_rmsnorm(x, scale_x)  # (B, M, dim_x) where M = N / cp_group_size
+
+        # Process visual features
+        qkv_x = self.qkv_x(x)  # (B, M, 3 * dim_x)
+        assert qkv_x.dtype == torch.bfloat16
+        qkv_x = cp.all_to_all_collect_tokens(qkv_x, self.num_heads)  # (3, B, N, local_h, head_dim)
+
+        # Process text features
+        y = modulated_rmsnorm(y, scale_y)  # (B, L, dim_y)
+        q_y, k_y, v_y = self.run_qkv_y(y)  # (B, L, local_heads, head_dim)
+        q_y = self.q_norm_y(q_y)
+        k_y = self.k_norm_y(k_y)
+
+        # Split qkv_x into q, k, v
+        q_x, k_x, v_x = qkv_x.unbind(0)  # (B, N, local_h, head_dim)
+        q_x = self.q_norm_x(q_x)
+        q_x = apply_rotary_emb_qk_real(q_x, rope_cos, rope_sin)
+        k_x = self.k_norm_x(k_x)
+        k_x = apply_rotary_emb_qk_real(k_x, rope_cos, rope_sin)
+
+        # Unite streams
+        qkv = unify_streams(
+            q_x,
+            k_x,
+            v_x,
+            q_y,
+            k_y,
+            v_y,
+            valid_token_indices,
+        )
+
+        return qkv
+
+    def flash_attention(self, qkv, cu_seqlens, max_seqlen_in_batch, total, local_dim):
+        with torch.autocast("cuda", enabled=False):
+            out: torch.Tensor = flash_varlen_qkvpacked_attn(
+                qkv,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen_in_batch,
+                dropout_p=0.0,
+                softmax_scale=self.softmax_scale,
+            )  # (total, local_heads, head_dim)
+            return out.view(total, local_dim)
+
+    def sdpa_attention(self, qkv):
+        q, k, v = rearrange(qkv, "(b s) t h d -> t b h s d", b=1)
+        with torch.autocast("cuda", enabled=False):
+            with sdpa_attn_ctx():
+                out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
+                return rearrange(out, "b h s d -> s (b h d)")
+
+    def sage_attention(self, qkv):
+        q, k, v = rearrange(qkv, "(b s) t h d -> t b h s d", b=1)
+        with torch.autocast("cuda", enabled=False):
+            out = sage_attn(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
+            return rearrange(out, "b h s d -> s (b h d)")
+
+    def comfy_attention(self, qkv):
+        q, k, v = rearrange(qkv, "(b s) t h d -> t b h s d", b=1)
+        with torch.autocast("cuda", enabled=False):
+            out = comfy_attn(q, k, v, heads=self.num_heads, skip_reshape=True)
+            return out.squeeze(0)
+
+    @torch.compiler.disable()
+    def run_attention(
+        self,
+        qkv: torch.Tensor,  # (total <= B * (N + L), 3, local_heads, head_dim)
+        *,
+        B: int,
+        L: int,
+        M: int,
+        cu_seqlens: torch.Tensor,
+        max_seqlen_in_batch: int,
+        valid_token_indices: torch.Tensor,
+    ):
+        _, cp_size = cp.get_cp_rank_size()
+        N = cp_size * M
+        assert self.num_heads % cp_size == 0
+        local_heads = self.num_heads // cp_size
+        local_dim = local_heads * self.head_dim
+        total = qkv.size(0)
+
+        if self.attention_mode != "flash":
+            assert B == 1, f"Non-flash attention only supports batch size 1, got {B}"
+
+        if self.attention_mode == "flash":
+            out = self.flash_attention(qkv, cu_seqlens, max_seqlen_in_batch, total, local_dim)
+        elif self.attention_mode == "sdpa":
+            out = self.sdpa_attention(qkv)
+        elif self.attention_mode == "sage":
+            out = self.sage_attention(qkv)
+        elif self.attention_mode == "comfy":
+            out = self.comfy_attention(qkv)
+
+        x, y = pad_and_split_xy(out, valid_token_indices, B, N, L, qkv.dtype)
+        assert x.size() == (B, N, local_dim)
+        assert y.size() == (B, L, local_dim)
+
+        x = x.view(B, N, local_heads, self.head_dim)
+        x = cp.all_to_all_collect_heads(x)  # (B, M, dim_x = num_heads * head_dim)
+        x = self.proj_x(x)  # (B, M, dim_x)
+
+        if cp.is_cp_active():
+            y = cp.all_gather(y)  # (cp_size * B, L, local_heads * head_dim)
+            y = rearrange(y, "(G B) L D -> B L (G D)", G=cp_size, D=local_dim)  # (B, L, dim_x)
+        y = self.proj_y(y)  # (B, L, dim_y)
+        return x, y
+
+    def forward(
+        self,
+        x: torch.Tensor,  # (B, N, dim_x)
+        y: torch.Tensor,  # (B, L, dim_y)
+        *,
+        scale_x: torch.Tensor,  # (B, dim_x), modulation for pre-RMSNorm.
+        scale_y: torch.Tensor,  # (B, dim_y), modulation for pre-RMSNorm.
+        packed_indices: Dict[str, torch.Tensor] = None,
+        **rope_rotation,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass of asymmetric multi-modal attention.
+
+        Args:
+            x: (B, N, dim_x) tensor for visual tokens
+            y: (B, L, dim_y) tensor of text token features
+            packed_indices: Dict with keys for Flash Attention
+            num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens
+
+        Returns:
+            x: (B, N, dim_x) tensor of visual tokens after multi-modal attention
+            y: (B, L, dim_y) tensor of text token features after multi-modal attention
+        """
+        B, L, _ = y.shape
+        _, M, _ = x.shape
+
+        # Predict a packed QKV tensor from visual and text features.
+        # Don't checkpoint the all_to_all.
+        qkv = self.prepare_qkv(
+            x=x,
+            y=y,
+            scale_x=scale_x,
+            scale_y=scale_y,
+            rope_cos=rope_rotation.get("rope_cos"),
+            rope_sin=rope_rotation.get("rope_sin"),
+            valid_token_indices=packed_indices["valid_token_indices_kv"],
+        )  # (total <= B * (N + L), 3, local_heads, head_dim)
+
+        x, y = self.run_attention(
+            qkv,
+            B=B,
+            L=L,
+            M=M,
+            cu_seqlens=packed_indices["cu_seqlens_kv"],
+            max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"],
+            valid_token_indices=packed_indices["valid_token_indices_kv"],
+        )
+        return x, y
+
+
+@torch.compile(disable=not COMPILE_MMDIT_BLOCK)
+class AsymmetricJointBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size_x: int,
+        hidden_size_y: int,
+        num_heads: int,
+        *,
+        mlp_ratio_x: float = 8.0,  # Ratio of hidden size to d_model for MLP for visual tokens.
+        mlp_ratio_y: float = 4.0,  # Ratio of hidden size to d_model for MLP for text tokens.
+        update_y: bool = True,  # Whether to update text tokens in this block.
+        device: Optional[torch.device] = None,
+        **block_kwargs,
+    ):
+        super().__init__()
+        self.update_y = update_y
+        self.hidden_size_x = hidden_size_x
+        self.hidden_size_y = hidden_size_y
+        self.mod_x = nn.Linear(hidden_size_x, 4 * hidden_size_x, device=device)
+        if self.update_y:
+            self.mod_y = nn.Linear(hidden_size_x, 4 * hidden_size_y, device=device)
+        else:
+            self.mod_y = nn.Linear(hidden_size_x, hidden_size_y, device=device)
+
+        # Self-attention:
+        self.attn = AsymmetricAttention(
+            hidden_size_x,
+            hidden_size_y,
+            num_heads=num_heads,
+            update_y=update_y,
+            device=device,
+            **block_kwargs,
+        )
+
+        # MLP.
+        mlp_hidden_dim_x = int(hidden_size_x * mlp_ratio_x)
+        assert mlp_hidden_dim_x == int(1536 * 8)
+        self.mlp_x = FeedForward(
+            in_features=hidden_size_x,
+            hidden_size=mlp_hidden_dim_x,
+            multiple_of=256,
+            ffn_dim_multiplier=None,
+            device=device,
+        )
+
+        # MLP for text not needed in last block.
+        if self.update_y:
+            mlp_hidden_dim_y = int(hidden_size_y * mlp_ratio_y)
+            self.mlp_y = FeedForward(
+                in_features=hidden_size_y,
+                hidden_size=mlp_hidden_dim_y,
+                multiple_of=256,
+                ffn_dim_multiplier=None,
+                device=device,
+            )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        c: torch.Tensor,
+        y: torch.Tensor,
+        **attn_kwargs,
+    ):
+        """Forward pass of a block.
+
+        Args:
+            x: (B, N, dim) tensor of visual tokens
+            c: (B, dim) tensor of conditioned features
+            y: (B, L, dim) tensor of text tokens
+            num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens
+
+        Returns:
+            x: (B, N, dim) tensor of visual tokens after block
+            y: (B, L, dim) tensor of text tokens after block
+        """
+        N = x.size(1)
+
+        c = F.silu(c)
+        mod_x = self.mod_x(c)
+        scale_msa_x, gate_msa_x, scale_mlp_x, gate_mlp_x = mod_x.chunk(4, dim=1)
+
+        mod_y = self.mod_y(c)
+        if self.update_y:
+            scale_msa_y, gate_msa_y, scale_mlp_y, gate_mlp_y = mod_y.chunk(4, dim=1)
+        else:
+            scale_msa_y = mod_y
+
+        # Self-attention block.
+        x_attn, y_attn = self.attn(
+            x,
+            y,
+            scale_x=scale_msa_x,
+            scale_y=scale_msa_y,
+            **attn_kwargs,
+        )
+
+        assert x_attn.size(1) == N
+        x = residual_tanh_gated_rmsnorm(x, x_attn, gate_msa_x)
+        if self.update_y:
+            y = residual_tanh_gated_rmsnorm(y, y_attn, gate_msa_y)
+
+        # MLP block.
+        x = self.ff_block_x(x, scale_mlp_x, gate_mlp_x)
+        if self.update_y:
+            y = self.ff_block_y(y, scale_mlp_y, gate_mlp_y)
+
+        return x, y
+
+    def ff_block_x(self, x, scale_x, gate_x):
+        x_mod = modulated_rmsnorm(x, scale_x)
+        x_res = self.mlp_x(x_mod)
+        x = residual_tanh_gated_rmsnorm(x, x_res, gate_x)  # Sandwich norm
+        return x
+
+    def ff_block_y(self, y, scale_y, gate_y):
+        y_mod = modulated_rmsnorm(y, scale_y)
+        y_res = self.mlp_y(y_mod)
+        y = residual_tanh_gated_rmsnorm(y, y_res, gate_y)  # Sandwich norm
+        return y
+
+
+@torch.compile(disable=not COMPILE_FINAL_LAYER)
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        patch_size,
+        out_channels,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, device=device)
+        self.mod = nn.Linear(hidden_size, 2 * hidden_size, device=device)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, device=device)
+
+    def forward(self, x, c):
+        c = F.silu(c)
+        shift, scale = self.mod(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class AsymmDiTJoint(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+
+    Ingests text embeddings instead of a label.
+    """
+
+    def __init__(
+        self,
+        *,
+        patch_size=2,
+        in_channels=4,
+        hidden_size_x=1152,
+        hidden_size_y=1152,
+        depth=48,
+        num_heads=16,
+        mlp_ratio_x=8.0,
+        mlp_ratio_y=4.0,
+        t5_feat_dim: int = 4096,
+        t5_token_length: int = 256,
+        patch_embed_bias: bool = True,
+        timestep_mlp_bias: bool = True,
+        timestep_scale: Optional[float] = None,
+        use_extended_posenc: bool = False,
+        rope_theta: float = 10000.0,
+        device: Optional[torch.device] = None,
+        **block_kwargs,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.hidden_size_x = hidden_size_x
+        self.hidden_size_y = hidden_size_y
+        self.head_dim = hidden_size_x // num_heads  # Head dimension and count is determined by visual.
+        self.use_extended_posenc = use_extended_posenc
+        self.t5_token_length = t5_token_length
+        self.t5_feat_dim = t5_feat_dim
+        self.rope_theta = rope_theta  # Scaling factor for frequency computation for temporal RoPE.
+
+        self.x_embedder = PatchEmbed(
+            patch_size=patch_size,
+            in_chans=in_channels,
+            embed_dim=hidden_size_x,
+            bias=patch_embed_bias,
+            device=device,
+        )
+        # Conditionings
+        # Timestep
+        self.t_embedder = TimestepEmbedder(hidden_size_x, bias=timestep_mlp_bias, timestep_scale=timestep_scale)
+
+        # Caption Pooling (T5)
+        self.t5_y_embedder = AttentionPool(t5_feat_dim, num_heads=8, output_dim=hidden_size_x, device=device)
+
+        # Dense Embedding Projection (T5)
+        self.t5_yproj = nn.Linear(t5_feat_dim, hidden_size_y, bias=True, device=device)
+
+        # Initialize pos_frequencies as an empty parameter.
+        self.pos_frequencies = nn.Parameter(torch.empty(3, self.num_heads, self.head_dim // 2, device=device))
+
+        # for depth 48:
+        #  b =  0: AsymmetricJointBlock, update_y=True
+        #  b =  1: AsymmetricJointBlock, update_y=True
+        #  ...
+        #  b = 46: AsymmetricJointBlock, update_y=True
+        #  b = 47: AsymmetricJointBlock, update_y=False. No need to update text features.
+        blocks = []
+        for b in range(depth):
+            # Joint multi-modal block
+            update_y = b < depth - 1
+            block = AsymmetricJointBlock(
+                hidden_size_x,
+                hidden_size_y,
+                num_heads,
+                mlp_ratio_x=mlp_ratio_x,
+                mlp_ratio_y=mlp_ratio_y,
+                update_y=update_y,
+                device=device,
+                **block_kwargs,
+            )
+
+            blocks.append(block)
+        self.blocks = nn.ModuleList(blocks)
+
+        self.final_layer = FinalLayer(hidden_size_x, patch_size, self.out_channels, device=device)
+
+    def embed_x(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (B, C=12, T, H, W) tensor of visual tokens
+
+        Returns:
+            x: (B, C=3072, N) tensor of visual tokens with positional embedding.
+        """
+        return self.x_embedder(x)  # Convert BcTHW to BCN
+
+    @torch.compile(disable=not COMPILE_MMDIT_BLOCK)
+    def prepare(
+        self,
+        x: torch.Tensor,
+        sigma: torch.Tensor,
+        t5_feat: torch.Tensor,
+        t5_mask: torch.Tensor,
+    ):
+        """Prepare input and conditioning embeddings."""
+
+        with torch.profiler.record_function("x_emb_pe"):
+            # Visual patch embeddings with positional encoding.
+            T, H, W = x.shape[-3:]
+            pH, pW = H // self.patch_size, W // self.patch_size
+            x = self.embed_x(x)  # (B, N, D), where N = T * H * W / patch_size ** 2
+            assert x.ndim == 3
+            B = x.size(0)
+
+        with torch.profiler.record_function("rope_cis"):
+            # Construct position array of size [N, 3].
+            # pos[:, 0] is the frame index for each location,
+            # pos[:, 1] is the row index for each location, and
+            # pos[:, 2] is the column index for each location.
+            pH, pW = H // self.patch_size, W // self.patch_size
+            N = T * pH * pW
+            assert x.size(1) == N
+            pos = create_position_matrix(T, pH=pH, pW=pW, device=x.device, dtype=torch.float32)  # (N, 3)
+            rope_cos, rope_sin = compute_mixed_rotation(
+                freqs=self.pos_frequencies, pos=pos
+            )  # Each are (N, num_heads, dim // 2)
+
+        with torch.profiler.record_function("t_emb"):
+            # Global vector embedding for conditionings.
+            c_t = self.t_embedder(1 - sigma)  # (B, D)
+
+        with torch.profiler.record_function("t5_pool"):
+            # Pool T5 tokens using attention pooler
+            # Note y_feat[1] contains T5 token features.
+            assert (
+                t5_feat.size(1) == self.t5_token_length
+            ), f"Expected L={self.t5_token_length}, got {t5_feat.shape} for y_feat."
+            t5_y_pool = self.t5_y_embedder(t5_feat, t5_mask)  # (B, D)
+            assert t5_y_pool.size(0) == B, f"Expected B={B}, got {t5_y_pool.shape} for t5_y_pool."
+
+        c = c_t + t5_y_pool
+
+        y_feat = self.t5_yproj(t5_feat)  # (B, L, t5_feat_dim) --> (B, L, D)
+
+        return x, c, y_feat, rope_cos, rope_sin
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        sigma: torch.Tensor,
+        y_feat: List[torch.Tensor],
+        y_mask: List[torch.Tensor],
+        packed_indices: Dict[str, torch.Tensor] = None,
+        rope_cos: torch.Tensor = None,
+        rope_sin: torch.Tensor = None,
+    ):
+        """Forward pass of DiT.
+
+        Args:
+            x: (B, C, T, H, W) tensor of spatial inputs (images or latent representations of images)
+            sigma: (B,) tensor of noise standard deviations
+            y_feat: List((B, L, y_feat_dim) tensor of caption token features. For SDXL text encoders: L=77, y_feat_dim=2048)
+            y_mask: List((B, L) boolean tensor indicating which tokens are not padding)
+            packed_indices: Dict with keys for Flash Attention. Result of compute_packed_indices.
+        """
+        B, _, T, H, W = x.shape
+
+        # Use EFFICIENT_ATTENTION backend for T5 pooling, since we have a mask.
+        # Have to call sdpa_kernel outside of a torch.compile region.
+        with sdpa_kernel(torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION):
+            x, c, y_feat, rope_cos, rope_sin = self.prepare(x, sigma, y_feat[0], y_mask[0])
+        del y_mask
+
+        cp_rank, cp_size = cp.get_cp_rank_size()
+        N = x.size(1)
+        M = N // cp_size
+        assert N % cp_size == 0, f"Visual sequence length ({x.shape[1]}) must be divisible by cp_size ({cp_size})."
+
+        if cp_size > 1:
+            x = x.narrow(1, cp_rank * M, M)
+
+            assert self.num_heads % cp_size == 0
+            local_heads = self.num_heads // cp_size
+            rope_cos = rope_cos.narrow(1, cp_rank * local_heads, local_heads)
+            rope_sin = rope_sin.narrow(1, cp_rank * local_heads, local_heads)
+
+        for i, block in enumerate(self.blocks):
+            x, y_feat = block(
+                x,
+                c,
+                y_feat,
+                rope_cos=rope_cos,
+                rope_sin=rope_sin,
+                packed_indices=packed_indices,
+            )  # (B, M, D), (B, L, D)
+        del y_feat  # Final layers don't use dense text features.
+
+        x = self.final_layer(x, c)  # (B, M, patch_size ** 2 * out_channels)
+
+        patch = x.size(2)
+        x = cp.all_gather(x)
+        x = rearrange(x, "(G B) M P -> B (G M) P", G=cp_size, P=patch)
+        x = rearrange(
+            x,
+            "B (T hp wp) (p1 p2 c) -> B c T (hp p1) (wp p2)",
+            T=T,
+            hp=H // self.patch_size,
+            wp=W // self.patch_size,
+            p1=self.patch_size,
+            p2=self.patch_size,
+            c=self.out_channels,
+        )
+
+        return x

src/genmo/mochi_preview/dit/joint_model/context_parallel.py

@@ -0,0 +1,155 @@
+import torch
+import torch.distributed as dist
+from einops import rearrange
+
+_CONTEXT_PARALLEL_GROUP = None
+_CONTEXT_PARALLEL_RANK = None
+_CONTEXT_PARALLEL_GROUP_SIZE = None
+_CONTEXT_PARALLEL_GROUP_RANKS = None
+
+
+def get_cp_rank_size():
+    if _CONTEXT_PARALLEL_GROUP:
+        return _CONTEXT_PARALLEL_RANK, _CONTEXT_PARALLEL_GROUP_SIZE
+    else:
+        return 0, 1
+
+
+def local_shard(x: torch.Tensor, dim: int = 2) -> torch.Tensor:
+    if not _CONTEXT_PARALLEL_GROUP:
+        return x
+
+    cp_rank, cp_size = get_cp_rank_size()
+    return x.tensor_split(cp_size, dim=dim)[cp_rank]
+
+
+def set_cp_group(cp_group, ranks, global_rank):
+    global _CONTEXT_PARALLEL_GROUP, _CONTEXT_PARALLEL_RANK, _CONTEXT_PARALLEL_GROUP_SIZE, _CONTEXT_PARALLEL_GROUP_RANKS
+    if _CONTEXT_PARALLEL_GROUP is not None:
+        raise RuntimeError("CP group already initialized.")
+    _CONTEXT_PARALLEL_GROUP = cp_group
+    _CONTEXT_PARALLEL_RANK = dist.get_rank(cp_group)
+    _CONTEXT_PARALLEL_GROUP_SIZE = dist.get_world_size(cp_group)
+    _CONTEXT_PARALLEL_GROUP_RANKS = ranks
+
+    assert _CONTEXT_PARALLEL_RANK == ranks.index(
+        global_rank
+    ), f"Rank mismatch: {global_rank} in {ranks} does not have position {_CONTEXT_PARALLEL_RANK} "
+    assert _CONTEXT_PARALLEL_GROUP_SIZE == len(
+        ranks
+    ), f"Group size mismatch: {_CONTEXT_PARALLEL_GROUP_SIZE} != len({ranks})"
+
+
+def get_cp_group():
+    if _CONTEXT_PARALLEL_GROUP is None:
+        raise RuntimeError("CP group not initialized")
+    return _CONTEXT_PARALLEL_GROUP
+
+
+def is_cp_active():
+    return _CONTEXT_PARALLEL_GROUP is not None
+
+
+class AllGatherIntoTensorFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, reduce_dtype, group: dist.ProcessGroup):
+        ctx.reduce_dtype = reduce_dtype
+        ctx.group = group
+        ctx.batch_size = x.size(0)
+        group_size = dist.get_world_size(group)
+
+        x = x.contiguous()
+        output = torch.empty(group_size * x.size(0), *x.shape[1:], dtype=x.dtype, device=x.device)
+        dist.all_gather_into_tensor(output, x, group=group)
+        return output
+
+
+def all_gather(tensor: torch.Tensor) -> torch.Tensor:
+    if not _CONTEXT_PARALLEL_GROUP:
+        return tensor
+
+    return AllGatherIntoTensorFunction.apply(tensor, torch.float32, _CONTEXT_PARALLEL_GROUP)
+
+
+@torch.compiler.disable()
+def _all_to_all_single(output, input, group):
+    # Disable compilation since torch compile changes contiguity.
+    assert input.is_contiguous(), "Input tensor must be contiguous."
+    assert output.is_contiguous(), "Output tensor must be contiguous."
+    return dist.all_to_all_single(output, input, group=group)
+
+
+class CollectTokens(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, qkv: torch.Tensor, group: dist.ProcessGroup, num_heads: int):
+        """Redistribute heads and receive tokens.
+
+        Args:
+            qkv: query, key or value. Shape: [B, M, 3 * num_heads * head_dim]
+
+        Returns:
+            qkv: shape: [3, B, N, local_heads, head_dim]
+
+        where M is the number of local tokens,
+        N = cp_size * M is the number of global tokens,
+        local_heads = num_heads // cp_size is the number of local heads.
+        """
+        ctx.group = group
+        ctx.num_heads = num_heads
+        cp_size = dist.get_world_size(group)
+        assert num_heads % cp_size == 0
+        ctx.local_heads = num_heads // cp_size
+
+        qkv = rearrange(
+            qkv,
+            "B M (qkv G h d) -> G M h B (qkv d)",
+            qkv=3,
+            G=cp_size,
+            h=ctx.local_heads,
+        ).contiguous()
+
+        output_chunks = torch.empty_like(qkv)
+        _all_to_all_single(output_chunks, qkv, group=group)
+
+        return rearrange(output_chunks, "G M h B (qkv d) -> qkv B (G M) h d", qkv=3)
+
+
+def all_to_all_collect_tokens(x: torch.Tensor, num_heads: int) -> torch.Tensor:
+    if not _CONTEXT_PARALLEL_GROUP:
+        # Move QKV dimension to the front.
+        #   B M (3 H d) -> 3 B M H d
+        B, M, _ = x.size()
+        x = x.view(B, M, 3, num_heads, -1)
+        return x.permute(2, 0, 1, 3, 4)
+
+    return CollectTokens.apply(x, _CONTEXT_PARALLEL_GROUP, num_heads)
+
+
+class CollectHeads(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, group: dist.ProcessGroup):
+        """Redistribute tokens and receive heads.
+
+        Args:
+            x: Output of attention. Shape: [B, N, local_heads, head_dim]
+
+        Returns:
+            Shape: [B, M, num_heads * head_dim]
+        """
+        ctx.group = group
+        ctx.local_heads = x.size(2)
+        ctx.head_dim = x.size(3)
+        group_size = dist.get_world_size(group)
+        x = rearrange(x, "B (G M) h D -> G h M B D", G=group_size).contiguous()
+        output = torch.empty_like(x)
+        _all_to_all_single(output, x, group=group)
+        del x
+        return rearrange(output, "G h M B D -> B M (G h D)")
+
+
+def all_to_all_collect_heads(x: torch.Tensor) -> torch.Tensor:
+    if not _CONTEXT_PARALLEL_GROUP:
+        # Merge heads.
+        return x.view(x.size(0), x.size(1), x.size(2) * x.size(3))
+
+    return CollectHeads.apply(x, _CONTEXT_PARALLEL_GROUP)

src/genmo/mochi_preview/dit/joint_model/layers.py

@@ -0,0 +1,176 @@
+import collections.abc
+import math
+from itertools import repeat
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+to_2tuple = _ntuple(2)
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        frequency_embedding_size: int = 256,
+        *,
+        bias: bool = True,
+        timestep_scale: Optional[float] = None,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=bias, device=device),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=bias, device=device),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+        self.timestep_scale = timestep_scale
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+        freqs.mul_(-math.log(max_period) / half).exp_()
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        if self.timestep_scale is not None:
+            t = t * self.timestep_scale
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class PooledCaptionEmbedder(nn.Module):
+    def __init__(
+        self,
+        caption_feature_dim: int,
+        hidden_size: int,
+        *,
+        bias: bool = True,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        self.caption_feature_dim = caption_feature_dim
+        self.hidden_size = hidden_size
+        self.mlp = nn.Sequential(
+            nn.Linear(caption_feature_dim, hidden_size, bias=bias, device=device),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=bias, device=device),
+        )
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_size: int,
+        multiple_of: int,
+        ffn_dim_multiplier: Optional[float],
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        # keep parameter count and computation constant compared to standard FFN
+        hidden_size = int(2 * hidden_size / 3)
+        # custom dim factor multiplier
+        if ffn_dim_multiplier is not None:
+            hidden_size = int(ffn_dim_multiplier * hidden_size)
+        hidden_size = multiple_of * ((hidden_size + multiple_of - 1) // multiple_of)
+
+        self.hidden_dim = hidden_size
+        self.w1 = nn.Linear(in_features, 2 * hidden_size, bias=False, device=device)
+        self.w2 = nn.Linear(hidden_size, in_features, bias=False, device=device)
+
+    def forward(self, x):
+        x, gate = self.w1(x).chunk(2, dim=-1)
+        x = self.w2(F.silu(x) * gate)
+        return x
+
+
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten: bool = True,
+        bias: bool = True,
+        dynamic_img_pad: bool = False,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        self.patch_size = to_2tuple(patch_size)
+        self.flatten = flatten
+        self.dynamic_img_pad = dynamic_img_pad
+
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=bias,
+            device=device,
+        )
+        assert norm_layer is None
+        self.norm = norm_layer(embed_dim, device=device) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, _C, T, H, W = x.shape
+        if not self.dynamic_img_pad:
+            assert (
+                H % self.patch_size[0] == 0
+            ), f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]})."
+            assert (
+                W % self.patch_size[1] == 0
+            ), f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]})."
+        else:
+            pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0]
+            pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1]
+            x = F.pad(x, (0, pad_w, 0, pad_h))
+
+        x = rearrange(x, "B C T H W -> (B T) C H W", B=B, T=T)
+        x = self.proj(x)
+
+        # Flatten temporal and spatial dimensions.
+        if not self.flatten:
+            raise NotImplementedError("Must flatten output.")
+        x = rearrange(x, "(B T) C H W -> B (T H W) C", B=B, T=T)
+
+        x = self.norm(x)
+        return x
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, hidden_size, eps=1e-5, device=None):
+        super().__init__()
+        self.eps = eps
+        self.weight = torch.nn.Parameter(torch.empty(hidden_size, device=device))
+        self.register_parameter("bias", None)
+
+    def forward(self, x):
+        x_fp32 = x.float()
+        x_normed = x_fp32 * torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x_normed * self.weight).type_as(x)

src/genmo/mochi_preview/dit/joint_model/mod_rmsnorm.py

@@ -0,0 +1,23 @@
+import torch
+
+
+class ModulatedRMSNorm(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, scale, eps=1e-6):
+        # Convert to fp32 for precision
+        x_fp32 = x.float()
+        scale_fp32 = scale.float()
+
+        # Compute RMS
+        mean_square = x_fp32.pow(2).mean(-1, keepdim=True)
+        inv_rms = torch.rsqrt(mean_square + eps)
+
+        # Normalize and modulate
+        x_normed = x_fp32 * inv_rms
+        x_modulated = x_normed * (1 + scale_fp32.unsqueeze(1))
+
+        return x_modulated.type_as(x)
+
+
+def modulated_rmsnorm(x, scale, eps=1e-6):
+    return ModulatedRMSNorm.apply(x, scale, eps)

src/genmo/mochi_preview/dit/joint_model/residual_tanh_gated_rmsnorm.py

@@ -0,0 +1,27 @@
+import torch
+
+
+class ResidualTanhGatedRMSNorm(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, x_res, gate, eps=1e-6):
+        # Convert to fp32 for precision
+        x_res_fp32 = x_res.float()
+
+        # Compute RMS
+        mean_square = x_res_fp32.pow(2).mean(-1, keepdim=True)
+        scale = torch.rsqrt(mean_square + eps)
+
+        # Apply tanh to gate
+        tanh_gate = torch.tanh(gate).unsqueeze(1)
+
+        # Normalize and apply gated scaling
+        x_normed = x_res_fp32 * scale * tanh_gate
+
+        # Apply residual connection
+        output = x + x_normed.type_as(x)
+
+        return output
+
+
+def residual_tanh_gated_rmsnorm(x, x_res, gate, eps=1e-6):
+    return ResidualTanhGatedRMSNorm.apply(x, x_res, gate, eps)

src/genmo/mochi_preview/dit/joint_model/rope_mixed.py

@@ -0,0 +1,88 @@
+import functools
+import math
+
+import torch
+
+
+def centers(start: float, stop, num, dtype=None, device=None):
+    """linspace through bin centers.
+
+    Args:
+        start (float): Start of the range.
+        stop (float): End of the range.
+        num (int): Number of points.
+        dtype (torch.dtype): Data type of the points.
+        device (torch.device): Device of the points.
+
+    Returns:
+        centers (Tensor): Centers of the bins. Shape: (num,).
+    """
+    edges = torch.linspace(start, stop, num + 1, dtype=dtype, device=device)
+    return (edges[:-1] + edges[1:]) / 2
+
+
+@functools.lru_cache(maxsize=1)
+def create_position_matrix(
+    T: int,
+    pH: int,
+    pW: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    *,
+    target_area: float = 36864,
+):
+    """
+    Args:
+        T: int - Temporal dimension
+        pH: int - Height dimension after patchify
+        pW: int - Width dimension after patchify
+
+    Returns:
+        pos: [T * pH * pW, 3] - position matrix
+    """
+    with torch.no_grad():
+        # Create 1D tensors for each dimension
+        t = torch.arange(T, dtype=dtype)
+
+        # Positionally interpolate to area 36864.
+        # (3072x3072 frame with 16x16 patches = 192x192 latents).
+        # This automatically scales rope positions when the resolution changes.
+        # We use a large target area so the model is more sensitive
+        # to changes in the learned pos_frequencies matrix.
+        scale = math.sqrt(target_area / (pW * pH))
+        w = centers(-pW * scale / 2, pW * scale / 2, pW)
+        h = centers(-pH * scale / 2, pH * scale / 2, pH)
+
+        # Use meshgrid to create 3D grids
+        grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
+
+        # Stack and reshape the grids.
+        pos = torch.stack([grid_t, grid_h, grid_w], dim=-1)  # [T, pH, pW, 3]
+        pos = pos.view(-1, 3)  # [T * pH * pW, 3]
+        pos = pos.to(dtype=dtype, device=device)
+
+    return pos
+
+
+def compute_mixed_rotation(
+    freqs: torch.Tensor,
+    pos: torch.Tensor,
+):
+    """
+    Project each 3-dim position into per-head, per-head-dim 1D frequencies.
+
+    Args:
+        freqs: [3, num_heads, num_freqs] - learned rotation frequency (for t, row, col) for each head position
+        pos: [N, 3] - position of each token
+        num_heads: int
+
+    Returns:
+        freqs_cos: [N, num_heads, num_freqs] - cosine components
+        freqs_sin: [N, num_heads, num_freqs] - sine components
+    """
+    with torch.autocast("cuda", enabled=False):
+        assert freqs.ndim == 3
+        freqs_sum = torch.einsum("Nd,dhf->Nhf", pos.to(freqs), freqs)
+        freqs_cos = torch.cos(freqs_sum)
+        freqs_sin = torch.sin(freqs_sum)
+    return freqs_cos, freqs_sin

src/genmo/mochi_preview/dit/joint_model/temporal_rope.py

@@ -0,0 +1,34 @@
+# Based on Llama3 Implementation.
+import torch
+
+
+def apply_rotary_emb_qk_real(
+    xqk: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor without complex numbers.
+
+    Args:
+        xqk (torch.Tensor): Query and/or Key tensors to apply rotary embeddings. Shape: (B, S, *, num_heads, D)
+                            Can be either just query or just key, or both stacked along some batch or * dim.
+        freqs_cos (torch.Tensor): Precomputed cosine frequency tensor.
+        freqs_sin (torch.Tensor): Precomputed sine frequency tensor.
+
+    Returns:
+        torch.Tensor: The input tensor with rotary embeddings applied.
+    """
+    assert xqk.dtype == torch.bfloat16
+    # Split the last dimension into even and odd parts
+    xqk_even = xqk[..., 0::2]
+    xqk_odd = xqk[..., 1::2]
+
+    # Apply rotation
+    cos_part = (xqk_even * freqs_cos - xqk_odd * freqs_sin).type_as(xqk)
+    sin_part = (xqk_even * freqs_sin + xqk_odd * freqs_cos).type_as(xqk)
+
+    # Interleave the results back into the original shape
+    out = torch.stack([cos_part, sin_part], dim=-1).flatten(-2)
+    assert out.dtype == torch.bfloat16
+    return out

src/genmo/mochi_preview/dit/joint_model/utils.py

@@ -0,0 +1,185 @@
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def pool_tokens(x: torch.Tensor, mask: torch.Tensor, *, keepdim=False) -> torch.Tensor:
+    """
+    Pool tokens in x using mask.
+
+    NOTE: We assume x does not require gradients.
+
+    Args:
+        x: (B, L, D) tensor of tokens.
+        mask: (B, L) boolean tensor indicating which tokens are not padding.
+
+    Returns:
+        pooled: (B, D) tensor of pooled tokens.
+    """
+    assert x.size(1) == mask.size(1)  # Expected mask to have same length as tokens.
+    assert x.size(0) == mask.size(0)  # Expected mask to have same batch size as tokens.
+    mask = mask[:, :, None].to(dtype=x.dtype)
+    mask = mask / mask.sum(dim=1, keepdim=True).clamp(min=1)
+    pooled = (x * mask).sum(dim=1, keepdim=keepdim)
+    return pooled
+
+
+class AttentionPool(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        output_dim: int = None,
+        device: Optional[torch.device] = None,
+    ):
+        """
+        Args:
+            spatial_dim (int): Number of tokens in sequence length.
+            embed_dim (int): Dimensionality of input tokens.
+            num_heads (int): Number of attention heads.
+            output_dim (int): Dimensionality of output tokens. Defaults to embed_dim.
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        self.to_kv = nn.Linear(embed_dim, 2 * embed_dim, device=device)
+        self.to_q = nn.Linear(embed_dim, embed_dim, device=device)
+        self.to_out = nn.Linear(embed_dim, output_dim or embed_dim, device=device)
+
+    def forward(self, x, mask):
+        """
+        Args:
+            x (torch.Tensor): (B, L, D) tensor of input tokens.
+            mask (torch.Tensor): (B, L) boolean tensor indicating which tokens are not padding.
+
+        NOTE: We assume x does not require gradients.
+
+        Returns:
+            x (torch.Tensor): (B, D) tensor of pooled tokens.
+        """
+        D = x.size(2)
+
+        # Construct attention mask, shape: (B, 1, num_queries=1, num_keys=1+L).
+        attn_mask = mask[:, None, None, :].bool()  # (B, 1, 1, L).
+        attn_mask = F.pad(attn_mask, (1, 0), value=True)  # (B, 1, 1, 1+L).
+
+        # Average non-padding token features. These will be used as the query.
+        x_pool = pool_tokens(x, mask, keepdim=True)  # (B, 1, D)
+
+        # Concat pooled features to input sequence.
+        x = torch.cat([x_pool, x], dim=1)  # (B, L+1, D)
+
+        # Compute queries, keys, values. Only the mean token is used to create a query.
+        kv = self.to_kv(x)  # (B, L+1, 2 * D)
+        q = self.to_q(x[:, 0])  # (B, D)
+
+        # Extract heads.
+        head_dim = D // self.num_heads
+        kv = kv.unflatten(2, (2, self.num_heads, head_dim))  # (B, 1+L, 2, H, head_dim)
+        kv = kv.transpose(1, 3)  # (B, H, 2, 1+L, head_dim)
+        k, v = kv.unbind(2)  # (B, H, 1+L, head_dim)
+        q = q.unflatten(1, (self.num_heads, head_dim))  # (B, H, head_dim)
+        q = q.unsqueeze(2)  # (B, H, 1, head_dim)
+
+        # Compute attention.
+        x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0)  # (B, H, 1, head_dim)
+
+        # Concatenate heads and run output.
+        x = x.squeeze(2).flatten(1, 2)  # (B, D = H * head_dim)
+        x = self.to_out(x)
+        return x
+
+
+class PadSplitXY(torch.autograd.Function):
+    """
+    Merge heads, pad and extract visual and text tokens,
+    and split along the sequence length.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        xy: torch.Tensor,
+        indices: torch.Tensor,
+        B: int,
+        N: int,
+        L: int,
+        dtype: torch.dtype,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            xy: Packed tokens. Shape: (total <= B * (N + L), num_heads * head_dim).
+            indices: Valid token indices out of unpacked tensor. Shape: (total,)
+
+        Returns:
+            x: Visual tokens. Shape: (B, N, num_heads * head_dim).
+            y: Text tokens. Shape: (B, L, num_heads * head_dim).
+        """
+        ctx.save_for_backward(indices)
+        ctx.B, ctx.N, ctx.L = B, N, L
+        D = xy.size(1)
+
+        # Pad sequences to (B, N + L, dim).
+        assert indices.ndim == 1
+        output = torch.zeros(B * (N + L), D, device=xy.device, dtype=dtype)
+        indices = indices.unsqueeze(1).expand(-1, D)  # (total,) -> (total, num_heads * head_dim)
+        output.scatter_(0, indices, xy)
+        xy = output.view(B, N + L, D)
+
+        # Split visual and text tokens along the sequence length.
+        return torch.tensor_split(xy, (N,), dim=1)
+
+
+def pad_and_split_xy(xy, indices, B, N, L, dtype) -> Tuple[torch.Tensor, torch.Tensor]:
+    return PadSplitXY.apply(xy, indices, B, N, L, dtype)
+
+
+class UnifyStreams(torch.autograd.Function):
+    """Unify visual and text streams."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        q_x: torch.Tensor,
+        k_x: torch.Tensor,
+        v_x: torch.Tensor,
+        q_y: torch.Tensor,
+        k_y: torch.Tensor,
+        v_y: torch.Tensor,
+        indices: torch.Tensor,
+    ):
+        """
+        Args:
+            q_x: (B, N, num_heads, head_dim)
+            k_x: (B, N, num_heads, head_dim)
+            v_x: (B, N, num_heads, head_dim)
+            q_y: (B, L, num_heads, head_dim)
+            k_y: (B, L, num_heads, head_dim)
+            v_y: (B, L, num_heads, head_dim)
+            indices: (total <= B * (N + L))
+
+        Returns:
+            qkv: (total <= B * (N + L), 3, num_heads, head_dim)
+        """
+        ctx.save_for_backward(indices)
+        B, N, num_heads, head_dim = q_x.size()
+        ctx.B, ctx.N, ctx.L = B, N, q_y.size(1)
+        D = num_heads * head_dim
+
+        q = torch.cat([q_x, q_y], dim=1)
+        k = torch.cat([k_x, k_y], dim=1)
+        v = torch.cat([v_x, v_y], dim=1)
+        qkv = torch.stack([q, k, v], dim=2).view(B * (N + ctx.L), 3, D)
+
+        indices = indices[:, None, None].expand(-1, 3, D)
+        qkv = torch.gather(qkv, 0, indices)  # (total, 3, num_heads * head_dim)
+        return qkv.unflatten(2, (num_heads, head_dim))
+
+
+def unify_streams(q_x, k_x, v_x, q_y, k_y, v_y, indices) -> torch.Tensor:
+    return UnifyStreams.apply(q_x, k_x, v_x, q_y, k_y, v_y, indices)

src/genmo/mochi_preview/pipelines.py

@@ -0,0 +1,658 @@
+import os
+import random
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from functools import partial
+from typing import Any, Dict, List, Literal, Optional, Union, cast
+
+import numpy as np
+import ray
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from safetensors.torch import load_file
+from torch import nn
+from torch.distributed.fsdp import (
+    BackwardPrefetch,
+    MixedPrecision,
+    ShardingStrategy,
+)
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+)
+from torch.distributed.fsdp.wrap import (
+    lambda_auto_wrap_policy,
+    transformer_auto_wrap_policy,
+)
+from transformers import T5EncoderModel, T5Tokenizer
+from transformers.models.t5.modeling_t5 import T5Block
+
+import genmo.mochi_preview.dit.joint_model.context_parallel as cp
+import genmo.mochi_preview.vae.cp_conv as cp_conv
+from genmo.mochi_preview.vae.model import Decoder, apply_tiled
+from genmo.lib.progress import get_new_progress_bar, progress_bar
+from genmo.lib.utils import Timer
+
+
+def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
+    if linear_steps is None:
+        linear_steps = num_steps // 2
+    linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
+    threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
+    quadratic_steps = num_steps - linear_steps
+    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
+    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
+    const = quadratic_coef * (linear_steps**2)
+    quadratic_sigma_schedule = [
+        quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps)
+    ]
+    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
+    sigma_schedule = [1.0 - x for x in sigma_schedule]
+    return sigma_schedule
+
+
+T5_MODEL = "google/t5-v1_1-xxl"
+MAX_T5_TOKEN_LENGTH = 256
+
+
+def setup_fsdp_sync(model, device_id, *, param_dtype, auto_wrap_policy) -> FSDP:
+    model = FSDP(
+        model,
+        sharding_strategy=ShardingStrategy.FULL_SHARD,
+        mixed_precision=MixedPrecision(
+            param_dtype=param_dtype,
+            reduce_dtype=torch.float32,
+            buffer_dtype=torch.float32,
+        ),
+        auto_wrap_policy=auto_wrap_policy,
+        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+        limit_all_gathers=True,
+        device_id=device_id,
+        sync_module_states=True,
+        use_orig_params=True,
+    )
+    torch.cuda.synchronize()
+    return model
+
+
+class ModelFactory(ABC):
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+
+    @abstractmethod
+    def get_model(self, *, local_rank: int, device_id: Union[int, Literal["cpu"]], world_size: int) -> Any:
+        if device_id == "cpu":
+            assert world_size == 1, "CPU offload only supports single-GPU inference"
+
+
+class T5ModelFactory(ModelFactory):
+    def __init__(self):
+        super().__init__()
+
+    def get_model(self, *, local_rank, device_id, world_size):
+        super().get_model(local_rank=local_rank, device_id=device_id, world_size=world_size)
+        model = T5EncoderModel.from_pretrained(T5_MODEL)
+        if world_size > 1:
+            model = setup_fsdp_sync(
+                model,
+                device_id=device_id,
+                param_dtype=torch.float32,
+                auto_wrap_policy=partial(
+                    transformer_auto_wrap_policy,
+                    transformer_layer_cls={
+                        T5Block,
+                    },
+                ),
+            )
+        elif isinstance(device_id, int):
+            model = model.to(torch.device(f"cuda:{device_id}"))  # type: ignore
+        return model.eval()
+
+
+class DitModelFactory(ModelFactory):
+    def __init__(self, *, model_path: str, model_dtype: str, attention_mode: Optional[str] = None):
+        if attention_mode is None:
+            from genmo.lib.attn_imports import flash_varlen_qkvpacked_attn # type: ignore
+
+            attention_mode = "sdpa" if flash_varlen_qkvpacked_attn is None else "flash"
+        print(f"Attention mode: {attention_mode}")
+        super().__init__(model_path=model_path, model_dtype=model_dtype, attention_mode=attention_mode)
+
+    def get_model(self, *, local_rank, device_id, world_size):
+        # TODO(ved): Set flag for torch.compile
+        from genmo.mochi_preview.dit.joint_model.asymm_models_joint import (
+            AsymmDiTJoint,
+        )
+
+        model: nn.Module = torch.nn.utils.skip_init(
+            AsymmDiTJoint,
+            depth=48,
+            patch_size=2,
+            num_heads=24,
+            hidden_size_x=3072,
+            hidden_size_y=1536,
+            mlp_ratio_x=4.0,
+            mlp_ratio_y=4.0,
+            in_channels=12,
+            qk_norm=True,
+            qkv_bias=False,
+            out_bias=True,
+            patch_embed_bias=True,
+            timestep_mlp_bias=True,
+            timestep_scale=1000.0,
+            t5_feat_dim=4096,
+            t5_token_length=256,
+            rope_theta=10000.0,
+            attention_mode=self.kwargs["attention_mode"],
+        )
+
+        if local_rank == 0:
+            # FSDP syncs weights from rank 0 to all other ranks
+            model.load_state_dict(load_file(self.kwargs["model_path"]))
+
+        if world_size > 1:
+            assert self.kwargs["model_dtype"] == "bf16", "FP8 is not supported for multi-GPU inference"
+            model = setup_fsdp_sync(
+                model,
+                device_id=device_id,
+                param_dtype=torch.bfloat16,
+                auto_wrap_policy=partial(
+                    lambda_auto_wrap_policy,
+                    lambda_fn=lambda m: m in model.blocks,
+                ),
+            )
+        elif isinstance(device_id, int):
+            model = model.to(torch.device(f"cuda:{device_id}"))
+        return model.eval()
+
+
+class DecoderModelFactory(ModelFactory):
+    def __init__(self, *, model_path: str, model_stats_path: str):
+        super().__init__(model_path=model_path, model_stats_path=model_stats_path)
+
+    def get_model(self, *, local_rank, device_id, world_size):
+        # TODO(ved): Set flag for torch.compile
+        # TODO(ved): Use skip_init
+        import json
+
+        decoder = Decoder(
+            out_channels=3,
+            base_channels=128,
+            channel_multipliers=[1, 2, 4, 6],
+            temporal_expansions=[1, 2, 3],
+            spatial_expansions=[2, 2, 2],
+            num_res_blocks=[3, 3, 4, 6, 3],
+            latent_dim=12,
+            has_attention=[False, False, False, False, False],
+            padding_mode="replicate",
+            output_norm=False,
+            nonlinearity="silu",
+            output_nonlinearity="silu",
+            causal=True,
+        )
+        # VAE is not FSDP-wrapped
+        state_dict = load_file(self.kwargs["model_path"])
+        decoder.load_state_dict(state_dict, strict=True)
+        device = torch.device(f"cuda:{device_id}") if isinstance(device_id, int) else "cpu"
+        decoder.eval().to(device)
+        vae_stats = json.load(open(self.kwargs["model_stats_path"]))
+        decoder.register_buffer("vae_mean", torch.tensor(vae_stats["mean"], device=device))
+        decoder.register_buffer("vae_std", torch.tensor(vae_stats["std"], device=device))
+        return decoder
+
+
+def get_conditioning(tokenizer, encoder, device, batch_inputs, *, prompt: str, negative_prompt: str):
+    if batch_inputs:
+        return dict(batched=get_conditioning_for_prompts(tokenizer, encoder, device, [prompt, negative_prompt]))
+    else:
+        cond_input = get_conditioning_for_prompts(tokenizer, encoder, device, [prompt])
+        null_input = get_conditioning_for_prompts(tokenizer, encoder, device, [negative_prompt])
+        return dict(cond=cond_input, null=null_input)
+
+
+def get_conditioning_for_prompts(tokenizer, encoder, device, prompts: List[str]):
+    assert len(prompts) in [1, 2]  # [neg] or [pos] or [pos, neg]
+    B = len(prompts)
+    t5_toks = tokenizer(
+        prompts,
+        padding="max_length",
+        truncation=True,
+        max_length=MAX_T5_TOKEN_LENGTH,
+        return_tensors="pt",
+        return_attention_mask=True,
+    )
+    caption_input_ids_t5 = t5_toks["input_ids"]
+    caption_attention_mask_t5 = t5_toks["attention_mask"].bool()
+    del t5_toks
+
+    assert caption_input_ids_t5.shape == (B, MAX_T5_TOKEN_LENGTH)
+    assert caption_attention_mask_t5.shape == (B, MAX_T5_TOKEN_LENGTH)
+
+    # Special-case empty negative prompt by zero-ing it
+    if prompts[-1] == "":
+        caption_input_ids_t5[-1] = 0
+        caption_attention_mask_t5[-1] = False
+
+    caption_input_ids_t5 = caption_input_ids_t5.to(device, non_blocking=True)
+    caption_attention_mask_t5 = caption_attention_mask_t5.to(device, non_blocking=True)
+
+    y_mask = [caption_attention_mask_t5]
+    y_feat = [encoder(caption_input_ids_t5, caption_attention_mask_t5).last_hidden_state.detach()]
+    # Sometimes returns a tensor, othertimes a tuple, not sure why
+    # See: https://huggingface.co/genmo/mochi-1-preview/discussions/3
+    assert tuple(y_feat[-1].shape) == (B, MAX_T5_TOKEN_LENGTH, 4096)
+    assert y_feat[-1].dtype == torch.float32
+
+    return dict(y_mask=y_mask, y_feat=y_feat)
+
+
+def compute_packed_indices(
+    device: torch.device, text_mask: torch.Tensor, num_latents: int
+) -> Dict[str, Union[torch.Tensor, int]]:
+    """
+    Based on https://github.com/Dao-AILab/flash-attention/blob/765741c1eeb86c96ee71a3291ad6968cfbf4e4a1/flash_attn/bert_padding.py#L60-L80
+
+    Args:
+        num_latents: Number of latent tokens
+        text_mask: (B, L) List of boolean tensor indicating which text tokens are not padding.
+
+    Returns:
+        packed_indices: Dict with keys for Flash Attention:
+            - valid_token_indices_kv: up to (B * (N + L),) tensor of valid token indices (non-padding)
+                                   in the packed sequence.
+            - cu_seqlens_kv: (B + 1,) tensor of cumulative sequence lengths in the packed sequence.
+            - max_seqlen_in_batch_kv: int of the maximum sequence length in the batch.
+    """
+    # Create an expanded token mask saying which tokens are valid across both visual and text tokens.
+    PATCH_SIZE = 2
+    num_visual_tokens = num_latents // (PATCH_SIZE**2)
+    assert num_visual_tokens > 0
+
+    mask = F.pad(text_mask, (num_visual_tokens, 0), value=True)  # (B, N + L)
+    seqlens_in_batch = mask.sum(dim=-1, dtype=torch.int32)  # (B,)
+    valid_token_indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten()  # up to (B * (N + L),)
+    assert valid_token_indices.size(0) >= text_mask.size(0) * num_visual_tokens  # At least (B * N,)
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+
+    return {
+        "cu_seqlens_kv": cu_seqlens.to(device, non_blocking=True),
+        "max_seqlen_in_batch_kv": cast(int, max_seqlen_in_batch),
+        "valid_token_indices_kv": valid_token_indices.to(device, non_blocking=True),
+    }
+
+
+def assert_eq(x, y, msg=None):
+    assert x == y, f"{msg or 'Assertion failed'}: {x} != {y}"
+
+
+def sample_model(device, dit, conditioning, **args):
+    random.seed(args["seed"])
+    np.random.seed(args["seed"])
+    torch.manual_seed(args["seed"])
+
+    generator = torch.Generator(device=device)
+    generator.manual_seed(args["seed"])
+
+    w, h, t = args["width"], args["height"], args["num_frames"]
+    sample_steps = args["num_inference_steps"]
+    cfg_schedule = args["cfg_schedule"]
+    sigma_schedule = args["sigma_schedule"]
+
+    assert_eq(len(cfg_schedule), sample_steps, "cfg_schedule must have length sample_steps")
+    assert_eq((t - 1) % 6, 0, "t - 1 must be divisible by 6")
+    assert_eq(
+        len(sigma_schedule),
+        sample_steps + 1,
+        "sigma_schedule must have length sample_steps + 1",
+    )
+
+    B = 1
+    SPATIAL_DOWNSAMPLE = 8
+    TEMPORAL_DOWNSAMPLE = 6
+    IN_CHANNELS = 12
+    latent_t = ((t - 1) // TEMPORAL_DOWNSAMPLE) + 1
+    latent_w, latent_h = w // SPATIAL_DOWNSAMPLE, h // SPATIAL_DOWNSAMPLE
+
+    z = torch.randn(
+        (B, IN_CHANNELS, latent_t, latent_h, latent_w),
+        device=device,
+        dtype=torch.float32,
+    )
+
+    num_latents = latent_t * latent_h * latent_w
+    cond_batched = cond_text = cond_null = None
+    if "cond" in conditioning:
+        cond_text = conditioning["cond"]
+        cond_null = conditioning["null"]
+        cond_text["packed_indices"] = compute_packed_indices(device, cond_text["y_mask"][0], num_latents)
+        cond_null["packed_indices"] = compute_packed_indices(device, cond_null["y_mask"][0], num_latents)
+    else:
+        cond_batched = conditioning["batched"]
+        cond_batched["packed_indices"] = compute_packed_indices(device, cond_batched["y_mask"][0], num_latents)
+        z = repeat(z, "b ... -> (repeat b) ...", repeat=2)
+
+    def model_fn(*, z, sigma, cfg_scale):
+        if cond_batched:
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                out = dit(z, sigma, **cond_batched)
+            out_cond, out_uncond = torch.chunk(out, chunks=2, dim=0)
+        else:
+            nonlocal cond_text, cond_null
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                out_cond = dit(z, sigma, **cond_text)
+                out_uncond = dit(z, sigma, **cond_null)
+        assert out_cond.shape == out_uncond.shape
+        return out_uncond + cfg_scale * (out_cond - out_uncond), out_cond
+
+    for i in get_new_progress_bar(range(0, sample_steps), desc="Sampling"):
+        sigma = sigma_schedule[i]
+        dsigma = sigma - sigma_schedule[i + 1]
+
+        # `pred` estimates `z_0 - eps`.
+        pred, output_cond = model_fn(
+            z=z,
+            sigma=torch.full([B] if cond_text else [B * 2], sigma, device=z.device),
+            cfg_scale=cfg_schedule[i],
+        )
+        pred = pred.to(z)
+        output_cond = output_cond.to(z)
+        z = z + dsigma * pred
+
+    return z[:B] if cond_batched else z
+
+
+def decoded_latents_to_frames(samples):
+    samples = samples.float()
+    samples = (samples + 1.0) / 2.0
+    samples.clamp_(0.0, 1.0)
+    frames = rearrange(samples, "b c t h w -> b t h w c")
+    return frames
+
+
+def decode_latents(decoder, z):
+    cp_rank, cp_size = cp.get_cp_rank_size()
+    z = z.tensor_split(cp_size, dim=2)[cp_rank]  # split along temporal dim
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        samples = decoder(z)
+    samples = cp_conv.gather_all_frames(samples)
+    return decoded_latents_to_frames(samples)
+
+
+@torch.inference_mode()
+def decode_latents_tiled_full(
+    decoder,
+    z,
+    *,
+    tile_sample_min_height: int = 240,
+    tile_sample_min_width: int = 424,
+    tile_overlap_factor_height: float = 0.1666,
+    tile_overlap_factor_width: float = 0.2,
+    auto_tile_size: bool = True,
+    frame_batch_size: int = 6,
+):
+    B, C, T, H, W = z.shape
+    assert frame_batch_size <= T, f"frame_batch_size must be <= T, got {frame_batch_size} > {T}"
+
+    tile_sample_min_height = tile_sample_min_height if not auto_tile_size else H // 2 * 8
+    tile_sample_min_width = tile_sample_min_width if not auto_tile_size else W // 2 * 8
+
+    tile_latent_min_height = int(tile_sample_min_height / 8)
+    tile_latent_min_width = int(tile_sample_min_width / 8)
+
+    def blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+                y / blend_extent
+            )
+        return b
+
+    def blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+        blend_extent = min(a.shape[4], b.shape[4], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+                x / blend_extent
+            )
+        return b
+
+    overlap_height = int(tile_latent_min_height * (1 - tile_overlap_factor_height))
+    overlap_width = int(tile_latent_min_width * (1 - tile_overlap_factor_width))
+    blend_extent_height = int(tile_sample_min_height * tile_overlap_factor_height)
+    blend_extent_width = int(tile_sample_min_width * tile_overlap_factor_width)
+    row_limit_height = tile_sample_min_height - blend_extent_height
+    row_limit_width = tile_sample_min_width - blend_extent_width
+
+    # Split z into overlapping tiles and decode them separately.
+    # The tiles have an overlap to avoid seams between tiles.
+    pbar = get_new_progress_bar(
+        desc="Decoding latent tiles",
+        total=len(range(0, H, overlap_height)) * len(range(0, W, overlap_width)) * len(range(T // frame_batch_size)),
+    )
+    rows = []
+    for i in range(0, H, overlap_height):
+        row = []
+        for j in range(0, W, overlap_width):
+            temporal = []
+            for k in range(T // frame_batch_size):
+                remaining_frames = T % frame_batch_size
+                start_frame = frame_batch_size * k + (0 if k == 0 else remaining_frames)
+                end_frame = frame_batch_size * (k + 1) + remaining_frames
+                tile = z[
+                    :,
+                    :,
+                    start_frame:end_frame,
+                    i : i + tile_latent_min_height,
+                    j : j + tile_latent_min_width,
+                ]
+                tile = decoder(tile)
+                temporal.append(tile)
+                pbar.update(1)
+            row.append(torch.cat(temporal, dim=2))
+        rows.append(row)
+
+    result_rows = []
+    for i, row in enumerate(rows):
+        result_row = []
+        for j, tile in enumerate(row):
+            # blend the above tile and the left tile
+            # to the current tile and add the current tile to the result row
+            if i > 0:
+                tile = blend_v(rows[i - 1][j], tile, blend_extent_height)
+            if j > 0:
+                tile = blend_h(row[j - 1], tile, blend_extent_width)
+            result_row.append(tile[:, :, :, :row_limit_height, :row_limit_width])
+        result_rows.append(torch.cat(result_row, dim=4))
+
+    return decoded_latents_to_frames(torch.cat(result_rows, dim=3))
+
+@torch.inference_mode()
+def decode_latents_tiled_spatial(
+    decoder,
+    z,
+    *,
+    num_tiles_w: int,
+    num_tiles_h: int,
+    overlap: int = 0,  # Number of pixel of overlap between adjacent tiles.
+    # Use a factor of 2 times the latent downsample factor.
+    min_block_size: int = 1,  # Minimum number of pixels in each dimension when subdividing.
+):
+    decoded = apply_tiled(decoder, z, num_tiles_w, num_tiles_h, overlap, min_block_size)
+    assert decoded is not None, f"Failed to decode latents with tiled spatial method"
+    return decoded
+
+@contextmanager
+def move_to_device(model: nn.Module, target_device):
+    og_device = next(model.parameters()).device
+    if og_device == target_device:
+        print(f"move_to_device is a no-op model is already on {target_device}")
+    else:
+        print(f"moving model from {og_device} -> {target_device}")
+
+    model.to(target_device)
+    yield
+    if og_device != target_device:
+        print(f"moving model from {target_device} -> {og_device}")
+    model.to(og_device)
+
+
+def t5_tokenizer():
+    return T5Tokenizer.from_pretrained(T5_MODEL, legacy=False)
+
+
+class MochiSingleGPUPipeline:
+    def __init__(
+        self,
+        *,
+        text_encoder_factory: ModelFactory,
+        dit_factory: ModelFactory,
+        decoder_factory: ModelFactory,
+        cpu_offload: Optional[bool] = False,
+        decode_type: str = "full",
+        decode_args: Optional[Dict[str, Any]] = None,
+    ):
+        self.device = torch.device("cuda:0")
+        self.tokenizer = t5_tokenizer()
+        t = Timer()
+        self.cpu_offload = cpu_offload
+        self.decode_args = decode_args or {}
+        self.decode_type = decode_type
+        init_id = "cpu" if cpu_offload else 0
+        with t("load_text_encoder"):
+            self.text_encoder = text_encoder_factory.get_model(
+                local_rank=0,
+                device_id=init_id,
+                world_size=1,
+            )
+        with t("load_dit"):
+            self.dit = dit_factory.get_model(local_rank=0, device_id=init_id, world_size=1)
+        with t("load_vae"):
+            self.decoder = decoder_factory.get_model(local_rank=0, device_id=init_id, world_size=1)
+        t.print_stats()
+
+    def __call__(self, batch_cfg, prompt, negative_prompt, **kwargs):
+        with progress_bar(type="tqdm"), torch.inference_mode():
+            print_max_memory = lambda: print(
+                f"Max memory reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB"
+            )
+            print_max_memory()
+            with move_to_device(self.text_encoder, self.device):
+                conditioning = get_conditioning(
+                    self.tokenizer,
+                    self.text_encoder,
+                    self.device,
+                    batch_cfg,
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                )
+            print_max_memory()
+            with move_to_device(self.dit, self.device):
+                latents = sample_model(self.device, self.dit, conditioning, **kwargs)
+            print_max_memory()
+            with move_to_device(self.decoder, self.device):
+                frames = (
+                    decode_latents_tiled_full(self.decoder, latents, **self.decode_args)
+                    if self.decode_type == "tiled_full"
+                    else
+                    decode_latents_tiled_spatial(self.decoder, latents, **self.decode_args)
+                    if self.decode_type == "tiled_spatial"
+                    else decode_latents(self.decoder, latents)
+                )
+            print_max_memory()
+            return frames.cpu().numpy()
+
+
+### ALL CODE BELOW HERE IS FOR MULTI-GPU MODE ###
+
+
+# In multi-gpu mode, all models must belong to a device which has a predefined context parallel group
+# So it doesn't make sense to work with models individually
+class MultiGPUContext:
+    def __init__(
+        self,
+        *,
+        text_encoder_factory,
+        dit_factory,
+        decoder_factory,
+        device_id,
+        local_rank,
+        world_size,
+    ):
+        t = Timer()
+        self.device = torch.device(f"cuda:{device_id}")
+        print(f"Initializing rank {local_rank+1}/{world_size}")
+        assert world_size > 1, f"Multi-GPU mode requires world_size > 1, got {world_size}"
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+        os.environ["MASTER_PORT"] = "29500"
+        with t("init_process_group"):
+            dist.init_process_group(
+                "nccl",
+                rank=local_rank,
+                world_size=world_size,
+                device_id=self.device,  # force non-lazy init
+            )
+        pg = dist.group.WORLD
+        cp.set_cp_group(pg, list(range(world_size)), local_rank)
+        distributed_kwargs = dict(local_rank=local_rank, device_id=device_id, world_size=world_size)
+        self.world_size = world_size
+        self.tokenizer = t5_tokenizer()
+        with t("load_text_encoder"):
+            self.text_encoder = text_encoder_factory.get_model(**distributed_kwargs)
+        with t("load_dit"):
+            self.dit = dit_factory.get_model(**distributed_kwargs)
+        with t("load_vae"):
+            self.decoder = decoder_factory.get_model(**distributed_kwargs)
+        self.local_rank = local_rank
+        t.print_stats()
+
+    def run(self, *, fn, **kwargs):
+        return fn(self, **kwargs)
+
+
+class MochiMultiGPUPipeline:
+    def __init__(
+        self,
+        *,
+        text_encoder_factory: ModelFactory,
+        dit_factory: ModelFactory,
+        decoder_factory: ModelFactory,
+        world_size: int,
+    ):
+        ray.init()
+        RemoteClass = ray.remote(MultiGPUContext)
+        self.ctxs = [
+            RemoteClass.options(num_gpus=1).remote(
+                text_encoder_factory=text_encoder_factory,
+                dit_factory=dit_factory,
+                decoder_factory=decoder_factory,
+                world_size=world_size,
+                device_id=0,
+                local_rank=i,
+            )
+            for i in range(world_size)
+        ]
+        for ctx in self.ctxs:
+            ray.get(ctx.__ray_ready__.remote())
+
+    def __call__(self, **kwargs):
+        def sample(ctx, *, batch_cfg, prompt, negative_prompt, **kwargs):
+            with progress_bar(type="ray_tqdm", enabled=ctx.local_rank == 0), torch.inference_mode():
+                conditioning = get_conditioning(
+                    ctx.tokenizer,
+                    ctx.text_encoder,
+                    ctx.device,
+                    batch_cfg,
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                )
+                latents = sample_model(ctx.device, ctx.dit, conditioning=conditioning, **kwargs)
+                if ctx.local_rank == 0:
+                    torch.save(latents, "latents.pt")
+                frames = decode_latents(ctx.decoder, latents)
+            return frames.cpu().numpy()
+
+        return ray.get([ctx.run.remote(fn=sample, **kwargs, show_progress=i == 0) for i, ctx in enumerate(self.ctxs)])[
+            0
+        ]

src/genmo/mochi_preview/vae/cp_conv.py

@@ -0,0 +1,152 @@
+from typing import Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+import genmo.mochi_preview.dit.joint_model.context_parallel as cp
+
+
+def cast_tuple(t, length=1):
+    return t if isinstance(t, tuple) else ((t,) * length)
+
+
+def cp_pass_frames(x: torch.Tensor, frames_to_send: int) -> torch.Tensor:
+    """
+    Forward pass that handles communication between ranks for inference.
+    Args:
+        x: Tensor of shape (B, C, T, H, W)
+        frames_to_send: int, number of frames to communicate between ranks
+    Returns:
+        output: Tensor of shape (B, C, T', H, W)
+    """
+    cp_rank, cp_world_size = cp.get_cp_rank_size()
+    if frames_to_send == 0 or cp_world_size == 1:
+        return x
+
+    group = cp.get_cp_group()
+    global_rank = dist.get_rank()
+
+    # Send to next rank
+    if cp_rank < cp_world_size - 1:
+        assert x.size(2) >= frames_to_send
+        tail = x[:, :, -frames_to_send:].contiguous()
+        dist.send(tail, global_rank + 1, group=group)
+
+    # Receive from previous rank
+    if cp_rank > 0:
+        B, C, _, H, W = x.shape
+        recv_buffer = torch.empty(
+            (B, C, frames_to_send, H, W),
+            dtype=x.dtype,
+            device=x.device,
+        )
+        dist.recv(recv_buffer, global_rank - 1, group=group)
+        x = torch.cat([recv_buffer, x], dim=2)
+
+    return x
+
+
+def _pad_to_max(x: torch.Tensor, max_T: int) -> torch.Tensor:
+    if max_T > x.size(2):
+        pad_T = max_T - x.size(2)
+        pad_dims = (0, 0, 0, 0, 0, pad_T)
+        return F.pad(x, pad_dims)
+    return x
+
+
+def gather_all_frames(x: torch.Tensor) -> torch.Tensor:
+    """
+    Gathers all frames from all processes for inference.
+    Args:
+        x: Tensor of shape (B, C, T, H, W)
+    Returns:
+        output: Tensor of shape (B, C, T_total, H, W)
+    """
+    cp_rank, cp_size = cp.get_cp_rank_size()
+    cp_group = cp.get_cp_group()
+
+    # Ensure the tensor is contiguous for collective operations
+    x = x.contiguous()
+
+    # Get the local time dimension size
+    local_T = x.size(2)
+    local_T_tensor = torch.tensor([local_T], device=x.device, dtype=torch.int64)
+
+    # Gather all T sizes from all processes
+    all_T = [torch.zeros(1, dtype=torch.int64, device=x.device) for _ in range(cp_size)]
+    dist.all_gather(all_T, local_T_tensor, group=cp_group)
+    all_T = [t.item() for t in all_T]
+
+    # Pad the tensor at the end of the time dimension to match max_T
+    max_T = max(all_T)
+    x = _pad_to_max(x, max_T).contiguous()
+
+    # Prepare a list to hold the gathered tensors
+    gathered_x = [torch.zeros_like(x).contiguous() for _ in range(cp_size)]
+
+    # Perform the all_gather operation
+    dist.all_gather(gathered_x, x, group=cp_group)
+
+    # Slice each gathered tensor back to its original T size
+    for idx, t_size in enumerate(all_T):
+        gathered_x[idx] = gathered_x[idx][:, :, :t_size]
+
+    return torch.cat(gathered_x, dim=2)
+
+
+def excessive_memory_usage(input: torch.Tensor, max_gb: float = 2.0) -> bool:
+    """Estimate memory usage based on input tensor size and data type."""
+    element_size = input.element_size()  # Size in bytes of each element
+    memory_bytes = input.numel() * element_size
+    memory_gb = memory_bytes / 1024**3
+    return memory_gb > max_gb
+
+
+class ContextParallelCausalConv3d(torch.nn.Conv3d):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size: Union[int, Tuple[int, int, int]],
+        stride: Union[int, Tuple[int, int, int]],
+        **kwargs,
+    ):
+        kernel_size = cast_tuple(kernel_size, 3)
+        stride = cast_tuple(stride, 3)
+        height_pad = (kernel_size[1] - 1) // 2
+        width_pad = (kernel_size[2] - 1) // 2
+
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=(1, 1, 1),
+            padding=(0, height_pad, width_pad),
+            **kwargs,
+        )
+
+    def forward(self, x: torch.Tensor):
+        cp_rank, cp_world_size = cp.get_cp_rank_size()
+
+        context_size = self.kernel_size[0] - 1
+        if cp_rank == 0:
+            mode = "constant" if self.padding_mode == "zeros" else self.padding_mode
+            x = F.pad(x, (0, 0, 0, 0, context_size, 0), mode=mode)
+
+        if cp_world_size == 1:
+            return super().forward(x)
+
+        if all(s == 1 for s in self.stride):
+            # Receive some frames from previous rank.
+            x = cp_pass_frames(x, context_size)
+            return super().forward(x)
+
+        # Less efficient implementation for strided convs.
+        # All gather x, infer and chunk.
+        x = gather_all_frames(x)  # [B, C, k - 1 + global_T, H, W]
+        x = super().forward(x)
+        x_chunks = x.tensor_split(cp_world_size, dim=2)
+        assert len(x_chunks) == cp_world_size
+        return x_chunks[cp_rank]

src/genmo/mochi_preview/vae/model.py

@@ -0,0 +1,808 @@
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+import genmo.mochi_preview.dit.joint_model.context_parallel as cp
+from genmo.mochi_preview.vae.cp_conv import cp_pass_frames, gather_all_frames
+
+
+def cast_tuple(t, length=1):
+    return t if isinstance(t, tuple) else ((t,) * length)
+
+
+class GroupNormSpatial(nn.GroupNorm):
+    """
+    GroupNorm applied per-frame.
+    """
+
+    def forward(self, x: torch.Tensor, *, chunk_size: int = 8):
+        B, C, T, H, W = x.shape
+        x = rearrange(x, "B C T H W -> (B T) C H W")
+        # Run group norm in chunks.
+        output = torch.empty_like(x)
+        for b in range(0, B * T, chunk_size):
+            output[b : b + chunk_size] = super().forward(x[b : b + chunk_size])
+        return rearrange(output, "(B T) C H W -> B C T H W", B=B, T=T)
+
+
+class SafeConv3d(torch.nn.Conv3d):
+    """
+    NOTE: No support for padding along time dimension.
+          Input must already be padded along time.
+    """
+
+    def forward(self, input):
+        memory_count = torch.prod(torch.tensor(input.shape)).item() * 2 / 1024**3
+        if memory_count > 2:
+            part_num = int(memory_count / 2) + 1
+            k = self.kernel_size[0]
+            input_idx = torch.arange(k - 1, input.size(2))
+            input_chunks_idx = torch.chunk(input_idx, part_num, dim=0)
+
+            # assert self.kernel_size == (3, 3, 3), f"kernel_size {self.kernel_size} != (3, 3, 3)"
+            assert self.stride[0] == 1, f"stride {self.stride}"
+            assert self.dilation[0] == 1, f"dilation {self.dilation}"
+            assert self.padding[0] == 0, f"padding {self.padding}"
+
+            # Comptue output size
+            assert not input.requires_grad
+            B, _, T_in, H_in, W_in = input.shape
+            output_size = (
+                B,
+                self.out_channels,
+                T_in - k + 1,
+                H_in // self.stride[1],
+                W_in // self.stride[2],
+            )
+            output = torch.empty(output_size, dtype=input.dtype, device=input.device)
+            for input_chunk_idx in input_chunks_idx:
+                input_s = input_chunk_idx[0] - k + 1
+                input_e = input_chunk_idx[-1] + 1
+                input_chunk = input[:, :, input_s:input_e, :, :]
+                output_chunk = super(SafeConv3d, self).forward(input_chunk)
+
+                output_s = input_s
+                output_e = output_s + output_chunk.size(2)
+                output[:, :, output_s:output_e, :, :] = output_chunk
+
+            return output
+        else:
+            return super(SafeConv3d, self).forward(input)
+
+
+class StridedSafeConv3d(torch.nn.Conv3d):
+    def forward(self, input, local_shard: bool = False):
+        assert self.stride[0] == self.kernel_size[0]
+        assert self.dilation[0] == 1
+        assert self.padding[0] == 0
+
+        kernel_size = self.kernel_size[0]
+        stride = self.stride[0]
+        T_in = input.size(2)
+        T_out = T_in // kernel_size
+
+        # Parallel implementation.
+        if local_shard:
+            idx = torch.arange(T_out)
+            idx = cp.local_shard(idx, dim=0)
+            start = idx.min() * stride
+            end = idx.max() * stride + kernel_size
+            local_input = input[:, :, start:end, :, :]
+            return torch.nn.Conv3d.forward(self, local_input)
+
+        raise NotImplementedError
+
+
+class ContextParallelConv3d(SafeConv3d):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size: Union[int, Tuple[int, int, int]],
+        stride: Union[int, Tuple[int, int, int]],
+        causal: bool = True,
+        context_parallel: bool = True,
+        **kwargs,
+    ):
+        self.causal = causal
+        self.context_parallel = context_parallel
+        kernel_size = cast_tuple(kernel_size, 3)
+        stride = cast_tuple(stride, 3)
+        height_pad = (kernel_size[1] - 1) // 2
+        width_pad = (kernel_size[2] - 1) // 2
+
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            dilation=(1, 1, 1),
+            padding=(0, height_pad, width_pad),
+            **kwargs,
+        )
+
+    def forward(self, x: torch.Tensor):
+        cp_rank, cp_world_size = cp.get_cp_rank_size()
+
+        # Compute padding amounts.
+        context_size = self.kernel_size[0] - 1
+        if self.causal:
+            pad_front = context_size
+            pad_back = 0
+        else:
+            pad_front = context_size // 2
+            pad_back = context_size - pad_front
+
+        # Apply padding.
+        assert self.padding_mode == "replicate"  # DEBUG
+        mode = "constant" if self.padding_mode == "zeros" else self.padding_mode
+        if self.context_parallel and cp_world_size == 1:
+            x = F.pad(x, (0, 0, 0, 0, pad_front, pad_back), mode=mode)
+        else:
+            if cp_rank == 0:
+                x = F.pad(x, (0, 0, 0, 0, pad_front, 0), mode=mode)
+            elif cp_rank == cp_world_size - 1 and pad_back:
+                x = F.pad(x, (0, 0, 0, 0, 0, pad_back), mode=mode)
+
+        if self.context_parallel and cp_world_size == 1:
+            return super().forward(x)
+
+        if self.stride[0] == 1:
+            # Receive some frames from previous rank.
+            x = cp_pass_frames(x, context_size)
+            return super().forward(x)
+
+        # Less efficient implementation for strided convs.
+        # All gather x, infer and chunk.
+        assert x.dtype == torch.bfloat16, f"Expected x to be of type torch.bfloat16, got {x.dtype}"
+
+        x = gather_all_frames(x)  # [B, C, k - 1 + global_T, H, W]
+        return StridedSafeConv3d.forward(self, x, local_shard=True)
+
+
+class Conv1x1(nn.Linear):
+    """*1x1 Conv implemented with a linear layer."""
+
+    def __init__(self, in_features: int, out_features: int, *args, **kwargs):
+        super().__init__(in_features, out_features, *args, **kwargs)
+
+    def forward(self, x: torch.Tensor):
+        """Forward pass.
+
+        Args:
+            x: Input tensor. Shape: [B, C, *] or [B, *, C].
+
+        Returns:
+            x: Output tensor. Shape: [B, C', *] or [B, *, C'].
+        """
+        x = x.movedim(1, -1)
+        x = super().forward(x)
+        x = x.movedim(-1, 1)
+        return x
+
+
+class DepthToSpaceTime(nn.Module):
+    def __init__(
+        self,
+        temporal_expansion: int,
+        spatial_expansion: int,
+    ):
+        super().__init__()
+        self.temporal_expansion = temporal_expansion
+        self.spatial_expansion = spatial_expansion
+
+    # When printed, this module should show the temporal and spatial expansion factors.
+    def extra_repr(self):
+        return f"texp={self.temporal_expansion}, sexp={self.spatial_expansion}"
+
+    def forward(self, x: torch.Tensor):
+        """Forward pass.
+
+        Args:
+            x: Input tensor. Shape: [B, C, T, H, W].
+
+        Returns:
+            x: Rearranged tensor. Shape: [B, C/(st*s*s), T*st, H*s, W*s].
+        """
+        x = rearrange(
+            x,
+            "B (C st sh sw) T H W -> B C (T st) (H sh) (W sw)",
+            st=self.temporal_expansion,
+            sh=self.spatial_expansion,
+            sw=self.spatial_expansion,
+        )
+
+        cp_rank, _ = cp.get_cp_rank_size()
+        if self.temporal_expansion > 1 and cp_rank == 0:
+            # Drop the first self.temporal_expansion - 1 frames.
+            # This is because we always want the 3x3x3 conv filter to only apply
+            # to the first frame, and the first frame doesn't need to be repeated.
+            assert all(x.shape)
+            x = x[:, :, self.temporal_expansion - 1 :]
+            assert all(x.shape)
+
+        return x
+
+
+def norm_fn(
+    in_channels: int,
+    affine: bool = True,
+):
+    return GroupNormSpatial(affine=affine, num_groups=32, num_channels=in_channels)
+
+
+class ResBlock(nn.Module):
+    """Residual block that preserves the spatial dimensions."""
+
+    def __init__(
+        self,
+        channels: int,
+        *,
+        affine: bool = True,
+        attn_block: Optional[nn.Module] = None,
+        padding_mode: str = "replicate",
+        causal: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+
+        assert causal
+        self.stack = nn.Sequential(
+            norm_fn(channels, affine=affine),
+            nn.SiLU(inplace=True),
+            ContextParallelConv3d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=(3, 3, 3),
+                stride=(1, 1, 1),
+                padding_mode=padding_mode,
+                bias=True,
+                causal=causal,
+            ),
+            norm_fn(channels, affine=affine),
+            nn.SiLU(inplace=True),
+            ContextParallelConv3d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=(3, 3, 3),
+                stride=(1, 1, 1),
+                padding_mode=padding_mode,
+                bias=True,
+                causal=causal,
+            ),
+        )
+
+        self.attn_block = attn_block if attn_block else nn.Identity()
+
+    def forward(self, x: torch.Tensor):
+        """Forward pass.
+
+        Args:
+            x: Input tensor. Shape: [B, C, T, H, W].
+        """
+        residual = x
+        x = self.stack(x)
+        x = x + residual
+        del residual
+
+        return self.attn_block(x)
+
+
+def prepare_for_attention(qkv: torch.Tensor, head_dim: int, qk_norm: bool = True):
+    """Prepare qkv tensor for attention and normalize qk.
+
+    Args:
+        qkv: Input tensor. Shape: [B, L, 3 * num_heads * head_dim].
+
+    Returns:
+        q, k, v: qkv tensor split into q, k, v. Shape: [B, num_heads, L, head_dim].
+    """
+    assert qkv.ndim == 3  # [B, L, 3 * num_heads * head_dim]
+    assert qkv.size(2) % (3 * head_dim) == 0
+    num_heads = qkv.size(2) // (3 * head_dim)
+    qkv = qkv.unflatten(2, (3, num_heads, head_dim))
+
+    q, k, v = qkv.unbind(2)  # [B, L, num_heads, head_dim]
+    q = q.transpose(1, 2)  # [B, num_heads, L, head_dim]
+    k = k.transpose(1, 2)  # [B, num_heads, L, head_dim]
+    v = v.transpose(1, 2)  # [B, num_heads, L, head_dim]
+
+    if qk_norm:
+        q = F.normalize(q, p=2, dim=-1)
+        k = F.normalize(k, p=2, dim=-1)
+
+        # Mixed precision can change the dtype of normed q/k to float32.
+        q = q.to(dtype=qkv.dtype)
+        k = k.to(dtype=qkv.dtype)
+
+    return q, k, v
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        head_dim: int = 32,
+        qkv_bias: bool = False,
+        out_bias: bool = True,
+        qk_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        self.head_dim = head_dim
+        self.num_heads = dim // head_dim
+        self.qk_norm = qk_norm
+
+        self.qkv = nn.Linear(dim, 3 * dim, bias=qkv_bias)
+        self.out = nn.Linear(dim, dim, bias=out_bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        *,
+        chunk_size=2**15,
+    ) -> torch.Tensor:
+        """Compute temporal self-attention.
+
+        Args:
+            x: Input tensor. Shape: [B, C, T, H, W].
+            chunk_size: Chunk size for large tensors.
+
+        Returns:
+            x: Output tensor. Shape: [B, C, T, H, W].
+        """
+        B, _, T, H, W = x.shape
+
+        if T == 1:
+            # No attention for single frame.
+            x = x.movedim(1, -1)  # [B, C, T, H, W] -> [B, T, H, W, C]
+            qkv = self.qkv(x)
+            _, _, x = qkv.chunk(3, dim=-1)  # Throw away queries and keys.
+            x = self.out(x)
+            return x.movedim(-1, 1)  # [B, T, H, W, C] -> [B, C, T, H, W]
+
+        # 1D temporal attention.
+        x = rearrange(x, "B C t h w -> (B h w) t C")
+        qkv = self.qkv(x)
+
+        # Input: qkv with shape [B, t, 3 * num_heads * head_dim]
+        # Output: x with shape [B, num_heads, t, head_dim]
+        q, k, v = prepare_for_attention(qkv, self.head_dim, qk_norm=self.qk_norm)
+
+        attn_kwargs = dict(
+            attn_mask=None,
+            dropout_p=0.0,
+            is_causal=True,
+            scale=self.head_dim**-0.5,
+        )
+
+        if q.size(0) <= chunk_size:
+            x = F.scaled_dot_product_attention(q, k, v, **attn_kwargs)  # [B, num_heads, t, head_dim]
+        else:
+            # Evaluate in chunks to avoid `RuntimeError: CUDA error: invalid configuration argument.`
+            # Chunks of 2**16 and up cause an error.
+            x = torch.empty_like(q)
+            for i in range(0, q.size(0), chunk_size):
+                qc = q[i : i + chunk_size]
+                kc = k[i : i + chunk_size]
+                vc = v[i : i + chunk_size]
+                chunk = F.scaled_dot_product_attention(qc, kc, vc, **attn_kwargs)
+                x[i : i + chunk_size].copy_(chunk)
+
+        assert x.size(0) == q.size(0)
+        x = x.transpose(1, 2)  # [B, t, num_heads, head_dim]
+        x = x.flatten(2)  # [B, t, num_heads * head_dim]
+
+        x = self.out(x)
+        x = rearrange(x, "(B h w) t C -> B C t h w", B=B, h=H, w=W)
+        return x
+
+
+class AttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        **attn_kwargs,
+    ) -> None:
+        super().__init__()
+        self.norm = norm_fn(dim)
+        self.attn = Attention(dim, **attn_kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.attn(self.norm(x))
+
+
+class CausalUpsampleBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_res_blocks: int,
+        *,
+        temporal_expansion: int = 2,
+        spatial_expansion: int = 2,
+        **block_kwargs,
+    ):
+        super().__init__()
+
+        blocks = []
+        for _ in range(num_res_blocks):
+            blocks.append(block_fn(in_channels, **block_kwargs))
+        self.blocks = nn.Sequential(*blocks)
+
+        self.temporal_expansion = temporal_expansion
+        self.spatial_expansion = spatial_expansion
+
+        # Change channels in the final convolution layer.
+        self.proj = Conv1x1(
+            in_channels,
+            out_channels * temporal_expansion * (spatial_expansion**2),
+        )
+
+        self.d2st = DepthToSpaceTime(temporal_expansion=temporal_expansion, spatial_expansion=spatial_expansion)
+
+    def forward(self, x):
+        x = self.blocks(x)
+        x = self.proj(x)
+        x = self.d2st(x)
+        return x
+
+
+def block_fn(channels, *, has_attention: bool = False, **block_kwargs):
+    attn_block = AttentionBlock(channels) if has_attention else None
+
+    return ResBlock(channels, affine=True, attn_block=attn_block, **block_kwargs)
+
+
+class DownsampleBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_res_blocks,
+        *,
+        temporal_reduction=2,
+        spatial_reduction=2,
+        **block_kwargs,
+    ):
+        """
+        Downsample block for the VAE encoder.
+
+        Args:
+            in_channels: Number of input channels.
+            out_channels: Number of output channels.
+            num_res_blocks: Number of residual blocks.
+            temporal_reduction: Temporal reduction factor.
+            spatial_reduction: Spatial reduction factor.
+        """
+        super().__init__()
+        layers = []
+
+        # Change the channel count in the strided convolution.
+        # This lets the ResBlock have uniform channel count,
+        # as in ConvNeXt.
+        assert in_channels != out_channels
+        layers.append(
+            ContextParallelConv3d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(temporal_reduction, spatial_reduction, spatial_reduction),
+                stride=(temporal_reduction, spatial_reduction, spatial_reduction),
+                padding_mode="replicate",
+                bias=True,
+            )
+        )
+
+        for _ in range(num_res_blocks):
+            layers.append(block_fn(out_channels, **block_kwargs))
+
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+def add_fourier_features(inputs: torch.Tensor, start=6, stop=8, step=1):
+    num_freqs = (stop - start) // step
+    assert inputs.ndim == 5
+    C = inputs.size(1)
+
+    # Create Base 2 Fourier features.
+    freqs = torch.arange(start, stop, step, dtype=inputs.dtype, device=inputs.device)
+    assert num_freqs == len(freqs)
+    w = torch.pow(2.0, freqs) * (2 * torch.pi)  # [num_freqs]
+    C = inputs.shape[1]
+    w = w.repeat(C)[None, :, None, None, None]  # [1, C * num_freqs, 1, 1, 1]
+
+    # Interleaved repeat of input channels to match w.
+    h = inputs.repeat_interleave(num_freqs, dim=1)  # [B, C * num_freqs, T, H, W]
+    # Scale channels by frequency.
+    h = w * h
+
+    return torch.cat(
+        [
+            inputs,
+            torch.sin(h),
+            torch.cos(h),
+        ],
+        dim=1,
+    )
+
+
+class FourierFeatures(nn.Module):
+    def __init__(self, start: int = 6, stop: int = 8, step: int = 1):
+        super().__init__()
+        self.start = start
+        self.stop = stop
+        self.step = step
+
+    def forward(self, inputs):
+        """Add Fourier features to inputs.
+
+        Args:
+            inputs: Input tensor. Shape: [B, C, T, H, W]
+
+        Returns:
+            h: Output tensor. Shape: [B, (1 + 2 * num_freqs) * C, T, H, W]
+        """
+        return add_fourier_features(inputs, self.start, self.stop, self.step)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        out_channels: int = 3,
+        latent_dim: int,
+        base_channels: int,
+        channel_multipliers: List[int],
+        num_res_blocks: List[int],
+        temporal_expansions: Optional[List[int]] = None,
+        spatial_expansions: Optional[List[int]] = None,
+        has_attention: List[bool],
+        output_norm: bool = True,
+        nonlinearity: str = "silu",
+        output_nonlinearity: str = "silu",
+        causal: bool = True,
+        **block_kwargs,
+    ):
+        super().__init__()
+        self.input_channels = latent_dim
+        self.base_channels = base_channels
+        self.channel_multipliers = channel_multipliers
+        self.num_res_blocks = num_res_blocks
+        self.output_nonlinearity = output_nonlinearity
+        assert nonlinearity == "silu"
+        assert causal
+
+        ch = [mult * base_channels for mult in channel_multipliers]
+        self.num_up_blocks = len(ch) - 1
+        assert len(num_res_blocks) == self.num_up_blocks + 2
+
+        blocks = []
+
+        first_block = [nn.Conv3d(latent_dim, ch[-1], kernel_size=(1, 1, 1))]  # Input layer.
+        # First set of blocks preserve channel count.
+        for _ in range(num_res_blocks[-1]):
+            first_block.append(
+                block_fn(
+                    ch[-1],
+                    has_attention=has_attention[-1],
+                    causal=causal,
+                    **block_kwargs,
+                )
+            )
+        blocks.append(nn.Sequential(*first_block))
+
+        assert len(temporal_expansions) == len(spatial_expansions) == self.num_up_blocks
+        assert len(num_res_blocks) == len(has_attention) == self.num_up_blocks + 2
+
+        upsample_block_fn = CausalUpsampleBlock
+
+        for i in range(self.num_up_blocks):
+            block = upsample_block_fn(
+                ch[-i - 1],
+                ch[-i - 2],
+                num_res_blocks=num_res_blocks[-i - 2],
+                has_attention=has_attention[-i - 2],
+                temporal_expansion=temporal_expansions[-i - 1],
+                spatial_expansion=spatial_expansions[-i - 1],
+                causal=causal,
+                **block_kwargs,
+            )
+            blocks.append(block)
+
+        assert not output_norm
+
+        # Last block. Preserve channel count.
+        last_block = []
+        for _ in range(num_res_blocks[0]):
+            last_block.append(block_fn(ch[0], has_attention=has_attention[0], causal=causal, **block_kwargs))
+        blocks.append(nn.Sequential(*last_block))
+
+        self.blocks = nn.ModuleList(blocks)
+        self.output_proj = Conv1x1(ch[0], out_channels)
+
+    def unnormalize_latents(
+        self,
+        z: torch.Tensor,
+        mean: torch.Tensor,
+        std: torch.Tensor,
+    ) -> torch.Tensor:
+        """Unnormalize latents. Useful for decoding DiT samples.
+
+        Args:
+            z (torch.Tensor): [B, C_z, T_z, H_z, W_z], float
+
+        Returns:
+            torch.Tensor: [B, C_z, T_z, H_z, W_z], float
+        """
+        mean = mean[:, None, None, None]
+        std = std[:, None, None, None]
+
+        assert z.ndim == 5
+        assert z.size(1) == mean.size(0) == std.size(0)
+        return z * std.to(z) + mean.to(z)
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x: Latent tensor. Shape: [B, input_channels, t, h, w]. Scaled [-1, 1].
+
+        Returns:
+            x: Reconstructed video tensor. Shape: [B, C, T, H, W]. Scaled to [-1, 1].
+               T + 1 = (t - 1) * 4.
+               H = h * 16, W = w * 16.
+        """
+        x = self.unnormalize_latents(x, self.vae_mean, self.vae_std)
+
+        for block in self.blocks:
+            x = block(x)
+
+        if self.output_nonlinearity == "silu":
+            x = F.silu(x, inplace=not self.training)
+        else:
+            assert not self.output_nonlinearity  # StyleGAN3 omits the to-RGB nonlinearity.
+
+        return self.output_proj(x).contiguous()
+
+
+def make_broadcastable(
+    tensor: torch.Tensor,
+    axis: int,
+    ndim: int,
+) -> torch.Tensor:
+    """
+    Reshapes the input tensor to have singleton dimensions in all axes except the specified axis.
+
+    Args:
+        tensor (torch.Tensor): The tensor to reshape. Typically 1D.
+        axis (int): The axis along which the tensor should retain its original size.
+        ndim (int): The total number of dimensions the reshaped tensor should have.
+
+    Returns:
+        torch.Tensor: The reshaped tensor with shape suitable for broadcasting.
+    """
+    if tensor.dim() != 1:
+        raise ValueError(f"Expected tensor to be 1D, but got {tensor.dim()}D tensor.")
+
+    axis = (axis + ndim) % ndim  # Ensure the axis is within the tensor dimensions
+    shape = [1] * ndim  # Start with all dimensions as 1
+    shape[axis] = tensor.size(0)  # Set the specified axis to the size of the tensor
+    return tensor.view(*shape)
+
+
+def blend(a: torch.Tensor, b: torch.Tensor, axis: int) -> torch.Tensor:
+    """
+    Blends two tensors `a` and `b` along the specified axis using linear interpolation.
+
+    Args:
+        a (torch.Tensor): The first tensor.
+        b (torch.Tensor): The second tensor. Must have the same shape as `a`.
+        axis (int): The axis along which to perform the blending.
+
+    Returns:
+        torch.Tensor: The blended tensor.
+    """
+    assert a.shape == b.shape, f"Tensors must have the same shape, got {a.shape} and {b.shape}"
+    steps = a.size(axis)
+
+    # Create a weight tensor that linearly interpolates from 0 to 1
+    start = 1 / (steps + 1)
+    end = steps / (steps + 1)
+    weight = torch.linspace(start, end, steps=steps, device=a.device, dtype=a.dtype)
+
+    # Make the weight tensor broadcastable across all dimensions
+    weight = make_broadcastable(weight, axis, a.dim())
+
+    # Perform the blending
+    return a * (1 - weight) + b * weight
+
+
+def blend_horizontal(a: torch.Tensor, b: torch.Tensor, overlap: int) -> torch.Tensor:
+    if overlap == 0:
+        return torch.cat([a, b], dim=-1)
+
+    assert a.size(-1) >= overlap
+    assert b.size(-1) >= overlap
+    a_left, a_overlap = a[..., :-overlap], a[..., -overlap:]
+    b_overlap, b_right = b[..., :overlap], b[..., overlap:]
+    return torch.cat([a_left, blend(a_overlap, b_overlap, -1), b_right], dim=-1)
+
+
+def blend_vertical(a: torch.Tensor, b: torch.Tensor, overlap: int) -> torch.Tensor:
+    if overlap == 0:
+        return torch.cat([a, b], dim=-2)
+
+    assert a.size(-2) >= overlap
+    assert b.size(-2) >= overlap
+    a_top, a_overlap = a[..., :-overlap, :], a[..., -overlap:, :]
+    b_overlap, b_bottom = b[..., :overlap, :], b[..., overlap:, :]
+    return torch.cat([a_top, blend(a_overlap, b_overlap, -2), b_bottom], dim=-2)
+
+
+def nearest_multiple(x: int, multiple: int) -> int:
+    return round(x / multiple) * multiple
+
+
+def apply_tiled(
+    fn: Callable[[torch.Tensor], torch.Tensor],
+    x: torch.Tensor,
+    num_tiles_w: int,
+    num_tiles_h: int,
+    overlap: int = 0,  # Number of pixel of overlap between adjacent tiles.
+    # Use a factor of 2 times the latent downsample factor.
+    min_block_size: int = 1,  # Minimum number of pixels in each dimension when subdividing.
+) -> Optional[torch.Tensor]:
+    if num_tiles_w == 1 and num_tiles_h == 1:
+        return fn(x)
+
+    assert num_tiles_w & (num_tiles_w - 1) == 0, f"num_tiles_w={num_tiles_w} must be a power of 2"
+    assert num_tiles_h & (num_tiles_h - 1) == 0, f"num_tiles_h={num_tiles_h} must be a power of 2"
+
+    H, W = x.shape[-2:]
+    assert H % min_block_size == 0
+    assert W % min_block_size == 0
+    ov = overlap // 2
+    assert ov % min_block_size == 0
+
+    if num_tiles_w >= 2:
+        # Subdivide horizontally.
+        half_W = nearest_multiple(W // 2, min_block_size)
+        left = x[..., :, : half_W + ov]
+        right = x[..., :, half_W - ov :]
+
+        assert num_tiles_w % 2 == 0, f"num_tiles_w={num_tiles_w} must be even"
+        left = apply_tiled(fn, left, num_tiles_w // 2, num_tiles_h, overlap, min_block_size)
+        right = apply_tiled(fn, right, num_tiles_w // 2, num_tiles_h, overlap, min_block_size)
+        if left is None or right is None:
+            return None
+
+        # If `fn` changed the resolution, adjust the overlap.
+        resample_factor = left.size(-1) / (half_W + ov)
+        out_overlap = int(overlap * resample_factor)
+
+        return blend_horizontal(left, right, out_overlap)
+
+    if num_tiles_h >= 2:
+        # Subdivide vertically.
+        half_H = nearest_multiple(H // 2, min_block_size)
+        top = x[..., : half_H + ov, :]
+        bottom = x[..., half_H - ov :, :]
+
+        assert num_tiles_h % 2 == 0, f"num_tiles_h={num_tiles_h} must be even"
+        top = apply_tiled(fn, top, num_tiles_w, num_tiles_h // 2, overlap, min_block_size)
+        bottom = apply_tiled(fn, bottom, num_tiles_w, num_tiles_h // 2, overlap, min_block_size)
+        if top is None or bottom is None:
+            return None
+
+        # If `fn` changed the resolution, adjust the overlap.
+        resample_factor = top.size(-2) / (half_H + ov)
+        out_overlap = int(overlap * resample_factor)
+
+        return blend_vertical(top, bottom, out_overlap)
+
+    raise ValueError(f"Invalid num_tiles_w={num_tiles_w} and num_tiles_h={num_tiles_h}")