diff --git a/docs/source/en/api/cache.md b/docs/source/en/api/cache.md index a5ed8751118d..5d0d16585013 100644 --- a/docs/source/en/api/cache.md +++ b/docs/source/en/api/cache.md @@ -46,3 +46,9 @@ Cache methods speedup diffusion transformers by storing and reusing intermediate [[autodoc]] MagCacheConfig [[autodoc]] apply_mag_cache + +## SeaCacheConfig + +[[autodoc]] SeaCacheConfig + +[[autodoc]] apply_sea_cache diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index ae8de9ca66c8..36f5f089339f 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -660,6 +660,36 @@ if result.action is not None: +## SeaCache + +SeaCache is disabled by default. Cosmos 3 supports enabling it explicitly with [`SeaCacheConfig`]. SeaCache reuses +transformer residuals when the Spectral-Evolution-Aware indicator changes slowly, reducing the number of full +transformer executions. Enable it on the transformer with scheduler metadata callbacks from the pipeline: + +```python +import torch +from diffusers import Cosmos3OmniPipeline, SeaCacheConfig + +pipe = Cosmos3OmniPipeline.from_pretrained( + "nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda" +) + +pipe.transformer.enable_cache( + SeaCacheConfig( + threshold=0.2, + max_consecutive_cached=2, + current_step_callback=lambda: pipe.current_step_index, + current_sigma_callback=lambda: pipe.current_sigma, + num_inference_steps_callback=lambda: pipe.num_timesteps, + ) +) +``` + +The same model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and +[`Cosmos3DistilledModularPipeline`]. SeaCache is approximate and can change generated outputs. Disable it with +`pipe.transformer.disable_cache()` when you need every denoising step to execute the full transformer. Cache state is +reset after each pipeline call, and conditional and unconditional guidance branches keep independent histories. + ## Context parallelism For long videos or high resolutions, a single forward pass can exceed the memory and latency budget of one GPU. Cosmos 3 supports **context parallelism (CP)** to shard the sequence dimension across multiple GPUs, splitting the attention computation so each device holds only a slice of the tokens. diff --git a/docs/source/en/optimization/cache.md b/docs/source/en/optimization/cache.md index 079f073b73f0..ad65fceeaada 100644 --- a/docs/source/en/optimization/cache.md +++ b/docs/source/en/optimization/cache.md @@ -68,6 +68,34 @@ config = FasterCacheConfig( pipeline.transformer.enable_cache(config) ``` +## SeaCache + +[SeaCache](https://huggingface.co/papers/2602.18993) compares Spectral-Evolution-Aware (SEA) indicators between +successive denoising steps. When the accumulated indicator change remains below a threshold, it skips the expensive +transformer block stack and predicts its output from cached residuals. + +SeaCache is disabled by default. Enable it on the transformer and provide callbacks for the active scheduler step, +sigma, and number of inference steps: + +```python +from diffusers import Cosmos3OmniPipeline, SeaCacheConfig + +pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano") +pipe.transformer.enable_cache( + SeaCacheConfig( + threshold=0.2, + max_consecutive_cached=2, + current_step_callback=lambda: pipe.current_step_index, + current_sigma_callback=lambda: pipe.current_sigma, + num_inference_steps_callback=lambda: pipe.num_timesteps, + ) +) +``` + +This model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and +[`Cosmos3DistilledModularPipeline`]. SeaCache is an approximate optimization and may change generated outputs. Call +`pipe.transformer.disable_cache()` when you need every denoising step to execute the full transformer. + ## FirstBlockCache [FirstBlock Cache](https://huggingface.co/docs/diffusers/main/en/api/cache#diffusers.FirstBlockCacheConfig) checks how much the early layers of the denoiser changes from one timestep to the next. If the change is small, the model skips the expensive later layers and reuses the previous output. diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 1fc34e6bdbf6..68789af008f5 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -204,6 +204,7 @@ "LayerSkipConfig", "MagCacheConfig", "PyramidAttentionBroadcastConfig", + "SeaCacheConfig", "SmoothedEnergyGuidanceConfig", "TaylorSeerCacheConfig", "TextKVCacheConfig", @@ -212,6 +213,7 @@ "apply_layer_skip", "apply_mag_cache", "apply_pyramid_attention_broadcast", + "apply_sea_cache", "apply_taylorseer_cache", "apply_text_kv_cache", ] @@ -1081,6 +1083,7 @@ LayerSkipConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, + SeaCacheConfig, SmoothedEnergyGuidanceConfig, TaylorSeerCacheConfig, TextKVCacheConfig, @@ -1089,6 +1092,7 @@ apply_layer_skip, apply_mag_cache, apply_pyramid_attention_broadcast, + apply_sea_cache, apply_taylorseer_cache, apply_text_kv_cache, ) diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index d999ab32d6d7..70399f9ce805 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -25,6 +25,7 @@ from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook from .mag_cache import MagCacheConfig, apply_mag_cache from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast + from .sea_cache import SeaCacheConfig, apply_sea_cache from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache from .tensor_parallel import apply_tensor_parallel diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index 9cbe5bc8108f..bfc6c59477c8 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -27,6 +27,8 @@ class TransformerBlockMetadata: return_hidden_states_index: int = None return_encoder_hidden_states_index: int = None hidden_states_argument_name: str = "hidden_states" + encoder_hidden_states_argument_name: str = "encoder_hidden_states" + hidden_states_norm_module_name: str = None _cls: Type = None _cached_parameter_indices: dict[str, int] = None @@ -174,6 +176,7 @@ def _register_transformer_blocks_metadata(): from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock from ..models.transformers.transformer_bria import BriaTransformerBlock from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock + from ..models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer from ..models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock from ..models.transformers.transformer_hunyuan_video import ( HunyuanVideoSingleTransformerBlock, @@ -230,6 +233,18 @@ def _register_transformer_blocks_metadata(): ), ) + # Cosmos 3 + TransformerBlockRegistry.register( + model_class=Cosmos3VLTextMoTDecoderLayer, + metadata=TransformerBlockMetadata( + return_hidden_states_index=1, + return_encoder_hidden_states_index=0, + hidden_states_argument_name="gen_seq", + encoder_hidden_states_argument_name="und_seq", + hidden_states_norm_module_name="input_layernorm_moe_gen", + ), + ) + # Flux TransformerBlockRegistry.register( model_class=FluxTransformerBlock, diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py new file mode 100644 index 000000000000..f277f90f0ce7 --- /dev/null +++ b/src/diffusers/hooks/sea_cache.py @@ -0,0 +1,992 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import math +from dataclasses import dataclass +from typing import Any, Callable, Literal + +import torch + +from ..utils import logging +from ..utils.torch_utils import unwrap_module +from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS +from ._helpers import TransformerBlockMetadata, TransformerBlockRegistry +from .hooks import BaseState, HookRegistry, ModelHook, StateManager + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +_SEA_CACHE_ROOT_HOOK = "sea_cache_root" +_SEA_CACHE_LEADER_BLOCK_HOOK = "sea_cache_leader_block" +_SEA_CACHE_BLOCK_HOOK = "sea_cache_block" +_SEA_CACHE_POST_NORM_HOOK = "sea_cache_post_norm" + + +@dataclass +class SeaCacheConfig: + r""" + Configuration for [SeaCache](https://huggingface.co/papers/2602.18993). + + SeaCache is disabled by default and only activates after this configuration is passed to + `transformer.enable_cache(config)`. + + SeaCache compares Spectral-Evolution-Aware (SEA) indicators between scheduler steps. If their accumulated relative + change stays below `threshold`, the expensive language-model hidden transform is replaced with a cached residual. + For Cosmos 3, the residual spans the decoder stack and final pathway normalization; input packing and modality + prediction heads still execute. + + Args: + threshold (`float`, defaults to `0.25`): + Accumulated relative-L1 budget. Larger values reuse the cache more often. + residual_order (`int`, defaults to `1`): + Order used to predict the generation-stream language-model residual. `0` directly reuses the most recent + residual and `1` linearly extrapolates from the two most recent full executions. + retention_steps (`int`, defaults to `1`): + Number of initial scheduler steps that always execute in full. + cache_end_steps (`int`, defaults to `1`): + Number of final scheduler steps that always execute in full. + max_consecutive_cached (`int`, defaults to `2`): + Maximum consecutive residual reuses per cache context before forcing a full execution. `0` disables the + limit. + power_exp (`float`, defaults to `3.0`): + Exponent of the SEA clean-signal power prior. SeaCache uses `3.0` for video features. + indicator_source (`str`, defaults to `"raw_vision_latents"`): + Feature source used to construct the SEA indicator. `"raw_vision_latents"` filters the complete raw vision + latent, including clean conditioning frames in I2V. `"first_block"` filters the timestep-modulated + pre-attention input of the first transformer block. Thresholds are not generally transferable between the + two sources. + current_step_callback (`Callable[[], int]`): + Callback returning the current scheduler step index. + current_sigma_callback (`Callable[[], float]`): + Callback returning the exact current scheduler sigma in `[0, 1]`. + num_inference_steps_callback (`Callable[[], int]`): + Callback returning the number of scheduler steps in the current pipeline call. + metadata_callback (`Callable`, *optional*): + Advanced model adapter returning a list of `(indices, (T, H, W))` entries that locate projected noisy + vision tokens in the generation stream for the `"first_block"` indicator. Cosmos 3 uses its native adapter + when this is omitted. + raw_vision_callback (`Callable`, *optional*): + Advanced model adapter returning raw vision latents with shape `(C, T, H, W)` for the + `"raw_vision_latents"` indicator. Cosmos 3 uses its native adapter when this is omitted. + + Example: + ```python + >>> from diffusers import Cosmos3OmniPipeline, SeaCacheConfig + + >>> pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano") + >>> pipe.transformer.enable_cache( + ... SeaCacheConfig( + ... current_step_callback=lambda: pipe.current_step_index, + ... current_sigma_callback=lambda: pipe.current_sigma, + ... num_inference_steps_callback=lambda: pipe.num_timesteps, + ... ) + ... ) + ``` + """ + + threshold: float = 0.25 + residual_order: int = 1 + retention_steps: int = 1 + cache_end_steps: int = 1 + max_consecutive_cached: int = 2 + power_exp: float = 3.0 + indicator_source: Literal["first_block", "raw_vision_latents"] = "raw_vision_latents" + current_step_callback: Callable[[], int] = None + current_sigma_callback: Callable[[], float] = None + num_inference_steps_callback: Callable[[], int] = None + metadata_callback: Callable[ + [torch.nn.Module, tuple[Any, ...], dict[str, Any]], + list[tuple[torch.Tensor, tuple[int, int, int]]] | None, + ] = None + raw_vision_callback: Callable[ + [torch.nn.Module, tuple[Any, ...], dict[str, Any]], + list[torch.Tensor] | None, + ] = None + + def __post_init__(self): + if not math.isfinite(self.threshold) or self.threshold < 0: + raise ValueError(f"`threshold` must be non-negative, got {self.threshold}.") + if self.residual_order not in (0, 1): + raise ValueError(f"`residual_order` must be 0 or 1, got {self.residual_order}.") + if self.retention_steps < 0: + raise ValueError(f"`retention_steps` must be non-negative, got {self.retention_steps}.") + if self.cache_end_steps < 0: + raise ValueError(f"`cache_end_steps` must be non-negative, got {self.cache_end_steps}.") + if ( + isinstance(self.max_consecutive_cached, bool) + or not isinstance(self.max_consecutive_cached, int) + or self.max_consecutive_cached < 0 + ): + raise ValueError( + f"`max_consecutive_cached` must be a non-negative integer, got {self.max_consecutive_cached!r}." + ) + if not math.isfinite(self.power_exp) or self.power_exp <= 0: + raise ValueError(f"`power_exp` must be positive, got {self.power_exp}.") + if self.indicator_source not in ("first_block", "raw_vision_latents"): + raise ValueError( + f"`indicator_source` must be 'first_block' or 'raw_vision_latents', got {self.indicator_source!r}." + ) + for name in ( + "current_step_callback", + "current_sigma_callback", + "num_inference_steps_callback", + "metadata_callback", + "raw_vision_callback", + ): + callback = getattr(self, name) + if callback is not None and not callable(callback): + raise TypeError(f"`{name}` must be callable or `None`.") + + +@dataclass +class _SeaCacheForwardMetadata: + step_index: int + sigma: float + num_inference_steps: int + vision_layout: list[tuple[torch.Tensor, tuple[int, int, int]]] | None = None + raw_vision: list[torch.Tensor] | None = None + + +class SeaCacheContextState(BaseState): + def __init__(self): + self.history: list[tuple[int, torch.Tensor, torch.Tensor]] = [] + self.gate_key: tuple[int, float] | None = None + self.gate_should_compute = True + self.previous_indicator: list[torch.Tensor] | None = None + self.accumulated_distance = 0.0 + self.consecutive_cached = 0 + self.skip_remaining = False + self.full_execution_pending = False + self.cacheable_execution = False + self.step_index: int | None = None + self.gen_input: torch.Tensor | None = None + self.und_output: torch.Tensor | None = None + self.cached_und_output: torch.Tensor | None = None + self.cached_gen_residual: torch.Tensor | None = None + + def reset_forward(self): + self.skip_remaining = False + self.full_execution_pending = False + self.cacheable_execution = False + self.step_index = None + self.gen_input = None + self.und_output = None + self.cached_und_output = None + self.cached_gen_residual = None + + def reset_trajectory(self): + self.history = [] + self.gate_key = None + self.gate_should_compute = True + self.previous_indicator = None + self.accumulated_distance = 0.0 + self.consecutive_cached = 0 + + def reset(self): + self.reset_trajectory() + self.reset_forward() + + +class SeaCacheSharedState: + def __init__(self): + self._warned_messages: set[str] = set() + self.reset() + + def reset(self): + self.forward_metadata: _SeaCacheForwardMetadata | None = None + + def warn_once(self, message: str): + if message not in self._warned_messages: + logger.warning(message) + self._warned_messages.add(message) + + def mark_fail_open(self, message: str): + self.warn_once(message) + + def resolve_gate( + self, + state: SeaCacheContextState, + metadata: _SeaCacheForwardMetadata, + indicator: list[torch.Tensor] | None, + config: SeaCacheConfig, + ) -> bool: + gate_key = (metadata.step_index, metadata.sigma) + if state.gate_key == gate_key: + return state.gate_should_compute + + is_non_adjacent = state.gate_key is not None and metadata.step_index != state.gate_key[0] + 1 + if is_non_adjacent: + self.mark_fail_open("SeaCache received non-adjacent scheduler steps; running full.") + state.reset_trajectory() + is_retained = metadata.step_index < config.retention_steps + is_in_cache_end = metadata.step_index >= metadata.num_inference_steps - config.cache_end_steps + is_first_observation = state.previous_indicator is None + is_max_consecutive = bool( + config.max_consecutive_cached and state.consecutive_cached >= config.max_consecutive_cached + ) + invalid_gate = is_non_adjacent or indicator is None + forced_compute = invalid_gate or is_retained or is_in_cache_end or is_first_observation or is_max_consecutive + candidate_accumulated_distance = 0.0 + + if forced_compute: + natural_should_compute = True + else: + if len(indicator) != len(state.previous_indicator) or not indicator: + distance = float("inf") + invalid_gate = True + else: + distance = 0.0 + for current, previous in zip(indicator, state.previous_indicator): + if ( + current.shape != previous.shape + or current.device != previous.device + or current.dtype != previous.dtype + ): + distance = float("inf") + invalid_gate = True + break + numerator = (current.float() - previous.float()).abs().mean() + denominator = previous.float().abs().mean() + 1e-16 + distance += float((numerator / denominator).detach().cpu()) + distance /= len(indicator) + + if not math.isfinite(distance): + invalid_gate = True + self.mark_fail_open("SeaCache indicator history changed shape, device, or dtype; running full.") + candidate_accumulated_distance = state.accumulated_distance + distance + natural_should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold + + should_compute = natural_should_compute + + if is_max_consecutive: + should_compute = True + + state.accumulated_distance = 0.0 if should_compute else candidate_accumulated_distance + + state.gate_key = gate_key + state.gate_should_compute = should_compute + state.previous_indicator = None if indicator is None else [value.detach() for value in indicator] + return should_compute + + +def _get_block_inputs( + metadata: TransformerBlockMetadata, args: tuple[Any, ...], kwargs: dict[str, Any] +) -> tuple[torch.Tensor, torch.Tensor | None]: + hidden_states = metadata._get_parameter_from_args_kwargs(metadata.hidden_states_argument_name, args, kwargs) + encoder_hidden_states = None + if metadata.return_encoder_hidden_states_index is not None: + encoder_hidden_states = metadata._get_parameter_from_args_kwargs( + metadata.encoder_hidden_states_argument_name, args, kwargs + ) + return hidden_states, encoder_hidden_states + + +def _build_block_output( + metadata: TransformerBlockMetadata, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, +) -> torch.Tensor | tuple[torch.Tensor, ...]: + if metadata.return_encoder_hidden_states_index is None: + return hidden_states + + output = [None] * (max(metadata.return_hidden_states_index, metadata.return_encoder_hidden_states_index) + 1) + output[metadata.return_hidden_states_index] = hidden_states + output[metadata.return_encoder_hidden_states_index] = encoder_hidden_states + return tuple(output) + + +def _get_block_outputs( + metadata: TransformerBlockMetadata, output: torch.Tensor | tuple[torch.Tensor, ...] +) -> tuple[torch.Tensor, torch.Tensor | None]: + if isinstance(output, tuple): + hidden_states = output[metadata.return_hidden_states_index] + encoder_hidden_states = ( + output[metadata.return_encoder_hidden_states_index] + if metadata.return_encoder_hidden_states_index is not None + else None + ) + return hidden_states, encoder_hidden_states + return output, None + + +def _record_full_execution( + config: SeaCacheConfig, + state: SeaCacheContextState, + gen_output: torch.Tensor, + und_output: torch.Tensor | None, +) -> None: + state.consecutive_cached = 0 + if ( + state.cacheable_execution + and state.step_index is not None + and state.gen_input is not None + and und_output is not None + and gen_output.shape == state.gen_input.shape + ): + state.history.append( + ( + state.step_index, + und_output.detach().clone(), + (gen_output - state.gen_input).detach().clone(), + ) + ) + state.history = state.history[-(config.residual_order + 1) :] + state.reset_forward() + + +def _prepare_cosmos3_vision_metadata( + module: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any] +) -> list[tuple[torch.Tensor, tuple[int, int, int]]] | None: + module = unwrap_module(module) + bound_arguments = inspect.signature(module.__class__.forward).bind_partial(module, *args, **kwargs).arguments + vision_tokens = bound_arguments.get("vision_tokens") + vision_token_shapes = bound_arguments.get("vision_token_shapes") + vision_sequence_indexes = bound_arguments.get("vision_sequence_indexes") + vision_timesteps = bound_arguments.get("vision_timesteps") + vision_noisy_frame_indexes = bound_arguments.get("vision_noisy_frame_indexes") + und_len = bound_arguments.get("und_len") + + if ( + not isinstance(vision_tokens, (list, tuple)) + or not isinstance(vision_token_shapes, (list, tuple)) + or not isinstance(vision_sequence_indexes, torch.Tensor) + or not isinstance(vision_timesteps, torch.Tensor) + or vision_timesteps.numel() == 0 + or not isinstance(vision_noisy_frame_indexes, (list, tuple)) + or und_len is None + or len(vision_tokens) != len(vision_token_shapes) + or len(vision_tokens) != len(vision_noisy_frame_indexes) + ): + return None + + vision_sequence_indexes = vision_sequence_indexes.flatten() + layout = [] + offset = 0 + for token, token_shape, noisy_frame_indexes in zip(vision_tokens, vision_token_shapes, vision_noisy_frame_indexes): + if ( + not isinstance(token, torch.Tensor) + or not isinstance(noisy_frame_indexes, torch.Tensor) + or len(token_shape) != 3 + ): + return None + + temporal, height, width = (int(value) for value in token_shape) + item_numel = temporal * height * width + item_indexes = vision_sequence_indexes[offset : offset + item_numel] + if item_indexes.numel() != item_numel: + return None + offset += item_numel + + noisy_frame_indexes = noisy_frame_indexes.flatten().to(device=item_indexes.device, dtype=torch.long) + if noisy_frame_indexes.numel() == 0: + continue + if torch.any(noisy_frame_indexes < 0) or torch.any(noisy_frame_indexes >= temporal): + return None + + item_indexes = item_indexes.reshape(temporal, height, width) + generation_indexes = item_indexes[noisy_frame_indexes].flatten() - int(und_len) + if torch.any(generation_indexes < 0): + return None + layout.append( + ( + generation_indexes, + (int(noisy_frame_indexes.numel()), height, width), + ) + ) + + if offset != vision_sequence_indexes.numel() or not layout: + return None + return layout + + +def _prepare_cosmos3_raw_vision_metadata( + module: torch.nn.Module, args: tuple[Any, ...], kwargs: dict[str, Any] +) -> list[torch.Tensor] | None: + module = unwrap_module(module) + bound_arguments = inspect.signature(module.__class__.forward).bind_partial(module, *args, **kwargs).arguments + vision_tokens = bound_arguments.get("vision_tokens") + vision_noisy_frame_indexes = bound_arguments.get("vision_noisy_frame_indexes") + + if ( + not isinstance(vision_tokens, (list, tuple)) + or not isinstance(vision_noisy_frame_indexes, (list, tuple)) + or len(vision_tokens) != len(vision_noisy_frame_indexes) + ): + return None + + raw_vision = [] + has_noisy_vision = False + for latent, noisy_frame_indexes in zip(vision_tokens, vision_noisy_frame_indexes): + if not isinstance(latent, torch.Tensor) or not isinstance(noisy_frame_indexes, torch.Tensor): + return None + if latent.ndim == 5: + if latent.shape[0] != 1: + return None + latent = latent.squeeze(0) + if latent.ndim != 4: + return None + + noisy_frame_indexes = noisy_frame_indexes.flatten().to(device=latent.device, dtype=torch.long) + if torch.any(noisy_frame_indexes < 0) or torch.any(noisy_frame_indexes >= latent.shape[1]): + return None + has_noisy_vision = has_noisy_vision or noisy_frame_indexes.numel() > 0 + raw_vision.append(latent) + + return raw_vision if raw_vision and has_noisy_vision else None + + +def _apply_sea_filter( + hidden_states: torch.Tensor, + sigma: float, + power_exp: float, +) -> torch.Tensor: + hidden_states_dtype = hidden_states.dtype + hidden_states = hidden_states.contiguous().float() + dimensions = (0, 1, 2) + spectrum = torch.fft.fftn(hidden_states, dim=dimensions) + + sigma = max(1e-6, min(1.0 - 1e-6, sigma)) + signal_scale = 1.0 - sigma + noise_scale = sigma + gain = None + for axis in dimensions: + frequencies = torch.fft.fftfreq(hidden_states.shape[axis], device=hidden_states.device, dtype=torch.float32) + clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16) + axis_gain = signal_scale * clean_power / (signal_scale**2 * clean_power + noise_scale**2 + 1e-16) + axis_shape = [1] * hidden_states.ndim + axis_shape[axis] = axis_gain.shape[0] + gain = axis_gain.reshape(axis_shape) if gain is None else gain * axis_gain.reshape(axis_shape) + + # SeaCache Eq. (7): density-normalize the combined spatiotemporal response + # to unit mean so cache distances are comparable across scheduler steps. + mean_gain = gain.mean() + if torch.isfinite(mean_gain) and mean_gain > 0: + gain = gain / mean_gain + return torch.fft.ifftn(spectrum * gain, dim=dimensions).real.to(hidden_states_dtype) + + +def _is_parameter_sharded(module: torch.nn.Module) -> bool: + """Whether a block is managed by a parameter-sharding runtime that SeaCache cannot safely bypass.""" + + for submodule in unwrap_module(module).modules(): + module_type = type(submodule) + if callable(getattr(submodule, "_get_fsdp_state", None)): + return True + if module_type.__name__ == "FullyShardedDataParallel" and module_type.__module__.startswith( + "torch.distributed.fsdp" + ): + return True + for parameter in submodule.parameters(recurse=False): + parameter_type = type(parameter) + if ( + parameter_type.__name__ == "FlatParameter" + and parameter_type.__module__.startswith("torch.distributed.fsdp") + ) or ( + parameter_type.__name__ == "DTensor" + and parameter_type.__module__.startswith("torch.distributed.tensor") + ): + return True + return False + + +class SeaCacheRootHook(ModelHook): + _is_stateful = True + + def __init__( + self, + config: SeaCacheConfig, + state_manager: StateManager, + shared_state: SeaCacheSharedState, + metadata_callback: Callable, + raw_vision_callback: Callable, + ): + super().__init__() + self.config = config + self.state_manager = state_manager + self.shared_state = shared_state + self.metadata_callback = metadata_callback + self.raw_vision_callback = raw_vision_callback + + def pre_forward(self, module: torch.nn.Module, *args, **kwargs): + self.shared_state.forward_metadata = None + if torch.is_grad_enabled(): + self.shared_state.mark_fail_open( + "SeaCache is inference-only; calls with autograd enabled run in fail-open mode." + ) + return args, kwargs + if self.state_manager._current_context is None: + self.shared_state.mark_fail_open( + "SeaCache requires a cache context for each transformer call; running in fail-open mode." + ) + return args, kwargs + callbacks = ( + self.config.current_step_callback, + self.config.current_sigma_callback, + self.config.num_inference_steps_callback, + ) + if any(callback is None for callback in callbacks): + self.shared_state.mark_fail_open( + "SeaCache is running in fail-open mode because scheduler step, sigma, and step-count callbacks are " + "required." + ) + return args, kwargs + + try: + step_index = self.config.current_step_callback() + sigma = self.config.current_sigma_callback() + num_inference_steps = self.config.num_inference_steps_callback() + if isinstance(step_index, torch.Tensor): + step_index = step_index.item() + if isinstance(sigma, torch.Tensor): + sigma = sigma.item() + if isinstance(num_inference_steps, torch.Tensor): + num_inference_steps = num_inference_steps.item() + step_index = int(step_index) + sigma = float(sigma) + num_inference_steps = int(num_inference_steps) + except (IndexError, TypeError, ValueError, RuntimeError) as error: + self.shared_state.mark_fail_open( + f"SeaCache scheduler metadata is unavailable; running in fail-open mode: {error}" + ) + return args, kwargs + + if ( + step_index < 0 + or num_inference_steps <= 0 + or step_index >= num_inference_steps + or not math.isfinite(sigma) + or not 0.0 <= sigma <= 1.0 + ): + self.shared_state.mark_fail_open( + "SeaCache scheduler metadata is invalid; expected a valid step index and exact sigma in [0, 1]." + ) + return args, kwargs + + vision_layout = None + raw_vision = None + try: + if self.config.indicator_source == "first_block": + vision_layout = ( + self.metadata_callback(module, args, kwargs) if self.metadata_callback is not None else None + ) + else: + raw_vision = ( + self.raw_vision_callback(module, args, kwargs) if self.raw_vision_callback is not None else None + ) + except (TypeError, ValueError, RuntimeError) as error: + self.shared_state.mark_fail_open( + f"SeaCache model metadata is unavailable; running in fail-open mode: {error}" + ) + return args, kwargs + if self.config.indicator_source == "first_block" and not vision_layout: + self.shared_state.mark_fail_open( + "SeaCache requires noisy vision tokens; action-only, sound-only, and conditioning-only calls run in " + "fail-open mode." + ) + return args, kwargs + if self.config.indicator_source == "raw_vision_latents" and not raw_vision: + self.shared_state.mark_fail_open( + "SeaCache requires raw noisy vision latents for the selected indicator source; action-only, sound-only, " + "conditioning-only, and unsupported model calls run in fail-open mode." + ) + return args, kwargs + + self.shared_state.forward_metadata = _SeaCacheForwardMetadata( + step_index=step_index, + sigma=sigma, + num_inference_steps=num_inference_steps, + vision_layout=vision_layout, + raw_vision=raw_vision, + ) + return args, kwargs + + def post_forward(self, module: torch.nn.Module, output: Any) -> Any: + self.shared_state.forward_metadata = None + if self.state_manager._current_context is not None: + self.state_manager.get_state().reset_forward() + return output + + def reset_state(self, module: torch.nn.Module): + self.state_manager.reset() + self.shared_state.reset() + return module + + +class SeaCacheLeaderBlockHook(ModelHook): + def __init__( + self, + config: SeaCacheConfig, + state_manager: StateManager, + shared_state: SeaCacheSharedState, + post_norm_boundary: bool = False, + ): + super().__init__() + self.config = config + self.state_manager = state_manager + self.shared_state = shared_state + self.post_norm_boundary = post_norm_boundary + self._metadata = None + self._normalization = None + + def initialize_hook(self, module: torch.nn.Module): + module = unwrap_module(module) + self._metadata = TransformerBlockRegistry.get(module.__class__) + if self._metadata.hidden_states_norm_module_name is not None: + self._normalization = getattr(module, self._metadata.hidden_states_norm_module_name) + return module + + def _build_indicator( + self, + hidden_states: torch.Tensor, + forward_metadata: _SeaCacheForwardMetadata, + ) -> list[torch.Tensor] | None: + if self.config.indicator_source == "raw_vision_latents": + if not forward_metadata.raw_vision: + return None + indicator = [] + for latent in forward_metadata.raw_vision: + raw_vision = latent.movedim(0, -1) + indicator.append( + _apply_sea_filter( + raw_vision, + sigma=forward_metadata.sigma, + power_exp=self.config.power_exp, + ).detach() + ) + return indicator + + if self._normalization is None: + return None + if not forward_metadata.vision_layout: + return None + normalized_hidden_states = self._normalization(hidden_states) + indicator = [] + for indexes, shape in forward_metadata.vision_layout: + indexes = indexes.to(device=normalized_hidden_states.device, dtype=torch.long) + if ( + indexes.numel() != math.prod(shape) + or torch.any(indexes < 0) + or torch.any(indexes >= normalized_hidden_states.shape[0]) + ): + return None + noisy_vision = normalized_hidden_states.index_select(0, indexes).reshape(*shape, -1) + indicator.append( + _apply_sea_filter( + noisy_vision, + sigma=forward_metadata.sigma, + power_exp=self.config.power_exp, + ).detach() + ) + return indicator + + @torch.compiler.disable + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + hidden_states, encoder_hidden_states = _get_block_inputs(self._metadata, args, kwargs) + context_is_set = self.state_manager._current_context is not None + state = self.state_manager.get_state() if context_is_set else None + if state is not None: + state.reset_forward() + state.full_execution_pending = True + state.gen_input = hidden_states + + forward_metadata = self.shared_state.forward_metadata + if state is None or forward_metadata is None: + return self.fn_ref.original_forward(*args, **kwargs) + + state.step_index = forward_metadata.step_index + state.cacheable_execution = True + indicator_error_reported = False + if _is_parameter_sharded(module): + self.shared_state.mark_fail_open( + "SeaCache cannot safely bypass parameter-sharded transformer blocks; running in fail-open mode." + ) + indicator = None + indicator_error_reported = True + else: + try: + indicator = self._build_indicator(hidden_states, forward_metadata) + except (TypeError, ValueError, RuntimeError) as error: + self.shared_state.mark_fail_open( + f"SeaCache could not construct its vision indicator; running in fail-open mode: {error}" + ) + indicator = None + indicator_error_reported = True + if indicator is None and not indicator_error_reported: + self.shared_state.mark_fail_open( + "SeaCache could not construct its vision indicator; running in fail-open mode." + ) + should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config) + + if should_compute or not state.history: + if not should_compute: + state.accumulated_distance = 0.0 + self.shared_state.mark_fail_open( + "SeaCache selected a cache hit without residual history; running in fail-open mode." + ) + return self.fn_ref.original_forward(*args, **kwargs) + + residual_history = state.history[-(self.config.residual_order + 1) :] + _, cached_und, cached_residual = residual_history[-1] + if ( + any( + residual.shape != hidden_states.shape + or residual.device != hidden_states.device + or residual.dtype != hidden_states.dtype + for _, _, residual in residual_history + ) + or encoder_hidden_states is None + or cached_und.shape != encoder_hidden_states.shape + or cached_und.device != encoder_hidden_states.device + or cached_und.dtype != encoder_hidden_states.dtype + ): + state.history = [] + state.accumulated_distance = 0.0 + self.shared_state.mark_fail_open( + "SeaCache residual history changed shape, device, or dtype; running in fail-open mode." + ) + return self.fn_ref.original_forward(*args, **kwargs) + + if self.config.residual_order == 1 and len(residual_history) >= 2: + previous_step, _, previous_residual = residual_history[-2] + latest_step, _, latest_residual = residual_history[-1] + if latest_step != previous_step: + step_scale = (forward_metadata.step_index - latest_step) / (latest_step - previous_step) + cached_residual = latest_residual + (latest_residual - previous_residual) * step_scale + + state.skip_remaining = True + state.full_execution_pending = False + state.cached_und_output = cached_und + state.cached_gen_residual = cached_residual + state.consecutive_cached += 1 + if self.post_norm_boundary: + return _build_block_output(self._metadata, hidden_states, encoder_hidden_states) + return _build_block_output(self._metadata, hidden_states + cached_residual, cached_und) + + +class SeaCacheBlockHook(ModelHook): + def __init__( + self, + config: SeaCacheConfig, + state_manager: StateManager, + shared_state: SeaCacheSharedState, + is_tail: bool = False, + post_norm_boundary: bool = False, + ): + super().__init__() + self.config = config + self.state_manager = state_manager + self.shared_state = shared_state + self.is_tail = is_tail + self.post_norm_boundary = post_norm_boundary + self._metadata = None + + def initialize_hook(self, module: torch.nn.Module): + self._metadata = TransformerBlockRegistry.get(unwrap_module(module).__class__) + return module + + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + if self.state_manager._current_context is None: + return self.fn_ref.original_forward(*args, **kwargs) + + state: SeaCacheContextState = self.state_manager.get_state() + if state.skip_remaining: + hidden_states, encoder_hidden_states = _get_block_inputs(self._metadata, args, kwargs) + return _build_block_output(self._metadata, hidden_states, encoder_hidden_states) + + output = self.fn_ref.original_forward(*args, **kwargs) + if not self.is_tail or state.skip_remaining or not state.full_execution_pending: + return output + if self.post_norm_boundary: + return output + + hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output) + _record_full_execution( + self.config, + state, + gen_output=hidden_states, + und_output=encoder_hidden_states, + ) + return output + + +class SeaCachePostNormHook(ModelHook): + def __init__( + self, + config: SeaCacheConfig, + state_manager: StateManager, + shared_state: SeaCacheSharedState, + pathway: Literal["und", "gen"], + ): + super().__init__() + self.config = config + self.state_manager = state_manager + self.shared_state = shared_state + self.pathway = pathway + + @torch.compiler.disable + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + if self.state_manager._current_context is None: + return self.fn_ref.original_forward(*args, **kwargs) + + state: SeaCacheContextState = self.state_manager.get_state() + if state.skip_remaining: + if self.pathway == "und": + if state.cached_und_output is not None: + return state.cached_und_output + elif state.gen_input is not None and state.cached_gen_residual is not None: + output = state.gen_input + state.cached_gen_residual + state.reset_forward() + return output + + self.shared_state.mark_fail_open( + "SeaCache post-normalization state is incomplete after the decoder stack was skipped." + ) + output = self.fn_ref.original_forward(*args, **kwargs) + state.reset_forward() + return output + + output = self.fn_ref.original_forward(*args, **kwargs) + if not state.full_execution_pending: + return output + if self.pathway == "und": + state.und_output = output + return output + + _record_full_execution( + self.config, + state, + gen_output=output, + und_output=state.und_output, + ) + return output + + +def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None: + r""" + Apply SeaCache to a supported transformer. + + The hook caches the transformer's expensive language-model hidden transform. For Cosmos 3, the cache stores + post-normalization understanding output and a generation residual from the pre-block input to the + post-normalization output. Modality prediction heads continue to run normally. Other model adapters fall back to + caching the complete repeated-block stack. + + Args: + module (`torch.nn.Module`): + Transformer module to cache. + config (`SeaCacheConfig`): + SeaCache configuration. + """ + from ..models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer + + unwrapped_module = unwrap_module(module) + is_cosmos3 = isinstance(unwrapped_module, Cosmos3OmniTransformer) + metadata_callback = config.metadata_callback + raw_vision_callback = config.raw_vision_callback + if metadata_callback is None or raw_vision_callback is None: + if is_cosmos3: + if metadata_callback is None: + metadata_callback = _prepare_cosmos3_vision_metadata + if raw_vision_callback is None: + raw_vision_callback = _prepare_cosmos3_raw_vision_metadata + + post_norm_modules = None + if is_cosmos3: + und_norm = getattr(unwrapped_module, "norm", None) + gen_norm = getattr(unwrapped_module, "norm_moe_gen", None) + if isinstance(und_norm, torch.nn.Module) and isinstance(gen_norm, torch.nn.Module): + post_norm_modules = (("und", "norm", und_norm), ("gen", "norm_moe_gen", gen_norm)) + else: + logger.warning( + "SeaCache could not locate the Cosmos 3 final pathway normalizations; falling back to the repeated-block " + "residual boundary." + ) + post_norm_boundary = post_norm_modules is not None + + blocks = [] + for name, submodule in unwrapped_module.named_children(): + if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList): + continue + blocks.extend((f"{name}.{index}", block) for index, block in enumerate(submodule)) + + if not blocks: + raise ValueError("SeaCache found no repeated transformer blocks on the model.") + + state_manager = StateManager(SeaCacheContextState) + shared_state = SeaCacheSharedState() + root_registry = HookRegistry.check_if_exists_or_initialize(module) + registrations: list[tuple[HookRegistry, str]] = [] + + def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: + registry = HookRegistry.check_if_exists_or_initialize(target) + registry.register_hook(hook, name) + registrations.append((registry, name)) + + try: + register_hook( + module, + SeaCacheRootHook( + config, + state_manager, + shared_state, + metadata_callback, + raw_vision_callback, + ), + _SEA_CACHE_ROOT_HOOK, + ) + + leader_name, leader = blocks[0] + logger.debug(f"Applying SeaCache leader hook to '{leader_name}'.") + register_hook( + leader, + SeaCacheLeaderBlockHook(config, state_manager, shared_state, post_norm_boundary=post_norm_boundary), + _SEA_CACHE_LEADER_BLOCK_HOOK, + ) + + for name, block in blocks[1:-1]: + logger.debug(f"Applying SeaCache identity hook to '{name}'.") + register_hook( + block, + SeaCacheBlockHook(config, state_manager, shared_state, post_norm_boundary=post_norm_boundary), + _SEA_CACHE_BLOCK_HOOK, + ) + + tail_name, tail = blocks[-1] + logger.debug(f"Applying SeaCache tail hook to '{tail_name}'.") + register_hook( + tail, + SeaCacheBlockHook( + config, + state_manager, + shared_state, + is_tail=True, + post_norm_boundary=post_norm_boundary, + ), + _SEA_CACHE_BLOCK_HOOK, + ) + if post_norm_modules is not None: + for pathway, name, norm_module in post_norm_modules: + logger.debug(f"Applying SeaCache post-normalization hook to '{name}'.") + register_hook( + norm_module, + SeaCachePostNormHook(config, state_manager, shared_state, pathway=pathway), + _SEA_CACHE_POST_NORM_HOOK, + ) + except Exception: + for registry, name in reversed(registrations): + registry.remove_hook(name, recurse=False) + root_registry._child_registries_cache = None + raise + + root_registry._child_registries_cache = None diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index 5aa189987ba2..2d7e0309db3b 100644 --- a/src/diffusers/models/cache_utils.py +++ b/src/diffusers/models/cache_utils.py @@ -28,6 +28,7 @@ class CacheMixin: - [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) - [FasterCache](https://huggingface.co/papers/2410.19355) - [FirstBlockCache](https://github.com/chengzeyi/ParaAttention/blob/7a266123671b55e7e5a2fe9af3121f07a36afc78/README.md#first-block-cache-our-dynamic-caching) + - [SeaCache](https://huggingface.co/papers/2602.18993) """ _cache_config = None @@ -41,11 +42,12 @@ def enable_cache(self, config) -> None: Enable caching techniques on the model. Args: - config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | TextKVCacheConfig`): + config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | SeaCacheConfig | TextKVCacheConfig`): The configuration for applying the caching technique. Currently supported caching techniques are: - [`~hooks.PyramidAttentionBroadcastConfig`] - [`~hooks.FasterCacheConfig`] - [`~hooks.FirstBlockCacheConfig`] + - [`~hooks.SeaCacheConfig`] - [`~hooks.TextKVCacheConfig`] Example: @@ -69,14 +71,17 @@ def enable_cache(self, config) -> None: from ..hooks import ( FasterCacheConfig, FirstBlockCacheConfig, + HookRegistry, MagCacheConfig, PyramidAttentionBroadcastConfig, + SeaCacheConfig, TaylorSeerCacheConfig, TextKVCacheConfig, apply_faster_cache, apply_first_block_cache, apply_mag_cache, apply_pyramid_attention_broadcast, + apply_sea_cache, apply_taylorseer_cache, apply_text_kv_cache, ) @@ -96,12 +101,15 @@ def enable_cache(self, config) -> None: apply_text_kv_cache(self, config) elif isinstance(config, PyramidAttentionBroadcastConfig): apply_pyramid_attention_broadcast(self, config) + elif isinstance(config, SeaCacheConfig): + apply_sea_cache(self, config) elif isinstance(config, TaylorSeerCacheConfig): apply_taylorseer_cache(self, config) else: raise ValueError(f"Cache config {type(config)} is not supported.") self._cache_config = config + HookRegistry.check_if_exists_or_initialize(self)._child_registries_cache = None def disable_cache(self) -> None: from ..hooks import ( @@ -110,6 +118,7 @@ def disable_cache(self) -> None: HookRegistry, MagCacheConfig, PyramidAttentionBroadcastConfig, + SeaCacheConfig, TaylorSeerCacheConfig, TextKVCacheConfig, ) @@ -117,6 +126,12 @@ def disable_cache(self) -> None: from ..hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK from ..hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK from ..hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK + from ..hooks.sea_cache import ( + _SEA_CACHE_BLOCK_HOOK, + _SEA_CACHE_LEADER_BLOCK_HOOK, + _SEA_CACHE_POST_NORM_HOOK, + _SEA_CACHE_ROOT_HOOK, + ) from ..hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK from ..hooks.text_kv_cache import _TEXT_KV_CACHE_BLOCK_HOOK, _TEXT_KV_CACHE_TRANSFORMER_HOOK @@ -139,12 +154,18 @@ def disable_cache(self) -> None: elif isinstance(self._cache_config, TextKVCacheConfig): registry.remove_hook(_TEXT_KV_CACHE_TRANSFORMER_HOOK, recurse=True) registry.remove_hook(_TEXT_KV_CACHE_BLOCK_HOOK, recurse=True) + elif isinstance(self._cache_config, SeaCacheConfig): + registry.remove_hook(_SEA_CACHE_POST_NORM_HOOK, recurse=True) + registry.remove_hook(_SEA_CACHE_BLOCK_HOOK, recurse=True) + registry.remove_hook(_SEA_CACHE_LEADER_BLOCK_HOOK, recurse=True) + registry.remove_hook(_SEA_CACHE_ROOT_HOOK, recurse=True) elif isinstance(self._cache_config, TaylorSeerCacheConfig): registry.remove_hook(_TAYLORSEER_CACHE_HOOK, recurse=True) else: raise ValueError(f"Cache config {type(self._cache_config)} is not supported.") self._cache_config = None + registry._child_registries_cache = None def _reset_stateful_cache(self, recurse: bool = True) -> None: from ..hooks import HookRegistry @@ -159,6 +180,7 @@ def cache_context(self, name: str): registry = HookRegistry.check_if_exists_or_initialize(self) registry._set_context(name) - yield - - registry._set_context(None) + try: + yield + finally: + registry._set_context(None) diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py index f7cfc317bc79..07c6bd12d0dc 100644 --- a/src/diffusers/models/transformers/transformer_cosmos3.py +++ b/src/diffusers/models/transformers/transformer_cosmos3.py @@ -23,6 +23,7 @@ from ...utils import BaseOutput from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn +from ..cache_utils import CacheMixin from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_utils import ModelMixin from ..normalization import RMSNorm @@ -370,7 +371,7 @@ def forward( return residual_und + mlp_out_und, residual_gen + mlp_out_gen -class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin): +class Cosmos3OmniTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin, CacheMixin): _supports_gradient_checkpointing = True _no_split_modules = ["Cosmos3VLTextMoTDecoderLayer"] _repeated_blocks = ["Cosmos3VLTextMoTDecoderLayer"] diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index 7bf431aa855b..7e9d83fa6316 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -127,7 +127,7 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - dtype = components.transformer.dtype + sampling_dtype = torch.float32 x0_tokens_vision = block_state.x0_tokens_vision if x0_tokens_vision is None: @@ -151,21 +151,26 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) block_state.fps_vision = float(block_state.fps) condition_frames = block_state.vision_condition_frames or [] - block_state.vision_condition_mask = torch.zeros((x0_tokens_vision.shape[2], 1, 1), device=device, dtype=dtype) + block_state.vision_condition_mask = torch.zeros( + (x0_tokens_vision.shape[2], 1, 1), device=device, dtype=sampling_dtype + ) for frame_idx in condition_frames: if 0 <= frame_idx < block_state.vision_condition_mask.shape[0]: block_state.vision_condition_mask[frame_idx, 0, 0] = 1.0 if block_state.latents is None: pure_noise = randn_tensor( - tuple(x0_tokens_vision.shape), generator=block_state.generator, device=device, dtype=dtype + tuple(x0_tokens_vision.shape), + generator=block_state.generator, + device=device, + dtype=sampling_dtype, ) block_state.latents = ( - block_state.vision_condition_mask * x0_tokens_vision.to(device=device, dtype=dtype) + block_state.vision_condition_mask * x0_tokens_vision.to(device=device, dtype=sampling_dtype) + (1.0 - block_state.vision_condition_mask) * pure_noise ) else: - block_state.latents = block_state.latents.to(device=device, dtype=dtype) + block_state.latents = block_state.latents.to(device=device, dtype=sampling_dtype) vision_condition_indexes = torch.nonzero( block_state.vision_condition_mask[:, 0, 0] > 0, as_tuple=False @@ -223,7 +228,7 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - dtype = components.transformer.dtype + sampling_dtype = torch.float32 if not components.transformer.config.sound_gen: raise ValueError("Sound generation requires a transformer trained with sound_gen=True.") @@ -233,19 +238,21 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) n_audio_samples = int(block_state.num_frames / block_state.fps * components.sound_sampling_rate) hop_size = components.sound_hop_size t_sound = (n_audio_samples + hop_size - 1) // hop_size - x0_tokens_sound = torch.zeros(sound_dim, t_sound, device=device, dtype=dtype) - block_state.sound_condition_mask = torch.zeros((x0_tokens_sound.shape[1], 1), device=device, dtype=dtype) + x0_tokens_sound = torch.zeros(sound_dim, t_sound, device=device, dtype=sampling_dtype) + block_state.sound_condition_mask = torch.zeros( + (x0_tokens_sound.shape[1], 1), device=device, dtype=sampling_dtype + ) if block_state.sound_latents is None: pure_noise = randn_tensor( - tuple(x0_tokens_sound.shape), generator=block_state.generator, device=device, dtype=dtype + tuple(x0_tokens_sound.shape), generator=block_state.generator, device=device, dtype=sampling_dtype ) block_state.sound_latents = ( block_state.sound_condition_mask.T * x0_tokens_sound + (1.0 - block_state.sound_condition_mask.T) * pure_noise ) else: - block_state.sound_latents = block_state.sound_latents.to(device=device, dtype=dtype) + block_state.sound_latents = block_state.sound_latents.to(device=device, dtype=sampling_dtype) block_state.sound_scheduler = copy.deepcopy(components.scheduler) @@ -320,7 +327,7 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - dtype = components.transformer.dtype + sampling_dtype = torch.float32 action = block_state.action if not components.transformer.config.action_gen: @@ -342,7 +349,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) raw_actions = action.raw_actions if raw_actions is None: raise ValueError("action_mode='forward_dynamics' requires an action tensor.") - raw_actions = raw_actions.to(device=device, dtype=dtype) + raw_actions = raw_actions.to(device=device, dtype=sampling_dtype) if raw_actions.shape[-1] > action_dim: raise ValueError( f"Cosmos3 action dimension {raw_actions.shape[-1]} exceeds model action_dim={action_dim}." @@ -363,7 +370,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) raw_actions = torch.cat([raw_actions, action_padding], dim=-1) x0_tokens_action = raw_actions else: - x0_tokens_action = torch.zeros(action_chunk_size, action_dim, device=device, dtype=dtype) + x0_tokens_action = torch.zeros(action_chunk_size, action_dim, device=device, dtype=sampling_dtype) if action.domain_name not in _EMBODIMENT_TO_DOMAIN_ID: raise ValueError( @@ -373,14 +380,16 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) torch.tensor([_EMBODIMENT_TO_DOMAIN_ID[action.domain_name]], dtype=torch.long, device=device) ] condition_frames = block_state.action_condition_frame_indexes or [] - block_state.action_condition_mask = torch.zeros((x0_tokens_action.shape[0], 1), device=device, dtype=dtype) + block_state.action_condition_mask = torch.zeros( + (x0_tokens_action.shape[0], 1), device=device, dtype=sampling_dtype + ) for frame_idx in condition_frames: if 0 <= frame_idx < block_state.action_condition_mask.shape[0]: block_state.action_condition_mask[frame_idx, 0] = 1.0 if block_state.action_latents is None: pure_noise = randn_tensor( - tuple(x0_tokens_action.shape), generator=block_state.generator, device=device, dtype=dtype + tuple(x0_tokens_action.shape), generator=block_state.generator, device=device, dtype=sampling_dtype ) block_state.action_latents = ( block_state.action_condition_mask * x0_tokens_action @@ -389,7 +398,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) if block_state.raw_action_dim_resolved is not None: block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0 else: - block_state.action_latents = block_state.action_latents.to(device=device, dtype=dtype) + block_state.action_latents = block_state.action_latents.to(device=device, dtype=sampling_dtype) block_state.action_scheduler = copy.deepcopy(components.scheduler) @@ -1039,20 +1048,22 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - dtype = components.transformer.dtype + sampling_dtype = torch.float32 tcf = components.vae_scale_factor_temporal - target_x0 = block_state.x0_tokens_vision.to(device=device) + target_x0 = block_state.x0_tokens_vision.to(device=device, dtype=sampling_dtype) current_conditional_frames = block_state.current_conditional_frames # Build the noisy target latents + conditioning mask from the clean target latents. latent_t = target_x0.shape[2] - condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype) + condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=sampling_dtype) latent_condition_frames = 0 if current_conditional_frames > 0: latent_condition_frames = (current_conditional_frames - 1) // tcf + 1 condition_mask[:latent_condition_frames] = 1.0 - noise = randn_tensor(tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=dtype) + noise = randn_tensor( + tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=sampling_dtype + ) block_state.latents = condition_mask * target_x0 + (1.0 - condition_mask) * noise block_state.velocity_mask = 1.0 - condition_mask block_state.condition_latents = condition_mask * target_x0 diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index eda37c8e99cf..ba6e06a15244 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -215,7 +215,10 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta transformer_kwargs = { name: value for name, value in transformer_kwargs.items() if name in transformer_args } - preds_vision, preds_sound, preds_action = components.transformer(**transformer_kwargs, return_dict=False) + with components.transformer.cache_context(pass_name): + preds_vision, preds_sound, preds_action = components.transformer( + **transformer_kwargs, return_dict=False + ) velocities[pass_name] = components._mask_velocity_predictions( preds_vision, preds_sound, @@ -227,8 +230,14 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta ) cond_velocity_vision, cond_velocity_sound, cond_velocity_action = velocities["cond"] + cond_velocity_vision = cond_velocity_vision.float() + cond_velocity_sound = cond_velocity_sound.float() if cond_velocity_sound is not None else None + cond_velocity_action = cond_velocity_action.float() if cond_velocity_action is not None else None if do_cfg: uncond_velocity_vision, uncond_velocity_sound, uncond_velocity_action = velocities["uncond"] + uncond_velocity_vision = uncond_velocity_vision.float() + uncond_velocity_sound = uncond_velocity_sound.float() if uncond_velocity_sound is not None else None + uncond_velocity_action = uncond_velocity_action.float() if uncond_velocity_action is not None else None block_state.velocity_vision = uncond_velocity_vision + block_state.guidance_scale * ( cond_velocity_vision - uncond_velocity_vision ) @@ -327,11 +336,14 @@ def intermediate_outputs(self) -> list[OutputParam]: @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + velocity_vision = block_state.velocity_vision.float() + latents = block_state.latents.float() + # Pass the generator so the scheduler's stochastic (SDE) re-noising is seedable/reproducible. block_state.latents = components.scheduler.step( - block_state.velocity_vision.unsqueeze(0), + velocity_vision.unsqueeze(0), t, - block_state.latents.unsqueeze(0), + latents.unsqueeze(0), generator=block_state.generator, return_dict=False, )[0].squeeze(0) @@ -468,13 +480,23 @@ def loop_inputs(self) -> list[InputParam]: @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) + components._current_step_index = None + components._current_sigma = None + components._num_timesteps = len(block_state.timesteps) with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: for i, t in enumerate(block_state.timesteps): + components._current_step_index = i + scheduler_sigmas = getattr(components.scheduler, "sigmas", None) + components._current_sigma = ( + scheduler_sigmas[i] if scheduler_sigmas is not None and i < len(scheduler_sigmas) else None + ) components, block_state = self.loop_step(components, block_state, i=i, t=t) if i == len(block_state.timesteps) - 1 or ( (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 ): progress_bar.update() + components._current_step_index = None + components._current_sigma = None self.set_block_state(state, block_state) return components, state @@ -706,21 +728,22 @@ def intermediate_outputs(self) -> list[OutputParam]: return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted (masked) transfer velocity.")] @staticmethod - def _forward(components, static, vision_tokens, vision_timesteps): - preds_vision, _, _ = components.transformer( - input_ids=static["input_ids"], - text_indexes=static["text_indexes"], - position_ids=static["position_ids"], - und_len=static["und_len"], - sequence_length=static["sequence_length"], - vision_tokens=vision_tokens, - vision_token_shapes=static["vision_token_shapes"], - vision_sequence_indexes=static["vision_sequence_indexes"], - vision_mse_loss_indexes=static["vision_mse_loss_indexes"], - vision_timesteps=vision_timesteps, - vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"], - return_dict=False, - ) + def _forward(components, static, vision_tokens, vision_timesteps, context_name): + with components.transformer.cache_context(context_name): + preds_vision, _, _ = components.transformer( + input_ids=static["input_ids"], + text_indexes=static["text_indexes"], + position_ids=static["position_ids"], + und_len=static["und_len"], + sequence_length=static["sequence_length"], + vision_tokens=vision_tokens, + vision_token_shapes=static["vision_token_shapes"], + vision_sequence_indexes=static["vision_sequence_indexes"], + vision_mse_loss_indexes=static["vision_mse_loss_indexes"], + vision_timesteps=vision_timesteps, + vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"], + return_dict=False, + ) return preds_vision[-1] @torch.no_grad() @@ -745,7 +768,11 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta uncond_full_static = denoiser_input_fields["uncond_full_static"] cond_full = self._forward( - components, cond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps + components, + cond_full_static, + block_state.vision_tokens_full, + block_state.vision_timesteps, + "cond", ) cond_no_control = None @@ -755,6 +782,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta cond_no_control_static, block_state.vision_tokens_target, block_state.vision_timesteps, + "cond_no_control", ) uncond_full = None @@ -764,8 +792,13 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta uncond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps, + "uncond", ) + cond_full = cond_full.float() + cond_no_control = cond_no_control.float() if cond_no_control is not None else None + uncond_full = uncond_full.float() if uncond_full is not None else None + if needs_control_cfg and needs_text_cfg: control_cond = cond_no_control + step_control * (cond_full - cond_no_control) velocity = uncond_full + step_guidance * (control_cond - uncond_full) @@ -840,7 +873,8 @@ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper): Runs the per-chunk transfer denoising loop over scheduler timesteps. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Inputs: timesteps (`Tensor`): diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index d7c6c846fc8e..d15a176a0824 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -53,11 +53,11 @@ # auto_docstring class Cosmos3TransferTextBlocks(SequentialPipelineBlocks): """ - Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) prompt in - transfer mode using the per-chunk frame count. + Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) prompt in transfer mode using the per-chunk frame count. Components: - video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) + video_processor (`VideoProcessor`) + text_tokenizer (`AutoTokenizer`) Inputs: control_videos (`dict`): @@ -121,10 +121,12 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): - Cosmos3TextEncoderStep runs otherwise. Components: - video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) + video_processor (`VideoProcessor`) + text_tokenizer (`AutoTokenizer`) Configs: - default_use_system_prompt (default: True) enable_safety_checker (default: True) + default_use_system_prompt (default: True) + enable_safety_checker (default: True) Inputs: control_videos (`dict`, *optional*): @@ -204,7 +206,8 @@ class Cosmos3AutoVaeEncoderStep(ConditionalPipelineBlocks): - when no action, image, or video conditioning is provided, this block is skipped. Components: - vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) Inputs: action (`CosmosActionCondition`, *optional*): @@ -313,7 +316,9 @@ class Cosmos3DecodeStep(SequentialPipelineBlocks): Decodes denoised latents into modality outputs. Components: - vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) Inputs: latents (`Tensor`): @@ -368,7 +373,8 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text-and-vision Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) @@ -439,7 +445,8 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, and sound Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) @@ -523,7 +530,8 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, and action Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) @@ -611,7 +619,8 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, sound, and action Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) @@ -707,13 +716,13 @@ def outputs(self): # auto_docstring class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): """ - Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a non-leaf - LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). Per-chunk - cross-carry (previous_output, output_chunks) lives on PipelineState. + Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState. Components: - vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer (`Cosmos3OmniTransformer`) scheduler - (`UniPCMultistepScheduler`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Inputs: chunk_id (`int`, *optional*, defaults to 0): @@ -842,6 +851,8 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) state.set("output_chunks", []) state.set("previous_output", None) for chunk_id in range(num_chunks): + if chunk_id > 0: + components.transformer._reset_stateful_cache() state.set("chunk_id", chunk_id) for _, block in self.sub_blocks.items(): components, state = block(components, state) @@ -854,8 +865,10 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks): Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop. Components: - transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler - (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + scheduler (`UniPCMultistepScheduler`) Inputs: cond_input_ids (`None`): @@ -973,8 +986,10 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): - vision runs otherwise. Components: - transformer (`Cosmos3OmniTransformer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler - (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + scheduler (`UniPCMultistepScheduler`) Configs: use_native_flow_schedule (default: False) @@ -1162,13 +1177,17 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): - `action_inverse_dynamics`: requires `prompt`, `action` Components: - video_processor (`VideoProcessor`) text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) transformer - (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) sound_tokenizer - (`Cosmos3AVAEAudioTokenizer`) + video_processor (`VideoProcessor`) + text_tokenizer (`AutoTokenizer`) + vae (`AutoencoderKLWan`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) Configs: - default_use_system_prompt (default: True) enable_safety_checker (default: True) use_native_flow_schedule - (default: False) + default_use_system_prompt (default: True) + enable_safety_checker (default: True) + use_native_flow_schedule (default: False) Inputs: control_videos (`dict`, *optional*): diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py index e168cc24d8cd..b6ec5d47eb8d 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py @@ -25,7 +25,8 @@ class Cosmos3DistilledAutoVaeEncoderStep(ConditionalPipelineBlocks): - when no image or video conditioning is provided, this block is skipped. Components: - vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) Inputs: video (`None`, *optional*): @@ -83,10 +84,12 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text-and-vision distilled Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`FlowMatchEulerDiscreteScheduler`) Configs: - is_distilled (default: True) distilled_sigmas (default: None) + is_distilled (default: True) + distilled_sigmas (default: None) Inputs: cond_input_ids (`None`): @@ -112,8 +115,8 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): num_inference_steps (`int`, *optional*): The number of denoising steps. guidance_scale (`float`, *optional*): - Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is - forced to 1.0. Passing a value other than 1.0 raises an error. + Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is forced to + 1.0. Passing a value other than 1.0 raises an error. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. @@ -161,11 +164,16 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): - `video2video`: requires `prompt`, `video` Components: - text_tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) transformer - (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) + text_tokenizer (`AutoTokenizer`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`FlowMatchEulerDiscreteScheduler`) Configs: - default_use_system_prompt (default: True) enable_safety_checker (default: True) is_distilled (default: True) + default_use_system_prompt (default: True) + enable_safety_checker (default: True) + is_distilled (default: True) distilled_sigmas (default: None) Inputs: @@ -179,7 +187,7 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): Width of the generated video or image in pixels. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. - use_system_prompt (`bool`, *optional*, defaults to True): + use_system_prompt (`bool | NoneType`, *optional*): Whether to prepend the Cosmos3 system prompt. add_resolution_template (`bool`, *optional*, defaults to True): Whether to add resolution metadata to the prompt. @@ -204,8 +212,8 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): num_inference_steps (`int`, *optional*): The number of denoising steps. guidance_scale (`float`, *optional*): - Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is - forced to 1.0. Passing a value other than 1.0 raises an error. + Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the scale is forced to + 1.0. Passing a value other than 1.0 raises an error. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. output_type (`str`, *optional*, defaults to pil): diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py index d6c09703c12e..fe6f3cbd721d 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py @@ -18,6 +18,28 @@ class Cosmos3OmniModularPipeline(ModularPipeline): inverse_image_resolution_template = "This image is not of {height}x{width} resolution." inverse_video_resolution_template = "This video is not of {height}x{width} resolution." + @property + def current_step_index(self): + return getattr(self, "_current_step_index", None) + + @property + def current_sigma(self): + return getattr(self, "_current_sigma", None) + + @property + def num_timesteps(self): + return getattr(self, "_num_timesteps", None) + + def __call__(self, *args, **kwargs): + transformer = getattr(self, "transformer", None) + try: + return super().__call__(*args, **kwargs) + finally: + self._current_step_index = None + self._current_sigma = None + if hasattr(transformer, "_reset_stateful_cache"): + transformer._reset_stateful_cache() + @property def vae_scale_factor_spatial(self): if getattr(self, "vae", None) is not None: diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 589e0ed3d6b0..9f22da1035bb 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -466,11 +466,10 @@ def __init__( self.inverse_image_resolution_template = "This image is not of {height}x{width} resolution." self.inverse_video_resolution_template = "This video is not of {height}x{width} resolution." - # Recommended quality-control negative prompts are documented in the Cosmos3 docs - # page (text2video / image2video). When the caller passes None we fall back to "". - # TODO YiYi & Daniel: fix for this use case in the base class def _get_execution_device(self) -> torch.device: + from ...hooks.group_offloading import _get_group_onload_device + # `self._execution_device` walks `self.components` and ultimately falls back to # `self.device`, which iterates modules in sorted order and ignores # `_exclude_from_cpu_offload`. With `safety_checker` registered, that path picks @@ -482,6 +481,11 @@ def _get_execution_device(self) -> torch.device: if not isinstance(component, torch.nn.Module): continue + try: + return _get_group_onload_device(component) + except ValueError: + pass + for module in component.modules(): hook = getattr(module, "_hf_hook", None) execution_device = getattr(hook, "execution_device", None) @@ -1292,6 +1296,14 @@ def _apply_video_safety_check(self, video: Any, output_type: str, device: torch. def current_timestep(self): return self._current_timestep + @property + def current_step_index(self): + return self._current_step_index + + @property + def current_sigma(self): + return self._current_sigma + @property def guidance_scale(self): return self._guidance_scale @@ -1486,6 +1498,8 @@ def __call__( ) self._current_timestep = None + self._current_step_index = None + self._current_sigma = None self._interrupt = False self._guidance_scale = guidance_scale @@ -1497,6 +1511,7 @@ def __call__( device = self._get_execution_device() dtype = self.transformer.dtype + sampling_dtype = torch.float32 if enable_safety_check and isinstance(self.safety_checker, CosmosSafetyChecker): self.safety_checker.to(device) @@ -1557,7 +1572,7 @@ def __call__( action_latents=action_latents, generator=generator, device=device, - dtype=dtype, + dtype=sampling_dtype, enable_sound=enable_sound, action=action, ) @@ -1691,6 +1706,8 @@ def __call__( continue self._current_timestep = t + self._current_step_index = i + self._current_sigma = self.scheduler.sigmas[i] timestep = t.item() # The transformer projections (proj_in / audio_proj_in) are bf16; cast the per-step @@ -1709,33 +1726,34 @@ def __call__( ) # --- Conditional pass --- - preds_vision, preds_sound, preds_action = self.transformer( - input_ids=cond_packed_static["input_ids"], - text_indexes=cond_packed_static["text_indexes"], - position_ids=cond_packed_static["position_ids"], - und_len=cond_packed_static["und_len"], - sequence_length=cond_packed_static["sequence_length"], - vision_tokens=[vision_tokens], - vision_token_shapes=cond_packed_static["vision_token_shapes"], - vision_sequence_indexes=cond_packed_static["vision_sequence_indexes"], - vision_mse_loss_indexes=cond_packed_static["vision_mse_loss_indexes"], - vision_timesteps=vision_timesteps, - vision_noisy_frame_indexes=cond_packed_static["vision_noisy_frame_indexes"], - sound_tokens=[sound_tokens] if sound_tokens is not None else None, - sound_token_shapes=cond_packed_static.get("sound_token_shapes"), - sound_sequence_indexes=cond_packed_static.get("sound_sequence_indexes"), - sound_mse_loss_indexes=cond_packed_static.get("sound_mse_loss_indexes"), - sound_timesteps=sound_timesteps, - sound_noisy_frame_indexes=cond_packed_static.get("sound_noisy_frame_indexes"), - action_tokens=[action_tokens] if action_tokens is not None else None, - action_token_shapes=cond_packed_static.get("action_token_shapes"), - action_sequence_indexes=cond_packed_static.get("action_sequence_indexes"), - action_mse_loss_indexes=cond_packed_static.get("action_mse_loss_indexes"), - action_timesteps=action_timesteps, - action_noisy_frame_indexes=cond_packed_static.get("action_noisy_frame_indexes"), - action_domain_ids=[action_domain_id] if action_domain_id is not None else None, - return_dict=False, - ) + with self.transformer.cache_context("cond"): + preds_vision, preds_sound, preds_action = self.transformer( + input_ids=cond_packed_static["input_ids"], + text_indexes=cond_packed_static["text_indexes"], + position_ids=cond_packed_static["position_ids"], + und_len=cond_packed_static["und_len"], + sequence_length=cond_packed_static["sequence_length"], + vision_tokens=[vision_tokens], + vision_token_shapes=cond_packed_static["vision_token_shapes"], + vision_sequence_indexes=cond_packed_static["vision_sequence_indexes"], + vision_mse_loss_indexes=cond_packed_static["vision_mse_loss_indexes"], + vision_timesteps=vision_timesteps, + vision_noisy_frame_indexes=cond_packed_static["vision_noisy_frame_indexes"], + sound_tokens=[sound_tokens] if sound_tokens is not None else None, + sound_token_shapes=cond_packed_static.get("sound_token_shapes"), + sound_sequence_indexes=cond_packed_static.get("sound_sequence_indexes"), + sound_mse_loss_indexes=cond_packed_static.get("sound_mse_loss_indexes"), + sound_timesteps=sound_timesteps, + sound_noisy_frame_indexes=cond_packed_static.get("sound_noisy_frame_indexes"), + action_tokens=[action_tokens] if action_tokens is not None else None, + action_token_shapes=cond_packed_static.get("action_token_shapes"), + action_sequence_indexes=cond_packed_static.get("action_sequence_indexes"), + action_mse_loss_indexes=cond_packed_static.get("action_mse_loss_indexes"), + action_timesteps=action_timesteps, + action_noisy_frame_indexes=cond_packed_static.get("action_noisy_frame_indexes"), + action_domain_ids=[action_domain_id] if action_domain_id is not None else None, + return_dict=False, + ) cond_v_vision, cond_v_sound, cond_v_action = self._mask_velocity_predictions( preds_vision, preds_sound, @@ -1749,33 +1767,34 @@ def __call__( # --- Unconditional pass (Skip if not using CFG) --- uncond_v_vision = uncond_v_sound = uncond_v_action = None if self.do_classifier_free_guidance: - preds_vision, preds_sound, preds_action = self.transformer( - input_ids=uncond_packed_static["input_ids"], - text_indexes=uncond_packed_static["text_indexes"], - position_ids=uncond_packed_static["position_ids"], - und_len=uncond_packed_static["und_len"], - sequence_length=uncond_packed_static["sequence_length"], - vision_tokens=[vision_tokens], - vision_token_shapes=uncond_packed_static["vision_token_shapes"], - vision_sequence_indexes=uncond_packed_static["vision_sequence_indexes"], - vision_mse_loss_indexes=uncond_packed_static["vision_mse_loss_indexes"], - vision_timesteps=vision_timesteps, - vision_noisy_frame_indexes=uncond_packed_static["vision_noisy_frame_indexes"], - sound_tokens=[sound_tokens] if sound_tokens is not None else None, - sound_token_shapes=uncond_packed_static.get("sound_token_shapes"), - sound_sequence_indexes=uncond_packed_static.get("sound_sequence_indexes"), - sound_mse_loss_indexes=uncond_packed_static.get("sound_mse_loss_indexes"), - sound_timesteps=sound_timesteps, - sound_noisy_frame_indexes=uncond_packed_static.get("sound_noisy_frame_indexes"), - action_tokens=[action_tokens] if action_tokens is not None else None, - action_token_shapes=uncond_packed_static.get("action_token_shapes"), - action_sequence_indexes=uncond_packed_static.get("action_sequence_indexes"), - action_mse_loss_indexes=uncond_packed_static.get("action_mse_loss_indexes"), - action_timesteps=action_timesteps, - action_noisy_frame_indexes=uncond_packed_static.get("action_noisy_frame_indexes"), - action_domain_ids=[action_domain_id] if action_domain_id is not None else None, - return_dict=False, - ) + with self.transformer.cache_context("uncond"): + preds_vision, preds_sound, preds_action = self.transformer( + input_ids=uncond_packed_static["input_ids"], + text_indexes=uncond_packed_static["text_indexes"], + position_ids=uncond_packed_static["position_ids"], + und_len=uncond_packed_static["und_len"], + sequence_length=uncond_packed_static["sequence_length"], + vision_tokens=[vision_tokens], + vision_token_shapes=uncond_packed_static["vision_token_shapes"], + vision_sequence_indexes=uncond_packed_static["vision_sequence_indexes"], + vision_mse_loss_indexes=uncond_packed_static["vision_mse_loss_indexes"], + vision_timesteps=vision_timesteps, + vision_noisy_frame_indexes=uncond_packed_static["vision_noisy_frame_indexes"], + sound_tokens=[sound_tokens] if sound_tokens is not None else None, + sound_token_shapes=uncond_packed_static.get("sound_token_shapes"), + sound_sequence_indexes=uncond_packed_static.get("sound_sequence_indexes"), + sound_mse_loss_indexes=uncond_packed_static.get("sound_mse_loss_indexes"), + sound_timesteps=sound_timesteps, + sound_noisy_frame_indexes=uncond_packed_static.get("sound_noisy_frame_indexes"), + action_tokens=[action_tokens] if action_tokens is not None else None, + action_token_shapes=uncond_packed_static.get("action_token_shapes"), + action_sequence_indexes=uncond_packed_static.get("action_sequence_indexes"), + action_mse_loss_indexes=uncond_packed_static.get("action_mse_loss_indexes"), + action_timesteps=action_timesteps, + action_noisy_frame_indexes=uncond_packed_static.get("action_noisy_frame_indexes"), + action_domain_ids=[action_domain_id] if action_domain_id is not None else None, + return_dict=False, + ) uncond_v_vision, uncond_v_sound, uncond_v_action = self._mask_velocity_predictions( preds_vision, preds_sound, @@ -1786,6 +1805,13 @@ def __call__( raw_action_dim=raw_action_dim_resolved, ) + cond_v_vision = cond_v_vision.float() + cond_v_sound = cond_v_sound.float() if cond_v_sound is not None else None + cond_v_action = cond_v_action.float() if cond_v_action is not None else None + uncond_v_vision = uncond_v_vision.float() if uncond_v_vision is not None else None + uncond_v_sound = uncond_v_sound.float() if uncond_v_sound is not None else None + uncond_v_action = uncond_v_action.float() if uncond_v_action is not None else None + # --- CFG combine + per-modality scheduler step --- # UniPC's multistep_uni_p_bh_update einsum ("k,bkc...->bc...") requires sample # to carry a batch dim; per-modality latents have no batch axis, so wrap for the step. @@ -1829,12 +1855,14 @@ def __call__( for key in callback_on_step_end_tensor_inputs: callback_kwargs[key] = locals()[key] callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) - latents = callback_outputs.pop("latents", latents) + latents = callback_outputs.pop("latents", latents).float() if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() self._current_timestep = None + self._current_step_index = None + self._current_sigma = None # 8. Postprocess + decode sound = self.decode_sound(sound_latents) if sound_latents is not None else None diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 1598814f835a..94700aa2689c 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -272,6 +272,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class SeaCacheConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class SmoothedEnergyGuidanceConfig(metaclass=DummyObject): _backends = ["torch"] @@ -337,6 +352,10 @@ def apply_pyramid_attention_broadcast(*args, **kwargs): requires_backends(apply_pyramid_attention_broadcast, ["torch"]) +def apply_sea_cache(*args, **kwargs): + requires_backends(apply_sea_cache, ["torch"]) + + def apply_taylorseer_cache(*args, **kwargs): requires_backends(apply_taylorseer_cache, ["torch"]) diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 760a3fac04e0..2d7d5ae23257 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -9,6 +9,8 @@ MagCacheTesterMixin, PyramidAttentionBroadcastConfigMixin, PyramidAttentionBroadcastTesterMixin, + SeaCacheConfigMixin, + SeaCacheTesterMixin, TaylorSeerCacheConfigMixin, TaylorSeerCacheTesterMixin, ) @@ -90,6 +92,8 @@ "NunchakuLiteTesterMixin", "PyramidAttentionBroadcastConfigMixin", "PyramidAttentionBroadcastTesterMixin", + "SeaCacheConfigMixin", + "SeaCacheTesterMixin", "TaylorSeerCacheConfigMixin", "TaylorSeerCacheTesterMixin", "QuantizationCompileTesterMixin", diff --git a/tests/models/testing_utils/cache.py b/tests/models/testing_utils/cache.py index 8357d34f3077..c8d757140348 100644 --- a/tests/models/testing_utils/cache.py +++ b/tests/models/testing_utils/cache.py @@ -23,12 +23,19 @@ FirstBlockCacheConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, + SeaCacheConfig, TaylorSeerCacheConfig, ) from diffusers.hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK from diffusers.hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK from diffusers.hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK from diffusers.hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK +from diffusers.hooks.sea_cache import ( + _SEA_CACHE_BLOCK_HOOK, + _SEA_CACHE_LEADER_BLOCK_HOOK, + _SEA_CACHE_POST_NORM_HOOK, + _SEA_CACHE_ROOT_HOOK, +) from diffusers.hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK from diffusers.models.cache_utils import CacheMixin @@ -433,6 +440,143 @@ def test_fbc_reset_stateful_cache(self): self._test_reset_stateful_cache() +@is_cache +class SeaCacheConfigMixin: + """ + Base mixin providing SeaCache config. + + Expected class attributes: + - model_class: The model class to test (must use CacheMixin) + """ + + SEA_CACHE_CONFIG = { + "threshold": 100.0, + "retention_steps": 0, + "cache_end_steps": 0, + } + + def _get_cache_config(self): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} + self._sea_cache_runtime = runtime + return SeaCacheConfig( + **self.SEA_CACHE_CONFIG, + current_step_callback=lambda: runtime["step"], + current_sigma_callback=lambda: runtime["sigma"], + num_inference_steps_callback=lambda: runtime["num_steps"], + ) + + def _get_hook_names(self): + return [ + _SEA_CACHE_ROOT_HOOK, + _SEA_CACHE_LEADER_BLOCK_HOOK, + _SEA_CACHE_BLOCK_HOOK, + _SEA_CACHE_POST_NORM_HOOK, + ] + + +@is_cache +class SeaCacheTesterMixin(SeaCacheConfigMixin, CacheTesterMixin): + """ + Mixin class for testing SeaCache on models. + + Expected methods to be implemented by subclasses: + - get_init_dict(): Returns dict of arguments to initialize the model + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Pytest mark: cache + Use `pytest -m "not cache"` to skip these tests + """ + + @staticmethod + def _unwrap_cache_output(output): + while isinstance(output, (list, tuple)): + output = output[0] + return output + + def _get_modified_cache_inputs(self): + inputs = self.get_dummy_inputs() + value = inputs[self.cache_input_key] + if isinstance(value, torch.Tensor): + inputs[self.cache_input_key] = value + 0.1 + else: + inputs[self.cache_input_key] = [tensor + 0.1 for tensor in value] + return inputs + + @torch.no_grad() + def _test_cache_inference(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model.enable_cache(self._get_cache_config()) + + with model.cache_context("sea_cache_test"): + model(**self.get_dummy_inputs(), return_dict=False) + + self._sea_cache_runtime.update(step=1, sigma=0.6) + modified_inputs = self._get_modified_cache_inputs() + with model.cache_context("sea_cache_test"): + output_with_cache = self._unwrap_cache_output(model(**modified_inputs, return_dict=False)) + + assert output_with_cache is not None + assert not torch.isnan(output_with_cache).any() + + model.disable_cache() + output_without_cache = self._unwrap_cache_output(model(**modified_inputs, return_dict=False)) + assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5) + + @torch.no_grad() + def _test_cache_context_manager(self, atol=1e-5, rtol=0): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model.enable_cache(self._get_cache_config()) + inputs = self.get_dummy_inputs() + + with model.cache_context("context_1"): + output_ctx1 = self._unwrap_cache_output(model(**inputs, return_dict=False)) + with model.cache_context("context_2"): + output_ctx2 = self._unwrap_cache_output(model(**inputs, return_dict=False)) + + assert_tensors_close( + output_ctx1, + output_ctx2, + atol=atol, + rtol=rtol, + msg="First pass in different cache contexts should produce the same output.", + ) + model.disable_cache() + + @torch.no_grad() + def _test_reset_stateful_cache(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model.enable_cache(self._get_cache_config()) + + with model.cache_context("sea_cache_test"): + model(**self.get_dummy_inputs(), return_dict=False) + model._reset_stateful_cache() + model.disable_cache() + + @require_cache_mixin + def test_sea_cache_enable_disable_state(self): + self._test_cache_enable_disable_state() + + @require_cache_mixin + def test_sea_cache_double_enable_raises_error(self): + self._test_cache_double_enable_raises_error() + + @require_cache_mixin + def test_sea_cache_hooks_registered(self): + self._test_cache_hooks_registered() + + @require_cache_mixin + def test_sea_cache_inference(self): + self._test_cache_inference() + + @require_cache_mixin + def test_sea_cache_context_manager(self): + self._test_cache_context_manager() + + @require_cache_mixin + def test_sea_cache_reset_stateful_cache(self): + self._test_reset_stateful_cache() + + @is_cache class FasterCacheConfigMixin: """ diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index 0ead8cb1a6bb..8a078b64ba76 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -16,11 +16,15 @@ import pytest import torch -from diffusers import Cosmos3OmniTransformer +from diffusers import Cosmos3OmniTransformer, SeaCacheConfig +from diffusers.hooks._helpers import TransformerBlockRegistry +from diffusers.hooks.sea_cache import _SEA_CACHE_ROOT_HOOK +from diffusers.models.cache_utils import CacheMixin from diffusers.models.transformers.transformer_cosmos3 import ( Cosmos3NemotronRMSNorm, Cosmos3OmniTransformerOutput, Cosmos3PackedMoTAttention, + Cosmos3VLTextMoTDecoderLayer, ) from diffusers.utils.torch_utils import randn_tensor @@ -30,6 +34,7 @@ BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, + SeaCacheTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -103,6 +108,134 @@ def output_shape(self) -> tuple[int, ...]: class TestCosmos3OmniTransformerModel(Cosmos3OmniTransformerTesterConfig, ModelTesterMixin): + @pytest.mark.parametrize("indicator_source", ["first_block", "raw_vision_latents"]) + def test_cosmos3_supports_sea_cache_without_changing_state_dict_keys(self, indicator_source): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + state_dict_keys = set(model.state_dict()) + norm_calls = {"und": 0, "gen": 0} + original_und_norm_forward = model.norm.forward + original_gen_norm_forward = model.norm_moe_gen.forward + + def counted_und_norm_forward(*args, **kwargs): + norm_calls["und"] += 1 + return original_und_norm_forward(*args, **kwargs) + + def counted_gen_norm_forward(*args, **kwargs): + norm_calls["gen"] += 1 + return original_gen_norm_forward(*args, **kwargs) + + model.norm.forward = counted_und_norm_forward + model.norm_moe_gen.forward = counted_gen_norm_forward + runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} + config = SeaCacheConfig( + threshold=100.0, + cache_end_steps=0, + indicator_source=indicator_source, + current_step_callback=lambda: runtime["step"], + current_sigma_callback=lambda: runtime["sigma"], + num_inference_steps_callback=lambda: runtime["num_steps"], + ) + + assert isinstance(model, CacheMixin) + model.enable_cache(config) + assert set(model.state_dict()) == state_dict_keys + + inputs = self.get_dummy_inputs() + with torch.no_grad(), model.cache_context("cond"): + model(**inputs) + assert norm_calls == {"und": 1, "gen": 1} + + runtime.update(step=1, sigma=0.6) + cached_inputs = self.get_dummy_inputs() + cached_inputs["vision_tokens"] = [cached_inputs["vision_tokens"][0] + 0.1] + with torch.no_grad(), model.cache_context("cond"): + output = model(**cached_inputs) + + # A hit bypasses both final pathway normalizations but still runs the prediction heads. + assert norm_calls == {"und": 1, "gen": 1} + assert output.sample[0].shape == self.output_shape + + model.disable_cache() + assert set(model.state_dict()) == state_dict_keys + with torch.no_grad(): + model(**self.get_dummy_inputs()) + assert norm_calls == {"und": 2, "gen": 2} + + def test_cosmos3_sea_cache_post_norm_boundary_numeric_semantics(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} + config = SeaCacheConfig( + threshold=100.0, + cache_end_steps=0, + current_step_callback=lambda: runtime["step"], + current_sigma_callback=lambda: runtime["sigma"], + num_inference_steps_callback=lambda: runtime["num_steps"], + ) + first_block_gen_inputs = [] + returned_und_outputs = [] + returned_gen_outputs = [] + projection_head_inputs = [] + projection_head_calls = 0 + + def capture_first_block_gen_input(module, args): + first_block_gen_inputs.append(args[1].detach().clone()) + + def capture_und_output(module, args, output): + returned_und_outputs.append(output.detach().clone()) + + def capture_gen_output(module, args, output): + returned_gen_outputs.append(output.detach().clone()) + + original_projection_forward = model.proj_out.forward + + def counted_projection_forward(hidden_states): + nonlocal projection_head_calls + projection_head_calls += 1 + projection_head_inputs.append(hidden_states.detach().clone()) + return original_projection_forward(hidden_states) + + model.layers[0].register_forward_pre_hook(capture_first_block_gen_input) + model.norm.register_forward_hook(capture_und_output) + model.norm_moe_gen.register_forward_hook(capture_gen_output) + model.proj_out.forward = counted_projection_forward + model.enable_cache(config) + + with torch.no_grad(), model.cache_context("cond"): + model(**self.get_dummy_inputs()) + + root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) + state = root_hook.state_manager._state_cache["cond"] + cached_step, cached_und_output, cached_gen_residual = state.history[-1] + assert cached_step == 0 + torch.testing.assert_close(cached_und_output, returned_und_outputs[0]) + torch.testing.assert_close( + cached_gen_residual, + returned_gen_outputs[0] - first_block_gen_inputs[0], + ) + + runtime.update(step=1, sigma=0.6) + cached_inputs = self.get_dummy_inputs() + cached_inputs["vision_tokens"] = [cached_inputs["vision_tokens"][0] + 0.1] + with torch.no_grad(), model.cache_context("cond"): + model(**cached_inputs) + + torch.testing.assert_close(returned_und_outputs[1], cached_und_output) + torch.testing.assert_close( + returned_gen_outputs[1], + first_block_gen_inputs[-1] + cached_gen_residual, + ) + torch.testing.assert_close(projection_head_inputs[1], returned_gen_outputs[1]) + assert projection_head_calls == 2 + + def test_cosmos3_decoder_layer_cache_metadata_tracks_generation_stream(self): + metadata = TransformerBlockRegistry.get(Cosmos3VLTextMoTDecoderLayer) + + assert metadata.return_hidden_states_index == 1 + assert metadata.return_encoder_hidden_states_index == 0 + assert metadata.hidden_states_argument_name == "gen_seq" + assert metadata.encoder_hidden_states_argument_name == "und_seq" + assert metadata.hidden_states_norm_module_name == "input_layernorm_moe_gen" + def test_output_format(self): model = self.model_class(**self.get_init_dict()).to(torch_device).eval() @@ -243,6 +376,10 @@ def test_cosmos3_nemotron_rms_norm_multiplies_in_float32(self): torch.testing.assert_close(norm(hidden_states), expected, rtol=0, atol=0) +class TestCosmos3OmniTransformerSeaCache(Cosmos3OmniTransformerTesterConfig, SeaCacheTesterMixin): + cache_input_key = "vision_tokens" + + class TestCosmos3OmniTransformerMemory(Cosmos3OmniTransformerTesterConfig, MemoryTesterMixin): @pytest.mark.skip("The transformer returns one tensor list per generated modality.") def test_layerwise_casting_training(self): diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index b5bf5fb04393..abe293471ac8 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -13,12 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import contextmanager +from unittest import mock + import numpy as np import pytest import torch from PIL import Image -from diffusers import ModularPipeline, UniPCMultistepScheduler +from diffusers import CosmosActionCondition, ModularPipeline, UniPCMultistepScheduler from diffusers.modular_pipelines import ( Cosmos3OmniBlocks, Cosmos3OmniModularPipeline, @@ -27,12 +30,16 @@ from diffusers.modular_pipelines.cosmos.before_denoise import ( Cosmos3ActionDenoiseInputStep, Cosmos3ActionPackSequenceStep, + Cosmos3ActionPrepareLatentsStep, Cosmos3SetTimestepsStep, Cosmos3SoundDenoiseInputStep, + Cosmos3SoundPrepareLatentsStep, + Cosmos3TransferPrepareLatentsStep, Cosmos3VisionDenoiseInputStep, Cosmos3VisionPackSequenceStep, ) from diffusers.modular_pipelines.cosmos.encoders import Cosmos3TextEncoderStep +from diffusers.modular_pipelines.cosmos.modular_blocks_cosmos3 import Cosmos3TransferChunkDenoiseStep from ...testing_utils import torch_device from ..testing_utils import ( @@ -167,6 +174,105 @@ def test_num_images_per_prompt(self): def test_float16_inference(self): pass + def test_transformer_cache_contexts_receive_exact_scheduler_metadata(self): + pipe = self.get_pipeline().to(torch_device) + observed = [] + + @contextmanager + def record_context(name): + observed.append((name, pipe.current_step_index, pipe.current_sigma)) + yield + + with mock.patch.object(pipe.transformer, "cache_context", side_effect=record_context): + pipe(**self.get_dummy_inputs(), output=self.output_name) + + assert [name for name, _, _ in observed] == ["cond", "uncond", "cond", "uncond"] + for call_index, (_, step_index, sigma) in enumerate(observed): + expected_step = call_index // 2 + assert step_index == expected_step + torch.testing.assert_close(sigma, pipe.scheduler.sigmas[expected_step]) + assert pipe.current_step_index is None + assert pipe.current_sigma is None + + def _get_sampling_state_block_pipe(self, block): + pipe = block.init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(torch_dtype=torch.bfloat16) + pipe.to(torch_device) + return pipe + + def test_sound_prepare_latents_uses_fp32(self): + pipe = self._get_sampling_state_block_pipe(Cosmos3SoundPrepareLatentsStep()) + + outputs = pipe( + num_frames=5, + fps=24.0, + generator=self.get_generator(0), + output=["sound_latents", "sound_condition_mask"], + ) + + assert outputs["sound_latents"].dtype == torch.float32 + assert outputs["sound_condition_mask"].dtype == torch.float32 + + def test_action_prepare_latents_uses_fp32(self): + pipe = self._get_sampling_state_block_pipe(Cosmos3ActionPrepareLatentsStep()) + action = CosmosActionCondition( + mode="policy", + chunk_size=2, + domain_name="av", + image=torch.zeros(3, 16, 16), + ) + + outputs = pipe( + action=action, + action_condition_frame_indexes=[], + generator=self.get_generator(0), + output=["action_latents", "action_condition_mask"], + ) + + assert outputs["action_latents"].dtype == torch.float32 + assert outputs["action_condition_mask"].dtype == torch.float32 + + def test_transfer_prepare_latents_uses_fp32(self): + pipe = self._get_sampling_state_block_pipe(Cosmos3TransferPrepareLatentsStep()) + + outputs = pipe( + x0_tokens_vision=torch.zeros(1, 4, 2, 2, 2), + current_conditional_frames=1, + generator=self.get_generator(0), + output=["latents", "velocity_mask", "condition_latents"], + ) + + assert outputs["latents"].dtype == torch.float32 + assert outputs["velocity_mask"].dtype == torch.float32 + assert outputs["condition_latents"].dtype == torch.float32 + + def test_transfer_chunks_reset_stateful_cache_at_boundaries(self): + block = Cosmos3TransferChunkDenoiseStep() + child_block = mock.Mock(side_effect=lambda components, state: (components, state)) + block.sub_blocks = {"child": child_block} + components = mock.Mock() + state = mock.Mock() + state.get.return_value = 3 + + block(components, state) + + assert child_block.call_count == 3 + assert components.transformer._reset_stateful_cache.call_count == 2 + assert [call.args for call in state.set.call_args_list if call.args[0] == "chunk_id"] == [ + ("chunk_id", 0), + ("chunk_id", 1), + ("chunk_id", 2), + ] + + def test_sampling_state_uses_fp32_for_modular_cfg_and_scheduler(self): + pipe = self.get_pipeline(dtype=torch.bfloat16).to(torch_device) + inputs = self.get_dummy_inputs() + + outputs = pipe(**inputs, output=["velocity_vision", "latents"]) + + assert outputs["velocity_vision"].dtype == torch.float32 + assert outputs["latents"].dtype == torch.float32 + def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self): pipe = self.get_pipeline() vae_encoder = pipe.blocks.sub_blocks["vae_encoder"] diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py index 97bb2b478989..86bc6ad4d6f2 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -19,6 +19,8 @@ from diffusers import ModularPipeline from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline +from diffusers.modular_pipelines.cosmos.before_denoise import Cosmos3VisionPrepareLatentsStep +from diffusers.modular_pipelines.cosmos.encoders import Cosmos3DistilledTextEncoderStep from ...testing_utils import torch_device from ..testing_utils import ( @@ -113,6 +115,56 @@ def test_declares_distilled_configs(self): pipe = self.pipeline_class() assert pipe.config.is_distilled is True assert pipe.config.distilled_sigmas is None + assert pipe.config.default_use_system_prompt is True + + def test_distilled_text_step_uses_system_prompt_config_fallback(self): + text_pipe = Cosmos3DistilledTextEncoderStep().init_pipeline(self.pretrained_model_name_or_path) + text_pipe.load_components() + text_pipe.disable_safety_checker() + + inputs = { + "prompt": "A small robot moves across a table.", + "num_frames": 5, + "height": 32, + "width": 32, + } + default_with_system_prompt = text_pipe(**inputs, output="cond_input_ids") + explicit_with_system_prompt = text_pipe(**inputs, use_system_prompt=True, output="cond_input_ids") + explicit_without_system_prompt = text_pipe(**inputs, use_system_prompt=False, output="cond_input_ids") + + text_pipe.update_components(default_use_system_prompt=False) + default_without_system_prompt = text_pipe(**inputs, output="cond_input_ids") + updated_with_system_prompt = text_pipe(**inputs, use_system_prompt=True, output="cond_input_ids") + updated_without_system_prompt = text_pipe(**inputs, use_system_prompt=False, output="cond_input_ids") + + assert default_with_system_prompt == explicit_with_system_prompt == updated_with_system_prompt + assert explicit_without_system_prompt == default_without_system_prompt == updated_without_system_prompt + assert len(default_with_system_prompt) > len(default_without_system_prompt) + + def test_prepare_vision_latents_uses_fp32(self): + prepare_pipe = Cosmos3VisionPrepareLatentsStep().init_pipeline(self.pretrained_model_name_or_path) + prepare_pipe.load_components(torch_dtype=torch.bfloat16) + prepare_pipe.to(torch_device) + + outputs = prepare_pipe( + num_frames=5, + height=32, + width=32, + fps=24.0, + generator=self.get_generator(0), + output=["latents", "vision_condition_mask"], + ) + + assert outputs["latents"].dtype == torch.float32 + assert outputs["vision_condition_mask"].dtype == torch.float32 + + def test_distilled_scheduler_uses_fp32_state(self): + pipe = self.get_pipeline(torch_dtype=torch.bfloat16).to(torch_device) + inputs = self.get_dummy_inputs() + + latents = pipe(**inputs, output=self.output_name) + + assert latents.dtype == torch.float32 def test_vae_encoder_rejects_image_and_video_together(self): vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"] diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index dfc1dd0a1ea1..fdcb8cc53931 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -13,6 +13,7 @@ # limitations under the License. import unittest +from contextlib import contextmanager from unittest import mock import numpy as np @@ -20,7 +21,12 @@ from PIL import Image from transformers import AutoTokenizer -from diffusers import AutoencoderKLWan, Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler +from diffusers import ( + AutoencoderKLWan, + Cosmos3OmniPipeline, + Cosmos3OmniTransformer, + UniPCMultistepScheduler, +) from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image from ...testing_utils import enable_full_determinism, torch_device @@ -118,6 +124,60 @@ def test_inference(self): self.assertEqual(video.shape, (1, 16, 16, 3)) + def test_fp32_sampling_state_keeps_transformer_inputs_in_model_dtype(self): + pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipeline.transformer.to(dtype=torch.float16) + pipeline.set_progress_bar_config(disable=None) + transformer_input_dtypes = [] + callback_latent_dtypes = [] + + def transformer_forward(**kwargs): + vision_tokens = kwargs["vision_tokens"][0] + transformer_input_dtypes.append(vision_tokens.dtype) + return ([torch.zeros_like(vision_tokens)], None, None) + + def callback_on_step_end(_pipeline, _step_index, _timestep, callback_kwargs): + callback_latent_dtypes.append(callback_kwargs["latents"].dtype) + return {"latents": callback_kwargs["latents"].to(torch.float16)} + + inputs = self.get_dummy_inputs(torch_device) + inputs.update( + output_type="latent", + enable_safety_check=False, + callback_on_step_end=callback_on_step_end, + ) + with mock.patch.object(pipeline.transformer, "forward", side_effect=transformer_forward): + latents = pipeline(**inputs).video + + assert transformer_input_dtypes + assert all(dtype == torch.float16 for dtype in transformer_input_dtypes) + assert callback_latent_dtypes + assert all(dtype == torch.float32 for dtype in callback_latent_dtypes) + assert latents.dtype == torch.float32 + + def test_transformer_cache_contexts_receive_exact_scheduler_metadata(self): + pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipeline.set_progress_bar_config(disable=None) + observed = [] + + @contextmanager + def record_context(name): + observed.append((name, pipeline.current_step_index, pipeline.current_sigma)) + yield + + inputs = self.get_dummy_inputs(torch_device) + inputs["guidance_scale"] = 2.0 + with mock.patch.object(pipeline.transformer, "cache_context", side_effect=record_context): + pipeline(**inputs) + + assert [name for name, _, _ in observed] == ["cond", "uncond", "cond", "uncond"] + for call_index, (_, step_index, sigma) in enumerate(observed): + expected_step = call_index // 2 + assert step_index == expected_step + torch.testing.assert_close(sigma, pipeline.scheduler.sigmas[expected_step]) + assert pipeline.current_step_index is None + assert pipeline.current_sigma is None + def test_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(self): components = self.get_dummy_components() components["default_use_system_prompt"] = False