From 01cf00622b7fde03526f78f097e91ded4ffac413 Mon Sep 17 00:00:00 2001 From: jiaxiangc Date: Fri, 23 Oct 2026 13:30:00 +0800 Subject: [PATCH 1/4] Add Hy-Nomas pipeline and PixelUMM decoder integration Co-authored-by: Cursor --- docs/source/en/_toctree.yml | 6 + .../en/api/models/hy_nomas_transformer3d.md | 30 + docs/source/en/api/models/pixelumm_decoder.md | 35 + docs/source/en/api/pipelines/hy_nomas.md | 235 ++++ .../community/hy_nomas_keyboard_overlay.py | 230 ++++ .../README.md | 2 +- scripts/convert_hy_nomas_to_diffusers.py | 146 +++ src/diffusers/__init__.py | 10 + src/diffusers/models/__init__.py | 4 + src/diffusers/models/attention_dispatch.py | 54 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoders/autoencoder_pixelumm.py | 1115 +++++++++++++++++ src/diffusers/models/modeling_utils.py | 1 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_hy_nomas.py | 908 ++++++++++++++ src/diffusers/modular_pipelines/__init__.py | 2 + .../modular_pipelines/hy_nomas/__init__.py | 60 + .../hy_nomas/before_denoise.py | 319 +++++ .../modular_pipelines/hy_nomas/decoders.py | 133 ++ .../modular_pipelines/hy_nomas/denoise.py | 326 +++++ .../modular_pipelines/hy_nomas/encoders.py | 204 +++ .../hy_nomas/modular_blocks_hy_nomas.py | 123 ++ .../hy_nomas/modular_pipeline.py | 57 + .../modular_pipelines/modular_pipeline.py | 1 + src/diffusers/utils/__init__.py | 1 + src/diffusers/utils/dummy_pt_objects.py | 30 + .../dummy_torch_and_transformers_objects.py | 45 + src/diffusers/utils/import_utils.py | 5 + .../autoencoders/test_models_pixelumm.py | 245 ++++ tests/models/test_attention_dispatch.py | 56 + .../test_models_transformer_hy_nomas.py | 404 ++++++ tests/modular_pipelines/hy_nomas/__init__.py | 0 .../test_modular_pipeline_hy_nomas.py | 339 +++++ .../others/test_hy_nomas_keyboard_overlay.py | 214 ++++ 34 files changed, 5341 insertions(+), 1 deletion(-) create mode 100644 docs/source/en/api/models/hy_nomas_transformer3d.md create mode 100644 docs/source/en/api/models/pixelumm_decoder.md create mode 100644 docs/source/en/api/pipelines/hy_nomas.md create mode 100644 examples/community/hy_nomas_keyboard_overlay.py create mode 100644 scripts/convert_hy_nomas_to_diffusers.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_pixelumm.py create mode 100644 src/diffusers/models/transformers/transformer_hy_nomas.py create mode 100644 src/diffusers/modular_pipelines/hy_nomas/__init__.py create mode 100644 src/diffusers/modular_pipelines/hy_nomas/before_denoise.py create mode 100644 src/diffusers/modular_pipelines/hy_nomas/decoders.py create mode 100644 src/diffusers/modular_pipelines/hy_nomas/denoise.py create mode 100644 src/diffusers/modular_pipelines/hy_nomas/encoders.py create mode 100644 src/diffusers/modular_pipelines/hy_nomas/modular_blocks_hy_nomas.py create mode 100644 src/diffusers/modular_pipelines/hy_nomas/modular_pipeline.py create mode 100644 tests/models/autoencoders/test_models_pixelumm.py create mode 100644 tests/models/transformers/test_models_transformer_hy_nomas.py create mode 100644 tests/modular_pipelines/hy_nomas/__init__.py create mode 100644 tests/modular_pipelines/hy_nomas/test_modular_pipeline_hy_nomas.py create mode 100644 tests/others/test_hy_nomas_keyboard_overlay.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 4384bc183da2..417eace2df3c 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -373,6 +373,8 @@ title: HunyuanVideo15Transformer3DModel - local: api/models/hunyuan_video_transformer_3d title: HunyuanVideoTransformer3DModel + - local: api/models/hy_nomas_transformer3d + title: HyNomasTransformer3DModel - local: api/models/ideogram4_transformer2d title: Ideogram4Transformer2DModel - local: api/models/transformer_joyimage_edit_plus @@ -523,6 +525,8 @@ title: LTX2VideoDiffusionDecoderModel - local: api/models/autoencoder_oobleck title: Oobleck AutoEncoder + - local: api/models/pixelumm_decoder + title: PixelUMMDecoder - local: api/models/autoencoder_tiny title: Tiny AutoEncoder - local: api/models/vq @@ -751,6 +755,8 @@ title: HunyuanVideo - local: api/pipelines/hunyuan_video15 title: HunyuanVideo1.5 + - local: api/pipelines/hy_nomas + title: Hy-Nomas - local: api/pipelines/kandinsky5_video title: Kandinsky 5.0 Video - local: api/pipelines/kandinsky6 diff --git a/docs/source/en/api/models/hy_nomas_transformer3d.md b/docs/source/en/api/models/hy_nomas_transformer3d.md new file mode 100644 index 000000000000..a621a6e62ca8 --- /dev/null +++ b/docs/source/en/api/models/hy_nomas_transformer3d.md @@ -0,0 +1,30 @@ + + +# HyNomasTransformer3DModel + +A double-stream transformer for segment autoregressive video generation. Video tokens use three-dimensional rotary +positions, text tokens retain their segment-local positions, and camera attention conditions each latent frame on +four pixel-frame extrinsics. The caller supplies the key/value cache explicitly. + +```python +import torch +from diffusers import HyNomasTransformer3DModel + +model = HyNomasTransformer3DModel.from_pretrained( + "/path/to/hy-nomas-diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 +) +``` + +## HyNomasTransformer3DModel + +[[autodoc]] HyNomasTransformer3DModel + - forward diff --git a/docs/source/en/api/models/pixelumm_decoder.md b/docs/source/en/api/models/pixelumm_decoder.md new file mode 100644 index 000000000000..59f80cd1e95c --- /dev/null +++ b/docs/source/en/api/models/pixelumm_decoder.md @@ -0,0 +1,35 @@ + + +# PixelUMMDecoder + +PixelUMM is a distilled diffusion decoder for 32-channel video latents, with spatial compression 32 and temporal +compression 4. It predicts clean pixels conditioned on the latent video and an independently sampled noise tensor. +It does not encode videos into latents. + +```python +import torch +from diffusers import PixelUMMDecoder + +decoder = PixelUMMDecoder.from_pretrained( + "/path/to/hy-nomas-diffusers", subfolder="decoder", torch_dtype=torch.bfloat16 +) +``` + +The decoder supports local spatial attention, overlapping temporal windows, optional spatial tiling, and explicit +streaming sessions. Noise is indexed by global latent frame, allowing overlapping windows to use consistent noise. + +## PixelUMMDecoder + +[[autodoc]] PixelUMMDecoder + - forward + - decode + - create_streaming_decode_session diff --git a/docs/source/en/api/pipelines/hy_nomas.md b/docs/source/en/api/pipelines/hy_nomas.md new file mode 100644 index 000000000000..1541c540bb77 --- /dev/null +++ b/docs/source/en/api/pipelines/hy_nomas.md @@ -0,0 +1,235 @@ + + +# Hy-Nomas + +Hy-Nomas generates a world one video segment at a time, conditioned on text and optional camera extrinsics. Its +distilled checkpoint uses eight flow-matching Euler steps, a Qwen2.5-VL text encoder, and a PixelUMM diffusion decoder. + +This implementation supports the released double-stream transformer, action-aware text visibility, camera attention, +and continuation through an explicit generation state. Text prompts can be plain strings or the segment JSON produced +by the model's prompt preprocessing. Action timestamps in JSON must already be relative to their segment. + +## Installation + +Install PyTorch 2.10.0 with CUDA support for your system, then run the following command from the Diffusers checkout: + +```bash +python -m pip install -e . "transformers==5.20.2" "accelerate==1.13.0" imageio imageio-ffmpeg +``` + +The text encoder requires `Qwen2_5_VLTextModel`, provided by the Transformers version above. The video export example +also requires `imageio` and `imageio-ffmpeg`. The base `pip install -e .` command does not install these model-specific +dependencies. Use a converted model directory as described below; the original Hy-Nomas code checkout is not required. + +## Convert the checkpoints + +Run the conversion script from a Diffusers checkout. The decoder input must be the original floating-point checkpoint. +The script checks parameter mappings and verifies the saved components by loading them again. + +```bash +python scripts/convert_hy_nomas_to_diffusers.py \ + --transformer_checkpoint /path/to/dit.safetensors \ + --decoder_checkpoint /path/to/pixelumm.safetensors \ + --text_encoder_path /path/to/Qwen2.5-VL-7B-Instruct \ + --output_path /path/to/hy-nomas-diffusers \ + --dtype bfloat16 +``` + +The output contains `modular_model_index.json` and standard component subdirectories. The converted pipeline loads +without importing the original repository or executing remote code. + +## Generate video + +```python +import torch + +from diffusers import ModularPipeline +from diffusers.utils import export_to_video + +pipe = ModularPipeline.from_pretrained("/path/to/hy-nomas-diffusers") +pipe.load_components(dtype=torch.bfloat16) +pipe.to(device="cuda", dtype=torch.bfloat16) + +result = pipe( + prompt=["A forest beside a river.", "The view moves along the river toward a waterfall."], + height=704, + width=1280, + num_frames=97, + segment_latent_lengths=[12, 13], + num_inference_steps=8, + generator=torch.Generator("cpu").manual_seed(42), + output=["videos", "session"], +) +export_to_video(result["videos"][0], "world.mp4", fps=24) +``` + +Pass the dtype to both loading and `to` to reproduce the original BF16 inference. The latter also casts the text +encoder's nonpersistent rotary-position buffers, which are recreated in FP32 when loading its weights. + +### Match reference attention + +The transformer defaults to PyTorch scaled dot product attention. On CUDA, the PixelUMM decoder uses compiled +FlexAttention for local attention, so its first call includes compilation. This default inference path does not require +MagiAttention or FlashAttention. To compare against the original BF16 implementation, install its compatible +MagiAttention, FlashAttention 2, and FlashAttention 3 dependencies and select the reference processors after loading +the transformer: + +```python +pipe.transformer.enable_reference_attention(text_backend="flash_varlen", deterministic=True) +``` + +This uses FlashAttention 2 for the unpadded text refiner, MagiAttention for action and camera visibility ranges, and +FlashAttention 3 for unrestricted history attention. MagiAttention is imported only when its backend runs; the default +processors do not require these optional dependencies. Use `disable_reference_attention()` or +`set_attention_backend("native")` to return to the standard processors. Set +`MAGI_ATTENTION_DETERMINISTIC_MODE=1` in the original implementation as well when comparing outputs. Use the same +checkpoint precision, noise timeline capacity, CPU generator seed, prompts, normalized camera inputs, and full history +on both sides. A matching seed alone does not establish numerical equivalence across different attention kernels. + +A list of prompts represents **successive segments of one video**, with batch size one. It is not a batch of independent +videos. `num_frames` is the number of new pixel frames; the latent frame count is `(num_frames - 1) // 4 + 1`. +`segment_latent_lengths` must sum to this count. If omitted, all segments receive equal lengths. +Height and width must be multiples of 32. For the original nominal 720p preset, the decoded height is 704 pixels. + +For visual quality checks, use segment lengths similar to the example. Very short segments and decoding windows can +produce oversharpened or oversaturated clips. The default PixelUMM decoding window contains 10 latent frames. + +Use `output_type="latent"` to skip decoding. The `latents` output has shape `(1, 32, T, H/32, W/32)`. +Video output types are `np` (batch, frames, height, width, channels), `pt` (batch, frames, channels, height, width), +and `pil`. + +## Camera control + +Pass `camera_extrinsics` with shape `(1, num_frames, 4, 4)`. These are **normalized camera-to-world matrices in the +model's training coordinate system**. Camera calibration, trajectory generation, and normalization must be performed +before calling the pipeline. Omitting the argument disables camera conditioning. + +```python +camera = torch.eye(4).repeat(1, 97, 1, 1) +camera[:, :, 0, 3] = torch.linspace(0, 0.2, 97) +videos = pipe( + prompt="A forest path.", + height=704, + width=1280, + num_frames=97, + camera_extrinsics=camera, + output="videos", +) +``` + +### Display recorded keyboard controls + +The `examples/community/hy_nomas_keyboard_overlay.py` example in the Diffusers checkout displays WASD and all four arrow +keys on an existing silent video. It runs on CPU without model weights or a Hy-Nomas checkout. The original video is +preserved and the example refuses to overwrite an existing output. + +```bash +pip install numpy pillow imageio imageio-ffmpeg +python examples/community/hy_nomas_keyboard_overlay.py \ + --video video.mp4 \ + --manifest replay_manifest.json \ + --output video_keyboard.mp4 +``` + +The replay manifest contains `sessions[].turns[].params`, with each turn specifying `num_frames`, `video_fps`, and +optional `camera_trajectory`. The trajectory can be a JSON object or a JSON string. Select a session with +`--session-index` when the manifest contains multiple sessions. The input video must contain exactly that session's +concatenated turns at the recorded FPS, without duplicated stable-prefix carry frames. + +For a single camera recording, replace `--manifest replay_manifest.json` with `--camera-trajectory camera.json`: + +```json +{ + "durationMs": 1000, + "keyEvents": [ + {"t": 100, "code": "KeyW", "pressed": true}, + {"t": 250, "code": "ArrowUp", "pressed": true}, + {"t": 600, "code": "ArrowUp", "pressed": false}, + {"t": 800, "code": "KeyW", "pressed": false} + ] +} +``` + +Event timestamps are milliseconds. Supported codes are `KeyW`, `KeyA`, `KeyS`, `KeyD`, `ArrowUp`, `ArrowDown`, +`ArrowLeft`, and `ArrowRight`. The video frame clock determines playback progress, which maps to the recording's +`durationMs`, matching camera playback. If `durationMs` is absent, timestamps use the turn's video clock directly. +States reset between turns. Segments without `keyEvents` show **NO KEY EVENTS** with every key released; camera motion +does not fabricate recorded input. This example displays recorded events and does not record live keyboard input or +convert keyboard events to model camera poses. Camera conditioning still uses `camera_extrinsics`. + +For a custom preview, import the example from the repository root and render frames directly: + +```python +from examples.community.hy_nomas_keyboard_overlay import KeyboardTimeline, draw_keyboard_overlay + +timeline = KeyboardTimeline([turn["params"] for turn in manifest["sessions"][0]["turns"]]) +pressed_keys, label = timeline.state_at_frame(frame_index) +preview = draw_keyboard_overlay(rgb_uint8_frame, pressed_keys, label) +``` + +The returned preview is a copy. Apply overlays after generation, and use the original frames for numerical comparisons. + +## Continue or branch a world + +Pass the returned `session` to another call. The pipeline forks the cache before adding new segments, so the same +state can be reused for alternative continuations. Keep the pipeline components, precision, device, and resolution +unchanged when reusing a state. + +```python +continued = pipe( + prompt="The river opens into a lake.", + height=704, + width=1280, + num_frames=49, + session=result["session"], + output=["videos", "session"], +) +export_to_video(continued["videos"][0], "continued-world.mp4", fps=24) +``` + +Both `videos` and `latents` include the accumulated history. Each call contributes exactly `num_frames` pixel frames; +padding at the end of a call's last segment is removed. Decoding runs over that history, which allows the temporal +decoder to revise the end of the previous clip. The session is a runtime object and is not part of `save_pretrained`. + +For identical noise between a combined rollout and incremental calls, set `noise_capacity` on the first call to the +total latent timeline capacity. Reuse that session without supplying a new generator. This allocates a fixed noise +timeline and rejects requests beyond its capacity. The decoder has independent noise; `decoder_generator` selects +its initial seed, with 42 as the default. PixelUMM does not advance this generator. + +For applications that need incremental pixel updates, `pipe.decoder.create_streaming_decode_session()` provides +`push_preview`, `finish`, and `rollback_last_push`. Preview updates identify the pixel offset to replace. + +## Reuse individual blocks + +```python +from diffusers.modular_pipelines.hy_nomas.encoders import HyNomasTextEncoderStep + +encoder = HyNomasTextEncoderStep().init_pipeline("/path/to/hy-nomas-diffusers") +encoder.load_components(dtype=torch.bfloat16) +encoder.to(device="cuda", dtype=torch.bfloat16) +encoded = encoder(prompt="A forest path.", output="prompt_embeds") +``` + +The distilled workflow uses guidance scale one. Training, original application servers, prompt expansion services, +and custom FP8 kernels remain outside this inference integration. + +## HyNomasModularPipeline + +[[autodoc]] HyNomasModularPipeline + +## HyNomasAutoBlocks + +[[autodoc]] HyNomasAutoBlocks + +## HyNomasGenerationState + +[[autodoc]] HyNomasGenerationState diff --git a/examples/community/hy_nomas_keyboard_overlay.py b/examples/community/hy_nomas_keyboard_overlay.py new file mode 100644 index 000000000000..10736b9f7220 --- /dev/null +++ b/examples/community/hy_nomas_keyboard_overlay.py @@ -0,0 +1,230 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# 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. + +"""Display recorded Hy-Nomas keyboard events on an exported, silent video. + +Install the example dependencies with `pip install numpy pillow imageio imageio-ffmpeg`. +The input JSON can be a replay manifest or a camera trajectory containing `keyEvents`. +No model weights or reference repository are needed. +""" + +import argparse +import json +import math +import os +import tempfile +from pathlib import Path + +import imageio.v2 as imageio +import numpy as np +from PIL import Image, ImageDraw, ImageFont + + +KEY_LAYOUT = { + "KeyW": (1, 0), + "KeyA": (0, 1), + "KeyS": (1, 1), + "KeyD": (2, 1), + "ArrowUp": (5, 0), + "ArrowLeft": (4, 1), + "ArrowDown": (5, 1), + "ArrowRight": (6, 1), +} + + +class KeyboardTimeline: + """Map concatenated turn frames to recorded key states. + + Each turn is a replay manifest's `params` dictionary with `num_frames`, `video_fps`, and optional + `camera_trajectory` (a dictionary or JSON string). Event `t` values are milliseconds. Like camera playback, + the recording's `durationMs` is mapped to the turn duration; without it, timestamps use the turn's video clock. + Key states reset at every turn boundary. Missing events never imply keys from camera motion. + """ + + def __init__(self, turns: list[dict]): + if not turns: + raise ValueError("At least one turn is required.") + self.fps = float(turns[0]["video_fps"]) + if not math.isfinite(self.fps) or self.fps <= 0: + raise ValueError("video_fps must be finite and positive.") + self.segments = [] + self.num_frames = 0 + for turn in turns: + num_frames = turn["num_frames"] + if not isinstance(num_frames, int) or isinstance(num_frames, bool) or num_frames <= 0: + raise ValueError("num_frames must be a positive integer.") + if not math.isclose(float(turn["video_fps"]), self.fps, rel_tol=1e-6): + raise ValueError("All turns must use the same video_fps.") + if turn.get("stable_prefix_carry_frames", 0): + raise ValueError("Use turns without duplicated stable-prefix carry frames.") + trajectory = turn.get("camera_trajectory") or {} + if isinstance(trajectory, str): + trajectory = json.loads(trajectory) + duration_ms = float(trajectory.get("durationMs", num_frames * 1000 / self.fps)) + if not math.isfinite(duration_ms) or duration_ms <= 0: + raise ValueError("durationMs must be finite and positive.") + events = trajectory.get("keyEvents", []) + for event in events: + if event["code"] not in KEY_LAYOUT: + raise ValueError(f"Unsupported key code: {event['code']}") + if not isinstance(event["pressed"], bool): + raise ValueError("Each key event's pressed value must be a boolean.") + if not math.isfinite(float(event["t"])) or float(event["t"]) < 0: + raise ValueError("Each key event's t must be finite, nonnegative milliseconds.") + self.segments.append( + { + "start_frame": self.num_frames, + "end_frame": self.num_frames + num_frames, + "duration_ms": duration_ms, + "events": sorted(events, key=lambda event: float(event["t"])), + } + ) + self.num_frames += num_frames + + def state_at_frame(self, frame_index: int) -> tuple[set[str], str]: + """Return pressed key codes and a provenance label for a zero-based video frame.""" + if not 0 <= frame_index < self.num_frames: + raise IndexError("Frame index is outside the keyboard timeline.") + for segment in self.segments: + if frame_index < segment["end_frame"]: + local_frame = frame_index - segment["start_frame"] + time_ms = local_frame * segment["duration_ms"] / (segment["end_frame"] - segment["start_frame"]) + pressed = set() + for event in segment["events"]: + if float(event["t"]) > time_ms: + break + if event["pressed"]: + pressed.add(event["code"]) + else: + pressed.discard(event["code"]) + label = "RECORDED KEY EVENTS" if segment["events"] else "NO KEY EVENTS" + return pressed, label + + +def draw_keyboard_overlay(frame: np.ndarray, pressed_keys: set[str], label: str) -> np.ndarray: + """Return an RGB uint8 copy with WASD and vector arrows; the input frame is preserved.""" + if frame.dtype != np.uint8 or frame.ndim != 3 or frame.shape[-1] != 3: + raise ValueError("frame must be an RGB uint8 array with shape (height, width, 3).") + panel = Image.new("RGBA", (216, 112)) + draw = ImageDraw.Draw(panel) + font = ImageFont.load_default() + draw.rounded_rectangle((0, 0, 215, 111), radius=7, fill=(12, 18, 30, 225)) + draw.text((10, 8), "MOVE", font=font, fill=(221, 230, 242)) + draw.text((122, 8), "LOOK", font=font, fill=(221, 230, 242)) + for code, (column, row) in KEY_LAYOUT.items(): + x, y = 10 + column * 28, 28 + row * 28 + active = code in pressed_keys + fill = (63, 226, 185, 255) if active else (47, 59, 77, 255) + foreground = (9, 28, 28, 255) if active else (221, 230, 242, 255) + draw.rounded_rectangle((x, y, x + 24, y + 24), radius=4, fill=fill) + if code.startswith("Key"): + draw.text((x + 12, y + 12), code[-1], font=font, anchor="mm", fill=foreground) + else: + points = np.array([[-5, -1], [0, -6], [5, -1], [2, -1], [2, 6], [-2, 6], [-2, -1]]) + if code == "ArrowDown": + points = -points + elif code == "ArrowLeft": + points = points[:, ::-1] + elif code == "ArrowRight": + points = -points[:, ::-1] + draw.polygon([tuple(point + [x + 12, y + 12]) for point in points], fill=foreground) + draw.text((10, 94), label, font=font, fill=(185, 202, 220)) + height, width = frame.shape[:2] + scale = min(width / 640, height / 352) + panel = panel.resize((max(1, round(216 * scale)), max(1, round(112 * scale))), Image.Resampling.LANCZOS) + margin = round(8 * scale) + output = Image.fromarray(frame.copy()) + output.paste(panel, (margin, height - panel.height - margin), panel) + return np.asarray(output) + + +def export_keyboard_video(video_path, metadata_path, output_path, session_index=0) -> dict: + """Stream a silent video to a new MP4 using a replay manifest or camera-trajectory JSON.""" + video_path, metadata_path, output_path = Path(video_path), Path(metadata_path), Path(output_path) + if output_path.resolve() == video_path.resolve(): + raise ValueError("The output must be a different file from the original video.") + if output_path.exists(): + raise FileExistsError(f"Output already exists: {output_path}") + if output_path.suffix.lower() != ".mp4": + raise ValueError("The output filename must end in .mp4.") + metadata = json.loads(metadata_path.read_text()) + with imageio.get_reader(str(video_path)) as reader: + fps = reader.get_meta_data()["fps"] + num_frames = reader.count_frames() + if "sessions" in metadata: + if not 0 <= session_index < len(metadata["sessions"]): + raise ValueError("session_index is outside the manifest's sessions.") + turns = [turn["params"] for turn in metadata["sessions"][session_index]["turns"]] + else: + if session_index != 0: + raise ValueError("session_index only applies to replay manifests.") + turns = [{"num_frames": num_frames, "video_fps": fps, "camera_trajectory": metadata}] + timeline = KeyboardTimeline(turns) + if timeline.num_frames != num_frames or not math.isclose(timeline.fps, fps, rel_tol=1e-6): + raise ValueError( + f"Video ({num_frames} frames, {fps} FPS) does not match the manifest " + f"({timeline.num_frames} frames, {timeline.fps} FPS)." + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_path = tempfile.mkstemp(suffix=".mp4", prefix=".keyboard-", dir=output_path.parent) + os.close(fd) + try: + with imageio.get_writer( + temporary_path, + fps=fps, + codec="libx264", + quality=8, + macro_block_size=1, + ffmpeg_log_level="error", + output_params=["-threads", "2"], + ) as writer: + written_frames = 0 + for frame_index, frame in enumerate(reader): + pressed, label = timeline.state_at_frame(frame_index) + writer.append_data(draw_keyboard_overlay(frame, pressed, label)) + written_frames += 1 + if written_frames != num_frames: + raise ValueError(f"Decoded {written_frames} frames; expected {num_frames}.") + os.replace(temporary_path, output_path) + finally: + if os.path.exists(temporary_path): + os.unlink(temporary_path) + return { + "input_video": str(video_path), + "output_video": str(output_path), + "fps": fps, + "num_frames": num_frames, + "recorded_key_events": sum(len(segment["events"]) for segment in timeline.segments), + "segments": [ + { + "start_frame": segment["start_frame"], + "end_frame": segment["end_frame"], + "key_source": "recorded" if segment["events"] else "none", + } + for segment in timeline.segments + ], + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--video", required=True, help="Original silent MP4, preserved without changes.") + metadata_group = parser.add_mutually_exclusive_group(required=True) + metadata_group.add_argument("--manifest", dest="metadata", help="Replay manifest JSON with sessions and turns.") + metadata_group.add_argument("--camera-trajectory", dest="metadata", help="Camera JSON containing keyEvents.") + parser.add_argument("--output", required=True, help="New MP4 filename; existing files are never overwritten.") + parser.add_argument("--session-index", type=int, default=0, help="Zero-based session index in a replay manifest.") + args = parser.parse_args() + report = export_keyboard_video(args.video, args.metadata, args.output, args.session_index) + print(json.dumps(report, indent=2)) diff --git a/examples/research_projects/multi_subject_dreambooth_inpainting/README.md b/examples/research_projects/multi_subject_dreambooth_inpainting/README.md index 3412de662f58..cb9d6ef8ba6a 100644 --- a/examples/research_projects/multi_subject_dreambooth_inpainting/README.md +++ b/examples/research_projects/multi_subject_dreambooth_inpainting/README.md @@ -73,7 +73,7 @@ accelerate launch train_multi_subject_dreambooth_inpaint.py \ ## 3. Results -A [![Weights & Biases](https://img.shields.io/badge/Weights%20&%20Biases-Report-blue)](https://wandb.ai/gzguevara/uncategorized/reports/Multi-Subject-Dreambooth-for-Inpainting--Vmlldzo2MzY5NDQ4?accessToken=y0nya2d7baguhbryxaikbfr1203amvn1jsmyl07vk122mrs7tnph037u1nqgse8t) is provided showing the training progress by every 50 steps. Note, the reported weights & biases run was performed on a A100 GPU with the following stetting: +A [![Weights & Biases](https://img.shields.io/badge/Weights%20&%20Biases-Report-blue)](https://wandb.ai/gzguevara/uncategorized/reports/Multi-Subject-Dreambooth-for-Inpainting--Vmlldzo2MzY5NDQ4) is provided showing the training progress by every 50 steps. Note, the reported weights & biases run was performed on a A100 GPU with the following stetting: ```bash accelerate launch train_multi_subject_dreambooth_inpaint.py \ diff --git a/scripts/convert_hy_nomas_to_diffusers.py b/scripts/convert_hy_nomas_to_diffusers.py new file mode 100644 index 000000000000..e531fe4cd282 --- /dev/null +++ b/scripts/convert_hy_nomas_to_diffusers.py @@ -0,0 +1,146 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import json +import re +from pathlib import Path + +import torch +from safetensors import safe_open +from transformers import AutoModel, AutoTokenizer + +from diffusers import ( + FlowMatchEulerDiscreteScheduler, + HyNomasModularPipeline, + HyNomasTransformer3DModel, + PixelUMMDecoder, +) + + +def convert_transformer_state_dict(state_dict): + converted = {} + for key, value in state_dict.items(): + key = re.sub(r"^(double_blocks\.\d+)\.((?:img|txt)_attn_)", r"\1.attn.\2", key) + key = re.sub( + r"^(txt_in\.individual_token_refiner\.blocks\.\d+)\.(self_attn_(?:qkv|proj))", + r"\1.attn.\2", + key, + ) + if key in converted: + raise ValueError(f"Duplicate converted transformer key: {key}") + converted[key] = value + return converted + + +def convert_transformer(state_dict, config=None, dtype=torch.float32): + with torch.device("meta"): + model = HyNomasTransformer3DModel(**(config or {})) + model.load_state_dict(convert_transformer_state_dict(state_dict), strict=True, assign=True) + return model.to(dtype=dtype).eval().requires_grad_(False) + + +def convert_decoder(state_dict, config=None, dtype=torch.float32): + unexpected = [key for key in state_dict if not key.startswith(("decoder.", "encoder."))] + if unexpected: + raise ValueError(f"Unexpected PixelUMM checkpoint keys: {unexpected}") + decoder_state_dict = {key: value for key, value in state_dict.items() if key.startswith("decoder.")} + if any(key.endswith((".fp8_weight", ".fp8_scale")) for key in decoder_state_dict): + raise ValueError("Use the original floating-point PixelUMM checkpoint for conversion, not its FP8 export.") + with torch.device("meta"): + model = PixelUMMDecoder(**(config or {})) + model.load_state_dict(decoder_state_dict, strict=True, assign=True) + return model.to(dtype=dtype).eval().requires_grad_(False) + + +def load_checkpoint(path, prefix=None): + with safe_open(path, framework="pt", device="cpu") as checkpoint: + return { + key: checkpoint.get_tensor(key) for key in checkpoint.keys() if prefix is None or key.startswith(prefix) + } + + +def convert_text_encoder(path, dtype=torch.bfloat16): + model = AutoModel.from_pretrained(path, dtype=dtype) + if hasattr(model, "language_model"): + model = model.language_model + if model.__class__.__name__ != "Qwen2_5_VLTextModel": + raise ValueError(f"Expected a Qwen2.5-VL text encoder, got {model.__class__.__name__}.") + return model.eval().requires_grad_(False) + + +def verify_saved_pipeline(pipeline, output_path, dtype): + reloaded = HyNomasModularPipeline.from_pretrained(output_path) + reloaded.load_components(dtype=dtype) + for component_name in ("transformer", "decoder", "text_encoder"): + original_state = getattr(pipeline, component_name).state_dict() + reloaded_state = getattr(reloaded, component_name).state_dict() + if original_state.keys() != reloaded_state.keys(): + raise ValueError(f"Saving and reloading changed {component_name} parameter names.") + for key, tensor in original_state.items(): + if not torch.equal(tensor, reloaded_state[key]): + raise ValueError(f"Saving and reloading changed {component_name}.{key}.") + if pipeline.tokenizer.get_vocab() != reloaded.tokenizer.get_vocab(): + raise ValueError("Saving and reloading changed the tokenizer vocabulary.") + original_scheduler_config = { + key: value for key, value in pipeline.scheduler.config.items() if not key.startswith("_") + } + reloaded_scheduler_config = { + key: value for key, value in reloaded.scheduler.config.items() if not key.startswith("_") + } + if original_scheduler_config != reloaded_scheduler_config: + raise ValueError("Saving and reloading changed the scheduler configuration.") + + +def main(args): + dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16}[args.dtype] + transformer_config = json.loads(Path(args.transformer_config).read_text()) if args.transformer_config else None + decoder_config = json.loads(Path(args.decoder_config).read_text()) if args.decoder_config else None + print("Converting the Hy-Nomas transformer...", flush=True) + transformer = convert_transformer(load_checkpoint(args.transformer_checkpoint), transformer_config, dtype) + print("Converting the PixelUMM decoder...", flush=True) + decoder = convert_decoder(load_checkpoint(args.decoder_checkpoint, prefix="decoder."), decoder_config, dtype) + print("Extracting the Qwen text encoder...", flush=True) + text_encoder = convert_text_encoder(args.text_encoder_path, dtype) + tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_path, padding_side="right") + pipeline = HyNomasModularPipeline() + pipeline.update_components( + transformer=transformer, + decoder=decoder, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=FlowMatchEulerDiscreteScheduler(shift=5.0), + ) + output_path = str(Path(args.output_path).resolve()) + print(f"Saving the pipeline to {output_path}...", flush=True) + pipeline.save_pretrained(output_path, max_shard_size="5GB", overwrite_modular_index=True) + print("Reloading the pipeline and checking every parameter...", flush=True) + verify_saved_pipeline(pipeline, output_path, dtype) + print(f"Saved and verified the Hy-Nomas pipeline at {output_path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Convert the Hy-Nomas checkpoints into a Diffusers modular pipeline.") + parser.add_argument( + "--transformer_checkpoint", required=True, help="Original floating-point DiT safetensors file." + ) + parser.add_argument( + "--decoder_checkpoint", required=True, help="Original floating-point PixelUMM safetensors file." + ) + parser.add_argument("--text_encoder_path", required=True, help="Qwen2.5-VL-7B-Instruct directory or Hub model ID.") + parser.add_argument("--output_path", required=True, help="Directory for the converted pipeline.") + parser.add_argument("--dtype", choices=("float32", "bfloat16"), default="bfloat16") + parser.add_argument("--transformer_config", help="Optional JSON model configuration for a compatible checkpoint.") + parser.add_argument("--decoder_config", help="Optional JSON model configuration for a compatible checkpoint.") + main(parser.parse_args()) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index ac677917242b..62cf68bda736 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -316,6 +316,7 @@ "HunyuanVideo15Transformer3DModel", "HunyuanVideoFramepackTransformer3DModel", "HunyuanVideoTransformer3DModel", + "HyNomasTransformer3DModel", "I2VGenXLUNet", "Ideogram4Transformer2DModel", "JoyImageEditPlusTransformer3DModel", @@ -354,6 +355,7 @@ "OvisImageTransformer2DModel", "ParallelConfig", "PixArtTransformer2DModel", + "PixelUMMDecoder", "PriorTransformer", "PRXTransformer2DModel", "QwenImage21Transformer2DModel", @@ -575,6 +577,9 @@ "HeliosPyramidModularPipeline", "HunyuanVideo15AutoBlocks", "HunyuanVideo15ModularPipeline", + "HyNomasAutoBlocks", + "HyNomasGenerationState", + "HyNomasModularPipeline", "Ideogram4AutoBlocks", "Ideogram4ModularPipeline", "Krea2AutoBlocks", @@ -1238,6 +1243,7 @@ HunyuanVideo15Transformer3DModel, HunyuanVideoFramepackTransformer3DModel, HunyuanVideoTransformer3DModel, + HyNomasTransformer3DModel, I2VGenXLUNet, Ideogram4Transformer2DModel, JoyImageEditPlusTransformer3DModel, @@ -1276,6 +1282,7 @@ OvisImageTransformer2DModel, ParallelConfig, PixArtTransformer2DModel, + PixelUMMDecoder, PriorTransformer, PRXTransformer2DModel, QwenImage21Transformer2DModel, @@ -1476,6 +1483,9 @@ HeliosPyramidModularPipeline, HunyuanVideo15AutoBlocks, HunyuanVideo15ModularPipeline, + HyNomasAutoBlocks, + HyNomasGenerationState, + HyNomasModularPipeline, Ideogram4AutoBlocks, Ideogram4ModularPipeline, Krea2AutoBlocks, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 3713cdfec22c..068529cbed27 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -58,6 +58,7 @@ _import_structure["autoencoders.autoencoder_longcat_audio_dit"] = ["LongCatAudioDiTVae"] _import_structure["autoencoders.autoencoder_mmaudio"] = ["MMAudioVAE"] _import_structure["autoencoders.autoencoder_oobleck"] = ["AutoencoderOobleck"] + _import_structure["autoencoders.autoencoder_pixelumm"] = ["PixelUMMDecoder"] _import_structure["autoencoders.autoencoder_rae"] = ["AutoencoderRAE"] _import_structure["autoencoders.autoencoder_rfd4_proteina"] = ["RFD4ProteinaAutoencoder"] _import_structure["autoencoders.autoencoder_same"] = ["AutoencoderSAME"] @@ -134,6 +135,7 @@ _import_structure["transformers.transformer_hunyuan_video15"] = ["HunyuanVideo15Transformer3DModel"] _import_structure["transformers.transformer_hunyuan_video_framepack"] = ["HunyuanVideoFramepackTransformer3DModel"] _import_structure["transformers.transformer_hunyuanimage"] = ["HunyuanImageTransformer2DModel"] + _import_structure["transformers.transformer_hy_nomas"] = ["HyNomasTransformer3DModel"] _import_structure["transformers.transformer_ideogram4"] = ["Ideogram4Transformer2DModel"] _import_structure["transformers.transformer_joyimage"] = ["JoyImageEditTransformer3DModel"] _import_structure["transformers.transformer_joyimage_edit_plus"] = ["JoyImageEditPlusTransformer3DModel"] @@ -227,6 +229,7 @@ MiniMaxMusic3Vocoder, MMAudioVAE, MMAudioVocoder, + PixelUMMDecoder, RFD4ProteinaAutoencoder, VQModel, ) @@ -288,6 +291,7 @@ HunyuanVideo15Transformer3DModel, HunyuanVideoFramepackTransformer3DModel, HunyuanVideoTransformer3DModel, + HyNomasTransformer3DModel, Ideogram4Transformer2DModel, JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 3d2ae36bd9a2..a66333f1d2f9 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -41,6 +41,7 @@ is_flash_attn_version, is_kernels_available, is_kernels_version, + is_magi_attention_available, is_sageattention_available, is_sageattention_version, is_torch_npu_available, @@ -69,6 +70,7 @@ _CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION) _CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available() _CAN_USE_SAGE_ATTN = is_sageattention_available() and is_sageattention_version(">=", _REQUIRED_SAGE_VERSION) +_CAN_USE_MAGI_ATTN = is_magi_attention_available() _CAN_USE_FLEX_ATTN = is_torch_version(">=", _REQUIRED_FLEX_VERSION) _CAN_USE_NPU_ATTN = is_torch_npu_available() _CAN_USE_XLA_ATTN = is_torch_xla_available() and is_torch_xla_version(">=", _REQUIRED_XLA_VERSION) @@ -232,6 +234,9 @@ class AttentionBackendName(str, Enum): # `aiter` (via the `kernels-community/aiter-flash-attn-ck` Hub kernel) AITER_FA2_HUB = "aiter_fa2_hub" + # `magi_attention` rectangle-range attention + MAGI = "magi" + # PyTorch native FLEX = "flex" NATIVE = "native" @@ -583,6 +588,12 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None f"Sage Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `sageattention>={_REQUIRED_SAGE_VERSION}`." ) + elif backend == AttentionBackendName.MAGI: + if not _CAN_USE_MAGI_ATTN: + raise RuntimeError( + "Magi Attention backend requires the optional `magi_attention` package. Please install it with `pip install magi_attention`." + ) + elif backend == AttentionBackendName.FLEX: if not _CAN_USE_FLEX_ATTN: raise RuntimeError( @@ -3219,6 +3230,45 @@ def _flash_varlen_attention_hub( return (out, lse) if return_lse else out +@_AttentionBackendRegistry.register( + AttentionBackendName.MAGI, + constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape], +) +def _magi_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + query_ranges: torch.Tensor, + key_ranges: torch.Tensor, + attn_type_map: torch.Tensor | None = None, + attn_mask: torch.Tensor | None = None, + scale: float | None = None, + is_causal: bool = False, + deterministic: bool = False, +) -> torch.Tensor: + if attn_mask is not None or is_causal: + raise ValueError("Magi attention uses explicit query/key ranges and attention types.") + _check_attention_backend_requirements(AttentionBackendName.MAGI) + try: + from magi_attention.functional import flex_flash_attn_func + except ImportError as error: + raise ImportError("The magi attention backend requires the optional magi_attention package.") from error + outputs = [] + for sample_query, sample_key, sample_value in zip(query, key, value): + output, _ = flex_flash_attn_func( + q=sample_query, + k=sample_key, + v=sample_value, + q_ranges=query_ranges, + k_ranges=key_ranges, + attn_type_map=attn_type_map, + softmax_scale=scale, + deterministic=deterministic, + ) + outputs.append(output) + return torch.stack(outputs) + + @_AttentionBackendRegistry.register( AttentionBackendName.FLASH_VARLEN, constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape], @@ -3233,6 +3283,7 @@ def _flash_varlen_attention( is_causal: bool = False, window_size: tuple[int, int] = (-1, -1), return_lse: bool = False, + deterministic: bool = False, _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: batch_size, seq_len_q, _, _ = query.shape @@ -3270,6 +3321,7 @@ def _flash_varlen_attention( causal=is_causal, window_size=window_size, return_attn_probs=return_lse, + deterministic=deterministic, ) out = out.unflatten(0, (batch_size, -1)) @@ -3523,6 +3575,7 @@ def _flash_varlen_attention_3( scale: float | None = None, is_causal: bool = False, return_lse: bool = False, + deterministic: bool = False, _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: batch_size, seq_len_q, _, _ = query.shape @@ -3558,6 +3611,7 @@ def _flash_varlen_attention_3( softmax_scale=scale, causal=is_causal, return_attn_probs=return_lse, + deterministic=deterministic, ) if isinstance(result, tuple): out, lse, *_ = result diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 1cbe189979c9..48ce1ea276e9 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -28,6 +28,7 @@ from .autoencoder_longcat_audio_dit import LongCatAudioDiTVae from .autoencoder_mmaudio import MMAudioVAE from .autoencoder_oobleck import AutoencoderOobleck +from .autoencoder_pixelumm import PixelUMMDecoder from .autoencoder_rae import AutoencoderRAE from .autoencoder_rfd4_proteina import RFD4ProteinaAutoencoder from .autoencoder_same import AutoencoderSAME diff --git a/src/diffusers/models/autoencoders/autoencoder_pixelumm.py b/src/diffusers/models/autoencoders/autoencoder_pixelumm.py new file mode 100644 index 000000000000..c8316b9beb80 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_pixelumm.py @@ -0,0 +1,1115 @@ +# Copyright 2026 The Hunyuan Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from contextlib import nullcontext +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import BaseOutput +from ...utils.accelerate_utils import apply_forward_hook +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_utils import ModelMixin +from .vae import DecoderOutput + + +TUBELET = (4, 16, 16) +TUBE_DIM = TUBELET[0] * TUBELET[1] * TUBELET[2] * 3 +LATENT_CHANNELS = 32 +# A DCAE latent cell is (4, 32, 32) pixels: 1 tubelet in time, 2 x 2 in space. +LATENT_TOKEN_GRID = 2 + + +@dataclass(frozen=True) +class PixelUMMArchitecture: + """Qwen3-8B (hidden 4096) with PixelUMM's SenseNova T/H/W MRoPE.""" + + hidden_size: int = 4096 + intermediate_size: int = 12288 + num_layers: int = 36 + num_heads: int = 32 + num_kv_heads: int = 8 + head_dim: int = 128 + rms_norm_eps: float = 1e-6 + rope_theta_t: float = 1_000_000.0 + rope_theta_hw: float = 10_000.0 + + def __post_init__(self): + if self.head_dim % 4 or self.num_heads % self.num_kv_heads: + raise ValueError("head_dim must be divisible by 4 and heads by kv heads") + + @property + def rope_sections(self) -> tuple[int, int, int]: + """Head-dim slices rotated by the T, H and W position respectively.""" + return self.head_dim // 2, self.head_dim // 4, self.head_dim // 4 + + +def patchify(video: torch.Tensor) -> torch.Tensor: + """[B,3,T,H,W] -> [B, (T/4)(H/16)(W/16), 3072]; token order T,H,W, values (t,p,q,c).""" + if video.ndim != 5 or video.shape[1] != 3: + raise ValueError(f"expected RGB video [B,3,T,H,W], got {tuple(video.shape)}") + pt, ph, pw = TUBELET + batch, channels, frames, height, width = video.shape + if frames % pt or height % ph or width % pw: + raise ValueError(f"video {tuple(video.shape[2:])} must be a multiple of {TUBELET}") + x = video.reshape(batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw) + return x.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(batch, -1, TUBE_DIM) + + +def unpatchify(tokens: torch.Tensor, grid: tuple[int, int, int]) -> torch.Tensor: + batch, count, dim = tokens.shape + nt, nh, nw = grid + pt, ph, pw = TUBELET + if count != nt * nh * nw or dim != TUBE_DIM: + raise ValueError(f"tokens {tuple(tokens.shape)} do not match grid {grid}") + x = tokens.reshape(batch, nt, nh, nw, pt, ph, pw, 3) + return x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(batch, 3, nt * pt, nh * ph, nw * pw) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def mrope_tables(arch: PixelUMMArchitecture, grid: tuple[int, int, int], device) -> tuple[torch.Tensor, torch.Tensor]: + """cos/sin [1, T*H*W, 1, head_dim] for a THW-ordered grid, laid out [T | H | W].""" + positions = torch.meshgrid(*(torch.arange(n, device=device) for n in grid), indexing="ij") + thetas = (arch.rope_theta_t, arch.rope_theta_hw, arch.rope_theta_hw) + cos_parts, sin_parts = [], [] + for position, dim, theta in zip(positions, arch.rope_sections, thetas): + inv_freq = theta ** (-torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim) + angles = position.flatten().float()[:, None] * inv_freq[None] + angles = torch.cat((angles, angles), dim=-1) + cos_parts.append(angles.cos()) + sin_parts.append(angles.sin()) + return ( + torch.cat(cos_parts, dim=-1)[None, :, None, :], + torch.cat(sin_parts, dim=-1)[None, :, None, :], + ) + + +def apply_mrope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, sections) -> torch.Tensor: + """x [B,N,H,D]; each section is rotated within itself (rotate_half per axis).""" + rotated = torch.cat([_rotate_half(part) for part in x.split(tuple(sections), dim=-1)], dim=-1) + return x * cos + rotated * sin + + +def local_attention_layout(grid, tile, device): + frames, height, width = grid + tiles_h, tiles_w = -(-height // tile), -(-width // tile) + if tiles_h <= 3 and tiles_w <= 3: + return None + time, row, column = torch.meshgrid(*(torch.arange(n, device=device) for n in grid), indexing="ij") + tile_h, tile_w = (row // tile).flatten(), (column // tile).flatten() + order = (((tile_h * tiles_w + tile_w) * frames + time.flatten()) * tile + (row % tile).flatten()) * tile + permutation = torch.argsort(order + (column % tile).flatten()) + ranges = [] + offset = 0 + for row_index in range(tiles_h): + for column_index in range(tiles_w): + count = frames * min(tile, height - row_index * tile) * min(tile, width - column_index * tile) + ranges.append((offset, offset + count)) + offset += count + neighborhoods = [] + for row_index in range(tiles_h): + for column_index in range(tiles_w): + center_h = min(max(row_index, 1), max(tiles_h - 2, 1)) + center_w = min(max(column_index, 1), max(tiles_w - 2, 1)) + neighbors = [ + ranges[h * tiles_w + w] + for h in range(max(0, center_h - 1), min(tiles_h, center_h + 2)) + for w in range(max(0, center_w - 1), min(tiles_w, center_w + 2)) + ] + neighborhoods.append((ranges[row_index * tiles_w + column_index], neighbors)) + block_mask = None + if torch.device(device).type == "cuda": + from torch.nn.attention.flex_attention import create_block_mask + + def mask_mod(batch, head, query_index, key_index): + # The permutation groups tiles row by row, including partial edge tiles. + row_tokens = frames * tile * width + query_row, key_row = query_index // row_tokens, key_index // row_tokens + query_height = (height - query_row * tile).clamp(1, tile) + key_height = (height - key_row * tile).clamp(1, tile) + query_column = (query_index % row_tokens) // (frames * query_height * tile) + key_column = (key_index % row_tokens) // (frames * key_height * tile) + center_row = query_row.clamp(1, max(tiles_h - 2, 1)) + center_column = query_column.clamp(1, max(tiles_w - 2, 1)) + return ((center_row - key_row).abs() <= 1) & ((center_column - key_column).abs() <= 1) + + count = frames * height * width + block_mask = create_block_mask( + mask_mod, None, None, count, count, device=device, _compile=not torch.compiler.is_compiling() + ) + return permutation, torch.argsort(permutation), (neighborhoods, block_mask) + + +def _pixelumm_flex_attention(query, key, value, block_mask, parallel_config): + return dispatch_attention_fn( + query, key, value, attn_mask=block_mask, enable_gqa=True, backend="flex", parallel_config=parallel_config + ) + + +class PixelUMMRMSNorm(nn.Module): + """Qwen3 RMSNorm: normalize in fp32, cast back to the input dtype, then scale.""" + + def __init__(self, dim: int, eps: float): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + return self.weight * x.to(dtype) + + +class PixelUMMAttnProcessor: + _attention_backend = None + _parallel_config = None + _compiled_flex_attention = None + + def __call__(self, attn, hidden_states, cos, sin, neighborhoods=None): + batch, count, _ = hidden_states.shape + arch = attn.arch + query = attn.q_proj(hidden_states).view(batch, count, arch.num_heads, arch.head_dim) + key = attn.k_proj(hidden_states).view(batch, count, arch.num_kv_heads, arch.head_dim) + value = attn.v_proj(hidden_states).view(batch, count, arch.num_kv_heads, arch.head_dim) + query_time, query_spatial = query.float().chunk(2, dim=-1) + key_time, key_spatial = key.float().chunk(2, dim=-1) + query = torch.cat((attn.q_norm_t(query_time), attn.q_norm_hw(query_spatial)), dim=-1) + key = torch.cat((attn.k_norm_t(key_time), attn.k_norm_hw(key_spatial)), dim=-1) + query = apply_mrope(query, cos, sin, arch.rope_sections).to(value.dtype) + key = apply_mrope(key, cos, sin, arch.rope_sections).to(value.dtype) + regions, block_mask = neighborhoods if neighborhoods is not None else (None, None) + if block_mask is not None and self._attention_backend in (None, "flex"): + if torch.compiler.is_compiling(): + hidden_states = _pixelumm_flex_attention(query, key, value, block_mask, self._parallel_config) + else: + if self._compiled_flex_attention is None: + self._compiled_flex_attention = torch.compile(_pixelumm_flex_attention, dynamic=False) + hidden_states = self._compiled_flex_attention(query, key, value, block_mask, self._parallel_config) + else: + repeat = arch.num_heads // arch.num_kv_heads + key = key.repeat_interleave(repeat, dim=2) + value = value.repeat_interleave(repeat, dim=2) + if regions is None: + hidden_states = dispatch_attention_fn( + query, key, value, backend=self._attention_backend, parallel_config=self._parallel_config + ) + else: + outputs = [] + for (start, end), neighbors in regions: + local_key = torch.cat([key[:, first:last] for first, last in neighbors], dim=1) + local_value = torch.cat([value[:, first:last] for first, last in neighbors], dim=1) + outputs.append( + dispatch_attention_fn( + query[:, start:end].contiguous(), + local_key, + local_value, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + ) + hidden_states = torch.cat(outputs, dim=1) + return attn.o_proj(hidden_states.reshape(batch, count, arch.num_heads * arch.head_dim)) + + +class PixelUMMAttention(nn.Module, AttentionModuleMixin): + _supports_qkv_fusion = False + _default_processor_cls = PixelUMMAttnProcessor + _available_processors = [PixelUMMAttnProcessor] + + def __init__(self, arch): + super().__init__() + self.arch = arch + hidden, heads, kv, dim = arch.hidden_size, arch.num_heads, arch.num_kv_heads, arch.head_dim + self.q_proj = nn.Linear(hidden, heads * dim, bias=False) + self.k_proj = nn.Linear(hidden, kv * dim, bias=False) + self.v_proj = nn.Linear(hidden, kv * dim, bias=False) + self.o_proj = nn.Linear(heads * dim, hidden, bias=False) + self.q_norm_t = PixelUMMRMSNorm(dim // 2, arch.rms_norm_eps) + self.q_norm_hw = PixelUMMRMSNorm(dim // 2, arch.rms_norm_eps) + self.k_norm_t = PixelUMMRMSNorm(dim // 2, arch.rms_norm_eps) + self.k_norm_hw = PixelUMMRMSNorm(dim // 2, arch.rms_norm_eps) + self.set_processor(PixelUMMAttnProcessor()) + + def forward(self, hidden_states, cos, sin, neighborhoods=None) -> torch.Tensor: + return self.processor(self, hidden_states, cos, sin, neighborhoods) + + +class PixelUMMMLP(nn.Module): + def __init__(self, arch: PixelUMMArchitecture): + super().__init__() + self.gate_proj = nn.Linear(arch.hidden_size, arch.intermediate_size, bias=False) + self.up_proj = nn.Linear(arch.hidden_size, arch.intermediate_size, bias=False) + self.down_proj = nn.Linear(arch.intermediate_size, arch.hidden_size, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class PixelUMMBlock(nn.Module): + def __init__(self, arch: PixelUMMArchitecture): + super().__init__() + self.input_norm = PixelUMMRMSNorm(arch.hidden_size, arch.rms_norm_eps) + self.attn = PixelUMMAttention(arch) + self.post_attn_norm = PixelUMMRMSNorm(arch.hidden_size, arch.rms_norm_eps) + self.mlp = PixelUMMMLP(arch) + + def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, neighborhoods=None) -> torch.Tensor: + x = x + self.attn(self.input_norm(x), cos, sin, neighborhoods) + return x + self.mlp(self.post_attn_norm(x)) + + +class PixelUMMHead(nn.Module): + """Normalize hidden states and project each tubelet to RGB pixels.""" + + def __init__(self, hidden: int): + super().__init__() + self.norm = nn.RMSNorm(hidden, eps=1e-6) + self.linear = nn.Linear(hidden, TUBE_DIM, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.norm(x)) + + +class PixelUMMDecoderModel(nn.Module): + """`forward(latent [B,32,L,h,w], xt [B,3,4L,32h,32w]) -> x0 prediction`.""" + + def __init__(self, arch: PixelUMMArchitecture, *, local_attention_tile: int = 0): + super().__init__() + if local_attention_tile < 0: + raise ValueError(f"local_attention_tile must be >= 0 tokens, got {local_attention_tile}") + self.arch = arch + self.local_attention_tile = int(local_attention_tile) + self.gradient_checkpointing = False + hidden = arch.hidden_size + self.embed = nn.Linear(TUBE_DIM, hidden) + self.cond_proj = nn.Linear(LATENT_CHANNELS, LATENT_TOKEN_GRID**2 * hidden) + self.layers = nn.ModuleList(PixelUMMBlock(arch) for _ in range(arch.num_layers)) + self.final_norm = PixelUMMRMSNorm(hidden, arch.rms_norm_eps) + self.head = PixelUMMHead(hidden) + + def forward(self, latent: torch.Tensor, xt: torch.Tensor) -> torch.Tensor: + if latent.ndim != 5 or latent.shape[1] != LATENT_CHANNELS: + raise ValueError(f"expected DCAE latent [B,{LATENT_CHANNELS},L,h,w], got {tuple(latent.shape)}") + batch, _, frames, height, width = latent.shape + pixel_shape = (batch, 3, frames * TUBELET[0], height * 32, width * 32) + if tuple(xt.shape) != pixel_shape: + raise ValueError(f"xt {tuple(xt.shape)} != expected {pixel_shape}") + dtype = latent.dtype + context = ( + torch.autocast(device_type=latent.device.type, dtype=dtype) + if dtype in (torch.bfloat16, torch.float16) + else nullcontext() + ) + grid = (frames, height * LATENT_TOKEN_GRID, width * LATENT_TOKEN_GRID) + with context: + tokens = self.embed(patchify(xt.to(dtype))) + condition = self.cond_proj(latent.to(dtype).permute(0, 2, 3, 4, 1)) + condition = condition.view( + batch, frames, height, width, LATENT_TOKEN_GRID, LATENT_TOKEN_GRID, self.arch.hidden_size + ) + condition = condition.permute(0, 1, 2, 4, 3, 5, 6).reshape(batch, -1, self.arch.hidden_size) + tokens = tokens + condition.to(tokens.device) + cos, sin = mrope_tables(self.arch, grid, tokens.device) + layout = ( + local_attention_layout(grid, self.local_attention_tile, tokens.device) + if self.local_attention_tile + else None + ) + neighborhoods = None + if layout is not None: + perm, inverse, neighborhoods = layout + tokens, cos, sin = tokens[:, perm], cos[:, perm], sin[:, perm] + for layer in self.layers: + if torch.is_grad_enabled() and self.gradient_checkpointing: + tokens = self._gradient_checkpointing_func(layer, tokens, cos, sin, neighborhoods) + else: + tokens = layer(tokens, cos, sin, neighborhoods) + tokens = self.final_norm(tokens) + if layout is not None: + tokens = tokens[:, inverse.to(tokens.device)] + # The source runs the head after leaving autocast (forward_features returns first). + return unpatchify(self.head(tokens), grid) + + +@dataclass +class StreamingPreviewUpdate(BaseOutput): + """One complete segment preview plus any corrected preceding tail.""" + + video: torch.Tensor + pixel_start: int + pixel_end: int + replace_pixel_frames: int + stable_pixel_frames: int + provisional_pixel_frames: int + + +def _window_starts(length: int, window: int, stride: int) -> list[int]: + if length < 1: + raise ValueError(f"length must be positive, got {length}") + if window < 1 or stride < 1 or stride > window: + raise ValueError(f"expected 1 <= stride <= window, got stride={stride}, window={window}") + if length <= window: + return [0] + starts = list(range(0, length - window + 1, stride)) + tail = length - window + if starts[-1] != tail: + starts.append(tail) + return starts + + +def _slice_along(tensor: torch.Tensor, dim: int, start: int, end: int) -> torch.Tensor: + slices = [slice(None)] * tensor.ndim + slices[dim] = slice(start, end) + return tensor[tuple(slices)] + + +def _merge_overlapping(parts: list[torch.Tensor], starts: list[int], *, dim: int) -> torch.Tensor: + if len(parts) != len(starts) or not parts: + raise ValueError("parts and starts must be non-empty and have equal length") + result = parts[0] + result_end = starts[0] + parts[0].shape[dim] + for part, start in zip(parts[1:], starts[1:]): + part_end = start + part.shape[dim] + overlap = result_end - start + if overlap < 0 or part_end <= result_end: + raise ValueError("overlap windows must be ordered and extend the range") + overlap = min(overlap, result.shape[dim], part.shape[dim]) + position = torch.arange(1, overlap + 1, device=part.device, dtype=torch.float32) + weight = 0.5 - 0.5 * torch.cos(math.pi * position / (overlap + 1)) + shape = [1] * part.ndim + shape[dim] = overlap + weight = weight.view(shape).to(dtype=part.dtype) + result_head = _slice_along(result, dim, 0, result.shape[dim] - overlap) + result_tail = _slice_along(result, dim, result.shape[dim] - overlap, result.shape[dim]) + part_head = _slice_along(part, dim, 0, overlap) + part_tail = _slice_along(part, dim, overlap, part.shape[dim]) + result = torch.cat([result_head, result_tail * (1.0 - weight) + part_head * weight, part_tail], dim=dim) + result_end = part_end + return result + + +class _TemporalBlender: + """Incremental twin of `_merge_overlapping(parts, starts, dim=2)` into one buffer. + + Same arithmetic as the source's `_merge_into_buffer`: one preallocated clip buffer, each window cosine-blended in + as soon as it is complete. + """ + + def __init__(self, total_frames: int): + self.total = int(total_frames) + self.out: torch.Tensor | None = None + self.end = 0 + + def add(self, part: torch.Tensor, start: int) -> None: + length = part.shape[2] + if self.out is None: + self.out = part.new_empty((*part.shape[:2], self.total, *part.shape[3:])) + self.out[:, :, :length] = part + self.end = length + return + overlap = self.end - start + if overlap < 0 or start + length <= self.end: + raise ValueError( + f"overlap windows must be ordered and extend the covered range: start={start}, result_end={self.end}, part_end={start + length}" + ) + position = torch.arange(1, overlap + 1, device=part.device, dtype=torch.float32) + weight = (0.5 - 0.5 * torch.cos(math.pi * position / (overlap + 1))).view(1, 1, -1, 1, 1) + weight = weight.to(dtype=part.dtype) + out = self.out + previous = out[:, :, start : self.end] + out[:, :, start : self.end] = previous * (1.0 - weight) + part[:, :, :overlap] * weight + out[:, :, self.end : start + length] = part[:, :, overlap:] + self.end = start + length + + def result(self) -> torch.Tensor: + if self.out is None or self.end != self.total: + raise RuntimeError(f"temporal merge incomplete: covered {self.end} of {self.total}") + return self.out + + +class PixelUMMDecoder(ModelMixin, ConfigMixin, AttentionMixin): + """ + PixelUMM DMD2 decoder for DCAE-V video latents. + + This decoder predicts RGB pixels from 32-channel latents with spatial compression 32 and temporal compression 4. + Local attention covers three neighboring spatial tiles along each axis and the complete temporal window. The model + supports standard [`ModelMixin.save_pretrained`] and [`ModelMixin.from_pretrained`] serialization. + + Parameters: + hidden_size (`int`, defaults to 4096): + Transformer hidden dimension. + intermediate_size (`int`, defaults to 12288): + Gated feed-forward intermediate dimension. + num_layers (`int`, defaults to 36): + Number of transformer layers. + num_heads (`int`, defaults to 32): + Number of query heads. + num_kv_heads (`int`, defaults to 8): + Number of key and value heads. + head_dim (`int`, defaults to 128): + Dimension of each attention head. + rms_norm_eps (`float`, defaults to 1e-6): + RMS normalization epsilon. + rope_theta_t (`float`, defaults to 1000000.0): + Rotary embedding frequency base for the temporal axis. + rope_theta_hw (`float`, defaults to 10000.0): + Rotary embedding frequency base for both spatial axes. + local_attention_tile_pixels (`int`, defaults to 128): + Spatial attention tile size in pixels. Zero enables full attention. + denoising_step_list (`tuple[float]`, defaults to `(0.999, 0.866, 0.634, 0.342)`): + Distilled noise schedule, in decreasing order. + num_steps (`int`, defaults to 1): + Number of entries to use from the distilled schedule. + seed (`int`, defaults to 42): + Pixel noise seed used when no generator is supplied. + temporal_window_latents (`int`, defaults to 10): + Number of latent frames decoded in each temporal window. + temporal_overlap_latents (`int`, defaults to 5): + Overlap between temporal windows, in latent frames. + decoder_tile_pixels (`int`, defaults to 0): + Spatial decode tile size in pixels. Zero decodes the full spatial field. + decoder_tile_overlap_pixels (`int`, defaults to 64): + Overlap between spatial decode tiles, in pixels. + tail_hold_latents (`int`, defaults to 1): + Number of tail latents withheld by preview consumers when a later segment will correct the tail. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = ["PixelUMMBlock"] + _repeated_blocks = ["PixelUMMBlock"] + _skip_layerwise_casting_patterns = ["embed", "cond_proj", "norm", "head"] + + spatial_compression_ratio = 32 + time_compression_ratio = 4 + temporal_compression_ratio = 4 + scaling_factor = 1.0 + + @register_to_config + def __init__( + self, + hidden_size: int = 4096, + intermediate_size: int = 12288, + num_layers: int = 36, + num_heads: int = 32, + num_kv_heads: int = 8, + head_dim: int = 128, + rms_norm_eps: float = 1e-6, + rope_theta_t: float = 1_000_000.0, + rope_theta_hw: float = 10_000.0, + local_attention_tile_pixels: int = 128, + denoising_step_list: tuple[float, ...] = (0.999, 0.866, 0.634, 0.342), + num_steps: int = 1, + seed: int = 42, + temporal_window_latents: int = 10, + temporal_overlap_latents: int = 5, + decoder_tile_pixels: int = 0, + decoder_tile_overlap_pixels: int = 64, + tail_hold_latents: int = 1, + ): + super().__init__() + if local_attention_tile_pixels < 0 or local_attention_tile_pixels % 16: + raise ValueError("local_attention_tile_pixels must be zero or a multiple of 16") + if temporal_window_latents < 1 or not 0 <= temporal_overlap_latents < temporal_window_latents: + raise ValueError("temporal overlap must be nonnegative and smaller than the window") + if decoder_tile_pixels < 0 or decoder_tile_pixels % 32: + raise ValueError("decoder_tile_pixels must be zero or a multiple of 32") + if decoder_tile_overlap_pixels < 0 or decoder_tile_overlap_pixels % 32: + raise ValueError("decoder_tile_overlap_pixels must be a nonnegative multiple of 32") + if decoder_tile_pixels and decoder_tile_overlap_pixels >= decoder_tile_pixels: + raise ValueError("decoder overlap must be smaller than the tile") + if tail_hold_latents < 0: + raise ValueError("tail_hold_latents must be nonnegative") + if not denoising_step_list or any(not 0 < value <= 1 for value in denoising_step_list): + raise ValueError("denoising_step_list must contain values in (0, 1]") + if any(first <= second for first, second in zip(denoising_step_list, denoising_step_list[1:])): + raise ValueError("denoising_step_list must be strictly decreasing") + if not 1 <= num_steps <= len(denoising_step_list): + raise ValueError("num_steps must be between 1 and the distilled schedule length") + arch = PixelUMMArchitecture( + hidden_size, + intermediate_size, + num_layers, + num_heads, + num_kv_heads, + head_dim, + rms_norm_eps, + rope_theta_t, + rope_theta_hw, + ) + self.decoder = PixelUMMDecoderModel(arch, local_attention_tile=local_attention_tile_pixels // 16) + self.temporal_window_latents = temporal_window_latents + self.temporal_overlap_latents = temporal_overlap_latents + self.tail_hold_latents = tail_hold_latents + + def forward( + self, latent: torch.Tensor, sample: torch.Tensor, return_dict: bool = True + ) -> DecoderOutput | tuple[torch.Tensor]: + """ + Predict a clean RGB video from a noisy video and its latent conditioning. + + Args: + latent (`torch.Tensor`): + Conditioning with shape `(batch, 32, frames, height, width)`. + sample (`torch.Tensor`): + Noisy RGB pixels with shape `(batch, 3, frames * 4, height * 32, width * 32)`. + return_dict (`bool`, defaults to `True`): + Return a [`~models.autoencoders.vae.DecoderOutput`] instead of a one-element tuple. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + Predicted clean pixels with the same shape as `sample`. + """ + output = self.decoder(latent, sample) + return DecoderOutput(sample=output) if return_dict else (output,) + + def _make_noise(self, latent, start, end, seed): + batch, _, _, height, width = latent.shape + frames = [] + for index in range(start, end): + generator = torch.Generator(device=latent.device).manual_seed((seed * 1_000_003 + index) % (2**64)) + frames.append( + torch.randn( + (batch, 3, 4, height * 32, width * 32), + generator=generator, + device=latent.device, + dtype=torch.float32, + ) + ) + return torch.cat(frames, dim=2) + + def _noise_tile(self, latent, seed): + cached_range = None + cached_noise = None + + def source(first, last, row, column, height, width): + nonlocal cached_range, cached_noise + if cached_range != (first, last): + cached_noise = self._make_noise(latent, first, last, seed) + cached_range = (first, last) + return cached_noise[..., row * 32 : (row + height) * 32, column * 32 : (column + width) * 32] + + return source + + def _predict_x0(self, latent, xt_source, seed, t_starts=None, pad_short=False): + frames, height, width = latent.shape[2:] + window = self.config.temporal_window_latents + stride = window - self.config.temporal_overlap_latents + t_starts = _window_starts(frames, window, stride) if t_starts is None else t_starts + length = min(window, frames) + cells = self.config.decoder_tile_pixels // 32 + if cells <= 0 or (height <= cells and width <= cells): + rows, columns, tile_h, tile_w = [0], [0], height, width + else: + spatial_stride = cells - self.config.decoder_tile_overlap_pixels // 32 + tile_h, tile_w = min(cells, height), min(cells, width) + rows = _window_starts(height, tile_h, min(spatial_stride, tile_h)) + columns = _window_starts(width, tile_w, min(spatial_stride, tile_w)) + pad_noise = None + if pad_short and frames >= window and t_starts[-1] + window > frames: + variance, mean = torch.var_mean(latent.detach().float(), dim=(2, 3, 4), correction=0, keepdim=True) + generator = torch.Generator(device=latent.device).manual_seed((seed * 1_000_003 + frames - 1) % (2**64)) + pad_noise = mean + variance.sqrt() * torch.randn( + (latent.shape[0], latent.shape[1], window, height, width), + generator=generator, + device=latent.device, + dtype=torch.float32, + ) + blender = _TemporalBlender((t_starts[-1] + length - t_starts[0]) * 4) + for first in t_starts: + row_outputs = [] + for row in rows: + tile_outputs = [] + for column in columns: + real = latent[:, :, first : first + length, row : row + tile_h, column : column + tile_w] + if pad_noise is not None and real.shape[2] < length: + fill = length - real.shape[2] + offset = (first // stride) % (window - fill + 1) + tile = pad_noise[:, :, offset : offset + fill, row : row + tile_h, column : column + tile_w] + real = torch.cat([real, tile.to(real.dtype)], dim=2) + noise = xt_source(first, first + length, row, column, tile_h, tile_w) + tile_outputs.append(self.decoder(real.contiguous(), noise.contiguous())) + row_outputs.append(_merge_overlapping(tile_outputs, [column * 32 for column in columns], dim=-1)) + decoded = _merge_overlapping(row_outputs, [row * 32 for row in rows], dim=-2) + blender.add(decoded, (first - t_starts[0]) * 4) + return blender.result() + + @torch.no_grad() + @apply_forward_hook + def decode( + self, + latents: torch.Tensor, + generator: torch.Generator | None = None, + return_dict: bool = True, + pixel_frames: int | None = None, + ) -> DecoderOutput | tuple[torch.Tensor]: + """ + Decode a latent clip with overlapping temporal windows and optional spatial tiles. + + Args: + latents (`torch.Tensor`): + Latent video with shape `(batch, 32, frames, height, width)`. + generator (`torch.Generator`, *optional*): + Its initial seed selects the pixel noise stream. Each absolute latent-frame index has its own seed, so + overlapping windows receive identical noise. The generator state is not advanced. + return_dict (`bool`, defaults to `True`): + Return a [`~models.autoencoders.vae.DecoderOutput`] instead of a one-element tuple. + pixel_frames (`int`, *optional*): + Keep this many leading decoded frames; otherwise return `frames * 4` frames. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + RGB pixels with shape `(batch, 3, pixel_frames, height * 32, width * 32)`. + """ + if latents.ndim != 5 or latents.shape[1] != 32 or min(latents.shape[2:]) < 1: + raise ValueError("latents must have nonempty shape (batch, 32, frames, height, width)") + seed = self.config.seed if generator is None else generator.initial_seed() + steps = self.config.denoising_step_list[: self.config.num_steps] + if len(steps) == 1: + video = self._predict_x0(latents, self._noise_tile(latents, seed), seed) + else: + state = self._make_noise(latents, 0, latents.shape[2], seed) + + def source(first, last, row, column, height, width): + return state[ + :, :, first * 4 : last * 4, row * 32 : (row + height) * 32, column * 32 : (column + width) * 32 + ] + + for index, current in enumerate(steps): + video = self._predict_x0(latents, source, seed) + if index + 1 < len(steps): + ratio = steps[index + 1] / current + for offset in range(0, state.shape[2], 40): + anchor = video[:, :, offset : offset + 40].to(device=state.device, dtype=torch.float32) + state[:, :, offset : offset + 40] = anchor + ratio * ( + state[:, :, offset : offset + 40] - anchor + ) + if pixel_frames is not None: + if not 1 <= pixel_frames <= video.shape[2]: + raise ValueError("pixel_frames must be positive and no larger than the decoded frame count") + video = video[:, :, :pixel_frames] + return DecoderOutput(sample=video) if return_dict else (video,) + + @torch.no_grad() + @apply_forward_hook + def _decode_windows( + self, latent: torch.Tensor, starts: list[int], seed: int, pad_short: bool = False + ) -> torch.Tensor: + latent = latent.to(device=self.device, dtype=self.dtype) + return self._predict_x0(latent, self._noise_tile(latent, seed), seed, t_starts=starts, pad_short=pad_short) + + def create_streaming_decode_session( + self, *, to_cpu: bool = True, retain_history: bool = True, generator: torch.Generator | None = None + ) -> "PixelUMMStreamingDecodeSession": + """ + Create independent state for append-only decoding, tail correction, and rollback. + + Streaming requires `num_steps=1`. Set `session.final_segment=False` before a nonfinal push to enable padded + tail correction. Returned previews include their replacement range on the visible pixel timeline. `to_cpu` + controls where returned pixels are stored. `retain_history` enables `slice_pixels` and `full_video`. The + optional generator selects the same absolute-frame noise stream as [`~PixelUMMDecoder.decode`]. + """ + if self.config.num_steps != 1: + raise ValueError("Streaming PixelUMM decoding requires num_steps=1") + seed = self.config.seed if generator is None else generator.initial_seed() + return PixelUMMStreamingDecodeSession(self, to_cpu=to_cpu, retain_history=retain_history, seed=seed) + + +class PixelUMMStreamingDecodeSession: + """Independent latent timeline for stable previews, tail correction, and rollback.""" + + def __init__(self, model, *, to_cpu: bool = True, retain_history: bool = True, seed: int = 42): + self.model = model + self.to_cpu = bool(to_cpu) + self.retain_history = bool(retain_history) + self.window = model.temporal_window_latents + self.stride = self.window - model.temporal_overlap_latents + self.temporal_ratio = model.time_compression_ratio + self.buffer: torch.Tensor | None = None + self.available_latent_frames = 0 + self.available_pixel_frames = 0 + self.next_window_start = 0 + self.pending: torch.Tensor | None = None + self.pending_pixel_start = 0 + self.committed: list[tuple[int, torch.Tensor]] = [] + self.committed_raw_pixel_frames = 0 + self.committed_pixel_frames = 0 + self.push_latent_frames: list[int] = [] + self.push_pixel_frames: list[int] = [] + self.push_stable_pixel_frames: list[int] = [] + self.finished = False + self.seed = seed + self.final_segment = True + + def _append_decoded_windows(self, video: torch.Tensor, *, first_window_start: int) -> None: + clip_start = first_window_start * self.temporal_ratio + if self.pending is None: + self.pending = video + self.pending_pixel_start = clip_start + return + relative_start = clip_start - self.pending_pixel_start + self.pending = _merge_overlapping([self.pending, video], [0, relative_start], dim=2) + + @property + def available_raw_pixel_frames(self) -> int: + return self.available_latent_frames * self.temporal_ratio + + def _segment_raw_visible_spans(self): + raw_start = 0 + visible_start = 0 + for latent_frames, pixel_frames in zip(self.push_latent_frames, self.push_pixel_frames): + raw_end = raw_start + latent_frames * self.temporal_ratio + visible_end = visible_start + pixel_frames + yield (raw_start, raw_start + pixel_frames, raw_end, visible_start, visible_end) + raw_start = raw_end + visible_start = visible_end + + def _visible_frames_before_raw(self, raw_end: int) -> int: + raw_end = max(0, min(int(raw_end), self.available_raw_pixel_frames)) + visible = 0 + for ( + raw_start, + raw_visible_end, + _raw_segment_end, + _visible_start, + visible_end, + ) in self._segment_raw_visible_spans(): + if raw_end <= raw_start: + break + visible += max(0, min(raw_end, raw_visible_end) - raw_start) + if raw_end < raw_visible_end: + break + if raw_end >= raw_visible_end: + visible = visible_end + return visible + + def _visible_slice_from_raw(self, video: torch.Tensor, *, raw_start: int, raw_end: int) -> torch.Tensor: + """Drop per-segment decoder padding from one raw-pixel tensor.""" + raw_start = int(raw_start) + raw_end = int(raw_end) + expected = raw_end - raw_start + if expected < 0 or int(video.shape[2]) != expected: + raise RuntimeError( + f"PixelUMM raw preview range mismatch: range=[{raw_start}, {raw_end}), video={video.shape[2]}" + ) + parts = [] + for ( + segment_raw_start, + segment_visible_raw_end, + _segment_raw_end, + _segment_visible_start, + _segment_visible_end, + ) in self._segment_raw_visible_spans(): + overlap_start = max(raw_start, segment_raw_start) + overlap_end = min(raw_end, segment_visible_raw_end) + if overlap_end <= overlap_start: + continue + parts.append(video[:, :, overlap_start - raw_start : overlap_end - raw_start]) + if parts: + return torch.cat(parts, dim=2) + return video[:, :, :0] + + def _provisional_window_starts(self) -> list[int]: + if self.model.tail_hold_latents and (not self.final_segment): + length = self.available_latent_frames + return [0] if length <= self.window else list(range(0, length, self.stride)) + return _window_starts(self.available_latent_frames, self.window, self.stride) + + def _decode_provisional_raw(self) -> torch.Tensor | None: + """Decode the exact current tail from the raw stable boundary onward.""" + raw_start = self.committed_raw_pixel_frames + raw_end = self.available_raw_pixel_frames + if raw_end <= raw_start: + assert self.buffer is not None + if self.pending is not None: + return self.pending[:, :, :0] + return self.buffer.new_empty((0,)) + regular_decoded_latent_end = ( + self.next_window_start - self.stride + self.window if self.next_window_start > 0 else 0 + ) + regular_decoded_raw_end = regular_decoded_latent_end * self.temporal_ratio + if raw_end <= regular_decoded_raw_end: + if self.pending is not None: + pending_end = self.pending_pixel_start + int(self.pending.shape[2]) + if not (self.pending_pixel_start <= raw_start and pending_end >= raw_end): + raise RuntimeError( + f"PixelUMM pending output does not cover the regular decoded range [{raw_start}, {raw_end}); pending=[{self.pending_pixel_start}, {pending_end})" + ) + return self.pending[ + :, :, raw_start - self.pending_pixel_start : raw_end - self.pending_pixel_start + ].contiguous() + raise RuntimeError("PixelUMM session is missing pending regular-window output") + assert self.buffer is not None + starts = self._provisional_window_starts() + committed_latent = raw_start // self.temporal_ratio + first_index = next( + (index for index, start in enumerate(starts) if start + self.window > committed_latent), len(starts) - 1 + ) + selected_starts = starts[first_index:] + decoded = self._decode_windows(selected_starts) + assert decoded is not None + decoded_raw_start = selected_starts[0] * self.temporal_ratio + local_start = raw_start - decoded_raw_start + local_end = raw_end - decoded_raw_start + if not 0 <= local_start <= local_end <= decoded.shape[2]: + raise RuntimeError( + f"PixelUMM provisional decode does not cover the requested raw range: decoded_start={decoded_raw_start}, requested=[{raw_start}, {raw_end}), decoded_frames={decoded.shape[2]}" + ) + return decoded[:, :, local_start:local_end].contiguous() + + def _commit_until(self, stable_raw_until: int) -> torch.Tensor | None: + stable_raw_until = int(stable_raw_until) + if stable_raw_until <= self.committed_raw_pixel_frames: + return None + if self.pending is None: + raise RuntimeError("PixelUMM session has no decoded pixels to commit") + visible_start = self.committed_pixel_frames + stable_local_end = stable_raw_until - self.pending_pixel_start + if not 0 < stable_local_end <= self.pending.shape[2]: + raise RuntimeError( + f"Invalid PixelUMM stable-prefix boundary: local_end={stable_local_end}, pending={self.pending.shape[2]}" + ) + stable_raw = self.pending[:, :, :stable_local_end].contiguous() + stable = self._visible_slice_from_raw(stable_raw, raw_start=self.pending_pixel_start, raw_end=stable_raw_until) + if self.to_cpu: + stable = stable.cpu() + if self.retain_history and stable.shape[2] > 0: + self.committed.append((visible_start, stable)) + self.pending = self.pending[:, :, stable_local_end:].contiguous() + self.pending_pixel_start = stable_raw_until + self.committed_raw_pixel_frames = stable_raw_until + self.committed_pixel_frames = self._visible_frames_before_raw(stable_raw_until) + return stable + + def _decode_ready_regular_windows(self) -> torch.Tensor | None: + starts = [] + while self.next_window_start + self.window <= self.available_latent_frames: + starts.append(self.next_window_start) + self.next_window_start += self.stride + if not starts: + return None + decoded = self._decode_windows(starts) + assert decoded is not None + self._append_decoded_windows(decoded, first_window_start=starts[0]) + tail_start = self.available_latent_frames - self.window + stable_latent_frames = ( + tail_start + 1 if tail_start >= 0 and tail_start % self.stride == 0 else max(0, tail_start) + ) + return self._commit_until(stable_latent_frames * self.temporal_ratio) + + @property + def stable_pixel_frames(self) -> int: + return self.committed_pixel_frames + + @property + def withheld_pixel_frames(self) -> int: + return self.available_pixel_frames - self.committed_pixel_frames + + def slice_pixels(self, start: int, end: int) -> torch.Tensor: + """Read a future-stable pixel range from retained history.""" + if not self.retain_history: + raise RuntimeError("PixelUMM streaming history retention is disabled for this session.") + if not 0 <= start < end <= self.committed_pixel_frames: + raise ValueError( + f"Requested PixelUMM pixels are not stable yet: range=[{start}, {end}), stable={self.committed_pixel_frames}." + ) + selected = [] + cursor = start + for part_start, part in self.committed: + part_end = part_start + int(part.shape[2]) + overlap_start = max(start, part_start) + overlap_end = min(end, part_end) + if overlap_end <= overlap_start: + continue + if overlap_start != cursor: + raise RuntimeError( + f"PixelUMM streaming output has a pixel gap: cursor={cursor}, next={overlap_start}." + ) + selected.append(part[:, :, overlap_start - part_start : overlap_end - part_start]) + cursor = overlap_end + if cursor != end or not selected: + raise RuntimeError( + f"PixelUMM streaming output does not cover requested range: [{start}, {end}), covered_to={cursor}." + ) + return torch.cat(selected, dim=2) + + @torch.no_grad() + def push(self, latent: torch.Tensor, *, pixel_frames: int | None = None) -> torch.Tensor: + """Append one arbitrary-length segment and return newly stable pixels.""" + if self.finished: + raise RuntimeError("Cannot push PixelUMM latents after finish().") + if latent.ndim != 5 or latent.shape[2] <= 0: + raise ValueError(f"expected non-empty [B,C,T,H,W], got {tuple(latent.shape)}") + latent_frames = int(latent.shape[2]) + segment_pixel_frames = latent_frames * self.temporal_ratio if pixel_frames is None else int(pixel_frames) + if segment_pixel_frames <= 0: + raise ValueError("pixel_frames must be positive") + min_pixel_frames = (latent_frames - 1) * self.temporal_ratio + 1 + max_pixel_frames = latent_frames * self.temporal_ratio + if not min_pixel_frames <= segment_pixel_frames <= max_pixel_frames: + raise ValueError( + f"PixelUMM streaming segment pixel length is inconsistent with its latent length: latent_frames={latent_frames}, pixel_frames={segment_pixel_frames}, expected_range=[{min_pixel_frames}, {max_pixel_frames}]." + ) + previous_stable = self.committed_pixel_frames + self.buffer = latent.contiguous() if self.buffer is None else torch.cat([self.buffer, latent], dim=2) + self.available_latent_frames += latent_frames + self.available_pixel_frames += segment_pixel_frames + self.push_latent_frames.append(latent_frames) + self.push_pixel_frames.append(segment_pixel_frames) + stable = self._decode_ready_regular_windows() + self.push_stable_pixel_frames.append(self.committed_pixel_frames) + if stable is None or self.committed_pixel_frames <= previous_stable: + assert self.buffer is not None + return self.buffer.new_empty((0,)) + return stable + + @torch.no_grad() + def push_preview(self, latent: torch.Tensor, *, pixel_frames: int | None = None) -> StreamingPreviewUpdate: + """Append a segment and publish its complete provisional preview. + + `push()` publishes only future-stable pixels. Interactive web playback instead needs the complete segment + immediately. This method additionally exposes the current provisional overlap. On the next push, the preceding + provisional tail is merged with the newly decoded cross-boundary PixelUMM window and returned again for + in-place replacement. + + Segment boundaries do not need to align to the PixelUMM stride. Callers copy the returned tensor over the + preceding provisional tail, then append the remaining new frames. + """ + previous_available = int(self.available_pixel_frames) + previous_stable = int(self.committed_pixel_frames) + stable = self.push(latent, pixel_frames=pixel_frames) + replace_pixel_frames = previous_available - previous_stable + pixel_start = previous_stable + pixel_end = int(self.available_pixel_frames) + provisional_pixel_frames = pixel_end - int(self.committed_pixel_frames) + provisional_raw = None + if provisional_pixel_frames > 0: + provisional_raw = self._decode_provisional_raw() + parts = [] + if stable.numel() > 0: + parts.append(stable) + if provisional_pixel_frames > 0: + assert provisional_raw is not None + pending = self._visible_slice_from_raw( + provisional_raw, raw_start=self.committed_raw_pixel_frames, raw_end=self.available_raw_pixel_frames + ) + if self.to_cpu: + pending = pending.cpu() + parts.append(pending) + if not parts: + assert self.buffer is not None + video = self.buffer.new_empty((0,)) + elif len(parts) == 1: + video = parts[0] + else: + video = torch.cat(parts, dim=2) + expected = pixel_end - pixel_start + if int(video.shape[2]) != expected: + raise RuntimeError( + f"PixelUMM preview update has the wrong pixel range: start={pixel_start}, end={pixel_end}, video={video.shape[2]}" + ) + return StreamingPreviewUpdate( + video=video, + pixel_start=pixel_start, + pixel_end=pixel_end, + replace_pixel_frames=replace_pixel_frames, + stable_pixel_frames=int(self.committed_pixel_frames), + provisional_pixel_frames=provisional_pixel_frames, + ) + + @torch.no_grad() + def finish(self) -> torch.Tensor: + """Publish the withheld tail and close the append-only timeline.""" + if self.finished: + raise RuntimeError("PixelUMM streaming decode is already finished.") + if self.buffer is None: + raise RuntimeError("Cannot finish an empty PixelUMM streaming session.") + previous_stable = self.committed_pixel_frames + provisional_raw = self._decode_provisional_raw() + assert provisional_raw is not None + stable = self._visible_slice_from_raw( + provisional_raw, raw_start=self.committed_raw_pixel_frames, raw_end=self.available_raw_pixel_frames + ) + if self.to_cpu: + stable = stable.cpu() + if self.retain_history and stable.shape[2] > 0: + self.committed.append((previous_stable, stable)) + if self.pending is not None: + self.pending = self.pending[:, :, :0] + self.pending_pixel_start = self.available_raw_pixel_frames + self.committed_raw_pixel_frames = self.available_raw_pixel_frames + self.committed_pixel_frames = self.available_pixel_frames + self.finished = True + if self.committed_pixel_frames <= previous_stable: + assert self.buffer is not None + return self.buffer.new_empty((0,)) + return stable + + @torch.no_grad() + def rollback_last_push(self) -> None: + """Discard the latest segment and rebuild the exact stable prefix.""" + if self.finished: + raise RuntimeError("Cannot roll back a finalized PixelUMM session.") + if not self.push_latent_frames: + raise RuntimeError("There is no PixelUMM streaming segment to roll back.") + latent_frames = self.push_latent_frames.pop() + self.push_pixel_frames.pop() + self.push_stable_pixel_frames.pop() + assert self.buffer is not None + keep = self.available_latent_frames - latent_frames + retained = self.buffer[:, :, :keep].contiguous() if keep > 0 else None + retained_latent_frames = list(self.push_latent_frames) + retained_pixel_frames = list(self.push_pixel_frames) + retained_stable_pixel_frames = list(self.push_stable_pixel_frames) + self.buffer = retained + self.available_latent_frames = keep + self.available_pixel_frames = sum(retained_pixel_frames) + self.next_window_start = 0 + self.pending = None + self.pending_pixel_start = 0 + self.committed = [] + self.committed_raw_pixel_frames = 0 + self.committed_pixel_frames = 0 + self.push_latent_frames = retained_latent_frames + self.push_pixel_frames = retained_pixel_frames + self.push_stable_pixel_frames = retained_stable_pixel_frames + if retained is not None: + self._decode_ready_regular_windows() + if self.push_stable_pixel_frames and self.push_stable_pixel_frames[-1] != self.committed_pixel_frames: + raise RuntimeError( + f"Rebuilt PixelUMM stable prefix differs after rollback: expected={self.push_stable_pixel_frames[-1]}, rebuilt={self.committed_pixel_frames}." + ) + + def full_video(self) -> torch.Tensor: + """Return the finalized full video from retained history.""" + if not self.finished: + raise RuntimeError("PixelUMM streaming decode is not finished") + if not self.retain_history: + raise RuntimeError("PixelUMM streaming history retention is disabled for this session.") + return self.slice_pixels(0, self.available_pixel_frames) + + def _decode_windows(self, starts): + video = self.model._decode_windows( + self.buffer, starts, self.seed, pad_short=bool(self.model.tail_hold_latents) and (not self.final_segment) + ) + return video.cpu() if self.to_cpu else video diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 469906613a5e..67635d9dcf25 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1515,6 +1515,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None "device_map": device_map, "offload_dir": offload_folder, "offload_index": offload_index, + "skip_keys": model._skip_keys, } dispatch_model(model, **device_map_kwargs) diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index dd7014096e54..d02e2c80a360 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -43,6 +43,7 @@ from .transformer_hunyuan_video15 import HunyuanVideo15Transformer3DModel from .transformer_hunyuan_video_framepack import HunyuanVideoFramepackTransformer3DModel from .transformer_hunyuanimage import HunyuanImageTransformer2DModel + from .transformer_hy_nomas import HyNomasTransformer3DModel from .transformer_ideogram4 import Ideogram4Transformer2DModel from .transformer_joyimage import JoyImageEditTransformer3DModel from .transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel diff --git a/src/diffusers/models/transformers/transformer_hy_nomas.py b/src/diffusers/models/transformers/transformer_hy_nomas.py new file mode 100644 index 000000000000..b73d7f087dff --- /dev/null +++ b/src/diffusers/models/transformers/transformer_hy_nomas.py @@ -0,0 +1,908 @@ +# Copyright 2026 The Hy-Nomas Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import AttentionBackendName, _check_attention_backend_requirements, dispatch_attention_fn +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import RMSNorm + + +def _apply_rotary_emb(hidden_states, rotary_emb): + cosine, sine = rotary_emb + real, imaginary = hidden_states.float().unflatten(-1, (-1, 2)).unbind(-1) + rotated = torch.stack((-imaginary, real), dim=-1).flatten(-2) + return (hidden_states.float() * cosine[None, :, None] + rotated * sine[None, :, None]).to(hidden_states.dtype) + + +class HyNomasLayerNorm(nn.LayerNorm): + def forward(self, hidden_states) -> torch.Tensor: + # The released model keeps normalization and subsequent modulation in float32. + return F.layer_norm( + hidden_states.float(), + self.normalized_shape, + None if self.weight is None else self.weight.float(), + None if self.bias is None else self.bias.float(), + self.eps, + ) + + +class HyNomasRotaryPosEmbed(nn.Module): + def __init__(self, axes_dim, theta): + super().__init__() + self.axes_dim = axes_dim + self.theta = theta + + def forward(self, num_frames, height, width, device, frame_offset=0) -> tuple[torch.Tensor, torch.Tensor]: + # Continuation positions in the released model use CPU trigonometry. Its rounding matters before + # casting the rotated query and key back to the model dtype. + computation_device = "cpu" if frame_offset else device + positions = torch.meshgrid( + torch.arange(frame_offset, frame_offset + num_frames, device=computation_device, dtype=torch.float32), + torch.arange(height, device=computation_device, dtype=torch.float32), + torch.arange(width, device=computation_device, dtype=torch.float32), + indexing="ij", + ) + cosines, sines = [], [] + for dim, position in zip(self.axes_dim, positions): + frequencies = 1.0 / (self.theta ** (torch.arange(0, dim, 2, device=computation_device).float() / dim)) + angles = torch.outer(position.flatten(), frequencies) + cosines.append(angles.cos().repeat_interleave(2, dim=-1)) + sines.append(angles.sin().repeat_interleave(2, dim=-1)) + return torch.cat(cosines, dim=-1).to(device), torch.cat(sines, dim=-1).to(device) + + +class HyNomasTimestepEmbedder(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.mlp = nn.Sequential(nn.Linear(256, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size)) + # CPU frequencies preserve the released model's rounding and are never cast with model weights. + self.frequencies = tuple( + torch.exp(-math.log(10000) * torch.arange(128, dtype=torch.float32, device="cpu") / 128).tolist() + ) + + def forward(self, timestep, dtype) -> torch.Tensor: + frequencies = torch.tensor(self.frequencies, dtype=torch.float32, device=timestep.device) + angles = timestep[:, None].float() * frequencies[None] + hidden_states = torch.cat((angles.cos(), angles.sin()), dim=-1) + return self.mlp(hidden_states.to(dtype)) + + +class HyNomasTextProjection(nn.Module): + def __init__(self, in_channels, hidden_size): + super().__init__() + self.linear_1 = nn.Linear(in_channels, hidden_size) + self.linear_2 = nn.Linear(hidden_size, hidden_size) + + def forward(self, hidden_states) -> torch.Tensor: + hidden_states = self.linear_1(hidden_states) + hidden_states = F.silu(hidden_states) + return self.linear_2(hidden_states) + + +class HyNomasMLP(nn.Module): + def __init__(self, hidden_size, mlp_ratio=4.0, act_fn="gelu"): + super().__init__() + self.fc1 = nn.Linear(hidden_size, int(hidden_size * mlp_ratio)) + self.act = nn.SiLU() if act_fn == "silu" else nn.GELU(approximate="tanh") + self.fc2 = nn.Linear(int(hidden_size * mlp_ratio), hidden_size) + + def forward(self, hidden_states) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.act(hidden_states) + return self.fc2(hidden_states) + + +class HyNomasTextAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__(self, attn, hidden_states, attention_mask=None): + query, key, value = attn.self_attn_qkv(hidden_states).unflatten(-1, (3, attn.heads, -1)).unbind(-3) + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + return attn.self_attn_proj(hidden_states.flatten(-2)) + + +def _hy_nomas_unpadded_attention(query, key, value, query_mask, key_mask, backend, deterministic=False): + outputs = torch.zeros_like(query) + for sample in range(query.shape[0]): + query_indices = torch.nonzero(query_mask[sample], as_tuple=True)[0] + key_indices = torch.nonzero(key_mask[sample], as_tuple=True)[0] + output = dispatch_attention_fn( + query[sample : sample + 1, query_indices], + key[sample : sample + 1, key_indices], + value[sample : sample + 1, key_indices], + backend=backend, + attention_kwargs={"deterministic": deterministic}, + ) + outputs[sample, query_indices] = output[0] + return outputs + + +class HyNomasFlashTextAttnProcessor: + """FlashAttention with the reference model's query and key padding removal.""" + + _attention_backend = "flash_varlen" + _parallel_config = None + + def __init__(self, backend="flash_varlen", deterministic=False): + if backend not in ("flash_varlen", "_flash_varlen_3"): + raise ValueError("Reference text attention requires flash_varlen or _flash_varlen_3.") + self._attention_backend = backend + self.deterministic = deterministic + + def __call__(self, attn, hidden_states, attention_mask=None): + query, key, value = attn.self_attn_qkv(hidden_states).unflatten(-1, (3, attn.heads, -1)).unbind(-3) + mask = attention_mask + if mask is None: + mask = torch.ones(query.shape[:2], dtype=torch.bool, device=query.device) + else: + mask = mask.reshape(query.shape[:2]) + output = _hy_nomas_unpadded_attention( + query, key, value, mask, mask, self._attention_backend, self.deterministic + ) + return attn.self_attn_proj(output.flatten(-2)) + + +class HyNomasTextAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = HyNomasTextAttnProcessor + _available_processors = [HyNomasTextAttnProcessor, HyNomasFlashTextAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, hidden_size, heads): + super().__init__() + self.heads = heads + self.self_attn_qkv = nn.Linear(hidden_size, 3 * hidden_size) + self.self_attn_proj = nn.Linear(hidden_size, hidden_size) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states, attention_mask=None) -> torch.Tensor: + return self.processor(self, hidden_states, attention_mask) + + +class HyNomasTokenRefinerBlock(nn.Module): + def __init__(self, hidden_size, heads): + super().__init__() + self.norm1 = HyNomasLayerNorm(hidden_size, eps=1e-6) + self.attn = HyNomasTextAttention(hidden_size, heads) + self.norm2 = HyNomasLayerNorm(hidden_size, eps=1e-6) + self.mlp = HyNomasMLP(hidden_size, act_fn="silu") + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size)) + + def forward(self, hidden_states, temb, attention_mask) -> torch.Tensor: + attention_gate, mlp_gate = self.adaLN_modulation(temb).chunk(2, dim=-1) + hidden_states = ( + hidden_states + + self.attn(self.norm1(hidden_states).to(hidden_states.dtype), attention_mask) * attention_gate[:, None] + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states).to(hidden_states.dtype)) * mlp_gate[:, None] + return hidden_states + + +class HyNomasIndividualTokenRefiner(nn.Module): + def __init__(self, hidden_size, heads): + super().__init__() + self.blocks = nn.ModuleList([HyNomasTokenRefinerBlock(hidden_size, heads) for _ in range(2)]) + + def forward(self, hidden_states, temb, attention_mask) -> torch.Tensor: + if attention_mask is not None: + attention_mask = attention_mask.clone().bool() + attention_mask[:, 0] = True + attention_mask = attention_mask[:, None, None, :] + for block in self.blocks: + hidden_states = block(hidden_states, temb, attention_mask) + return hidden_states + + +class HyNomasTokenRefiner(nn.Module): + def __init__(self, in_channels, hidden_size, heads): + super().__init__() + self.input_embedder = nn.Linear(in_channels, hidden_size) + self.t_embedder = HyNomasTimestepEmbedder(hidden_size) + self.c_embedder = HyNomasTextProjection(in_channels, hidden_size) + self.individual_token_refiner = HyNomasIndividualTokenRefiner(hidden_size, heads) + + def forward(self, hidden_states, timestep, attention_mask=None) -> torch.Tensor: + temb = self.t_embedder(timestep, hidden_states.dtype) + if attention_mask is None: + context = hidden_states.mean(dim=1) + else: + mask = attention_mask.to(hidden_states.dtype).unsqueeze(-1) + context = (hidden_states * mask).sum(dim=1, dtype=torch.float32) / mask.sum( + dim=1, dtype=torch.float32 + ).clamp_min(1) + temb = temb + self.c_embedder(context.to(hidden_states.dtype)) + hidden_states = self.input_embedder(hidden_states) + return self.individual_token_refiner(hidden_states, temb, attention_mask) + + +class HyNomasAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn, + hidden_states, + encoder_hidden_states, + rotary_emb, + encoder_attention_mask=None, + kv_cache=None, + update_cache=False, + action_text_offsets=None, + action_token_windows=None, + cache_action_start=None, + ): + query = attn.img_attn_q(hidden_states).unflatten(-1, (attn.heads, -1)) + key = attn.img_attn_k(hidden_states).unflatten(-1, (attn.heads, -1)) + value = attn.img_attn_v(hidden_states).unflatten(-1, (attn.heads, -1)) + query = _apply_rotary_emb(attn.img_attn_q_norm(query).to(value.dtype), rotary_emb) + key = _apply_rotary_emb(attn.img_attn_k_norm(key).to(value.dtype), rotary_emb) + encoder_query = attn.txt_attn_q(encoder_hidden_states).unflatten(-1, (attn.heads, -1)) + encoder_key = attn.txt_attn_k(encoder_hidden_states).unflatten(-1, (attn.heads, -1)) + encoder_value = attn.txt_attn_v(encoder_hidden_states).unflatten(-1, (attn.heads, -1)) + encoder_query = attn.txt_attn_q_norm(encoder_query).to(encoder_value.dtype) + encoder_key = attn.txt_attn_k_norm(encoder_key).to(encoder_value.dtype) + query = torch.cat((query, encoder_query), dim=1) + current_key = torch.cat((key, encoder_key), dim=1) + current_value = torch.cat((value, encoder_value), dim=1) + batch_size, video_length = hidden_states.shape[:2] + text_length = encoder_hidden_states.shape[1] + text_mask = encoder_attention_mask + if text_mask is None: + text_mask = torch.ones(batch_size, text_length, dtype=torch.bool, device=query.device) + current_mask = torch.cat( + (torch.ones(batch_size, video_length, dtype=torch.bool, device=query.device), text_mask.bool()), dim=1 + ) + history_length = 0 + key, value, key_mask = current_key, current_value, current_mask + if kv_cache is not None and "key" in kv_cache: + history_length = kv_cache["key"].shape[1] + key = torch.cat((kv_cache["key"], current_key), dim=1) + value = torch.cat((kv_cache["value"], current_value), dim=1) + key_mask = torch.cat((kv_cache["mask"], current_mask), dim=1) + + output = self._compute_attention( + query, + key, + value, + key_mask, + text_mask, + history_length, + video_length, + text_length, + encoder_attention_mask, + action_text_offsets, + action_token_windows, + ) + output = output * current_mask[:, :, None, None].to(output.dtype) + + if update_cache: + cache_mask = current_mask.clone() + if cache_action_start is not None: + cache_mask[:, video_length + cache_action_start :] = False + elif action_text_offsets is not None: + for batch_index, offsets in enumerate(action_text_offsets): + if offsets is not None: + cache_mask[batch_index, video_length + offsets[0] :] = False + if history_length: + cache_mask = torch.cat((kv_cache["mask"], cache_mask), dim=1) + kv_cache.update(key=key.detach(), value=value.detach(), mask=cache_mask) + + output = output.flatten(-2) + hidden_states, encoder_hidden_states = output.split((video_length, text_length), dim=1) + return attn.img_attn_proj(hidden_states), attn.txt_attn_proj(encoder_hidden_states) + + def _compute_attention( + self, + query, + key, + value, + key_mask, + text_mask, + history_length, + video_length, + text_length, + encoder_attention_mask, + action_text_offsets, + action_token_windows, + ): + if action_text_offsets is None: + attention_mask = ( + key_mask[:, None, None, :] if encoder_attention_mask is not None or history_length else None + ) + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + else: + # Each query group shares a visible key set; no quadratic dense mask is needed. + outputs = [] + for batch_index, offsets in enumerate(action_text_offsets): + sample_output = torch.zeros_like(query[batch_index : batch_index + 1]) + sample_text_mask = text_mask[batch_index].bool() + text_positions = sample_text_mask.cumsum(dim=0) - 1 + valid_text = sample_text_mask.sum() + if offsets is None: + groups = [(0, video_length + text_length, None)] + else: + groups = list(action_token_windows) + groups.append((video_length, video_length + min(offsets[0], text_length), -1)) + groups.extend( + (video_length + min(start, text_length), video_length + min(end, text_length), index) + for index, (start, end) in enumerate(zip(offsets[:-1], offsets[1:])) + ) + for start, end, action_index in groups: + if end <= start: + continue + visible_text = sample_text_mask + if offsets is not None: + visible_text = text_positions < offsets[0] + if action_index >= 0: + visible_text = visible_text | ( + (text_positions >= offsets[action_index]) + & (text_positions < offsets[action_index + 1]) + ) + visible_text = visible_text & sample_text_mask + attention_mask = torch.cat((key_mask[batch_index, : history_length + video_length], visible_text))[ + None, None, None, : + ] + group_output = dispatch_attention_fn( + query[batch_index : batch_index + 1, start:end], + key[batch_index : batch_index + 1], + value[batch_index : batch_index + 1], + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + if offsets is not None: + query_mask = torch.arange(start, end, device=query.device) < video_length + valid_text + group_output = group_output * query_mask[None, :, None, None].to(group_output.dtype) + sample_output[:, start:end] = group_output + outputs.append(sample_output) + output = torch.cat(outputs, dim=0) + return output + + +class HyNomasMagiAttnProcessor(HyNomasAttnProcessor): + """Reference query/key packing and rectangle ordering for MagiAttention.""" + + def __init__(self, deterministic=False, flash_deterministic=False): + self.deterministic = deterministic + self.flash_deterministic = flash_deterministic + + def _compute_attention( + self, + query, + key, + value, + key_mask, + text_mask, + history_length, + video_length, + text_length, + encoder_attention_mask, + action_text_offsets, + action_token_windows, + ): + query_mask = key_mask[:, history_length:] + if action_text_offsets is None or all(offsets is None for offsets in action_text_offsets): + return _hy_nomas_unpadded_attention( + query, key, value, query_mask, key_mask, "_flash_varlen_3", self.flash_deterministic + ) + output = torch.zeros_like(query) + for sample, offsets in enumerate(action_text_offsets): + history_indices = torch.nonzero(key_mask[sample, :history_length], as_tuple=True)[0] + text_indices = torch.nonzero(text_mask[sample], as_tuple=True)[0] + query_indices = torch.cat((torch.arange(video_length, device=query.device), video_length + text_indices)) + key_indices = torch.cat((history_indices, history_length + query_indices)) + visible_history = history_indices.numel() + valid_text = text_indices.numel() + query_ranges, key_ranges = [], [] + if offsets is None: + query_ranges.append([0, query_indices.numel()]) + key_ranges.append([0, key_indices.numel()]) + else: + offsets = [min(offset, valid_text) for offset in offsets] + groups = [ + (max(0, min(start, video_length)), max(0, min(end, video_length)), action) + for start, end, action in action_token_windows + ] + groups.append((video_length, video_length + offsets[0], -1)) + groups.extend( + (video_length + start, video_length + end, action) + for action, (start, end) in enumerate(zip(offsets[:-1], offsets[1:])) + ) + for start, end, action in groups: + if end <= start: + continue + ranges = [] + if visible_history: + ranges.append([0, visible_history]) + ranges.append([visible_history, visible_history + video_length]) + text_start = visible_history + video_length + if offsets[0]: + ranges.append([text_start, text_start + offsets[0]]) + if 0 <= action < len(offsets) - 1 and offsets[action + 1] > offsets[action]: + ranges.append([text_start + offsets[action], text_start + offsets[action + 1]]) + query_ranges.extend([[start, end]] * len(ranges)) + key_ranges.extend(ranges) + attention_kwargs = { + "query_ranges": torch.tensor(query_ranges, dtype=torch.int32, device=query.device), + "key_ranges": torch.tensor(key_ranges, dtype=torch.int32, device=query.device), + "attn_type_map": torch.zeros(len(query_ranges), dtype=torch.int32, device=query.device), + "deterministic": self.deterministic, + } + sample_output = dispatch_attention_fn( + query[sample : sample + 1, query_indices], + key[sample : sample + 1, key_indices], + value[sample : sample + 1, key_indices], + backend="magi", + attention_kwargs=attention_kwargs, + ) + output[sample, query_indices] = sample_output[0] + return output + + +class HyNomasAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = HyNomasAttnProcessor + _available_processors = [HyNomasAttnProcessor, HyNomasMagiAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, hidden_size, heads): + super().__init__() + self.heads = heads + for prefix in ("img", "txt"): + for projection in ("q", "k", "v", "proj"): + setattr(self, f"{prefix}_attn_{projection}", nn.Linear(hidden_size, hidden_size)) + setattr(self, f"{prefix}_attn_q_norm", RMSNorm(hidden_size // heads, eps=1e-6)) + setattr(self, f"{prefix}_attn_k_norm", RMSNorm(hidden_size // heads, eps=1e-6)) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states, encoder_hidden_states, rotary_emb, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: + return self.processor(self, hidden_states, encoder_hidden_states, rotary_emb, **kwargs) + + +class HyNomasModulation(nn.Module): + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size * 6) + + def forward(self, hidden_states) -> torch.Tensor: + return self.linear(F.silu(hidden_states)) + + +class HyNomasTransformerBlock(nn.Module): + def __init__(self, hidden_size, heads, mlp_ratio): + super().__init__() + self.img_mod = HyNomasModulation(hidden_size) + self.img_norm1 = HyNomasLayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) + self.img_norm2 = HyNomasLayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) + self.img_mlp = HyNomasMLP(hidden_size, mlp_ratio) + self.txt_mod = HyNomasModulation(hidden_size) + self.txt_norm1 = HyNomasLayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) + self.txt_norm2 = HyNomasLayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) + self.txt_mlp = HyNomasMLP(hidden_size, mlp_ratio) + self.attn = HyNomasAttention(hidden_size, heads) + + def forward( + self, hidden_states, encoder_hidden_states, temb, rotary_emb, attention_kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + shift, scale, gate, mlp_shift, mlp_scale, mlp_gate = self.img_mod(temb).chunk(6, dim=-1) + text_shift, text_scale, text_gate, text_mlp_shift, text_mlp_scale, text_mlp_gate = self.txt_mod(temb).chunk( + 6, dim=-1 + ) + norm_hidden_states = self.img_norm1(hidden_states) * (1 + scale[:, None]) + shift[:, None] + norm_encoder_hidden_states = ( + self.txt_norm1(encoder_hidden_states) * (1 + text_scale[:, None]) + text_shift[:, None] + ) + attention_output, encoder_attention_output = self.attn( + norm_hidden_states.to(hidden_states.dtype), + norm_encoder_hidden_states.to(encoder_hidden_states.dtype), + rotary_emb, + **attention_kwargs, + ) + hidden_states = hidden_states + attention_output * gate[:, None] + hidden_states = ( + hidden_states + + self.img_mlp( + (self.img_norm2(hidden_states) * (1 + mlp_scale[:, None]) + mlp_shift[:, None]).to(hidden_states.dtype) + ) + * mlp_gate[:, None] + ) + encoder_hidden_states = encoder_hidden_states + encoder_attention_output * text_gate[:, None] + encoder_hidden_states = ( + encoder_hidden_states + + self.txt_mlp( + (self.txt_norm2(encoder_hidden_states) * (1 + text_mlp_scale[:, None]) + text_mlp_shift[:, None]).to( + encoder_hidden_states.dtype + ) + ) + * text_mlp_gate[:, None] + ) + return hidden_states, encoder_hidden_states + + +class HyNomasCameraAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__(self, attn, hidden_states, camera_hidden_states, rotary_emb, camera_attention_mask=None): + batch_size, num_frames = camera_hidden_states.shape[:2] + query = attn.img_attn_q(hidden_states).unflatten(-1, (attn.heads, -1)) + camera_hidden_states = attn.keyboard_embed(camera_hidden_states) + key, value = attn.keyboard_attn_kv(camera_hidden_states).unflatten(-1, (2, attn.heads, -1)).unbind(-3) + query = attn.key_attn_q_norm(query).to(value.dtype) + key = attn.key_attn_k_norm(key).to(value.dtype) + tokens_per_frame = hidden_states.shape[1] // num_frames + query_rotary_emb = tuple(component.repeat_interleave(tokens_per_frame, dim=0) for component in rotary_emb) + query = _apply_rotary_emb(query, query_rotary_emb) + key = _apply_rotary_emb(key, rotary_emb) + output = self._compute_attention(query, key, value, num_frames, tokens_per_frame) + output = output.flatten(-2).reshape(batch_size, -1, attn.camera_attention_dim) + output = attn.proj_keyboard(output) + if camera_attention_mask is not None: + output = output * camera_attention_mask.reshape(batch_size, 1, 1).to(output.dtype) + return hidden_states + output + + def _compute_attention(self, query, key, value, num_frames, tokens_per_frame): + query = query.unflatten(1, (num_frames, tokens_per_frame)).flatten(0, 1) + key = key.flatten(0, 1).unsqueeze(1) + value = value.flatten(0, 1).unsqueeze(1) + output = dispatch_attention_fn( + query, key, value, backend=self._attention_backend, parallel_config=self._parallel_config + ) + return output + + +class HyNomasMagiCameraAttnProcessor(HyNomasCameraAttnProcessor): + """Reference frame-local camera range attention.""" + + def _compute_attention(self, query, key, value, num_frames, tokens_per_frame): + if num_frames == 1: + return dispatch_attention_fn(query, key, value, backend="flash") + starts = torch.arange(num_frames, dtype=torch.int32, device=query.device) + frame_ranges = torch.stack((starts, starts + 1), dim=-1) + return dispatch_attention_fn( + query, + key, + value, + backend="magi", + attention_kwargs={ + "query_ranges": frame_ranges * tokens_per_frame, + "key_ranges": frame_ranges, + "attn_type_map": torch.zeros(num_frames, dtype=torch.int32, device=query.device), + }, + ) + + +class HyNomasCameraAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = HyNomasCameraAttnProcessor + _available_processors = [HyNomasCameraAttnProcessor, HyNomasMagiCameraAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, hidden_size, camera_embed_dim, camera_attention_dim, heads): + super().__init__() + self.heads = heads + self.camera_attention_dim = camera_attention_dim + self.keyboard_embed = nn.Sequential( + nn.Linear(64, camera_embed_dim), nn.SiLU(), nn.Linear(camera_embed_dim, camera_embed_dim) + ) + self.key_attn_q_norm = RMSNorm(camera_attention_dim // heads, eps=1e-6) + self.key_attn_k_norm = RMSNorm(camera_attention_dim // heads, eps=1e-6) + self.img_attn_q = nn.Linear(hidden_size, camera_attention_dim, bias=False) + self.keyboard_attn_kv = nn.Linear(camera_embed_dim, 2 * camera_attention_dim, bias=False) + self.proj_keyboard = nn.Linear(camera_attention_dim, hidden_size, bias=False) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states, camera_hidden_states, rotary_emb, camera_attention_mask=None) -> torch.Tensor: + return self.processor(self, hidden_states, camera_hidden_states, rotary_emb, camera_attention_mask) + + +class HyNomasPatchEmbed(nn.Module): + def __init__(self, in_channels, hidden_size): + super().__init__() + self.proj = nn.Conv3d(in_channels, hidden_size, kernel_size=1, stride=1) + + def forward(self, hidden_states) -> torch.Tensor: + return self.proj(hidden_states).flatten(2).transpose(1, 2).contiguous() + + +class HyNomasFinalLayer(nn.Module): + def __init__(self, hidden_size, out_channels): + super().__init__() + self.norm_final = HyNomasLayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) + self.linear = nn.Linear(hidden_size, out_channels) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size)) + + def forward(self, hidden_states, temb) -> torch.Tensor: + dtype = hidden_states.dtype + shift, scale = self.adaLN_modulation(temb).chunk(2, dim=-1) + hidden_states = self.norm_final(hidden_states) * (1 + scale[:, None]) + shift[:, None] + return self.linear(hidden_states.to(dtype)) + + +class HyNomasTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin): + r"""The Hy-Nomas segment autoregressive video transformer with frame-local camera conditioning. + + Args: + in_channels (`int`, defaults to `32`): Number of latent channels. + out_channels (`int`, defaults to `32`): Number of output channels. + num_layers (`int`, defaults to `54`): Number of double stream blocks. + num_attention_heads (`int`, defaults to `16`): Number of attention heads. + attention_head_dim (`int`, defaults to `128`): Dimension of an attention head. + text_embed_dim (`int`, defaults to `3584`): Dimension of text encoder features. + mlp_ratio (`float`, defaults to `4.0`): MLP expansion ratio. + rope_axes_dim (`tuple[int]`, defaults to `(16, 56, 56)`): Temporal and spatial RoPE dimensions. + rope_theta (`float`, defaults to `256.0`): RoPE frequency base. + camera_embed_dim (`int`, defaults to `128`): Camera embedding dimension. + camera_attention_dim (`int`, defaults to `1024`): Camera attention projection dimension. + camera_num_attention_heads (`int`, defaults to `16`): Number of camera attention heads. + camera_rope_axes_dim (`tuple[int]`, defaults to `(8, 28, 28)`): Camera RoPE dimensions. + """ + + _supports_gradient_checkpointing = True + _no_split_modules = [ + "HyNomasTransformerBlock", + "HyNomasTokenRefiner", + "HyNomasCameraAttention", + "HyNomasFinalLayer", + ] + _repeated_blocks = ["HyNomasTransformerBlock"] + _skip_layerwise_casting_patterns = ["img_in", "txt_in", "time_in", "norm"] + _skip_keys = ["kv_cache"] + + @register_to_config + def __init__( + self, + in_channels: int = 32, + out_channels: int = 32, + num_layers: int = 54, + num_attention_heads: int = 16, + attention_head_dim: int = 128, + text_embed_dim: int = 3584, + mlp_ratio: float = 4.0, + rope_axes_dim: tuple[int, int, int] = (16, 56, 56), + rope_theta: float = 256.0, + camera_embed_dim: int = 128, + camera_attention_dim: int = 1024, + camera_num_attention_heads: int = 16, + camera_rope_axes_dim: tuple[int, int, int] = (8, 28, 28), + ): + super().__init__() + if sum(rope_axes_dim) != attention_head_dim or any(dim % 2 for dim in rope_axes_dim): + raise ValueError("Even rope_axes_dim entries must sum to attention_head_dim.") + if camera_attention_dim != camera_num_attention_heads * sum(camera_rope_axes_dim): + raise ValueError("camera_attention_dim must match camera_num_attention_heads * sum(camera_rope_axes_dim).") + hidden_size = num_attention_heads * attention_head_dim + self.img_in = HyNomasPatchEmbed(in_channels, hidden_size) + self.txt_in = HyNomasTokenRefiner(text_embed_dim, hidden_size, num_attention_heads) + self.time_in = HyNomasTimestepEmbedder(hidden_size) + self.rotary_embed = HyNomasRotaryPosEmbed(rope_axes_dim, rope_theta) + self.camera_rotary_embed = HyNomasRotaryPosEmbed(camera_rope_axes_dim, rope_theta) + self.double_blocks = nn.ModuleList( + [HyNomasTransformerBlock(hidden_size, num_attention_heads, mlp_ratio) for _ in range(num_layers)] + ) + self.action_modules = nn.ModuleList( + [ + HyNomasCameraAttention(hidden_size, camera_embed_dim, camera_attention_dim, camera_num_attention_heads) + for _ in range(num_layers) + ] + ) + self.extrinsics_embedding_camctrl = nn.Linear(64, 64) + self.c2ws_hidden_states_layer1 = nn.Linear(64, 64) + self.c2ws_hidden_states_layer2 = nn.Linear(64, 64) + self.final_layer = HyNomasFinalLayer(hidden_size, out_channels) + self.gradient_checkpointing = False + + def enable_reference_attention(self, text_backend="flash_varlen", deterministic=False, flash_deterministic=False): + """Use optional Magi/FlashAttention kernels with the released model's token packing. + + This runtime processor choice is not saved in the model configuration. Call this method after loading. + `text_backend` selects `flash_varlen` (FlashAttention 2) or `_flash_varlen_3` (FlashAttention 3). Main action + attention uses Magi, and unrestricted cached attention uses FlashAttention 3. + """ + if text_backend not in ("flash_varlen", "_flash_varlen_3"): + raise ValueError("Reference text attention requires flash_varlen or _flash_varlen_3.") + # Both FlashAttention versions are required: the one-frame camera path uses FA2, and + # unrestricted cached attention uses FA3. Magi itself remains lazily imported at dispatch. + _check_attention_backend_requirements(AttentionBackendName.FLASH_VARLEN) + _check_attention_backend_requirements(AttentionBackendName._FLASH_VARLEN_3) + _check_attention_backend_requirements(AttentionBackendName.MAGI) + for module in self.modules(): + if isinstance(module, HyNomasTextAttention): + module.set_processor(HyNomasFlashTextAttnProcessor(text_backend, flash_deterministic)) + elif isinstance(module, HyNomasAttention): + module.set_processor(HyNomasMagiAttnProcessor(deterministic, flash_deterministic)) + elif isinstance(module, HyNomasCameraAttention): + module.set_processor(HyNomasMagiCameraAttnProcessor()) + + def disable_reference_attention(self) -> None: + """Restore the standard attention processors after enabling reference attention.""" + reference_processors = ( + HyNomasFlashTextAttnProcessor, + HyNomasMagiAttnProcessor, + HyNomasMagiCameraAttnProcessor, + ) + for module in self.modules(): + if isinstance(module, AttentionModuleMixin) and isinstance(module.processor, reference_processors): + module.set_processor(module._default_processor_cls()) + + def set_attention_backend(self, backend: str) -> None: + """Restore the standard processors and select a Diffusers attention backend.""" + self.disable_reference_attention() + super().set_attention_backend(backend) + + def reset_attention_backend(self) -> None: + """Restore the standard processors and reset their attention backend.""" + self.disable_reference_attention() + super().reset_attention_backend() + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_attention_mask: torch.Tensor | None = None, + kv_cache: list[dict[str, torch.Tensor]] | None = None, + update_cache: bool = False, + frame_offset: int = 0, + camera_extrinsics: torch.Tensor | None = None, + camera_attention_mask: torch.Tensor | None = None, + action_text_offsets: list[list[int] | None] | None = None, + action_frame_windows: list[tuple[int, int, int]] | None = None, + cache_action_start: int | None = None, + return_dict: bool = True, + ) -> Transformer2DModelOutput | tuple[torch.Tensor]: + r"""Predict flow velocity for one segment. + + Args: + hidden_states (`torch.Tensor`): + Noisy video latents of shape `(batch, channels, frames, height, width)`. + encoder_hidden_states (`torch.Tensor`): + Text features of shape `(batch, text_tokens, text_embed_dim)`. + timestep (`torch.Tensor`): + Diffusion timesteps of shape `(batch,)` in the reference model's 0–1000 scale. + encoder_attention_mask (`torch.Tensor`, *optional*): + Boolean text validity mask of shape `(batch, text_tokens)`. + kv_cache (`list[dict[str, torch.Tensor]]`, *optional*): + Caller-owned history, with one dictionary per transformer block. Empty dictionaries start a new + session. Populated entries contain `key`, `value`, and `mask` tensors. The model stores no session + state. + update_cache (`bool`, defaults to `False`): + Append detached keys, values, and a visibility mask to `kv_cache` after attention. + frame_offset (`int`, defaults to `0`): + Number of previously generated latent frames, used for video rotary positions. + camera_extrinsics (`torch.Tensor`, *optional*): + Normalized poses of shape `(batch, pixel_frames, 4, 4)`. Four consecutive poses form one latent frame; + the first pose is repeated at the front when necessary. + camera_attention_mask (`torch.Tensor`, *optional*): + Per-sample camera gate of shape `(batch,)`. + action_text_offsets (`list[list[int] | None]`, *optional*): + Head and action text boundaries for each sample; `None` entries use unrestricted current text. + action_frame_windows (`list[tuple[int, int, int]]`, *optional*): + `(start_frame, end_frame, action_index)` intervals covering this segment, with `-1` denoting a gap. + cache_action_start (`int`, *optional*): + Text boundary after which tokens become invisible to later segments, independently of action windows. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`] or a tuple. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple[torch.Tensor]`: + Predicted flow velocity, with the same video geometry as `hidden_states`. + """ + batch_size, _, num_frames, height, width = hidden_states.shape + if kv_cache is not None and len(kv_cache) != len(self.double_blocks): + raise ValueError("kv_cache must contain one dictionary per transformer block.") + if update_cache and kv_cache is None: + raise ValueError("update_cache=True requires a caller-owned kv_cache.") + if update_cache and self.gradient_checkpointing and torch.is_grad_enabled(): + raise ValueError("Cache writes cannot be combined with gradient checkpointing.") + if action_text_offsets is not None and ( + len(action_text_offsets) != batch_size or action_frame_windows is None + ): + raise ValueError("Action offsets require one entry per sample and action_frame_windows.") + if cache_action_start is not None and not 0 <= cache_action_start <= encoder_hidden_states.shape[1]: + raise ValueError("cache_action_start must be a valid text boundary or None.") + if kv_cache is not None: + cache_lengths = [] + for layer_cache in kv_cache: + if not layer_cache: + cache_lengths.append(0) + continue + if not {"key", "value", "mask"}.issubset(layer_cache): + raise ValueError("Each nonempty cache entry must contain key, value, and mask tensors.") + key, value, mask = layer_cache["key"], layer_cache["value"], layer_cache["mask"] + expected_shape = ( + batch_size, + key.shape[1], + self.config.num_attention_heads, + self.config.attention_head_dim, + ) + if key.shape != expected_shape or value.shape != expected_shape or mask.shape != expected_shape[:2]: + raise ValueError("kv_cache tensor shapes must match the model and input batch size.") + cache_lengths.append(key.shape[1]) + if len(set(cache_lengths)) != 1: + raise ValueError("All kv_cache layers must contain the same number of tokens.") + rotary_emb = self.rotary_embed(num_frames, height, width, hidden_states.device, frame_offset) + hidden_states = self.img_in(hidden_states) + temb = self.time_in(timestep, hidden_states.dtype) + encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep, encoder_attention_mask) + camera_hidden_states = None + if camera_extrinsics is not None: + camera_extrinsics = camera_extrinsics.to(device=hidden_states.device, dtype=hidden_states.dtype) + padding = (-camera_extrinsics.shape[1]) % 4 + if padding: + camera_extrinsics = torch.cat( + (camera_extrinsics[:, :1].repeat(1, padding, 1, 1), camera_extrinsics), dim=1 + ) + if camera_extrinsics.shape[1] != num_frames * 4: + raise ValueError("Camera poses must align with the segment's latent frames after front padding.") + camera_hidden_states = self.extrinsics_embedding_camctrl( + camera_extrinsics.reshape(batch_size, num_frames, 64) + ) + camera_hidden_states = camera_hidden_states + self.c2ws_hidden_states_layer2( + F.silu(self.c2ws_hidden_states_layer1(camera_hidden_states)) + ).to(camera_hidden_states.device) + camera_rotary_emb = self.camera_rotary_embed(num_frames, 1, 1, hidden_states.device) + action_token_windows = None + if action_frame_windows is not None: + action_token_windows = [ + (start * height * width, end * height * width, index) for start, end, index in action_frame_windows + ] + for index, block in enumerate(self.double_blocks): + attention_kwargs = { + "encoder_attention_mask": encoder_attention_mask, + "kv_cache": kv_cache[index] if kv_cache is not None else None, + "update_cache": update_cache, + "action_text_offsets": action_text_offsets, + "action_token_windows": action_token_windows, + "cache_action_start": cache_action_start, + } + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, temb, rotary_emb, attention_kwargs + ) + else: + hidden_states, encoder_hidden_states = block( + hidden_states, encoder_hidden_states, temb, rotary_emb, attention_kwargs=attention_kwargs + ) + if camera_hidden_states is not None: + hidden_states = self.action_modules[index]( + hidden_states, camera_hidden_states, camera_rotary_emb, camera_attention_mask + ) + hidden_states = self.final_layer(hidden_states, temb) + output = hidden_states.transpose(1, 2).reshape(batch_size, self.config.out_channels, num_frames, height, width) + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 48ccfcc98125..6b73238c36e0 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -67,6 +67,7 @@ "HeliosPyramidDistilledModularPipeline", "HeliosPyramidModularPipeline", ] + _import_structure["hy_nomas"] = ["HyNomasAutoBlocks", "HyNomasGenerationState", "HyNomasModularPipeline"] _import_structure["flux"] = [ "FluxAutoBlocks", "FluxModularPipeline", @@ -203,6 +204,7 @@ HunyuanVideo15AutoBlocks, HunyuanVideo15ModularPipeline, ) + from .hy_nomas import HyNomasAutoBlocks, HyNomasGenerationState, HyNomasModularPipeline from .ideogram4 import ( Ideogram4AutoBlocks, Ideogram4ModularPipeline, diff --git a/src/diffusers/modular_pipelines/hy_nomas/__init__.py b/src/diffusers/modular_pipelines/hy_nomas/__init__.py new file mode 100644 index 000000000000..791fa408a07d --- /dev/null +++ b/src/diffusers/modular_pipelines/hy_nomas/__init__.py @@ -0,0 +1,60 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_hy_nomas"] = ["HyNomasAutoBlocks"] + _import_structure["modular_pipeline"] = ["HyNomasGenerationState", "HyNomasModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_hy_nomas import HyNomasAutoBlocks + from .modular_pipeline import HyNomasGenerationState, HyNomasModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/modular_pipelines/hy_nomas/before_denoise.py b/src/diffusers/modular_pipelines/hy_nomas/before_denoise.py new file mode 100644 index 000000000000..11b74c0e7243 --- /dev/null +++ b/src/diffusers/modular_pipelines/hy_nomas/before_denoise.py @@ -0,0 +1,319 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import torch + +from ...models.transformers.transformer_hy_nomas import HyNomasTransformer3DModel +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import HyNomasGenerationState, HyNomasModularPipeline + + +def _timestamp(value): + parts = str(value).split(":") + if len(parts) > 3: + raise ValueError(f"Invalid action timestamp: {value}") + return sum(float(part) * 60**i for i, part in enumerate(reversed(parts))) + + +def _action_windows(prompt, frames, fps): + try: + obj = json.loads(prompt) + except (ValueError, TypeError): + return None + segments = obj.get("segments") if isinstance(obj, dict) else None + if not isinstance(segments, list) or not segments or not isinstance(segments[0], dict): + return None + actions = segments[0].get("contained_actions", []) + if isinstance(actions, dict): + actions = list(actions.values()) + if not isinstance(actions, list) or not actions: + return None + spans = [] + for index, action in enumerate(actions): + if not isinstance(action, dict) or "start_time" not in action or "end_time" not in action: + return None + start, end = [ + max(0, min(frames, int(max(_timestamp(action[key]), 0) * fps) // 4)) for key in ("start_time", "end_time") + ] + if end > start: + spans.append((start, end, index)) + if not spans: + return None + spans.sort() + cursor, windows = 0, [] + for start, end, index in spans: + if start < cursor: + raise ValueError("Action time windows must not overlap.") + if start > cursor: + windows.append((cursor, start, -1)) + windows.append((start, end, index)) + cursor = end + if cursor < frames: + windows.append((cursor, frames, -1)) + return None if len(windows) == len(actions) == 1 else windows + + +# auto_docstring +class HyNomasPrepareStep(ModularPipelineBlocks): + """ + Validate segment lengths and camera poses, fork the world state, and prepare rollout noise. + + Components: + transformer (`HyNomasTransformer3DModel`) + + Inputs: + prompts (`list`): + Ordered prompts for successive segments of one video, rather than a batch of independent videos. + height (`int`, *optional*, defaults to 704): + Output height in pixels; must be a positive multiple of 32 and match the existing session. + width (`int`, *optional*, defaults to 1280): + Output width in pixels; must be a positive multiple of 32 and match the existing session. + num_frames (`int`, *optional*, defaults to 97): + Exact number of new pixel frames to generate in this call, excluding all prior session history. + segment_latent_lengths (`list`, *optional*): + Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4). + camera_extrinsics (`Tensor`, *optional*): + Normalized camera-to-world matrices in the training coordinate system, shaped (1, num_frames, 4, 4), for + the new rollout only. + fps (`float`, *optional*, defaults to 24.0): + Frames per second used to convert structured action timestamps into latent-frame intervals. + generator (`Generator`, *optional*): + Generator for rollout noise and cache-write noise; omitted values restore the session generator or start + from seed 42. + latents (`Tensor`, *optional*): + Optional initial noise for the new rollout, shaped (1, channels, ceil(num_frames / 4), height / 32, width + / 32). + noise_capacity (`int`, *optional*): + Total latent-frame capacity reserved when creating a fixed session noise timeline, including future + continuation segments. + session (`HyNomasGenerationState`, *optional*): + Earlier segment history to continue or branch; this call forks the state before appending new cache + entries and latents. + + Outputs: + session (`HyNomasGenerationState`): + Forked generation state containing prior segment history and the optional fixed noise timeline. + initial_latents (`Tensor`): + Initial noise for the new rollout, shaped (1, channels, new_latent_frames, height / 32, width / 32). + segment_latent_lengths (`list`): + Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4). + segment_cameras (`list`): + Normalized camera-to-world poses for each segment, shaped (1, latent_frames * 4, 4, 4), or None without + camera conditioning. + action_frame_windows (`list`): + Per-segment (start, end, action_index) intervals in latent frames; -1 denotes an action gap, and None + disables action-specific windows. + generator (`Generator`): + Generator initialized or restored for this rollout and its cache-write noise. + segment_pixel_lengths (`list`): + Exact visible pixel-frame count of each new segment after removing decoder padding from the final + segment. + """ + + model_name = "hy-nomas" + + @property + def description(self): + return "Validate segment lengths and camera poses, fork the world state, and prepare rollout noise." + + @property + def expected_components(self): + return [ComponentSpec("transformer", HyNomasTransformer3DModel)] + + @property + def inputs(self): + return [ + InputParam( + "prompts", + type_hint=list[str], + description="Ordered prompts for successive segments of one video, rather than a batch of independent videos.", + required=True, + ), + InputParam( + "height", + type_hint=int, + description="Output height in pixels; must be a positive multiple of 32 and match the existing session.", + default=704, + ), + InputParam( + "width", + type_hint=int, + description="Output width in pixels; must be a positive multiple of 32 and match the existing session.", + default=1280, + ), + InputParam( + "num_frames", + type_hint=int, + description="Exact number of new pixel frames to generate in this call, excluding all prior session history.", + default=97, + ), + InputParam( + "segment_latent_lengths", + type_hint=list[int], + description="Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4).", + ), + InputParam( + "camera_extrinsics", + type_hint=torch.Tensor, + description="Normalized camera-to-world matrices in the training coordinate system, shaped (1, num_frames, 4, 4), for the new rollout only.", + ), + InputParam( + "fps", + type_hint=float, + description="Frames per second used to convert structured action timestamps into latent-frame intervals.", + default=24.0, + ), + InputParam( + "generator", + type_hint=torch.Generator, + description="Generator for rollout noise and cache-write noise; omitted values restore the session generator or start from seed 42.", + ), + InputParam( + "latents", + type_hint=torch.Tensor, + description="Optional initial noise for the new rollout, shaped (1, channels, ceil(num_frames / 4), height / 32, width / 32).", + ), + InputParam( + "noise_capacity", + type_hint=int, + description="Total latent-frame capacity reserved when creating a fixed session noise timeline, including future continuation segments.", + ), + InputParam( + "session", + type_hint=HyNomasGenerationState, + description="Earlier segment history to continue or branch; this call forks the state before appending new cache entries and latents.", + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "session", + type_hint=HyNomasGenerationState, + description="Forked generation state containing prior segment history and the optional fixed noise timeline.", + ), + OutputParam( + "initial_latents", + type_hint=torch.Tensor, + description="Initial noise for the new rollout, shaped (1, channels, new_latent_frames, height / 32, width / 32).", + ), + OutputParam( + "segment_latent_lengths", + type_hint=list[int], + description="Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4).", + ), + OutputParam( + "segment_cameras", + type_hint=list[torch.Tensor | None], + description="Normalized camera-to-world poses for each segment, shaped (1, latent_frames * 4, 4, 4), or None without camera conditioning.", + ), + OutputParam( + "action_frame_windows", + type_hint=list[list[tuple[int, int, int]] | None], + description="Per-segment (start, end, action_index) intervals in latent frames; -1 denotes an action gap, and None disables action-specific windows.", + ), + OutputParam( + "generator", + type_hint=torch.Generator, + description="Generator initialized or restored for this rollout and its cache-write noise.", + ), + OutputParam( + "segment_pixel_lengths", + type_hint=list[int], + description="Exact visible pixel-frame count of each new segment after removing decoder padding from the final segment.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state) -> tuple[HyNomasModularPipeline, PipelineState]: + s = self.get_block_state(state) + if min(s.height, s.width) < 32 or s.height % 32 or s.width % 32: + raise ValueError("height and width must be positive multiples of 32.") + if not isinstance(s.num_frames, int) or s.num_frames < 1: + raise ValueError("num_frames must be a positive integer.") + if s.fps <= 0: + raise ValueError("fps must be positive.") + frames = (s.num_frames - 1) // 4 + 1 + count = len(s.prompts) + if s.segment_latent_lengths is None: + if frames % count: + raise ValueError( + "Provide segment_latent_lengths when latent frames do not divide equally among prompts." + ) + s.segment_latent_lengths = [frames // count] * count + lengths = s.segment_latent_lengths + if len(lengths) != count or any(not isinstance(x, int) or x < 1 for x in lengths) or sum(lengths) != frames: + raise ValueError( + "segment_latent_lengths must contain one positive length per prompt and sum to latent frames." + ) + session = ( + s.session.fork() + if s.session is not None + else HyNomasGenerationState( + kv_cache=[{} for _ in range(components.transformer.config.num_layers)], height=s.height, width=s.width + ) + ) + if (session.height, session.width) != (s.height, s.width): + raise ValueError("A continuation must use the same height and width as its session.") + if s.generator is None: + s.generator = torch.Generator(device=session.generator_device) + if session.generator_state is None: + s.generator.manual_seed(42) + else: + s.generator.set_state(session.generator_state) + if isinstance(s.generator, list): + raise ValueError("Hy-Nomas generates one world per call; provide a single generator.") + shape = (1, components.transformer.config.in_channels, frames, s.height // 32, s.width // 32) + device, dtype = components._execution_device, components.transformer.dtype + if s.latents is not None: + if tuple(s.latents.shape) != shape: + raise ValueError(f"latents must have shape {shape}.") + noise = s.latents.to(device=device, dtype=dtype) + elif session.noise_bank is not None or s.noise_capacity is not None: + if session.noise_bank is None: + if s.noise_capacity < frames: + raise ValueError("noise_capacity must cover all requested latent frames.") + session.noise_bank = randn_tensor( + (shape[0], shape[1], s.noise_capacity, *shape[3:]), + generator=s.generator, + device=device, + dtype=dtype, + ) + end = session.frame_offset + frames + if end > session.noise_bank.shape[2]: + raise ValueError("The request exceeds the session's fixed noise_capacity.") + noise = session.noise_bank[:, :, session.frame_offset : end] + else: + noise = randn_tensor(shape, generator=s.generator, device=device, dtype=dtype) + cameras = [None] * count + if s.camera_extrinsics is not None: + camera = s.camera_extrinsics + if tuple(camera.shape) != (1, s.num_frames, 4, 4) or not torch.isfinite(camera).all(): + raise ValueError("camera_extrinsics must be finite with shape (1, num_frames, 4, 4).") + camera = camera.to(device=device, dtype=dtype) + # The original camera embedder pads at the beginning to a whole group of four pixel frames. + padding = frames * 4 - s.num_frames + camera = torch.cat([camera[:, :1].expand(-1, padding, -1, -1), camera], dim=1) + cameras = list(camera.split([length * 4 for length in lengths], dim=1)) + s.session, s.initial_latents, s.segment_cameras = session, noise, cameras + s.segment_pixel_lengths = [length * 4 for length in lengths] + s.segment_pixel_lengths[-1] -= frames * 4 - s.num_frames + s.action_frame_windows = [_action_windows(prompt, length, s.fps) for prompt, length in zip(s.prompts, lengths)] + self.set_block_state(state, s) + return components, state diff --git a/src/diffusers/modular_pipelines/hy_nomas/decoders.py b/src/diffusers/modular_pipelines/hy_nomas/decoders.py new file mode 100644 index 000000000000..882c448e7508 --- /dev/null +++ b/src/diffusers/modular_pipelines/hy_nomas/decoders.py @@ -0,0 +1,133 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import PIL +import torch + +from ...configuration_utils import FrozenDict +from ...models.autoencoders.autoencoder_pixelumm import PixelUMMDecoder +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import HyNomasGenerationState, HyNomasModularPipeline + + +# auto_docstring +class HyNomasDecodeStep(ModularPipelineBlocks): + """ + Decode the accumulated world latents with the distilled PixelUMM diffusion decoder. + + Components: + decoder (`PixelUMMDecoder`) video_processor (`VideoProcessor`) + + Inputs: + latents (`Tensor`): + Denoised latents for the accumulated world, including all earlier segments, shaped (1, 32, latent_frames, + latent_height, latent_width). + session (`HyNomasGenerationState`, *optional*): + History carrying each segment's exact visible pixel length so its decoder padding can be removed; omitted + states use (latent_frames - 1) * 4 + 1 frames. + output_type (`str`, *optional*, defaults to np): + Output representation: np for NumPy arrays, pt for tensors, pil for frame images, or latent to skip + decoding. + decoder_generator (`Generator`, *optional*): + Independent PixelUMM noise generator; only its initial seed is used, without advancing its state. If + omitted, the decoder configuration seed is used (42 by default). + + Outputs: + videos (`Tensor | ndarray | list`): + Complete accumulated video including prior history: pt has shape (1, frames, 3, height, width), np has + shape (1, frames, height, width, 3), and pil is a nested list of frames. The latent output preserves the + input latent tensor. + """ + + model_name = "hy-nomas" + + @property + def description(self): + return "Decode the accumulated world latents with the distilled PixelUMM diffusion decoder." + + @property + def expected_components(self): + return [ + ComponentSpec("decoder", PixelUMMDecoder), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 32, "do_normalize": False}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self): + return [ + InputParam( + "latents", + type_hint=torch.Tensor, + description="Denoised latents for the accumulated world, including all earlier segments, shaped (1, 32, latent_frames, latent_height, latent_width).", + required=True, + ), + InputParam( + "session", + type_hint=HyNomasGenerationState, + description="History carrying each segment's exact visible pixel length so its decoder padding can be removed; omitted states use (latent_frames - 1) * 4 + 1 frames.", + ), + InputParam( + "output_type", + type_hint=str, + description="Output representation: np for NumPy arrays, pt for tensors, pil for frame images, or latent to skip decoding.", + default="np", + ), + InputParam( + "decoder_generator", + type_hint=torch.Generator, + description="Independent PixelUMM noise generator; only its initial seed is used, without advancing its state. If omitted, the decoder configuration seed is used (42 by default).", + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "videos", + type_hint=torch.Tensor | np.ndarray | list[list[PIL.Image.Image]], + description="Complete accumulated video including prior history: pt has shape (1, frames, 3, height, width), np has shape (1, frames, height, width, 3), and pil is a nested list of frames. The latent output preserves the input latent tensor.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state) -> tuple[HyNomasModularPipeline, PipelineState]: + s = self.get_block_state(state) + if s.output_type == "latent": + s.videos = s.latents + else: + video = components.decoder.decode( + s.latents.to(device=components._execution_device, dtype=components.decoder.dtype), + generator=s.decoder_generator, + return_dict=False, + )[0] + if s.session is None: + video = video[:, :, : (s.latents.shape[2] - 1) * 4 + 1] + else: + chunks = video.split([chunk.shape[2] * 4 for chunk in s.session.latent_chunks], dim=2) + video = torch.cat( + [chunk[:, :, :length] for chunk, length in zip(chunks, s.session.pixel_lengths)], dim=2 + ) + # Preserve the released decoder's arithmetic precision before converting output formats. + video = (video * 0.5 + 0.5).clamp(0, 1) + s.videos = components.video_processor.postprocess_video(video, output_type=s.output_type) + self.set_block_state(state, s) + return components, state diff --git a/src/diffusers/modular_pipelines/hy_nomas/denoise.py b/src/diffusers/modular_pipelines/hy_nomas/denoise.py new file mode 100644 index 000000000000..b2d7bede1d14 --- /dev/null +++ b/src/diffusers/modular_pipelines/hy_nomas/denoise.py @@ -0,0 +1,326 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ...models.transformers.transformer_hy_nomas import HyNomasTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import HyNomasGenerationState, HyNomasModularPipeline + + +# auto_docstring +class HyNomasSegmentDenoiseStep(ModularPipelineBlocks): + """ + Run flow-matching Euler steps for one segment, with earlier segments supplied as cached keys and values. + + Components: + transformer (`HyNomasTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + initial_latents (`Tensor`): + Initial noise for all new segments, shaped (1, channels, new_latent_frames, latent_height, latent_width). + segment_latent_lengths (`list`): + Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4). + segment_pixel_lengths (`list`): + Exact visible pixel-frame count of each new segment after removing decoder padding from the final + segment. + prompt_embeds (`list`): + Qwen hidden states for each segment, shaped (1, max_sequence_length, hidden_size), with the chat prefix + removed. + prompt_attention_masks (`list`): + Per-segment text masks shaped (1, max_sequence_length), with one for valid tokens and zero for padding. + segment_cameras (`list`): + Normalized camera-to-world poses for each segment, shaped (1, latent_frames * 4, 4, 4), or None without + camera conditioning. + action_text_offsets (`list`): + Per-segment action-token boundaries in the cropped text sequence, or None for prompts without structured + actions. + action_frame_windows (`list`): + Per-segment (start, end, action_index) intervals in latent frames; -1 denotes an action gap, and None + disables action-specific windows. + cache_action_starts (`list`): + Per-segment text boundary after which tokens are hidden from future segments; None keeps all valid text + visible. + session (`HyNomasGenerationState`): + Forked world state supplying cached keys and values from earlier segments, including prior calls. + num_inference_steps (`int`, *optional*, defaults to 8): + Number of flow-matching Euler denoising steps for each new segment. + cache_noise_sigma (`float`, *optional*, defaults to 0.206897): + Noise level in [0, 1] applied to a completed segment before committing its transformer keys and values to + history. + generator (`Generator`, *optional*): + Rollout generator used to draw the noise for each segment cache write. + + Outputs: + segment_latents (`Tensor`): + Denoised latents for the current segment, shaped (1, channels, segment_latent_frames, latent_height, + latent_width). + session (`HyNomasGenerationState`): + World state updated with the completed segment, its visible pixel length, and its keys and values. + """ + + model_name = "hy-nomas" + + @property + def description(self): + return ( + "Run flow-matching Euler steps for one segment, with earlier segments supplied as cached keys and values." + ) + + @property + def expected_components(self): + return [ + ComponentSpec("transformer", HyNomasTransformer3DModel), + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def inputs(self): + return [ + InputParam( + "initial_latents", + type_hint=torch.Tensor, + description="Initial noise for all new segments, shaped (1, channels, new_latent_frames, latent_height, latent_width).", + required=True, + ), + InputParam( + "segment_latent_lengths", + type_hint=list[int], + description="Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4).", + required=True, + ), + InputParam( + "segment_pixel_lengths", + type_hint=list[int], + description="Exact visible pixel-frame count of each new segment after removing decoder padding from the final segment.", + required=True, + ), + InputParam( + "prompt_embeds", + type_hint=list[torch.Tensor], + description="Qwen hidden states for each segment, shaped (1, max_sequence_length, hidden_size), with the chat prefix removed.", + required=True, + ), + InputParam( + "prompt_attention_masks", + type_hint=list[torch.Tensor], + description="Per-segment text masks shaped (1, max_sequence_length), with one for valid tokens and zero for padding.", + required=True, + ), + InputParam( + "segment_cameras", + type_hint=list[torch.Tensor | None], + description="Normalized camera-to-world poses for each segment, shaped (1, latent_frames * 4, 4, 4), or None without camera conditioning.", + required=True, + ), + InputParam( + "action_text_offsets", + type_hint=list[list[int] | None], + description="Per-segment action-token boundaries in the cropped text sequence, or None for prompts without structured actions.", + required=True, + ), + InputParam( + "action_frame_windows", + type_hint=list[list[tuple[int, int, int]] | None], + description="Per-segment (start, end, action_index) intervals in latent frames; -1 denotes an action gap, and None disables action-specific windows.", + required=True, + ), + InputParam( + "cache_action_starts", + type_hint=list[int | None], + description="Per-segment text boundary after which tokens are hidden from future segments; None keeps all valid text visible.", + required=True, + ), + InputParam( + "session", + type_hint=HyNomasGenerationState, + description="Forked world state supplying cached keys and values from earlier segments, including prior calls.", + required=True, + ), + InputParam( + "num_inference_steps", + type_hint=int, + description="Number of flow-matching Euler denoising steps for each new segment.", + default=8, + ), + InputParam( + "cache_noise_sigma", + type_hint=float, + description="Noise level in [0, 1] applied to a completed segment before committing its transformer keys and values to history.", + default=0.206897, + ), + InputParam( + "generator", + type_hint=torch.Generator, + description="Rollout generator used to draw the noise for each segment cache write.", + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "segment_latents", + type_hint=torch.Tensor, + description="Denoised latents for the current segment, shaped (1, channels, segment_latent_frames, latent_height, latent_width).", + ), + OutputParam( + "session", + type_hint=HyNomasGenerationState, + description="World state updated with the completed segment, its visible pixel length, and its keys and values.", + ), + ] + + @torch.no_grad() + def __call__(self, components, s, segment_index) -> tuple[HyNomasModularPipeline, BlockState]: + if s.num_inference_steps < 1: + raise ValueError("num_inference_steps must be positive.") + if not 0 <= s.cache_noise_sigma <= 1: + raise ValueError("cache_noise_sigma must lie in [0, 1].") + start = sum(s.segment_latent_lengths[:segment_index]) + end = start + s.segment_latent_lengths[segment_index] + latents = s.initial_latents[:, :, start:end].clone() + # The released solver uses a grid ending at 1 / steps, followed by zero. + sigmas = torch.linspace(1, 0, s.num_inference_steps + 1)[:-1].tolist() + components.scheduler.set_timesteps(sigmas=sigmas, device=latents.device) + mask = s.prompt_attention_masks[segment_index] + if mask is not None and mask.all(): + mask = None + model_inputs = { + "encoder_hidden_states": s.prompt_embeds[segment_index].to( + device=components._execution_device, dtype=components.transformer.dtype + ), + "encoder_attention_mask": None if mask is None else mask.to(components._execution_device), + "kv_cache": s.session.kv_cache, + "frame_offset": s.session.frame_offset, + "camera_extrinsics": s.segment_cameras[segment_index], + "action_text_offsets": ( + [s.action_text_offsets[segment_index]] if s.action_frame_windows[segment_index] is not None else None + ), + "action_frame_windows": s.action_frame_windows[segment_index], + "cache_action_start": s.cache_action_starts[segment_index], + } + for timestep in components.scheduler.timesteps: + velocity = components.transformer( + hidden_states=latents, timestep=timestep.expand(1), **model_inputs, return_dict=False + )[0] + latents = components.scheduler.step(velocity.float(), timestep, latents.float(), return_dict=False)[0].to( + latents.dtype + ) + s.segment_latents = latents + cache_latents = latents + sigma = s.cache_noise_sigma + if sigma: + noise = randn_tensor(latents.shape, generator=s.generator, device=latents.device, dtype=latents.dtype) + cache_latents = ((1 - sigma) * latents.float() + sigma * noise.float()).to(latents.dtype) + components.transformer( + hidden_states=cache_latents, + timestep=torch.full((1,), sigma * 1000, device=latents.device, dtype=torch.float32), + update_cache=True, + **model_inputs, + return_dict=False, + ) + s.session.latent_chunks.append(latents) + s.session.pixel_lengths.append(s.segment_pixel_lengths[segment_index]) + s.session.frame_offset += latents.shape[2] + return components, s + + +# auto_docstring +class HyNomasSegmentLoop(LoopSequentialPipelineBlocks): + """ + Generate successive segments and commit each completed segment to the caller's world history. + + Components: + transformer (`HyNomasTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + initial_latents (`Tensor`): + Initial noise for all new segments, shaped (1, channels, new_latent_frames, latent_height, latent_width). + segment_latent_lengths (`list`): + Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4). + segment_pixel_lengths (`list`): + Exact visible pixel-frame count of each new segment after removing decoder padding from the final + segment. + prompt_embeds (`list`): + Qwen hidden states for each segment, shaped (1, max_sequence_length, hidden_size), with the chat prefix + removed. + prompt_attention_masks (`list`): + Per-segment text masks shaped (1, max_sequence_length), with one for valid tokens and zero for padding. + segment_cameras (`list`): + Normalized camera-to-world poses for each segment, shaped (1, latent_frames * 4, 4, 4), or None without + camera conditioning. + action_text_offsets (`list`): + Per-segment action-token boundaries in the cropped text sequence, or None for prompts without structured + actions. + action_frame_windows (`list`): + Per-segment (start, end, action_index) intervals in latent frames; -1 denotes an action gap, and None + disables action-specific windows. + cache_action_starts (`list`): + Per-segment text boundary after which tokens are hidden from future segments; None keeps all valid text + visible. + session (`HyNomasGenerationState`): + Forked world state supplying cached keys and values from earlier segments, including prior calls. + num_inference_steps (`int`, *optional*, defaults to 8): + Number of flow-matching Euler denoising steps for each new segment. + cache_noise_sigma (`float`, *optional*, defaults to 0.206897): + Noise level in [0, 1] applied to a completed segment before committing its transformer keys and values to + history. + generator (`Generator`, *optional*): + Rollout generator used to draw the noise for each segment cache write. + + Outputs: + segment_latents (`Tensor`): + Denoised latents for the current segment, shaped (1, channels, segment_latent_frames, latent_height, + latent_width). + session (`HyNomasGenerationState`): + World state updated with the completed segment, its visible pixel length, and its keys and values. + """ + + model_name = "hy-nomas" + block_classes = [HyNomasSegmentDenoiseStep] + block_names = ["denoise_segment"] + + @property + def description(self): + return "Generate successive segments and commit each completed segment to the caller's world history." + + @property + def loop_intermediate_outputs(self): + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="Concatenated denoised latents for the complete world history, including prior calls and every new segment.", + ), + OutputParam( + "session", + type_hint=HyNomasGenerationState, + description="Updated world history and generator state, ready for continuation or branching in a later call.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state) -> tuple[HyNomasModularPipeline, PipelineState]: + s = self.get_block_state(state) + for index in range(len(s.segment_latent_lengths)): + components, s = self.loop_step(components, s, segment_index=index) + s.latents = torch.cat(s.session.latent_chunks, dim=2) + s.session.generator_state = s.generator.get_state().clone() + s.session.generator_device = str(s.generator.device) + self.set_block_state(state, s) + return components, state diff --git a/src/diffusers/modular_pipelines/hy_nomas/encoders.py b/src/diffusers/modular_pipelines/hy_nomas/encoders.py new file mode 100644 index 000000000000..5a4e4e21449a --- /dev/null +++ b/src/diffusers/modular_pipelines/hy_nomas/encoders.py @@ -0,0 +1,204 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import torch +from transformers import AutoTokenizer, Qwen2_5_VLTextModel + +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import HyNomasModularPipeline + + +SYSTEM_PROMPT = ( + "You are a helpful assistant. Describe the video by detailing the following aspects: " + " 1. The main content and theme of the video. " + " 2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects. " + " 3. Actions, events, behaviors temporal relationships, physical movement changes of the objects. " + " 4. background environment, light, style and atmosphere. " + " 5. camera angles, movements, and transitions used in the video." +) + + +def _tokenize(tokenizer, prompt, max_length=None): + kwargs = {} if max_length is None else {"truncation": True, "max_length": max_length, "padding": "max_length"} + return tokenizer.apply_chat_template( + [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + **kwargs, + ) + + +def _action_variants(prompt): + try: + obj = json.loads(prompt) + except (ValueError, TypeError): + return [] + segments = obj.get("segments") if isinstance(obj, dict) else None + if not isinstance(segments, list) or not segments or not isinstance(segments[0], dict): + return [] + actions = segments[0].get("contained_actions") + if not isinstance(actions, (dict, list)) or not actions: + return [] + variants = [] + for count in range(len(actions)): + segment = dict(segments[0]) + if count == 0: + segment.pop("contained_actions") + elif isinstance(actions, dict): + segment["contained_actions"] = dict(list(actions.items())[:count]) + else: + segment["contained_actions"] = actions[:count] + variants.append( + json.dumps( + {key: [segment] if key == "segments" else value for key, value in obj.items()}, ensure_ascii=False + ) + ) + return variants + + +# auto_docstring +class HyNomasTextEncoderStep(ModularPipelineBlocks): + """ + Encode each segment with the Qwen video template and locate its action tokens. + + Components: + text_encoder (`Qwen2_5_VLTextModel`) tokenizer (`AutoTokenizer`) + + Inputs: + prompt (`str | list`): + One prompt or an ordered list of prompts describing successive segments of a single video. + max_sequence_length (`int`, *optional*, defaults to 1280): + Number of text tokens retained per segment after removing the Qwen chat prefix; shorter sequences are + padded. + + Outputs: + prompts (`list`): + Ordered prompts for successive segments of one video, rather than a batch of independent videos. + prompt_embeds (`list`): + Qwen hidden states for each segment, shaped (1, max_sequence_length, hidden_size), with the chat prefix + removed. + prompt_attention_masks (`list`): + Per-segment text masks shaped (1, max_sequence_length), with one for valid tokens and zero for padding. + action_text_offsets (`list`): + Per-segment action-token boundaries in the cropped text sequence, or None for prompts without structured + actions. + cache_action_starts (`list`): + Per-segment text boundary after which tokens are hidden from future segments; None keeps all valid text + visible. + """ + + model_name = "hy-nomas" + + @property + def description(self): + return "Encode each segment with the Qwen video template and locate its action tokens." + + @property + def expected_components(self): + return [ComponentSpec("text_encoder", Qwen2_5_VLTextModel), ComponentSpec("tokenizer", AutoTokenizer)] + + @property + def inputs(self): + return [ + InputParam( + "prompt", + type_hint=str | list[str], + description="One prompt or an ordered list of prompts describing successive segments of a single video.", + required=True, + ), + InputParam( + "max_sequence_length", + type_hint=int, + description="Number of text tokens retained per segment after removing the Qwen chat prefix; shorter sequences are padded.", + default=1280, + ), + ] + + @property + def intermediate_outputs(self): + return [ + OutputParam( + "prompts", + type_hint=list[str], + description="Ordered prompts for successive segments of one video, rather than a batch of independent videos.", + ), + OutputParam( + "prompt_embeds", + type_hint=list[torch.Tensor], + description="Qwen hidden states for each segment, shaped (1, max_sequence_length, hidden_size), with the chat prefix removed.", + ), + OutputParam( + "prompt_attention_masks", + type_hint=list[torch.Tensor], + description="Per-segment text masks shaped (1, max_sequence_length), with one for valid tokens and zero for padding.", + ), + OutputParam( + "action_text_offsets", + type_hint=list[list[int] | None], + description="Per-segment action-token boundaries in the cropped text sequence, or None for prompts without structured actions.", + ), + OutputParam( + "cache_action_starts", + type_hint=list[int | None], + description="Per-segment text boundary after which tokens are hidden from future segments; None keeps all valid text visible.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state) -> tuple[HyNomasModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else block_state.prompt + if not prompts or not all(isinstance(prompt, str) for prompt in prompts): + raise ValueError("prompt must be a string or a nonempty list of segment strings.") + if block_state.max_sequence_length < 1: + raise ValueError("max_sequence_length must be positive.") + tokenizer = components.tokenizer + marker = tokenizer("<|im_start|>user\n", add_special_tokens=False)["input_ids"] + ids = _tokenize(tokenizer, "")["input_ids"][0].tolist() + crop = next( + (i + len(marker) for i in range(len(ids) - len(marker) + 1) if ids[i : i + len(marker)] == marker), None + ) + if crop is None: + raise ValueError("The tokenizer chat template must contain the Qwen user marker.") + length = block_state.max_sequence_length + embeds, masks, offsets, starts = [], [], [], [] + for prompt in prompts: + tokens = _tokenize(tokenizer, prompt, length + crop).to(components._execution_device) + encoded = components.text_encoder(**tokens, output_hidden_states=True, use_cache=False) + embeds.append(encoded.hidden_states[-3][:, crop:]) + masks.append(tokens["attention_mask"][:, crop:]) + full_ids = tokens["input_ids"][0].tolist() + boundaries = [] + for variant in _action_variants(prompt): + variant_ids = _tokenize(tokenizer, variant, length + crop)["input_ids"][0].tolist() + common = next( + (i for i, (a, b) in enumerate(zip(full_ids, variant_ids)) if a != b), + min(len(full_ids), len(variant_ids)), + ) + boundaries.append(max(boundaries[-1] if boundaries else 0, min(length - 1, common - crop))) + starts.append(boundaries[0] if boundaries else None) + # The training mask's final boundary includes the uncropped valid length, capped at the context length. + offsets.append(boundaries + [min(length, int(tokens["attention_mask"].sum()))] if boundaries else None) + block_state.prompts = prompts + block_state.prompt_embeds = embeds + block_state.prompt_attention_masks = masks + block_state.action_text_offsets = offsets + block_state.cache_action_starts = starts + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/hy_nomas/modular_blocks_hy_nomas.py b/src/diffusers/modular_pipelines/hy_nomas/modular_blocks_hy_nomas.py new file mode 100644 index 000000000000..2dacfd437b14 --- /dev/null +++ b/src/diffusers/modular_pipelines/hy_nomas/modular_blocks_hy_nomas.py @@ -0,0 +1,123 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ..modular_pipeline import SequentialPipelineBlocks +from .before_denoise import HyNomasPrepareStep +from .decoders import HyNomasDecodeStep +from .denoise import HyNomasSegmentLoop +from .encoders import HyNomasTextEncoderStep + + +# auto_docstring +class HyNomasAutoBlocks(SequentialPipelineBlocks): + """ + Encode segment prompts, generate world latents with camera control, and decode video using PixelUMM. + + Components: + text_encoder (`Qwen2_5_VLTextModel`) tokenizer (`AutoTokenizer`) transformer (`HyNomasTransformer3DModel`) + scheduler (`FlowMatchEulerDiscreteScheduler`) decoder (`PixelUMMDecoder`) video_processor (`VideoProcessor`) + + Inputs: + prompt (`str | list`): + One prompt or an ordered list of prompts describing successive segments of a single video. + max_sequence_length (`int`, *optional*, defaults to 1280): + Number of text tokens retained per segment after removing the Qwen chat prefix; shorter sequences are + padded. + height (`int`, *optional*, defaults to 704): + Output height in pixels; must be a positive multiple of 32 and match the existing session. + width (`int`, *optional*, defaults to 1280): + Output width in pixels; must be a positive multiple of 32 and match the existing session. + num_frames (`int`, *optional*, defaults to 97): + Exact number of new pixel frames to generate in this call, excluding all prior session history. + segment_latent_lengths (`list`, *optional*): + Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4). + camera_extrinsics (`Tensor`, *optional*): + Normalized camera-to-world matrices in the training coordinate system, shaped (1, num_frames, 4, 4), for + the new rollout only. + fps (`float`, *optional*, defaults to 24.0): + Frames per second used to convert structured action timestamps into latent-frame intervals. + generator (`Generator`, *optional*): + Generator for rollout noise and cache-write noise; omitted values restore the session generator or start + from seed 42. + latents (`Tensor`, *optional*): + Optional initial noise for the new rollout, shaped (1, channels, ceil(num_frames / 4), height / 32, width + / 32). + noise_capacity (`int`, *optional*): + Total latent-frame capacity reserved when creating a fixed session noise timeline, including future + continuation segments. + session (`HyNomasGenerationState`, *optional*): + Earlier segment history to continue or branch; this call forks the state before appending new cache + entries and latents. + num_inference_steps (`int`, *optional*, defaults to 8): + Number of flow-matching Euler denoising steps for each new segment. + cache_noise_sigma (`float`, *optional*, defaults to 0.206897): + Noise level in [0, 1] applied to a completed segment before committing its transformer keys and values to + history. + output_type (`str`, *optional*, defaults to np): + Output representation: np for NumPy arrays, pt for tensors, pil for frame images, or latent to skip + decoding. + decoder_generator (`Generator`, *optional*): + Independent PixelUMM noise generator; only its initial seed is used, without advancing its state. If + omitted, the decoder configuration seed is used (42 by default). + + Outputs: + prompts (`list`): + Ordered prompts for successive segments of one video, rather than a batch of independent videos. + prompt_embeds (`list`): + Qwen hidden states for each segment, shaped (1, max_sequence_length, hidden_size), with the chat prefix + removed. + prompt_attention_masks (`list`): + Per-segment text masks shaped (1, max_sequence_length), with one for valid tokens and zero for padding. + action_text_offsets (`list`): + Per-segment action-token boundaries in the cropped text sequence, or None for prompts without structured + actions. + cache_action_starts (`list`): + Per-segment text boundary after which tokens are hidden from future segments; None keeps all valid text + visible. + session (`HyNomasGenerationState`): + Forked generation state containing prior segment history and the optional fixed noise timeline. + initial_latents (`Tensor`): + Initial noise for the new rollout, shaped (1, channels, new_latent_frames, height / 32, width / 32). + segment_latent_lengths (`list`): + Number of latent frames in each new segment, one per prompt; the total must equal ceil(num_frames / 4). + segment_cameras (`list`): + Normalized camera-to-world poses for each segment, shaped (1, latent_frames * 4, 4, 4), or None without + camera conditioning. + action_frame_windows (`list`): + Per-segment (start, end, action_index) intervals in latent frames; -1 denotes an action gap, and None + disables action-specific windows. + generator (`Generator`): + Generator initialized or restored for this rollout and its cache-write noise. + segment_pixel_lengths (`list`): + Exact visible pixel-frame count of each new segment after removing decoder padding from the final + segment. + segment_latents (`Tensor`): + Denoised latents for the current segment, shaped (1, channels, segment_latent_frames, latent_height, + latent_width). + latents (`Tensor`): + Concatenated denoised latents for the complete world history, including prior calls and every new + segment. + videos (`Tensor | ndarray | list`): + Complete accumulated video including prior history: pt has shape (1, frames, 3, height, width), np has + shape (1, frames, height, width, 3), and pil is a nested list of frames. The latent output preserves the + input latent tensor. + """ + + model_name = "hy-nomas" + block_classes = [HyNomasTextEncoderStep, HyNomasPrepareStep, HyNomasSegmentLoop, HyNomasDecodeStep] + block_names = ["text_encoder", "prepare", "denoise", "decode"] + + @property + def description(self): + return "Encode segment prompts, generate world latents with camera control, and decode video using PixelUMM." diff --git a/src/diffusers/modular_pipelines/hy_nomas/modular_pipeline.py b/src/diffusers/modular_pipelines/hy_nomas/modular_pipeline.py new file mode 100644 index 000000000000..cb0b45a6b841 --- /dev/null +++ b/src/diffusers/modular_pipelines/hy_nomas/modular_pipeline.py @@ -0,0 +1,57 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field, replace + +import torch + +from ..modular_pipeline import ModularPipeline + + +@dataclass +class HyNomasGenerationState: + """History for one world. Pass the returned state as `session` to continue generation. + + Tensors belong to the execution device. A pipeline call forks the state before writing new cache entries, so an + earlier state can also be reused to branch the world. Do not mutate tensors in a state while it is in use. + """ + + kv_cache: list[dict[str, torch.Tensor]] = field(default_factory=list) + latent_chunks: list[torch.Tensor] = field(default_factory=list) + pixel_lengths: list[int] = field(default_factory=list) + frame_offset: int = 0 + height: int = 0 + width: int = 0 + noise_bank: torch.Tensor | None = None + generator_state: torch.Tensor | None = None + generator_device: str = "cpu" + + def fork(self): + return replace( + self, + kv_cache=[dict(entry) for entry in self.kv_cache], + latent_chunks=list(self.latent_chunks), + pixel_lengths=list(self.pixel_lengths), + ) + + +class HyNomasModularPipeline(ModularPipeline): + """Modular inference for the distilled Hy-Nomas text and camera conditioned world model. + + A list of prompts describes successive segments of a single video. `segment_latent_lengths` assigns their + durations; each latent frame corresponds to four pixel frames. `num_frames` is the length of the new rollout. + Camera matrices are normalized camera-to-world extrinsics in the training coordinate system. + """ + + default_blocks_name = "HyNomasAutoBlocks" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 28e485eddbc1..7d338eefa86d 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -209,6 +209,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("z-image", _create_default_map_fn("ZImageModularPipeline")), ("cosmos3-omni", _create_default_map_fn("Cosmos3OmniModularPipeline")), ("helios", _create_default_map_fn("HeliosModularPipeline")), + ("hy-nomas", _create_default_map_fn("HyNomasModularPipeline")), ("helios-pyramid", _helios_pyramid_map_fn), ("hunyuan-video-1.5", _create_default_map_fn("HunyuanVideo15ModularPipeline")), ("echo", _create_default_map_fn("EchoModularPipeline")), diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index c25edb7f7c52..a21f438b2495 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -91,6 +91,7 @@ is_kernels_version, is_kornia_available, is_librosa_available, + is_magi_attention_available, is_matplotlib_available, is_nltk_available, is_note_seq_available, diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 22e653a403cc..856c476fe23d 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1594,6 +1594,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class HyNomasTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class I2VGenXLUNet(metaclass=DummyObject): _backends = ["torch"] @@ -2164,6 +2179,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class PixelUMMDecoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class PriorTransformer(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index c9bd3b496d60..10678c4ccc93 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -482,6 +482,51 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class HyNomasAutoBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class HyNomasGenerationState(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class HyNomasModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Ideogram4AutoBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index c0752f36d503..a5b96282dac1 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -213,6 +213,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b _nltk_available, _nltk_version = _is_package_available("nltk") _cosmos_guardrail_available, _cosmos_guardrail_version = _is_package_available("cosmos_guardrail") _sageattention_available, _sageattention_version = _is_package_available("sageattention") +_magi_attention_available = _is_package_available("magi_attention")[0] _flash_attn_available, _flash_attn_version = _is_package_available("flash_attn") _flashinfer_available, _flashinfer_version = _is_package_available("flashinfer", get_dist_name=True) _flash_attn_3_available, _flash_attn_3_version = _is_package_available("flash_attn_3") @@ -416,6 +417,10 @@ def is_sageattention_available(): return _sageattention_available +def is_magi_attention_available(): + return _magi_attention_available + + def is_flash_attn_available(): return _flash_attn_available diff --git a/tests/models/autoencoders/test_models_pixelumm.py b/tests/models/autoencoders/test_models_pixelumm.py new file mode 100644 index 000000000000..4b490ed62101 --- /dev/null +++ b/tests/models/autoencoders/test_models_pixelumm.py @@ -0,0 +1,245 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import PixelUMMDecoder +from diffusers.models.autoencoders.autoencoder_pixelumm import local_attention_layout +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import ( + enable_full_determinism, + is_torch_compile, + require_torch_gpu, + require_torch_multi_gpu, + torch_device, +) +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class PixelUMMDecoderTesterConfig(BaseModelTesterConfig): + main_input_name = "sample" + # The tiny model is dominated by the fixed RGB projections; reserve room for their offload staging. + model_split_percents = [0.95, 0.97] + + @property + def model_class(self): + return PixelUMMDecoder + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self): + return { + "hidden_size": 16, + # Keep the blocks large enough for balanced placement to split the fixed-size RGB projections. + "intermediate_size": 512, + "num_layers": 2, + "num_heads": 2, + "num_kv_heads": 1, + "head_dim": 8, + "local_attention_tile_pixels": 0, + "temporal_window_latents": 3, + "temporal_overlap_latents": 1, + } + + def get_dummy_inputs(self, height=1, width=1): + return { + "latent": randn_tensor((1, 32, 1, height, width), generator=self.generator, device=torch_device), + "sample": randn_tensor((1, 3, 4, height * 32, width * 32), generator=self.generator, device=torch_device), + } + + @property + def input_shape(self): + return (3, 4, 32, 32) + + @property + def output_shape(self): + return (3, 4, 32, 32) + + +class TestPixelUMMDecoderModel(PixelUMMDecoderTesterConfig, ModelTesterMixin): + @torch.no_grad() + def test_decode_generator_and_crop(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + latent = self.get_dummy_inputs()["latent"] + first = model.decode(latent, generator=torch.Generator().manual_seed(7)).sample + second = model.decode(latent, generator=torch.Generator().manual_seed(7), pixel_frames=3, return_dict=False)[0] + torch.testing.assert_close(first[:, :, :3], second, atol=0, rtol=0) + other = model.decode(latent, generator=torch.Generator().manual_seed(8)).sample + assert not torch.equal(first, other) + large_seed = model.decode(latent, generator=torch.Generator().manual_seed(2**64 - 1)).sample + assert large_seed.isfinite().all() + with pytest.raises(ValueError, match="pixel_frames"): + model.decode(latent, pixel_frames=5) + + @torch.no_grad() + def test_decode_without_temporal_or_spatial_overlap(self): + config = self.get_init_dict() + config.update( + temporal_window_latents=2, + temporal_overlap_latents=0, + decoder_tile_pixels=32, + decoder_tile_overlap_pixels=0, + ) + model = self.model_class(**config).to(torch_device).eval() + latent = randn_tensor((1, 32, 4, 2, 2), generator=self.generator, device=torch_device) + first = model.decode(latent).sample + assert first.shape == (1, 3, 16, 64, 64) + assert first.isfinite().all() + latent[:, :, 2:, 1:, 1:] += 5 + second = model.decode(latent).sample + torch.testing.assert_close(first[:, :, :8], second[:, :, :8], atol=0, rtol=0) + torch.testing.assert_close(first[..., :32, :32], second[..., :32, :32], atol=0, rtol=0) + assert not torch.equal(first[:, :, 8:, 32:, 32:], second[:, :, 8:, 32:, 32:]) + + @torch.no_grad() + def test_streaming_matches_whole_clip_with_segment_padding(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + latent = randn_tensor((1, 32, 7, 1, 1), generator=self.generator, device=torch_device) + complete = model.decode(latent).sample + session = model.create_streaming_decode_session(to_cpu=False) + first = session.push_preview(latent[:, :, :4], pixel_frames=15) + second = session.push_preview(latent[:, :, 4:], pixel_frames=10) + preview = torch.cat((first.video[:, :, : second.pixel_start], second.video), dim=2) + expected = torch.cat((complete[:, :, :15], complete[:, :, 16:26]), dim=2) + torch.testing.assert_close(preview, expected, atol=1e-5, rtol=1e-5) + session.finish() + torch.testing.assert_close(session.full_video(), expected, atol=1e-5, rtol=1e-5) + + @torch.no_grad() + def test_streaming_tail_correction_and_rollback(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + latent = randn_tensor((1, 32, 8, 1, 1), generator=self.generator, device=torch_device) + session = model.create_streaming_decode_session(to_cpu=False) + session.final_segment = False + first = session.push_preview(latent[:, :, :4]) + session.push_preview(latent[:, :, 4:6]) + session.rollback_last_push() + session.final_segment = True + update = session.push_preview(latent[:, :, 4:]) + preview = torch.cat((first.video[:, :, : update.pixel_start], update.video), dim=2) + expected = model.decode(latent).sample + torch.testing.assert_close(preview, expected, atol=1e-5, rtol=1e-5) + session.finish() + torch.testing.assert_close(session.full_video(), expected, atol=1e-5, rtol=1e-5) + with pytest.raises(RuntimeError, match="after finish"): + session.push(latent) + + @torch.no_grad() + def test_local_attention_excludes_distant_spatial_tiles(self): + config = self.get_init_dict() + config["local_attention_tile_pixels"] = 16 + config["num_layers"] = 1 + config["head_dim"] = 16 + model = self.model_class(**config).to(torch_device).eval() + inputs = self.get_dummy_inputs(height=4, width=4) + first = model(**inputs).sample + inputs["sample"] = inputs["sample"].clone() + inputs["sample"][..., -16:, -16:] += 5 + second = model(**inputs).sample + torch.testing.assert_close(first[..., :16, :16], second[..., :16, :16], atol=0, rtol=0) + assert not torch.equal(first[..., -16:, -16:], second[..., -16:, -16:]) + + @require_torch_gpu + @pytest.mark.parametrize("grid,tile", [((3, 8, 8), 2), ((2, 9, 5), 2), ((2, 2, 9), 2), ((1, 9, 2), 2)]) + def test_local_attention_mask_matches_tile_neighborhoods(self, grid, tile): + _, _, (neighborhoods, block_mask) = local_attention_layout(grid, tile, torch_device) + count = grid[0] * grid[1] * grid[2] + indices = torch.arange(count, device=torch_device) + actual = block_mask.mask_mod(None, None, indices[:, None], indices[None, :]) + expected = torch.zeros((count, count), dtype=torch.bool, device=torch_device) + for (start, end), neighbors in neighborhoods: + for first, last in neighbors: + expected[start:end, first:last] = True + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @torch.no_grad() + def test_multistep_decode(self): + config = self.get_init_dict() + config["num_steps"] = 2 + model = self.model_class(**config).to(torch_device).eval() + output = model.decode(self.get_dummy_inputs()["latent"]).sample + assert output.shape == (1, *self.output_shape) + assert output.isfinite().all() + with pytest.raises(ValueError, match="num_steps=1"): + model.create_streaming_decode_session() + + +class TestPixelUMMDecoderMemory(PixelUMMDecoderTesterConfig, MemoryTesterMixin): + @require_torch_multi_gpu + @pytest.mark.parametrize("local_attention_tile_pixels,num_steps", [(0, 1), (0, 2), (16, 1)]) + @torch.no_grad() + def test_device_map_embeddings_and_local_layout(self, local_attention_tile_pixels, num_steps, tmp_path): + torch.manual_seed(0) + config = self.get_init_dict() + config.update( + hidden_size=32, head_dim=16, local_attention_tile_pixels=local_attention_tile_pixels, num_steps=num_steps + ) + model = self.model_class(**config).to("cuda:0").eval() + inputs = self.get_dummy_inputs(height=2, width=2) + expected = model(**inputs).sample + expected_decoded = model.decode(inputs["latent"]).sample + model.cpu().save_pretrained(tmp_path) + device_map = { + "decoder.embed": "cuda:1", + "decoder.cond_proj": "cuda:0", + "decoder.layers.0": "cuda:0", + "decoder.layers.1": "cuda:1", + "decoder.final_norm": "cuda:0", + "decoder.head": "cuda:1", + } + loaded = self.model_class.from_pretrained(tmp_path, device_map=device_map) + assert loaded.decoder.embed.weight.device == torch.device("cuda:1") + assert loaded.decoder.cond_proj.weight.device == torch.device("cuda:0") + torch.testing.assert_close(loaded(**inputs).sample, expected, atol=1e-5, rtol=0) + actual_decoded = loaded.decode(inputs["latent"]).sample + torch.testing.assert_close(actual_decoded.cpu(), expected_decoded.cpu(), atol=1e-5, rtol=0) + + +class TestPixelUMMDecoderTorchCompile(PixelUMMDecoderTesterConfig, TorchCompileTesterMixin): + @is_torch_compile + @require_torch_gpu + @torch.no_grad() + def test_compile_local_attention(self): + config = self.get_init_dict() + config.update(local_attention_tile_pixels=16, num_layers=1, head_dim=16) + model = self.model_class(**config).to(torch_device).eval() + inputs = self.get_dummy_inputs(height=2, width=2) + expected = model(**inputs).sample + compiled = torch.compile(model, fullgraph=True) + actual = compiled(**inputs).sample + torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4) + + +class TestPixelUMMDecoderTraining(PixelUMMDecoderTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + super().test_gradient_checkpointing_is_applied(expected_set={"PixelUMMDecoderModel"}) + + +class TestPixelUMMDecoderAttention(PixelUMMDecoderTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index 478fb2853c0f..4e143fe7e8a8 100644 --- a/tests/models/test_attention_dispatch.py +++ b/tests/models/test_attention_dispatch.py @@ -22,9 +22,11 @@ import torch.multiprocessing as mp import torch.nn.functional as F +from diffusers.models import attention_dispatch from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig from diffusers.models.attention_dispatch import attention_backend as attention_backend_ctx from diffusers.models.attention_dispatch import dispatch_attention_fn +from diffusers.utils import is_magi_attention_available from ..testing_utils import ( is_attention, @@ -42,6 +44,60 @@ GRAD_RTOL = 2e-2 +@is_attention +class TestMagiAttentionBackend: + def test_missing_dependency(self, monkeypatch): + monkeypatch.setattr(attention_dispatch, "_CAN_USE_MAGI_ATTN", False) + with pytest.raises(RuntimeError, match="optional `magi_attention`"): + with attention_backend_ctx("magi"): + pass + query = torch.zeros(1, 2, 1, 64, dtype=torch.bfloat16) + ranges = torch.tensor([[0, 2]], dtype=torch.int32) + with pytest.raises(RuntimeError, match="optional `magi_attention`"): + dispatch_attention_fn( + query, + query, + query, + backend="magi", + attention_kwargs={"query_ranges": ranges, "key_ranges": ranges}, + ) + + @require_torch_accelerator + @pytest.mark.skipif( + not is_magi_attention_available() or torch_device != "cuda", reason="MagiAttention requires CUDA." + ) + @torch.no_grad() + def test_batched_rectangle_visibility(self): + torch.manual_seed(0) + query = torch.randn(2, 4, 2, 64, device=torch_device, dtype=torch.bfloat16) + key = torch.randn(2, 5, 2, 64, device=torch_device, dtype=torch.bfloat16) + value = torch.randn_like(key) + actual = dispatch_attention_fn( + query, + key, + value, + backend="magi", + attention_kwargs={ + "query_ranges": torch.tensor([[0, 2], [2, 4]], device=torch_device, dtype=torch.int32), + "key_ranges": torch.tensor([[0, 3], [1, 5]], device=torch_device, dtype=torch.int32), + "attn_type_map": torch.zeros(2, device=torch_device, dtype=torch.int32), + "deterministic": True, + }, + ) + expected = torch.cat( + [ + F.scaled_dot_product_attention( + query[:, start:end].float().transpose(1, 2), + key[:, first:last].float().transpose(1, 2), + value[:, first:last].float().transpose(1, 2), + ).transpose(1, 2) + for start, end, first, last in [(0, 2, 0, 3), (2, 4, 1, 5)] + ], + dim=1, + ).to(query.dtype) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + + def _attention_backward_parity_worker(rank, world_size, master_port, cp_dict, attention_backend, return_dict): """Op-level worker: check `dispatch_attention_fn` gradients against a single-process reference. diff --git a/tests/models/transformers/test_models_transformer_hy_nomas.py b/tests/models/transformers/test_models_transformer_hy_nomas.py new file mode 100644 index 000000000000..4b586d85f4a2 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_hy_nomas.py @@ -0,0 +1,404 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import math + +import pytest +import torch + +from diffusers import HyNomasTransformer3DModel +from diffusers.models.attention_dispatch import attention_backend +from diffusers.models.transformers.transformer_hy_nomas import ( + HyNomasAttnProcessor, + HyNomasMagiAttnProcessor, + HyNomasRotaryPosEmbed, + HyNomasTimestepEmbedder, + HyNomasTokenRefiner, +) +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import ( + enable_full_determinism, + is_torch_compile, + require_torch_gpu, + require_torch_multi_gpu, + torch_device, +) +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class HyNomasTransformerTesterConfig(BaseModelTesterConfig): + main_input_name = "hidden_states" + + @property + def model_class(self): + return HyNomasTransformer3DModel + + def get_init_dict(self): + return { + "in_channels": 4, + "out_channels": 4, + "num_layers": 2, + "num_attention_heads": 2, + "attention_head_dim": 12, + "text_embed_dim": 16, + "mlp_ratio": 2.0, + "rope_axes_dim": (4, 4, 4), + "camera_embed_dim": 8, + "camera_attention_dim": 24, + "camera_num_attention_heads": 2, + "camera_rope_axes_dim": (4, 4, 4), + } + + def get_dummy_inputs(self, height=2, width=2): + generator = torch.Generator("cpu").manual_seed(0) + return { + "hidden_states": randn_tensor((1, 4, 2, height, width), generator=generator, device=torch_device), + "encoder_hidden_states": randn_tensor((1, 5, 16), generator=generator, device=torch_device), + "timestep": torch.tensor([500.0], device=torch_device), + "camera_extrinsics": torch.eye(4, device=torch_device).repeat(1, 5, 1, 1), + } + + @property + def input_shape(self): + return (4, 2, 2, 2) + + @property + def output_shape(self): + return (4, 2, 2, 2) + + +class TestHyNomasTransformerModel(HyNomasTransformerTesterConfig, ModelTesterMixin): + @pytest.mark.parametrize("reset_method", ["disable_reference_attention", "reset_attention_backend", "native"]) + @torch.no_grad() + def test_reference_attention_can_restore_standard_inference(self, monkeypatch, reset_method): + monkeypatch.setattr("diffusers.models.attention_dispatch._CAN_USE_FLASH_ATTN", True) + monkeypatch.setattr("diffusers.models.attention_dispatch._CAN_USE_FLASH_ATTN_3", True) + monkeypatch.setattr("diffusers.models.attention_dispatch._CAN_USE_MAGI_ATTN", True) + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + original_processors = {name: type(processor) for name, processor in model.attn_processors.items()} + inputs = self.get_dummy_inputs() + with attention_backend("native"): + expected = model(**inputs).sample + model.enable_reference_attention(deterministic=True) + assert all( + type(processor) != original_processors[name] for name, processor in model.attn_processors.items() + ) + if reset_method == "native": + model.set_attention_backend("native") + else: + getattr(model, reset_method)() + assert {name: type(processor) for name, processor in model.attn_processors.items()} == original_processors + torch.testing.assert_close(model(**inputs).sample, expected, atol=0, rtol=0) + + @pytest.mark.parametrize( + "unavailable_backend,error", + [ + ("_CAN_USE_FLASH_ATTN", "Flash Attention backend"), + ("_CAN_USE_FLASH_ATTN_3", "Flash Attention 3 backend"), + ("_CAN_USE_MAGI_ATTN", "Magi Attention backend"), + ], + ) + def test_reference_attention_reports_missing_dependency(self, monkeypatch, unavailable_backend, error): + monkeypatch.setattr("diffusers.models.attention_dispatch._CAN_USE_FLASH_ATTN", True) + monkeypatch.setattr("diffusers.models.attention_dispatch._CAN_USE_FLASH_ATTN_3", True) + monkeypatch.setattr("diffusers.models.attention_dispatch._CAN_USE_MAGI_ATTN", True) + monkeypatch.setattr(f"diffusers.models.attention_dispatch.{unavailable_backend}", False) + model = self.model_class(**self.get_init_dict()) + processors = model.attn_processors + with pytest.raises(RuntimeError, match=error): + model.enable_reference_attention() + assert model.attn_processors == processors + + @torch.no_grad() + def test_reference_action_ranges_match_dense_visibility(self, monkeypatch): + def dense_range_attention(query, key, value, backend, attention_kwargs): + assert backend == "magi" + mask = torch.zeros(query.shape[1], key.shape[1], dtype=torch.bool) + for (start, end), (key_start, key_end) in zip( + attention_kwargs["query_ranges"], attention_kwargs["key_ranges"] + ): + mask[start:end, key_start:key_end] = True + return torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), attn_mask=mask + ).transpose(1, 2) + + monkeypatch.setattr( + "diffusers.models.transformers.transformer_hy_nomas.dispatch_attention_fn", dense_range_attention + ) + generator = torch.Generator().manual_seed(17) + query = torch.randn(1, 10, 2, 8, generator=generator) + key = torch.randn(1, 13, 2, 8, generator=generator) + value = torch.randn(1, 13, 2, 8, generator=generator) + text_mask = torch.tensor([[True, True, True, True, True, False]]) + key_mask = torch.tensor([[True, False, True, True, True, True, True, True, True, True, True, True, False]]) + # History 0/2, current video 3:7, head text 7:9, actions 9 and 10:12. + expected_mask = torch.zeros(10, 13, dtype=torch.bool) + expected_mask[:9, [0, 2, 3, 4, 5, 6, 7, 8]] = True + expected_mask[[0, 1, 6], 9] = True + expected_mask[[2, 3, 7, 8], 10:12] = True + expected = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), attn_mask=expected_mask + ).transpose(1, 2) + actual = HyNomasMagiAttnProcessor()._compute_attention( + query, key, value, key_mask, text_mask, 3, 4, 6, text_mask, [[2, 3, 5]], [(0, 2, 0), (2, 4, 1)] + ) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) + + @require_torch_gpu + @torch.no_grad() + def test_bfloat16_preserves_autocast_normalization_precision(self): + model = self.model_class(**self.get_init_dict()).to(torch_device, dtype=torch.bfloat16).eval() + inputs = { + name: value.to(torch.bfloat16) if value.is_floating_point() else value + for name, value in self.get_dummy_inputs().items() + } + # Original inference uses CUDA autocast. Native inference must preserve its float32 + # LayerNorm/modulation boundary without requiring callers to enable autocast. + with torch.autocast("cuda", dtype=torch.bfloat16): + expected = model(**inputs).sample + actual = model(**inputs).sample + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + @require_torch_gpu + @torch.no_grad() + def test_masked_context_preserves_autocast_reduction_precision(self): + refiner = HyNomasTokenRefiner(16, 24, 2).to(torch_device, dtype=torch.bfloat16).eval() + generator = torch.Generator("cpu").manual_seed(42) + hidden_states = randn_tensor((1, 1280, 16), generator=generator, device=torch_device).to(torch.bfloat16) * 100 + mask = torch.arange(1280, device=torch_device)[None] < 251 + captured = [] + hook = refiner.c_embedder.register_forward_pre_hook(lambda module, args: captured.append(args[0].clone())) + refiner(hidden_states, torch.tensor([1000.0], device=torch_device), mask) + hook.remove() + # CUDA autocast promotes sum, including its division denominator, to float32. The linear + # layer then casts the quotient once. Reducing directly in bfloat16 loses information first. + with torch.autocast("cuda", dtype=torch.bfloat16): + weights = mask.to(hidden_states.dtype).unsqueeze(-1) + expected = (hidden_states * weights).sum(1) / weights.sum(1) + torch.testing.assert_close(captured[0], expected.to(torch.bfloat16), atol=0, rtol=0) + + @require_torch_gpu + @torch.no_grad() + def test_timestep_frequencies_survive_weight_cast(self): + embedder = HyNomasTimestepEmbedder(24).to(torch_device, dtype=torch.bfloat16) + captured = [] + hook = embedder.mlp[0].register_forward_pre_hook(lambda module, args: captured.append(args[0].clone())) + timestep = torch.tensor([1000.0, 916.6667, 500.0], device=torch_device) + embedder(timestep, torch.bfloat16) + hook.remove() + # The published sinusoid uses CPU float32 frequencies before its GPU trigonometry. + frequencies = torch.exp(-math.log(10000.0) * torch.arange(128, dtype=torch.float32) / 128) + angles = timestep[:, None] * frequencies.to(torch_device)[None] + expected = torch.cat((angles.cos(), angles.sin()), -1).to(torch.bfloat16) + torch.testing.assert_close(captured[0], expected, atol=0, rtol=0) + + @require_torch_gpu + @pytest.mark.parametrize("frame_offset,num_frames", [(36, 72), (108, 30), (138, 72), (210, 36)]) + def test_continuation_rotary_uses_cpu_rounding(self, frame_offset, num_frames): + rotary = HyNomasRotaryPosEmbed((16, 56, 56), 256) + expected = rotary(num_frames, 22, 40, "cpu", frame_offset) + actual = rotary(num_frames, 22, 40, torch_device, frame_offset) + for expected_component, actual_component in zip(expected, actual): + torch.testing.assert_close(actual_component.cpu(), expected_component, atol=0, rtol=0) + + @torch.no_grad() + def test_cache_is_explicit_and_reads_do_not_mutate(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + cache = [{} for _ in range(model.config.num_layers)] + baseline = model(**inputs).sample + first = model(**inputs, kv_cache=cache, update_cache=True, cache_action_start=2).sample + torch.testing.assert_close(baseline, first) + assert cache[0]["key"].shape[1] == 13 + assert cache[0]["mask"][0, :10].all() + assert not cache[0]["mask"][0, 10:].any() + before = copy.deepcopy(cache) + continued = model(**inputs, kv_cache=cache, frame_offset=2).sample + assert not torch.equal(continued, baseline) + for previous, current in zip(before, cache): + for key in previous: + torch.testing.assert_close(previous[key], current[key]) + model(**inputs, kv_cache=cache, update_cache=True, frame_offset=2) + assert cache[0]["key"].shape[1] == 26 + torch.testing.assert_close(model(**inputs).sample, baseline) + + @torch.no_grad() + def test_camera_off_matches_no_camera(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + disabled = model(**inputs, camera_attention_mask=torch.zeros(1, device=torch_device)).sample + inputs.pop("camera_extrinsics") + torch.testing.assert_close(disabled, model(**inputs).sample) + + @torch.no_grad() + def test_text_padding_does_not_change_output(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs["encoder_attention_mask"] = torch.tensor([[1, 1, 1, 0, 0]], device=torch_device) + expected = model(**inputs).sample + inputs["encoder_hidden_states"][:, 3:] += 100 + torch.testing.assert_close(expected, model(**inputs).sample) + + @torch.no_grad() + def test_action_groups_match_unrestricted_attention_when_all_text_visible(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + expected = model(**inputs).sample + actual = model(**inputs, action_text_offsets=[[0, 5]], action_frame_windows=[(0, 2, 0)]).sample + torch.testing.assert_close(expected, actual, atol=1e-6, rtol=1e-5) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @torch.no_grad() + def test_action_attention_matches_explicit_visibility(self, dtype): + generator = torch.Generator().manual_seed(17) + query = torch.randn(1, 10, 2, 8, generator=generator, dtype=dtype) + key = torch.randn(1, 13, 2, 8, generator=generator, dtype=dtype) + value = torch.randn(1, 13, 2, 8, generator=generator, dtype=dtype) + text_mask = torch.tensor([[True, False, True, True, True, False]]) + key_mask = torch.cat((torch.tensor([[True, False, True, True, True, True, True]]), text_mask), dim=1) + # Visible history: 0/2; video: 3:7; head text: 7; actions: 9 and 10:12. + expected_mask = torch.zeros(10, 13, dtype=torch.bool) + expected_mask[:8, [0, 2, 3, 4, 5, 6, 7]] = True + expected_mask[[0, 1, 5], 9] = True + expected_mask[[2, 3, 6, 7], 10:12] = True + expected = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), attn_mask=expected_mask + ).transpose(1, 2) + actual = HyNomasAttnProcessor()._compute_attention( + query, key, value, key_mask, text_mask, 3, 4, 6, text_mask, [[1, 2, 6]], [(0, 2, 0), (2, 4, 1)] + ) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) + + @torch.no_grad() + def test_hidden_cached_text_has_no_effect(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + cache = [{} for _ in range(model.config.num_layers)] + model(**inputs, kv_cache=cache, update_cache=True, cache_action_start=2) + expected = model(**inputs, kv_cache=cache, frame_offset=2).sample + for layer in cache: + layer["key"][:, 10:] += 100 + layer["value"][:, 10:] += 100 + torch.testing.assert_close(expected, model(**inputs, kv_cache=cache, frame_offset=2).sample) + + def test_cache_validation(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + with pytest.raises(ValueError, match="caller-owned"): + model(**self.get_dummy_inputs(), update_cache=True) + with pytest.raises(ValueError, match="one dictionary"): + model(**self.get_dummy_inputs(), kv_cache=[{}]) + + +class TestHyNomasTransformerMemory(HyNomasTransformerTesterConfig, MemoryTesterMixin): + @require_torch_multi_gpu + @pytest.mark.parametrize("camera_device", [0, 1]) + @torch.no_grad() + def test_device_map_camera_residual(self, camera_device, tmp_path): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).to("cuda:0").eval() + inputs = self.get_dummy_inputs() + inputs.update( + encoder_attention_mask=torch.tensor([[1, 1, 1, 0, 0]], device="cuda:0", dtype=torch.bool), + action_text_offsets=[[1, 2, 3]], + action_frame_windows=[(0, 1, 0), (1, 2, 1)], + cache_action_start=1, + ) + expected_cache = [{} for _ in range(model.config.num_layers)] + expected = model(**inputs, kv_cache=expected_cache, update_cache=True).sample + expected_continuation = model(**inputs, kv_cache=expected_cache, frame_offset=2).sample + model.cpu().save_pretrained(tmp_path) + device_map = { + name: "cuda:0" for name, _ in model.named_children() if name not in ("double_blocks", "action_modules") + } + device_map["extrinsics_embedding_camctrl"] = f"cuda:{camera_device}" + device_map["c2ws_hidden_states_layer2"] = f"cuda:{1 - camera_device}" + device_map["final_layer"] = "cuda:1" + for index in range(model.config.num_layers): + device_map[f"double_blocks.{index}"] = f"cuda:{index % 2}" + device_map[f"action_modules.{index}"] = f"cuda:{1 - index % 2}" + loaded = self.model_class.from_pretrained(tmp_path, device_map=device_map) + assert loaded.extrinsics_embedding_camctrl.weight.device == torch.device(f"cuda:{camera_device}") + assert loaded.c2ws_hidden_states_layer2.weight.device == torch.device(f"cuda:{1 - camera_device}") + actual_cache = [{} for _ in range(model.config.num_layers)] + actual = loaded(**inputs, kv_cache=actual_cache, update_cache=True).sample + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=0) + actual_continuation = loaded(**inputs, kv_cache=actual_cache, frame_offset=2).sample + torch.testing.assert_close(actual_continuation, expected_continuation, atol=1e-5, rtol=0) + for expected_layer, actual_layer in zip(expected_cache, actual_cache): + for name in expected_layer: + torch.testing.assert_close(actual_layer[name].cpu(), expected_layer[name].cpu(), atol=1e-5, rtol=0) + + +class TestHyNomasTransformerTorchCompile(HyNomasTransformerTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(2, 2), (2, 4), (4, 4)] + + +@is_torch_compile +class TestHyNomasTransformerActionCompile(HyNomasTransformerTesterConfig): + @pytest.mark.parametrize("cache_mode", ["none", "read", "write"]) + @torch.no_grad() + def test_action_masks_and_history_compile(self, cache_mode): + torch.compiler.reset() + model = self.model_class(**self.get_init_dict()).cpu().eval() + inputs = {name: value.cpu() for name, value in self.get_dummy_inputs().items()} + history = None + if cache_mode != "none": + history = [{} for _ in range(model.config.num_layers)] + model(**inputs, kv_cache=history, update_cache=True, cache_action_start=1) + inputs.update( + action_text_offsets=[[1, 3, 5]], + action_frame_windows=[(0, 1, 0), (1, 2, 1)], + frame_offset=0 if history is None else 2, + update_cache=cache_mode == "write", + return_dict=False, + ) + compiled = torch.compile(model, backend="eager", fullgraph=True) + try: + with torch._dynamo.config.patch(error_on_recompile=True): + for mask in ([1, 1, 1, 0, 0], [1, 0, 1, 1, 0], [0, 0, 0, 0, 0], [1, 1, 1, 1, 1]): + inputs["encoder_attention_mask"] = torch.tensor([mask], dtype=torch.bool) + eager_cache, compiled_cache = copy.deepcopy(history), copy.deepcopy(history) + expected = model(**inputs, kv_cache=eager_cache)[0] + actual = compiled(**inputs, kv_cache=compiled_cache)[0] + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if history is not None: + for eager_layer, compiled_layer in zip(eager_cache, compiled_cache): + for key in eager_layer: + torch.testing.assert_close(compiled_layer[key], eager_layer[key], atol=0, rtol=0) + finally: + torch.compiler.reset() + + +class TestHyNomasTransformerTraining(HyNomasTransformerTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + super().test_gradient_checkpointing_is_applied(expected_set={"HyNomasTransformer3DModel"}) + + +class TestHyNomasTransformerAttention(HyNomasTransformerTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/modular_pipelines/hy_nomas/__init__.py b/tests/modular_pipelines/hy_nomas/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/hy_nomas/test_modular_pipeline_hy_nomas.py b/tests/modular_pipelines/hy_nomas/test_modular_pipeline_hy_nomas.py new file mode 100644 index 000000000000..5393f021b769 --- /dev/null +++ b/tests/modular_pipelines/hy_nomas/test_modular_pipeline_hy_nomas.py @@ -0,0 +1,339 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import tempfile + +import pytest +import torch +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import PreTrainedTokenizerFast, Qwen2_5_VLTextConfig, Qwen2_5_VLTextModel + +from diffusers import ( + FlowMatchEulerDiscreteScheduler, + HyNomasAutoBlocks, + HyNomasModularPipeline, + HyNomasTransformer3DModel, + ModularPipeline, + PixelUMMDecoder, +) +from diffusers.modular_pipelines.hy_nomas.before_denoise import HyNomasPrepareStep +from diffusers.modular_pipelines.hy_nomas.encoders import HyNomasTextEncoderStep + +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, + ModularWorkflowTesterMixin, +) + + +def make_pipeline(): + torch.manual_seed(0) + tokenizer_model = Tokenizer( + WordLevel( + { + "[UNK]": 0, + "[PAD]": 1, + "<|im_start|>": 2, + "<|im_end|>": 3, + "system": 4, + "user": 5, + "assistant": 6, + "A": 7, + "forest": 8, + "river": 9, + }, + unk_token="[UNK]", + ) + ) + tokenizer_model.pre_tokenizer = Whitespace() + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tokenizer_model, + unk_token="[UNK]", + pad_token="[PAD]", + additional_special_tokens=["<|im_start|>", "<|im_end|>"], + ) + tokenizer.chat_template = ( + "{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] " + "+ '<|im_end|>\\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}" + ) + torch.manual_seed(0) + text_encoder = Qwen2_5_VLTextModel( + Qwen2_5_VLTextConfig( + vocab_size=len(tokenizer), + bos_token_id=2, + eos_token_id=3, + pad_token_id=1, + hidden_size=16, + intermediate_size=24, + num_hidden_layers=3, + num_attention_heads=2, + num_key_value_heads=1, + rope_scaling={"type": "mrope", "mrope_section": [1, 1, 2]}, + ) + ).eval() + torch.manual_seed(0) + transformer = HyNomasTransformer3DModel( + in_channels=32, + out_channels=32, + num_layers=2, + num_attention_heads=2, + attention_head_dim=12, + text_embed_dim=16, + mlp_ratio=2, + rope_axes_dim=(4, 4, 4), + camera_embed_dim=8, + camera_attention_dim=24, + camera_num_attention_heads=2, + camera_rope_axes_dim=(4, 4, 4), + ).eval() + torch.manual_seed(0) + decoder = PixelUMMDecoder( + hidden_size=16, + intermediate_size=24, + num_layers=1, + num_heads=2, + num_kv_heads=1, + head_dim=8, + local_attention_tile_pixels=0, + temporal_window_latents=3, + temporal_overlap_latents=1, + ).eval() + pipe = HyNomasModularPipeline() + pipe.update_components( + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + decoder=decoder, + scheduler=FlowMatchEulerDiscreteScheduler(shift=5), + ) + return pipe + + +def inputs(**kwargs): + return { + "prompt": "A forest", + "height": 32, + "width": 32, + "num_frames": 5, + "max_sequence_length": 16, + "num_inference_steps": 2, + **kwargs, + } + + +@pytest.fixture(scope="module") +def tiny_hy_nomas_checkpoint(tmp_path_factory): + directory = tmp_path_factory.mktemp("tiny-hy-nomas-modular") + pipe = make_pipeline() + pipe.save_pretrained(str(directory), overwrite_modular_index=True) + return str(directory) + + +@pytest.fixture(scope="class", autouse=True) +def use_local_checkpoint(request, tiny_hy_nomas_checkpoint): + request.cls.pretrained_model_name_or_path = tiny_hy_nomas_checkpoint + + +class HyNomasModularPipelineTesterConfig(BaseModularPipelineTesterConfig): + pipeline_class = HyNomasModularPipeline + pipeline_blocks_class = HyNomasAutoBlocks + pretrained_model_name_or_path = None + expected_latents_shape = (1, 32, 2, 1, 1) + params = frozenset(["prompt", "height", "width", "num_frames", "camera_extrinsics", "session"]) + optional_params = frozenset(["num_inference_steps", "latents", "output_type"]) + batch_params = frozenset() + not_params = frozenset(["negative_prompt", "guidance_scale", "num_images_per_prompt"]) + output_name = "videos" + expected_workflow_blocks = { + None: [ + ("text_encoder", "HyNomasTextEncoderStep"), + ("prepare", "HyNomasPrepareStep"), + ("denoise", "HyNomasSegmentLoop"), + ("decode", "HyNomasDecodeStep"), + ], + } + + def get_dummy_inputs(self, seed=0): + # CPU normal sampling depends on dtype; precision tests must start from the same latent values. + latents = torch.randn(self.expected_latents_shape, generator=self.get_generator(seed), dtype=torch.float32) + return inputs(output_type="pt", latents=latents, generator=self.get_generator(seed)) + + +class TestHyNomasModularPipeline(HyNomasModularPipelineTesterConfig, ModularPipelineTesterMixin): + @pytest.mark.skip( + reason="Prompt lists are successive segments of one world; independent-world batching is unsupported." + ) + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip( + reason="Prompt lists are successive segments of one world; independent-world batching is unsupported." + ) + def test_inference_batch_single_identical(self): + pass + + def test_encode_denoise_decode(self): + pipe = self.get_pipeline() + result = pipe(**inputs(), output=["videos", "latents", "session"]) + assert result["latents"].shape == (1, 32, 2, 1, 1) + assert result["videos"].shape == (1, 5, 32, 32, 3) + assert result["videos"].min() >= 0 + assert result["videos"].max() <= 1 + assert result["session"].frame_offset == 2 + second = pipe(**inputs(), output="latents") + torch.testing.assert_close(result["latents"], second, atol=0, rtol=0) + + def test_sequential_calls_match_fixed_timeline(self): + pipe = self.get_pipeline() + kwargs = {"noise_capacity": 8, "output_type": "latent"} + combined = pipe( + **inputs(prompt=["A forest", "A river"], num_frames=13, segment_latent_lengths=[2, 2], **kwargs), + output=["latents", "session"], + ) + first = pipe(**inputs(**kwargs), output="session") + cache_before = first.kv_cache[0]["key"].clone() + continued = pipe(**inputs(prompt="A river", session=first, **kwargs), output=["latents", "session"]) + torch.testing.assert_close(combined["latents"], continued["latents"], atol=0, rtol=0) + torch.testing.assert_close(first.kv_cache[0]["key"], cache_before, atol=0, rtol=0) + assert first.frame_offset == 2 + assert continued["session"].frame_offset == 4 + assert combined["session"].kv_cache[0]["key"].shape == continued["session"].kv_cache[0]["key"].shape + + def test_save_reload(self): + pipe = self.get_pipeline() + pipe.update_components(scheduler=FlowMatchEulerDiscreteScheduler(shift=7)) + expected = pipe(**inputs(output_type="latent"), output="latents") + with tempfile.TemporaryDirectory() as directory: + pipe.save_pretrained(str(directory), overwrite_modular_index=True) + loaded = ModularPipeline.from_pretrained(directory) + loaded.load_components() + actual = loaded(**inputs(output_type="latent"), output="latents") + assert isinstance(loaded, HyNomasModularPipeline) + assert loaded.scheduler.config.shift == 7 + torch.testing.assert_close(expected, actual, atol=1e-6, rtol=1e-6) + + @pytest.mark.parametrize("num_frames", [4, 6]) + def test_arbitrary_pixel_lengths(self, num_frames): + result = self.get_pipeline()(**inputs(num_frames=num_frames), output="videos") + assert result.shape == (1, num_frames, 32, 32, 3) + + def test_continuation_keeps_segment_pixel_lengths(self): + pipe = self.get_pipeline() + first = pipe(**inputs(), output="session") + result = pipe(**inputs(prompt="A river", session=first), output=["videos", "session"]) + assert result["videos"].shape == (1, 10, 32, 32, 3) + assert result["session"].frame_offset == 4 + + def test_camera_changes_generation(self): + pipe = self.get_pipeline() + baseline = pipe(**inputs(output_type="latent"), output="latents") + camera = torch.eye(4).repeat(1, 5, 1, 1) + conditioned = pipe(**inputs(camera_extrinsics=camera, output_type="latent"), output="latents") + assert not torch.equal(conditioned, baseline) + + def test_structured_actions_flow_into_segment_cache(self): + prompt = json.dumps( + { + "segments": [ + { + "visual_description_long": "A forest", + "contained_actions": [ + {"start_time": 0, "end_time": 0.25, "action_description": "A forest"}, + {"start_time": 0.25, "end_time": 0.5, "action_description": "A river"}, + ], + } + ] + } + ) + pipe = self.get_pipeline() + result = pipe( + **inputs( + prompt=[prompt, "A river"], + num_frames=16, + segment_latent_lengths=[2, 2], + fps=16, + max_sequence_length=128, + output_type="latent", + ), + output=["latents", "session", "action_text_offsets", "cache_action_starts", "action_frame_windows"], + ) + assert torch.isfinite(result["latents"]).all() + assert result["action_frame_windows"][0] == [(0, 1, 0), (1, 2, 1)] + assert len(result["action_text_offsets"][0]) == 3 + assert result["action_text_offsets"][1] is None + boundary = result["cache_action_starts"][0] + assert 0 < boundary < 128 + for layer in result["session"].kv_cache: + assert layer["mask"][0, : 2 + boundary].all() + assert not layer["mask"][0, 2 + boundary : 130].any() + + def test_text_encoder_is_reusable(self): + pipe = self.get_pipeline() + encoder = HyNomasTextEncoderStep().init_pipeline() + encoder.update_components(text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer) + encoded = encoder(prompt="A forest", max_sequence_length=16, output="prompt_embeds") + full = pipe(**inputs(output_type="latent"), output="prompt_embeds") + torch.testing.assert_close(encoded[0], full[0], atol=0, rtol=0) + + @pytest.mark.parametrize( + "kwargs,match", + [ + ({"height": 31}, "multiples of 32"), + ({"num_frames": 0}, "positive"), + ({"segment_latent_lengths": [1]}, "sum"), + ({"noise_capacity": 1}, "noise_capacity"), + ({"camera_extrinsics": torch.eye(4)}, "camera_extrinsics"), + ], + ) + def test_invalid_inputs(self, kwargs, match): + with pytest.raises(ValueError, match=match): + self.get_pipeline()(**inputs(**kwargs)) + + def test_action_windows(self): + prepare = HyNomasPrepareStep().init_pipeline() + prepare.update_components(transformer=self.get_pipeline().transformer) + + def run(actions): + return prepare( + prompts=[json.dumps({"segments": [{"contained_actions": actions}]})], + height=32, + width=32, + num_frames=32, + fps=16, + output="action_frame_windows", + )[0] + + actions = [{"start_time": 0, "end_time": 0.5}, {"start_time": 1, "end_time": 1.5}] + assert run(actions) == [(0, 2, 0), (2, 4, -1), (4, 6, 1), (6, 8, -1)] + assert run([]) is None + actions[1]["start_time"] = 0.25 + with pytest.raises(ValueError, match="overlap"): + run(actions) + + +class TestHyNomasModularPipelineLoading(HyNomasModularPipelineTesterConfig, ModularLoadingTesterMixin): + pass + + +class TestHyNomasModularPipelineWorkflow(HyNomasModularPipelineTesterConfig, ModularWorkflowTesterMixin): + pass + + +class TestHyNomasModularPipelineMemory(HyNomasModularPipelineTesterConfig, ModularMemoryTesterMixin): + pass diff --git a/tests/others/test_hy_nomas_keyboard_overlay.py b/tests/others/test_hy_nomas_keyboard_overlay.py new file mode 100644 index 000000000000..94a14720bb7d --- /dev/null +++ b/tests/others/test_hy_nomas_keyboard_overlay.py @@ -0,0 +1,214 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import hashlib +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + + +imageio = pytest.importorskip("imageio.v2") +pytest.importorskip("imageio_ffmpeg") + +from examples.community.hy_nomas_keyboard_overlay import ( # noqa: E402 + KEY_LAYOUT, + KeyboardTimeline, + draw_keyboard_overlay, + export_keyboard_video, +) + + +def make_turn(events=None, num_frames=60, fps=60, duration_ms=1000): + return { + "num_frames": num_frames, + "video_fps": fps, + "camera_trajectory": {"durationMs": duration_ms, "keyEvents": events or []}, + } + + +def test_subframe_events_simultaneous_keys_and_release(): + timeline = KeyboardTimeline( + [ + make_turn( + [ + {"t": 50, "code": "KeyW", "pressed": False}, + {"t": 16.3, "code": "KeyW", "pressed": True}, + {"t": 16.3, "code": "ArrowUp", "pressed": True}, + {"t": 66.7, "code": "ArrowUp", "pressed": False}, + ] + ) + ] + ) + assert timeline.state_at_frame(0) == (set(), "RECORDED KEY EVENTS") + assert timeline.state_at_frame(1)[0] == {"KeyW", "ArrowUp"} + assert timeline.state_at_frame(2)[0] == {"KeyW", "ArrowUp"} + assert timeline.state_at_frame(3)[0] == {"ArrowUp"} + assert timeline.state_at_frame(4)[0] == {"ArrowUp"} + assert timeline.state_at_frame(5)[0] == set() + + +def test_recording_duration_maps_to_video_clock_and_turn_boundary_resets(): + first = make_turn([{"t": 250, "code": "ArrowDown", "pressed": True}], num_frames=48, fps=24) + first["camera_trajectory"] = json.dumps(first["camera_trajectory"]) + second = make_turn(num_frames=24, fps=24) + second["camera_trajectory"]["frames"] = [{"position": {"x": 100, "y": 0, "z": 50}}] + timeline = KeyboardTimeline([first, second]) + assert timeline.state_at_frame(11)[0] == set() + assert timeline.state_at_frame(12)[0] == {"ArrowDown"} + assert timeline.state_at_frame(47)[0] == {"ArrowDown"} + assert timeline.state_at_frame(48) == (set(), "NO KEY EVENTS") + with pytest.raises(IndexError): + timeline.state_at_frame(72) + + +def test_missing_recording_duration_uses_video_fps(): + turn = make_turn([{"t": 100, "code": "KeyD", "pressed": True}], num_frames=24, fps=24) + del turn["camera_trajectory"]["durationMs"] + timeline = KeyboardTimeline([turn]) + assert timeline.state_at_frame(2)[0] == set() + assert timeline.state_at_frame(3)[0] == {"KeyD"} + + +@pytest.mark.parametrize("code", KEY_LAYOUT) +def test_each_key_highlights_only_its_own_box_without_modifying_input(code): + frame = np.full((352, 640, 3), 100, dtype=np.uint8) + original = frame.copy() + events = [{"t": 0, "code": code, "pressed": True}, {"t": 500, "code": code, "pressed": False}] + timeline = KeyboardTimeline([make_turn(events)]) + active = draw_keyboard_overlay(frame, *timeline.state_at_frame(0)) + idle = draw_keyboard_overlay(frame, *timeline.state_at_frame(30)) + difference = np.any(active != idle, axis=-1) + column, row = KEY_LAYOUT[code] + x, y = 8 + 10 + column * 28, 352 - 112 - 8 + 28 + row * 28 + assert difference[y : y + 25, x : x + 25].sum() > 300 + difference[y : y + 25, x : x + 25] = False + assert not difference.any() + np.testing.assert_array_equal(frame, original) + np.testing.assert_array_equal(active[:200], frame[:200]) + + +def test_four_arrows_are_distinct_vector_glyphs(): + frame = draw_keyboard_overlay(np.zeros((352, 640, 3), dtype=np.uint8), set(), "NO KEY EVENTS") + glyphs = [] + for code in ["ArrowUp", "ArrowDown", "ArrowLeft", "ArrowRight"]: + column, row = KEY_LAYOUT[code] + x, y = 8 + 10 + column * 28, 352 - 112 - 8 + 28 + row * 28 + glyph = frame[y + 5 : y + 20, x + 5 : x + 20, 0] > 150 + assert glyph.sum() > 40 + glyphs.append(glyph) + assert len({glyph.tobytes() for glyph in glyphs}) == 4 + + +@pytest.mark.parametrize( + "event", + [ + {"t": -1, "code": "KeyW", "pressed": True}, + {"t": float("nan"), "code": "KeyW", "pressed": True}, + {"t": 0, "code": "KeyW", "pressed": "false"}, + {"t": 0, "code": "Space", "pressed": True}, + ], +) +def test_invalid_recorded_events_raise(event): + with pytest.raises(ValueError): + KeyboardTimeline([make_turn([event])]) + + +@pytest.fixture +def original_video(tmp_path): + video = tmp_path / "original.mp4" + with imageio.get_writer(str(video), fps=4, macro_block_size=1, codec="libx264") as writer: + for _ in range(4): + writer.append_data(np.full((352, 640, 3), 100, dtype=np.uint8)) + return video + + +def test_cli_exports_real_events_to_separate_video(original_video, tmp_path): + metadata = tmp_path / "replay.json" + turn = make_turn( + [{"t": 250, "code": "ArrowDown", "pressed": True}, {"t": 750, "code": "ArrowDown", "pressed": False}], + num_frames=4, + fps=4, + ) + metadata.write_text(json.dumps({"sessions": [{"turns": [{"params": turn}]}]})) + output = tmp_path / "keyboard.mp4" + original_sha256 = hashlib.sha256(original_video.read_bytes()).hexdigest() + script = Path(__file__).resolve().parents[2] / "examples/community/hy_nomas_keyboard_overlay.py" + result = subprocess.run( + [ + sys.executable, + str(script), + "--video", + str(original_video), + "--manifest", + str(metadata), + "--output", + str(output), + ], + capture_output=True, + text=True, + check=True, + ) + report = json.loads(result.stdout) + assert report["num_frames"] == 4 + assert report["fps"] == 4 + assert report["recorded_key_events"] == 2 + assert hashlib.sha256(original_video.read_bytes()).hexdigest() == original_sha256 + with imageio.get_reader(str(output)) as reader: + assert reader.count_frames() == 4 + assert reader.get_meta_data()["fps"] == 4 + assert reader.get_meta_data()["size"] == (640, 352) + green_pixel_counts = [] + for frame in reader: + frame = np.asarray(frame) + green_pixel_counts.append(((frame[..., 1] > 180) & (frame[..., 0] < 120)).sum()) + assert green_pixel_counts[0] < 10 + assert green_pixel_counts[1] > 200 + assert green_pixel_counts[2] > 200 + assert green_pixel_counts[3] < 10 + + +def test_camera_json_without_events_stays_idle(original_video, tmp_path): + metadata = tmp_path / "camera.json" + metadata.write_text(json.dumps({"durationMs": 1000, "frames": [{"position": [1, 2, 3]}]})) + report = export_keyboard_video(original_video, metadata, tmp_path / "idle.mp4") + assert report["recorded_key_events"] == 0 + assert report["segments"][0]["key_source"] == "none" + + +def test_export_rejects_mismatched_manifest_before_creating_output(original_video, tmp_path): + metadata = tmp_path / "manifest.json" + metadata.write_text(json.dumps({"sessions": [{"turns": [{"params": make_turn()}]}]})) + output = tmp_path / "mismatched.mp4" + with pytest.raises(ValueError, match="does not match the manifest"): + export_keyboard_video(original_video, metadata, output) + assert not output.exists() + + +def test_export_never_overwrites_original_or_existing_output(tmp_path): + original = tmp_path / "original.mp4" + original.write_bytes(b"original") + alias = tmp_path / "alias.mp4" + alias.symlink_to(original) + with pytest.raises(ValueError, match="different file"): + export_keyboard_video(original, tmp_path / "unused.json", alias) + output = tmp_path / "existing.mp4" + output.write_bytes(b"existing") + with pytest.raises(FileExistsError): + export_keyboard_video(original, tmp_path / "unused.json", output) + assert original.read_bytes() == b"original" + assert output.read_bytes() == b"existing" From ec0886c0d4e6b071e58c0fd8bd797f724d8307a5 Mon Sep 17 00:00:00 2001 From: jiaxiangc Date: Fri, 23 Oct 2026 17:51:42 +0800 Subject: [PATCH 2/4] Update Hy-Nomas installation and inference examples --- .../en/api/models/hy_nomas_transformer3d.md | 2 +- docs/source/en/api/models/pixelumm_decoder.md | 4 +- docs/source/en/api/pipelines/hy_nomas.md | 229 ++++++++++++++---- setup.py | 7 + src/diffusers/dependency_versions_table.py | 2 + 5 files changed, 194 insertions(+), 50 deletions(-) diff --git a/docs/source/en/api/models/hy_nomas_transformer3d.md b/docs/source/en/api/models/hy_nomas_transformer3d.md index a621a6e62ca8..e28d561ea97d 100644 --- a/docs/source/en/api/models/hy_nomas_transformer3d.md +++ b/docs/source/en/api/models/hy_nomas_transformer3d.md @@ -20,7 +20,7 @@ import torch from diffusers import HyNomasTransformer3DModel model = HyNomasTransformer3DModel.from_pretrained( - "/path/to/hy-nomas-diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 + "tencent/Hy-Nomas", subfolder="transformer", torch_dtype=torch.bfloat16 ) ``` diff --git a/docs/source/en/api/models/pixelumm_decoder.md b/docs/source/en/api/models/pixelumm_decoder.md index 59f80cd1e95c..9814cd2039c3 100644 --- a/docs/source/en/api/models/pixelumm_decoder.md +++ b/docs/source/en/api/models/pixelumm_decoder.md @@ -19,9 +19,7 @@ It does not encode videos into latents. import torch from diffusers import PixelUMMDecoder -decoder = PixelUMMDecoder.from_pretrained( - "/path/to/hy-nomas-diffusers", subfolder="decoder", torch_dtype=torch.bfloat16 -) +decoder = PixelUMMDecoder.from_pretrained("tencent/Hy-Nomas", subfolder="decoder", torch_dtype=torch.bfloat16) ``` The decoder supports local spatial attention, overlapping temporal windows, optional spatial tiling, and explicit diff --git a/docs/source/en/api/pipelines/hy_nomas.md b/docs/source/en/api/pipelines/hy_nomas.md index 1541c540bb77..18cbceacf75a 100644 --- a/docs/source/en/api/pipelines/hy_nomas.md +++ b/docs/source/en/api/pipelines/hy_nomas.md @@ -12,7 +12,7 @@ specific language governing permissions and limitations under the License. --> # Hy-Nomas Hy-Nomas generates a world one video segment at a time, conditioned on text and optional camera extrinsics. Its -distilled checkpoint uses eight flow-matching Euler steps, a Qwen2.5-VL text encoder, and a PixelUMM diffusion decoder. +distilled checkpoint uses flow-matching Euler sampling, a Qwen2.5-VL text encoder, and a PixelUMM diffusion decoder. This implementation supports the released double-stream transformer, action-aware text visibility, camera attention, and continuation through an explicit generation state. Text prompts can be plain strings or the segment JSON produced @@ -23,12 +23,11 @@ by the model's prompt preprocessing. Action timestamps in JSON must already be r Install PyTorch 2.10.0 with CUDA support for your system, then run the following command from the Diffusers checkout: ```bash -python -m pip install -e . "transformers==5.20.2" "accelerate==1.13.0" imageio imageio-ffmpeg +pip install -e ".[hy-nomas]" ``` -The text encoder requires `Qwen2_5_VLTextModel`, provided by the Transformers version above. The video export example -also requires `imageio` and `imageio-ffmpeg`. The base `pip install -e .` command does not install these model-specific -dependencies. Use a converted model directory as described below; the original Hy-Nomas code checkout is not required. +The `hy-nomas` extra installs the Transformers text encoder, Accelerate loading support, and the `imageio` and +`imageio-ffmpeg` video export dependencies. The original Hy-Nomas code checkout is not required. ## Convert the checkpoints @@ -40,41 +39,147 @@ python scripts/convert_hy_nomas_to_diffusers.py \ --transformer_checkpoint /path/to/dit.safetensors \ --decoder_checkpoint /path/to/pixelumm.safetensors \ --text_encoder_path /path/to/Qwen2.5-VL-7B-Instruct \ - --output_path /path/to/hy-nomas-diffusers \ + --output_path /path/to/Hy-Nomas \ --dtype bfloat16 ``` The output contains `modular_model_index.json` and standard component subdirectories. The converted pipeline loads -without importing the original repository or executing remote code. +without importing the original repository or executing remote code. The examples below use `tencent/Hy-Nomas` as the +model ID; to use a local conversion, replace it with `/path/to/Hy-Nomas`. -## Generate video +## Generate video from a prompt + +Use a plain string to describe one segment. This example generates 120 frames at 24 FPS with four denoising steps and +returns a session for subsequent interaction. ```python import torch -from diffusers import ModularPipeline +from diffusers import HyNomasModularPipeline from diffusers.utils import export_to_video -pipe = ModularPipeline.from_pretrained("/path/to/hy-nomas-diffusers") +pipe = HyNomasModularPipeline.from_pretrained("tencent/Hy-Nomas") pipe.load_components(dtype=torch.bfloat16) pipe.to(device="cuda", dtype=torch.bfloat16) -result = pipe( - prompt=["A forest beside a river.", "The view moves along the river toward a waterfall."], +first = pipe( + prompt=( + "A small orange fox stands on a stone path in a luminous forest. " + "Blue fireflies drift between the trees. The fox walks slowly along the path." + ), height=704, width=1280, - num_frames=97, - segment_latent_lengths=[12, 13], - num_inference_steps=8, + num_frames=120, + fps=24, + num_inference_steps=4, generator=torch.Generator("cpu").manual_seed(42), output=["videos", "session"], ) -export_to_video(result["videos"][0], "world.mp4", fps=24) +export_to_video(first["videos"][0], "hy_nomas_video.mp4", fps=24, quality=None) ``` Pass the dtype to both loading and `to` to reproduce the original BF16 inference. The latter also casts the text encoder's nonpersistent rotary-position buffers, which are recreated in FP32 when loading its weights. +### Structured prompts and multiple segments + +Structured prompts preserve scene descriptions and action timing. Serialize each segment with `json.dumps`; `prompt` +accepts a string or a list of strings, not a Python dictionary. Each serialized object must contain exactly one entry +in `segments`. A list describes successive segments of one video, with batch size one. + +The following two prompts use the `forest_world` scene and action descriptions from the model assets' +`examples/prompts.json`, with each segment extended to five seconds. Keep action timestamps relative to each segment: +the second segment's 5–10 second interval becomes 0–5 seconds. Prompt enhancement and splitting a complete prompt file +into segments happen before the pipeline call. + +```python +import json + +prompts = [ + { + "Style": "Cinematic fantasy", + "perspective_view": "Third-person view", + "subject_list": [ + { + "subject_id": "fox", + "display_name": "a small orange fox", + "subject_attribute": "Orange fur, white chest, and a fluffy tail.", + } + ], + "segments": [ + { + "visual_description_long": ( + "A small orange fox stands on a stone path in a luminous forest. " + "Blue fireflies drift between the trees." + ), + "contained_actions": [ + { + "start_time": "00:00.000", + "end_time": "00:05.000", + "action_description": "The fox walks slowly along the path.", + } + ], + } + ], + }, + { + "segments": [ + { + "description_change": "Glowing blue flowers open beside the path as the fox continues forward.", + "contained_actions": [ + { + "start_time": "00:00.000", + "end_time": "00:05.000", + "action_description": "The fox passes the flowers and looks around.", + } + ], + } + ] + }, +] + +videos = pipe( + prompt=[json.dumps(prompt) for prompt in prompts], + height=704, + width=1280, + num_frames=240, + segment_latent_lengths=[30, 30], + fps=24, + num_inference_steps=4, + generator=torch.Generator("cpu").manual_seed(42), + output="videos", +) +export_to_video(videos[0], "hy_nomas_video.mp4", fps=24, quality=None) +``` + +Each segment contains 30 latent frames, corresponding to 120 pixel frames or five seconds at 24 FPS. To generate only +the first segment, pass `prompt=json.dumps(prompts[0])`, `num_frames=120`, and `segment_latent_lengths=[30]`. + +This adapted `forest_world` case covers ten seconds. Given absolute segment boundaries of 0, 5, and 10 seconds, +the original Hy-Nomas event-based inference derives 240 frames and pixel boundaries `[0, 120, 240]`. Use +`num_frames=240` and `segment_latent_lengths=[30, 30]` to generate the same duration in one pipeline call. + +| Timestamp location | First segment | Second segment | +|---|---|---| +| Source JSON for this adapted example, relative to the full video | 0–5 seconds | 5–10 seconds | +| Actions in each model-ready prompt, relative to its segment | 0–5 seconds | 0–5 seconds | + +The original preprocessor subtracts the segment start time from its action timestamps. The pipeline expects these +segment-relative action times; absolute segment boundaries are supplied separately through frame counts. + +`num_frames` is the number of new pixel frames; the latent frame count is `(num_frames - 1) // 4 + 1`. +`segment_latent_lengths` must sum to this count. If omitted, all segments receive equal lengths, and the latent count +must be divisible by the number of prompts. The pipeline uses `fps` to map action timestamps to latent frames; keep +it equal to the FPS passed to `export_to_video`. It does not infer the rollout length from JSON timestamps. +Height and width must be multiples of 32. For the original nominal 720p preset, the decoded height is 704 pixels. + +The default PixelUMM decoding window contains 10 latent frames. These examples use five-second segments; output +quality also depends on the prompt, seed, and decoding context. + +Use `output_type="latent"` to skip decoding. The `latents` output has shape `(1, 32, T, H/32, W/32)`. +Video output types are `np` (batch, frames, height, width, channels), `pt` (batch, frames, channels, height, width), +and `pil`. + ### Match reference attention The transformer defaults to PyTorch scaled dot product attention. On CUDA, the PixelUMM decoder uses compiled @@ -95,42 +200,62 @@ processors do not require these optional dependencies. Use `disable_reference_at checkpoint precision, noise timeline capacity, CPU generator seed, prompts, normalized camera inputs, and full history on both sides. A matching seed alone does not establish numerical equivalence across different attention kernels. -A list of prompts represents **successive segments of one video**, with batch size one. It is not a batch of independent -videos. `num_frames` is the number of new pixel frames; the latent frame count is `(num_frames - 1) // 4 + 1`. -`segment_latent_lengths` must sum to this count. If omitted, all segments receive equal lengths. -Height and width must be multiples of 32. For the original nominal 720p preset, the decoded height is 704 pixels. +Enable these processors before running the adapted `forest_world` example: use the two structured prompts above in +one 240-frame call with `segment_latent_lengths=[30, 30]`, 24 FPS, four denoising steps, and CPU seed 42. +Two separate 120-frame calls sample differently shaped noise tensors; matching their +seed and combined duration does not reproduce the single-call noise. For interactive equivalence, reserve the full +60-latent timeline with `noise_capacity=60` on the first call, then continue its session without a new generator. -For visual quality checks, use segment lengths similar to the example. Very short segments and decoding windows can -produce oversharpened or oversaturated clips. The default PixelUMM decoding window contains 10 latent frames. +To match the original MP4 encoding quality as well, export with: -Use `output_type="latent"` to skip decoding. The `latents` output has shape `(1, 32, T, H/32, W/32)`. -Video output types are `np` (batch, frames, height, width, channels), `pt` (batch, frames, channels, height, width), -and `pil`. +```python +export_to_video(videos[0], "hy_nomas_video.mp4", fps=24, quality=None) +``` + +The original exporter leaves x264 at its default CRF 23, while `export_to_video` normally passes `quality=5`, which +maps to CRF 25. This encoding setting is separate from model inference. Pixel matching requires the reference +attention kernels; the default PyTorch attention path can produce different floating-point results. -## Camera control +## Generate video from a prompt and camera -Pass `camera_extrinsics` with shape `(1, num_frames, 4, 4)`. These are **normalized camera-to-world matrices in the -model's training coordinate system**. Camera calibration, trajectory generation, and normalization must be performed -before calling the pipeline. Omitting the argument disables camera conditioning. +Use the same `prompt` input and add `camera_extrinsics`, shaped `(1, num_frames, 4, 4)`. Omitting it disables camera +conditioning; no additional switch is needed. Reuse the loaded `pipe` and the structured `prompts` above: ```python -camera = torch.eye(4).repeat(1, 97, 1, 1) -camera[:, :, 0, 3] = torch.linspace(0, 0.2, 97) +# Normalized relative poses for forward motion, including first-pose padding. +camera = torch.eye(4).repeat(1, 120, 1, 1) +camera[:, 4:, 2, 3] = 1.0 + videos = pipe( - prompt="A forest path.", + prompt=json.dumps(prompts[0]), + camera_extrinsics=camera, height=704, width=1280, - num_frames=97, - camera_extrinsics=camera, + num_frames=120, + fps=24, + num_inference_steps=4, + generator=torch.Generator("cpu").manual_seed(42), output="videos", ) +export_to_video(videos[0], "hy_nomas_camera.mp4", fps=24, quality=None) ``` +These inputs follow the model's training normalization. With framewise normalization enabled in the reference +preprocessing, matrices represent relative motion between neighboring camera poses, rather than absolute world +positions. In this example, positive Z translation specifies normalized forward motion. The first four identity +matrices include the repeated initial pose required by this 120-frame layout; they do not describe four independently +recorded stationary frames. + +The pipeline encodes the supplied matrices without estimating camera motion or normalizing an absolute trajectory. +For custom motion, prepare camera-to-world poses in the training coordinate system and apply the reference +normalization, including segment boundaries and initial-pose padding, before calling the pipeline. Supply camera +matrices only for the new frames when continuing a session. The same camera input works with a plain text prompt. + ### Display recorded keyboard controls -The `examples/community/hy_nomas_keyboard_overlay.py` example in the Diffusers checkout displays WASD and all four arrow -keys on an existing silent video. It runs on CPU without model weights or a Hy-Nomas checkout. The original video is -preserved and the example refuses to overwrite an existing output. +The `examples/community/hy_nomas_keyboard_overlay.py` example in the Diffusers checkout displays WASD and all four +arrow keys on an existing silent video. It runs on CPU without model weights or a Hy-Nomas checkout. The original video +is preserved and the example refuses to overwrite an existing output. ```bash pip install numpy pillow imageio imageio-ffmpeg @@ -178,24 +303,36 @@ preview = draw_keyboard_overlay(rgb_uint8_frame, pressed_keys, label) The returned preview is a copy. Apply overlays after generation, and use the original frames for numerical comparisons. -## Continue or branch a world +## Multi-turn interaction -Pass the returned `session` to another call. The pipeline forks the cache before adding new segments, so the same -state can be reused for alternative continuations. Keep the pipeline components, precision, device, and resolution -unchanged when reusing a state. +Pass the returned `session` to the next call to continue the same world with a new prompt. For example, continue the +five-second `first` result from the plain text example with the next scene: ```python continued = pipe( - prompt="The river opens into a lake.", + prompt=( + "Glowing blue flowers open beside the path as the fox continues forward. " + "The fox passes the flowers and looks around." + ), height=704, width=1280, - num_frames=49, - session=result["session"], + num_frames=120, + fps=24, + num_inference_steps=4, + session=first["session"], output=["videos", "session"], ) -export_to_video(continued["videos"][0], "continued-world.mp4", fps=24) +export_to_video(continued["videos"][0], "hy_nomas_continued.mp4", fps=24, quality=None) ``` +The result contains both turns, totaling ten seconds. A structured continuation can instead use +`prompt=json.dumps(prompts[1])`. Its action timestamps remain relative to the new segment. Pass +`session=continued["session"]` for the next turn; a continuation restores the session's generator state when no new +generator is supplied. + +The pipeline forks the cache before adding new segments, so an earlier state can be reused for alternative +continuations. Keep the pipeline components, precision, device, and resolution unchanged when reusing a state. + Both `videos` and `latents` include the accumulated history. Each call contributes exactly `num_frames` pixel frames; padding at the end of a call's last segment is removed. Decoding runs over that history, which allows the temporal decoder to revise the end of the previous clip. The session is a runtime object and is not part of `save_pretrained`. @@ -213,7 +350,7 @@ For applications that need incremental pixel updates, `pipe.decoder.create_strea ```python from diffusers.modular_pipelines.hy_nomas.encoders import HyNomasTextEncoderStep -encoder = HyNomasTextEncoderStep().init_pipeline("/path/to/hy-nomas-diffusers") +encoder = HyNomasTextEncoderStep().init_pipeline("tencent/Hy-Nomas") encoder.load_components(dtype=torch.bfloat16) encoder.to(device="cuda", dtype=torch.bfloat16) encoded = encoder(prompt="A forest path.", output="prompt_embeds") diff --git a/setup.py b/setup.py index f490209e99f2..cee510adb7c1 100644 --- a/setup.py +++ b/setup.py @@ -102,6 +102,8 @@ "hf-doc-builder>=0.3.0", "huggingface-hub>=1.32.0,<3.0", "requests-mock==1.10.0", + "imageio", + "imageio-ffmpeg", "importlib_metadata", "invisible-watermark>=0.2.0", "isort>=5.5.4", @@ -245,6 +247,11 @@ def run(self): "phonemizer", ) extras["torch"] = deps_list("torch", "accelerate") +extras["hy-nomas"] = [ + "torch>=2.10.0", + "transformers>=5.20.2", + "accelerate>=1.13.0", +] + deps_list("imageio", "imageio-ffmpeg") extras["bitsandbytes"] = deps_list("bitsandbytes", "accelerate") extras["gguf"] = deps_list("gguf", "accelerate") diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index da572914dc57..39ba25b5f257 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -10,6 +10,8 @@ "hf-doc-builder": "hf-doc-builder>=0.3.0", "huggingface-hub": "huggingface-hub>=1.32.0,<3.0", "requests-mock": "requests-mock==1.10.0", + "imageio": "imageio", + "imageio-ffmpeg": "imageio-ffmpeg", "importlib_metadata": "importlib_metadata", "invisible-watermark": "invisible-watermark>=0.2.0", "isort": "isort>=5.5.4", From 12722ef11acf8e335f96cfd4faea17de78cd8909 Mon Sep 17 00:00:00 2001 From: jiaxiangc Date: Sat, 24 Oct 2026 10:41:15 +0800 Subject: [PATCH 3/4] docs: update Hy-Nomas Diffusers model ID --- docs/source/en/api/models/hy_nomas_transformer3d.md | 2 +- docs/source/en/api/models/pixelumm_decoder.md | 2 +- docs/source/en/api/pipelines/hy_nomas.md | 10 +++++----- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/source/en/api/models/hy_nomas_transformer3d.md b/docs/source/en/api/models/hy_nomas_transformer3d.md index e28d561ea97d..a2dc6ce5cb4a 100644 --- a/docs/source/en/api/models/hy_nomas_transformer3d.md +++ b/docs/source/en/api/models/hy_nomas_transformer3d.md @@ -20,7 +20,7 @@ import torch from diffusers import HyNomasTransformer3DModel model = HyNomasTransformer3DModel.from_pretrained( - "tencent/Hy-Nomas", subfolder="transformer", torch_dtype=torch.bfloat16 + "tencent/Hy-Nomas-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 ) ``` diff --git a/docs/source/en/api/models/pixelumm_decoder.md b/docs/source/en/api/models/pixelumm_decoder.md index 9814cd2039c3..ab40ebaffad7 100644 --- a/docs/source/en/api/models/pixelumm_decoder.md +++ b/docs/source/en/api/models/pixelumm_decoder.md @@ -19,7 +19,7 @@ It does not encode videos into latents. import torch from diffusers import PixelUMMDecoder -decoder = PixelUMMDecoder.from_pretrained("tencent/Hy-Nomas", subfolder="decoder", torch_dtype=torch.bfloat16) +decoder = PixelUMMDecoder.from_pretrained("tencent/Hy-Nomas-Diffusers", subfolder="decoder", torch_dtype=torch.bfloat16) ``` The decoder supports local spatial attention, overlapping temporal windows, optional spatial tiling, and explicit diff --git a/docs/source/en/api/pipelines/hy_nomas.md b/docs/source/en/api/pipelines/hy_nomas.md index 18cbceacf75a..cb3cd62e92ed 100644 --- a/docs/source/en/api/pipelines/hy_nomas.md +++ b/docs/source/en/api/pipelines/hy_nomas.md @@ -39,13 +39,13 @@ python scripts/convert_hy_nomas_to_diffusers.py \ --transformer_checkpoint /path/to/dit.safetensors \ --decoder_checkpoint /path/to/pixelumm.safetensors \ --text_encoder_path /path/to/Qwen2.5-VL-7B-Instruct \ - --output_path /path/to/Hy-Nomas \ + --output_path /path/to/Hy-Nomas-Diffusers \ --dtype bfloat16 ``` The output contains `modular_model_index.json` and standard component subdirectories. The converted pipeline loads -without importing the original repository or executing remote code. The examples below use `tencent/Hy-Nomas` as the -model ID; to use a local conversion, replace it with `/path/to/Hy-Nomas`. +without importing the original repository or executing remote code. The examples below use `tencent/Hy-Nomas-Diffusers` +as the model ID; to use a local conversion, replace it with `/path/to/Hy-Nomas-Diffusers`. ## Generate video from a prompt @@ -58,7 +58,7 @@ import torch from diffusers import HyNomasModularPipeline from diffusers.utils import export_to_video -pipe = HyNomasModularPipeline.from_pretrained("tencent/Hy-Nomas") +pipe = HyNomasModularPipeline.from_pretrained("tencent/Hy-Nomas-Diffusers") pipe.load_components(dtype=torch.bfloat16) pipe.to(device="cuda", dtype=torch.bfloat16) @@ -350,7 +350,7 @@ For applications that need incremental pixel updates, `pipe.decoder.create_strea ```python from diffusers.modular_pipelines.hy_nomas.encoders import HyNomasTextEncoderStep -encoder = HyNomasTextEncoderStep().init_pipeline("tencent/Hy-Nomas") +encoder = HyNomasTextEncoderStep().init_pipeline("tencent/Hy-Nomas-Diffusers") encoder.load_components(dtype=torch.bfloat16) encoder.to(device="cuda", dtype=torch.bfloat16) encoded = encoder(prompt="A forest path.", output="prompt_embeds") From 1a0aaa32baa8ab02beaa376fec26f93817393562 Mon Sep 17 00:00:00 2001 From: jiaxiangc Date: Sun, 25 Oct 2026 00:03:57 +0800 Subject: [PATCH 4/4] docs: clarify Hy-Nomas reference attention setup --- docs/source/en/api/pipelines/hy_nomas.md | 77 ++++++++++++++++++++---- 1 file changed, 65 insertions(+), 12 deletions(-) diff --git a/docs/source/en/api/pipelines/hy_nomas.md b/docs/source/en/api/pipelines/hy_nomas.md index cb3cd62e92ed..5b969deaa53f 100644 --- a/docs/source/en/api/pipelines/hy_nomas.md +++ b/docs/source/en/api/pipelines/hy_nomas.md @@ -20,16 +20,65 @@ by the model's prompt preprocessing. Action timestamps in JSON must already be r ## Installation -Install PyTorch 2.10.0 with CUDA support for your system, then run the following command from the Diffusers checkout: +The examples below enable reference attention to reproduce the original Hy-Nomas computation. This mode requires +FlashAttention 2, FlashAttention 3, and MagiAttention. The `hy-nomas` extra installs PyTorch, Transformers, Accelerate, +and the video export dependencies, but does not install these three CUDA attention libraries. + +### Reuse the Hy-Nomas environment + +If you have already prepared the [Hy-Nomas environment](https://github.com/Tencent-Hunyuan/Hy-Nomas#environment-setup), +reuse the Python environment containing its attention kernels, including the preinstalled environment inside the +Hy-Nomas Docker image. Run the following command from the Diffusers checkout; no separate attention installation is +needed when all three libraries are already available: ```bash pip install -e ".[hy-nomas]" ``` -The `hy-nomas` extra installs the Transformers text encoder, Accelerate loading support, and the `imageio` and -`imageio-ffmpeg` video export dependencies. The original Hy-Nomas code checkout is not required. +Run both implementations with the same PyTorch/CUDA and attention library builds when comparing pixels. Installing +Diffusers in a separate environment does not automatically make the original environment's CUDA kernels available. + +### Install the attention libraries with pip + +For a new environment, prepare Linux, Python 3.12+, the CUDA 13 Toolkit with `nvcc`, a C++ compiler, and an NVIDIA GPU +supported by the reference kernels. The following public source versions follow the +[Hy-Nomas Dockerfile](https://github.com/Tencent-Hunyuan/Hy-Nomas/blob/main/docker/Dockerfile): + +| Library | Source version | +|---|---| +| FlashAttention 2 | `v2.8.3.post2` | +| FlashAttention 3 | `v2.8.3.post2`, `hopper/` subdirectory (`3.0.0b1`) | +| MagiAttention | `v1.1.1` | + +Install PyTorch and the build/runtime dependencies before compiling the attention libraries: + +```bash +pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cu130 +pip install "setuptools>=78.1.1" wheel "packaging>=24.2" ninja psutil einops \ + versioningit==3.3.0 \ + "nvidia-cutlass-dsl[cu13]==4.8.0" nvidia-nvshmem-cu13==3.8.0 +``` -## Convert the checkpoints +Then install all three attention libraries and Diffusers from the Diffusers checkout: + +```bash +FLASH_ATTENTION_FORCE_BUILD=TRUE MAGI_ATTENTION_PREBUILD_FFA=0 MAX_JOBS=8 NVCC_THREADS=4 \ +pip install --no-build-isolation \ + "flash-attn @ git+https://github.com/Dao-AILab/flash-attention.git@v2.8.3.post2" \ + "flash-attn-3 @ git+https://github.com/Dao-AILab/flash-attention.git@v2.8.3.post2#subdirectory=hopper" \ + "magi-attention @ git+https://github.com/SandAI-org/MagiAttention.git@v1.1.1" \ + -e ".[hy-nomas]" +``` + +These are source-build instructions, not the recorded versions of every published Docker image or the earlier parity +test environment. Use the resulting environment for both implementations; installing the libraries alone does not +establish pixel equivalence. The converted weights do not require importing the original Hy-Nomas code. + +## Optional: convert original checkpoints + +Skip this section when using `tencent/Hy-Nomas-Diffusers`. The repository includes all required model components, +including the text encoder and tokenizer. No separate Qwen checkpoint download is required. The bundled text encoder +is distinct from the optional Nomas-PE prompt enhancement model. Run the conversion script from a Diffusers checkout. The decoder input must be the original floating-point checkpoint. The script checks parameter mappings and verifies the saved components by loading them again. @@ -49,6 +98,10 @@ as the model ID; to use a local conversion, replace it with `/path/to/Hy-Nomas-D ## Generate video from a prompt +The examples below place all pipeline components on a single CUDA device in BF16, without CPU offloading or +quantization. Peak GPU memory usage depends on the resolution, segment length, accumulated history, and attention +backend. + Use a plain string to describe one segment. This example generates 120 frames at 24 FPS with four denoising steps and returns a session for subsequent interaction. @@ -61,6 +114,7 @@ from diffusers.utils import export_to_video pipe = HyNomasModularPipeline.from_pretrained("tencent/Hy-Nomas-Diffusers") pipe.load_components(dtype=torch.bfloat16) pipe.to(device="cuda", dtype=torch.bfloat16) +pipe.transformer.enable_reference_attention(text_backend="flash_varlen", deterministic=True) first = pipe( prompt=( @@ -182,20 +236,19 @@ and `pil`. ### Match reference attention -The transformer defaults to PyTorch scaled dot product attention. On CUDA, the PixelUMM decoder uses compiled -FlexAttention for local attention, so its first call includes compilation. This default inference path does not require -MagiAttention or FlashAttention. To compare against the original BF16 implementation, install its compatible -MagiAttention, FlashAttention 2, and FlashAttention 3 dependencies and select the reference processors after loading -the transformer: +The examples above already enable reference attention after loading the transformer. Complete the +[installation steps](#installation) before running them. The current `enable_reference_attention()` API checks for all +three attention packages, even when a particular prompt does not execute every backend: ```python pipe.transformer.enable_reference_attention(text_backend="flash_varlen", deterministic=True) ``` This uses FlashAttention 2 for the unpadded text refiner, MagiAttention for action and camera visibility ranges, and -FlashAttention 3 for unrestricted history attention. MagiAttention is imported only when its backend runs; the default -processors do not require these optional dependencies. Use `disable_reference_attention()` or -`set_attention_backend("native")` to return to the standard processors. Set +FlashAttention 3 for unrestricted history attention. MagiAttention is imported only when its backend runs. Calling +`disable_reference_attention()` or `set_attention_backend("native")` restores the default PyTorch attention processors; +that path does not preserve pixel equivalence with the original implementation. On CUDA, the PixelUMM decoder uses +compiled FlexAttention for local attention, so its first call includes compilation. Set `MAGI_ATTENTION_DETERMINISTIC_MODE=1` in the original implementation as well when comparing outputs. Use the same checkpoint precision, noise timeline capacity, CPU generator seed, prompts, normalized camera inputs, and full history on both sides. A matching seed alone does not establish numerical equivalence across different attention kernels.