diff --git a/LICENSE b/LICENSE index 1e0105f47..22556846a 100644 --- a/LICENSE +++ b/LICENSE @@ -19,6 +19,11 @@ FlashDreams is open-source software. Its licensing posture is: is licensed under the Zlib license. Full text: LICENSES/Zlib.txt. + * MiniMax H3 codec and conditioning ports under + integrations_v2/minimax_h3/impl/ retain their upstream Apache-2.0 + attribution. Full text: LICENSES/Apache-2.0.txt; source details + are recorded in THIRD-PARTY-NOTICES. + Each first-party source file carries an inline SPDX license identifier; the REUSE 3.3 manifest at REUSE.toml fills the gaps for files (configuration, assets, generated outputs) that cannot diff --git a/NOTICE b/NOTICE index 8903fdebb..7c4cc9112 100644 --- a/NOTICE +++ b/NOTICE @@ -15,6 +15,10 @@ LICENSES/: cudaraster/framework/3rdparty/lodepng/{lodepng.h,lodepng.cpp} Zlib (see LICENSES/Zlib.txt) +The MiniMax/HuggingFace H3 codec and conditioning ports under +integrations_v2/minimax_h3/impl/ retain their upstream Apache-2.0 +attribution (see LICENSES/Apache-2.0.txt). + Third-party software attributions, source-level redistribution disclosures, and the full per-dependency license inventory are documented in THIRD-PARTY-NOTICES at the repository root. diff --git a/REUSE.toml b/REUSE.toml index 0c8a7f83d..ca3798b9d 100644 --- a/REUSE.toml +++ b/REUSE.toml @@ -66,6 +66,21 @@ SPDX-FileCopyrightText = [ ] SPDX-License-Identifier = "Apache-2.0" +# --- MiniMax/HuggingFace inference mathematics adapted to native APIs. --- +[[annotations]] +path = [ + "integrations_v2/minimax_h3/impl/video_vae.py", + "integrations_v2/minimax_h3/impl/audio_encoder.py", + "integrations_v2/minimax_h3/impl/conditioning.py", + "integrations_v2/minimax_h3/impl/layout.py", +] +precedence = "aggregate" +SPDX-FileCopyrightText = [ + "Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.", + "Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved.", +] +SPDX-License-Identifier = "Apache-2.0" + # --- Generated protobuf stubs. The .proto sources are authoritative; the # *_pb2.py / *_pb2.pyi / *_pb2_grpc.py outputs inherit the same license # because the regeneration script (compile_protos.sh) is project-owned. --- diff --git a/THIRD-PARTY-NOTICES b/THIRD-PARTY-NOTICES index 641e9ebe3..d0de3dbc5 100644 --- a/THIRD-PARTY-NOTICES +++ b/THIRD-PARTY-NOTICES @@ -147,6 +147,30 @@ The following third-party source is physically present in this repository (not just consumed at runtime). Each file retains its original copyright notice inline; license texts are reproduced under LICENSES/. +-------------------------------------------------------------------------------- +MiniMax H3 native codecs and conditioning +-------------------------------------------------------------------------------- + +Path: integrations_v2/minimax_h3/impl/{video_vae,audio_encoder,conditioning,layout}.py +License: Apache-2.0 (see LICENSES/Apache-2.0.txt) + Copyright 2026 The MiniMax and HuggingFace Teams. All rights reserved. +Upstream: https://github.com/huggingface/diffusers/tree/175fe6b2419a01db9c2ceabd01ec37d2c0305fc2/src/diffusers + +The inference mathematics are adapted from the MiniMax H3 implementation. +FlashDreams replaces pipeline blocks, model mixins, attention processors, +checkpoint loading, and runtime orchestration with its native APIs. The +audio port contains only reference encoding, not generated-audio decoding. +Diffusers is not a runtime dependency of this integration. Model weights +are not redistributed. Native checkpoint assets are pinned to +MiniMaxAI/MiniMax-H3 revision 42ed227ee7df40d41602854ae760620d6eb651fe. + +The integration dynamically uses Transformers and Accelerate (Apache-2.0; +https://github.com/huggingface/transformers and +https://github.com/huggingface/accelerate), PyAV (BSD-3-Clause; +https://github.com/PyAV-Org/PyAV), and Pillow (MIT-CMU; +https://github.com/python-pillow/Pillow), alongside the shared runtime +dependencies listed above. These dependencies are not vendored. + -------------------------------------------------------------------------------- HPG-2011 NVIDIA CUDA Rasterizer port -------------------------------------------------------------------------------- @@ -176,8 +200,7 @@ License: Zlib (see LICENSES/Zlib.txt) Upstream: https://lodev.org/lodepng/ LodePNG is a PNG codec implementation embedded by the upstream -cudaraster framework. It is the only non-NVIDIA source physically -redistributed in this repository. The inline notice in lodepng.h +cudaraster framework. The inline notice in lodepng.h preserves the upstream copyright notice as required by the Zlib license; LICENSES/Zlib.txt reproduces the license text in full. diff --git a/apps/t2v/t2v/application.py b/apps/t2v/t2v/application.py index 2c722ce2d..e48402e95 100644 --- a/apps/t2v/t2v/application.py +++ b/apps/t2v/t2v/application.py @@ -180,6 +180,15 @@ def create_session(self, session_desc: SessionDesc) -> ISession: if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()): raise ValueError("A session prompt must be non-empty text.") self._validate_frame_size(session_desc, pipeline) + cache_kwargs = self._cache_initialization_kwargs(session_desc) + if cache_kwargs: + return self.session_type( + pipeline, + prompt, + session_desc, + config.total_blocks, + cache_init_kwargs=cache_kwargs, + ) return self.session_type(pipeline, prompt, session_desc, config.total_blocks) def close(self) -> None: @@ -193,6 +202,13 @@ def close(self) -> None: ## Integration hooks + def _cache_initialization_kwargs(self, session_desc: SessionDesc) -> dict[str, Any]: + """Return model-specific request inputs retained across session resets. + + Standard text, image, height, and width inputs remain framework-owned. + """ + return {} + def _configure_argument_parser(self, parser: argparse.ArgumentParser) -> None: """Add arguments this integration takes beyond the shared ones.""" diff --git a/apps/t2v/t2v/session.py b/apps/t2v/t2v/session.py index c168897ca..2ac64a6c9 100644 --- a/apps/t2v/t2v/session.py +++ b/apps/t2v/t2v/session.py @@ -3,7 +3,7 @@ """One text-to-video rollout: a prompt in, a chunk of frames per step out.""" -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from flashdreams.api_v2.loop import IModelLoop @@ -25,6 +25,9 @@ class T2VModelState: image: Any = None """Optional first-frame tensor for image-conditioned generation.""" + cache_init_kwargs: dict[str, Any] = field(default_factory=dict) + """Integration-specific request inputs reused when resetting the cache.""" + blocks_generated: int = 0 cache: Any = None @@ -91,6 +94,7 @@ def __init__( total_blocks: int, *, image: Any = None, + cache_init_kwargs: dict[str, Any] | None = None, ) -> None: """ Args: @@ -101,12 +105,20 @@ def __init__( against what the model can produce. total_blocks: Blocks this rollout generates before it is finished. image: Optional first-frame tensor retained across session resets. + cache_init_kwargs: Additional request inputs, excluding standard + ``text``, ``image``, ``height``, and ``width`` arguments. """ self._pipeline = pipeline self._prompt = prompt self._session_desc = session_desc self._total_blocks = total_blocks self._image = image + self._cache_init_kwargs = dict(cache_init_kwargs or {}) + reserved = {"text", "image", "height", "width"} & self._cache_init_kwargs.keys() + if reserved: + raise ValueError( + f"Cache inputs cannot override framework arguments: {sorted(reserved)}" + ) def init(self) -> None: """Encode the prompt and prepare the rollout's cache. @@ -119,6 +131,7 @@ def init(self) -> None: session_desc=self._session_desc, total_blocks=self._total_blocks, image=self._image, + cache_init_kwargs=dict(self._cache_init_kwargs), ) if state.prompt is not None: state.cache = _new_cache(state) @@ -139,10 +152,16 @@ def _new_cache(state: T2VModelState) -> Any: """Encode the prompt into a cache for one rollout.""" if state.prompt is None: raise RuntimeError("Cannot initialize a text-to-video cache without a prompt.") + reserved = {"text", "image", "height", "width"} & state.cache_init_kwargs.keys() + if reserved: + raise ValueError( + f"Cache inputs cannot override framework arguments: {sorted(reserved)}" + ) if state.image is not None: return state.pipeline.initialize_cache( text=[state.prompt], image=state.image, + **state.cache_init_kwargs, ) ratio = state.pipeline.decoder.spatial_compression_ratio return state.pipeline.initialize_cache( @@ -150,4 +169,5 @@ def _new_cache(state: T2VModelState) -> Any: image=None, height=state.session_desc.video_height // ratio, width=state.session_desc.video_width // ratio, + **state.cache_init_kwargs, ) diff --git a/apps/t2v/tests/test_session.py b/apps/t2v/tests/test_session.py index e31c40102..39f211082 100644 --- a/apps/t2v/tests/test_session.py +++ b/apps/t2v/tests/test_session.py @@ -22,6 +22,33 @@ pytestmark = pytest.mark.ci_cpu + +def test_request_kwargs_survive_reset_and_cannot_override_standard_fields(): + from t2v.session import T2VSession + + pipeline = _stand_in() + session = T2VSession( + pipeline, + _PROMPT, + _session_desc(), + 1, + cache_init_kwargs={"reference_paths": ("first", "last")}, + ) + session.init() + session.model_loop.reset() + assert len(pipeline.caches) == 2 + assert pipeline.caches[0]["reference_paths"] == ("first", "last") + assert pipeline.caches[1]["reference_paths"] == ("first", "last") + with pytest.raises(ValueError, match="framework arguments"): + T2VSession( + pipeline, + _PROMPT, + _session_desc(), + 1, + cache_init_kwargs={"text": ["override"]}, + ) + + _WIDTH = 128 """Frame width the stand-in generates. Not square, so a transposed frame cannot pass unnoticed, and a whole number of latents across.""" diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py b/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py index ca8766268..f257ae0da 100644 --- a/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/__init__.py @@ -271,7 +271,7 @@ def compute_kv( def forward( self, x: Tensor, - kv_cache: KVCacheT, + kv_cache: KVCacheT | None = None, rope_freqs: Tensor | None = None, ) -> Tensor: """Apply the configured attention type to ``x`` and ``kv_cache``. @@ -281,7 +281,8 @@ def forward( Args: x: Query tokens, shape ``[..., L, query_dim]``. - kv_cache: Streaming cache for self-attention or precomputed static + kv_cache: ``None`` for cacheless bidirectional self-attention, + streaming cache for self-attention or precomputed static cache for cross-attention. A streaming cache must already be in its current-chunk update phase. rope_freqs: Optional positional data. Before-cache RoPE expects the diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py b/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py index 7e2118c16..6115840fd 100644 --- a/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py @@ -24,8 +24,6 @@ from enum import Enum import torch -from torch import Tensor, nn - from flashdreams.accelerated.common.non_persistent_linear import ( NonPersistentLinear, ) @@ -37,9 +35,9 @@ RoPEScope, RoPEStyle, ) -from flashdreams.accelerated.multi_head_attention.cudnn import ( - native_cudnn_fp8_sdpa, - torch_cudnn_sdpa, +from flashdreams.accelerated.multi_head_attention.sdpa import ( + SDPABackend, + scaled_dot_product_attention, ) from flashdreams.accelerated.multi_head_attention.triton import ( flash_attention_2, @@ -57,16 +55,7 @@ ) from flashdreams.core.attention import BlockKVCache from flashdreams.core.attention.rope_kernel import apply_rotary_pos_emb - - -class SDPABackend(str, Enum): - """Scaled-dot-product attention implementation.""" - - CUDNN = "cudnn" - """Use Torch cuDNN for FP16/BF16 and native cuDNN Frontend for FP8.""" - - FA2 = "fa2" - """Use Triton FlashAttention2 (FA2).""" +from torch import Tensor, nn class QKVFusionOption(str, Enum): @@ -160,6 +149,8 @@ def __post_init__(self) -> None: ) if not isinstance(self.use_tma, bool): raise TypeError(f"use_tma must be a bool; got {self.use_tma!r}") + if self.sdpa_backend is SDPABackend.TORCH and self.quantization.quantized_sdpa: + raise ValueError("Torch SDPA does not support quantized SDPA") class OptimizedMultiHeadAttention(MultiHeadAttention[BlockKVCache]): @@ -316,6 +307,11 @@ def _new_quantized_projection( dtype, ) + @torch.no_grad() + def refresh_derived_weights(self) -> None: + """Refresh fused and quantized projections after in-place weight edits.""" + self._refresh_derived_weights() + @torch.no_grad() def _refresh_derived_weights(self, *args: object) -> None: """Rebuild fused projection modules from checkpoint parameters. @@ -535,7 +531,7 @@ def compute_kv( def forward( self, x: Tensor, - kv_cache: BlockKVCache, + kv_cache: BlockKVCache | None = None, rope_freqs: Tensor | None = None, ) -> Tensor: """Apply self- or cross-attention using the configured cache lifecycle. @@ -546,7 +542,8 @@ def forward( Args: x: Query tokens shaped ``[..., L, Q]``. - kv_cache: Prepared rolling cache for self-attention or precomputed + kv_cache: ``None`` for cacheless self-attention, or prepared rolling + cache for self-attention or precomputed static cache for cross-attention. rope_freqs: Optional positional data. Before-cache RoPE expects the current chunk. After-cache RoPE expects positions covering the @@ -556,27 +553,37 @@ def forward( Returns: Output-projected tokens with the same shape and dtype as ``x``. """ - query_rope_freqs, key_rope_freqs = self._slice_rope_freqs( - rope_freqs, kv_cache, x.shape[-2] - ) - if self.attention_type is AttentionType.SELF_ATTENTION: - query = self._update_kv_and_compute_query(x, kv_cache, query_rope_freqs) + if kv_cache is None: + if self.attention_type is not AttentionType.SELF_ATTENTION: + raise ValueError("cross-attention requires a K/V cache") + self._validate_tokens(x, self.attention_config.query_dim, "x") + if self.qkv_fusion_option is QKVFusionOption.FULL: + query, key, value = self._project_qkv(x) + query = self._apply_qk_norm(query, self.query_norm) + key = self._apply_qk_norm(key, self.key_norm) + else: + query = self._project_query(x) + key, value = self._project_kv(x) + if rope_freqs is not None: + query = self._apply_rope(query, rope_freqs) + key = self._apply_rope(key, rope_freqs) else: - query = self._compute_query(x, query_rope_freqs) - self._validate_cache(kv_cache, x) - - # ``cached_k/v`` expose only the valid prefix while a rolling cache fills, - # and the complete fixed-size buffer after it reaches steady state. - key = kv_cache.cached_k() - if ( - self.attention_config.rope_config is not None - and self.attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE - and key_rope_freqs is not None - ): - # The shared RoPE kernel is in-place; keep cache storage unrotated so - # rolling positions can be applied again on the next attention call. - key = self._apply_rope(key.to(x.dtype, copy=True), key_rope_freqs) - value = kv_cache.cached_v() + query_rope_freqs, key_rope_freqs = self._slice_rope_freqs( + rope_freqs, kv_cache, x.shape[-2] + ) + if self.attention_type is AttentionType.SELF_ATTENTION: + query = self._update_kv_and_compute_query(x, kv_cache, query_rope_freqs) + else: + query = self._compute_query(x, query_rope_freqs) + self._validate_cache(kv_cache, x) + key = kv_cache.cached_k() + if ( + self.attention_config.rope_config is not None + and self.attention_config.rope_config.scope is RoPEScope.AFTER_KV_CACHE + and key_rope_freqs is not None + ): + key = self._apply_rope(key.to(x.dtype, copy=True), key_rope_freqs) + value = kv_cache.cached_v() if self.optimized_impl_config.quantization.quantized_sdpa: query = query.to(torch.float8_e4m3fn) key = key.to(torch.float8_e4m3fn) @@ -730,33 +737,14 @@ def _attention( Returns: Attention output with shape ``[B, L, H, D]``. """ - if self.sdpa_backend is SDPABackend.CUDNN: - # The module and Triton kernel use token-major ``[B, L/S, H, D]``. - # PyTorch SDPA instead interprets its two middle axes as ``[H, L/S]``. - # These transposes change only shape/stride metadata. - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - - # PyTorch's public dispatcher rejects FP8 inputs, so use a cuDNN - # Frontend FP8 graph for e4m3 attention. - if query.dtype is torch.float8_e4m3fn: - output = native_cudnn_fp8_sdpa(query, key, value) - else: - output = torch_cudnn_sdpa(query, key, value) - - # Restore the module-wide ``[B, L, H, D]`` contract for head merging. - output = output.transpose(1, 2) - return output if output_dtype is None else output.to(output_dtype) - - attention = ( - flash_attention_2_tma - if self.use_tma and is_tma_flash_attention_supported(query, key, value) - else flash_attention_2 + return scaled_dot_product_attention( + query, + key, + value, + backend=self.sdpa_backend, + use_tma=self.use_tma, + output_dtype=output_dtype, ) - if output_dtype is None: - return attention(query, key, value) - return attention(query, key, value, output_dtype=output_dtype) def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: """Apply the shared RoPE kernel to token-major head features. @@ -771,12 +759,24 @@ def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: rope_config = self.attention_config.rope_config if rope_config is None: return x - return apply_rotary_pos_emb( - x, + if rope_freqs.ndim != 4: + raise ValueError("rope_freqs must have shape [L, 1, 1, rotary_dim]") + rotary_dim = rope_freqs.shape[-1] + if rotary_dim <= 0 or rotary_dim % 2 or rotary_dim > x.shape[-1]: + raise ValueError( + "RoPE width must be positive, even and no larger than head_dim" + ) + if tuple(rope_freqs.shape) != (x.shape[1], 1, 1, rotary_dim): + raise ValueError("rope_freqs must have shape [L, 1, 1, rotary_dim]") + if rope_freqs.device != x.device: + raise RuntimeError("rope_freqs and x must be on the same device") + apply_rotary_pos_emb( + x[..., :rotary_dim], rope_freqs, interleaved=rope_config.style is RoPEStyle.INTERLEAVED, inplace=True, ) + return x # ------------------------ Validation ------------------------ # @@ -787,7 +787,7 @@ def _validate_cuda_device(self, device: torch.device | str) -> None: device: Device used by attention inputs and cache storage. Raises: - RuntimeError: The CUDA device predates Hopper. + RuntimeError: The device predates Ampere, or FP8 requires Hopper. """ device = torch.device(device) if device.type != "cuda": @@ -797,9 +797,19 @@ def _validate_cuda_device(self, device: torch.device | str) -> None: ) if self._validated_cuda_device_index == device_index: return - if torch.cuda.get_device_capability(device_index)[0] < 9: + quantization = self.optimized_impl_config.quantization + fp8 = (torch.float8_e4m3fn, torch.float8_e5m2) + needs_hopper = ( + quantization.quantized_sdpa + or quantization.projection in fp8 + or quantization.output_projection in fp8 + ) + if torch.cuda.get_device_capability(device_index)[0] < ( + 9 if needs_hopper else 8 + ): raise RuntimeError( - "OptimizedMultiHeadAttention requires compute capability 9.0 or newer" + "OptimizedMultiHeadAttention requires compute capability " + + ("9.0 for FP8" if needs_hopper else "8.0 or newer") ) self._validated_cuda_device_index = device_index @@ -813,7 +823,8 @@ def _validate_tokens(self, x: Tensor, feature_dim: int, name: str) -> None: Raises: ValueError: ``x`` lacks sequence/feature axes or has the wrong width. - RuntimeError: ``x`` is not CUDA FP16/BF16 or the GPU predates Hopper. + RuntimeError: ``x`` is not CUDA FP16/BF16 or the device does not + support the configured precision. """ if x.ndim < 2: raise ValueError( @@ -974,7 +985,18 @@ def _validate_fused_update_inputs( if self.attention_config.rope_config is not None and rope_freqs is not None: # RoPE coefficients cover this ``L``-token chunk and broadcast across # flattened batches and ``H`` heads inside the shared kernel. - expected_rope_shape = (x.shape[-2], 1, 1, self.attention_config.head_dim) + if rope_freqs.ndim != 4: + raise ValueError("rope_freqs must have shape [L, 1, 1, rotary_dim]") + rotary_dim = rope_freqs.shape[-1] + if ( + rotary_dim <= 0 + or rotary_dim % 2 + or rotary_dim > self.attention_config.head_dim + ): + raise ValueError( + "RoPE width must be positive, even and no larger than head_dim" + ) + expected_rope_shape = (x.shape[-2], 1, 1, rotary_dim) if tuple(rope_freqs.shape) != expected_rope_shape: raise ValueError( f"rope_freqs must have shape {expected_rope_shape}; " @@ -1151,11 +1173,11 @@ def _project_output(self, x: Tensor) -> Tensor: __all__ = [ + "OptimizedImplConfig", + "OptimizedMultiHeadAttention", "QKVFusionOption", "QuantizationOption", "SDPABackend", - "OptimizedImplConfig", - "OptimizedMultiHeadAttention", "flash_attention_2", "flash_attention_2_tma", "is_tma_flash_attention_supported", diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/sdpa.py b/flashdreams/flashdreams/accelerated/multi_head_attention/sdpa.py new file mode 100644 index 000000000..7717613a2 --- /dev/null +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/sdpa.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Shared scaled-dot-product attention over projected token-major tensors.""" + +from enum import Enum + +import torch +import torch.nn.functional as F +from torch import Tensor + + +class SDPABackend(str, Enum): + """Scaled-dot-product attention implementation.""" + + TORCH = "torch" + """Use PyTorch's native device- and dtype-aware dispatcher.""" + + CUDNN = "cudnn" + """Use Torch cuDNN or native cuDNN Frontend for FP8.""" + + FA2 = "fa2" + """Use Triton FlashAttention2.""" + + +def scaled_dot_product_attention( + query: Tensor, + key: Tensor, + value: Tensor, + *, + is_causal: bool = False, + backend: SDPABackend = SDPABackend.TORCH, + use_tma: bool = False, + output_dtype: torch.dtype | None = None, +) -> Tensor: + """Attend to projected Q/K/V in ``[B, L/S, H, D]`` layout. + + Args: + query: Projected query tokens. + key: Projected key tokens. + value: Projected value tokens. + is_causal: Apply a causal mask; supported by the native Torch backend. + backend: Implementation policy, independent of model projections. + use_tma: Prefer TMA when the FA2 backend and hardware support it. + output_dtype: Optional output storage dtype. + + Returns: + Attention output in token-major layout. + + Raises: + ValueError: Tensor geometry or backend policy is invalid. + """ + if not isinstance(backend, SDPABackend): + raise TypeError(f"unsupported SDPA backend: {backend!r}") + if any(x.ndim != 4 for x in (query, key, value)): + raise ValueError("Q/K/V must have shape [B, L/S, H, D]") + if ( + key.shape != value.shape + or query.shape[0] != key.shape[0] + or query.shape[2:] != key.shape[2:] + ): + raise ValueError("Q/K/V batch, head and feature dimensions must match") + if is_causal and backend is not SDPABackend.TORCH: + raise ValueError("causal attention requires the Torch SDPA backend") + if backend is SDPABackend.TORCH: + if query.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + raise ValueError("Torch SDPA does not support FP8 inputs") + output = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + dropout_p=0.0, + is_causal=is_causal, + ).transpose(1, 2) + elif backend is SDPABackend.CUDNN: + from flashdreams.accelerated.multi_head_attention.cudnn import ( + native_cudnn_fp8_sdpa, + torch_cudnn_sdpa, + ) + + attention = ( + native_cudnn_fp8_sdpa + if query.dtype is torch.float8_e4m3fn + else torch_cudnn_sdpa + ) + output = attention( + query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) + ).transpose(1, 2) + else: + from flashdreams.accelerated.multi_head_attention.triton import ( + flash_attention_2, + flash_attention_2_tma, + is_tma_flash_attention_supported, + ) + + attention = ( + flash_attention_2_tma + if use_tma and is_tma_flash_attention_supported(query, key, value) + else flash_attention_2 + ) + return attention(query, key, value, output_dtype=output_dtype) + return output if output_dtype is None else output.to(output_dtype) diff --git a/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py b/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py index e28001905..8dd0b8889 100644 --- a/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py +++ b/flashdreams/flashdreams/accelerated/multi_head_attention/torch.py @@ -20,9 +20,6 @@ from abc import abstractmethod import torch -import torch.nn.functional as F -from torch import Tensor, nn - from flashdreams.accelerated.multi_head_attention import ( AttentionConfig, AttentionType, @@ -31,7 +28,11 @@ RoPEScope, RoPEStyle, ) +from flashdreams.accelerated.multi_head_attention.sdpa import ( + scaled_dot_product_attention, +) from flashdreams.core.attention import BlockKVCache +from torch import Tensor, nn class TorchMultiHeadAttention(MultiHeadAttention[BlockKVCache]): @@ -178,14 +179,15 @@ def compute_kv( def forward( self, x: Tensor, - kv_cache: BlockKVCache, + kv_cache: BlockKVCache | None = None, rope_freqs: Tensor | None = None, ) -> Tensor: """Apply self- or cross-attention using the configured cache lifecycle. Args: x: Query tokens, shape ``[..., L, query_dim]``. - kv_cache: Prepared rolling cache for self-attention or precomputed + kv_cache: ``None`` for cacheless self-attention, or prepared rolling + cache for self-attention or precomputed static cache for cross-attention. rope_freqs: Optional positional data. Before-cache RoPE expects the current chunk. After-cache RoPE expects positions covering the @@ -195,6 +197,17 @@ def forward( Returns: Output-projected tokens with the same shape as ``x``. """ + if kv_cache is None: + if self.attention_type is not AttentionType.SELF_ATTENTION: + raise ValueError("cross-attention requires a K/V cache") + query = self._project_query(x) + key, value = self._project_kv(x) + if rope_freqs is not None: + query = self._apply_rope(query, rope_freqs) + key = self._apply_rope(key, rope_freqs) + output = self._output_projection(self._attention(query, key, value)) + return output.reshape(x.shape[:-2] + output.shape[-2:]) + query_rope_freqs, key_rope_freqs = self._slice_rope_freqs( rope_freqs, kv_cache, x.shape[-2] ) @@ -467,12 +480,17 @@ def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: """ if self.attention_config.rope_config is None: return x - if x.shape[-1] % 2 != 0: - raise ValueError(f"RoPE requires an even head_dim; got {x.shape[-1]}") + if rope_freqs.ndim != 4: + raise ValueError("rope_freqs must have shape [L, 1, 1, rotary_dim]") + rotary_dim = rope_freqs.shape[-1] + if rotary_dim <= 0 or rotary_dim % 2 or rotary_dim > x.shape[-1]: + raise ValueError( + "RoPE width must be positive, even and no larger than head_dim" + ) # RoPE lookup shape is ``[L, 1, 1, D]`` for an input shaped # ``[..., L, H, D]``. - expected_shape = (x.shape[-3], 1, 1, x.shape[-1]) + expected_shape = (x.shape[-3], 1, 1, rotary_dim) if tuple(rope_freqs.shape) != expected_shape: raise ValueError( f"rope_freqs must have shape {expected_shape}; " @@ -482,8 +500,9 @@ def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: # Broadcast positions over leading dimensions and heads: # ``[L, 1, 1, D] -> [..., L, 1, D]``. freqs = rope_freqs[:, 0, 0, :].reshape( - (1,) * (x.ndim - 3) + (x.shape[-3], 1, x.shape[-1]) + (1,) * (x.ndim - 3) + (x.shape[-3], 1, rotary_dim) ) + prefix, tail = x[..., :rotary_dim], x[..., rotary_dim:] # Materialize rotation coefficients in activation precision so the # elementwise rotation neither promotes projected tensors nor cache data. @@ -491,14 +510,16 @@ def _apply_rope(self, x: Tensor, rope_freqs: Tensor) -> Tensor: sin_freqs = torch.sin(freqs).to(dtype=x.dtype) if self.attention_config.rope_config.style is RoPEStyle.INTERLEAVED: # Rotate adjacent feature pairs; shape stays ``[..., L, H, D]``. - rotated = torch.stack((-x[..., 1::2], x[..., 0::2]), dim=-1).flatten(-2) + rotated = torch.stack( + (-prefix[..., 1::2], prefix[..., 0::2]), dim=-1 + ).flatten(-2) else: # Rotate matching half-split features; shape stays ``[..., L, H, D]``. - first, second = x.chunk(2, dim=-1) + first, second = prefix.chunk(2, dim=-1) rotated = torch.cat((-second, first), dim=-1) # Apply the elementwise complex rotation: ``[..., L, H, D]``. - return x * cos_freqs + rotated * sin_freqs + return torch.cat((prefix * cos_freqs + rotated * sin_freqs, tail), dim=-1) def _attention(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor: """Apply non-causal scaled dot-product attention over visible K/V. @@ -514,23 +535,7 @@ def _attention(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor: # Move heads before tokens for SDPA: # Q ``[..., L, H, D] -> [..., H, L, D]`` and # K/V ``[..., S, H, D] -> [..., H, S, D]``. - query_heads = query.transpose(-3, -2) - key_heads = key.transpose(-3, -2) - value_heads = value.transpose(-3, -2) - - # Let PyTorch dispatch the available SDPA backend so the reference works - # on CPU and CUDA. Cache visibility defines the allowed context, while - # zero dropout and a non-causal mask make inference deterministic. - output = F.scaled_dot_product_attention( - query_heads, - key_heads, - value_heads, - dropout_p=0.0, - is_causal=False, - ) - - # Restore token-major layout: ``[..., H, L, D] -> [..., L, H, D]``. - return output.transpose(-3, -2) + return scaled_dot_product_attention(query, key, value) def _output_projection(self, x: Tensor) -> Tensor: """Concatenate attention heads and project back to query features. diff --git a/flashdreams/flashdreams/core/checkpoint/load.py b/flashdreams/flashdreams/core/checkpoint/load.py index 4432e1be5..00355543b 100644 --- a/flashdreams/flashdreams/core/checkpoint/load.py +++ b/flashdreams/flashdreams/core/checkpoint/load.py @@ -765,12 +765,17 @@ def _copy_checkpoint_tensor(destination: torch.Tensor, source: torch.Tensor) -> def _stream_safetensors_into_model( model: torch.nn.Module, path: str, + include_prefixes: tuple[str, ...] | None = None, ) -> torch.nn.Module: """Copy a safetensors checkpoint into a model with bounded host residency.""" model_state = model.state_dict() with safe_open(path, framework="pt", device="cpu") as source: - checkpoint_keys = set(source.keys()) + checkpoint_keys = { + key + for key in source.keys() + if include_prefixes is None or key.startswith(include_prefixes) + } model_keys = set(model_state) missing = sorted(model_keys - checkpoint_keys) unexpected = sorted(checkpoint_keys - model_keys) @@ -941,6 +946,7 @@ def _stream_sharded_safetensors_index_into_model( *, model: torch.nn.Module, checkpoint_min_free_gb: float | None, + include_prefixes: tuple[str, ...] | None = None, ) -> torch.nn.Module | None: """Stream a safetensors index checkpoint into ``model`` without merging.""" if checkpoint_path.startswith("s3://"): @@ -991,6 +997,7 @@ def _stream_sharded_safetensors_index_into_model( f"Invalid or empty weight_map in safetensors index: {index_local}" ) + weight_map = _select_checkpoint_prefixes(weight_map, model, include_prefixes) unique_shards = sorted(set(weight_map.values())) shard_to_path = _parallel_hf_hub_download_shards( repo_id=repo_id, @@ -1022,6 +1029,7 @@ def resolve_shard_path(shard_file: str) -> str: raise ValueError( f"Invalid or empty weight_map in safetensors index: {checkpoint_path}" ) + weight_map = _select_checkpoint_prefixes(weight_map, model, include_prefixes) base_dir = os.path.dirname(os.path.abspath(checkpoint_path)) def resolve_shard_path(shard_file: str) -> str: @@ -1034,6 +1042,29 @@ def resolve_shard_path(shard_file: str) -> str: ) +def _select_checkpoint_prefixes( + weight_map: dict[str, str], + model: torch.nn.Module, + include_prefixes: tuple[str, ...] | None, +) -> dict[str, str]: + """Select and validate component keys before downloading their shards.""" + if include_prefixes is None: + return weight_map + selected = { + name: shard + for name, shard in weight_map.items() + if name.startswith(include_prefixes) + } + model_keys = set(model.state_dict()) + if set(selected) != model_keys: + raise RuntimeError( + "Selected checkpoint components do not match model: " + f"missing={sorted(model_keys - selected.keys())[:20]}, " + f"unexpected={sorted(selected.keys() - model_keys)[:20]}" + ) + return selected + + def _resolve_streamable_safetensors_path( checkpoint_path: str, *, @@ -1126,6 +1157,8 @@ def load_checkpoint( map_location: str | torch.device = "cpu", check_success: bool = False, checkpoint_min_free_gb: float | None = None, + *, + include_prefixes: tuple[str, ...] | None = None, ) -> dict[str, torch.Tensor]: ... @@ -1139,6 +1172,8 @@ def load_checkpoint( map_location: str | torch.device = "cpu", check_success: bool = False, checkpoint_min_free_gb: float | None = None, + *, + include_prefixes: tuple[str, ...] | None = None, ) -> torch.nn.Module: ... @@ -1151,6 +1186,8 @@ def load_checkpoint( map_location: str | torch.device = "cpu", check_success: bool = False, checkpoint_min_free_gb: float | None = None, + *, + include_prefixes: tuple[str, ...] | None = None, ) -> dict[str, torch.Tensor] | torch.nn.Module: """Load checkpoints from S3, local disk, or Hugging Face. @@ -1172,6 +1209,9 @@ def load_checkpoint( checkpoint_min_free_gb: Optional first-run free-space requirement in GiB for Hugging Face checkpoint downloads. The ``FLASHDREAMS_MIN_CACHE_FREE_GB`` environment override still wins. + include_prefixes: Explicit component prefixes, including trailing dots. + Selected keys keep their names and must exactly match ``model``. + Only local/Hugging Face safetensors model loads support selection. Returns: State dict if ``model`` is ``None``, otherwise ``model`` with weights @@ -1186,6 +1226,26 @@ def load_checkpoint( >>> state = load_checkpoint("s3://bucket/foo.safetensors") >>> model = load_checkpoint("s3://bucket/dcp_dir/", model=my_model) """ + if include_prefixes is not None: + if not include_prefixes or any( + not prefix or not prefix.endswith(".") for prefix in include_prefixes + ): + raise ValueError( + "include_prefixes must contain nonempty module prefixes ending in '.'" + ) + if ( + model is None + or checkpoint_path.startswith("s3://") + or not ( + _is_sharded_safetensors_index_checkpoint(checkpoint_path) + or _get_checkpoint_extension(checkpoint_path) == ".safetensors" + ) + or checkpoint_type == "distributed" + ): + raise ValueError( + "Prefix selection requires a model and local/Hugging Face safetensors" + ) + # Auto-detect checkpoint type if checkpoint_type == "auto": if _is_sharded_safetensors_index_checkpoint(checkpoint_path): @@ -1204,6 +1264,7 @@ def load_checkpoint( checkpoint_path, model=model, checkpoint_min_free_gb=checkpoint_min_free_gb, + include_prefixes=include_prefixes, ) if streamed_model is not None: logger.info(f"Streamed checkpoint into model: {checkpoint_path}") @@ -1214,7 +1275,7 @@ def load_checkpoint( checkpoint_min_free_gb=checkpoint_min_free_gb, ) if stream_path is not None: - _stream_safetensors_into_model(model, stream_path) + _stream_safetensors_into_model(model, stream_path, include_prefixes) logger.info(f"Streamed checkpoint into model: {checkpoint_path}") return model state_dict = load_single_checkpoint( diff --git a/flashdreams/flashdreams/infra/acceleration/__init__.py b/flashdreams/flashdreams/infra/acceleration/__init__.py index c92c244cc..9f469611e 100644 --- a/flashdreams/flashdreams/infra/acceleration/__init__.py +++ b/flashdreams/flashdreams/infra/acceleration/__init__.py @@ -13,6 +13,7 @@ move_tensors_to_cpu, release_one_shot_encoder_references, run_one_shot_encoder_stage, + run_one_shot_stage, setup_one_shot_encoder, ) from flashdreams.infra.acceleration.frame_prefetch import ( @@ -57,6 +58,7 @@ "release_one_shot_encoder_references", "run_prewarm_sequence", "run_one_shot_encoder_stage", + "run_one_shot_stage", "run_timed_prewarm", "setup_one_shot_encoder", ] diff --git a/flashdreams/flashdreams/infra/acceleration/encoder_lifecycle.py b/flashdreams/flashdreams/infra/acceleration/encoder_lifecycle.py index 5fd405b85..376b8e50b 100644 --- a/flashdreams/flashdreams/infra/acceleration/encoder_lifecycle.py +++ b/flashdreams/flashdreams/infra/acceleration/encoder_lifecycle.py @@ -134,14 +134,14 @@ def move_tensors_to_cpu(value: Any, *, torch_module: Any | None = None) -> Any: return value -def run_one_shot_encoder_stage( +def run_one_shot_stage( stage: Callable[[], Any], *, release: Callable[[], Any] | None = None, cpu_result: bool = True, torch_module: Any | None = None, ) -> Any: - """Run an encoder-only stage under ``no_grad`` and release encoders after it.""" + """Run a stage under ``no_grad`` and release owned modules even on failure.""" torch = torch_module if torch_module is not None else _maybe_import_torch() no_grad = getattr(torch, "no_grad", None) context = no_grad() if callable(no_grad) else nullcontext() @@ -157,6 +157,10 @@ def run_one_shot_encoder_stage( del release_result +run_one_shot_encoder_stage = run_one_shot_stage +"""Backward-compatible name for encoder-only callers.""" + + def _maybe_import_torch() -> Any | None: try: return importlib.import_module("torch") diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/__init__.py b/flashdreams/flashdreams/infra/diffusion/scheduler/__init__.py index 6df37c806..03b8a3705 100644 --- a/flashdreams/flashdreams/infra/diffusion/scheduler/__init__.py +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/__init__.py @@ -20,6 +20,14 @@ Scheduler, SchedulerConfig, ) +from flashdreams.infra.diffusion.scheduler.data_flow_euler import ( + DataFlowEulerScheduler, + DataFlowEulerSchedulerConfig, +) +from flashdreams.infra.diffusion.scheduler.synchronized import ( + StepScheduler, + sample_synchronized, +) from flashdreams.infra.diffusion.scheduler.fm import ( FlowMatchScheduler, FlowMatchSchedulerConfig, @@ -34,6 +42,10 @@ ) __all__ = [ + "DataFlowEulerScheduler", + "DataFlowEulerSchedulerConfig", + "StepScheduler", + "sample_synchronized", "FlowPredictor", "Scheduler", "SchedulerConfig", diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/data_flow_euler.py b/flashdreams/flashdreams/infra/diffusion/scheduler/data_flow_euler.py new file mode 100644 index 000000000..6ed9a3475 --- /dev/null +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/data_flow_euler.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Data-ward rectified-flow Euler sampling on a shifted endpoint-inclusive grid.""" + +from dataclasses import dataclass, field +import math + +import torch +from torch import Tensor + +from flashdreams.infra.diffusion.scheduler.base import ( + FlowPredictor, + Scheduler, + SchedulerConfig, +) +from flashdreams.infra.diffusion.scheduler.synchronized import sample_synchronized + + +@dataclass(kw_only=True) +class DataFlowEulerSchedulerConfig(SchedulerConfig): + """Shifted sigma grid with a velocity pointing toward clean data.""" + + _target: type["DataFlowEulerScheduler"] = field( + default_factory=lambda: DataFlowEulerScheduler + ) + num_inference_steps: int = 30 + """Number of grid points; there are one fewer model evaluations.""" + shift: float = 12.0 + """Positive rational warp of the endpoint-inclusive sigma grid.""" + + +class DataFlowEulerScheduler(Scheduler): + """Blend the current sample and its data prediction in sigma space.""" + + def __init__(self, config: DataFlowEulerSchedulerConfig) -> None: + super().__init__(config) + if config.num_inference_steps < 2: + raise ValueError("num_inference_steps must be at least 2 grid points") + if not math.isfinite(config.shift) or config.shift <= 0: + raise ValueError("shift must be finite and positive") + base = torch.linspace( + 1.0, 0.0, config.num_inference_steps, dtype=torch.float32, device="cpu" + ) + sigmas = torch.unique_consecutive( + config.shift * base / (1 + (config.shift - 1) * base) + ) + self.register_buffer("sigmas", sigmas, persistent=False) + self.register_buffer("timesteps", 1.0 - sigmas[:-1], persistent=False) + + def _apply(self, fn, recurse=True): + grids = {name: getattr(self, name) for name in ("sigmas", "timesteps")} + super()._apply(fn, recurse=recurse) + for name, original in grids.items(): + setattr(self, name, original.to(device=getattr(self, name).device)) + return self + + def step(self, sample: Tensor, flow: Tensor, index: int) -> Tensor: + """Advance one step, preserving data-ward blend rounding.""" + time = self.timesteps[index].to(sample.device, sample.dtype) + denoised = sample + (1 - time) * flow + compute_dtype = ( + torch.float32 + if sample.dtype in (torch.float16, torch.bfloat16) + else sample.dtype + ) + ratio = self.sigmas[index + 1].to(sample.device, compute_dtype) / self.sigmas[ + index + ].to(sample.device, compute_dtype) + return ( + ratio * sample.to(compute_dtype) + (1 - ratio) * denoised.to(compute_dtype) + ).to(sample.dtype) + + @torch.no_grad() + def sample( + self, + initial_noise: Tensor, + predict_flow: FlowPredictor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Sample one stream through the shared synchronized loop.""" + del rng + return sample_synchronized( + (initial_noise,), + (self,), + lambda samples, times: (predict_flow(samples[0], times[0]),), + )[0] + + def add_noise( + self, clean_input: Tensor, timestep: Tensor, rng: torch.Generator | None = None + ) -> Tensor: + """Mix clean data and Gaussian noise under the data-ward time convention.""" + noise = torch.randn( + clean_input.shape, + dtype=clean_input.dtype, + device=clean_input.device, + generator=rng, + ) + time = timestep.to(clean_input.device, clean_input.dtype) + while time.ndim < clean_input.ndim: + time = time.unsqueeze(-1) + return time * clean_input + (1 - time) * noise diff --git a/flashdreams/flashdreams/infra/diffusion/scheduler/synchronized.py b/flashdreams/flashdreams/infra/diffusion/scheduler/synchronized.py new file mode 100644 index 000000000..8d94fe3f0 --- /dev/null +++ b/flashdreams/flashdreams/infra/diffusion/scheduler/synchronized.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Synchronized sampling of coupled diffusion streams.""" + +from collections.abc import Callable +from typing import Protocol + +from torch import Tensor + + +class StepScheduler(Protocol): + """Indexed deterministic schedule for one member of a joint prediction.""" + + timesteps: Tensor + """One-dimensional grid with one entry per model evaluation.""" + + def step(self, sample: Tensor, flow: Tensor, index: int) -> Tensor: + """Advance one sample using the prediction at ``index``.""" + ... + + +def sample_synchronized( + initial_samples: tuple[Tensor, ...], + schedulers: tuple[StepScheduler, ...], + predict_flow: Callable[ + [tuple[Tensor, ...], tuple[Tensor, ...]], tuple[Tensor, ...] + ], +) -> tuple[Tensor, ...]: + """Advance coupled streams with exactly one joint prediction per step. + + Args: + initial_samples: Generated samples, excluding immutable conditioning. + schedulers: One schedule per sample, all with the same number of steps. + predict_flow: Joint predictor returning one flow of each sample's shape. + + Returns: + Final samples in their original order, shapes, devices, and dtypes. + + Raises: + ValueError: Stream counts, schedule lengths, or tensor geometry disagree. + """ + if not initial_samples or len(initial_samples) != len(schedulers): + raise ValueError("Provide a nonempty sample tuple and one scheduler per sample") + grids = tuple(scheduler.timesteps for scheduler in schedulers) + if any(grid.ndim != 1 or grid.numel() == 0 for grid in grids): + raise ValueError("Schedules must be nonempty one-dimensional timestep grids") + if any(grid.numel() != grids[0].numel() for grid in grids): + raise ValueError("Coupled schedules must have equal lengths") + samples = initial_samples + for index in range(grids[0].numel()): + times = tuple( + grid[index].to(sample.device) + for grid, sample in zip(grids, samples, strict=True) + ) + flows = predict_flow(samples, times) + if not isinstance(flows, tuple) or len(flows) != len(samples): + raise ValueError("Joint predictor must return one flow per sample") + for sample, flow in zip(samples, flows, strict=True): + if flow.shape != sample.shape or flow.device != sample.device: + raise ValueError("Predicted flows must match sample shapes and devices") + if not flow.is_floating_point(): + raise ValueError("Predicted flows must be floating-point tensors") + advanced = tuple( + scheduler.step(sample, flow, index) + for scheduler, sample, flow in zip(schedulers, samples, flows, strict=True) + ) + if any( + new.shape != old.shape or new.device != old.device or new.dtype != old.dtype + for new, old in zip(advanced, samples, strict=True) + ): + raise ValueError("Schedulers must preserve sample shape, device, and dtype") + samples = advanced + return samples diff --git a/flashdreams/flashdreams/infra/encoder/text/qwen3_vl.py b/flashdreams/flashdreams/infra/encoder/text/qwen3_vl.py new file mode 100644 index 000000000..251f5d460 --- /dev/null +++ b/flashdreams/flashdreams/infra/encoder/text/qwen3_vl.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Headless Qwen3-VL hidden-state conditioning through Transformers.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + +from flashdreams.infra.encoder import Encoder, EncoderConfig + + +@dataclass(kw_only=True) +class Qwen3VLEncoderConfig(EncoderConfig): + """Checkpoint source and raw hidden-state selection.""" + + _target: type["Qwen3VLEncoder"] = field(default_factory=lambda: Qwen3VLEncoder) + model_name: str = "MiniMaxAI/MiniMax-H3" + """Repository or local checkpoint directory.""" + revision: str | None = None + """Pinned checkpoint revision when using the Hub.""" + cache_dir: str | None = None + """Optional Hugging Face cache directory.""" + hidden_layer: int = 50 + """Raw intermediate hidden state; must precede the final normalized state.""" + dtype: torch.dtype = torch.bfloat16 + """Weight and output precision.""" + subfolder: str = "text_encoder" + """Checkpoint partition containing the conditioner.""" + processor_subfolder: str = "processor" + """Checkpoint partition containing Qwen's media processor.""" + tokenizer_subfolder: str = "tokenizer" + """Checkpoint partition containing the presentation tokenizer.""" + local_files_only: bool = False + """Require all checkpoint files to be present locally.""" + loading_device: str | None = None + """Place checkpoint tensors directly on this device to avoid a CPU weight copy.""" + + +class Qwen3VLEncoder(Encoder): + """Selected Qwen3-VL state without a language-model projection.""" + + def __init__(self, config: Qwen3VLEncoderConfig) -> None: + super().__init__(config) + from transformers import Qwen2TokenizerFast, Qwen3VLModel, Qwen3VLProcessor + + kwargs = dict( + revision=config.revision, + cache_dir=config.cache_dir, + local_files_only=config.local_files_only, + ) + self.model = ( + Qwen3VLModel.from_pretrained( + config.model_name, + subfolder=config.subfolder, + dtype=config.dtype, + device_map=config.loading_device, + **kwargs, + ) + .eval() + .requires_grad_(False) + ) + self.processor = Qwen3VLProcessor.from_pretrained( + config.model_name, subfolder=config.processor_subfolder, **kwargs + ) + self.tokenizer = Qwen2TokenizerFast.from_pretrained( + config.model_name, subfolder=config.tokenizer_subfolder, **kwargs + ) + if ( + not 0 + <= config.hidden_layer + < self.model.config.text_config.num_hidden_layers + ): + raise ValueError( + "Qwen conditioning must select a raw intermediate hidden layer" + ) + + @torch.no_grad() + def forward(self, input: dict[str, Any]) -> torch.Tensor: + """Encode token IDs and optional processor-produced vision tensors.""" + device = self.model.device + token_ids = input["token_ids"] + input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) + modalities = torch.tensor( + self.processor.create_mm_token_type_ids([token_ids]), + dtype=torch.long, + device=device, + ) + vision = { + name: value.to(device=device, dtype=self.model.dtype) + if name.startswith("pixel_") + else value.to(device=device) + for name, value in input.get("vision_inputs", {}).items() + } + output = self.model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + mm_token_type_ids=modalities, + use_cache=False, + output_hidden_states=True, + **vision, + ) + return output.hidden_states[self.config.hidden_layer].to( + dtype=self.config.dtype + ) diff --git a/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py b/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py index 5413e68ec..e323a9ed2 100644 --- a/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py +++ b/flashdreams/tests/accelerated/multi_head_attention/test_mha_optimized.py @@ -19,8 +19,6 @@ import pytest import torch -from torch import Tensor - from flashdreams.accelerated.multi_head_attention import ( AttentionConfig, AttentionType, @@ -43,6 +41,7 @@ DTYPE_MAX, ) from flashdreams.core.attention import BlockKVCache +from torch import Tensor pytestmark = pytest.mark.ci_gpu @@ -577,7 +576,7 @@ def test_mha_optimized_quantized_projections_match_torch( ) @pytest.mark.parametrize("rope_scope", tuple(RoPEScope), ids=lambda value: value.value) @pytest.mark.parametrize( - "sdpa_backend", tuple(SDPABackend), ids=lambda value: value.value + "sdpa_backend", (SDPABackend.CUDNN, SDPABackend.FA2), ids=lambda value: value.value ) @pytest.mark.parametrize("use_tma", (False, True), ids=("no-tma", "tma")) @torch.inference_mode() @@ -631,6 +630,45 @@ def test_mha_optimized_quantized_sdpa_matches_torch( ) +@torch.inference_mode() +@pytest.mark.parametrize("fusion", (QKVFusionOption.NONE, QKVFusionOption.FULL)) +@pytest.mark.parametrize("rotary_dim", (48, 96, 128)) +def test_cacheless_partial_rope_on_ampere( + cuda_device: torch.device, + fusion: QKVFusionOption, + rotary_dim: int, +) -> None: + """Compare cacheless native attention without Hopper-only features.""" + if torch.cuda.get_device_capability(cuda_device)[0] < 8: + pytest.skip("Ampere or newer required") + config = AttentionConfig( + query_dim=128, + n_heads=2, + head_dim=128, + rope_config=RoPEConfig(style=RoPEStyle.SPLIT), + ) + reference = _TorchMHA(AttentionType.SELF_ATTENTION, config) + actual = _OptimizedMHA( + AttentionType.SELF_ATTENTION, + config, + OptimizedImplConfig( + qkv_fusion_option=fusion, + sdpa_backend=SDPABackend.TORCH, + use_tma=False, + ), + ) + actual.load_state_dict(reference.state_dict(), strict=True) + reference.to(device=cuda_device, dtype=torch.bfloat16).eval() + actual.to(device=cuda_device, dtype=torch.bfloat16).eval() + x = torch.randn(1, 7, 128, device=cuda_device, dtype=torch.bfloat16) + half = torch.randn(7, 1, 1, rotary_dim // 2, device=cuda_device) + frequencies = torch.cat((half, half), dim=-1) + _assert_close( + actual(x, rope_freqs=frequencies), reference(x, rope_freqs=frequencies) + ) + _assert_close(actual(x), reference(x)) + + @torch.inference_mode() def test_native_cudnn_fp8_sdpa_returns_independent_outputs( cuda_device: torch.device, diff --git a/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py b/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py index 6a0e80ed4..e3ace0ea3 100644 --- a/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py +++ b/flashdreams/tests/accelerated/multi_head_attention/test_mha_torch.py @@ -20,8 +20,6 @@ import pytest import torch import torch.nn.functional as F -from torch import Tensor, nn - from flashdreams.accelerated.multi_head_attention import ( AttentionConfig, AttentionType, @@ -31,10 +29,92 @@ RoPEStyle, ) from flashdreams.accelerated.multi_head_attention.torch import TorchMultiHeadAttention +from torch import Tensor, nn pytestmark = pytest.mark.ci_cpu +@pytest.mark.parametrize("rope_style", tuple(RoPEStyle)) +def test_cacheless_partial_rope_matches_manual_attention(rope_style: RoPEStyle) -> None: + """Preserve unrotated features and cacheless full-sequence semantics.""" + module = _IdentityMHA(rope_style) + x = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) / 10 + frequencies = torch.tensor([0.1, 0.4, 0.9]).reshape(3, 1, 1, 1).expand(3, 1, 1, 2) + heads = x.unsqueeze(2) + rotated = torch.cat( + ( + _apply_rope(heads[..., :2], frequencies, rope_style), + heads[..., 2:], + ), + dim=-1, + ) + actual_rotated = module._apply_rope(heads, frequencies) + torch.testing.assert_close(actual_rotated, rotated) + assert torch.equal(actual_rotated[..., 2:], heads[..., 2:]) + expected = ( + F.scaled_dot_product_attention( + rotated.transpose(1, 2), + rotated.transpose(1, 2), + heads.transpose(1, 2), + ) + .transpose(1, 2) + .flatten(-2) + ) + torch.testing.assert_close(module(x, rope_freqs=frequencies), expected) + torch.testing.assert_close( + module(x), + F.scaled_dot_product_attention( + heads.transpose(1, 2), + heads.transpose(1, 2), + heads.transpose(1, 2), + ) + .transpose(1, 2) + .flatten(-2), + ) + + +def test_cacheless_cross_attention_rejected() -> None: + module = _IdentityMHA(RoPEStyle.SPLIT, AttentionType.CROSS_ATTENTION) + with pytest.raises(ValueError, match="requires a K/V cache"): + module(torch.zeros(1, 2, 4)) + + +def test_partial_rope_rounds_coefficients_before_half_precision_products() -> None: + """Preserve video-VAE coefficient precision and the untouched feature tail.""" + module = _IdentityMHA(RoPEStyle.SPLIT) + x = torch.tensor([0.17, -0.41, 0.83, 0.32], dtype=torch.float16).reshape(1, 1, 1, 4) + freqs = torch.tensor([0.71, 0.71]).reshape(1, 1, 1, 2) + prefix = x[..., :2] + rotated = torch.cat((-prefix[..., 1:], prefix[..., :1]), dim=-1) + expected = prefix * freqs.cos().half() + rotated * freqs.sin().half() + actual = module._apply_rope(x, freqs) + assert torch.equal(actual[..., :2], expected) + assert torch.equal(actual[..., 2:], x[..., 2:]) + + +def test_projected_causal_sdpa_matches_torch() -> None: + """Keep causal audio attention and FP32 computation on the shared path.""" + from flashdreams.accelerated.multi_head_attention.sdpa import ( + SDPABackend, + scaled_dot_product_attention, + ) + + query = torch.randn(2, 5, 3, 8) + expected = F.scaled_dot_product_attention( + query.transpose(1, 2), + query.transpose(1, 2), + query.transpose(1, 2), + is_causal=True, + ).transpose(1, 2) + torch.testing.assert_close( + scaled_dot_product_attention(query, query, query, is_causal=True), expected + ) + with pytest.raises(ValueError, match="causal attention requires"): + scaled_dot_product_attention( + query, query, query, is_causal=True, backend=SDPABackend.FA2 + ) + + class _IdentityMHA(TorchMultiHeadAttention): """Provide identity projections for direct attention comparisons.""" diff --git a/flashdreams/tests/test_checkpoint_loading.py b/flashdreams/tests/test_checkpoint_loading.py index f3b129bd1..70116cbbd 100644 --- a/flashdreams/tests/test_checkpoint_loading.py +++ b/flashdreams/tests/test_checkpoint_loading.py @@ -17,6 +17,92 @@ pytestmark = pytest.mark.ci_cpu +@pytest.mark.parametrize("sharded", [False, True]) +def test_scoped_model_load_is_strict_and_ignores_unselected_shards(tmp_path, sharded): + """Select codec components without opening unrelated weight shards.""" + checkpoint_load = importlib.import_module("flashdreams.core.checkpoint.load") + model = torch.nn.Module() + model.decoder = torch.nn.Linear(2, 2, bias=False) + expected = torch.arange(4, dtype=torch.float32).reshape(2, 2) + shard = tmp_path / "decoder.safetensors" + if sharded: + save_safetensors_file({"decoder.weight": expected}, shard) + checkpoint = tmp_path / "model.safetensors.index.json" + checkpoint.write_text( + json.dumps( + { + "weight_map": { + "decoder.weight": shard.name, + "encoder.weight": "absent-unused-shard.safetensors", + } + } + ) + ) + else: + checkpoint = shard + save_safetensors_file( + {"decoder.weight": expected, "encoder.weight": torch.zeros(1)}, shard + ) + checkpoint_load.load_checkpoint( + str(checkpoint), model=model, include_prefixes=("decoder.",) + ) + torch.testing.assert_close(model.decoder.weight, expected) + with pytest.raises(RuntimeError, match="match"): + checkpoint_load.load_checkpoint( + str(checkpoint), model=model, include_prefixes=("encoder.",) + ) + with pytest.raises(RuntimeError, match="match"): + checkpoint_load.load_checkpoint(str(checkpoint), model=model) + + +def test_scoped_remote_load_filters_before_downloading(monkeypatch, tmp_path): + """Download only selected shards from a Hub index.""" + module = importlib.import_module("flashdreams.core.checkpoint.load") + index = tmp_path / "model.safetensors.index.json" + index.write_text( + json.dumps( + { + "weight_map": { + "decoder.weight": "wanted.safetensors", + "encoder.weight": "unwanted.safetensors", + } + } + ) + ) + shard = tmp_path / "wanted.safetensors" + save_safetensors_file({"decoder.weight": torch.ones(2, 2)}, shard) + monkeypatch.setattr(module, "hf_hub_download", lambda **kwargs: str(index)) + monkeypatch.setattr( + module, "_preflight_checkpoint_cache_requirement", lambda **kwargs: None + ) + monkeypatch.setattr(module, "_preflight_hf_cache", lambda **kwargs: 0) + downloaded = [] + + def fetch(**kwargs): + downloaded.extend(kwargs["shard_files"]) + return {"wanted.safetensors": str(shard)} + + monkeypatch.setattr(module, "_parallel_hf_hub_download_shards", fetch) + model = torch.nn.Module() + model.decoder = torch.nn.Linear(2, 2, bias=False) + module.load_checkpoint( + "https://huggingface.co/example/model/blob/abc/model.safetensors.index.json", + model=model, + include_prefixes=("decoder.",), + ) + assert downloaded == ["wanted.safetensors"] + + +@pytest.mark.parametrize("prefixes", [(), ("",), ("decoder",)]) +def test_scoped_load_rejects_ambiguous_prefixes(prefixes): + from flashdreams.core.checkpoint.load import load_checkpoint + + with pytest.raises(ValueError, match="prefixes"): + load_checkpoint( + "unused.safetensors", model=torch.nn.Linear(2, 2), include_prefixes=prefixes + ) + + def test_local_safetensors_uses_file_backed_loader( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/flashdreams/tests/test_rope_kernel.py b/flashdreams/tests/test_rope_kernel.py index fa5f8f71e..420a1d8fe 100644 --- a/flashdreams/tests/test_rope_kernel.py +++ b/flashdreams/tests/test_rope_kernel.py @@ -33,10 +33,9 @@ import pytest import torch -from torch import Tensor - from flashdreams.core.attention.rope import apply_rope_freqs from flashdreams.core.attention.rope_kernel import apply_rotary_pos_emb +from torch import Tensor def _load_te_apply_rope() -> Callable[..., Tensor] | None: @@ -191,6 +190,44 @@ def test_zero_freqs_is_identity(cuda_device): torch.testing.assert_close(out, x) +@pytest.mark.parametrize("x_dtype", _DTYPES) +def test_partial_prefix_view_matches_torch( + cuda_device: torch.device, + x_dtype: torch.dtype, +) -> None: + """A sliced head prefix rotates in place without touching trailing channels.""" + B, S, H, D = 2, 67, 7, 128 + rotary_dim = 96 + x = torch.randn( + B, + S, + H, + D, + device=cuda_device, + dtype=x_dtype, + ) + frequencies = _expanded_freqs( + S, rotary_dim, interleaved=False, device=cuda_device, seed=7 + ) + cos = frequencies[:, 0, 0].cos().to(x_dtype)[None, :, None, :] + sin = frequencies[:, 0, 0].sin().to(x_dtype)[None, :, None, :] + + rotary = x[..., :rotary_dim] + first, second = rotary.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + expected = x.clone() + expected[..., :rotary_dim] = rotary * cos + rotated * sin + + actual = x.clone() + prefix = actual[..., :rotary_dim] + returned = apply_rope_freqs(prefix, frequencies) + + atol, rtol = _parity_tol(x_dtype) + assert returned.data_ptr() == prefix.data_ptr() + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + assert torch.equal(actual[..., rotary_dim:], x[..., rotary_dim:]) + + @_requires_te def test_non_contiguous_x(cuda_device): """The kernel respects arbitrary strides on the B / S / H axes.""" diff --git a/flashdreams/tests/test_synchronized_sampling.py b/flashdreams/tests/test_synchronized_sampling.py new file mode 100644 index 000000000..de9131ee8 --- /dev/null +++ b/flashdreams/tests/test_synchronized_sampling.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU contracts for synchronized data-ward flow sampling.""" + +import pytest +import torch + +from flashdreams.infra.diffusion.scheduler import ( + DataFlowEulerSchedulerConfig, + sample_synchronized, +) + +pytestmark = pytest.mark.ci_cpu + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_joint_matches_scalar_reference(dtype): + samples = ( + torch.linspace(-1, 1, 12).reshape(3, 4).to(dtype), + torch.zeros(2, 3, dtype=dtype), + ) + schedulers = tuple( + DataFlowEulerSchedulerConfig(shift=shift).setup() for shift in (12, 3) + ) + calls = [] + + def predict(values, times): + calls.append(times) + return tuple(value.float() * 0.1 + 0.2 for value in values) + + result = sample_synchronized(samples, schedulers, predict) + assert len(calls) == 29 + for initial, scheduler, actual in zip(samples, schedulers, result, strict=True): + expected = initial + for index, timestep in enumerate(scheduler.timesteps): + flow = expected.float() * 0.1 + 0.2 + denoised = expected + (1 - timestep.to(dtype)) * flow + ratio = scheduler.sigmas[index + 1] / scheduler.sigmas[index] + expected = (ratio * expected.float() + (1 - ratio) * denoised.float()).to( + dtype + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + scalar = scheduler.sample( + initial, lambda value, time: value.float() * 0.1 + 0.2 + ) + torch.testing.assert_close(actual, scalar, rtol=0, atol=0) + + +def test_mismatched_schedules_fail_before_prediction(): + schedulers = tuple( + DataFlowEulerSchedulerConfig(num_inference_steps=n).setup() for n in (3, 4) + ) + with pytest.raises(ValueError, match="equal lengths"): + sample_synchronized( + (torch.zeros(1), torch.zeros(1)), + schedulers, + lambda *_: pytest.fail("must not predict"), + ) + + +def test_reject_invalid_prediction_shape(): + scheduler = DataFlowEulerSchedulerConfig(num_inference_steps=2).setup() + with pytest.raises(ValueError, match="shapes"): + sample_synchronized( + (torch.zeros(2),), (scheduler,), lambda *_: (torch.zeros(3),) + ) + + +def test_grid_survives_dtype_conversion(): + scheduler = DataFlowEulerSchedulerConfig().setup() + expected = scheduler.sigmas.clone() + scheduler.to(torch.bfloat16) + assert scheduler.timesteps.dtype == torch.float32 + torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=0) diff --git a/integrations_v2/README.md b/integrations_v2/README.md index 049e11003..4d6c5a5d7 100644 --- a/integrations_v2/README.md +++ b/integrations_v2/README.md @@ -35,6 +35,8 @@ follows is already done for you. `apps/v2v` video-to-video application. - `null_model` — not an application. A v1 pipeline the framework tests use as a fixture. +- `minimax_h3` — native joint video/audio inference with video-only T2VA, + first/last-keyframe, and ordered-reference bindings to `apps/t2v`. ## The layout diff --git a/integrations_v2/minimax_h3/README.md b/integrations_v2/minimax_h3/README.md new file mode 100644 index 000000000..804963c4f --- /dev/null +++ b/integrations_v2/minimax_h3/README.md @@ -0,0 +1,80 @@ + + + +# MiniMax H3 + +Native, video-only MiniMax H3 inference through FlashDreams v2. The three +workflows share FlashDreams' T2V application, synchronized diffusion sampler, +accelerated attention, and scoped checkpoint loader. There is no Diffusers +runtime dependency or fallback pipeline. + +## Run + +```bash +uv sync --package flashdreams-minimax-h3 --extra dev --inexact +uv run --no-sync flashdreams-run-v2 t2v-minimax-h3-t2va --output-path clip.mp4 -- \ + --prompt "A cat surfing" --duration 5 --steps 30 --seed 42 +uv run --no-sync flashdreams-run-v2 t2v-minimax-h3-fl2va --output-path clip.mp4 -- \ + --prompt "The camera moves through the scene" --image-path first.png --last-image-path last.png +uv run --no-sync flashdreams-run-v2 t2v-minimax-h3-ref2va --output-path clip.mp4 -- \ + --prompt "A cinematic scene" --reference image:subject.png --reference audio:reference.wav +``` + +Runtime options precede `--`; H3 options follow it. Use `-- --help` without +loading any weights. Resolution defaults to 768×768; runtime +`--pixel-width`/`--pixel-height` must be multiples of 32 and have aspect ratio +between 1:4 and 4:1. FPS is fixed at 24. Duration is 5–15 seconds, rounded +up to H3's `17n+5` frame grid, without exceeding 15 seconds. One block generates +the complete video; 30 scheduler points mean 29 joint model predictions. + +FL2VA accepts first, last, or both keyframes. REF2VA keeps reference order and +supports images, videos, and audio (at least one visual reference). Audio +conditioning and audio denoising remain active, but generated audio is not +decoded or muxed into the output. `--mode webrtc` uses the shared prompt UI. + +LoRA options are `--lora PATH_OR_REPO`, `--lora-weight-name FILE`, and +`--lora-scale NUMBER`. Musubi conversion is retained; a new network is loaded +for each request, so LoRA weights do not accumulate across resets. + +## Memory and acceleration + +Stage-scoped loading is always enabled: conditioning, denoising, and decoding +do not retain each other's weights. This replaces the old `--low-ram` switch. +Only requested codec subtrees are constructed and loaded. The joint transformer +still needs to fit on one GPU; blockwise CPU offload and distributed inference +are not implemented. A100 80 GB is the initial validation target, not a measured +guarantee that every resolution/duration fits. + +Default attention is FlashDreams `OptimizedMultiHeadAttention`, native BF16 +SDPA, no quantization, TMA, or duplicated QKV weights. `--attention torch` selects +FlashDreams' reference implementation. `--fuse-qkv` opts into additional fused +weight storage (roughly 11 GiB for the transformer). `--compile` uses the shared +compile helper and is opt-in. These options are **unvalidated performance +candidates**, not published speedups. CUDA graphs and quantization are deferred. + +The video codec retains FP32 weights, FP16 decode autocast, reference tiling and +overlap, and precision-preserving Torch attention. Audio references use FP32 +encoding. Checkpoint assets are pinned to +`MiniMaxAI/MiniMax-H3@42ed227ee7df40d41602854ae760620d6eb651fe`; `--model-id` accepts +a compatible local snapshot or repository, and `--revision` overrides the pin. +Transformers/Accelerate remain dependencies for headless Qwen3-VL loading; they +do not orchestrate H3 diffusion. + +## Migration and checks + +The old `flashdreams-run minimax-h3-*` commands and recovery checkpoint files +are not used. Replace them with the three v2 commands above. Existing output +and recovery files are left untouched. No background checkpoint writer exists. + +```bash +PYTHONPATH=flashdreams:apps/t2v:integrations_v2 .venv/bin/python -m pytest \ + integrations_v2/minimax_h3/tests -m ci_cpu +``` + +Native CPU checks require no checkpoint downloads. Optional reference checks +use an already-installed pinned Diffusers oracle and cached checkpoint headers; +they are skipped when unavailable. Diffusers is not installed by this package. +GPU tests are separate (`-m ci_gpu`); real-checkpoint generation/parity must be +requested explicitly. CPU correctness checks do not establish GPU video parity +or a speedup. Stage timings and peak allocated GPU memory are reported through +the shared v2 `--stats-path` output. diff --git a/integrations_v2/minimax_h3/__init__.py b/integrations_v2/minimax_h3/__init__.py new file mode 100644 index 000000000..4fd63b3c4 --- /dev/null +++ b/integrations_v2/minimax_h3/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native MiniMax H3 video generation for the FlashDreams v2 runtime.""" diff --git a/integrations_v2/minimax_h3/apps/__init__.py b/integrations_v2/minimax_h3/apps/__init__.py new file mode 100644 index 000000000..c93a64165 --- /dev/null +++ b/integrations_v2/minimax_h3/apps/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 application bindings.""" diff --git a/integrations_v2/minimax_h3/apps/t2v/README.md b/integrations_v2/minimax_h3/apps/t2v/README.md new file mode 100644 index 000000000..1eb97f90d --- /dev/null +++ b/integrations_v2/minimax_h3/apps/t2v/README.md @@ -0,0 +1,5 @@ +# MiniMax H3 T2V application + +Use `flashdreams-run-v2 t2v-minimax-h3-t2va`, +`t2v-minimax-h3-fl2va`, or `t2v-minimax-h3-ref2va`. +See the [integration guide](../../README.md) for inputs, migration, and validation. diff --git a/integrations_v2/minimax_h3/apps/t2v/__init__.py b/integrations_v2/minimax_h3/apps/t2v/__init__.py new file mode 100644 index 000000000..4afedde68 --- /dev/null +++ b/integrations_v2/minimax_h3/apps/t2v/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 bindings for the shared T2V application.""" diff --git a/integrations_v2/minimax_h3/apps/t2v/adapter.py b/integrations_v2/minimax_h3/apps/t2v/adapter.py new file mode 100644 index 000000000..41c2ce2fd --- /dev/null +++ b/integrations_v2/minimax_h3/apps/t2v/adapter.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 workflows over the shared FlashDreams v2 T2V application.""" + +import argparse +from dataclasses import replace +from pathlib import Path +from typing import Any + +from t2v import T2VApplication, T2VApplicationDefaults + +from flashdreams.accelerated.multi_head_attention.optimized import QKVFusionOption +from flashdreams.api_v2.application import IApplication +from flashdreams.infra.config import derive_config +from flashdreams.runtime_v2.session_desc import SessionDesc +from minimax_h3.config import ( + PIPELINE_MINIMAX_H3_FL2VA, + PIPELINE_MINIMAX_H3_REF2VA, + PIPELINE_MINIMAX_H3_T2VA, +) +from minimax_h3.impl.constants import FPS, align_num_frames, validate_canvas +from minimax_h3.impl.pipeline import MiniMaxH3PipelineConfig +from minimax_h3.impl.references import parse_reference_specs + + +class MiniMaxH3Application(T2VApplication): + """One-block joint audio/video inference with video-only v2 output.""" + + def __init__( + self, pipeline_config: MiniMaxH3PipelineConfig = PIPELINE_MINIMAX_H3_T2VA + ) -> None: + super().__init__( + defaults=T2VApplicationDefaults( + pipeline_config=pipeline_config, + total_blocks=1, + pixel_width=768, + pixel_height=768, + fps=FPS, + ) + ) + self._request_inputs: dict[str, Any] = {} + + def _configure_argument_parser(self, parser: argparse.ArgumentParser) -> None: + parser.add_argument("--duration", type=float, default=5.0) + parser.add_argument( + "--steps", + type=int, + default=30, + help="Scheduler grid points (30 means 29 joint predictions).", + ) + parser.add_argument("--image-path", type=Path) + parser.add_argument("--last-image-path", type=Path) + parser.add_argument( + "--reference", + action="append", + default=[], + help="Ordered image:path, video:path, or audio:path input.", + ) + parser.add_argument("--lora") + parser.add_argument("--lora-weight-name") + parser.add_argument("--lora-scale", type=float, default=1.0) + parser.add_argument( + "--attention", choices=("optimized", "torch"), default="optimized" + ) + parser.add_argument( + "--fuse-qkv", + action="store_true", + help="Opt in to extra fused-QKV weight storage; increases peak memory.", + ) + parser.add_argument("--model-id", default=self.pipeline_config.model_id) + parser.add_argument("--revision", default=self.pipeline_config.revision) + + def _apply_parsed_arguments(self, args: argparse.Namespace) -> None: + align_num_frames(args.duration) + if args.steps < 2: + raise ValueError("--steps must be at least 2 scheduler points") + if not 0 <= args.lora_scale <= 4: + raise ValueError("--lora-scale must be between 0 and 4") + references = parse_reference_specs(args.reference) if args.reference else () + workflow = self.pipeline_config.workflow + if workflow == "t2va" and ( + args.image_path or args.last_image_path or references + ): + raise ValueError("t2va does not accept keyframes or references") + if workflow == "fl2va" and ( + not (args.image_path or args.last_image_path) or references + ): + raise ValueError("fl2va requires first/last keyframes and no references") + if workflow == "ref2va" and ( + not references or args.image_path or args.last_image_path + ): + raise ValueError("ref2va requires ordered references and no keyframes") + for path in (args.image_path, args.last_image_path): + if path is not None and not path.is_file(): + raise FileNotFoundError(path) + self._request_inputs = dict( + duration=args.duration, + image_path=args.image_path, + last_image_path=args.last_image_path, + references=references, + lora=args.lora, + lora_weight_name=args.lora_weight_name, + lora_scale=args.lora_scale, + ) + config = self.pipeline_config + optimized = replace( + config.transformer.optimized_impl, + qkv_fusion_option=QKVFusionOption.FULL + if args.fuse_qkv + else QKVFusionOption.NONE, + ) + self._pipeline_config = replace( + config, + model_id=args.model_id, + revision=args.revision, + transformer=replace( + config.transformer, + attention_backend=args.attention, + optimized_impl=optimized, + ), + scheduler=replace(config.scheduler, num_inference_steps=args.steps), + audio_scheduler=replace( + config.audio_scheduler, num_inference_steps=args.steps + ), + ) + + def _cache_initialization_kwargs(self, session_desc: SessionDesc) -> dict[str, Any]: + return dict(self._request_inputs) + + def _validate_total_blocks(self, total_blocks: int) -> None: + if total_blocks != 1: + raise ValueError("MiniMax H3 generates its complete clip in one block") + + def _validate_frame_size(self, session_desc: SessionDesc, pipeline: Any) -> None: + validate_canvas(session_desc.video_width, session_desc.video_height) + if session_desc.frames_per_second_for_step != FPS: + raise ValueError("MiniMax H3 requires 24 fps") + + def _apply_compile_override(self, pipeline_config: Any, enabled: bool) -> Any: + return derive_config(pipeline_config, compile_network=enabled) + + def _apply_seed_override(self, pipeline_config: Any, seed: int) -> Any: + return derive_config(pipeline_config, seed=seed) + + +def create_app() -> IApplication: + """Create the prompt-only H3 application without loading weights.""" + return MiniMaxH3Application() + + +def create_app_fl2va() -> IApplication: + """Create the first/last-keyframe H3 application.""" + return MiniMaxH3Application(PIPELINE_MINIMAX_H3_FL2VA) + + +def create_app_ref2va() -> IApplication: + """Create the ordered-reference H3 application.""" + return MiniMaxH3Application(PIPELINE_MINIMAX_H3_REF2VA) diff --git a/integrations_v2/minimax_h3/config.py b/integrations_v2/minimax_h3/config.py new file mode 100644 index 000000000..303f692a7 --- /dev/null +++ b/integrations_v2/minimax_h3/config.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native MiniMax H3 workflow configurations.""" + +from flashdreams.infra.config import derive_config +from minimax_h3.impl.pipeline import MiniMaxH3PipelineConfig + +PIPELINE_MINIMAX_H3_T2VA = MiniMaxH3PipelineConfig( + name="minimax-h3-t2va", workflow="t2va" +) +PIPELINE_MINIMAX_H3_FL2VA = derive_config( + PIPELINE_MINIMAX_H3_T2VA, name="minimax-h3-fl2va", workflow="fl2va" +) +PIPELINE_MINIMAX_H3_REF2VA = derive_config( + PIPELINE_MINIMAX_H3_T2VA, name="minimax-h3-ref2va", workflow="ref2va" +) + +MINIMAX_H3_CONFIGS = { + config.name: config + for config in ( + PIPELINE_MINIMAX_H3_T2VA, + PIPELINE_MINIMAX_H3_FL2VA, + PIPELINE_MINIMAX_H3_REF2VA, + ) +} diff --git a/integrations_v2/minimax_h3/impl/__init__.py b/integrations_v2/minimax_h3/impl/__init__.py new file mode 100644 index 000000000..01a8560e8 --- /dev/null +++ b/integrations_v2/minimax_h3/impl/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 model components and request conditioning.""" diff --git a/integrations_v2/minimax_h3/impl/audio_encoder.py b/integrations_v2/minimax_h3/impl/audio_encoder.py new file mode 100644 index 000000000..6c825e85b --- /dev/null +++ b/integrations_v2/minimax_h3/impl/audio_encoder.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright 2025 The MiniMax authors and The HuggingFace Team. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""MiniMax H3 reference-audio encoder without decoder or posterior wrappers.""" + +# Adapted from Diffusers' MiniMax H3 audio autoencoder: encoder-only modules +# with native checkpoint names and shared FlashDreams causal attention. + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass, fields +from typing import Any + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from torch.nn.utils import weight_norm + +from flashdreams.accelerated.multi_head_attention.sdpa import ( + scaled_dot_product_attention, +) + + +@dataclass(frozen=True, kw_only=True) +class AudioEncoderConfig: + """Encoder-only checkpoint geometry and per-channel latent normalization.""" + + encoder_dim: int = 64 + """Initial waveform-encoder feature width.""" + encoder_rates: tuple[int, ...] = (2, 4, 4, 5, 5) + """Strides of the channel-doubling encoder blocks.""" + latent_dim: int = 2048 + """Encoder trunk width before the attention projection.""" + latent_channels: int = 32 + """Channels in the raw audio latent domain.""" + num_attention_heads: int = 8 + """Causal projection heads, mean-pooled before latent projection.""" + sampling_rate: int = 32000 + """Reference waveform sample rate.""" + latents_mean: tuple[float, ...] = (0.0,) * 32 + """Raw latent means applied by the conditioning facade.""" + latents_std: tuple[float, ...] = (1.0,) * 32 + """Raw latent standard deviations applied by the conditioning facade.""" + + @classmethod + def from_dict(cls, values: Mapping[str, Any]) -> AudioEncoderConfig: + """Read a full audio VAE config while excluding decoder-only settings.""" + names = {item.name for item in fields(cls)} + decoder_names = { + "decoder_dim", + "decoder_rates", + "decoder_kernel_sizes", + "resblock_kernel_sizes", + "resblock_dilation_sizes", + } + unknown = ( + {key for key in values if not key.startswith("_")} - names - decoder_names + ) + if unknown: + raise ValueError( + f"Unknown audio VAE configuration fields: {sorted(unknown)}" + ) + return cls(**{key: value for key, value in values.items() if key in names}) + + @property + def hop_length(self) -> int: + """Return waveform samples per encoded latent.""" + return math.prod(self.encoder_rates) + + +class _Snake(nn.Module): + """DAC encoder activation with checkpoint-native per-channel frequency.""" + + def __init__(self, channels: int): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x: Tensor) -> Tensor: + """Add the learned periodic residual to the waveform features.""" + return x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2) + + +class _ResidualUnit(nn.Module): + """Weight-normalized DAC residual unit.""" + + def __init__(self, dim: int, dilation: int): + super().__init__() + self.block = nn.Sequential( + _Snake(dim), + weight_norm( + nn.Conv1d(dim, dim, 7, dilation=dilation, padding=3 * dilation) + ), + _Snake(dim), + weight_norm(nn.Conv1d(dim, dim, 1)), + ) + + def forward(self, x: Tensor) -> Tensor: + """Add the residual, center-cropping its shortcut when required.""" + residual = self.block(x) + pad = (x.shape[-1] - residual.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return x + residual + + +class _EncoderBlock(nn.Module): + """Three residual units followed by strided channel doubling.""" + + def __init__(self, dim: int, stride: int): + super().__init__() + self.block = nn.Sequential( + *[_ResidualUnit(dim // 2, dilation) for dilation in (1, 3, 9)], + _Snake(dim // 2), + weight_norm( + nn.Conv1d( + dim // 2, + dim, + 2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + ), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class _Encoder(nn.Module): + """Mono waveform encoder retaining checkpoint-native sequential indices.""" + + def __init__(self, config: AudioEncoderConfig): + super().__init__() + dim = config.encoder_dim + blocks: list[nn.Module] = [weight_norm(nn.Conv1d(1, dim, 7, padding=3))] + for stride in config.encoder_rates: + dim *= 2 + blocks.append(_EncoderBlock(dim, stride)) + blocks.extend( + [_Snake(dim), weight_norm(nn.Conv1d(dim, config.latent_dim, 3, padding=1))] + ) + self.block = nn.Sequential(*blocks) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class _CausalAttention(nn.Module): + """Causal attention with head-mean and adaptive feature pooling.""" + + def __init__(self, in_dim: int, out_dim: int, num_heads: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = in_dim // num_heads + self.out_dim = out_dim + self.qkv = nn.Linear(in_dim, in_dim * 3, bias=False) + self.q_bias = nn.Parameter(torch.zeros(in_dim)) + self.v_bias = nn.Parameter(torch.zeros(in_dim)) + self.register_buffer("zero_k_bias", torch.zeros(in_dim)) + self.proj = nn.Linear(out_dim, out_dim) + + def forward(self, x: Tensor) -> Tensor: + """Attend causally, then pool heads and feature width independently.""" + b, length, _ = x.shape + qkv = F.linear( + x, self.qkv.weight, torch.cat([self.q_bias, self.zero_k_bias, self.v_bias]) + ) + q, k, v = ( + qkv.reshape(b, length, 3, self.num_heads, self.head_dim) + .permute(2, 0, 1, 3, 4) + .unbind(0) + ) + x = scaled_dot_product_attention(q, k, v, is_causal=True) + x = F.adaptive_avg_pool1d(x.mean(dim=2), self.out_dim) + return self.proj(x) + + +class _GeGLU(nn.Module): + """Pre-normalized GeGLU projection with checkpoint-native names.""" + + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.w0 = nn.Linear(dim, hidden_dim) + self.w1 = nn.Linear(dim, hidden_dim) + self.w2 = nn.Linear(hidden_dim, dim) + + def forward(self, x: Tensor) -> Tensor: + x = self.norm(x) + return self.w2(F.gelu(self.w0(x), approximate="tanh") * self.w1(x)) + + +class _AttentionProjection(nn.Module): + """Residual causal-attention projection from trunk to latent width.""" + + def __init__(self, in_dim: int, out_dim: int, num_heads: int): + super().__init__() + self.norm1 = nn.LayerNorm(in_dim) + self.attn = _CausalAttention(in_dim, out_dim, num_heads) + self.proj = nn.Linear(in_dim, out_dim) + self.norm3 = nn.LayerNorm(in_dim) + self.norm2 = nn.LayerNorm(out_dim) + self.mlp = _GeGLU(out_dim, out_dim * 2) + + def forward(self, x: Tensor) -> Tensor: + x = self.proj(self.norm3(x)) + self.attn(self.norm1(x)) + return x + self.mlp(self.norm2(x)) + + +class MiniMaxH3AudioEncoder(nn.Module): + """FP32 reference-audio encoder returning only the consumed posterior mean.""" + + def __init__(self, config: AudioEncoderConfig): + super().__init__() + if any( + value <= 0 + for value in ( + config.encoder_dim, + config.latent_dim, + config.latent_channels, + config.num_attention_heads, + config.sampling_rate, + ) + ): + raise ValueError( + "Audio encoder widths, head count and sampling rate must be positive" + ) + if ( + len(config.latents_mean) != config.latent_channels + or len(config.latents_std) != config.latent_channels + ): + raise ValueError("Audio latent normalization must match latent_channels") + if any(not math.isfinite(value) for value in config.latents_mean) or any( + not math.isfinite(value) or value <= 0 for value in config.latents_std + ): + raise ValueError( + "Audio latent means must be finite and standard deviations positive" + ) + if ( + config.latent_dim % config.latent_channels + or config.latent_dim % config.num_attention_heads + ): + raise ValueError( + "Audio trunk width must be divisible by latent channels and attention heads" + ) + if not config.encoder_rates or any(rate <= 0 for rate in config.encoder_rates): + raise ValueError("Audio encoder strides must be positive") + self.config = config + self.encoder = _Encoder(config) + self.pre_block = _AttentionProjection( + config.latent_dim, config.latent_channels, config.num_attention_heads + ) + self.mean_proj = nn.Conv1d(config.latent_channels, config.latent_channels, 1) + + def encode(self, waveform: Tensor) -> Tensor: + """Return raw posterior means for mono waveforms shaped ``[B,1,S]``. + + Stereo references occupy two batch items. Right-pad to a whole encoder + hop; the caller applies the checkpoint's per-channel mean and std. + """ + if ( + waveform.ndim != 3 + or waveform.shape[1] != 1 + or any(size <= 0 for size in waveform.shape) + ): + raise ValueError( + f"Expected nonempty mono waveform [B,1,S], got {tuple(waveform.shape)}" + ) + if next(self.parameters()).dtype != torch.float32: + raise ValueError("H3 audio encoder weights must remain float32") + with torch.autocast(device_type=waveform.device.type, enabled=False): + waveform = F.pad( + waveform.float(), (0, (-waveform.shape[-1]) % self.config.hop_length) + ) + x = self.encoder(waveform) + x = self.pre_block(x.transpose(1, 2)).transpose(1, 2) + return self.mean_proj(x) + + def forward(self, waveform: Tensor) -> Tensor: + """Return raw audio posterior means.""" + return self.encode(waveform) diff --git a/integrations_v2/minimax_h3/impl/conditioning.py b/integrations_v2/minimax_h3/impl/conditioning.py new file mode 100644 index 000000000..d3ffef194 --- /dev/null +++ b/integrations_v2/minimax_h3/impl/conditioning.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 MiniMax and HuggingFace Teams +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""MiniMax H3 presentation packing and staged native conditioning.""" + +from __future__ import annotations + +import math +from collections.abc import Callable +from pathlib import Path +from typing import Any +import numpy as np +import torch +from PIL import Image, ImageOps +from flashdreams.infra.acceleration.encoder_lifecycle import ( + run_one_shot_stage, + collect_and_release_cuda_memory, +) +from .references import MiniMaxH3Reference, MiniMaxH3ReferenceSpec, load_references + +MINIMAX_H3_MIN_ASPECT_RATIO = 1 / 4 +MINIMAX_H3_MAX_ASPECT_RATIO = 4 + + +def _stage(factory: Callable, operation: Callable) -> Any: + holder = [] + + def compute(): + holder.append(factory()) + return operation(holder[0]) + + def release(): + holder.clear() + collect_and_release_cuda_memory() + + return run_one_shot_stage(compute, release=release) + + +def normalize_references( + references: list[MiniMaxH3Reference], num_frames: int +) -> list[MiniMaxH3Reference]: + """Normalize reference media at the released image, video and audio rates.""" + normalized = [] + for ref in references: + audio = ref.audio + if audio is not None: + if ref.sample_rate != 32000: + raise ValueError("Reference loader must resample audio to 32000 Hz") + audio = audio.float()[:, : int(num_frames / 24 * 32000)] + if audio.shape[0] == 1: + audio = audio.expand(2, -1).contiguous() + if audio.shape[0] != 2 or not audio.shape[1]: + raise ValueError( + "Reference audio must contain nonempty mono or stereo samples" + ) + if ref.kind == "image": + image = ref.image + width, height = image.size + if not 1 / 4 <= width / height <= 4: + raise ValueError( + "Reference image aspect ratio must be between 1:4 and 4:1" + ) + scale = 2048 / min(width, height) + size = ( + max(32, round(width * scale / 32) * 32), + max(32, round(height * scale / 32) * 32), + ) + normalized.append( + MiniMaxH3Reference( + kind="image", image=image.resize(size, Image.Resampling.LANCZOS) + ) + ) + elif ref.kind == "video": + normalized.append( + MiniMaxH3Reference( + kind="video", + frames=_normalize_video_condition( + ref.frames, ref.fps, num_frames, 32, 768, 768 * 1344, 24 + ), + fps=24, + audio=audio, + sample_rate=32000 if audio is not None else None, + ) + ) + else: + normalized.append( + MiniMaxH3Reference(kind="audio", audio=audio, sample_rate=32000) + ) + return normalized + + +def prepare_keyframes( + image_path: Path | None, last_image_path: Path | None, width: int, height: int +) -> tuple[list[Image.Image], tuple[str, ...]]: + """Stretch the geometry anchor and cover-crop its optional follower.""" + frames, anchors = [], [] + for anchor, path in (("first", image_path), ("last", last_image_path)): + if path is None: + continue + with Image.open(path) as source: + frame = ImageOps.exif_transpose(source).convert("RGB") + if not frames: + frame = frame.resize((width, height), Image.Resampling.LANCZOS) + else: + scale = max(width / frame.width, height / frame.height) + size = ( + max(width, round(frame.width * scale)), + max(height, round(frame.height * scale)), + ) + left, top = (size[0] - width) // 2, (size[1] - height) // 2 + frame = frame.resize(size, Image.Resampling.LANCZOS).crop( + (left, top, left + width, top + height) + ) + frames.append(frame) + anchors.append(anchor) + return frames, tuple(anchors) + + +def encode_visual_condition(encoder: Any, pixels: torch.Tensor) -> torch.Tensor: + """Sample seed-42 visual conditioning and apply the released fp16 rounding.""" + mean = pixels.new_tensor((0.485, 0.456, 0.406), dtype=torch.float32).view( + 1, 3, 1, 1, 1 + ) + std = pixels.new_tensor((0.229, 0.224, 0.225), dtype=torch.float32).view( + 1, 3, 1, 1, 1 + ) + pixels = (pixels.float() / 255 - mean) / std + latents = encoder.sample(pixels, generator=torch.Generator().manual_seed(42)) + latents = latents.half().float().cpu() + mean = torch.tensor(encoder.config.latents_mean).view(1, -1, 1, 1, 1) + std = torch.tensor(encoder.config.latents_std).view(1, -1, 1, 1, 1) + return (latents - mean) / std + + +def condition_request( + *, + prompt: str, + workflow: str, + width: int, + height: int, + num_frames: int, + qwen_encoder_factory: Callable, + image_path: Path | None = None, + last_image_path: Path | None = None, + references: tuple[MiniMaxH3ReferenceSpec, ...] = (), + video_encoder_factory: Callable | None = None, + audio_encoder_factory: Callable | None = None, + device: str | torch.device = "cpu", +) -> dict[str, Any]: + """Prepare one request with only one heavyweight encoder resident at a time.""" + if not isinstance(prompt, str) or not prompt.strip(): + raise ValueError("Prompt must be nonempty text") + if min(width, height) <= 0 or width % 32 or height % 32: + raise ValueError("Canvas dimensions must be positive multiples of 32") + if not 1 / 4 <= width / height <= 4: + raise ValueError("Canvas aspect ratio must be between 1:4 and 4:1") + if num_frames % 17 != 5 or not 5 <= num_frames / 24 <= 15: + raise ValueError( + "Frame count must be 17*n+5 and duration between 5 and 15 seconds" + ) + if workflow not in {"t2va", "fl2va", "ref2va"}: + raise ValueError(f"Unsupported H3 workflow: {workflow}") + if workflow != "fl2va" and (image_path is not None or last_image_path is not None): + raise ValueError("Keyframes require fl2va") + if workflow != "ref2va" and references: + raise ValueError("Ordered references require ref2va") + keyframes, anchors = prepare_keyframes(image_path, last_image_path, width, height) + if workflow == "fl2va" and not keyframes: + raise ValueError("fl2va requires a first or last keyframe") + refs = ( + normalize_references(load_references(references), num_frames) + if workflow == "ref2va" + else [] + ) + visual_refs = ( + [MiniMaxH3Reference(kind="image", image=image) for image in keyframes] + if keyframes + else refs + ) + + def encode_video(encoder): + conditions = [] + for reference in visual_refs: + if reference.kind == "image": + pixels = torch.from_numpy(np.array(reference.image)).permute(2, 0, 1)[ + None, :, None + ] + elif reference.kind == "video": + count = max(1, (len(reference.frames) - 5) // 17) * 17 + 5 + pixels = torch.from_numpy(reference.frames[:count].copy()).permute( + 3, 0, 1, 2 + )[None] + else: + continue + conditions.append(encode_visual_condition(encoder, pixels.to(device))) + return conditions + + has_visual = any(ref.kind in {"image", "video"} for ref in visual_refs) + if has_visual and video_encoder_factory is None: + raise ValueError("Visual conditioning requires a video encoder factory") + conditions = _stage(video_encoder_factory, encode_video) if has_visual else [] + + def encode_audio(encoder): + mean = torch.tensor(encoder.config.latents_mean).view(1, 1, -1) + std = torch.tensor(encoder.config.latents_std).view(1, 1, -1) + return [ + ( + ( + encoder.encode(ref.audio.to(device)[:, None]) + .float() + .cpu() + .transpose(1, 2) + - mean + ) + / std + ).reshape(-1, 32) + for ref in refs + if ref.has_audio + ] + + has_audio = any(ref.has_audio for ref in refs) + if has_audio and audio_encoder_factory is None: + raise ValueError("Audio references require an audio encoder factory") + audio_conditions = _stage(audio_encoder_factory, encode_audio) if has_audio else [] + + def encode_text(encoder): + vision, image_counts, video_counts, timestamps = _gather_vision_features( + encoder.processor, visual_refs, 24 + ) + ids, tags = _build_presentation( + encoder.tokenizer, + prompt, + visual_refs, + image_counts, + video_counts, + timestamps, + ) + embeddings = encoder({"token_ids": ids, "vision_inputs": vision}) + return { + "prompt_embeds": embeddings, + "text_token_tags": torch.tensor(tags, dtype=torch.long), + } + + return { + **_stage(qwen_encoder_factory, encode_text), + "condition_latents": conditions, + "audio_condition_latents": audio_conditions, + "height": height, + "width": width, + "num_frames": num_frames, + "keyframe_anchors": anchors, + "normalized_references": refs, + } + + +def resolve_canvas_size( + aspect_width: float, + aspect_height: float, + canvas_multiple: int, + short_edge: int, + max_pixels: int, + min_aspect_ratio: float = MINIMAX_H3_MIN_ASPECT_RATIO, + max_aspect_ratio: float = MINIMAX_H3_MAX_ASPECT_RATIO, +) -> tuple[int, int]: + """Resolve a display aspect ratio into a MiniMax-H3 canvas.""" + if aspect_width <= 0 or aspect_height <= 0: + raise ValueError( + f"The aspect ratio must be positive, got {aspect_width}:{aspect_height}." + ) + + ratio = aspect_width / aspect_height + if not min_aspect_ratio <= ratio <= max_aspect_ratio: + raise ValueError( + f"MiniMax-H3 supports aspect ratios from 1:{1 / min_aspect_ratio:g} to {max_aspect_ratio:g}:1, got " + f"{aspect_width}:{aspect_height} ({ratio:g})." + ) + + if ratio >= 1.0: + width, height = short_edge * ratio, float(short_edge) + else: + width, height = float(short_edge), short_edge / ratio + + area = width * height + if area > max_pixels: + scale = (max_pixels / area) ** 0.5 + width, height = width * scale, height * scale + + multiple = canvas_multiple + return max(multiple, round(height / multiple) * multiple), max( + multiple, round(width / multiple) * multiple + ) + + +def _normalize_video_condition( + frames, + fps: float, + num_frames: int, + canvas_multiple: int, + canvas_short_edge: int, + canvas_max_pixels: int, + target_fps: float, +) -> np.ndarray: + """Normalize a video reference's frames: any accepted layout, onto `uint8` at 24 fps, truncated to the generated""" + # Any accepted layout onto `uint8` THWC. A `torch.Tensor` is channels-first, as everywhere else in + # diffusers, and a `np.ndarray` channels-last; floating point values are read over `[0, 1]`. + if isinstance(frames, list): + frames = np.stack([np.asarray(frame.convert("RGB")) for frame in frames]) + if isinstance(frames, torch.Tensor): + frames = frames.movedim(-3, -1).cpu().numpy() + frames = np.asarray(frames) + if frames.dtype != np.uint8: + frames = (frames * 255.0).round().clip(0, 255).astype(np.uint8) + if frames.ndim != 4 or frames.shape[3] != 3: + raise ValueError( + f"A reference video must be `(num_frames, height, width, 3)` RGB frames, got {tuple(frames.shape)}." + ) + + # Onto MiniMax-H3's 24 fps grid: every frame is held until the slot of the next one, and the last one until + # the slot the stream's end rounds to. + if not math.isfinite(fps) or fps <= 0: + raise ValueError( + f"A reference video must have a positive frame rate, got {fps}." + ) + if fps != target_fps: + scale = target_fps / fps + slots = np.floor(np.arange(frames.shape[0]) * scale + 0.5).astype(np.int64) + frames = np.repeat( + frames, + np.diff(slots, append=math.floor(frames.shape[0] * scale + 0.5)), + axis=0, + ) + + # Truncated to the generated frame count and put on the canvas of its *own* aspect ratio — the same rule the + # target canvas follows, unlike an image reference. + frames = frames[:num_frames] + if not len(frames): + raise ValueError("Reference video is too short to contain a frame at 24 fps") + height, width = resolve_canvas_size( + frames.shape[2], + frames.shape[1], + canvas_multiple, + canvas_short_edge, + canvas_max_pixels, + ) + if frames.shape[1:3] == (height, width): + return frames + return np.stack( + [ + np.asarray( + Image.fromarray(frame).resize((width, height), Image.Resampling.LANCZOS) + ) + for frame in frames + ] + ) + + +def _sample_video_condition_frames( + frames: np.ndarray, fps: float, sample_fps: float, temporal_patch: int +) -> tuple[list[np.ndarray], list[float]]: + """Sample the frames the conditioner sees from a normalized reference video, and label their vision blocks.""" + stride = fps / sample_fps + indices, cursor = [], 0.0 + while round(cursor) < frames.shape[0]: + if not indices or round(cursor) > indices[-1]: + indices.append(round(cursor)) + cursor += stride + if len(indices) < temporal_patch: + minimum = round((temporal_patch - 1) * stride) + 1 + raise ValueError( + f"A reference video is read at {sample_fps:g} fps and its sampled frames are merged in groups of " + f"{temporal_patch}, so it must run at least {minimum} frames at {fps:g} fps " + f"({minimum / fps:.2g} seconds), got {frames.shape[0]}." + ) + + timestamps = [index / sample_fps for index in range(len(indices))] + timestamps += [timestamps[-1]] * (-len(timestamps) % temporal_patch) + block_timestamps = [ + (timestamps[index] + timestamps[index + temporal_patch - 1]) / 2 + for index in range(0, len(timestamps), temporal_patch) + ] + return [frames[index] for index in indices], block_timestamps + + +def _gather_vision_features( + processor, references: list[MiniMaxH3Reference], fps: float +) -> tuple[dict, list[int], list[int], list[list[float]]]: + """Run the references' pixels through the conditioner's processors, batched per modality.""" + merge_size = processor.image_processor.merge_size**2 + vision_inputs = {} + + image_token_counts = [] + images = [reference.image for reference in references if reference.kind == "image"] + if images: + image_features = processor.image_processor(images=images, return_tensors="pt") + vision_inputs["pixel_values"] = image_features["pixel_values"] + vision_inputs["image_grid_thw"] = image_features["image_grid_thw"] + image_token_counts = [ + int(grid.prod()) // merge_size for grid in image_features["image_grid_thw"] + ] + + video_block_token_counts, video_block_timestamps = [], [] + videos = [reference for reference in references if reference.kind == "video"] + if videos: + temporal_patch = processor.video_processor.temporal_patch_size + sampled = [ + _sample_video_condition_frames(reference.frames, fps, 2.0, temporal_patch) + for reference in videos + ] + video_block_timestamps = [timestamps for _, timestamps in sampled] + video_features = processor.video_processor( + videos=[np.stack(frames) for frames, _ in sampled], + do_sample_frames=False, + return_tensors="pt", + ) + vision_inputs["pixel_values_videos"] = video_features["pixel_values_videos"] + vision_inputs["video_grid_thw"] = video_features["video_grid_thw"] + video_block_token_counts = [ + int(grid[1]) * int(grid[2]) // merge_size + for grid in video_features["video_grid_thw"] + ] + for timestamps, grid in zip( + video_block_timestamps, video_features["video_grid_thw"] + ): + if int(grid[0]) != len(timestamps): + raise ValueError( + f"The processor merged a reference video into {int(grid[0])} vision blocks, but MiniMax-H3 " + f"labels {len(timestamps)} of them." + ) + + return ( + vision_inputs, + image_token_counts, + video_block_token_counts, + video_block_timestamps, + ) + + +def _build_presentation( + tokenizer, + prompt: str, + references: list[MiniMaxH3Reference], + image_token_counts: list[int], + video_block_token_counts: list[int], + video_block_timestamps: list[list[float]], + text_tag: int = 1, + video_tag: int = 0, +) -> tuple[list[int], list[int]]: + """Tokenize MiniMax-H3's presentation of a `ref2va` request.""" + + def text(value: str) -> tuple[list[int], list[int]]: + token_ids = tokenizer(value, add_special_tokens=False)["input_ids"] + return token_ids, [text_tag] * len(token_ids) + + def vision(pad_token: str, num_tokens: int) -> tuple[list[int], list[int]]: + token_ids = ( + [tokenizer.convert_tokens_to_ids("<|vision_start|>")] + + [tokenizer.convert_tokens_to_ids(pad_token)] * num_tokens + + [tokenizer.convert_tokens_to_ids("<|vision_end|>")] + ) + return token_ids, [video_tag] * len(token_ids) + + token_ids, token_tags = [], [] + + def emit(segment: tuple[list[int], list[int]]) -> None: + token_ids.extend(segment[0]) + token_tags.extend(segment[1]) + + counts = {"image": 0, "video": 0, "audio": 0} + for reference in references: + if reference.has_audio: + counts["audio"] += 1 + emit(text(f"