From 575abcef7df3c57de4f0940989cf02b94df7b034 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 03:17:48 -0700 Subject: [PATCH 01/20] Add optional FP32 sampling state for Cosmos3 --- .../cosmos/before_denoise.py | 35 ++++++++-- .../modular_pipelines/cosmos/denoise.py | 16 ++++- .../modular_blocks_cosmos3_distilled.py | 14 +++- .../pipelines/cosmos/pipeline_cosmos3_omni.py | 19 +++++- ...test_modular_pipeline_cosmos3_distilled.py | 67 ++++++++++++++++++- tests/pipelines/cosmos/test_cosmos3.py | 31 +++++++++ 6 files changed, 170 insertions(+), 12 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index 7bf431aa855b..298312e450c9 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -68,6 +68,10 @@ def description(self) -> str: def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ConfigSpec(name="default_use_fp32_sampling_state", default=False)] + @property def inputs(self) -> list[InputParam]: return [ @@ -98,6 +102,15 @@ def inputs(self) -> list[InputParam]: description="Pre-generated noisy vision latents.", ), InputParam.template("generator"), + InputParam( + name="use_fp32_sampling_state", + type_hint=bool | None, + default=None, + description=( + "Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the " + "pipeline's `default_use_fp32_sampling_state` config." + ), + ), ] @property @@ -121,13 +134,20 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=torch.Tensor, description="Clean encoded vision latents used to re-anchor image conditioning each step.", ), + OutputParam( + "use_fp32_sampling_state", + type_hint=bool, + description="Whether vision sampling state is kept in float32.", + ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - dtype = components.transformer.dtype + if block_state.use_fp32_sampling_state is None: + block_state.use_fp32_sampling_state = components.config.default_use_fp32_sampling_state + sampling_dtype = torch.float32 if block_state.use_fp32_sampling_state else components.transformer.dtype x0_tokens_vision = block_state.x0_tokens_vision if x0_tokens_vision is None: @@ -151,21 +171,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 diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index eda37c8e99cf..ee28ea935883 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -319,6 +319,12 @@ def inputs(self) -> list[InputParam]: description="Indexes of conditioned vision latent frames; non-empty for image-to-video.", ), InputParam.template("generator"), + InputParam( + name="use_fp32_sampling_state", + type_hint=bool, + default=False, + description="Whether to keep the distilled vision scheduler state in float32.", + ), ] @property @@ -327,11 +333,17 @@ 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 + latents = block_state.latents + if block_state.use_fp32_sampling_state: + velocity_vision = velocity_vision.float() + latents = 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) 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..a0de9459e5ec 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py @@ -86,7 +86,8 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) Configs: - is_distilled (default: True) distilled_sigmas (default: None) + default_use_fp32_sampling_state (default: False) is_distilled (default: True) distilled_sigmas (default: + None) Inputs: cond_input_ids (`None`): @@ -109,6 +110,9 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the pipeline's + `default_use_fp32_sampling_state` config. num_inference_steps (`int`, *optional*): The number of denoising steps. guidance_scale (`float`, *optional*): @@ -165,8 +169,9 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) Configs: - default_use_system_prompt (default: True) enable_safety_checker (default: True) is_distilled (default: True) - distilled_sigmas (default: None) + default_use_system_prompt (default: True) enable_safety_checker (default: True) + default_use_fp32_sampling_state (default: False) is_distilled (default: True) distilled_sigmas (default: + None) Inputs: prompt (`str`): @@ -201,6 +206,9 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the pipeline's + `default_use_fp32_sampling_state` config. num_inference_steps (`int`, *optional*): The number of denoising steps. guidance_scale (`float`, *optional*): diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 0d0851310578..d793484fd8b7 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -1341,6 +1341,7 @@ def __call__( add_resolution_template: bool = True, add_duration_template: bool = True, enable_safety_check: bool = True, + use_fp32_sampling_state: bool = False, ) -> Cosmos3OmniPipelineOutput: r""" Run the Cosmos 3 omni pipeline end-to-end: encode the (optional) conditioning image/video, denoise vision and @@ -1436,6 +1437,11 @@ def __call__( When `True` and a `CosmosSafetyChecker` is attached, runs the text guardrail on the prompt before generation and the video guardrail on the decoded frames. Set to `False` to skip both for this call; the checker remains loaded for subsequent calls. + use_fp32_sampling_state (`bool`, *optional*, defaults to `False`): + When `True`, keeps vision, sound, and action denoising latents plus classifier-free-guidance arithmetic + in `torch.float32`. Transformer inputs are still cast to the transformer's dtype before each forward. + This improves sampling-state precision at the cost of additional memory and preserves the existing + model-dtype behavior when disabled. Returns: [`Cosmos3OmniPipelineOutput`] or `tuple`: @@ -1498,6 +1504,7 @@ def __call__( device = self._get_execution_device() dtype = self.transformer.dtype + sampling_dtype = torch.float32 if use_fp32_sampling_state else dtype if enable_safety_check and isinstance(self.safety_checker, CosmosSafetyChecker): self.safety_checker.to(device) @@ -1558,7 +1565,7 @@ def __call__( action_latents=action_latents, generator=generator, device=device, - dtype=dtype, + dtype=sampling_dtype, enable_sound=enable_sound, action=action, ) @@ -1787,6 +1794,14 @@ def __call__( raw_action_dim=raw_action_dim_resolved, ) + if use_fp32_sampling_state: + 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. @@ -1831,6 +1846,8 @@ def __call__( callback_kwargs[key] = locals()[key] callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) latents = callback_outputs.pop("latents", latents) + if use_fp32_sampling_state: + latents = latents.float() if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() 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..11e6900d12a1 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -17,8 +17,12 @@ import torch from PIL import Image -from diffusers import ModularPipeline +from diffusers import FlowMatchEulerDiscreteScheduler, ModularPipeline from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline +from diffusers.modular_pipelines.cosmos.before_denoise import Cosmos3VisionPrepareLatentsStep +from diffusers.modular_pipelines.cosmos.denoise import Cosmos3DistilledVisionLoopSchedulerStep +from diffusers.modular_pipelines.cosmos.encoders import Cosmos3DistilledTextEncoderStep +from diffusers.modular_pipelines.modular_pipeline import BlockState, PipelineState from ...testing_utils import torch_device from ..testing_utils import ( @@ -113,6 +117,67 @@ 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 + assert pipe.config.default_use_fp32_sampling_state is False + @pytest.mark.parametrize( + ("config_enabled", "input_enabled", "expected_dtype"), + [ + (False, None, torch.bfloat16), + (True, None, torch.float32), + (True, False, torch.bfloat16), + ], + ) + def test_prepare_vision_latents_fp32_sampling_state( + self, config_enabled, input_enabled, expected_dtype + ): + components = BlockState( + _execution_device=torch.device("cpu"), + transformer=BlockState(dtype=torch.bfloat16), + config=BlockState(default_use_fp32_sampling_state=config_enabled), + vae_scale_factor_spatial=16, + vae_scale_factor_temporal=4, + num_channels_latents=4, + ) + state = PipelineState() + state.set("x0_tokens_vision", None) + state.set("vision_condition_frames", None) + state.set("num_frames", 5) + state.set("height", 32) + state.set("width", 32) + state.set("fps", 24.0) + state.set("latents", None) + state.set("generator", torch.Generator("cpu").manual_seed(0)) + state.set("use_fp32_sampling_state", input_enabled) + + Cosmos3VisionPrepareLatentsStep()(components, state) + + assert state.get("latents").dtype == expected_dtype + assert state.get("vision_condition_mask").dtype == expected_dtype + assert state.get("use_fp32_sampling_state") is (config_enabled if input_enabled is None else input_enabled) + + @pytest.mark.parametrize( + ("use_fp32_sampling_state", "expected_dtype"), + [(False, torch.bfloat16), (True, torch.float32)], + ) + def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_dtype): + scheduler = FlowMatchEulerDiscreteScheduler(stochastic_sampling=True) + scheduler.set_timesteps(sigmas=[1.0, 0.5]) + latents = torch.zeros((1, 2, 1, 1, 1), dtype=torch.bfloat16) + block_state = BlockState( + latents=latents, + velocity_vision=torch.zeros_like(latents), + vision_condition_mask=torch.zeros((1, 1, 1), dtype=expected_dtype), + vision_conditioning_latents=None, + vision_condition_indexes_for_pack=[], + generator=torch.Generator("cpu").manual_seed(0), + use_fp32_sampling_state=use_fp32_sampling_state, + ) + + Cosmos3DistilledVisionLoopSchedulerStep()( + BlockState(scheduler=scheduler), block_state, i=0, t=scheduler.timesteps[0] + ) + + assert block_state.latents.dtype == expected_dtype 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 86bd371c83c5..9dd8f8d0ae25 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -110,6 +110,37 @@ def test_inference(self): assert video[0].shape == self.output_shape + 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, + use_fp32_sampling_state=True, + 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_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(self): components = self.get_dummy_components() components["default_use_system_prompt"] = False From 5e2e41c8bb464d5b8b40face8e9bf5ae23819e76 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 03:43:50 -0700 Subject: [PATCH 02/20] Add cache execution contexts to Cosmos3 pipelines --- src/diffusers/models/cache_utils.py | 9 +- .../transformers/transformer_cosmos3.py | 3 +- .../modular_pipelines/cosmos/denoise.py | 60 +++++--- .../cosmos/modular_pipeline.py | 24 ++++ .../pipelines/cosmos/pipeline_cosmos3_omni.py | 133 +++++++++++------- .../cosmos/test_modular_pipeline_cosmos3.py | 23 +++ tests/pipelines/cosmos/test_cosmos3.py | 25 ++++ 7 files changed, 202 insertions(+), 75 deletions(-) diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index 3c51eb7ba63a..5e7c632ac180 100644 --- a/src/diffusers/models/cache_utils.py +++ b/src/diffusers/models/cache_utils.py @@ -108,6 +108,7 @@ def enable_cache(self, config) -> None: HookRegistry.check_if_exists_or_initialize(self).invalidate_child_registries_cache() self._cache_config = config + HookRegistry.check_if_exists_or_initialize(self)._child_registries_cache = None def disable_cache(self) -> None: from ..hooks import ( @@ -154,6 +155,7 @@ def disable_cache(self) -> None: registry.invalidate_child_registries_cache() self._cache_config = None + registry._child_registries_cache = None def _reset_stateful_cache(self, recurse: bool = True) -> None: from ..hooks import HookRegistry @@ -168,6 +170,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/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index ee28ea935883..ce6ba0a4bee5 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -1,4 +1,5 @@ import inspect +from contextlib import nullcontext import torch @@ -14,6 +15,11 @@ from .modular_pipeline import Cosmos3OmniModularPipeline +def _cache_context(transformer: torch.nn.Module, name: str): + cache_context = getattr(transformer, "cache_context", None) + return cache_context(name) if callable(cache_context) else nullcontext() + + class Cosmos3VisionLoopPrepareStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @@ -215,7 +221,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 _cache_context(components.transformer, 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, @@ -480,13 +489,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 @@ -718,21 +737,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 _cache_context(components.transformer, 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() @@ -757,7 +777,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 @@ -767,6 +791,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 @@ -776,6 +801,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta uncond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps, + "uncond", ) if needs_control_cfg and needs_text_cfg: diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py index d6c09703c12e..5efba85de64f 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py @@ -18,6 +18,30 @@ 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) + if hasattr(transformer, "_reset_stateful_cache"): + transformer._reset_stateful_cache() + 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 d793484fd8b7..e5d1c630f130 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -16,6 +16,7 @@ import json import math from collections.abc import Iterable +from contextlib import nullcontext from dataclasses import dataclass from typing import Any, Callable, Literal @@ -42,6 +43,11 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +def _cache_context(transformer: torch.nn.Module, name: str): + cache_context = getattr(transformer, "cache_context", None) + return cache_context(name) if callable(cache_context) else nullcontext() + + if is_cosmos_guardrail_available(): from cosmos_guardrail import CosmosSafetyChecker else: @@ -1293,6 +1299,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 @@ -1449,6 +1463,9 @@ def __call__( `sound` (`torch.Tensor` of shape `[C, N]`, or `None` when `enable_sound=False`). Otherwise a tuple `(video, sound)` with the same fields. """ + if hasattr(self.transformer, "_reset_stateful_cache"): + self.transformer._reset_stateful_cache() + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs @@ -1493,6 +1510,8 @@ def __call__( ) self._current_timestep = None + self._current_step_index = None + self._current_sigma = None self._interrupt = False self._guidance_scale = guidance_scale @@ -1699,6 +1718,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 @@ -1717,33 +1738,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 _cache_context(self.transformer, "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, @@ -1757,33 +1779,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 _cache_context(self.transformer, "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, @@ -1853,6 +1876,8 @@ def __call__( 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/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index b5bf5fb04393..2750838897ab 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -13,6 +13,9 @@ # 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 @@ -167,6 +170,26 @@ 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 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/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index 9dd8f8d0ae25..9cd6a8dcdab4 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -12,6 +12,7 @@ # 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 @@ -141,6 +142,30 @@ def callback_on_step_end(_pipeline, _step_index, _timestep, callback_kwargs): 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 From d6998328b8ea3aed66bf0a8e90432b26fcfd347a Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 03:51:52 -0700 Subject: [PATCH 03/20] Add SeaCache support for Cosmos3 --- src/diffusers/__init__.py | 6 + src/diffusers/hooks/__init__.py | 1 + src/diffusers/hooks/_helpers.py | 15 + src/diffusers/hooks/sea_cache.py | 1111 +++++++++++++++++ src/diffusers/models/cache_utils.py | 31 +- src/diffusers/utils/dummy_pt_objects.py | 23 + tests/hooks/test_sea_cache.py | 443 +++++++ .../test_models_transformer_cosmos3.py | 142 ++- 8 files changed, 1770 insertions(+), 2 deletions(-) create mode 100644 src/diffusers/hooks/sea_cache.py create mode 100644 tests/hooks/test_sea_cache.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 3ed956c75e49..687b06486ade 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -204,6 +204,7 @@ "LayerSkipConfig", "MagCacheConfig", "PyramidAttentionBroadcastConfig", + "SeaCacheConfig", "SmoothedEnergyGuidanceConfig", "TaylorSeerCacheConfig", "TextKVCacheConfig", @@ -212,8 +213,10 @@ "apply_layer_skip", "apply_mag_cache", "apply_pyramid_attention_broadcast", + "apply_sea_cache", "apply_taylorseer_cache", "apply_text_kv_cache", + "get_sea_cache_stats", ] ) _import_structure["image_processor"] = [ @@ -1083,6 +1086,7 @@ LayerSkipConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, + SeaCacheConfig, SmoothedEnergyGuidanceConfig, TaylorSeerCacheConfig, TextKVCacheConfig, @@ -1091,8 +1095,10 @@ apply_layer_skip, apply_mag_cache, apply_pyramid_attention_broadcast, + apply_sea_cache, apply_taylorseer_cache, apply_text_kv_cache, + get_sea_cache_stats, ) from .image_processor import ( InpaintProcessor, diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index d999ab32d6d7..2e9c7395c2b1 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, get_sea_cache_stats 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..c5d243a32721 --- /dev/null +++ b/src/diffusers/hooks/sea_cache.py @@ -0,0 +1,1111 @@ +# 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 copy +import inspect +import math +import time +from dataclasses import dataclass +from typing import Any, Callable, Literal, Sequence + +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 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.35`): + Accumulated relative-L1 budget. Larger values reuse the cache more often. + residual_order (`int`, defaults to `0`): + 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. + 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 `"first_block"`): + Feature source used to construct the SEA indicator. `"first_block"` follows the published method and + filters the timestep-modulated pre-attention input of the first transformer block. + `"raw_vision_latents"` is an opt-in lower-cost approximation that filters the complete raw vision latent, + including clean conditioning frames in I2V. 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. + gate_schedule (`Sequence[bool]`, *optional*): + Per-step full-compute decisions to replay for controlled residual-prediction ablations. SEA still evaluates + its natural decision and reports mismatches; invalid or unsafe calls always fail open to full compute. + + 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.35 + residual_order: int = 0 + retention_steps: int = 1 + cache_end_steps: int = 1 + power_exp: float = 3.0 + indicator_source: Literal["first_block", "raw_vision_latents"] = "first_block" + 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 + gate_schedule: Sequence[bool] = 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 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`.") + if self.gate_schedule is not None: + self.gate_schedule = tuple(self.gate_schedule) + if any(not isinstance(value, bool) for value in self.gate_schedule): + raise TypeError("`gate_schedule` must contain only boolean values.") + + +@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.skip_remaining = False + self.full_execution_pending = False + self.cacheable_execution = False + self.step_index: int | None = None + self.und_input: torch.Tensor | 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 + self.full_started_at: float | None = None + + def reset_forward(self): + self.skip_remaining = False + self.full_execution_pending = False + self.cacheable_execution = False + self.step_index = None + self.und_input = None + self.gen_input = None + self.und_output = None + self.cached_und_output = None + self.cached_gen_residual = None + self.full_started_at = None + + def reset(self): + self.history = [] + self.gate_key = None + self.gate_should_compute = True + self.previous_indicator = None + self.accumulated_distance = 0.0 + 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 + self.transformer_calls = 0 + self.gate_evaluations = 0 + self.gate_full_decisions = 0 + self.gate_skip_decisions = 0 + self.gate_trace: list[bool] = [] + self.gate_schedule_mismatches = 0 + self.num_full_steps = 0 + self.num_cached_steps = 0 + self.fail_open_calls = 0 + self.indicator_seconds = 0.0 + self.decision_seconds = 0.0 + self.full_seconds = 0.0 + self.branch_full_executions: dict[str, int] = {} + self.branch_reuses: dict[str, int] = {} + + 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.fail_open_calls += 1 + 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 + + self.gate_evaluations += 1 + 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.") + 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 + 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 + 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 config.gate_schedule is not None: + schedule_is_valid = len(config.gate_schedule) == metadata.num_inference_steps + scheduled_compute = bool(config.gate_schedule[metadata.step_index]) if schedule_is_valid else True + can_replay_skip = not (invalid_gate or is_retained or is_in_cache_end or is_first_observation) + if not schedule_is_valid: + self.mark_fail_open( + "SeaCache gate schedule length does not match the number of inference steps; running full." + ) + should_compute = True + elif scheduled_compute or can_replay_skip: + if natural_should_compute != scheduled_compute: + self.gate_schedule_mismatches += 1 + should_compute = scheduled_compute + else: + self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.") + 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] + self.gate_trace.append(should_compute) + if should_compute: + self.gate_full_decisions += 1 + else: + self.gate_skip_decisions += 1 + 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_manager: StateManager, + shared_state: SeaCacheSharedState, + state: SeaCacheContextState, + gen_output: torch.Tensor, + und_output: torch.Tensor | None, +) -> None: + shared_state.num_full_steps += 1 + branch = state_manager._current_context + shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1 + if state.full_started_at is not None: + shared_state.full_seconds += time.perf_counter() - state.full_started_at + 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, + residual_boundary: str, + ): + 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 + self.residual_boundary = residual_boundary + self._last_stats: dict[str, Any] | None = None + + @property + def num_full_steps(self) -> int: + return self.shared_state.num_full_steps + + @property + def num_cached_steps(self) -> int: + return self.shared_state.num_cached_steps + + def stats(self) -> dict[str, Any]: + if self.shared_state.transformer_calls == 0 and self._last_stats is not None: + return copy.deepcopy(self._last_stats) + + persistent_cache_bytes = 0 + per_branch = {} + for branch, state in sorted(self.state_manager._state_cache.items()): + if state.previous_indicator is not None: + persistent_cache_bytes += sum( + value.numel() * value.element_size() for value in state.previous_indicator + ) + for _, und_output, gen_residual in state.history: + persistent_cache_bytes += und_output.numel() * und_output.element_size() + persistent_cache_bytes += gen_residual.numel() * gen_residual.element_size() + per_branch[branch] = { + "full_calls": self.shared_state.branch_full_executions.get(branch, 0), + "reuse_calls": self.shared_state.branch_reuses.get(branch, 0), + } + + opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps + return { + "indicator_source": self.config.indicator_source, + "residual_order": self.config.residual_order, + "residual_boundary": self.residual_boundary, + "transformer_calls": self.shared_state.transformer_calls, + "gate_evaluations": self.shared_state.gate_evaluations, + "gate_full_decisions": self.shared_state.gate_full_decisions, + "gate_skip_decisions": self.shared_state.gate_skip_decisions, + "gate_trace": list(self.shared_state.gate_trace), + "gate_schedule_replayed": self.config.gate_schedule is not None, + "gate_schedule_mismatches": self.shared_state.gate_schedule_mismatches, + "actual_full_executions": self.shared_state.num_full_steps, + "actual_reuses": self.shared_state.num_cached_steps, + "actual_reuse_rate": (self.shared_state.num_cached_steps / opportunities if opportunities else 0.0), + "fail_open_calls": self.shared_state.fail_open_calls, + "sea_indicator_seconds": self.shared_state.indicator_seconds, + "sea_decision_seconds": self.shared_state.decision_seconds, + "sea_seconds": (self.shared_state.indicator_seconds + self.shared_state.decision_seconds), + "full_seconds": self.shared_state.full_seconds, + "timing_note": "host wall time; CUDA work is asynchronous", + "persistent_cache_bytes": persistent_cache_bytes, + "branch_full_executions": dict(sorted(self.shared_state.branch_full_executions.items())), + "branch_reuses": dict(sorted(self.shared_state.branch_reuses.items())), + "per_branch": per_branch, + } + + def pre_forward(self, module: torch.nn.Module, *args, **kwargs): + self.shared_state.forward_metadata = None + self.shared_state.transformer_calls += 1 + 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): + if self.shared_state.transformer_calls > 0: + self._last_stats = self.stats() + 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 + state.und_input = encoder_hidden_states + + forward_metadata = self.shared_state.forward_metadata + if state is None or forward_metadata is None: + if state is not None: + state.full_started_at = time.perf_counter() + return self.fn_ref.original_forward(*args, **kwargs) + + state.step_index = forward_metadata.step_index + state.cacheable_execution = True + indicator_started = time.perf_counter() + 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 + self.shared_state.indicator_seconds += time.perf_counter() - indicator_started + 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." + ) + decision_started = time.perf_counter() + should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config) + self.shared_state.decision_seconds += time.perf_counter() - decision_started + + 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." + ) + state.full_started_at = time.perf_counter() + 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." + ) + state.full_started_at = time.perf_counter() + 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 + self.shared_state.num_cached_steps += 1 + branch = self.state_manager._current_context + self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 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, + self.state_manager, + self.shared_state, + 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, + self.state_manager, + self.shared_state, + 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 + residual_boundary = "post_language_model_norm" if post_norm_boundary else "repeated_block_stack" + + 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, + residual_boundary, + ), + _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 + + +def get_sea_cache_stats(module: torch.nn.Module) -> dict[str, Any]: + """Return statistics for the SeaCache instance currently attached to ``module``.""" + + registry = getattr(module, "_diffusers_hook", None) + root_hook = registry.get_hook(_SEA_CACHE_ROOT_HOOK) if registry is not None else None + if not isinstance(root_hook, SeaCacheRootHook): + raise ValueError("SeaCache is not enabled on this module.") + return root_hook.stats() diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index 5e7c632ac180..9a62ace4192c 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: @@ -72,12 +74,14 @@ def enable_cache(self, config) -> None: 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, ) @@ -97,6 +101,8 @@ 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: @@ -117,6 +123,7 @@ def disable_cache(self) -> None: HookRegistry, MagCacheConfig, PyramidAttentionBroadcastConfig, + SeaCacheConfig, TaylorSeerCacheConfig, TextKVCacheConfig, ) @@ -124,6 +131,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 @@ -146,6 +159,11 @@ 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: @@ -157,6 +175,17 @@ def disable_cache(self) -> None: self._cache_config = None registry._child_registries_cache = None + def get_cache_stats(self) -> dict: + """Return statistics for the currently enabled cache implementation.""" + + from ..hooks import SeaCacheConfig, get_sea_cache_stats + + if self._cache_config is None: + raise ValueError("Caching is not enabled.") + if isinstance(self._cache_config, SeaCacheConfig): + return get_sea_cache_stats(self) + raise NotImplementedError(f"Cache statistics are not available for {type(self._cache_config).__name__}.") + def _reset_stateful_cache(self, recurse: bool = True) -> None: from ..hooks import HookRegistry diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 1598814f835a..324c92f3d507 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"]) @@ -345,6 +364,10 @@ def apply_text_kv_cache(*args, **kwargs): requires_backends(apply_text_kv_cache, ["torch"]) +def get_sea_cache_stats(*args, **kwargs): + requires_backends(get_sea_cache_stats, ["torch"]) + + class InpaintProcessor(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py new file mode 100644 index 000000000000..85db1fbc22c1 --- /dev/null +++ b/tests/hooks/test_sea_cache.py @@ -0,0 +1,443 @@ +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import SeaCacheConfig +from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry +from diffusers.hooks.hooks import HookRegistry, ModelHook +from diffusers.hooks.sea_cache import ( + _SEA_CACHE_BLOCK_HOOK, + _SEA_CACHE_LEADER_BLOCK_HOOK, + _SEA_CACHE_ROOT_HOOK, + _apply_sea_filter, +) +from diffusers.models.cache_utils import CacheMixin + + +class CountingIdentity(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, hidden_states): + self.calls += 1 + return hidden_states + + +class DummySeaBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.indicator_norm = CountingIdentity() + self.calls = 0 + + def forward(self, und_seq, gen_seq, rotary_emb=None): + self.calls += 1 + return und_seq + 1, gen_seq * 2 + + +class DummySeaTransformer(torch.nn.Module, CacheMixin): + def __init__(self, num_layers=3): + super().__init__() + self.layers = torch.nn.ModuleList([DummySeaBlock() for _ in range(num_layers)]) + + def forward(self, hidden_states): + shape = hidden_states.shape + gen_seq = hidden_states.reshape(-1, shape[-1]) + und_seq = torch.zeros(1, shape[-1], device=hidden_states.device, dtype=hidden_states.dtype) + for layer in self.layers: + und_seq, gen_seq = layer(und_seq, gen_seq) + return und_seq, gen_seq.reshape(shape) + + +@pytest.fixture(autouse=True) +def register_dummy_sea_block(): + TransformerBlockRegistry.register( + DummySeaBlock, + 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="indicator_norm", + ), + ) + + +def _metadata_callback(module, args, kwargs): + hidden_states = kwargs.get("hidden_states", args[0] if args else None) + if hidden_states is None: + return None + temporal, height, width = hidden_states.shape[:3] + indexes = torch.arange(temporal * height * width, device=hidden_states.device) + return [(indexes, (temporal, height, width))] + + +def _make_config(runtime, **kwargs): + config_kwargs = { + "threshold": 100.0, + "retention_steps": 1, + "cache_end_steps": 0, + "current_step_callback": lambda: runtime["step"], + "current_sigma_callback": lambda: runtime["sigma"], + "num_inference_steps_callback": lambda: runtime["num_steps"], + "metadata_callback": _metadata_callback, + } + config_kwargs.update(kwargs) + return SeaCacheConfig(**config_kwargs) + + +def _get_root_hook(model): + return model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) + + +@torch.no_grad() +def test_sea_cache_uses_independent_gates_and_histories_per_context(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} + model = DummySeaTransformer() + model.enable_cache(_make_config(runtime, threshold=0.5)) + root_hook = _get_root_hook(model) + + first_input = torch.ones(2, 2, 2, 1) + with model.cache_context("cond"): + _, first_output = model(first_input) + torch.testing.assert_close(first_output, first_input * 8) + assert [block.calls for block in model.layers] == [1, 1, 1] + + with model.cache_context("uncond"): + _, first_uncond_output = model(first_input) + torch.testing.assert_close(first_uncond_output, first_input * 8) + assert [block.calls for block in model.layers] == [2, 2, 2] + + runtime.update(step=1) + with model.cache_context("cond"): + _, cached_output = model(first_input) + torch.testing.assert_close(cached_output, first_input * 8) + assert [block.calls for block in model.layers] == [2, 2, 2] + + changed_uncond_input = first_input * 100 + with model.cache_context("uncond"): + _, uncond_output = model(changed_uncond_input) + torch.testing.assert_close(uncond_output, changed_uncond_input * 8) + assert [block.calls for block in model.layers] == [3, 3, 3] + assert model.layers[0].indicator_norm.calls == 4 + assert root_hook.num_full_steps == 3 + assert root_hook.num_cached_steps == 1 + stats = model.get_cache_stats() + assert stats["indicator_source"] == "first_block" + assert stats["residual_order"] == 0 + assert stats["residual_boundary"] == "repeated_block_stack" + assert stats["transformer_calls"] == 4 + assert stats["gate_evaluations"] == 4 + assert stats["gate_trace"] == [True, True, False, True] + assert stats["branch_full_executions"] == {"cond": 1, "uncond": 2} + assert stats["branch_reuses"] == {"cond": 1} + assert stats["persistent_cache_bytes"] > 0 + + model._reset_stateful_cache() + assert root_hook.num_full_steps == 0 + assert root_hook.num_cached_steps == 0 + archived_stats = model.get_cache_stats() + assert archived_stats["transformer_calls"] == 4 + assert archived_stats["gate_trace"] == [True, True, False, True] + + model.disable_cache() + assert not model.is_cache_enabled + + +@torch.no_grad() +def test_sea_cache_raw_vision_indicator_includes_conditioning_frames(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + model = DummySeaTransformer() + + def raw_vision(module, args, kwargs): + hidden_states = kwargs.get("hidden_states", args[0] if args else None) + latent = hidden_states.permute(3, 0, 1, 2) + return [latent] + + model.enable_cache( + _make_config( + runtime, + threshold=1e-6, + indicator_source="raw_vision_latents", + raw_vision_callback=raw_vision, + ) + ) + + first_input = torch.tensor([1.0, 2.0]).reshape(2, 1, 1, 1) + with model.cache_context("cond"): + model(first_input) + + # The noisy frame stays fixed, but changing the conditioning frame changes the complete raw-latent indicator. + runtime.update(step=1, sigma=0.9) + second_input = torch.tensor([100.0, 2.0]).reshape(2, 1, 1, 1) + with model.cache_context("cond"): + model(second_input) + + stats = model.get_cache_stats() + assert stats["indicator_source"] == "raw_vision_latents" + assert stats["gate_trace"] == [True, True] + assert stats["actual_full_executions"] == 2 + assert stats["actual_reuses"] == 0 + assert model.layers[0].indicator_norm.calls == 0 + + +@torch.no_grad() +def test_sea_cache_residual_order_one_uses_actual_full_step_history(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} + config = _make_config(runtime, residual_order=1, threshold=0.0) + model = DummySeaTransformer() + model.enable_cache(config) + + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) + + runtime.update(step=1, sigma=0.7) + with model.cache_context("cond"): + model(torch.full((1, 1, 1, 1), 2.0)) + + config.threshold = 100.0 + runtime.update(step=2, sigma=0.5) + with model.cache_context("cond"): + _, output = model(torch.full((1, 1, 1, 1), 3.0)) + + # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2. + torch.testing.assert_close(output, torch.full_like(output, 24.0)) + root_hook = _get_root_hook(model) + assert root_hook.num_full_steps == 2 + assert root_hook.num_cached_steps == 1 + + +@torch.no_grad() +def test_sea_cache_replays_gate_schedule_and_reports_natural_mismatches(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} + model = DummySeaTransformer() + model.enable_cache( + _make_config( + runtime, + threshold=0.0, + gate_schedule=(True, False, True), + ) + ) + + for step, sigma in enumerate((0.9, 0.6, 0.3)): + runtime.update(step=step, sigma=sigma) + with model.cache_context("cond"): + model(torch.full((1, 1, 1, 1), float(step + 1))) + + stats = model.get_cache_stats() + assert stats["gate_trace"] == [True, False, True] + assert stats["gate_schedule_replayed"] + assert stats["gate_schedule_mismatches"] == 1 + assert stats["actual_full_executions"] == 2 + assert stats["actual_reuses"] == 1 + + +@torch.no_grad() +def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_uncached_call(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + model = DummySeaTransformer() + + # Pipeline cache contexts may be entered before a cache is enabled (for example, during the baseline). + with model.cache_context("baseline"): + model(torch.ones(1, 1, 1, 1)) + + model.enable_cache(_make_config(runtime)) + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) + + assert model.get_cache_stats()["branch_full_executions"] == {"cond": 1} + + +@torch.no_grad() +def test_sea_cache_fails_open_without_vision_metadata(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + config = _make_config(runtime) + config.metadata_callback = lambda module, args, kwargs: None + model = DummySeaTransformer() + model.enable_cache(config) + + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) + runtime.update(step=1, sigma=0.5) + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) + + assert [block.calls for block in model.layers] == [2, 2, 2] + root_hook = _get_root_hook(model) + assert root_hook.num_full_steps == 2 + assert root_hook.num_cached_steps == 0 + + +@torch.no_grad() +def test_sea_cache_fails_open_for_non_adjacent_steps_and_shape_changes(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} + model = DummySeaTransformer() + model.enable_cache(_make_config(runtime, residual_order=1)) + + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) + + runtime.update(step=2, sigma=0.5) + with model.cache_context("cond"): + model(torch.ones(1, 2, 1, 1)) + + runtime.update(step=3, sigma=0.2) + with model.cache_context("cond"): + model(torch.ones(1, 2, 1, 1)) + + root_hook = _get_root_hook(model) + assert root_hook.num_full_steps == 3 + assert root_hook.num_cached_steps == 0 + assert [block.calls for block in model.layers] == [3, 3, 3] + + +def test_sea_cache_is_inference_only_and_fails_open_with_autograd(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + model = DummySeaTransformer() + model.enable_cache(_make_config(runtime)) + + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1, requires_grad=True)) + runtime.update(step=1, sigma=0.5) + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1, requires_grad=True)) + + root_hook = _get_root_hook(model) + assert root_hook.num_full_steps == 2 + assert root_hook.num_cached_steps == 0 + + +@torch.no_grad() +def test_sea_filter_density_normalizes_gain_to_unit_mean(): + impulse = torch.zeros(2, 3, 4, 1) + impulse[0, 0, 0, 0] = 1 + + filtered = _apply_sea_filter(impulse, sigma=0.5, power_exp=3.0) + recovered_gain = torch.fft.fftn(filtered.float(), dim=(0, 1, 2)) + + torch.testing.assert_close( + recovered_gain.real.mean(), + torch.tensor(1.0), + atol=1e-5, + rtol=1e-5, + ) + torch.testing.assert_close( + recovered_gain.imag, + torch.zeros_like(recovered_gain.imag), + atol=1e-5, + rtol=0, + ) + + +@pytest.mark.parametrize("sigma", [0.0, 1.0]) +def test_sea_filter_is_finite_at_scheduler_endpoints(sigma): + filtered = _apply_sea_filter(torch.randn(2, 2, 2, 4), sigma=sigma, power_exp=2.0) + + assert torch.isfinite(filtered).all() + assert filtered.abs().sum() > 0 + + +@torch.no_grad() +def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables_cleanly(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + model = DummySeaTransformer(num_layers=1) + model.enable_cache(_make_config(runtime)) + + with model.cache_context("cond"): + _, first_output = model(torch.ones(1, 1, 1, 1)) + torch.testing.assert_close(first_output, torch.full_like(first_output, 2.0)) + + runtime.update(step=1, sigma=0.5) + with model.cache_context("cond"): + _, cached_output = model(torch.full((1, 1, 1, 1), 2.0)) + torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0)) + assert model.layers[0].calls == 1 + assert model.get_cache_stats()["actual_full_executions"] == 1 + assert model.get_cache_stats()["actual_reuses"] == 1 + + model.disable_cache() + + _, output = model(torch.ones(1, 1, 1, 1)) + + torch.testing.assert_close(output, torch.full_like(output, 2.0)) + assert model.layers[0].calls == 2 + assert model.layers[0]._diffusers_hook.hooks == {} + + +@torch.no_grad() +def test_sea_cache_fails_open_for_parameter_sharded_blocks(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + model = DummySeaTransformer() + model.layers[0]._get_fsdp_state = lambda: object() + model.enable_cache(_make_config(runtime)) + + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) + runtime.update(step=1, sigma=0.5) + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) + + stats = model.get_cache_stats() + assert stats["actual_full_executions"] == 2 + assert stats["actual_reuses"] == 0 + assert stats["fail_open_calls"] == 2 + assert [block.calls for block in model.layers] == [2, 2, 2] + + +def test_sea_cache_failed_enable_rolls_back_only_new_hooks(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + model = DummySeaTransformer() + existing_hook = ModelHook() + middle_registry = HookRegistry.check_if_exists_or_initialize(model.layers[1]) + middle_registry.register_hook(existing_hook, _SEA_CACHE_BLOCK_HOOK) + + with pytest.raises(ValueError, match="already exists"): + model.enable_cache(_make_config(runtime)) + + assert not model.is_cache_enabled + assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is None + assert model.layers[0]._diffusers_hook.get_hook(_SEA_CACHE_LEADER_BLOCK_HOOK) is None + assert middle_registry.get_hook(_SEA_CACHE_BLOCK_HOOK) is existing_hook + assert not hasattr(model.layers[2], "_diffusers_hook") + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"threshold": -1.0}, "threshold"), + ({"residual_order": 2}, "residual_order"), + ({"retention_steps": -1}, "retention_steps"), + ({"cache_end_steps": -1}, "cache_end_steps"), + ({"power_exp": 0.0}, "power_exp"), + ({"indicator_source": "raw"}, "indicator_source"), + ({"threshold": float("nan")}, "threshold"), + ({"power_exp": float("inf")}, "power_exp"), + ], +) +def test_sea_cache_config_validation(kwargs, message): + with pytest.raises(ValueError, match=message): + SeaCacheConfig(**kwargs) + + +def test_sea_cache_gate_schedule_validation(): + with pytest.raises(TypeError, match="gate_schedule"): + SeaCacheConfig(gate_schedule=(True, 1)) + + +@pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"]) +def test_sea_cache_callback_validation(callback_name): + with pytest.raises(TypeError, match=callback_name): + SeaCacheConfig(**{callback_name: 1}) diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index 0ead8cb1a6bb..977ea6b52a1a 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 @@ -103,6 +107,142 @@ 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) + + root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) + # 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 + assert root_hook.num_full_steps == 1 + assert root_hook.num_cached_steps == 1 + stats = model.get_cache_stats() + assert stats["indicator_source"] == indicator_source + assert stats["residual_boundary"] == "post_language_model_norm" + + 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[0].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].input_layernorm_moe_gen.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 + assert root_hook.num_full_steps == 1 + assert root_hook.num_cached_steps == 1 + + 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() From cabada1b8f08c7f5491943ee099b4f5030ead41a Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 04:16:09 -0700 Subject: [PATCH 04/20] Cap consecutive SeaCache reuses --- src/diffusers/hooks/sea_cache.py | 41 +++++++++++++++++++++++++++++--- tests/hooks/test_sea_cache.py | 41 ++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index c5d243a32721..264ba4545059 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -47,7 +47,7 @@ class SeaCacheConfig: prediction heads still execute. Args: - threshold (`float`, defaults to `0.35`): + threshold (`float`, defaults to `0.25`): Accumulated relative-L1 budget. Larger values reuse the cache more often. residual_order (`int`, defaults to `0`): Order used to predict the generation-stream language-model residual. `0` directly reuses the most recent @@ -56,6 +56,9 @@ class SeaCacheConfig: 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 `"first_block"`): @@ -96,10 +99,11 @@ class SeaCacheConfig: ``` """ - threshold: float = 0.35 + threshold: float = 0.25 residual_order: int = 0 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"] = "first_block" current_step_callback: Callable[[], int] = None @@ -124,6 +128,14 @@ def __post_init__(self): 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"): @@ -162,6 +174,9 @@ def __init__(self): self.gate_should_compute = True self.previous_indicator: list[torch.Tensor] | None = None self.accumulated_distance = 0.0 + self.consecutive_cached = 0 + self.max_consecutive_cached_observed = 0 + self.max_consecutive_forced_full = 0 self.skip_remaining = False self.full_execution_pending = False self.cacheable_execution = False @@ -191,6 +206,9 @@ def reset(self): self.gate_should_compute = True self.previous_indicator = None self.accumulated_distance = 0.0 + self.consecutive_cached = 0 + self.max_consecutive_cached_observed = 0 + self.max_consecutive_forced_full = 0 self.reset_forward() @@ -243,8 +261,11 @@ def resolve_gate( 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 + 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: @@ -294,6 +315,10 @@ def resolve_gate( self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.") should_compute = True + if is_max_consecutive: + should_compute = True + state.max_consecutive_forced_full += 1 + state.accumulated_distance = 0.0 if should_compute else candidate_accumulated_distance state.gate_key = gate_key @@ -356,6 +381,7 @@ def _record_full_execution( und_output: torch.Tensor | None, ) -> None: shared_state.num_full_steps += 1 + state.consecutive_cached = 0 branch = state_manager._current_context shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1 if state.full_started_at is not None: @@ -579,6 +605,8 @@ def stats(self) -> dict[str, Any]: per_branch[branch] = { "full_calls": self.shared_state.branch_full_executions.get(branch, 0), "reuse_calls": self.shared_state.branch_reuses.get(branch, 0), + "max_consecutive_cached_observed": state.max_consecutive_cached_observed, + "max_consecutive_forced_full": state.max_consecutive_forced_full, } opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps @@ -586,6 +614,11 @@ def stats(self) -> dict[str, Any]: "indicator_source": self.config.indicator_source, "residual_order": self.config.residual_order, "residual_boundary": self.residual_boundary, + "max_consecutive_cached": self.config.max_consecutive_cached, + "max_consecutive_cached_observed": max( + (stats["max_consecutive_cached_observed"] for stats in per_branch.values()), default=0 + ), + "max_consecutive_forced_full": sum(stats["max_consecutive_forced_full"] for stats in per_branch.values()), "transformer_calls": self.shared_state.transformer_calls, "gate_evaluations": self.shared_state.gate_evaluations, "gate_full_decisions": self.shared_state.gate_full_decisions, @@ -870,6 +903,8 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): state.full_execution_pending = False state.cached_und_output = cached_und state.cached_gen_residual = cached_residual + state.consecutive_cached += 1 + state.max_consecutive_cached_observed = max(state.max_consecutive_cached_observed, state.consecutive_cached) self.shared_state.num_cached_steps += 1 branch = self.state_manager._current_context self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 1 diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py index 85db1fbc22c1..d6b3798d548a 100644 --- a/tests/hooks/test_sea_cache.py +++ b/tests/hooks/test_sea_cache.py @@ -157,6 +157,37 @@ def test_sea_cache_uses_independent_gates_and_histories_per_context(): assert not model.is_cache_enabled +@torch.no_grad() +def test_sea_cache_max_consecutive_cached_forces_full_per_context(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 8} + model = DummySeaTransformer() + model.enable_cache( + _make_config( + runtime, + threshold=100.0, + retention_steps=0, + cache_end_steps=0, + max_consecutive_cached=2, + ) + ) + + for step in range(runtime["num_steps"]): + runtime.update(step=step, sigma=0.9 - step * 0.1) + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) + + stats = model.get_cache_stats() + assert stats["gate_trace"] == [True, False, False, True, False, False, True, False] + assert stats["actual_full_executions"] == 3 + assert stats["actual_reuses"] == 5 + assert stats["max_consecutive_cached"] == 2 + assert stats["max_consecutive_cached_observed"] == 2 + assert stats["max_consecutive_forced_full"] == 2 + assert stats["per_branch"]["cond"]["max_consecutive_cached_observed"] == 2 + assert stats["per_branch"]["cond"]["max_consecutive_forced_full"] == 2 + assert [block.calls for block in model.layers] == [3, 3, 3] + + @torch.no_grad() def test_sea_cache_raw_vision_indicator_includes_conditioning_frames(): runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} @@ -414,6 +445,13 @@ def test_sea_cache_failed_enable_rolls_back_only_new_hooks(): assert not hasattr(model.layers[2], "_diffusers_hook") +def test_sea_cache_config_defaults(): + config = SeaCacheConfig() + + assert config.threshold == 0.25 + assert config.max_consecutive_cached == 2 + + @pytest.mark.parametrize( ("kwargs", "message"), [ @@ -421,6 +459,9 @@ def test_sea_cache_failed_enable_rolls_back_only_new_hooks(): ({"residual_order": 2}, "residual_order"), ({"retention_steps": -1}, "retention_steps"), ({"cache_end_steps": -1}, "cache_end_steps"), + ({"max_consecutive_cached": -1}, "max_consecutive_cached"), + ({"max_consecutive_cached": 1.5}, "max_consecutive_cached"), + ({"max_consecutive_cached": True}, "max_consecutive_cached"), ({"power_exp": 0.0}, "power_exp"), ({"indicator_source": "raw"}, "indicator_source"), ({"threshold": float("nan")}, "threshold"), From a9524981f3ea8e98adbd56c510a03ee6398e35ea Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 04:47:29 -0700 Subject: [PATCH 05/20] Align SeaCache defaults with Cosmos3 inference --- src/diffusers/__init__.py | 2 - src/diffusers/hooks/__init__.py | 2 +- src/diffusers/hooks/sea_cache.py | 186 ++---------------- src/diffusers/models/cache_utils.py | 11 -- src/diffusers/utils/dummy_pt_objects.py | 4 - tests/hooks/test_sea_cache.py | 95 +-------- .../test_models_transformer_cosmos3.py | 12 +- 7 files changed, 24 insertions(+), 288 deletions(-) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 687b06486ade..d6ece77742a0 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -216,7 +216,6 @@ "apply_sea_cache", "apply_taylorseer_cache", "apply_text_kv_cache", - "get_sea_cache_stats", ] ) _import_structure["image_processor"] = [ @@ -1098,7 +1097,6 @@ apply_sea_cache, apply_taylorseer_cache, apply_text_kv_cache, - get_sea_cache_stats, ) from .image_processor import ( InpaintProcessor, diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 2e9c7395c2b1..70399f9ce805 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -25,7 +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, get_sea_cache_stats + 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/sea_cache.py b/src/diffusers/hooks/sea_cache.py index 264ba4545059..827648709518 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -12,12 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import inspect import math -import time from dataclasses import dataclass -from typing import Any, Callable, Literal, Sequence +from typing import Any, Callable, Literal import torch @@ -49,7 +47,7 @@ class SeaCacheConfig: Args: threshold (`float`, defaults to `0.25`): Accumulated relative-L1 budget. Larger values reuse the cache more often. - residual_order (`int`, defaults to `0`): + 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`): @@ -61,12 +59,11 @@ class SeaCacheConfig: 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 `"first_block"`): - Feature source used to construct the SEA indicator. `"first_block"` follows the published method and - filters the timestep-modulated pre-attention input of the first transformer block. - `"raw_vision_latents"` is an opt-in lower-cost approximation that filters the complete raw vision latent, - including clean conditioning frames in I2V. Thresholds are not generally transferable between the two - sources. + 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]`): @@ -80,10 +77,6 @@ class SeaCacheConfig: 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. - gate_schedule (`Sequence[bool]`, *optional*): - Per-step full-compute decisions to replay for controlled residual-prediction ablations. SEA still evaluates - its natural decision and reports mismatches; invalid or unsafe calls always fail open to full compute. - Example: ```python >>> from diffusers import Cosmos3OmniPipeline, SeaCacheConfig @@ -100,12 +93,12 @@ class SeaCacheConfig: """ threshold: float = 0.25 - residual_order: int = 0 + 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"] = "first_block" + 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 @@ -117,7 +110,6 @@ class SeaCacheConfig: [torch.nn.Module, tuple[Any, ...], dict[str, Any]], list[torch.Tensor] | None, ] = None - gate_schedule: Sequence[bool] = None def __post_init__(self): if not math.isfinite(self.threshold) or self.threshold < 0: @@ -152,10 +144,6 @@ def __post_init__(self): callback = getattr(self, name) if callback is not None and not callable(callback): raise TypeError(f"`{name}` must be callable or `None`.") - if self.gate_schedule is not None: - self.gate_schedule = tuple(self.gate_schedule) - if any(not isinstance(value, bool) for value in self.gate_schedule): - raise TypeError("`gate_schedule` must contain only boolean values.") @dataclass @@ -175,8 +163,6 @@ def __init__(self): self.previous_indicator: list[torch.Tensor] | None = None self.accumulated_distance = 0.0 self.consecutive_cached = 0 - self.max_consecutive_cached_observed = 0 - self.max_consecutive_forced_full = 0 self.skip_remaining = False self.full_execution_pending = False self.cacheable_execution = False @@ -186,7 +172,6 @@ def __init__(self): self.und_output: torch.Tensor | None = None self.cached_und_output: torch.Tensor | None = None self.cached_gen_residual: torch.Tensor | None = None - self.full_started_at: float | None = None def reset_forward(self): self.skip_remaining = False @@ -198,7 +183,6 @@ def reset_forward(self): self.und_output = None self.cached_und_output = None self.cached_gen_residual = None - self.full_started_at = None def reset(self): self.history = [] @@ -207,8 +191,6 @@ def reset(self): self.previous_indicator = None self.accumulated_distance = 0.0 self.consecutive_cached = 0 - self.max_consecutive_cached_observed = 0 - self.max_consecutive_forced_full = 0 self.reset_forward() @@ -219,20 +201,6 @@ def __init__(self): def reset(self): self.forward_metadata: _SeaCacheForwardMetadata | None = None - self.transformer_calls = 0 - self.gate_evaluations = 0 - self.gate_full_decisions = 0 - self.gate_skip_decisions = 0 - self.gate_trace: list[bool] = [] - self.gate_schedule_mismatches = 0 - self.num_full_steps = 0 - self.num_cached_steps = 0 - self.fail_open_calls = 0 - self.indicator_seconds = 0.0 - self.decision_seconds = 0.0 - self.full_seconds = 0.0 - self.branch_full_executions: dict[str, int] = {} - self.branch_reuses: dict[str, int] = {} def warn_once(self, message: str): if message not in self._warned_messages: @@ -240,7 +208,6 @@ def warn_once(self, message: str): self._warned_messages.add(message) def mark_fail_open(self, message: str): - self.fail_open_calls += 1 self.warn_once(message) def resolve_gate( @@ -254,7 +221,6 @@ def resolve_gate( if state.gate_key == gate_key: return state.gate_should_compute - self.gate_evaluations += 1 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.") @@ -269,7 +235,7 @@ def resolve_gate( candidate_accumulated_distance = 0.0 if forced_compute: - natural_should_compute = True + should_compute = True else: if len(indicator) != len(state.previous_indicator) or not indicator: distance = float("inf") @@ -294,41 +260,13 @@ def resolve_gate( 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 config.gate_schedule is not None: - schedule_is_valid = len(config.gate_schedule) == metadata.num_inference_steps - scheduled_compute = bool(config.gate_schedule[metadata.step_index]) if schedule_is_valid else True - can_replay_skip = not (invalid_gate or is_retained or is_in_cache_end or is_first_observation) - if not schedule_is_valid: - self.mark_fail_open( - "SeaCache gate schedule length does not match the number of inference steps; running full." - ) - should_compute = True - elif scheduled_compute or can_replay_skip: - if natural_should_compute != scheduled_compute: - self.gate_schedule_mismatches += 1 - should_compute = scheduled_compute - else: - self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.") - should_compute = True - - if is_max_consecutive: - should_compute = True - state.max_consecutive_forced_full += 1 + should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold 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] - self.gate_trace.append(should_compute) - if should_compute: - self.gate_full_decisions += 1 - else: - self.gate_skip_decisions += 1 return should_compute @@ -374,18 +312,11 @@ def _get_block_outputs( def _record_full_execution( config: SeaCacheConfig, - state_manager: StateManager, - shared_state: SeaCacheSharedState, state: SeaCacheContextState, gen_output: torch.Tensor, und_output: torch.Tensor | None, ) -> None: - shared_state.num_full_steps += 1 state.consecutive_cached = 0 - branch = state_manager._current_context - shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1 - if state.full_started_at is not None: - shared_state.full_seconds += time.perf_counter() - state.full_started_at if ( state.cacheable_execution and state.step_index is not None @@ -569,7 +500,6 @@ def __init__( shared_state: SeaCacheSharedState, metadata_callback: Callable, raw_vision_callback: Callable, - residual_boundary: str, ): super().__init__() self.config = config @@ -577,73 +507,9 @@ def __init__( self.shared_state = shared_state self.metadata_callback = metadata_callback self.raw_vision_callback = raw_vision_callback - self.residual_boundary = residual_boundary - self._last_stats: dict[str, Any] | None = None - - @property - def num_full_steps(self) -> int: - return self.shared_state.num_full_steps - - @property - def num_cached_steps(self) -> int: - return self.shared_state.num_cached_steps - - def stats(self) -> dict[str, Any]: - if self.shared_state.transformer_calls == 0 and self._last_stats is not None: - return copy.deepcopy(self._last_stats) - - persistent_cache_bytes = 0 - per_branch = {} - for branch, state in sorted(self.state_manager._state_cache.items()): - if state.previous_indicator is not None: - persistent_cache_bytes += sum( - value.numel() * value.element_size() for value in state.previous_indicator - ) - for _, und_output, gen_residual in state.history: - persistent_cache_bytes += und_output.numel() * und_output.element_size() - persistent_cache_bytes += gen_residual.numel() * gen_residual.element_size() - per_branch[branch] = { - "full_calls": self.shared_state.branch_full_executions.get(branch, 0), - "reuse_calls": self.shared_state.branch_reuses.get(branch, 0), - "max_consecutive_cached_observed": state.max_consecutive_cached_observed, - "max_consecutive_forced_full": state.max_consecutive_forced_full, - } - - opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps - return { - "indicator_source": self.config.indicator_source, - "residual_order": self.config.residual_order, - "residual_boundary": self.residual_boundary, - "max_consecutive_cached": self.config.max_consecutive_cached, - "max_consecutive_cached_observed": max( - (stats["max_consecutive_cached_observed"] for stats in per_branch.values()), default=0 - ), - "max_consecutive_forced_full": sum(stats["max_consecutive_forced_full"] for stats in per_branch.values()), - "transformer_calls": self.shared_state.transformer_calls, - "gate_evaluations": self.shared_state.gate_evaluations, - "gate_full_decisions": self.shared_state.gate_full_decisions, - "gate_skip_decisions": self.shared_state.gate_skip_decisions, - "gate_trace": list(self.shared_state.gate_trace), - "gate_schedule_replayed": self.config.gate_schedule is not None, - "gate_schedule_mismatches": self.shared_state.gate_schedule_mismatches, - "actual_full_executions": self.shared_state.num_full_steps, - "actual_reuses": self.shared_state.num_cached_steps, - "actual_reuse_rate": (self.shared_state.num_cached_steps / opportunities if opportunities else 0.0), - "fail_open_calls": self.shared_state.fail_open_calls, - "sea_indicator_seconds": self.shared_state.indicator_seconds, - "sea_decision_seconds": self.shared_state.decision_seconds, - "sea_seconds": (self.shared_state.indicator_seconds + self.shared_state.decision_seconds), - "full_seconds": self.shared_state.full_seconds, - "timing_note": "host wall time; CUDA work is asynchronous", - "persistent_cache_bytes": persistent_cache_bytes, - "branch_full_executions": dict(sorted(self.shared_state.branch_full_executions.items())), - "branch_reuses": dict(sorted(self.shared_state.branch_reuses.items())), - "per_branch": per_branch, - } def pre_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.forward_metadata = None - self.shared_state.transformer_calls += 1 if torch.is_grad_enabled(): self.shared_state.mark_fail_open( "SeaCache is inference-only; calls with autograd enabled run in fail-open mode." @@ -742,8 +608,6 @@ def post_forward(self, module: torch.nn.Module, output: Any) -> Any: return output def reset_state(self, module: torch.nn.Module): - if self.shared_state.transformer_calls > 0: - self._last_stats = self.stats() self.state_manager.reset() self.shared_state.reset() return module @@ -829,13 +693,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): forward_metadata = self.shared_state.forward_metadata if state is None or forward_metadata is None: - if state is not None: - state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) state.step_index = forward_metadata.step_index state.cacheable_execution = True - indicator_started = time.perf_counter() indicator_error_reported = False if _is_parameter_sharded(module): self.shared_state.mark_fail_open( @@ -852,14 +713,11 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): ) indicator = None indicator_error_reported = True - self.shared_state.indicator_seconds += time.perf_counter() - indicator_started 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." ) - decision_started = time.perf_counter() should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config) - self.shared_state.decision_seconds += time.perf_counter() - decision_started if should_compute or not state.history: if not should_compute: @@ -867,7 +725,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.mark_fail_open( "SeaCache selected a cache hit without residual history; running in fail-open mode." ) - state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) residual_history = state.history[-(self.config.residual_order + 1) :] @@ -889,7 +746,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.mark_fail_open( "SeaCache residual history changed shape, device, or dtype; running in fail-open mode." ) - state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) if self.config.residual_order == 1 and len(residual_history) >= 2: @@ -904,10 +760,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): state.cached_und_output = cached_und state.cached_gen_residual = cached_residual state.consecutive_cached += 1 - state.max_consecutive_cached_observed = max(state.max_consecutive_cached_observed, state.consecutive_cached) - self.shared_state.num_cached_steps += 1 - branch = self.state_manager._current_context - self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 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) @@ -952,8 +804,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output) _record_full_execution( self.config, - self.state_manager, - self.shared_state, state, gen_output=hidden_states, und_output=encoder_hidden_states, @@ -1006,8 +856,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): _record_full_execution( self.config, - self.state_manager, - self.shared_state, state, gen_output=output, und_output=state.und_output, @@ -1055,7 +903,6 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None: "residual boundary." ) post_norm_boundary = post_norm_modules is not None - residual_boundary = "post_language_model_norm" if post_norm_boundary else "repeated_block_stack" blocks = [] for name, submodule in unwrapped_module.named_children(): @@ -1085,7 +932,6 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: shared_state, metadata_callback, raw_vision_callback, - residual_boundary, ), _SEA_CACHE_ROOT_HOOK, ) @@ -1134,13 +980,3 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: raise root_registry._child_registries_cache = None - - -def get_sea_cache_stats(module: torch.nn.Module) -> dict[str, Any]: - """Return statistics for the SeaCache instance currently attached to ``module``.""" - - registry = getattr(module, "_diffusers_hook", None) - root_hook = registry.get_hook(_SEA_CACHE_ROOT_HOOK) if registry is not None else None - if not isinstance(root_hook, SeaCacheRootHook): - raise ValueError("SeaCache is not enabled on this module.") - return root_hook.stats() diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index 9a62ace4192c..c31971384cb8 100644 --- a/src/diffusers/models/cache_utils.py +++ b/src/diffusers/models/cache_utils.py @@ -175,17 +175,6 @@ def disable_cache(self) -> None: self._cache_config = None registry._child_registries_cache = None - def get_cache_stats(self) -> dict: - """Return statistics for the currently enabled cache implementation.""" - - from ..hooks import SeaCacheConfig, get_sea_cache_stats - - if self._cache_config is None: - raise ValueError("Caching is not enabled.") - if isinstance(self._cache_config, SeaCacheConfig): - return get_sea_cache_stats(self) - raise NotImplementedError(f"Cache statistics are not available for {type(self._cache_config).__name__}.") - def _reset_stateful_cache(self, recurse: bool = True) -> None: from ..hooks import HookRegistry diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 324c92f3d507..94700aa2689c 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -364,10 +364,6 @@ def apply_text_kv_cache(*args, **kwargs): requires_backends(apply_text_kv_cache, ["torch"]) -def get_sea_cache_stats(*args, **kwargs): - requires_backends(get_sea_cache_stats, ["torch"]) - - class InpaintProcessor(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py index d6b3798d548a..d68885420137 100644 --- a/tests/hooks/test_sea_cache.py +++ b/tests/hooks/test_sea_cache.py @@ -90,6 +90,7 @@ def _make_config(runtime, **kwargs): "threshold": 100.0, "retention_steps": 1, "cache_end_steps": 0, + "indicator_source": "first_block", "current_step_callback": lambda: runtime["step"], "current_sigma_callback": lambda: runtime["sigma"], "num_inference_steps_callback": lambda: runtime["num_steps"], @@ -99,16 +100,11 @@ def _make_config(runtime, **kwargs): return SeaCacheConfig(**config_kwargs) -def _get_root_hook(model): - return model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) - - @torch.no_grad() def test_sea_cache_uses_independent_gates_and_histories_per_context(): runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} model = DummySeaTransformer() model.enable_cache(_make_config(runtime, threshold=0.5)) - root_hook = _get_root_hook(model) first_input = torch.ones(2, 2, 2, 1) with model.cache_context("cond"): @@ -133,25 +129,11 @@ def test_sea_cache_uses_independent_gates_and_histories_per_context(): torch.testing.assert_close(uncond_output, changed_uncond_input * 8) assert [block.calls for block in model.layers] == [3, 3, 3] assert model.layers[0].indicator_norm.calls == 4 - assert root_hook.num_full_steps == 3 - assert root_hook.num_cached_steps == 1 - stats = model.get_cache_stats() - assert stats["indicator_source"] == "first_block" - assert stats["residual_order"] == 0 - assert stats["residual_boundary"] == "repeated_block_stack" - assert stats["transformer_calls"] == 4 - assert stats["gate_evaluations"] == 4 - assert stats["gate_trace"] == [True, True, False, True] - assert stats["branch_full_executions"] == {"cond": 1, "uncond": 2} - assert stats["branch_reuses"] == {"cond": 1} - assert stats["persistent_cache_bytes"] > 0 model._reset_stateful_cache() - assert root_hook.num_full_steps == 0 - assert root_hook.num_cached_steps == 0 - archived_stats = model.get_cache_stats() - assert archived_stats["transformer_calls"] == 4 - assert archived_stats["gate_trace"] == [True, True, False, True] + with model.cache_context("cond"): + model(first_input) + assert [block.calls for block in model.layers] == [4, 4, 4] model.disable_cache() assert not model.is_cache_enabled @@ -176,15 +158,6 @@ def test_sea_cache_max_consecutive_cached_forces_full_per_context(): with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) - stats = model.get_cache_stats() - assert stats["gate_trace"] == [True, False, False, True, False, False, True, False] - assert stats["actual_full_executions"] == 3 - assert stats["actual_reuses"] == 5 - assert stats["max_consecutive_cached"] == 2 - assert stats["max_consecutive_cached_observed"] == 2 - assert stats["max_consecutive_forced_full"] == 2 - assert stats["per_branch"]["cond"]["max_consecutive_cached_observed"] == 2 - assert stats["per_branch"]["cond"]["max_consecutive_forced_full"] == 2 assert [block.calls for block in model.layers] == [3, 3, 3] @@ -217,11 +190,7 @@ def raw_vision(module, args, kwargs): with model.cache_context("cond"): model(second_input) - stats = model.get_cache_stats() - assert stats["indicator_source"] == "raw_vision_latents" - assert stats["gate_trace"] == [True, True] - assert stats["actual_full_executions"] == 2 - assert stats["actual_reuses"] == 0 + assert [block.calls for block in model.layers] == [2, 2, 2] assert model.layers[0].indicator_norm.calls == 0 @@ -246,34 +215,7 @@ def test_sea_cache_residual_order_one_uses_actual_full_step_history(): # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2. torch.testing.assert_close(output, torch.full_like(output, 24.0)) - root_hook = _get_root_hook(model) - assert root_hook.num_full_steps == 2 - assert root_hook.num_cached_steps == 1 - - -@torch.no_grad() -def test_sea_cache_replays_gate_schedule_and_reports_natural_mismatches(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} - model = DummySeaTransformer() - model.enable_cache( - _make_config( - runtime, - threshold=0.0, - gate_schedule=(True, False, True), - ) - ) - - for step, sigma in enumerate((0.9, 0.6, 0.3)): - runtime.update(step=step, sigma=sigma) - with model.cache_context("cond"): - model(torch.full((1, 1, 1, 1), float(step + 1))) - - stats = model.get_cache_stats() - assert stats["gate_trace"] == [True, False, True] - assert stats["gate_schedule_replayed"] - assert stats["gate_schedule_mismatches"] == 1 - assert stats["actual_full_executions"] == 2 - assert stats["actual_reuses"] == 1 + assert [block.calls for block in model.layers] == [2, 2, 2] @torch.no_grad() @@ -289,7 +231,7 @@ def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_unca with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) - assert model.get_cache_stats()["branch_full_executions"] == {"cond": 1} + assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is not None @torch.no_grad() @@ -307,9 +249,6 @@ def test_sea_cache_fails_open_without_vision_metadata(): model(torch.ones(1, 1, 1, 1)) assert [block.calls for block in model.layers] == [2, 2, 2] - root_hook = _get_root_hook(model) - assert root_hook.num_full_steps == 2 - assert root_hook.num_cached_steps == 0 @torch.no_grad() @@ -329,9 +268,6 @@ def test_sea_cache_fails_open_for_non_adjacent_steps_and_shape_changes(): with model.cache_context("cond"): model(torch.ones(1, 2, 1, 1)) - root_hook = _get_root_hook(model) - assert root_hook.num_full_steps == 3 - assert root_hook.num_cached_steps == 0 assert [block.calls for block in model.layers] == [3, 3, 3] @@ -346,9 +282,7 @@ def test_sea_cache_is_inference_only_and_fails_open_with_autograd(): with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1, requires_grad=True)) - root_hook = _get_root_hook(model) - assert root_hook.num_full_steps == 2 - assert root_hook.num_cached_steps == 0 + assert [block.calls for block in model.layers] == [2, 2, 2] @torch.no_grad() @@ -396,8 +330,6 @@ def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables _, cached_output = model(torch.full((1, 1, 1, 1), 2.0)) torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0)) assert model.layers[0].calls == 1 - assert model.get_cache_stats()["actual_full_executions"] == 1 - assert model.get_cache_stats()["actual_reuses"] == 1 model.disable_cache() @@ -421,10 +353,6 @@ def test_sea_cache_fails_open_for_parameter_sharded_blocks(): with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) - stats = model.get_cache_stats() - assert stats["actual_full_executions"] == 2 - assert stats["actual_reuses"] == 0 - assert stats["fail_open_calls"] == 2 assert [block.calls for block in model.layers] == [2, 2, 2] @@ -449,7 +377,9 @@ def test_sea_cache_config_defaults(): config = SeaCacheConfig() assert config.threshold == 0.25 + assert config.residual_order == 1 assert config.max_consecutive_cached == 2 + assert config.indicator_source == "raw_vision_latents" @pytest.mark.parametrize( @@ -473,11 +403,6 @@ def test_sea_cache_config_validation(kwargs, message): SeaCacheConfig(**kwargs) -def test_sea_cache_gate_schedule_validation(): - with pytest.raises(TypeError, match="gate_schedule"): - SeaCacheConfig(gate_schedule=(True, 1)) - - @pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"]) def test_sea_cache_callback_validation(callback_name): with pytest.raises(TypeError, match=callback_name): diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index 977ea6b52a1a..c03d39f2b4c1 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -150,15 +150,9 @@ def counted_gen_norm_forward(*args, **kwargs): with torch.no_grad(), model.cache_context("cond"): output = model(**cached_inputs) - root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) # 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 - assert root_hook.num_full_steps == 1 - assert root_hook.num_cached_steps == 1 - stats = model.get_cache_stats() - assert stats["indicator_source"] == indicator_source - assert stats["residual_boundary"] == "post_language_model_norm" model.disable_cache() assert set(model.state_dict()) == state_dict_keys @@ -183,7 +177,7 @@ def test_cosmos3_sea_cache_post_norm_boundary_numeric_semantics(self): projection_head_calls = 0 def capture_first_block_gen_input(module, args): - first_block_gen_inputs.append(args[0].detach().clone()) + first_block_gen_inputs.append(args[1].detach().clone()) def capture_und_output(module, args, output): returned_und_outputs.append(output.detach().clone()) @@ -199,7 +193,7 @@ def counted_projection_forward(hidden_states): projection_head_inputs.append(hidden_states.detach().clone()) return original_projection_forward(hidden_states) - model.layers[0].input_layernorm_moe_gen.register_forward_pre_hook(capture_first_block_gen_input) + 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 @@ -231,8 +225,6 @@ def counted_projection_forward(hidden_states): ) torch.testing.assert_close(projection_head_inputs[1], returned_gen_outputs[1]) assert projection_head_calls == 2 - assert root_hook.num_full_steps == 1 - assert root_hook.num_cached_steps == 1 def test_cosmos3_decoder_layer_cache_metadata_tracks_generation_stream(self): metadata = TransformerBlockRegistry.get(Cosmos3VLTextMoTDecoderLayer) From 61ac99b732ee2d94e41764f4e28890aeb20a3a7f Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 05:27:43 -0700 Subject: [PATCH 06/20] Enable SeaCache by default for Cosmos3 --- .../cosmos/modular_pipeline.py | 13 +++++ .../pipelines/cosmos/pipeline_cosmos3_omni.py | 53 +++++++++++++++++++ .../cosmos/test_modular_pipeline_cosmos3.py | 17 +++++- tests/pipelines/cosmos/test_cosmos3.py | 28 +++++++++- 4 files changed, 109 insertions(+), 2 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py index 5efba85de64f..5121e331306e 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py @@ -30,7 +30,20 @@ def current_sigma(self): def num_timesteps(self): return getattr(self, "_num_timesteps", None) + def _prepare_sea_cache_config(self, config=None): + return Cosmos3OmniPipeline._prepare_sea_cache_config(self, config) + + def enable_sea_cache(self, config=None): + return Cosmos3OmniPipeline.enable_sea_cache(self, config) + + def disable_sea_cache(self): + return Cosmos3OmniPipeline.disable_sea_cache(self) + + def _maybe_enable_sea_cache(self): + return Cosmos3OmniPipeline._maybe_enable_sea_cache(self) + def __call__(self, *args, **kwargs): + self._maybe_enable_sea_cache() transformer = getattr(self, "transformer", None) if hasattr(transformer, "_reset_stateful_cache"): transformer._reset_stateful_cache() diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index e5d1c630f130..f999ab27cc72 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -27,6 +27,7 @@ from transformers import AutoTokenizer, BatchEncoding from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...hooks import SeaCacheConfig from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan from ...models.modeling_utils import get_parameter_device @@ -475,9 +476,55 @@ def __init__( # Recommended quality-control negative prompts are documented in the Cosmos3 docs # page (text2video / image2video). When the caller passes None we fall back to "". + self._is_sea_cache_enabled = True + + def _prepare_sea_cache_config(self, config: SeaCacheConfig | None = None) -> SeaCacheConfig: + config = SeaCacheConfig() if config is None else config + config.current_step_callback = lambda: self.current_step_index + config.current_sigma_callback = lambda: self.current_sigma + config.num_inference_steps_callback = lambda: self.num_timesteps + return config + + def enable_sea_cache(self, config: SeaCacheConfig | None = None) -> None: + """Enable SeaCache for subsequent pipeline calls.""" + transformer = getattr(self, "transformer", None) + if transformer is None: + raise ValueError("SeaCache requires a loaded transformer.") + current_config = getattr(transformer, "_cache_config", None) + if current_config is not None: + if isinstance(current_config, SeaCacheConfig): + self._prepare_sea_cache_config(current_config) + self._is_sea_cache_enabled = True + return + raise ValueError( + f"Caching is already enabled with {type(current_config).__name__}. Disable it before enabling SeaCache." + ) + + transformer.enable_cache(self._prepare_sea_cache_config(config)) + self._is_sea_cache_enabled = True + + def disable_sea_cache(self) -> None: + """Disable the default SeaCache optimization for subsequent pipeline calls.""" + self._is_sea_cache_enabled = False + transformer = getattr(self, "transformer", None) + if isinstance(getattr(transformer, "_cache_config", None), SeaCacheConfig): + transformer.disable_cache() + + def _maybe_enable_sea_cache(self) -> None: + transformer = getattr(self, "transformer", None) + if transformer is None or not getattr(self, "_is_sea_cache_enabled", True): + return + if isinstance(getattr(transformer, "_cache_config", None), SeaCacheConfig): + self._prepare_sea_cache_config(transformer._cache_config) + return + if transformer.is_cache_enabled: + return + self.enable_sea_cache() # 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 @@ -489,6 +536,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) @@ -1463,6 +1515,7 @@ def __call__( `sound` (`torch.Tensor` of shape `[C, N]`, or `None` when `enable_sound=False`). Otherwise a tuple `(video, sound)` with the same fields. """ + self._maybe_enable_sea_cache() if hasattr(self.transformer, "_reset_stateful_cache"): self.transformer._reset_stateful_cache() diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index 2750838897ab..395f32d9cc86 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -21,7 +21,7 @@ import torch from PIL import Image -from diffusers import ModularPipeline, UniPCMultistepScheduler +from diffusers import ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler from diffusers.modular_pipelines import ( Cosmos3OmniBlocks, Cosmos3OmniModularPipeline, @@ -189,6 +189,21 @@ def record_context(name): torch.testing.assert_close(sigma, pipe.scheduler.sigmas[expected_step]) assert pipe.current_step_index is None assert pipe.current_sigma is None + assert isinstance(pipe.transformer._cache_config, SeaCacheConfig) + + def test_default_sea_cache_can_be_disabled_and_reenabled(self): + pipe = self.get_pipeline() + pipe._maybe_enable_sea_cache() + + assert isinstance(pipe.transformer._cache_config, SeaCacheConfig) + pipe.disable_sea_cache() + assert not pipe.transformer.is_cache_enabled + + pipe._maybe_enable_sea_cache() + assert not pipe.transformer.is_cache_enabled + + pipe.enable_sea_cache() + assert isinstance(pipe.transformer._cache_config, SeaCacheConfig) def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self): pipe = self.get_pipeline() diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index 9cd6a8dcdab4..4946ebeebb14 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -21,7 +21,13 @@ from PIL import Image from transformers import AutoTokenizer -from diffusers import AutoencoderKLWan, Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler +from diffusers import ( + AutoencoderKLWan, + Cosmos3OmniPipeline, + Cosmos3OmniTransformer, + SeaCacheConfig, + UniPCMultistepScheduler, +) from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image from ...testing_utils import enable_full_determinism, torch_device @@ -111,6 +117,26 @@ def test_inference(self): assert video[0].shape == self.output_shape + def test_default_sea_cache_can_be_disabled_and_reenabled(self): + components = self.get_dummy_components() + pipeline = self.pipeline_class(**components) + pipeline._maybe_enable_sea_cache() + + self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig) + pipeline.disable_sea_cache() + self.assertFalse(pipeline.transformer.is_cache_enabled) + + pipeline._maybe_enable_sea_cache() + self.assertFalse(pipeline.transformer.is_cache_enabled) + + pipeline.enable_sea_cache() + self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig) + + replacement_pipeline = self.pipeline_class(**components) + replacement_pipeline._current_step_index = 3 + replacement_pipeline._maybe_enable_sea_cache() + self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 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) From d89c5eec4531c303315e2aaff20c0c46b553ab58 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 05:45:42 -0700 Subject: [PATCH 07/20] Document SeaCache for Cosmos3 --- docs/source/en/api/cache.md | 6 ++++ docs/source/en/api/pipelines/cosmos3.md | 38 +++++++++++++++++++++++++ docs/source/en/optimization/cache.md | 35 +++++++++++++++++++++++ 3 files changed, 79 insertions(+) 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..1430dbbaa4bb 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -660,6 +660,44 @@ if result.action is not None: +## SeaCache + +Cosmos 3 inference uses [`SeaCacheConfig`] by default. SeaCache reuses transformer residuals when the +Spectral-Evolution-Aware indicator changes slowly, reducing the number of full transformer executions. The default +configuration filters raw vision latents, linearly extrapolates cached residuals, uses a `0.25` threshold, and forces a +full execution after at most two consecutive cached steps. + +SeaCache is approximate and can change generated outputs. Disable it before inference when you need every denoising +step to execute the full transformer, for example when producing an uncached baseline: + +```python +pipe.disable_sea_cache() +result = pipe( + prompt=prompt, + num_frames=189, + height=720, + width=1280, +) +``` + +The same method works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and +[`Cosmos3DistilledModularPipeline`]. Reenable the default configuration with `pipe.enable_sea_cache()`, or pass a +custom configuration: + +```python +from diffusers import SeaCacheConfig + +pipe.enable_sea_cache( + SeaCacheConfig( + threshold=0.2, + max_consecutive_cached=2, + ) +) +``` + +The pipeline supplies SeaCache with the active scheduler step, sigma, and number of inference steps. Cache state is +reset for 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..4db41753507f 100644 --- a/docs/source/en/optimization/cache.md +++ b/docs/source/en/optimization/cache.md @@ -68,6 +68,41 @@ 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. + +[`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and [`Cosmos3DistilledModularPipeline`] enable SeaCache +automatically when inference starts. The default [`SeaCacheConfig`] filters raw vision latents, linearly extrapolates +the cached residual, uses `threshold=0.25`, and allows at most two consecutive cached steps before forcing a full +transformer execution. + +SeaCache is an approximate optimization and may change generated outputs. Disable it for full transformer execution, +such as when measuring an uncached baseline: + +```python +from diffusers import Cosmos3OmniPipeline + +pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano") +pipe.disable_sea_cache() +``` + +Call [`~Cosmos3OmniPipeline.enable_sea_cache`] to reenable it or provide a custom configuration. The pipeline fills in +the scheduler callbacks required by SeaCache. + +```python +from diffusers import SeaCacheConfig + +pipe.enable_sea_cache( + SeaCacheConfig( + threshold=0.2, + max_consecutive_cached=2, + ) +) +``` + ## 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. From 5cf0e5314b81a23f98c2acd717ce359112ba7453 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 07:14:16 -0700 Subject: [PATCH 08/20] Use FP32 sampling state by default for Cosmos3 --- docs/source/en/api/pipelines/cosmos3.md | 6 ++++++ src/diffusers/modular_pipelines/cosmos/before_denoise.py | 2 +- .../cosmos/modular_blocks_cosmos3_distilled.py | 4 ++-- src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py | 8 ++++---- .../cosmos/test_modular_pipeline_cosmos3_distilled.py | 3 ++- tests/pipelines/cosmos/test_cosmos3.py | 1 - 6 files changed, 15 insertions(+), 9 deletions(-) diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index 1430dbbaa4bb..84752ace65b8 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -698,6 +698,12 @@ pipe.enable_sea_cache( The pipeline supplies SeaCache with the active scheduler step, sigma, and number of inference steps. Cache state is reset for each pipeline call, and conditional and unconditional guidance branches keep independent histories. +## Sampling precision + +Cosmos 3 keeps denoising latents and classifier-free-guidance arithmetic in `torch.float32` by default while casting +transformer inputs to the model dtype. Pass `use_fp32_sampling_state=False` to an inference call to keep the sampling +state in the model dtype instead. + ## 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/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index 298312e450c9..a717e97518f9 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -70,7 +70,7 @@ def expected_components(self) -> list[ComponentSpec]: @property def expected_configs(self) -> list[ConfigSpec]: - return [ConfigSpec(name="default_use_fp32_sampling_state", default=False)] + return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)] @property def inputs(self) -> list[InputParam]: 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 a0de9459e5ec..53b140b2960b 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py @@ -86,7 +86,7 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) Configs: - default_use_fp32_sampling_state (default: False) is_distilled (default: True) distilled_sigmas (default: + default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default: None) Inputs: @@ -170,7 +170,7 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): Configs: default_use_system_prompt (default: True) enable_safety_checker (default: True) - default_use_fp32_sampling_state (default: False) is_distilled (default: True) distilled_sigmas (default: + default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default: None) Inputs: diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index f999ab27cc72..08893cfd5b00 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -1407,7 +1407,7 @@ def __call__( add_resolution_template: bool = True, add_duration_template: bool = True, enable_safety_check: bool = True, - use_fp32_sampling_state: bool = False, + use_fp32_sampling_state: bool = True, ) -> Cosmos3OmniPipelineOutput: r""" Run the Cosmos 3 omni pipeline end-to-end: encode the (optional) conditioning image/video, denoise vision and @@ -1503,11 +1503,11 @@ def __call__( When `True` and a `CosmosSafetyChecker` is attached, runs the text guardrail on the prompt before generation and the video guardrail on the decoded frames. Set to `False` to skip both for this call; the checker remains loaded for subsequent calls. - use_fp32_sampling_state (`bool`, *optional*, defaults to `False`): + use_fp32_sampling_state (`bool`, *optional*, defaults to `True`): When `True`, keeps vision, sound, and action denoising latents plus classifier-free-guidance arithmetic in `torch.float32`. Transformer inputs are still cast to the transformer's dtype before each forward. - This improves sampling-state precision at the cost of additional memory and preserves the existing - model-dtype behavior when disabled. + This improves sampling-state precision at the cost of additional memory. Set it to `False` to keep + sampling state in the model dtype. Returns: [`Cosmos3OmniPipelineOutput`] or `tuple`: 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 11e6900d12a1..566343115ca8 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -118,7 +118,8 @@ def test_declares_distilled_configs(self): assert pipe.config.is_distilled is True assert pipe.config.distilled_sigmas is None assert pipe.config.default_use_system_prompt is True - assert pipe.config.default_use_fp32_sampling_state is False + assert pipe.config.default_use_fp32_sampling_state is True + @pytest.mark.parametrize( ("config_enabled", "input_enabled", "expected_dtype"), [ diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index 4946ebeebb14..335398402d09 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -157,7 +157,6 @@ def callback_on_step_end(_pipeline, _step_index, _timestep, callback_kwargs): inputs.update( output_type="latent", enable_safety_check=False, - use_fp32_sampling_state=True, callback_on_step_end=callback_on_step_end, ) with mock.patch.object(pipeline.transformer, "forward", side_effect=transformer_forward): From 68bc0068044444a5a0cbd8e3e9848c38a926ee35 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 10:56:32 -0700 Subject: [PATCH 09/20] Fix SeaCache reconfiguration, transfer state isolation, and FP32 sampling --- src/diffusers/hooks/sea_cache.py | 182 +++++++++++++++++- .../cosmos/before_denoise.py | 102 ++++++++-- .../modular_pipelines/cosmos/denoise.py | 26 ++- .../cosmos/modular_blocks_cosmos3.py | 105 +++++++--- .../modular_blocks_cosmos3_distilled.py | 44 +++-- .../pipelines/cosmos/pipeline_cosmos3_omni.py | 15 +- tests/hooks/test_sea_cache.py | 133 +++++++++++-- .../cosmos/test_modular_pipeline_cosmos3.py | 119 +++++++++++- ...test_modular_pipeline_cosmos3_distilled.py | 103 +++++----- tests/pipelines/cosmos/test_cosmos3.py | 21 ++ 10 files changed, 711 insertions(+), 139 deletions(-) diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index 827648709518..93e2164dfccc 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import inspect import math +import time from dataclasses import dataclass -from typing import Any, Callable, Literal +from typing import Any, Callable, Literal, Sequence import torch @@ -72,11 +74,15 @@ class SeaCacheConfig: 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. + 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. + gate_schedule (`Sequence[bool]`, *optional*): + Per-step full-compute decisions to replay for controlled residual-prediction ablations. SEA still evaluates + its natural decision and reports mismatches; invalid or unsafe calls always fail open to full compute. + Example: ```python >>> from diffusers import Cosmos3OmniPipeline, SeaCacheConfig @@ -110,6 +116,7 @@ class SeaCacheConfig: [torch.nn.Module, tuple[Any, ...], dict[str, Any]], list[torch.Tensor] | None, ] = None + gate_schedule: Sequence[bool] = None def __post_init__(self): if not math.isfinite(self.threshold) or self.threshold < 0: @@ -144,6 +151,10 @@ def __post_init__(self): callback = getattr(self, name) if callback is not None and not callable(callback): raise TypeError(f"`{name}` must be callable or `None`.") + if self.gate_schedule is not None: + self.gate_schedule = tuple(self.gate_schedule) + if any(not isinstance(value, bool) for value in self.gate_schedule): + raise TypeError("`gate_schedule` must contain only boolean values.") @dataclass @@ -163,34 +174,41 @@ def __init__(self): self.previous_indicator: list[torch.Tensor] | None = None self.accumulated_distance = 0.0 self.consecutive_cached = 0 + self.max_consecutive_cached_observed = 0 + self.max_consecutive_forced_full = 0 self.skip_remaining = False self.full_execution_pending = False self.cacheable_execution = False self.step_index: int | None = None - self.und_input: torch.Tensor | 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 + self.full_started_at: float | None = None def reset_forward(self): self.skip_remaining = False self.full_execution_pending = False self.cacheable_execution = False self.step_index = None - self.und_input = None self.gen_input = None self.und_output = None self.cached_und_output = None self.cached_gen_residual = None + self.full_started_at = None - def reset(self): + 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.max_consecutive_cached_observed = 0 + self.max_consecutive_forced_full = 0 self.reset_forward() @@ -201,6 +219,20 @@ def __init__(self): def reset(self): self.forward_metadata: _SeaCacheForwardMetadata | None = None + self.transformer_calls = 0 + self.gate_evaluations = 0 + self.gate_full_decisions = 0 + self.gate_skip_decisions = 0 + self.gate_trace: list[bool] = [] + self.gate_schedule_mismatches = 0 + self.num_full_steps = 0 + self.num_cached_steps = 0 + self.fail_open_calls = 0 + self.indicator_seconds = 0.0 + self.decision_seconds = 0.0 + self.full_seconds = 0.0 + self.branch_full_executions: dict[str, int] = {} + self.branch_reuses: dict[str, int] = {} def warn_once(self, message: str): if message not in self._warned_messages: @@ -208,6 +240,7 @@ def warn_once(self, message: str): self._warned_messages.add(message) def mark_fail_open(self, message: str): + self.fail_open_calls += 1 self.warn_once(message) def resolve_gate( @@ -221,9 +254,11 @@ def resolve_gate( if state.gate_key == gate_key: return state.gate_should_compute + self.gate_evaluations += 1 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 @@ -235,7 +270,7 @@ def resolve_gate( candidate_accumulated_distance = 0.0 if forced_compute: - should_compute = True + natural_should_compute = True else: if len(indicator) != len(state.previous_indicator) or not indicator: distance = float("inf") @@ -260,13 +295,41 @@ def resolve_gate( invalid_gate = True self.mark_fail_open("SeaCache indicator history changed shape, device, or dtype; running full.") candidate_accumulated_distance = state.accumulated_distance + distance - should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold + natural_should_compute = invalid_gate or candidate_accumulated_distance >= config.threshold + + should_compute = natural_should_compute + + if config.gate_schedule is not None: + schedule_is_valid = len(config.gate_schedule) == metadata.num_inference_steps + scheduled_compute = bool(config.gate_schedule[metadata.step_index]) if schedule_is_valid else True + can_replay_skip = not (invalid_gate or is_retained or is_in_cache_end or is_first_observation) + if not schedule_is_valid: + self.mark_fail_open( + "SeaCache gate schedule length does not match the number of inference steps; running full." + ) + should_compute = True + elif scheduled_compute or can_replay_skip: + if natural_should_compute != scheduled_compute: + self.gate_schedule_mismatches += 1 + should_compute = scheduled_compute + else: + self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.") + should_compute = True + + if is_max_consecutive: + should_compute = True + state.max_consecutive_forced_full += 1 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] + self.gate_trace.append(should_compute) + if should_compute: + self.gate_full_decisions += 1 + else: + self.gate_skip_decisions += 1 return should_compute @@ -312,11 +375,18 @@ def _get_block_outputs( def _record_full_execution( config: SeaCacheConfig, + state_manager: StateManager, + shared_state: SeaCacheSharedState, state: SeaCacheContextState, gen_output: torch.Tensor, und_output: torch.Tensor | None, ) -> None: + shared_state.num_full_steps += 1 state.consecutive_cached = 0 + branch = state_manager._current_context + shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1 + if state.full_started_at is not None: + shared_state.full_seconds += time.perf_counter() - state.full_started_at if ( state.cacheable_execution and state.step_index is not None @@ -500,6 +570,7 @@ def __init__( shared_state: SeaCacheSharedState, metadata_callback: Callable, raw_vision_callback: Callable, + residual_boundary: str, ): super().__init__() self.config = config @@ -507,9 +578,73 @@ def __init__( self.shared_state = shared_state self.metadata_callback = metadata_callback self.raw_vision_callback = raw_vision_callback + self.residual_boundary = residual_boundary + self._last_stats: dict[str, Any] | None = None + + @property + def num_full_steps(self) -> int: + return self.shared_state.num_full_steps + + @property + def num_cached_steps(self) -> int: + return self.shared_state.num_cached_steps + + def stats(self) -> dict[str, Any]: + if self.shared_state.transformer_calls == 0 and self._last_stats is not None: + return copy.deepcopy(self._last_stats) + + persistent_cache_bytes = 0 + per_branch = {} + for branch, state in sorted(self.state_manager._state_cache.items()): + if state.previous_indicator is not None: + persistent_cache_bytes += sum( + value.numel() * value.element_size() for value in state.previous_indicator + ) + for _, und_output, gen_residual in state.history: + persistent_cache_bytes += und_output.numel() * und_output.element_size() + persistent_cache_bytes += gen_residual.numel() * gen_residual.element_size() + per_branch[branch] = { + "full_calls": self.shared_state.branch_full_executions.get(branch, 0), + "reuse_calls": self.shared_state.branch_reuses.get(branch, 0), + "max_consecutive_cached_observed": state.max_consecutive_cached_observed, + "max_consecutive_forced_full": state.max_consecutive_forced_full, + } + + opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps + return { + "indicator_source": self.config.indicator_source, + "residual_order": self.config.residual_order, + "residual_boundary": self.residual_boundary, + "max_consecutive_cached": self.config.max_consecutive_cached, + "max_consecutive_cached_observed": max( + (stats["max_consecutive_cached_observed"] for stats in per_branch.values()), default=0 + ), + "max_consecutive_forced_full": sum(stats["max_consecutive_forced_full"] for stats in per_branch.values()), + "transformer_calls": self.shared_state.transformer_calls, + "gate_evaluations": self.shared_state.gate_evaluations, + "gate_full_decisions": self.shared_state.gate_full_decisions, + "gate_skip_decisions": self.shared_state.gate_skip_decisions, + "gate_trace": list(self.shared_state.gate_trace), + "gate_schedule_replayed": self.config.gate_schedule is not None, + "gate_schedule_mismatches": self.shared_state.gate_schedule_mismatches, + "actual_full_executions": self.shared_state.num_full_steps, + "actual_reuses": self.shared_state.num_cached_steps, + "actual_reuse_rate": (self.shared_state.num_cached_steps / opportunities if opportunities else 0.0), + "fail_open_calls": self.shared_state.fail_open_calls, + "sea_indicator_seconds": self.shared_state.indicator_seconds, + "sea_decision_seconds": self.shared_state.decision_seconds, + "sea_seconds": (self.shared_state.indicator_seconds + self.shared_state.decision_seconds), + "full_seconds": self.shared_state.full_seconds, + "timing_note": "host wall time; CUDA work is asynchronous", + "persistent_cache_bytes": persistent_cache_bytes, + "branch_full_executions": dict(sorted(self.shared_state.branch_full_executions.items())), + "branch_reuses": dict(sorted(self.shared_state.branch_reuses.items())), + "per_branch": per_branch, + } def pre_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.forward_metadata = None + self.shared_state.transformer_calls += 1 if torch.is_grad_enabled(): self.shared_state.mark_fail_open( "SeaCache is inference-only; calls with autograd enabled run in fail-open mode." @@ -608,6 +743,8 @@ def post_forward(self, module: torch.nn.Module, output: Any) -> Any: return output def reset_state(self, module: torch.nn.Module): + if self.shared_state.transformer_calls > 0: + self._last_stats = self.stats() self.state_manager.reset() self.shared_state.reset() return module @@ -689,14 +826,16 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): state.reset_forward() state.full_execution_pending = True state.gen_input = hidden_states - state.und_input = encoder_hidden_states forward_metadata = self.shared_state.forward_metadata if state is None or forward_metadata is None: + if state is not None: + state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) state.step_index = forward_metadata.step_index state.cacheable_execution = True + indicator_started = time.perf_counter() indicator_error_reported = False if _is_parameter_sharded(module): self.shared_state.mark_fail_open( @@ -713,11 +852,14 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): ) indicator = None indicator_error_reported = True + self.shared_state.indicator_seconds += time.perf_counter() - indicator_started 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." ) + decision_started = time.perf_counter() should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config) + self.shared_state.decision_seconds += time.perf_counter() - decision_started if should_compute or not state.history: if not should_compute: @@ -725,6 +867,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.mark_fail_open( "SeaCache selected a cache hit without residual history; running in fail-open mode." ) + state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) residual_history = state.history[-(self.config.residual_order + 1) :] @@ -746,6 +889,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.mark_fail_open( "SeaCache residual history changed shape, device, or dtype; running in fail-open mode." ) + state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) if self.config.residual_order == 1 and len(residual_history) >= 2: @@ -760,6 +904,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): state.cached_und_output = cached_und state.cached_gen_residual = cached_residual state.consecutive_cached += 1 + state.max_consecutive_cached_observed = max(state.max_consecutive_cached_observed, state.consecutive_cached) + self.shared_state.num_cached_steps += 1 + branch = self.state_manager._current_context + self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 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) @@ -804,6 +952,8 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output) _record_full_execution( self.config, + self.state_manager, + self.shared_state, state, gen_output=hidden_states, und_output=encoder_hidden_states, @@ -856,6 +1006,8 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): _record_full_execution( self.config, + self.state_manager, + self.shared_state, state, gen_output=output, und_output=state.und_output, @@ -903,6 +1055,7 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None: "residual boundary." ) post_norm_boundary = post_norm_modules is not None + residual_boundary = "post_language_model_norm" if post_norm_boundary else "repeated_block_stack" blocks = [] for name, submodule in unwrapped_module.named_children(): @@ -932,6 +1085,7 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: shared_state, metadata_callback, raw_vision_callback, + residual_boundary, ), _SEA_CACHE_ROOT_HOOK, ) @@ -980,3 +1134,13 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: raise root_registry._child_registries_cache = None + + +def get_sea_cache_stats(module: torch.nn.Module) -> dict[str, Any]: + """Return statistics for the SeaCache instance currently attached to ``module``.""" + + registry = getattr(module, "_diffusers_hook", None) + root_hook = registry.get_hook(_SEA_CACHE_ROOT_HOOK) if registry is not None else None + if not isinstance(root_hook, SeaCacheRootHook): + raise ValueError("SeaCache is not enabled on this module.") + return root_hook.stats() diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index a717e97518f9..decbe325782b 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -12,6 +12,15 @@ from .modular_pipeline import Cosmos3OmniModularPipeline +def _resolve_sampling_dtype( + components: Cosmos3OmniModularPipeline, use_fp32_sampling_state: bool | None +) -> tuple[bool, torch.dtype]: + if use_fp32_sampling_state is None: + use_fp32_sampling_state = components.config.default_use_fp32_sampling_state + sampling_dtype = torch.float32 if use_fp32_sampling_state else components.transformer.dtype + return use_fp32_sampling_state, sampling_dtype + + class Cosmos3PrepareTextSegmentsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @@ -107,7 +116,8 @@ def inputs(self) -> list[InputParam]: type_hint=bool | None, default=None, description=( - "Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the " + "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. " + "If unset, uses the " "pipeline's `default_use_fp32_sampling_state` config." ), ), @@ -145,9 +155,9 @@ 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 - if block_state.use_fp32_sampling_state is None: - block_state.use_fp32_sampling_state = components.config.default_use_fp32_sampling_state - sampling_dtype = torch.float32 if block_state.use_fp32_sampling_state else components.transformer.dtype + block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype( + components, block_state.use_fp32_sampling_state + ) x0_tokens_vision = block_state.x0_tokens_vision if x0_tokens_vision is None: @@ -216,6 +226,10 @@ def expected_components(self) -> list[ComponentSpec]: ComponentSpec("scheduler", UniPCMultistepScheduler), ] + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)] + @property def inputs(self) -> list[InputParam]: return [ @@ -228,6 +242,16 @@ def inputs(self) -> list[InputParam]: description="Pre-generated noisy sound latents.", ), InputParam.template("generator"), + InputParam( + name="use_fp32_sampling_state", + type_hint=bool | None, + default=None, + description=( + "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. " + "If unset, uses the " + "pipeline's `default_use_fp32_sampling_state` config." + ), + ), ] @property @@ -248,7 +272,9 @@ 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 + block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype( + components, block_state.use_fp32_sampling_state + ) if not components.transformer.config.sound_gen: raise ValueError("Sound generation requires a transformer trained with sound_gen=True.") @@ -258,19 +284,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) @@ -292,6 +320,10 @@ def expected_components(self) -> list[ComponentSpec]: ComponentSpec("scheduler", UniPCMultistepScheduler), ] + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)] + @property def inputs(self) -> list[InputParam]: return [ @@ -314,6 +346,16 @@ def inputs(self) -> list[InputParam]: description="Pre-generated noisy action latents.", ), InputParam.template("generator"), + InputParam( + name="use_fp32_sampling_state", + type_hint=bool | None, + default=None, + description=( + "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. " + "If unset, uses the " + "pipeline's `default_use_fp32_sampling_state` config." + ), + ), ] @property @@ -345,7 +387,9 @@ 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 + block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype( + components, block_state.use_fp32_sampling_state + ) action = block_state.action if not components.transformer.config.action_gen: @@ -367,7 +411,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}." @@ -388,7 +432,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( @@ -398,14 +442,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 @@ -414,7 +460,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) @@ -1021,6 +1067,10 @@ def description(self) -> str: def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + @property + def expected_configs(self) -> list[ConfigSpec]: + return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)] + @property def inputs(self) -> list[InputParam]: return [ @@ -1037,6 +1087,16 @@ def inputs(self) -> list[InputParam]: description="Number of pixel frames used to seed this chunk's target.", ), InputParam.template("generator"), + InputParam( + name="use_fp32_sampling_state", + type_hint=bool | None, + default=None, + description=( + "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. " + "If unset, uses the " + "pipeline's `default_use_fp32_sampling_state` config." + ), + ), ] @property @@ -1064,20 +1124,24 @@ 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 + block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype( + components, block_state.use_fp32_sampling_state + ) 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 ce6ba0a4bee5..5b06b4f910a2 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -172,6 +172,12 @@ def inputs(self) -> list[InputParam]: default=6.0, description="Scale for classifier-free guidance.", ), + InputParam( + name="use_fp32_sampling_state", + type_hint=bool, + default=False, + description="Whether to keep velocity and guidance arithmetic in float32.", + ), ] @property @@ -236,6 +242,10 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta ) cond_velocity_vision, cond_velocity_sound, cond_velocity_action = velocities["cond"] + if block_state.use_fp32_sampling_state: + 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"] block_state.velocity_vision = uncond_velocity_vision + block_state.guidance_scale * ( @@ -730,6 +740,12 @@ def inputs(self) -> list[InputParam]: default=None, description="Timestep interval [lo, hi] over which control guidance is active (None = always).", ), + InputParam( + name="use_fp32_sampling_state", + type_hint=bool, + default=False, + description="Whether to keep velocity and guidance arithmetic in float32.", + ), ] @property @@ -804,6 +820,11 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta "uncond", ) + if block_state.use_fp32_sampling_state: + 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) @@ -878,7 +899,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`): @@ -905,6 +927,8 @@ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper): Timestep interval [lo, hi] over which text guidance is active (None = always). control_guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which control guidance is active (None = always). + use_fp32_sampling_state (`bool`, *optional*, defaults to False): + Whether to keep velocity and guidance arithmetic in float32. latents (`Tensor`): Noisy target latents to update. condition_latents (`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..ee02e192da21 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,9 +373,11 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text-and-vision Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Configs: + default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -394,6 +401,9 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): @@ -439,9 +449,11 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, and sound Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Configs: + default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -465,6 +477,9 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. sound_latents (`Tensor`, *optional*): @@ -523,9 +538,11 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, and action Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Configs: + default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -549,6 +566,9 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. action (`CosmosActionCondition`): @@ -611,9 +631,11 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text, vision, sound, and action Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) Configs: + default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -637,6 +659,9 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. sound_latents (`Tensor`, *optional*): @@ -707,13 +732,16 @@ 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`) + + Configs: + default_use_fp32_sampling_state (default: True) Inputs: chunk_id (`int`, *optional*, defaults to 0): @@ -740,6 +768,9 @@ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. cond_text_segment (`dict`): Conditional text segment. uncond_text_segment (`dict`): @@ -842,6 +873,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 +887,13 @@ 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`) + + Configs: + default_use_fp32_sampling_state (default: True) Inputs: cond_input_ids (`None`): @@ -886,6 +924,9 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. num_inference_steps (`int`): @@ -973,10 +1014,13 @@ 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: + default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -1008,6 +1052,9 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. num_inference_steps (`int`): @@ -1162,13 +1209,18 @@ 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) + default_use_fp32_sampling_state (default: True) + use_native_flow_schedule (default: False) Inputs: control_videos (`dict`, *optional*): @@ -1213,6 +1265,9 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Number of frames the first chunk reuses from the input video. generator (`Generator`, *optional*): Torch generator for deterministic generation. + use_fp32_sampling_state (`bool | NoneType`, *optional*): + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *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 53b140b2960b..6841d2f3ff4f 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,11 +84,13 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): Runs the text-and-vision distilled Cosmos3 denoising workflow. Components: - transformer (`Cosmos3OmniTransformer`) scheduler (`FlowMatchEulerDiscreteScheduler`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`FlowMatchEulerDiscreteScheduler`) Configs: - default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default: - None) + default_use_fp32_sampling_state (default: True) + is_distilled (default: True) + distilled_sigmas (default: None) Inputs: cond_input_ids (`None`): @@ -111,13 +114,13 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): generator (`Generator`, *optional*): Torch generator for deterministic generation. use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the pipeline's - `default_use_fp32_sampling_state` config. + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. 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. @@ -165,13 +168,18 @@ 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) - default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default: - None) + default_use_system_prompt (default: True) + enable_safety_checker (default: True) + default_use_fp32_sampling_state (default: True) + is_distilled (default: True) + distilled_sigmas (default: None) Inputs: prompt (`str`): @@ -184,7 +192,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. @@ -207,13 +215,13 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): generator (`Generator`, *optional*): Torch generator for deterministic generation. use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep vision latents, masks, and scheduler state in float32. If unset, uses the pipeline's - `default_use_fp32_sampling_state` config. + Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the + pipeline's `default_use_fp32_sampling_state` config. 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/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 08893cfd5b00..3847f06b761c 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -493,12 +493,15 @@ def enable_sea_cache(self, config: SeaCacheConfig | None = None) -> None: current_config = getattr(transformer, "_cache_config", None) if current_config is not None: if isinstance(current_config, SeaCacheConfig): - self._prepare_sea_cache_config(current_config) - self._is_sea_cache_enabled = True - return - raise ValueError( - f"Caching is already enabled with {type(current_config).__name__}. Disable it before enabling SeaCache." - ) + if config is None or config is current_config: + self._prepare_sea_cache_config(current_config) + self._is_sea_cache_enabled = True + return + transformer.disable_cache() + else: + raise ValueError( + f"Caching is already enabled with {type(current_config).__name__}. Disable it before enabling SeaCache." + ) transformer.enable_cache(self._prepare_sea_cache_config(config)) self._is_sea_cache_enabled = True diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py index d68885420137..895bf299ccf4 100644 --- a/tests/hooks/test_sea_cache.py +++ b/tests/hooks/test_sea_cache.py @@ -100,11 +100,16 @@ def _make_config(runtime, **kwargs): return SeaCacheConfig(**config_kwargs) +def _get_root_hook(model): + return model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) + + @torch.no_grad() def test_sea_cache_uses_independent_gates_and_histories_per_context(): runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} model = DummySeaTransformer() model.enable_cache(_make_config(runtime, threshold=0.5)) + root_hook = _get_root_hook(model) first_input = torch.ones(2, 2, 2, 1) with model.cache_context("cond"): @@ -129,11 +134,25 @@ def test_sea_cache_uses_independent_gates_and_histories_per_context(): torch.testing.assert_close(uncond_output, changed_uncond_input * 8) assert [block.calls for block in model.layers] == [3, 3, 3] assert model.layers[0].indicator_norm.calls == 4 + assert root_hook.num_full_steps == 3 + assert root_hook.num_cached_steps == 1 + stats = model.get_cache_stats() + assert stats["indicator_source"] == "first_block" + assert stats["residual_order"] == 1 + assert stats["residual_boundary"] == "repeated_block_stack" + assert stats["transformer_calls"] == 4 + assert stats["gate_evaluations"] == 4 + assert stats["gate_trace"] == [True, True, False, True] + assert stats["branch_full_executions"] == {"cond": 1, "uncond": 2} + assert stats["branch_reuses"] == {"cond": 1} + assert stats["persistent_cache_bytes"] > 0 model._reset_stateful_cache() - with model.cache_context("cond"): - model(first_input) - assert [block.calls for block in model.layers] == [4, 4, 4] + assert root_hook.num_full_steps == 0 + assert root_hook.num_cached_steps == 0 + archived_stats = model.get_cache_stats() + assert archived_stats["transformer_calls"] == 4 + assert archived_stats["gate_trace"] == [True, True, False, True] model.disable_cache() assert not model.is_cache_enabled @@ -158,6 +177,15 @@ def test_sea_cache_max_consecutive_cached_forces_full_per_context(): with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) + stats = model.get_cache_stats() + assert stats["gate_trace"] == [True, False, False, True, False, False, True, False] + assert stats["actual_full_executions"] == 3 + assert stats["actual_reuses"] == 5 + assert stats["max_consecutive_cached"] == 2 + assert stats["max_consecutive_cached_observed"] == 2 + assert stats["max_consecutive_forced_full"] == 2 + assert stats["per_branch"]["cond"]["max_consecutive_cached_observed"] == 2 + assert stats["per_branch"]["cond"]["max_consecutive_forced_full"] == 2 assert [block.calls for block in model.layers] == [3, 3, 3] @@ -190,7 +218,11 @@ def raw_vision(module, args, kwargs): with model.cache_context("cond"): model(second_input) - assert [block.calls for block in model.layers] == [2, 2, 2] + stats = model.get_cache_stats() + assert stats["indicator_source"] == "raw_vision_latents" + assert stats["gate_trace"] == [True, True] + assert stats["actual_full_executions"] == 2 + assert stats["actual_reuses"] == 0 assert model.layers[0].indicator_norm.calls == 0 @@ -215,7 +247,34 @@ def test_sea_cache_residual_order_one_uses_actual_full_step_history(): # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2. torch.testing.assert_close(output, torch.full_like(output, 24.0)) - assert [block.calls for block in model.layers] == [2, 2, 2] + root_hook = _get_root_hook(model) + assert root_hook.num_full_steps == 2 + assert root_hook.num_cached_steps == 1 + + +@torch.no_grad() +def test_sea_cache_replays_gate_schedule_and_reports_natural_mismatches(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} + model = DummySeaTransformer() + model.enable_cache( + _make_config( + runtime, + threshold=0.0, + gate_schedule=(True, False, True), + ) + ) + + for step, sigma in enumerate((0.9, 0.6, 0.3)): + runtime.update(step=step, sigma=sigma) + with model.cache_context("cond"): + model(torch.full((1, 1, 1, 1), float(step + 1))) + + stats = model.get_cache_stats() + assert stats["gate_trace"] == [True, False, True] + assert stats["gate_schedule_replayed"] + assert stats["gate_schedule_mismatches"] == 1 + assert stats["actual_full_executions"] == 2 + assert stats["actual_reuses"] == 1 @torch.no_grad() @@ -231,7 +290,7 @@ def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_unca with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) - assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is not None + assert model.get_cache_stats()["branch_full_executions"] == {"cond": 1} @torch.no_grad() @@ -249,26 +308,54 @@ def test_sea_cache_fails_open_without_vision_metadata(): model(torch.ones(1, 1, 1, 1)) assert [block.calls for block in model.layers] == [2, 2, 2] + root_hook = _get_root_hook(model) + assert root_hook.num_full_steps == 2 + assert root_hook.num_cached_steps == 0 @torch.no_grad() -def test_sea_cache_fails_open_for_non_adjacent_steps_and_shape_changes(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} +def test_sea_cache_non_adjacent_steps_start_a_new_residual_trajectory(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + config = _make_config(runtime, threshold=0.0, retention_steps=0, residual_order=1) model = DummySeaTransformer() - model.enable_cache(_make_config(runtime, residual_order=1)) + model.enable_cache(config) with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) + runtime.update(step=1, sigma=0.5) + with model.cache_context("cond"): + model(torch.full((1, 1, 1, 1), 100.0)) - runtime.update(step=2, sigma=0.5) + config.threshold = 100.0 + runtime.update(step=0, sigma=0.9) with model.cache_context("cond"): - model(torch.ones(1, 2, 1, 1)) + model(torch.full((1, 1, 1, 1), 2.0)) + runtime.update(step=1, sigma=0.5) + with model.cache_context("cond"): + _, output = model(torch.full((1, 1, 1, 1), 2.0)) + + torch.testing.assert_close(output, torch.full_like(output, 16.0)) + assert [block.calls for block in model.layers] == [3, 3, 3] + assert model.get_cache_stats()["gate_trace"] == [True, True, True, False] + + +@torch.no_grad() +def test_sea_cache_fails_open_for_shape_changes(): + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + model = DummySeaTransformer() + model.enable_cache(_make_config(runtime, residual_order=1)) + + with model.cache_context("cond"): + model(torch.ones(1, 1, 1, 1)) - runtime.update(step=3, sigma=0.2) + runtime.update(step=1, sigma=0.5) with model.cache_context("cond"): model(torch.ones(1, 2, 1, 1)) - assert [block.calls for block in model.layers] == [3, 3, 3] + root_hook = _get_root_hook(model) + assert root_hook.num_full_steps == 2 + assert root_hook.num_cached_steps == 0 + assert [block.calls for block in model.layers] == [2, 2, 2] def test_sea_cache_is_inference_only_and_fails_open_with_autograd(): @@ -276,12 +363,17 @@ def test_sea_cache_is_inference_only_and_fails_open_with_autograd(): model = DummySeaTransformer() model.enable_cache(_make_config(runtime)) - with model.cache_context("cond"): + with torch.enable_grad(), model.cache_context("cond"): model(torch.ones(1, 1, 1, 1, requires_grad=True)) runtime.update(step=1, sigma=0.5) - with model.cache_context("cond"): + with torch.enable_grad(), model.cache_context("cond"): model(torch.ones(1, 1, 1, 1, requires_grad=True)) + root_hook = _get_root_hook(model) + assert root_hook.shared_state.transformer_calls == 2 + assert root_hook.shared_state.fail_open_calls == 2 + assert root_hook.num_full_steps == 2 + assert root_hook.num_cached_steps == 0 assert [block.calls for block in model.layers] == [2, 2, 2] @@ -330,6 +422,8 @@ def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables _, cached_output = model(torch.full((1, 1, 1, 1), 2.0)) torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0)) assert model.layers[0].calls == 1 + assert model.get_cache_stats()["actual_full_executions"] == 1 + assert model.get_cache_stats()["actual_reuses"] == 1 model.disable_cache() @@ -353,6 +447,10 @@ def test_sea_cache_fails_open_for_parameter_sharded_blocks(): with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) + stats = model.get_cache_stats() + assert stats["actual_full_executions"] == 2 + assert stats["actual_reuses"] == 0 + assert stats["fail_open_calls"] == 2 assert [block.calls for block in model.layers] == [2, 2, 2] @@ -403,6 +501,11 @@ def test_sea_cache_config_validation(kwargs, message): SeaCacheConfig(**kwargs) +def test_sea_cache_gate_schedule_validation(): + with pytest.raises(TypeError, match="gate_schedule"): + SeaCacheConfig(gate_schedule=(True, 1)) + + @pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"]) def test_sea_cache_callback_validation(callback_name): with pytest.raises(TypeError, match=callback_name): diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index 395f32d9cc86..d3afd2e33a97 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -21,7 +21,7 @@ import torch from PIL import Image -from diffusers import ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler +from diffusers import CosmosActionCondition, ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler from diffusers.modular_pipelines import ( Cosmos3OmniBlocks, Cosmos3OmniModularPipeline, @@ -30,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 ( @@ -205,6 +209,119 @@ def test_default_sea_cache_can_be_disabled_and_reenabled(self): pipe.enable_sea_cache() assert isinstance(pipe.transformer._cache_config, SeaCacheConfig) + def test_default_sea_cache_reuses_transformer_execution(self): + pipe = self.get_pipeline().to(torch_device) + scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, use_flow_sigmas=True) + pipe.update_components(scheduler=scheduler, use_native_flow_schedule=True) + pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6)) + inputs = self.get_dummy_inputs() + inputs["num_inference_steps"] = 4 + + pipe(**inputs, output=self.output_name) + + stats = pipe.transformer.get_cache_stats() + assert stats["actual_reuses"] > 0 + assert stats["actual_full_executions"] > 0 + + 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 + + @pytest.mark.parametrize( + ("use_fp32_sampling_state", "expected_dtype"), + [(False, torch.bfloat16), (True, torch.float32)], + ) + def test_sound_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype): + pipe = self._get_sampling_state_block_pipe(Cosmos3SoundPrepareLatentsStep()) + + outputs = pipe( + num_frames=5, + fps=24.0, + generator=self.get_generator(0), + use_fp32_sampling_state=use_fp32_sampling_state, + output=["sound_latents", "sound_condition_mask"], + ) + + assert outputs["sound_latents"].dtype == expected_dtype + assert outputs["sound_condition_mask"].dtype == expected_dtype + + @pytest.mark.parametrize( + ("use_fp32_sampling_state", "expected_dtype"), + [(False, torch.bfloat16), (True, torch.float32)], + ) + def test_action_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype): + 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), + use_fp32_sampling_state=use_fp32_sampling_state, + output=["action_latents", "action_condition_mask"], + ) + + assert outputs["action_latents"].dtype == expected_dtype + assert outputs["action_condition_mask"].dtype == expected_dtype + + @pytest.mark.parametrize( + ("use_fp32_sampling_state", "expected_dtype"), + [(False, torch.bfloat16), (True, torch.float32)], + ) + def test_transfer_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype): + 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), + use_fp32_sampling_state=use_fp32_sampling_state, + output=["latents", "velocity_mask", "condition_latents"], + ) + + assert outputs["latents"].dtype == expected_dtype + assert outputs["velocity_mask"].dtype == expected_dtype + assert outputs["condition_latents"].dtype == expected_dtype + + 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), + ] + + @pytest.mark.parametrize( + ("use_fp32_sampling_state", "expected_dtype"), + [(False, torch.bfloat16), (True, torch.float32)], + ) + def test_sampling_state_controls_modular_cfg_and_scheduler_dtype(self, use_fp32_sampling_state, expected_dtype): + pipe = self.get_pipeline(dtype=torch.bfloat16).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["use_fp32_sampling_state"] = use_fp32_sampling_state + + outputs = pipe(**inputs, output=["velocity_vision", "latents"]) + + assert outputs["velocity_vision"].dtype == expected_dtype + assert outputs["latents"].dtype == expected_dtype + 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 566343115ca8..eb94dc80b010 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -17,12 +17,10 @@ import torch from PIL import Image -from diffusers import FlowMatchEulerDiscreteScheduler, ModularPipeline +from diffusers import ModularPipeline, SeaCacheConfig from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline from diffusers.modular_pipelines.cosmos.before_denoise import Cosmos3VisionPrepareLatentsStep -from diffusers.modular_pipelines.cosmos.denoise import Cosmos3DistilledVisionLoopSchedulerStep from diffusers.modular_pipelines.cosmos.encoders import Cosmos3DistilledTextEncoderStep -from diffusers.modular_pipelines.modular_pipeline import BlockState, PipelineState from ...testing_utils import torch_device from ..testing_utils import ( @@ -120,6 +118,30 @@ def test_declares_distilled_configs(self): assert pipe.config.default_use_system_prompt is True assert pipe.config.default_use_fp32_sampling_state 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) + @pytest.mark.parametrize( ("config_enabled", "input_enabled", "expected_dtype"), [ @@ -128,57 +150,48 @@ def test_declares_distilled_configs(self): (True, False, torch.bfloat16), ], ) - def test_prepare_vision_latents_fp32_sampling_state( - self, config_enabled, input_enabled, expected_dtype - ): - components = BlockState( - _execution_device=torch.device("cpu"), - transformer=BlockState(dtype=torch.bfloat16), - config=BlockState(default_use_fp32_sampling_state=config_enabled), - vae_scale_factor_spatial=16, - vae_scale_factor_temporal=4, - num_channels_latents=4, + def test_prepare_vision_latents_fp32_sampling_state(self, config_enabled, input_enabled, expected_dtype): + prepare_pipe = Cosmos3VisionPrepareLatentsStep().init_pipeline(self.pretrained_model_name_or_path) + prepare_pipe.load_components(torch_dtype=torch.bfloat16) + prepare_pipe.update_components(default_use_fp32_sampling_state=config_enabled) + prepare_pipe.to(torch_device) + + outputs = prepare_pipe( + num_frames=5, + height=32, + width=32, + fps=24.0, + generator=self.get_generator(0), + use_fp32_sampling_state=input_enabled, + output=["latents", "vision_condition_mask", "use_fp32_sampling_state"], ) - state = PipelineState() - state.set("x0_tokens_vision", None) - state.set("vision_condition_frames", None) - state.set("num_frames", 5) - state.set("height", 32) - state.set("width", 32) - state.set("fps", 24.0) - state.set("latents", None) - state.set("generator", torch.Generator("cpu").manual_seed(0)) - state.set("use_fp32_sampling_state", input_enabled) - - Cosmos3VisionPrepareLatentsStep()(components, state) - - assert state.get("latents").dtype == expected_dtype - assert state.get("vision_condition_mask").dtype == expected_dtype - assert state.get("use_fp32_sampling_state") is (config_enabled if input_enabled is None else input_enabled) + + assert outputs["latents"].dtype == expected_dtype + assert outputs["vision_condition_mask"].dtype == expected_dtype + assert outputs["use_fp32_sampling_state"] is (config_enabled if input_enabled is None else input_enabled) @pytest.mark.parametrize( ("use_fp32_sampling_state", "expected_dtype"), [(False, torch.bfloat16), (True, torch.float32)], ) def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_dtype): - scheduler = FlowMatchEulerDiscreteScheduler(stochastic_sampling=True) - scheduler.set_timesteps(sigmas=[1.0, 0.5]) - latents = torch.zeros((1, 2, 1, 1, 1), dtype=torch.bfloat16) - block_state = BlockState( - latents=latents, - velocity_vision=torch.zeros_like(latents), - vision_condition_mask=torch.zeros((1, 1, 1), dtype=expected_dtype), - vision_conditioning_latents=None, - vision_condition_indexes_for_pack=[], - generator=torch.Generator("cpu").manual_seed(0), - use_fp32_sampling_state=use_fp32_sampling_state, - ) + pipe = self.get_pipeline(torch_dtype=torch.bfloat16).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["use_fp32_sampling_state"] = use_fp32_sampling_state - Cosmos3DistilledVisionLoopSchedulerStep()( - BlockState(scheduler=scheduler), block_state, i=0, t=scheduler.timesteps[0] - ) + latents = pipe(**inputs, output=self.output_name) + + assert latents.dtype == expected_dtype + + def test_default_sea_cache_reuses_transformer_execution(self): + pipe = self.get_pipeline().to(torch_device) + pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6)) + + pipe(**self.get_dummy_inputs(), output=self.output_name) - assert block_state.latents.dtype == expected_dtype + stats = pipe.transformer.get_cache_stats() + assert stats["actual_reuses"] > 0 + assert stats["actual_full_executions"] > 0 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 335398402d09..ef1372f385b7 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -137,6 +137,27 @@ def test_default_sea_cache_can_be_disabled_and_reenabled(self): replacement_pipeline._maybe_enable_sea_cache() self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3) + replacement_config = SeaCacheConfig(threshold=0.1) + replacement_pipeline.enable_sea_cache(replacement_config) + self.assertIs(replacement_pipeline.transformer._cache_config, replacement_config) + self.assertEqual(replacement_pipeline.transformer._cache_config.threshold, 0.1) + self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3) + + def test_default_sea_cache_reuses_transformer_execution(self): + pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipeline.set_progress_bar_config(disable=None) + pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True) + pipeline.register_to_config(use_native_flow_schedule=True) + pipeline.enable_sea_cache(SeaCacheConfig(threshold=1e6)) + inputs = self.get_dummy_inputs(torch_device) + inputs.update(num_inference_steps=4, output_type="latent", enable_safety_check=False) + + pipeline(**inputs) + + stats = pipeline.transformer.get_cache_stats() + self.assertGreater(stats["actual_reuses"], 0) + self.assertGreater(stats["actual_full_executions"], 0) + 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) From 94078191a2988489d318e77b765e3f54cb11f518 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 31 Aug 2026 13:22:33 -0700 Subject: [PATCH 10/20] Remove SeaCache stats and ablation instrumentation --- src/diffusers/hooks/sea_cache.py | 159 +----------------- tests/hooks/test_sea_cache.py | 92 +--------- .../cosmos/test_modular_pipeline_cosmos3.py | 13 +- ...test_modular_pipeline_cosmos3_distilled.py | 13 +- tests/pipelines/cosmos/test_cosmos3.py | 13 +- tests/testing_utils.py | 32 ++++ 6 files changed, 58 insertions(+), 264 deletions(-) diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index 93e2164dfccc..19d9ce0817f6 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -12,12 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import inspect import math -import time from dataclasses import dataclass -from typing import Any, Callable, Literal, Sequence +from typing import Any, Callable, Literal import torch @@ -79,9 +77,6 @@ class SeaCacheConfig: 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. - gate_schedule (`Sequence[bool]`, *optional*): - Per-step full-compute decisions to replay for controlled residual-prediction ablations. SEA still evaluates - its natural decision and reports mismatches; invalid or unsafe calls always fail open to full compute. Example: ```python @@ -116,7 +111,6 @@ class SeaCacheConfig: [torch.nn.Module, tuple[Any, ...], dict[str, Any]], list[torch.Tensor] | None, ] = None - gate_schedule: Sequence[bool] = None def __post_init__(self): if not math.isfinite(self.threshold) or self.threshold < 0: @@ -151,10 +145,6 @@ def __post_init__(self): callback = getattr(self, name) if callback is not None and not callable(callback): raise TypeError(f"`{name}` must be callable or `None`.") - if self.gate_schedule is not None: - self.gate_schedule = tuple(self.gate_schedule) - if any(not isinstance(value, bool) for value in self.gate_schedule): - raise TypeError("`gate_schedule` must contain only boolean values.") @dataclass @@ -174,8 +164,6 @@ def __init__(self): self.previous_indicator: list[torch.Tensor] | None = None self.accumulated_distance = 0.0 self.consecutive_cached = 0 - self.max_consecutive_cached_observed = 0 - self.max_consecutive_forced_full = 0 self.skip_remaining = False self.full_execution_pending = False self.cacheable_execution = False @@ -184,7 +172,6 @@ def __init__(self): self.und_output: torch.Tensor | None = None self.cached_und_output: torch.Tensor | None = None self.cached_gen_residual: torch.Tensor | None = None - self.full_started_at: float | None = None def reset_forward(self): self.skip_remaining = False @@ -195,7 +182,6 @@ def reset_forward(self): self.und_output = None self.cached_und_output = None self.cached_gen_residual = None - self.full_started_at = None def reset_trajectory(self): self.history = [] @@ -207,8 +193,6 @@ def reset_trajectory(self): def reset(self): self.reset_trajectory() - self.max_consecutive_cached_observed = 0 - self.max_consecutive_forced_full = 0 self.reset_forward() @@ -219,20 +203,6 @@ def __init__(self): def reset(self): self.forward_metadata: _SeaCacheForwardMetadata | None = None - self.transformer_calls = 0 - self.gate_evaluations = 0 - self.gate_full_decisions = 0 - self.gate_skip_decisions = 0 - self.gate_trace: list[bool] = [] - self.gate_schedule_mismatches = 0 - self.num_full_steps = 0 - self.num_cached_steps = 0 - self.fail_open_calls = 0 - self.indicator_seconds = 0.0 - self.decision_seconds = 0.0 - self.full_seconds = 0.0 - self.branch_full_executions: dict[str, int] = {} - self.branch_reuses: dict[str, int] = {} def warn_once(self, message: str): if message not in self._warned_messages: @@ -240,7 +210,6 @@ def warn_once(self, message: str): self._warned_messages.add(message) def mark_fail_open(self, message: str): - self.fail_open_calls += 1 self.warn_once(message) def resolve_gate( @@ -254,7 +223,6 @@ def resolve_gate( if state.gate_key == gate_key: return state.gate_should_compute - self.gate_evaluations += 1 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.") @@ -299,37 +267,14 @@ def resolve_gate( should_compute = natural_should_compute - if config.gate_schedule is not None: - schedule_is_valid = len(config.gate_schedule) == metadata.num_inference_steps - scheduled_compute = bool(config.gate_schedule[metadata.step_index]) if schedule_is_valid else True - can_replay_skip = not (invalid_gate or is_retained or is_in_cache_end or is_first_observation) - if not schedule_is_valid: - self.mark_fail_open( - "SeaCache gate schedule length does not match the number of inference steps; running full." - ) - should_compute = True - elif scheduled_compute or can_replay_skip: - if natural_should_compute != scheduled_compute: - self.gate_schedule_mismatches += 1 - should_compute = scheduled_compute - else: - self.mark_fail_open("SeaCache gate schedule requested an unsafe cache hit; running full.") - should_compute = True - if is_max_consecutive: should_compute = True - state.max_consecutive_forced_full += 1 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] - self.gate_trace.append(should_compute) - if should_compute: - self.gate_full_decisions += 1 - else: - self.gate_skip_decisions += 1 return should_compute @@ -375,18 +320,11 @@ def _get_block_outputs( def _record_full_execution( config: SeaCacheConfig, - state_manager: StateManager, - shared_state: SeaCacheSharedState, state: SeaCacheContextState, gen_output: torch.Tensor, und_output: torch.Tensor | None, ) -> None: - shared_state.num_full_steps += 1 state.consecutive_cached = 0 - branch = state_manager._current_context - shared_state.branch_full_executions[branch] = shared_state.branch_full_executions.get(branch, 0) + 1 - if state.full_started_at is not None: - shared_state.full_seconds += time.perf_counter() - state.full_started_at if ( state.cacheable_execution and state.step_index is not None @@ -570,7 +508,6 @@ def __init__( shared_state: SeaCacheSharedState, metadata_callback: Callable, raw_vision_callback: Callable, - residual_boundary: str, ): super().__init__() self.config = config @@ -578,73 +515,9 @@ def __init__( self.shared_state = shared_state self.metadata_callback = metadata_callback self.raw_vision_callback = raw_vision_callback - self.residual_boundary = residual_boundary - self._last_stats: dict[str, Any] | None = None - - @property - def num_full_steps(self) -> int: - return self.shared_state.num_full_steps - - @property - def num_cached_steps(self) -> int: - return self.shared_state.num_cached_steps - - def stats(self) -> dict[str, Any]: - if self.shared_state.transformer_calls == 0 and self._last_stats is not None: - return copy.deepcopy(self._last_stats) - - persistent_cache_bytes = 0 - per_branch = {} - for branch, state in sorted(self.state_manager._state_cache.items()): - if state.previous_indicator is not None: - persistent_cache_bytes += sum( - value.numel() * value.element_size() for value in state.previous_indicator - ) - for _, und_output, gen_residual in state.history: - persistent_cache_bytes += und_output.numel() * und_output.element_size() - persistent_cache_bytes += gen_residual.numel() * gen_residual.element_size() - per_branch[branch] = { - "full_calls": self.shared_state.branch_full_executions.get(branch, 0), - "reuse_calls": self.shared_state.branch_reuses.get(branch, 0), - "max_consecutive_cached_observed": state.max_consecutive_cached_observed, - "max_consecutive_forced_full": state.max_consecutive_forced_full, - } - - opportunities = self.shared_state.num_full_steps + self.shared_state.num_cached_steps - return { - "indicator_source": self.config.indicator_source, - "residual_order": self.config.residual_order, - "residual_boundary": self.residual_boundary, - "max_consecutive_cached": self.config.max_consecutive_cached, - "max_consecutive_cached_observed": max( - (stats["max_consecutive_cached_observed"] for stats in per_branch.values()), default=0 - ), - "max_consecutive_forced_full": sum(stats["max_consecutive_forced_full"] for stats in per_branch.values()), - "transformer_calls": self.shared_state.transformer_calls, - "gate_evaluations": self.shared_state.gate_evaluations, - "gate_full_decisions": self.shared_state.gate_full_decisions, - "gate_skip_decisions": self.shared_state.gate_skip_decisions, - "gate_trace": list(self.shared_state.gate_trace), - "gate_schedule_replayed": self.config.gate_schedule is not None, - "gate_schedule_mismatches": self.shared_state.gate_schedule_mismatches, - "actual_full_executions": self.shared_state.num_full_steps, - "actual_reuses": self.shared_state.num_cached_steps, - "actual_reuse_rate": (self.shared_state.num_cached_steps / opportunities if opportunities else 0.0), - "fail_open_calls": self.shared_state.fail_open_calls, - "sea_indicator_seconds": self.shared_state.indicator_seconds, - "sea_decision_seconds": self.shared_state.decision_seconds, - "sea_seconds": (self.shared_state.indicator_seconds + self.shared_state.decision_seconds), - "full_seconds": self.shared_state.full_seconds, - "timing_note": "host wall time; CUDA work is asynchronous", - "persistent_cache_bytes": persistent_cache_bytes, - "branch_full_executions": dict(sorted(self.shared_state.branch_full_executions.items())), - "branch_reuses": dict(sorted(self.shared_state.branch_reuses.items())), - "per_branch": per_branch, - } def pre_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.forward_metadata = None - self.shared_state.transformer_calls += 1 if torch.is_grad_enabled(): self.shared_state.mark_fail_open( "SeaCache is inference-only; calls with autograd enabled run in fail-open mode." @@ -743,8 +616,6 @@ def post_forward(self, module: torch.nn.Module, output: Any) -> Any: return output def reset_state(self, module: torch.nn.Module): - if self.shared_state.transformer_calls > 0: - self._last_stats = self.stats() self.state_manager.reset() self.shared_state.reset() return module @@ -829,13 +700,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): forward_metadata = self.shared_state.forward_metadata if state is None or forward_metadata is None: - if state is not None: - state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) state.step_index = forward_metadata.step_index state.cacheable_execution = True - indicator_started = time.perf_counter() indicator_error_reported = False if _is_parameter_sharded(module): self.shared_state.mark_fail_open( @@ -852,14 +720,11 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): ) indicator = None indicator_error_reported = True - self.shared_state.indicator_seconds += time.perf_counter() - indicator_started 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." ) - decision_started = time.perf_counter() should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config) - self.shared_state.decision_seconds += time.perf_counter() - decision_started if should_compute or not state.history: if not should_compute: @@ -867,7 +732,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.mark_fail_open( "SeaCache selected a cache hit without residual history; running in fail-open mode." ) - state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) residual_history = state.history[-(self.config.residual_order + 1) :] @@ -889,7 +753,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.mark_fail_open( "SeaCache residual history changed shape, device, or dtype; running in fail-open mode." ) - state.full_started_at = time.perf_counter() return self.fn_ref.original_forward(*args, **kwargs) if self.config.residual_order == 1 and len(residual_history) >= 2: @@ -904,10 +767,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): state.cached_und_output = cached_und state.cached_gen_residual = cached_residual state.consecutive_cached += 1 - state.max_consecutive_cached_observed = max(state.max_consecutive_cached_observed, state.consecutive_cached) - self.shared_state.num_cached_steps += 1 - branch = self.state_manager._current_context - self.shared_state.branch_reuses[branch] = self.shared_state.branch_reuses.get(branch, 0) + 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) @@ -952,8 +811,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): hidden_states, encoder_hidden_states = _get_block_outputs(self._metadata, output) _record_full_execution( self.config, - self.state_manager, - self.shared_state, state, gen_output=hidden_states, und_output=encoder_hidden_states, @@ -1006,8 +863,6 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): _record_full_execution( self.config, - self.state_manager, - self.shared_state, state, gen_output=output, und_output=state.und_output, @@ -1055,7 +910,6 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None: "residual boundary." ) post_norm_boundary = post_norm_modules is not None - residual_boundary = "post_language_model_norm" if post_norm_boundary else "repeated_block_stack" blocks = [] for name, submodule in unwrapped_module.named_children(): @@ -1085,7 +939,6 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: shared_state, metadata_callback, raw_vision_callback, - residual_boundary, ), _SEA_CACHE_ROOT_HOOK, ) @@ -1134,13 +987,3 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: raise root_registry._child_registries_cache = None - - -def get_sea_cache_stats(module: torch.nn.Module) -> dict[str, Any]: - """Return statistics for the SeaCache instance currently attached to ``module``.""" - - registry = getattr(module, "_diffusers_hook", None) - root_hook = registry.get_hook(_SEA_CACHE_ROOT_HOOK) if registry is not None else None - if not isinstance(root_hook, SeaCacheRootHook): - raise ValueError("SeaCache is not enabled on this module.") - return root_hook.stats() diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py index 895bf299ccf4..5a31edcd564f 100644 --- a/tests/hooks/test_sea_cache.py +++ b/tests/hooks/test_sea_cache.py @@ -100,16 +100,11 @@ def _make_config(runtime, **kwargs): return SeaCacheConfig(**config_kwargs) -def _get_root_hook(model): - return model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) - - @torch.no_grad() def test_sea_cache_uses_independent_gates_and_histories_per_context(): runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} model = DummySeaTransformer() model.enable_cache(_make_config(runtime, threshold=0.5)) - root_hook = _get_root_hook(model) first_input = torch.ones(2, 2, 2, 1) with model.cache_context("cond"): @@ -134,25 +129,8 @@ def test_sea_cache_uses_independent_gates_and_histories_per_context(): torch.testing.assert_close(uncond_output, changed_uncond_input * 8) assert [block.calls for block in model.layers] == [3, 3, 3] assert model.layers[0].indicator_norm.calls == 4 - assert root_hook.num_full_steps == 3 - assert root_hook.num_cached_steps == 1 - stats = model.get_cache_stats() - assert stats["indicator_source"] == "first_block" - assert stats["residual_order"] == 1 - assert stats["residual_boundary"] == "repeated_block_stack" - assert stats["transformer_calls"] == 4 - assert stats["gate_evaluations"] == 4 - assert stats["gate_trace"] == [True, True, False, True] - assert stats["branch_full_executions"] == {"cond": 1, "uncond": 2} - assert stats["branch_reuses"] == {"cond": 1} - assert stats["persistent_cache_bytes"] > 0 model._reset_stateful_cache() - assert root_hook.num_full_steps == 0 - assert root_hook.num_cached_steps == 0 - archived_stats = model.get_cache_stats() - assert archived_stats["transformer_calls"] == 4 - assert archived_stats["gate_trace"] == [True, True, False, True] model.disable_cache() assert not model.is_cache_enabled @@ -177,15 +155,7 @@ def test_sea_cache_max_consecutive_cached_forces_full_per_context(): with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) - stats = model.get_cache_stats() - assert stats["gate_trace"] == [True, False, False, True, False, False, True, False] - assert stats["actual_full_executions"] == 3 - assert stats["actual_reuses"] == 5 - assert stats["max_consecutive_cached"] == 2 - assert stats["max_consecutive_cached_observed"] == 2 - assert stats["max_consecutive_forced_full"] == 2 - assert stats["per_branch"]["cond"]["max_consecutive_cached_observed"] == 2 - assert stats["per_branch"]["cond"]["max_consecutive_forced_full"] == 2 + # With max_consecutive_cached=2, every third step forces a full execution: 3 full, 5 cached. assert [block.calls for block in model.layers] == [3, 3, 3] @@ -218,11 +188,7 @@ def raw_vision(module, args, kwargs): with model.cache_context("cond"): model(second_input) - stats = model.get_cache_stats() - assert stats["indicator_source"] == "raw_vision_latents" - assert stats["gate_trace"] == [True, True] - assert stats["actual_full_executions"] == 2 - assert stats["actual_reuses"] == 0 + assert [block.calls for block in model.layers] == [2, 2, 2] assert model.layers[0].indicator_norm.calls == 0 @@ -247,34 +213,7 @@ def test_sea_cache_residual_order_one_uses_actual_full_step_history(): # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2. torch.testing.assert_close(output, torch.full_like(output, 24.0)) - root_hook = _get_root_hook(model) - assert root_hook.num_full_steps == 2 - assert root_hook.num_cached_steps == 1 - - -@torch.no_grad() -def test_sea_cache_replays_gate_schedule_and_reports_natural_mismatches(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} - model = DummySeaTransformer() - model.enable_cache( - _make_config( - runtime, - threshold=0.0, - gate_schedule=(True, False, True), - ) - ) - - for step, sigma in enumerate((0.9, 0.6, 0.3)): - runtime.update(step=step, sigma=sigma) - with model.cache_context("cond"): - model(torch.full((1, 1, 1, 1), float(step + 1))) - - stats = model.get_cache_stats() - assert stats["gate_trace"] == [True, False, True] - assert stats["gate_schedule_replayed"] - assert stats["gate_schedule_mismatches"] == 1 - assert stats["actual_full_executions"] == 2 - assert stats["actual_reuses"] == 1 + assert [block.calls for block in model.layers] == [2, 2, 2] @torch.no_grad() @@ -290,7 +229,7 @@ def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_unca with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) - assert model.get_cache_stats()["branch_full_executions"] == {"cond": 1} + assert [block.calls for block in model.layers] == [2, 2, 2] @torch.no_grad() @@ -308,9 +247,6 @@ def test_sea_cache_fails_open_without_vision_metadata(): model(torch.ones(1, 1, 1, 1)) assert [block.calls for block in model.layers] == [2, 2, 2] - root_hook = _get_root_hook(model) - assert root_hook.num_full_steps == 2 - assert root_hook.num_cached_steps == 0 @torch.no_grad() @@ -336,7 +272,6 @@ def test_sea_cache_non_adjacent_steps_start_a_new_residual_trajectory(): torch.testing.assert_close(output, torch.full_like(output, 16.0)) assert [block.calls for block in model.layers] == [3, 3, 3] - assert model.get_cache_stats()["gate_trace"] == [True, True, True, False] @torch.no_grad() @@ -352,9 +287,6 @@ def test_sea_cache_fails_open_for_shape_changes(): with model.cache_context("cond"): model(torch.ones(1, 2, 1, 1)) - root_hook = _get_root_hook(model) - assert root_hook.num_full_steps == 2 - assert root_hook.num_cached_steps == 0 assert [block.calls for block in model.layers] == [2, 2, 2] @@ -369,11 +301,6 @@ def test_sea_cache_is_inference_only_and_fails_open_with_autograd(): with torch.enable_grad(), model.cache_context("cond"): model(torch.ones(1, 1, 1, 1, requires_grad=True)) - root_hook = _get_root_hook(model) - assert root_hook.shared_state.transformer_calls == 2 - assert root_hook.shared_state.fail_open_calls == 2 - assert root_hook.num_full_steps == 2 - assert root_hook.num_cached_steps == 0 assert [block.calls for block in model.layers] == [2, 2, 2] @@ -422,8 +349,6 @@ def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables _, cached_output = model(torch.full((1, 1, 1, 1), 2.0)) torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0)) assert model.layers[0].calls == 1 - assert model.get_cache_stats()["actual_full_executions"] == 1 - assert model.get_cache_stats()["actual_reuses"] == 1 model.disable_cache() @@ -447,10 +372,6 @@ def test_sea_cache_fails_open_for_parameter_sharded_blocks(): with model.cache_context("cond"): model(torch.ones(1, 1, 1, 1)) - stats = model.get_cache_stats() - assert stats["actual_full_executions"] == 2 - assert stats["actual_reuses"] == 0 - assert stats["fail_open_calls"] == 2 assert [block.calls for block in model.layers] == [2, 2, 2] @@ -501,11 +422,6 @@ def test_sea_cache_config_validation(kwargs, message): SeaCacheConfig(**kwargs) -def test_sea_cache_gate_schedule_validation(): - with pytest.raises(TypeError, match="gate_schedule"): - SeaCacheConfig(gate_schedule=(True, 1)) - - @pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"]) def test_sea_cache_callback_validation(callback_name): with pytest.raises(TypeError, match=callback_name): diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index d3afd2e33a97..b4f9c78d6622 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -22,6 +22,7 @@ from PIL import Image from diffusers import CosmosActionCondition, ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler +from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer from diffusers.modular_pipelines import ( Cosmos3OmniBlocks, Cosmos3OmniModularPipeline, @@ -41,7 +42,7 @@ 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 count_transformer_cache_reuse, torch_device from ..testing_utils import ( BaseModularPipelineTesterConfig, ModularLoadingTesterMixin, @@ -213,15 +214,15 @@ def test_default_sea_cache_reuses_transformer_execution(self): pipe = self.get_pipeline().to(torch_device) scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, use_flow_sigmas=True) pipe.update_components(scheduler=scheduler, use_native_flow_schedule=True) - pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6)) inputs = self.get_dummy_inputs() inputs["num_inference_steps"] = 4 - pipe(**inputs, output=self.output_name) + with count_transformer_cache_reuse(pipe.transformer, Cosmos3VLTextMoTDecoderLayer) as counts: + pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6)) + pipe(**inputs, output=self.output_name) - stats = pipe.transformer.get_cache_stats() - assert stats["actual_reuses"] > 0 - assert stats["actual_full_executions"] > 0 + assert counts["cached_steps"] > 0 + assert counts["full_steps"] > 0 def _get_sampling_state_block_pipe(self, block): pipe = block.init_pipeline(self.pretrained_model_name_or_path) 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 eb94dc80b010..705ca74fd4d3 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -18,11 +18,12 @@ from PIL import Image from diffusers import ModularPipeline, SeaCacheConfig +from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer 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 count_transformer_cache_reuse, torch_device from ..testing_utils import ( BaseModularPipelineTesterConfig, ModularLoadingTesterMixin, @@ -185,13 +186,13 @@ def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_ def test_default_sea_cache_reuses_transformer_execution(self): pipe = self.get_pipeline().to(torch_device) - pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6)) - pipe(**self.get_dummy_inputs(), output=self.output_name) + with count_transformer_cache_reuse(pipe.transformer, Cosmos3VLTextMoTDecoderLayer) as counts: + pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6)) + pipe(**self.get_dummy_inputs(), output=self.output_name) - stats = pipe.transformer.get_cache_stats() - assert stats["actual_reuses"] > 0 - assert stats["actual_full_executions"] > 0 + assert counts["cached_steps"] > 0 + assert counts["full_steps"] > 0 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 ef1372f385b7..2ae7e6f37aa9 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -28,9 +28,10 @@ SeaCacheConfig, UniPCMultistepScheduler, ) +from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import count_transformer_cache_reuse, enable_full_determinism, torch_device from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin @@ -148,15 +149,15 @@ def test_default_sea_cache_reuses_transformer_execution(self): pipeline.set_progress_bar_config(disable=None) pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True) pipeline.register_to_config(use_native_flow_schedule=True) - pipeline.enable_sea_cache(SeaCacheConfig(threshold=1e6)) inputs = self.get_dummy_inputs(torch_device) inputs.update(num_inference_steps=4, output_type="latent", enable_safety_check=False) - pipeline(**inputs) + with count_transformer_cache_reuse(pipeline.transformer, Cosmos3VLTextMoTDecoderLayer) as counts: + pipeline.enable_sea_cache(SeaCacheConfig(threshold=1e6)) + pipeline(**inputs) - stats = pipeline.transformer.get_cache_stats() - self.assertGreater(stats["actual_reuses"], 0) - self.assertGreater(stats["actual_full_executions"], 0) + self.assertGreater(counts["cached_steps"], 0) + self.assertGreater(counts["full_steps"], 0) def test_fp32_sampling_state_keeps_transformer_inputs_in_model_dtype(self): pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device) diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 9da89e626198..e99350d87dc3 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -197,6 +197,38 @@ def numpy_cosine_similarity_distance(a, b): return distance +@contextmanager +def count_transformer_cache_reuse(transformer, block_cls): + """Count real block executions and total transformer calls while running under a cache. + + ``block_cls.forward`` is patched to count real executions (skipped blocks reuse cached outputs and never run), + and a forward pre-hook counts every transformer call. On exit the yielded dict is populated with ``full_steps`` + (transformer calls that recomputed the blocks) and ``cached_steps`` (calls served from the cache). + + The patch must be active before the cache is enabled, so enable the cache inside the ``with`` block. + """ + counts = {"block_calls": 0, "transformer_calls": 0, "full_steps": 0, "cached_steps": 0} + original_forward = block_cls.forward + + def counting_forward(self, *args, **kwargs): + counts["block_calls"] += 1 + return original_forward(self, *args, **kwargs) + + def counting_pre_hook(module, args): + counts["transformer_calls"] += 1 + + block_cls.forward = counting_forward + handle = transformer.register_forward_pre_hook(counting_pre_hook) + try: + yield counts + finally: + handle.remove() + block_cls.forward = original_forward + num_blocks = sum(1 for module in transformer.modules() if isinstance(module, block_cls)) + counts["full_steps"] = counts["block_calls"] // num_blocks if num_blocks else 0 + counts["cached_steps"] = counts["transformer_calls"] - counts["full_steps"] + + def check_if_dicts_are_equal(dict1, dict2): dict1, dict2 = dict1.copy(), dict2.copy() From 1dd8308c0abcc5ca35f395f5c825ab0dd0172810 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Wed, 2 Sep 2026 07:30:33 -0700 Subject: [PATCH 11/20] Move SeaCache to transformer-level and disable by default --- docs/source/en/api/pipelines/cosmos3.md | 35 ++++-------- docs/source/en/optimization/cache.md | 29 ++++------ src/diffusers/hooks/sea_cache.py | 3 + .../cosmos/modular_pipeline.py | 15 ----- .../pipelines/cosmos/pipeline_cosmos3_omni.py | 55 ------------------- .../cosmos/test_modular_pipeline_cosmos3.py | 34 +----------- ...test_modular_pipeline_cosmos3_distilled.py | 15 +---- tests/pipelines/cosmos/test_cosmos3.py | 45 +-------------- tests/testing_utils.py | 32 ----------- 9 files changed, 30 insertions(+), 233 deletions(-) diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index 84752ace65b8..dcfb7167db90 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -662,41 +662,28 @@ if result.action is not None: ## SeaCache -Cosmos 3 inference uses [`SeaCacheConfig`] by default. SeaCache reuses transformer residuals when the -Spectral-Evolution-Aware indicator changes slowly, reducing the number of full transformer executions. The default -configuration filters raw vision latents, linearly extrapolates cached residuals, uses a `0.25` threshold, and forces a -full execution after at most two consecutive cached steps. - -SeaCache is approximate and can change generated outputs. Disable it before inference when you need every denoising -step to execute the full transformer, for example when producing an uncached baseline: - -```python -pipe.disable_sea_cache() -result = pipe( - prompt=prompt, - num_frames=189, - height=720, - width=1280, -) -``` - -The same method works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and -[`Cosmos3DistilledModularPipeline`]. Reenable the default configuration with `pipe.enable_sea_cache()`, or pass a -custom configuration: +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 from diffusers import SeaCacheConfig -pipe.enable_sea_cache( +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 pipeline supplies SeaCache with the active scheduler step, sigma, and number of inference steps. Cache state is -reset for each pipeline call, and conditional and unconditional guidance branches keep independent histories. +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. ## Sampling precision diff --git a/docs/source/en/optimization/cache.md b/docs/source/en/optimization/cache.md index 4db41753507f..ad65fceeaada 100644 --- a/docs/source/en/optimization/cache.md +++ b/docs/source/en/optimization/cache.md @@ -74,35 +74,28 @@ pipeline.transformer.enable_cache(config) 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. -[`Cosmos3OmniPipeline`], [`Cosmos3OmniModularPipeline`], and [`Cosmos3DistilledModularPipeline`] enable SeaCache -automatically when inference starts. The default [`SeaCacheConfig`] filters raw vision latents, linearly extrapolates -the cached residual, uses `threshold=0.25`, and allows at most two consecutive cached steps before forcing a full -transformer execution. - -SeaCache is an approximate optimization and may change generated outputs. Disable it for full transformer execution, -such as when measuring an uncached baseline: +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 +from diffusers import Cosmos3OmniPipeline, SeaCacheConfig pipe = Cosmos3OmniPipeline.from_pretrained("nvidia/Cosmos3-Nano") -pipe.disable_sea_cache() -``` - -Call [`~Cosmos3OmniPipeline.enable_sea_cache`] to reenable it or provide a custom configuration. The pipeline fills in -the scheduler callbacks required by SeaCache. - -```python -from diffusers import SeaCacheConfig - -pipe.enable_sea_cache( +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/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index 19d9ce0817f6..f277f90f0ce7 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -39,6 +39,9 @@ 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 diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py index 5121e331306e..fe6f3cbd721d 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py @@ -30,23 +30,8 @@ def current_sigma(self): def num_timesteps(self): return getattr(self, "_num_timesteps", None) - def _prepare_sea_cache_config(self, config=None): - return Cosmos3OmniPipeline._prepare_sea_cache_config(self, config) - - def enable_sea_cache(self, config=None): - return Cosmos3OmniPipeline.enable_sea_cache(self, config) - - def disable_sea_cache(self): - return Cosmos3OmniPipeline.disable_sea_cache(self) - - def _maybe_enable_sea_cache(self): - return Cosmos3OmniPipeline._maybe_enable_sea_cache(self) - def __call__(self, *args, **kwargs): - self._maybe_enable_sea_cache() transformer = getattr(self, "transformer", None) - if hasattr(transformer, "_reset_stateful_cache"): - transformer._reset_stateful_cache() try: return super().__call__(*args, **kwargs) finally: diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 3847f06b761c..05e0ac669eed 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -27,7 +27,6 @@ from transformers import AutoTokenizer, BatchEncoding from ...callbacks import MultiPipelineCallbacks, PipelineCallback -from ...hooks import SeaCacheConfig from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan from ...models.modeling_utils import get_parameter_device @@ -474,56 +473,6 @@ 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 "". - self._is_sea_cache_enabled = True - - def _prepare_sea_cache_config(self, config: SeaCacheConfig | None = None) -> SeaCacheConfig: - config = SeaCacheConfig() if config is None else config - config.current_step_callback = lambda: self.current_step_index - config.current_sigma_callback = lambda: self.current_sigma - config.num_inference_steps_callback = lambda: self.num_timesteps - return config - - def enable_sea_cache(self, config: SeaCacheConfig | None = None) -> None: - """Enable SeaCache for subsequent pipeline calls.""" - transformer = getattr(self, "transformer", None) - if transformer is None: - raise ValueError("SeaCache requires a loaded transformer.") - current_config = getattr(transformer, "_cache_config", None) - if current_config is not None: - if isinstance(current_config, SeaCacheConfig): - if config is None or config is current_config: - self._prepare_sea_cache_config(current_config) - self._is_sea_cache_enabled = True - return - transformer.disable_cache() - else: - raise ValueError( - f"Caching is already enabled with {type(current_config).__name__}. Disable it before enabling SeaCache." - ) - - transformer.enable_cache(self._prepare_sea_cache_config(config)) - self._is_sea_cache_enabled = True - - def disable_sea_cache(self) -> None: - """Disable the default SeaCache optimization for subsequent pipeline calls.""" - self._is_sea_cache_enabled = False - transformer = getattr(self, "transformer", None) - if isinstance(getattr(transformer, "_cache_config", None), SeaCacheConfig): - transformer.disable_cache() - - def _maybe_enable_sea_cache(self) -> None: - transformer = getattr(self, "transformer", None) - if transformer is None or not getattr(self, "_is_sea_cache_enabled", True): - return - if isinstance(getattr(transformer, "_cache_config", None), SeaCacheConfig): - self._prepare_sea_cache_config(transformer._cache_config) - return - if transformer.is_cache_enabled: - return - self.enable_sea_cache() - # 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 @@ -1518,10 +1467,6 @@ def __call__( `sound` (`torch.Tensor` of shape `[C, N]`, or `None` when `enable_sound=False`). Otherwise a tuple `(video, sound)` with the same fields. """ - self._maybe_enable_sea_cache() - if hasattr(self.transformer, "_reset_stateful_cache"): - self.transformer._reset_stateful_cache() - if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index b4f9c78d6622..013940fd6970 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -21,8 +21,7 @@ import torch from PIL import Image -from diffusers import CosmosActionCondition, ModularPipeline, SeaCacheConfig, UniPCMultistepScheduler -from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer +from diffusers import CosmosActionCondition, ModularPipeline, UniPCMultistepScheduler from diffusers.modular_pipelines import ( Cosmos3OmniBlocks, Cosmos3OmniModularPipeline, @@ -42,7 +41,7 @@ from diffusers.modular_pipelines.cosmos.encoders import Cosmos3TextEncoderStep from diffusers.modular_pipelines.cosmos.modular_blocks_cosmos3 import Cosmos3TransferChunkDenoiseStep -from ...testing_utils import count_transformer_cache_reuse, torch_device +from ...testing_utils import torch_device from ..testing_utils import ( BaseModularPipelineTesterConfig, ModularLoadingTesterMixin, @@ -194,35 +193,6 @@ def record_context(name): torch.testing.assert_close(sigma, pipe.scheduler.sigmas[expected_step]) assert pipe.current_step_index is None assert pipe.current_sigma is None - assert isinstance(pipe.transformer._cache_config, SeaCacheConfig) - - def test_default_sea_cache_can_be_disabled_and_reenabled(self): - pipe = self.get_pipeline() - pipe._maybe_enable_sea_cache() - - assert isinstance(pipe.transformer._cache_config, SeaCacheConfig) - pipe.disable_sea_cache() - assert not pipe.transformer.is_cache_enabled - - pipe._maybe_enable_sea_cache() - assert not pipe.transformer.is_cache_enabled - - pipe.enable_sea_cache() - assert isinstance(pipe.transformer._cache_config, SeaCacheConfig) - - def test_default_sea_cache_reuses_transformer_execution(self): - pipe = self.get_pipeline().to(torch_device) - scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, use_flow_sigmas=True) - pipe.update_components(scheduler=scheduler, use_native_flow_schedule=True) - inputs = self.get_dummy_inputs() - inputs["num_inference_steps"] = 4 - - with count_transformer_cache_reuse(pipe.transformer, Cosmos3VLTextMoTDecoderLayer) as counts: - pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6)) - pipe(**inputs, output=self.output_name) - - assert counts["cached_steps"] > 0 - assert counts["full_steps"] > 0 def _get_sampling_state_block_pipe(self, block): pipe = block.init_pipeline(self.pretrained_model_name_or_path) 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 705ca74fd4d3..354b9b08288f 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -17,13 +17,12 @@ import torch from PIL import Image -from diffusers import ModularPipeline, SeaCacheConfig -from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer +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 count_transformer_cache_reuse, torch_device +from ...testing_utils import torch_device from ..testing_utils import ( BaseModularPipelineTesterConfig, ModularLoadingTesterMixin, @@ -184,16 +183,6 @@ def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_ assert latents.dtype == expected_dtype - def test_default_sea_cache_reuses_transformer_execution(self): - pipe = self.get_pipeline().to(torch_device) - - with count_transformer_cache_reuse(pipe.transformer, Cosmos3VLTextMoTDecoderLayer) as counts: - pipe.enable_sea_cache(SeaCacheConfig(threshold=1e6)) - pipe(**self.get_dummy_inputs(), output=self.output_name) - - assert counts["cached_steps"] > 0 - assert counts["full_steps"] > 0 - def test_vae_encoder_rejects_image_and_video_together(self): vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"] vae_pipe = vae_encoder.init_pipeline(self.pretrained_model_name_or_path) diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index 2ae7e6f37aa9..c08caf346319 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -25,13 +25,11 @@ AutoencoderKLWan, Cosmos3OmniPipeline, Cosmos3OmniTransformer, - SeaCacheConfig, UniPCMultistepScheduler, ) -from diffusers.models.transformers.transformer_cosmos3 import Cosmos3VLTextMoTDecoderLayer from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image -from ...testing_utils import count_transformer_cache_reuse, enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin @@ -118,47 +116,6 @@ def test_inference(self): assert video[0].shape == self.output_shape - def test_default_sea_cache_can_be_disabled_and_reenabled(self): - components = self.get_dummy_components() - pipeline = self.pipeline_class(**components) - pipeline._maybe_enable_sea_cache() - - self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig) - pipeline.disable_sea_cache() - self.assertFalse(pipeline.transformer.is_cache_enabled) - - pipeline._maybe_enable_sea_cache() - self.assertFalse(pipeline.transformer.is_cache_enabled) - - pipeline.enable_sea_cache() - self.assertIsInstance(pipeline.transformer._cache_config, SeaCacheConfig) - - replacement_pipeline = self.pipeline_class(**components) - replacement_pipeline._current_step_index = 3 - replacement_pipeline._maybe_enable_sea_cache() - self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3) - - replacement_config = SeaCacheConfig(threshold=0.1) - replacement_pipeline.enable_sea_cache(replacement_config) - self.assertIs(replacement_pipeline.transformer._cache_config, replacement_config) - self.assertEqual(replacement_pipeline.transformer._cache_config.threshold, 0.1) - self.assertEqual(replacement_pipeline.transformer._cache_config.current_step_callback(), 3) - - def test_default_sea_cache_reuses_transformer_execution(self): - pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - pipeline.set_progress_bar_config(disable=None) - pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config, use_flow_sigmas=True) - pipeline.register_to_config(use_native_flow_schedule=True) - inputs = self.get_dummy_inputs(torch_device) - inputs.update(num_inference_steps=4, output_type="latent", enable_safety_check=False) - - with count_transformer_cache_reuse(pipeline.transformer, Cosmos3VLTextMoTDecoderLayer) as counts: - pipeline.enable_sea_cache(SeaCacheConfig(threshold=1e6)) - pipeline(**inputs) - - self.assertGreater(counts["cached_steps"], 0) - self.assertGreater(counts["full_steps"], 0) - 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) diff --git a/tests/testing_utils.py b/tests/testing_utils.py index e99350d87dc3..9da89e626198 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -197,38 +197,6 @@ def numpy_cosine_similarity_distance(a, b): return distance -@contextmanager -def count_transformer_cache_reuse(transformer, block_cls): - """Count real block executions and total transformer calls while running under a cache. - - ``block_cls.forward`` is patched to count real executions (skipped blocks reuse cached outputs and never run), - and a forward pre-hook counts every transformer call. On exit the yielded dict is populated with ``full_steps`` - (transformer calls that recomputed the blocks) and ``cached_steps`` (calls served from the cache). - - The patch must be active before the cache is enabled, so enable the cache inside the ``with`` block. - """ - counts = {"block_calls": 0, "transformer_calls": 0, "full_steps": 0, "cached_steps": 0} - original_forward = block_cls.forward - - def counting_forward(self, *args, **kwargs): - counts["block_calls"] += 1 - return original_forward(self, *args, **kwargs) - - def counting_pre_hook(module, args): - counts["transformer_calls"] += 1 - - block_cls.forward = counting_forward - handle = transformer.register_forward_pre_hook(counting_pre_hook) - try: - yield counts - finally: - handle.remove() - block_cls.forward = original_forward - num_blocks = sum(1 for module in transformer.modules() if isinstance(module, block_cls)) - counts["full_steps"] = counts["block_calls"] // num_blocks if num_blocks else 0 - counts["cached_steps"] = counts["transformer_calls"] - counts["full_steps"] - - def check_if_dicts_are_equal(dict1, dict2): dict1, dict2 = dict1.copy(), dict2.copy() From c403025fcdb03eb71211dcf93f66b52c66bc7362 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Wed, 2 Sep 2026 08:00:36 -0700 Subject: [PATCH 12/20] Align SeaCache and tests with model-level cache patterns --- .../modular_pipelines/cosmos/denoise.py | 10 +- .../pipelines/cosmos/pipeline_cosmos3_omni.py | 10 +- tests/hooks/test_sea_cache.py | 428 ------------------ tests/models/testing_utils/__init__.py | 4 + tests/models/testing_utils/cache.py | 144 ++++++ .../test_models_transformer_cosmos3.py | 5 + 6 files changed, 157 insertions(+), 444 deletions(-) delete mode 100644 tests/hooks/test_sea_cache.py diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index 5b06b4f910a2..8cf4b377198f 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -1,5 +1,4 @@ import inspect -from contextlib import nullcontext import torch @@ -15,11 +14,6 @@ from .modular_pipeline import Cosmos3OmniModularPipeline -def _cache_context(transformer: torch.nn.Module, name: str): - cache_context = getattr(transformer, "cache_context", None) - return cache_context(name) if callable(cache_context) else nullcontext() - - class Cosmos3VisionLoopPrepareStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @@ -227,7 +221,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta transformer_kwargs = { name: value for name, value in transformer_kwargs.items() if name in transformer_args } - with _cache_context(components.transformer, pass_name): + with components.transformer.cache_context(pass_name): preds_vision, preds_sound, preds_action = components.transformer( **transformer_kwargs, return_dict=False ) @@ -754,7 +748,7 @@ def intermediate_outputs(self) -> list[OutputParam]: @staticmethod def _forward(components, static, vision_tokens, vision_timesteps, context_name): - with _cache_context(components.transformer, context_name): + with components.transformer.cache_context(context_name): preds_vision, _, _ = components.transformer( input_ids=static["input_ids"], text_indexes=static["text_indexes"], diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 05e0ac669eed..19657ae0fe33 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -16,7 +16,6 @@ import json import math from collections.abc import Iterable -from contextlib import nullcontext from dataclasses import dataclass from typing import Any, Callable, Literal @@ -43,11 +42,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def _cache_context(transformer: torch.nn.Module, name: str): - cache_context = getattr(transformer, "cache_context", None) - return cache_context(name) if callable(cache_context) else nullcontext() - - if is_cosmos_guardrail_available(): from cosmos_guardrail import CosmosSafetyChecker else: @@ -1739,7 +1733,7 @@ def __call__( ) # --- Conditional pass --- - with _cache_context(self.transformer, "cond"): + 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"], @@ -1780,7 +1774,7 @@ 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: - with _cache_context(self.transformer, "uncond"): + 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"], diff --git a/tests/hooks/test_sea_cache.py b/tests/hooks/test_sea_cache.py deleted file mode 100644 index 5a31edcd564f..000000000000 --- a/tests/hooks/test_sea_cache.py +++ /dev/null @@ -1,428 +0,0 @@ -# Copyright 2026 HuggingFace Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import pytest -import torch - -from diffusers import SeaCacheConfig -from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry -from diffusers.hooks.hooks import HookRegistry, ModelHook -from diffusers.hooks.sea_cache import ( - _SEA_CACHE_BLOCK_HOOK, - _SEA_CACHE_LEADER_BLOCK_HOOK, - _SEA_CACHE_ROOT_HOOK, - _apply_sea_filter, -) -from diffusers.models.cache_utils import CacheMixin - - -class CountingIdentity(torch.nn.Module): - def __init__(self): - super().__init__() - self.calls = 0 - - def forward(self, hidden_states): - self.calls += 1 - return hidden_states - - -class DummySeaBlock(torch.nn.Module): - def __init__(self): - super().__init__() - self.indicator_norm = CountingIdentity() - self.calls = 0 - - def forward(self, und_seq, gen_seq, rotary_emb=None): - self.calls += 1 - return und_seq + 1, gen_seq * 2 - - -class DummySeaTransformer(torch.nn.Module, CacheMixin): - def __init__(self, num_layers=3): - super().__init__() - self.layers = torch.nn.ModuleList([DummySeaBlock() for _ in range(num_layers)]) - - def forward(self, hidden_states): - shape = hidden_states.shape - gen_seq = hidden_states.reshape(-1, shape[-1]) - und_seq = torch.zeros(1, shape[-1], device=hidden_states.device, dtype=hidden_states.dtype) - for layer in self.layers: - und_seq, gen_seq = layer(und_seq, gen_seq) - return und_seq, gen_seq.reshape(shape) - - -@pytest.fixture(autouse=True) -def register_dummy_sea_block(): - TransformerBlockRegistry.register( - DummySeaBlock, - 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="indicator_norm", - ), - ) - - -def _metadata_callback(module, args, kwargs): - hidden_states = kwargs.get("hidden_states", args[0] if args else None) - if hidden_states is None: - return None - temporal, height, width = hidden_states.shape[:3] - indexes = torch.arange(temporal * height * width, device=hidden_states.device) - return [(indexes, (temporal, height, width))] - - -def _make_config(runtime, **kwargs): - config_kwargs = { - "threshold": 100.0, - "retention_steps": 1, - "cache_end_steps": 0, - "indicator_source": "first_block", - "current_step_callback": lambda: runtime["step"], - "current_sigma_callback": lambda: runtime["sigma"], - "num_inference_steps_callback": lambda: runtime["num_steps"], - "metadata_callback": _metadata_callback, - } - config_kwargs.update(kwargs) - return SeaCacheConfig(**config_kwargs) - - -@torch.no_grad() -def test_sea_cache_uses_independent_gates_and_histories_per_context(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} - model = DummySeaTransformer() - model.enable_cache(_make_config(runtime, threshold=0.5)) - - first_input = torch.ones(2, 2, 2, 1) - with model.cache_context("cond"): - _, first_output = model(first_input) - torch.testing.assert_close(first_output, first_input * 8) - assert [block.calls for block in model.layers] == [1, 1, 1] - - with model.cache_context("uncond"): - _, first_uncond_output = model(first_input) - torch.testing.assert_close(first_uncond_output, first_input * 8) - assert [block.calls for block in model.layers] == [2, 2, 2] - - runtime.update(step=1) - with model.cache_context("cond"): - _, cached_output = model(first_input) - torch.testing.assert_close(cached_output, first_input * 8) - assert [block.calls for block in model.layers] == [2, 2, 2] - - changed_uncond_input = first_input * 100 - with model.cache_context("uncond"): - _, uncond_output = model(changed_uncond_input) - torch.testing.assert_close(uncond_output, changed_uncond_input * 8) - assert [block.calls for block in model.layers] == [3, 3, 3] - assert model.layers[0].indicator_norm.calls == 4 - - model._reset_stateful_cache() - - model.disable_cache() - assert not model.is_cache_enabled - - -@torch.no_grad() -def test_sea_cache_max_consecutive_cached_forces_full_per_context(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 8} - model = DummySeaTransformer() - model.enable_cache( - _make_config( - runtime, - threshold=100.0, - retention_steps=0, - cache_end_steps=0, - max_consecutive_cached=2, - ) - ) - - for step in range(runtime["num_steps"]): - runtime.update(step=step, sigma=0.9 - step * 0.1) - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - - # With max_consecutive_cached=2, every third step forces a full execution: 3 full, 5 cached. - assert [block.calls for block in model.layers] == [3, 3, 3] - - -@torch.no_grad() -def test_sea_cache_raw_vision_indicator_includes_conditioning_frames(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - model = DummySeaTransformer() - - def raw_vision(module, args, kwargs): - hidden_states = kwargs.get("hidden_states", args[0] if args else None) - latent = hidden_states.permute(3, 0, 1, 2) - return [latent] - - model.enable_cache( - _make_config( - runtime, - threshold=1e-6, - indicator_source="raw_vision_latents", - raw_vision_callback=raw_vision, - ) - ) - - first_input = torch.tensor([1.0, 2.0]).reshape(2, 1, 1, 1) - with model.cache_context("cond"): - model(first_input) - - # The noisy frame stays fixed, but changing the conditioning frame changes the complete raw-latent indicator. - runtime.update(step=1, sigma=0.9) - second_input = torch.tensor([100.0, 2.0]).reshape(2, 1, 1, 1) - with model.cache_context("cond"): - model(second_input) - - assert [block.calls for block in model.layers] == [2, 2, 2] - assert model.layers[0].indicator_norm.calls == 0 - - -@torch.no_grad() -def test_sea_cache_residual_order_one_uses_actual_full_step_history(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} - config = _make_config(runtime, residual_order=1, threshold=0.0) - model = DummySeaTransformer() - model.enable_cache(config) - - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - - runtime.update(step=1, sigma=0.7) - with model.cache_context("cond"): - model(torch.full((1, 1, 1, 1), 2.0)) - - config.threshold = 100.0 - runtime.update(step=2, sigma=0.5) - with model.cache_context("cond"): - _, output = model(torch.full((1, 1, 1, 1), 3.0)) - - # Full residuals are 7 and 14 at steps 0 and 1, so linear extrapolation predicts 21 at step 2. - torch.testing.assert_close(output, torch.full_like(output, 24.0)) - assert [block.calls for block in model.layers] == [2, 2, 2] - - -@torch.no_grad() -def test_cache_context_registry_is_refreshed_when_cache_is_enabled_after_an_uncached_call(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - model = DummySeaTransformer() - - # Pipeline cache contexts may be entered before a cache is enabled (for example, during the baseline). - with model.cache_context("baseline"): - model(torch.ones(1, 1, 1, 1)) - - model.enable_cache(_make_config(runtime)) - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - - assert [block.calls for block in model.layers] == [2, 2, 2] - - -@torch.no_grad() -def test_sea_cache_fails_open_without_vision_metadata(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - config = _make_config(runtime) - config.metadata_callback = lambda module, args, kwargs: None - model = DummySeaTransformer() - model.enable_cache(config) - - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - runtime.update(step=1, sigma=0.5) - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - - assert [block.calls for block in model.layers] == [2, 2, 2] - - -@torch.no_grad() -def test_sea_cache_non_adjacent_steps_start_a_new_residual_trajectory(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - config = _make_config(runtime, threshold=0.0, retention_steps=0, residual_order=1) - model = DummySeaTransformer() - model.enable_cache(config) - - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - runtime.update(step=1, sigma=0.5) - with model.cache_context("cond"): - model(torch.full((1, 1, 1, 1), 100.0)) - - config.threshold = 100.0 - runtime.update(step=0, sigma=0.9) - with model.cache_context("cond"): - model(torch.full((1, 1, 1, 1), 2.0)) - runtime.update(step=1, sigma=0.5) - with model.cache_context("cond"): - _, output = model(torch.full((1, 1, 1, 1), 2.0)) - - torch.testing.assert_close(output, torch.full_like(output, 16.0)) - assert [block.calls for block in model.layers] == [3, 3, 3] - - -@torch.no_grad() -def test_sea_cache_fails_open_for_shape_changes(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - model = DummySeaTransformer() - model.enable_cache(_make_config(runtime, residual_order=1)) - - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - - runtime.update(step=1, sigma=0.5) - with model.cache_context("cond"): - model(torch.ones(1, 2, 1, 1)) - - assert [block.calls for block in model.layers] == [2, 2, 2] - - -def test_sea_cache_is_inference_only_and_fails_open_with_autograd(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - model = DummySeaTransformer() - model.enable_cache(_make_config(runtime)) - - with torch.enable_grad(), model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1, requires_grad=True)) - runtime.update(step=1, sigma=0.5) - with torch.enable_grad(), model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1, requires_grad=True)) - - assert [block.calls for block in model.layers] == [2, 2, 2] - - -@torch.no_grad() -def test_sea_filter_density_normalizes_gain_to_unit_mean(): - impulse = torch.zeros(2, 3, 4, 1) - impulse[0, 0, 0, 0] = 1 - - filtered = _apply_sea_filter(impulse, sigma=0.5, power_exp=3.0) - recovered_gain = torch.fft.fftn(filtered.float(), dim=(0, 1, 2)) - - torch.testing.assert_close( - recovered_gain.real.mean(), - torch.tensor(1.0), - atol=1e-5, - rtol=1e-5, - ) - torch.testing.assert_close( - recovered_gain.imag, - torch.zeros_like(recovered_gain.imag), - atol=1e-5, - rtol=0, - ) - - -@pytest.mark.parametrize("sigma", [0.0, 1.0]) -def test_sea_filter_is_finite_at_scheduler_endpoints(sigma): - filtered = _apply_sea_filter(torch.randn(2, 2, 2, 4), sigma=sigma, power_exp=2.0) - - assert torch.isfinite(filtered).all() - assert filtered.abs().sum() > 0 - - -@torch.no_grad() -def test_sea_cache_single_block_supports_full_and_cached_execution_then_disables_cleanly(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - model = DummySeaTransformer(num_layers=1) - model.enable_cache(_make_config(runtime)) - - with model.cache_context("cond"): - _, first_output = model(torch.ones(1, 1, 1, 1)) - torch.testing.assert_close(first_output, torch.full_like(first_output, 2.0)) - - runtime.update(step=1, sigma=0.5) - with model.cache_context("cond"): - _, cached_output = model(torch.full((1, 1, 1, 1), 2.0)) - torch.testing.assert_close(cached_output, torch.full_like(cached_output, 3.0)) - assert model.layers[0].calls == 1 - - model.disable_cache() - - _, output = model(torch.ones(1, 1, 1, 1)) - - torch.testing.assert_close(output, torch.full_like(output, 2.0)) - assert model.layers[0].calls == 2 - assert model.layers[0]._diffusers_hook.hooks == {} - - -@torch.no_grad() -def test_sea_cache_fails_open_for_parameter_sharded_blocks(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - model = DummySeaTransformer() - model.layers[0]._get_fsdp_state = lambda: object() - model.enable_cache(_make_config(runtime)) - - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - runtime.update(step=1, sigma=0.5) - with model.cache_context("cond"): - model(torch.ones(1, 1, 1, 1)) - - assert [block.calls for block in model.layers] == [2, 2, 2] - - -def test_sea_cache_failed_enable_rolls_back_only_new_hooks(): - runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} - model = DummySeaTransformer() - existing_hook = ModelHook() - middle_registry = HookRegistry.check_if_exists_or_initialize(model.layers[1]) - middle_registry.register_hook(existing_hook, _SEA_CACHE_BLOCK_HOOK) - - with pytest.raises(ValueError, match="already exists"): - model.enable_cache(_make_config(runtime)) - - assert not model.is_cache_enabled - assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is None - assert model.layers[0]._diffusers_hook.get_hook(_SEA_CACHE_LEADER_BLOCK_HOOK) is None - assert middle_registry.get_hook(_SEA_CACHE_BLOCK_HOOK) is existing_hook - assert not hasattr(model.layers[2], "_diffusers_hook") - - -def test_sea_cache_config_defaults(): - config = SeaCacheConfig() - - assert config.threshold == 0.25 - assert config.residual_order == 1 - assert config.max_consecutive_cached == 2 - assert config.indicator_source == "raw_vision_latents" - - -@pytest.mark.parametrize( - ("kwargs", "message"), - [ - ({"threshold": -1.0}, "threshold"), - ({"residual_order": 2}, "residual_order"), - ({"retention_steps": -1}, "retention_steps"), - ({"cache_end_steps": -1}, "cache_end_steps"), - ({"max_consecutive_cached": -1}, "max_consecutive_cached"), - ({"max_consecutive_cached": 1.5}, "max_consecutive_cached"), - ({"max_consecutive_cached": True}, "max_consecutive_cached"), - ({"power_exp": 0.0}, "power_exp"), - ({"indicator_source": "raw"}, "indicator_source"), - ({"threshold": float("nan")}, "threshold"), - ({"power_exp": float("inf")}, "power_exp"), - ], -) -def test_sea_cache_config_validation(kwargs, message): - with pytest.raises(ValueError, match=message): - SeaCacheConfig(**kwargs) - - -@pytest.mark.parametrize("callback_name", ["metadata_callback", "raw_vision_callback"]) -def test_sea_cache_callback_validation(callback_name): - with pytest.raises(TypeError, match=callback_name): - SeaCacheConfig(**{callback_name: 1}) 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 c03d39f2b4c1..8a078b64ba76 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -34,6 +34,7 @@ BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, + SeaCacheTesterMixin, TorchCompileTesterMixin, TrainingTesterMixin, ) @@ -375,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): From 0a943d3fbdf6c1c1bcaa050112164dfe252cc6bd Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Thu, 3 Sep 2026 04:13:02 -0700 Subject: [PATCH 13/20] Make Cosmos3 sampling state always FP32 --- docs/source/en/api/pipelines/cosmos3.md | 6 -- .../cosmos/before_denoise.py | 86 +------------------ .../modular_pipelines/cosmos/denoise.py | 44 +++------- .../cosmos/modular_blocks_cosmos3.py | 36 -------- .../modular_blocks_cosmos3_distilled.py | 8 -- .../pipelines/cosmos/pipeline_cosmos3_omni.py | 25 ++---- .../cosmos/test_modular_pipeline_cosmos3.py | 46 +++------- ...test_modular_pipeline_cosmos3_distilled.py | 29 ++----- 8 files changed, 42 insertions(+), 238 deletions(-) diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index dcfb7167db90..5a382c4743a8 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -685,12 +685,6 @@ The same model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModula `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. -## Sampling precision - -Cosmos 3 keeps denoising latents and classifier-free-guidance arithmetic in `torch.float32` by default while casting -transformer inputs to the model dtype. Pass `use_fp32_sampling_state=False` to an inference call to keep the sampling -state in the model dtype instead. - ## 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/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index decbe325782b..7e9d83fa6316 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -12,15 +12,6 @@ from .modular_pipeline import Cosmos3OmniModularPipeline -def _resolve_sampling_dtype( - components: Cosmos3OmniModularPipeline, use_fp32_sampling_state: bool | None -) -> tuple[bool, torch.dtype]: - if use_fp32_sampling_state is None: - use_fp32_sampling_state = components.config.default_use_fp32_sampling_state - sampling_dtype = torch.float32 if use_fp32_sampling_state else components.transformer.dtype - return use_fp32_sampling_state, sampling_dtype - - class Cosmos3PrepareTextSegmentsStep(ModularPipelineBlocks): model_name = "cosmos3-omni" @@ -77,10 +68,6 @@ def description(self) -> str: def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Cosmos3OmniTransformer)] - @property - def expected_configs(self) -> list[ConfigSpec]: - return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)] - @property def inputs(self) -> list[InputParam]: return [ @@ -111,16 +98,6 @@ def inputs(self) -> list[InputParam]: description="Pre-generated noisy vision latents.", ), InputParam.template("generator"), - InputParam( - name="use_fp32_sampling_state", - type_hint=bool | None, - default=None, - description=( - "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. " - "If unset, uses the " - "pipeline's `default_use_fp32_sampling_state` config." - ), - ), ] @property @@ -144,20 +121,13 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=torch.Tensor, description="Clean encoded vision latents used to re-anchor image conditioning each step.", ), - OutputParam( - "use_fp32_sampling_state", - type_hint=bool, - description="Whether vision sampling state is kept in float32.", - ), ] @torch.no_grad() def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype( - components, block_state.use_fp32_sampling_state - ) + sampling_dtype = torch.float32 x0_tokens_vision = block_state.x0_tokens_vision if x0_tokens_vision is None: @@ -226,10 +196,6 @@ def expected_components(self) -> list[ComponentSpec]: ComponentSpec("scheduler", UniPCMultistepScheduler), ] - @property - def expected_configs(self) -> list[ConfigSpec]: - return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)] - @property def inputs(self) -> list[InputParam]: return [ @@ -242,16 +208,6 @@ def inputs(self) -> list[InputParam]: description="Pre-generated noisy sound latents.", ), InputParam.template("generator"), - InputParam( - name="use_fp32_sampling_state", - type_hint=bool | None, - default=None, - description=( - "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. " - "If unset, uses the " - "pipeline's `default_use_fp32_sampling_state` config." - ), - ), ] @property @@ -272,9 +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 - block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype( - components, block_state.use_fp32_sampling_state - ) + sampling_dtype = torch.float32 if not components.transformer.config.sound_gen: raise ValueError("Sound generation requires a transformer trained with sound_gen=True.") @@ -320,10 +274,6 @@ def expected_components(self) -> list[ComponentSpec]: ComponentSpec("scheduler", UniPCMultistepScheduler), ] - @property - def expected_configs(self) -> list[ConfigSpec]: - return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)] - @property def inputs(self) -> list[InputParam]: return [ @@ -346,16 +296,6 @@ def inputs(self) -> list[InputParam]: description="Pre-generated noisy action latents.", ), InputParam.template("generator"), - InputParam( - name="use_fp32_sampling_state", - type_hint=bool | None, - default=None, - description=( - "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. " - "If unset, uses the " - "pipeline's `default_use_fp32_sampling_state` config." - ), - ), ] @property @@ -387,9 +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 - block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype( - components, block_state.use_fp32_sampling_state - ) + sampling_dtype = torch.float32 action = block_state.action if not components.transformer.config.action_gen: @@ -1067,10 +1005,6 @@ def description(self) -> str: def expected_components(self) -> list[ComponentSpec]: return [ComponentSpec("transformer", Cosmos3OmniTransformer)] - @property - def expected_configs(self) -> list[ConfigSpec]: - return [ConfigSpec(name="default_use_fp32_sampling_state", default=True)] - @property def inputs(self) -> list[InputParam]: return [ @@ -1087,16 +1021,6 @@ def inputs(self) -> list[InputParam]: description="Number of pixel frames used to seed this chunk's target.", ), InputParam.template("generator"), - InputParam( - name="use_fp32_sampling_state", - type_hint=bool | None, - default=None, - description=( - "Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. " - "If unset, uses the " - "pipeline's `default_use_fp32_sampling_state` config." - ), - ), ] @property @@ -1124,9 +1048,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 - block_state.use_fp32_sampling_state, sampling_dtype = _resolve_sampling_dtype( - components, block_state.use_fp32_sampling_state - ) + sampling_dtype = torch.float32 tcf = components.vae_scale_factor_temporal target_x0 = block_state.x0_tokens_vision.to(device=device, dtype=sampling_dtype) diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index 8cf4b377198f..ba6e06a15244 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -166,12 +166,6 @@ def inputs(self) -> list[InputParam]: default=6.0, description="Scale for classifier-free guidance.", ), - InputParam( - name="use_fp32_sampling_state", - type_hint=bool, - default=False, - description="Whether to keep velocity and guidance arithmetic in float32.", - ), ] @property @@ -236,12 +230,14 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta ) cond_velocity_vision, cond_velocity_sound, cond_velocity_action = velocities["cond"] - if block_state.use_fp32_sampling_state: - 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 + 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 ) @@ -332,12 +328,6 @@ def inputs(self) -> list[InputParam]: description="Indexes of conditioned vision latent frames; non-empty for image-to-video.", ), InputParam.template("generator"), - InputParam( - name="use_fp32_sampling_state", - type_hint=bool, - default=False, - description="Whether to keep the distilled vision scheduler state in float32.", - ), ] @property @@ -346,11 +336,8 @@ 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 - latents = block_state.latents - if block_state.use_fp32_sampling_state: - velocity_vision = velocity_vision.float() - latents = latents.float() + 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( @@ -734,12 +721,6 @@ def inputs(self) -> list[InputParam]: default=None, description="Timestep interval [lo, hi] over which control guidance is active (None = always).", ), - InputParam( - name="use_fp32_sampling_state", - type_hint=bool, - default=False, - description="Whether to keep velocity and guidance arithmetic in float32.", - ), ] @property @@ -814,10 +795,9 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta "uncond", ) - if block_state.use_fp32_sampling_state: - 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 + 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) @@ -921,8 +901,6 @@ class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper): Timestep interval [lo, hi] over which text guidance is active (None = always). control_guidance_interval (`tuple`, *optional*): Timestep interval [lo, hi] over which control guidance is active (None = always). - use_fp32_sampling_state (`bool`, *optional*, defaults to False): - Whether to keep velocity and guidance arithmetic in float32. latents (`Tensor`): Noisy target latents to update. condition_latents (`Tensor`): diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index ee02e192da21..d15a176a0824 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -377,7 +377,6 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): scheduler (`UniPCMultistepScheduler`) Configs: - default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -401,9 +400,6 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *optional*): @@ -453,7 +449,6 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): scheduler (`UniPCMultistepScheduler`) Configs: - default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -477,9 +472,6 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. sound_latents (`Tensor`, *optional*): @@ -542,7 +534,6 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): scheduler (`UniPCMultistepScheduler`) Configs: - default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -566,9 +557,6 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. action (`CosmosActionCondition`): @@ -635,7 +623,6 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): scheduler (`UniPCMultistepScheduler`) Configs: - default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -659,9 +646,6 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. sound_latents (`Tensor`, *optional*): @@ -740,9 +724,6 @@ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) - Configs: - default_use_fp32_sampling_state (default: True) - Inputs: chunk_id (`int`, *optional*, defaults to 0): Index of the current chunk. @@ -768,9 +749,6 @@ class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. cond_text_segment (`dict`): Conditional text segment. uncond_text_segment (`dict`): @@ -892,9 +870,6 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks): video_processor (`VideoProcessor`) scheduler (`UniPCMultistepScheduler`) - Configs: - default_use_fp32_sampling_state (default: True) - Inputs: cond_input_ids (`None`): Token IDs for the conditional prompt. @@ -924,9 +899,6 @@ class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. num_inference_steps (`int`): @@ -1020,7 +992,6 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): scheduler (`UniPCMultistepScheduler`) Configs: - default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -1052,9 +1023,6 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): Number of frames each later chunk reuses from the previous chunk's tail. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. num_inference_steps (`int`): @@ -1219,7 +1187,6 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Configs: default_use_system_prompt (default: True) enable_safety_checker (default: True) - default_use_fp32_sampling_state (default: True) use_native_flow_schedule (default: False) Inputs: @@ -1265,9 +1232,6 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Number of frames the first chunk reuses from the input video. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`): The number of denoising steps. **denoiser_input_fields (`None`, *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 6841d2f3ff4f..b6ec5d47eb8d 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3_distilled.py @@ -88,7 +88,6 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): scheduler (`FlowMatchEulerDiscreteScheduler`) Configs: - default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default: None) @@ -113,9 +112,6 @@ class Cosmos3DistilledVisionCoreDenoiseStep(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`, *optional*): The number of denoising steps. guidance_scale (`float`, *optional*): @@ -177,7 +173,6 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): Configs: default_use_system_prompt (default: True) enable_safety_checker (default: True) - default_use_fp32_sampling_state (default: True) is_distilled (default: True) distilled_sigmas (default: None) @@ -214,9 +209,6 @@ class Cosmos3DistilledBlocks(SequentialPipelineBlocks): Pre-generated noisy vision latents. generator (`Generator`, *optional*): Torch generator for deterministic generation. - use_fp32_sampling_state (`bool | NoneType`, *optional*): - Whether to keep denoising latents, masks, guidance arithmetic, and scheduler state in float32. If unset, uses the - pipeline's `default_use_fp32_sampling_state` config. num_inference_steps (`int`, *optional*): The number of denoising steps. guidance_scale (`float`, *optional*): diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 19657ae0fe33..08408bde25a6 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -1353,7 +1353,6 @@ def __call__( add_resolution_template: bool = True, add_duration_template: bool = True, enable_safety_check: bool = True, - use_fp32_sampling_state: bool = True, ) -> Cosmos3OmniPipelineOutput: r""" Run the Cosmos 3 omni pipeline end-to-end: encode the (optional) conditioning image/video, denoise vision and @@ -1449,11 +1448,6 @@ def __call__( When `True` and a `CosmosSafetyChecker` is attached, runs the text guardrail on the prompt before generation and the video guardrail on the decoded frames. Set to `False` to skip both for this call; the checker remains loaded for subsequent calls. - use_fp32_sampling_state (`bool`, *optional*, defaults to `True`): - When `True`, keeps vision, sound, and action denoising latents plus classifier-free-guidance arithmetic - in `torch.float32`. Transformer inputs are still cast to the transformer's dtype before each forward. - This improves sampling-state precision at the cost of additional memory. Set it to `False` to keep - sampling state in the model dtype. Returns: [`Cosmos3OmniPipelineOutput`] or `tuple`: @@ -1518,7 +1512,7 @@ def __call__( device = self._get_execution_device() dtype = self.transformer.dtype - sampling_dtype = torch.float32 if use_fp32_sampling_state else dtype + sampling_dtype = torch.float32 if enable_safety_check and isinstance(self.safety_checker, CosmosSafetyChecker): self.safety_checker.to(device) @@ -1812,13 +1806,12 @@ def __call__( raw_action_dim=raw_action_dim_resolved, ) - if use_fp32_sampling_state: - 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 + 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 @@ -1863,9 +1856,7 @@ 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) - if use_fp32_sampling_state: - latents = latents.float() + 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() diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index 013940fd6970..abe293471ac8 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -200,29 +200,20 @@ def _get_sampling_state_block_pipe(self, block): pipe.to(torch_device) return pipe - @pytest.mark.parametrize( - ("use_fp32_sampling_state", "expected_dtype"), - [(False, torch.bfloat16), (True, torch.float32)], - ) - def test_sound_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype): + 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), - use_fp32_sampling_state=use_fp32_sampling_state, output=["sound_latents", "sound_condition_mask"], ) - assert outputs["sound_latents"].dtype == expected_dtype - assert outputs["sound_condition_mask"].dtype == expected_dtype + assert outputs["sound_latents"].dtype == torch.float32 + assert outputs["sound_condition_mask"].dtype == torch.float32 - @pytest.mark.parametrize( - ("use_fp32_sampling_state", "expected_dtype"), - [(False, torch.bfloat16), (True, torch.float32)], - ) - def test_action_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype): + def test_action_prepare_latents_uses_fp32(self): pipe = self._get_sampling_state_block_pipe(Cosmos3ActionPrepareLatentsStep()) action = CosmosActionCondition( mode="policy", @@ -235,31 +226,25 @@ def test_action_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, ex action=action, action_condition_frame_indexes=[], generator=self.get_generator(0), - use_fp32_sampling_state=use_fp32_sampling_state, output=["action_latents", "action_condition_mask"], ) - assert outputs["action_latents"].dtype == expected_dtype - assert outputs["action_condition_mask"].dtype == expected_dtype + assert outputs["action_latents"].dtype == torch.float32 + assert outputs["action_condition_mask"].dtype == torch.float32 - @pytest.mark.parametrize( - ("use_fp32_sampling_state", "expected_dtype"), - [(False, torch.bfloat16), (True, torch.float32)], - ) - def test_transfer_prepare_latents_sampling_dtype(self, use_fp32_sampling_state, expected_dtype): + 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), - use_fp32_sampling_state=use_fp32_sampling_state, output=["latents", "velocity_mask", "condition_latents"], ) - assert outputs["latents"].dtype == expected_dtype - assert outputs["velocity_mask"].dtype == expected_dtype - assert outputs["condition_latents"].dtype == expected_dtype + 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() @@ -279,19 +264,14 @@ def test_transfer_chunks_reset_stateful_cache_at_boundaries(self): ("chunk_id", 2), ] - @pytest.mark.parametrize( - ("use_fp32_sampling_state", "expected_dtype"), - [(False, torch.bfloat16), (True, torch.float32)], - ) - def test_sampling_state_controls_modular_cfg_and_scheduler_dtype(self, use_fp32_sampling_state, expected_dtype): + 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() - inputs["use_fp32_sampling_state"] = use_fp32_sampling_state outputs = pipe(**inputs, output=["velocity_vision", "latents"]) - assert outputs["velocity_vision"].dtype == expected_dtype - assert outputs["latents"].dtype == expected_dtype + 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() 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 354b9b08288f..86bc6ad4d6f2 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -116,7 +116,6 @@ def test_declares_distilled_configs(self): assert pipe.config.is_distilled is True assert pipe.config.distilled_sigmas is None assert pipe.config.default_use_system_prompt is True - assert pipe.config.default_use_fp32_sampling_state is True def test_distilled_text_step_uses_system_prompt_config_fallback(self): text_pipe = Cosmos3DistilledTextEncoderStep().init_pipeline(self.pretrained_model_name_or_path) @@ -142,18 +141,9 @@ def test_distilled_text_step_uses_system_prompt_config_fallback(self): assert explicit_without_system_prompt == default_without_system_prompt == updated_without_system_prompt assert len(default_with_system_prompt) > len(default_without_system_prompt) - @pytest.mark.parametrize( - ("config_enabled", "input_enabled", "expected_dtype"), - [ - (False, None, torch.bfloat16), - (True, None, torch.float32), - (True, False, torch.bfloat16), - ], - ) - def test_prepare_vision_latents_fp32_sampling_state(self, config_enabled, input_enabled, expected_dtype): + 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.update_components(default_use_fp32_sampling_state=config_enabled) prepare_pipe.to(torch_device) outputs = prepare_pipe( @@ -162,26 +152,19 @@ def test_prepare_vision_latents_fp32_sampling_state(self, config_enabled, input_ width=32, fps=24.0, generator=self.get_generator(0), - use_fp32_sampling_state=input_enabled, - output=["latents", "vision_condition_mask", "use_fp32_sampling_state"], + output=["latents", "vision_condition_mask"], ) - assert outputs["latents"].dtype == expected_dtype - assert outputs["vision_condition_mask"].dtype == expected_dtype - assert outputs["use_fp32_sampling_state"] is (config_enabled if input_enabled is None else input_enabled) + assert outputs["latents"].dtype == torch.float32 + assert outputs["vision_condition_mask"].dtype == torch.float32 - @pytest.mark.parametrize( - ("use_fp32_sampling_state", "expected_dtype"), - [(False, torch.bfloat16), (True, torch.float32)], - ) - def test_distilled_scheduler_fp32_state(self, use_fp32_sampling_state, expected_dtype): + def test_distilled_scheduler_uses_fp32_state(self): pipe = self.get_pipeline(torch_dtype=torch.bfloat16).to(torch_device) inputs = self.get_dummy_inputs() - inputs["use_fp32_sampling_state"] = use_fp32_sampling_state latents = pipe(**inputs, output=self.output_name) - assert latents.dtype == expected_dtype + assert latents.dtype == torch.float32 def test_vae_encoder_rejects_image_and_video_together(self): vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"] From 2342b7262a2bf628111ba00630d5f80df476829a Mon Sep 17 00:00:00 2001 From: root Date: Thu, 3 Sep 2026 06:35:03 -0700 Subject: [PATCH 14/20] Add pipe definition to the docs snippet --- docs/source/en/api/pipelines/cosmos3.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index 5a382c4743a8..36f5f089339f 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -667,7 +667,12 @@ transformer residuals when the Spectral-Evolution-Aware indicator changes slowly transformer executions. Enable it on the transformer with scheduler metadata callbacks from the pipeline: ```python -from diffusers import SeaCacheConfig +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( From 211d1d1df8165f8b86e3a8f7752d668481b742ef Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Sat, 5 Sep 2026 04:41:27 -0700 Subject: [PATCH 15/20] Remove redundant hook registry cache resets --- src/diffusers/models/cache_utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index c31971384cb8..4bb8ed9c4140 100644 --- a/src/diffusers/models/cache_utils.py +++ b/src/diffusers/models/cache_utils.py @@ -114,7 +114,6 @@ def enable_cache(self, config) -> None: HookRegistry.check_if_exists_or_initialize(self).invalidate_child_registries_cache() self._cache_config = config - HookRegistry.check_if_exists_or_initialize(self)._child_registries_cache = None def disable_cache(self) -> None: from ..hooks import ( @@ -173,7 +172,6 @@ def disable_cache(self) -> None: registry.invalidate_child_registries_cache() self._cache_config = None - registry._child_registries_cache = None def _reset_stateful_cache(self, recurse: bool = True) -> None: from ..hooks import HookRegistry From 48dadff619cb9a7df28a69c54be22597291678f2 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 7 Sep 2026 05:48:43 -0700 Subject: [PATCH 16/20] Use raw vision latents for SeaCache indicators --- docs/source/en/api/pipelines/cosmos3.md | 3 +- docs/source/en/optimization/cache.md | 3 +- src/diffusers/hooks/_helpers.py | 2 - src/diffusers/hooks/sea_cache.py | 169 ++---------------- .../test_models_transformer_cosmos3.py | 5 +- 5 files changed, 20 insertions(+), 162 deletions(-) diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index 36f5f089339f..c8b56aa8ecf9 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -664,7 +664,8 @@ if result.action is not None: 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: +transformer executions. It computes the indicator from the raw vision latents, including clean conditioning frames for +image-to-video generation. Enable it on the transformer with scheduler metadata callbacks from the pipeline: ```python import torch diff --git a/docs/source/en/optimization/cache.md b/docs/source/en/optimization/cache.md index ad65fceeaada..8fe3915a52cd 100644 --- a/docs/source/en/optimization/cache.md +++ b/docs/source/en/optimization/cache.md @@ -72,7 +72,8 @@ pipeline.transformer.enable_cache(config) [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. +transformer block stack and predicts its output from cached residuals. The indicator is computed from the raw vision +latents, including clean conditioning frames for image-to-video generation. SeaCache is disabled by default. Enable it on the transformer and provide callbacks for the active scheduler step, sigma, and number of inference steps: diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index bfc6c59477c8..8c1ad1a8c472 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -28,7 +28,6 @@ class TransformerBlockMetadata: 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 @@ -241,7 +240,6 @@ def _register_transformer_blocks_metadata(): 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", ), ) diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index f277f90f0ce7..43a2f75a11c9 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -62,24 +62,15 @@ class SeaCacheConfig: 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. + Advanced model adapter returning raw vision latents with shape `(C, T, H, W)`. Cosmos 3 uses its native + adapter when this is omitted. Example: ```python @@ -102,14 +93,9 @@ class SeaCacheConfig: 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, @@ -134,15 +120,10 @@ def __post_init__(self): ) 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) @@ -155,7 +136,6 @@ 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 @@ -346,71 +326,6 @@ def _record_full_execution( 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: @@ -509,14 +424,12 @@ def __init__( 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): @@ -574,31 +487,18 @@ def pre_forward(self, module: torch.nn.Module, *args, **kwargs): ) 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 - ) + 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: + if not raw_vision: self.shared_state.mark_fail_open( - "SeaCache requires raw noisy vision latents for the selected indicator source; action-only, sound-only, " + "SeaCache requires raw vision latents containing at least one noisy frame; action-only, sound-only, " "conditioning-only, and unsupported model calls run in fail-open mode." ) return args, kwargs @@ -607,7 +507,6 @@ def pre_forward(self, module: torch.nn.Module, *args, **kwargs): step_index=step_index, sigma=sigma, num_inference_steps=num_inference_steps, - vision_layout=vision_layout, raw_vision=raw_vision, ) return args, kwargs @@ -638,53 +537,21 @@ def __init__( 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: + def _build_indicator(self, forward_metadata: _SeaCacheForwardMetadata) -> list[torch.Tensor] | None: + if not forward_metadata.raw_vision: 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) + for latent in forward_metadata.raw_vision: + raw_vision = latent.movedim(0, -1) indicator.append( _apply_sea_filter( - noisy_vision, + raw_vision, sigma=forward_metadata.sigma, power_exp=self.config.power_exp, ).detach() @@ -716,7 +583,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): indicator_error_reported = True else: try: - indicator = self._build_indicator(hidden_states, forward_metadata) + indicator = self._build_indicator(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}" @@ -892,14 +759,9 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None: 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 + if raw_vision_callback is None and is_cosmos3: + raw_vision_callback = _prepare_cosmos3_raw_vision_metadata post_norm_modules = None if is_cosmos3: @@ -940,7 +802,6 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: config, state_manager, shared_state, - metadata_callback, raw_vision_callback, ), _SEA_CACHE_ROOT_HOOK, diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index 8a078b64ba76..43bdb052ca5e 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -108,8 +108,7 @@ 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): + def test_cosmos3_supports_sea_cache_without_changing_state_dict_keys(self): 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} @@ -130,7 +129,6 @@ def counted_gen_norm_forward(*args, **kwargs): 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"], @@ -234,7 +232,6 @@ def test_cosmos3_decoder_layer_cache_metadata_tracks_generation_stream(self): 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() From 09e11b78f6579582d79914bfa83e34564fb2eab0 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 7 Sep 2026 07:12:36 -0700 Subject: [PATCH 17/20] Make Cosmos3 SeaCache safe for regional compilation --- src/diffusers/hooks/sea_cache.py | 241 ++++++++++++++---- .../transformers/transformer_cosmos3.py | 33 ++- .../test_models_transformer_cosmos3.py | 84 +++++- 3 files changed, 292 insertions(+), 66 deletions(-) diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index 43a2f75a11c9..616d4939a482 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -155,6 +155,7 @@ def __init__(self): self.und_output: torch.Tensor | None = None self.cached_und_output: torch.Tensor | None = None self.cached_gen_residual: torch.Tensor | None = None + self.gen_output: torch.Tensor | None = None def reset_forward(self): self.skip_remaining = False @@ -165,6 +166,7 @@ def reset_forward(self): self.und_output = None self.cached_und_output = None self.cached_gen_residual = None + self.gen_output = None def reset_trajectory(self): self.history = [] @@ -392,6 +394,22 @@ def _apply_sea_filter( return torch.fft.ifftn(spectrum * gain, dim=dimensions).real.to(hidden_states_dtype) +def _build_indicator( + config: SeaCacheConfig, + forward_metadata: _SeaCacheForwardMetadata, +) -> list[torch.Tensor] | None: + if not forward_metadata.raw_vision: + return None + return [ + _apply_sea_filter( + latent.movedim(0, -1), + sigma=forward_metadata.sigma, + power_exp=config.power_exp, + ).detach() + for latent in forward_metadata.raw_vision + ] + + def _is_parameter_sharded(module: torch.nn.Module) -> bool: """Whether a block is managed by a parameter-sharding runtime that SeaCache cannot safely bypass.""" @@ -425,15 +443,38 @@ def __init__( state_manager: StateManager, shared_state: SeaCacheSharedState, raw_vision_callback: Callable, + use_stack_boundary: bool = False, ): super().__init__() self.config = config self.state_manager = state_manager self.shared_state = shared_state self.raw_vision_callback = raw_vision_callback + self.use_stack_boundary = use_stack_boundary + + def initialize_hook(self, module: torch.nn.Module): + if not self.use_stack_boundary: + return module + + unwrapped_module = unwrap_module(module) + if not hasattr(unwrapped_module, "layers") or not unwrapped_module.layers: + raise ValueError("SeaCache requires Cosmos 3 to expose a non-empty decoder stack.") + unwrapped_module._sea_cache_prepare_decoder_stack = self.prepare_decoder_stack + unwrapped_module._sea_cache_record_decoder_stack = self.record_decoder_stack + return module + + def deinitalize_hook(self, module: torch.nn.Module): + if self.use_stack_boundary: + unwrapped_module = unwrap_module(module) + for name in ("_sea_cache_prepare_decoder_stack", "_sea_cache_record_decoder_stack"): + if hasattr(unwrapped_module, name): + delattr(unwrapped_module, name) + return module def pre_forward(self, module: torch.nn.Module, *args, **kwargs): self.shared_state.forward_metadata = None + if self.state_manager._current_context is not None: + self.state_manager.get_state().reset_forward() if torch.is_grad_enabled(): self.shared_state.mark_fail_open( "SeaCache is inference-only; calls with autograd enabled run in fail-open mode." @@ -511,10 +552,126 @@ def pre_forward(self, module: torch.nn.Module, *args, **kwargs): ) return args, kwargs + def prepare_decoder_stack( + self, + module: torch.nn.Module, + und_seq: torch.Tensor, + gen_seq: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, bool]: + """Resolve the Cosmos 3 gate eagerly before entering compiled decoder blocks.""" + forward_metadata = self.shared_state.forward_metadata + if self.state_manager._current_context is None or forward_metadata is None: + return und_seq, gen_seq, True + + state: SeaCacheContextState = self.state_manager.get_state() + state.full_execution_pending = True + state.cacheable_execution = True + state.step_index = forward_metadata.step_index + state.gen_input = gen_seq + + indicator_error_reported = False + unwrapped_module = unwrap_module(module) + layers = getattr(unwrapped_module, "layers", None) + if getattr(unwrapped_module, "_cp_shard_fn", None) is not None: + self.shared_state.mark_fail_open( + "SeaCache cannot safely bypass the Cosmos 3 decoder stack with context parallelism; " + "running in fail-open mode." + ) + indicator = None + indicator_error_reported = True + elif not layers or _is_parameter_sharded(layers): + 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 = _build_indicator(self.config, 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 und_seq, gen_seq, True + + residual_history = state.history[-(self.config.residual_order + 1) :] + _, cached_und, cached_residual = residual_history[-1] + if ( + any( + residual.shape != gen_seq.shape or residual.device != gen_seq.device or residual.dtype != gen_seq.dtype + for _, _, residual in residual_history + ) + or cached_und.shape != und_seq.shape + or cached_und.device != und_seq.device + or cached_und.dtype != und_seq.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 und_seq, gen_seq, True + + 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 + return cached_und, gen_seq + cached_residual, False + + def record_decoder_stack( + self, + module: torch.nn.Module, + und_seq: torch.Tensor, + gen_seq: torch.Tensor, + ) -> None: + """Record post-normalization Cosmos 3 decoder outputs after a full step.""" + if self.state_manager._current_context is None: + return + state: SeaCacheContextState = self.state_manager.get_state() + if state.full_execution_pending: + state.und_output = und_seq + state.gen_output = gen_seq + 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() + state: SeaCacheContextState = self.state_manager.get_state() + if ( + self.use_stack_boundary + and state.full_execution_pending + and state.und_output is not None + and state.gen_output is not None + ): + _record_full_execution( + self.config, + state, + gen_output=state.gen_output, + und_output=state.und_output, + ) + else: + state.reset_forward() return output def reset_state(self, module: torch.nn.Module): @@ -544,19 +701,7 @@ def initialize_hook(self, module: torch.nn.Module): return module def _build_indicator(self, forward_metadata: _SeaCacheForwardMetadata) -> list[torch.Tensor] | None: - 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 + return _build_indicator(self.config, forward_metadata) @torch.compiler.disable def new_forward(self, module: torch.nn.Module, *args, **kwargs): @@ -745,7 +890,7 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None: 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 understanding output and a generation residual from the decoder-stack 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. @@ -803,47 +948,49 @@ def register_hook(target: torch.nn.Module, hook: ModelHook, name: str) -> None: state_manager, shared_state, raw_vision_callback, + use_stack_boundary=is_cosmos3, ), _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}'.") + if not is_cosmos3: + leader_name, leader = blocks[0] + logger.debug(f"Applying SeaCache leader hook to '{leader_name}'.") register_hook( - block, - SeaCacheBlockHook(config, state_manager, shared_state, post_norm_boundary=post_norm_boundary), - _SEA_CACHE_BLOCK_HOOK, + leader, + SeaCacheLeaderBlockHook(config, state_manager, shared_state, post_norm_boundary=post_norm_boundary), + _SEA_CACHE_LEADER_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}'.") + for name, block in blocks[1:-1]: + logger.debug(f"Applying SeaCache identity hook to '{name}'.") register_hook( - norm_module, - SeaCachePostNormHook(config, state_manager, shared_state, pathway=pathway), - _SEA_CACHE_POST_NORM_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) diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py index 07c6bd12d0dc..2902aa38f914 100644 --- a/src/diffusers/models/transformers/transformer_cosmos3.py +++ b/src/diffusers/models/transformers/transformer_cosmos3.py @@ -799,15 +799,30 @@ def forward( if self._cp_shard_fn is not None: und_seq, gen_seq, rotary_emb = self._cp_shard_fn(und_seq, gen_seq, rotary_emb) - for decoder_layer in self.layers: - if torch.is_grad_enabled() and self.gradient_checkpointing: - und_seq, gen_seq = self._gradient_checkpointing_func( - decoder_layer.__call__, und_seq, gen_seq, rotary_emb - ) - else: - und_seq, gen_seq = decoder_layer(und_seq, gen_seq, rotary_emb) - und_out = self.norm(und_seq) - gen_out = self.norm_moe_gen(gen_seq) + run_decoder_stack = True + sea_cache_prepare = getattr(self, "_sea_cache_prepare_decoder_stack", None) + if sea_cache_prepare is not None: + und_seq, gen_seq, run_decoder_stack = sea_cache_prepare(self, und_seq, gen_seq) + + if run_decoder_stack: + for decoder_layer in self.layers: + if torch.is_grad_enabled() and self.gradient_checkpointing: + und_seq, gen_seq = self._gradient_checkpointing_func( + decoder_layer.__call__, und_seq, gen_seq, rotary_emb + ) + else: + und_seq, gen_seq = decoder_layer(und_seq, gen_seq, rotary_emb) + + und_out = self.norm(und_seq) + gen_out = self.norm_moe_gen(gen_seq) + + sea_cache_record = getattr(self, "_sea_cache_record_decoder_stack", None) + if sea_cache_record is not None: + sea_cache_record(self, und_out, gen_out) + else: + # The SeaCache stack boundary returns cached post-normalization states. + und_out = und_seq + gen_out = gen_seq # Optional context-parallelism gather seam: re-gather the full per-pathway # sequence on every rank (and drop the padding) before the global-index decode diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index 43bdb052ca5e..61e6cd656e93 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -137,6 +137,7 @@ def counted_gen_norm_forward(*args, **kwargs): assert isinstance(model, CacheMixin) model.enable_cache(config) assert set(model.state_dict()) == state_dict_keys + assert model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) is not None inputs = self.get_dummy_inputs() with torch.no_grad(), model.cache_context("cond"): @@ -149,12 +150,14 @@ def counted_gen_norm_forward(*args, **kwargs): 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. + # A hit bypasses the decoder stack and final pathway norms while prediction heads stay fresh. 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 + assert not hasattr(model, "_sea_cache_prepare_decoder_stack") + assert not hasattr(model, "_sea_cache_record_decoder_stack") with torch.no_grad(): model(**self.get_dummy_inputs()) assert norm_calls == {"und": 2, "gen": 2} @@ -169,14 +172,20 @@ def test_cosmos3_sea_cache_post_norm_boundary_numeric_semantics(self): current_sigma_callback=lambda: runtime["sigma"], num_inference_steps_callback=lambda: runtime["num_steps"], ) - first_block_gen_inputs = [] + decoder_gen_inputs = [] + prepared_stack_outputs = [] + und_norm_inputs = [] + gen_norm_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_norm_input(module, args): + und_norm_inputs.append(args[0].detach().clone()) + + def capture_gen_norm_input(module, args): + gen_norm_inputs.append(args[0].detach().clone()) def capture_und_output(module, args, output): returned_und_outputs.append(output.detach().clone()) @@ -192,12 +201,25 @@ def counted_projection_forward(hidden_states): 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_pre_hook(capture_und_norm_input) + model.norm_moe_gen.register_forward_pre_hook(capture_gen_norm_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) + original_prepare_decoder_stack = model._sea_cache_prepare_decoder_stack + + def capture_prepare_decoder_stack(module, und_seq, gen_seq): + decoder_gen_inputs.append(gen_seq.detach().clone()) + prepared_und, prepared_gen, should_compute = original_prepare_decoder_stack(module, und_seq, gen_seq) + prepared_stack_outputs.append( + (prepared_und.detach().clone(), prepared_gen.detach().clone(), should_compute) + ) + return prepared_und, prepared_gen, should_compute + + model._sea_cache_prepare_decoder_stack = capture_prepare_decoder_stack + with torch.no_grad(), model.cache_context("cond"): model(**self.get_dummy_inputs()) @@ -208,7 +230,7 @@ def counted_projection_forward(hidden_states): 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], + returned_gen_outputs[0] - decoder_gen_inputs[0], ) runtime.update(step=1, sigma=0.6) @@ -217,14 +239,56 @@ def counted_projection_forward(hidden_states): with torch.no_grad(), model.cache_context("cond"): model(**cached_inputs) - torch.testing.assert_close(returned_und_outputs[1], cached_und_output) + assert len(und_norm_inputs) == len(gen_norm_inputs) == 1 + assert prepared_stack_outputs[1][2] is False + torch.testing.assert_close(prepared_stack_outputs[1][0], cached_und_output) torch.testing.assert_close( - returned_gen_outputs[1], - first_block_gen_inputs[-1] + cached_gen_residual, + prepared_stack_outputs[1][1], + decoder_gen_inputs[1] + cached_gen_residual, ) - torch.testing.assert_close(projection_head_inputs[1], returned_gen_outputs[1]) + torch.testing.assert_close(projection_head_inputs[1], prepared_stack_outputs[1][1]) assert projection_head_calls == 2 + @pytest.mark.skipif(not hasattr(torch, "compile"), reason="torch.compile is unavailable") + def test_cosmos3_sea_cache_regional_compile_fullgraph_without_recompile(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + runtime = {"step": 0, "sigma": 0.9, "num_steps": 4} + model.enable_cache( + SeaCacheConfig( + threshold=100.0, + cache_end_steps=0, + max_consecutive_cached=1, + current_step_callback=lambda: runtime["step"], + current_sigma_callback=lambda: runtime["sigma"], + num_inference_steps_callback=lambda: runtime["num_steps"], + ) + ) + model.compile_repeated_blocks(backend="eager", fullgraph=True) + + with torch.no_grad(): + with model.cache_context("cond"): + first = model(**self.get_dummy_inputs()) + + with torch.no_grad(), torch._dynamo.config.patch(error_on_recompile=True): + 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 model.cache_context("cond"): + cached = model(**cached_inputs) + root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) + assert root_hook.state_manager._state_cache["cond"].consecutive_cached == 1 + + runtime.update(step=2, sigma=0.3) + refreshed_inputs = self.get_dummy_inputs() + refreshed_inputs["vision_tokens"] = [refreshed_inputs["vision_tokens"][0] + 0.2] + with model.cache_context("cond"): + refreshed = model(**refreshed_inputs) + + assert first.sample[0].shape == self.output_shape + assert cached.sample[0].shape == self.output_shape + assert refreshed.sample[0].shape == self.output_shape + assert root_hook.state_manager._state_cache["cond"].history[-1][0] == 2 + def test_cosmos3_decoder_layer_cache_metadata_tracks_generation_stream(self): metadata = TransformerBlockRegistry.get(Cosmos3VLTextMoTDecoderLayer) From 98877311607b698e0745b0084b949d436b47466f Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 7 Sep 2026 09:09:56 -0700 Subject: [PATCH 18/20] Fix SeaCache with Cosmos3 model parallelism --- src/diffusers/hooks/sea_cache.py | 151 ++++++++++++------ .../test_models_transformer_cosmos3.py | 131 +++++++++++++++ 2 files changed, 233 insertions(+), 49 deletions(-) diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index 616d4939a482..71244fb2217b 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -410,28 +410,49 @@ def _build_indicator( ] -def _is_parameter_sharded(module: torch.nn.Module) -> bool: - """Whether a block is managed by a parameter-sharding runtime that SeaCache cannot safely bypass.""" - +def _parameter_sharding_types(module: torch.nn.Module) -> tuple[bool, bool]: + """Return whether a module contains FSDP-managed and DTensor parameters.""" + has_fsdp = False + has_dtensor = False for submodule in unwrap_module(module).modules(): module_type = type(submodule) if callable(getattr(submodule, "_get_fsdp_state", None)): - return True + has_fsdp = True if module_type.__name__ == "FullyShardedDataParallel" and module_type.__module__.startswith( "torch.distributed.fsdp" ): - return True + has_fsdp = 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") + if parameter_type.__name__ == "FlatParameter" and parameter_type.__module__.startswith( + "torch.distributed.fsdp" ): - return True - return False + has_fsdp = True + elif parameter_type.__name__ == "DTensor" and parameter_type.__module__.startswith( + "torch.distributed.tensor" + ): + has_dtensor = True + return has_fsdp, has_dtensor + + +def _is_parameter_sharded(module: torch.nn.Module) -> bool: + """Whether a block is managed by a parameter-sharding runtime that SeaCache cannot safely bypass.""" + return any(_parameter_sharding_types(module)) + + +def _synchronize_compute_decision(should_compute: bool, device: torch.device) -> tuple[bool, bool]: + """Synchronize a cache decision across the active distributed world. + + The two votes distinguish unanimous full/skip decisions from disagreement. Disagreement always resolves to full + execution and resets cache trajectories. + """ + votes = torch.tensor( + [int(should_compute), int(not should_compute)], + dtype=torch.int32, + device=device, + ) + torch.distributed.all_reduce(votes, op=torch.distributed.ReduceOp.MAX) + return bool(votes[0].item()), bool(votes[0].item() and votes[1].item()) class SeaCacheRootHook(ModelHook): @@ -559,8 +580,24 @@ def prepare_decoder_stack( gen_seq: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, bool]: """Resolve the Cosmos 3 gate eagerly before entering compiled decoder blocks.""" + unwrapped_module = unwrap_module(module) + layers = getattr(unwrapped_module, "layers", None) + uses_context_parallel = getattr(unwrapped_module, "_cp_shard_fn", None) is not None + has_fsdp, has_dtensor = _parameter_sharding_types(layers) if layers else (False, False) + uses_supported_parallelism = uses_context_parallel or has_dtensor + synchronize_decision = ( + uses_supported_parallelism + and not has_fsdp + and torch.distributed.is_available() + and torch.distributed.is_initialized() + ) + forward_metadata = self.shared_state.forward_metadata if self.state_manager._current_context is None or forward_metadata is None: + if synchronize_decision: + _, disagreed = _synchronize_compute_decision(True, gen_seq.device) + if disagreed and self.state_manager._current_context is not None: + self.state_manager.get_state().reset_trajectory() return und_seq, gen_seq, True state: SeaCacheContextState = self.state_manager.get_state() @@ -570,21 +607,23 @@ def prepare_decoder_stack( state.gen_input = gen_seq indicator_error_reported = False - unwrapped_module = unwrap_module(module) - layers = getattr(unwrapped_module, "layers", None) - if getattr(unwrapped_module, "_cp_shard_fn", None) is not None: + if has_fsdp: self.shared_state.mark_fail_open( - "SeaCache cannot safely bypass the Cosmos 3 decoder stack with context parallelism; " - "running in fail-open mode." + "SeaCache cannot safely bypass FSDP-managed transformer blocks; running in fail-open mode." ) indicator = None indicator_error_reported = True - elif not layers or _is_parameter_sharded(layers): + elif uses_supported_parallelism and not synchronize_decision: self.shared_state.mark_fail_open( - "SeaCache cannot safely bypass parameter-sharded transformer blocks; running in fail-open mode." + "SeaCache requires an initialized distributed process group for Cosmos 3 context/tensor parallelism; " + "running in fail-open mode." ) indicator = None indicator_error_reported = True + elif not layers: + self.shared_state.mark_fail_open("SeaCache could not locate the Cosmos 3 decoder stack; running full.") + indicator = None + indicator_error_reported = True else: try: indicator = _build_indicator(self.config, forward_metadata) @@ -599,39 +638,53 @@ def prepare_decoder_stack( "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 und_seq, gen_seq, True + gate_should_compute = self.shared_state.resolve_gate(state, forward_metadata, indicator, self.config) + should_compute = gate_should_compute + cached_und = cached_residual = None - residual_history = state.history[-(self.config.residual_order + 1) :] - _, cached_und, cached_residual = residual_history[-1] - if ( - any( - residual.shape != gen_seq.shape or residual.device != gen_seq.device or residual.dtype != gen_seq.dtype - for _, _, residual in residual_history - ) - or cached_und.shape != und_seq.shape - or cached_und.device != und_seq.device - or cached_und.dtype != und_seq.dtype - ): - state.history = [] - state.accumulated_distance = 0.0 + if not should_compute and not state.history: + should_compute = True self.shared_state.mark_fail_open( - "SeaCache residual history changed shape, device, or dtype; running in fail-open mode." + "SeaCache selected a cache hit without residual history; running in fail-open mode." ) - return und_seq, gen_seq, True + elif not should_compute: + residual_history = state.history[-(self.config.residual_order + 1) :] + _, cached_und, cached_residual = residual_history[-1] + if ( + any( + residual.shape != gen_seq.shape + or residual.device != gen_seq.device + or residual.dtype != gen_seq.dtype + for _, _, residual in residual_history + ) + or cached_und.shape != und_seq.shape + or cached_und.device != und_seq.device + or cached_und.dtype != und_seq.dtype + ): + state.history = [] + should_compute = True + self.shared_state.mark_fail_open( + "SeaCache residual history changed shape, device, or dtype; running in fail-open mode." + ) + elif 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 + + disagreed = False + if synchronize_decision: + should_compute, disagreed = _synchronize_compute_decision(should_compute, gen_seq.device) + + if disagreed: + state.reset_trajectory() + elif should_compute and not gate_should_compute: + state.accumulated_distance = 0.0 + state.gate_should_compute = True - 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 + if should_compute: + return und_seq, gen_seq, True state.skip_remaining = True state.full_execution_pending = False diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index 61e6cd656e93..f19ffb4dfd7a 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -17,6 +17,7 @@ import torch from diffusers import Cosmos3OmniTransformer, SeaCacheConfig +from diffusers.hooks import sea_cache as sea_cache_module from diffusers.hooks._helpers import TransformerBlockRegistry from diffusers.hooks.sea_cache import _SEA_CACHE_ROOT_HOOK from diffusers.models.cache_utils import CacheMixin @@ -249,6 +250,136 @@ def capture_prepare_decoder_stack(module, und_seq, gen_seq): torch.testing.assert_close(projection_head_inputs[1], prepared_stack_outputs[1][1]) assert projection_head_calls == 2 + @pytest.mark.parametrize("parallelism", ["cp", "tp"]) + def test_cosmos3_sea_cache_reuses_with_model_parallel_consensus(self, monkeypatch, parallelism): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + if parallelism == "cp": + model._cp_shard_fn = lambda und_seq, gen_seq, rotary_emb: (und_seq, gen_seq, rotary_emb) + model._cp_gather_fn = lambda und_out, gen_out: (und_out, gen_out) + else: + monkeypatch.setattr(sea_cache_module, "_parameter_sharding_types", lambda _module: (False, True)) + + votes = [] + + def all_reduce(decision, *, op): + assert op == torch.distributed.ReduceOp.MAX + votes.append(decision.tolist()) + + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + + layer_calls = 0 + + def count_layer_calls(_module, _args, _output): + nonlocal layer_calls + layer_calls += 1 + + model.layers[0].register_forward_hook(count_layer_calls) + runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} + model.enable_cache( + 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"], + ) + ) + + with torch.no_grad(), model.cache_context("cond"): + model(**self.get_dummy_inputs()) + 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) + + root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) + assert layer_calls == 1 + assert root_hook.state_manager._state_cache["cond"].consecutive_cached == 1 + assert votes == [[1, 0], [0, 1]] + + def test_cosmos3_sea_cache_parallel_disagreement_forces_full_execution(self, monkeypatch): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model._cp_shard_fn = lambda und_seq, gen_seq, rotary_emb: (und_seq, gen_seq, rotary_emb) + model._cp_gather_fn = lambda und_out, gen_out: (und_out, gen_out) + reduction_count = 0 + + def all_reduce(decision, *, op): + nonlocal reduction_count + assert op == torch.distributed.ReduceOp.MAX + reduction_count += 1 + if reduction_count == 2: + decision.fill_(1) + + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + + layer_calls = 0 + + def count_layer_calls(_module, _args, _output): + nonlocal layer_calls + layer_calls += 1 + + model.layers[0].register_forward_hook(count_layer_calls) + runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} + model.enable_cache( + 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"], + ) + ) + + with torch.no_grad(), model.cache_context("cond"): + model(**self.get_dummy_inputs()) + runtime.update(step=1, sigma=0.6) + 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"] + assert reduction_count == 2 + assert layer_calls == 2 + assert state.consecutive_cached == 0 + assert [step for step, _, _ in state.history] == [1] + + def test_cosmos3_sea_cache_parallelism_without_process_group_fails_open(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + model._cp_shard_fn = lambda und_seq, gen_seq, rotary_emb: (und_seq, gen_seq, rotary_emb) + model._cp_gather_fn = lambda und_out, gen_out: (und_out, gen_out) + layer_calls = 0 + + def count_layer_calls(_module, _args, _output): + nonlocal layer_calls + layer_calls += 1 + + model.layers[0].register_forward_hook(count_layer_calls) + runtime = {"step": 0, "sigma": 0.9, "num_steps": 3} + model.enable_cache( + 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"], + ) + ) + + with torch.no_grad(), model.cache_context("cond"): + model(**self.get_dummy_inputs()) + runtime.update(step=1, sigma=0.6) + with torch.no_grad(), model.cache_context("cond"): + model(**self.get_dummy_inputs()) + + root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) + assert layer_calls == 2 + assert root_hook.state_manager._state_cache["cond"].consecutive_cached == 0 + @pytest.mark.skipif(not hasattr(torch, "compile"), reason="torch.compile is unavailable") def test_cosmos3_sea_cache_regional_compile_fullgraph_without_recompile(self): model = self.model_class(**self.get_init_dict()).to(torch_device).eval() From 0af0814a2db5ee0b9e3e63217aa2d7c2082c781e Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 7 Sep 2026 09:36:57 -0700 Subject: [PATCH 19/20] Support single-stream models in SeaCache --- src/diffusers/hooks/sea_cache.py | 24 ++++++++---- tests/models/testing_utils/cache.py | 60 +++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 7 deletions(-) diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index 71244fb2217b..be12730990f8 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -141,7 +141,7 @@ class _SeaCacheForwardMetadata: class SeaCacheContextState(BaseState): def __init__(self): - self.history: list[tuple[int, torch.Tensor, torch.Tensor]] = [] + self.history: list[tuple[int, torch.Tensor | None, torch.Tensor]] = [] self.gate_key: tuple[int, float] | None = None self.gate_should_compute = True self.previous_indicator: list[torch.Tensor] | None = None @@ -308,19 +308,20 @@ def _record_full_execution( state: SeaCacheContextState, gen_output: torch.Tensor, und_output: torch.Tensor | None, + requires_und_output: bool = True, ) -> 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 (und_output is not None or not requires_und_output) and gen_output.shape == state.gen_input.shape ): state.history.append( ( state.step_index, - und_output.detach().clone(), + und_output.detach().clone() if und_output is not None else None, (gen_output - state.gen_input).detach().clone(), ) ) @@ -804,6 +805,17 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): residual_history = state.history[-(self.config.residual_order + 1) :] _, cached_und, cached_residual = residual_history[-1] + expects_auxiliary_stream = self._metadata.return_encoder_hidden_states_index is not None + if expects_auxiliary_stream: + auxiliary_stream_mismatch = encoder_hidden_states is None or cached_und is None + else: + auxiliary_stream_mismatch = cached_und is not None + if expects_auxiliary_stream and encoder_hidden_states is not None and cached_und is not None: + auxiliary_stream_mismatch = auxiliary_stream_mismatch or ( + cached_und.shape != encoder_hidden_states.shape + or cached_und.device != encoder_hidden_states.device + or cached_und.dtype != encoder_hidden_states.dtype + ) if ( any( residual.shape != hidden_states.shape @@ -811,10 +823,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): 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 + or auxiliary_stream_mismatch ): state.history = [] state.accumulated_distance = 0.0 @@ -882,6 +891,7 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): state, gen_output=hidden_states, und_output=encoder_hidden_states, + requires_und_output=self._metadata.return_encoder_hidden_states_index is not None, ) return output diff --git a/tests/models/testing_utils/cache.py b/tests/models/testing_utils/cache.py index c8d757140348..e43412b88491 100644 --- a/tests/models/testing_utils/cache.py +++ b/tests/models/testing_utils/cache.py @@ -26,6 +26,7 @@ SeaCacheConfig, TaylorSeerCacheConfig, ) +from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry 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 @@ -552,6 +553,62 @@ def _test_reset_stateful_cache(self): model._reset_stateful_cache() model.disable_cache() + @torch.no_grad() + def _test_single_stream_cache_inference(self): + class SingleStreamBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, hidden_states): + self.calls += 1 + return hidden_states * 2 + 1 + + class SingleStreamTransformer(torch.nn.Module, CacheMixin): + def __init__(self): + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([SingleStreamBlock(), SingleStreamBlock()]) + + def forward(self, hidden_states, raw_vision): + for block in self.transformer_blocks: + hidden_states = block(hidden_states) + return hidden_states + + TransformerBlockRegistry.register( + SingleStreamBlock, + TransformerBlockMetadata(return_hidden_states_index=0), + ) + runtime = {"step": 0, "sigma": 0.9, "num_steps": 2} + model = SingleStreamTransformer().eval() + model.enable_cache( + SeaCacheConfig( + threshold=100.0, + residual_order=0, + retention_steps=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"], + raw_vision_callback=lambda module, args, kwargs: [kwargs["raw_vision"]], + ) + ) + + hidden_states = torch.zeros(1, 2, 3) + raw_vision = torch.ones(2, 1, 2, 2) + with model.cache_context("single_stream"): + model(hidden_states=hidden_states, raw_vision=raw_vision) + + root_hook = model._diffusers_hook.get_hook(_SEA_CACHE_ROOT_HOOK) + assert root_hook.state_manager._state_cache["single_stream"].history[-1][1] is None + assert [block.calls for block in model.transformer_blocks] == [1, 1] + + runtime.update(step=1, sigma=0.6) + with model.cache_context("single_stream"): + output = model(hidden_states=hidden_states + 0.25, raw_vision=raw_vision + 0.01) + + assert [block.calls for block in model.transformer_blocks] == [1, 1] + torch.testing.assert_close(output, torch.full_like(output, 3.25)) + @require_cache_mixin def test_sea_cache_enable_disable_state(self): self._test_cache_enable_disable_state() @@ -576,6 +633,9 @@ def test_sea_cache_context_manager(self): def test_sea_cache_reset_stateful_cache(self): self._test_reset_stateful_cache() + def test_sea_cache_single_stream_inference(self): + self._test_single_stream_cache_inference() + @is_cache class FasterCacheConfigMixin: From 14bba9fdd009fb489f57510cc1e25224f0b96838 Mon Sep 17 00:00:00 2001 From: Yuliya Zhautouskaya Date: Mon, 7 Sep 2026 11:20:43 -0700 Subject: [PATCH 20/20] Add Wan T2V support to SeaCache --- docs/source/en/api/pipelines/cosmos3.md | 10 +++++ docs/source/en/optimization/cache.md | 17 ++++++++ src/diffusers/hooks/sea_cache.py | 41 ++++++++++++++----- .../test_models_transformer_wan.py | 5 +++ 4 files changed, 62 insertions(+), 11 deletions(-) diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index c8b56aa8ecf9..b3879b658b64 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -691,6 +691,16 @@ The same model-level API works with [`Cosmos3OmniPipeline`], [`Cosmos3OmniModula `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. +Cosmos 3 keeps the SeaCache gate outside its repeated decoder layers, so it is compatible with regional compilation. +Compile the layers after enabling the cache: + +```python +pipe.transformer.compile_repeated_blocks(fullgraph=True) +``` + +SeaCache also supports the Cosmos 3 Ulysses context-parallel and DTensor-based tensor-parallel helpers documented +below. Cache decisions are synchronized across ranks. Full-model compilation and other model-sharding strategies are not claimed. + ## 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 8fe3915a52cd..fc5eac7cfb46 100644 --- a/docs/source/en/optimization/cache.md +++ b/docs/source/en/optimization/cache.md @@ -75,6 +75,23 @@ successive denoising steps. When the accumulated indicator change remains below transformer block stack and predicts its output from cached residuals. The indicator is computed from the raw vision latents, including clean conditioning frames for image-to-video generation. +The implementation provides built-in adapters for the following models: + +- **Cosmos 3** is the primary optimized and benchmarked integration. It caches the complete decoder stack through a + post-normalization boundary and supports eager inference and regional compilation. +- **Wan T2V** uses the generic repeated-block path in eager mode. This integration demonstrates how another + single-stream video transformer can provide raw vision latents to SeaCache; it is not a claim that the same cache + parameters are optimal for Wan or that other Wan variants are supported. + +Other video transformers can integrate with the generic path when they use `CacheMixin`, expose a recognized repeated +block list, and register the block input/output layout in `TransformerBlockRegistry`. The pipeline must enter a +`cache_context` for every transformer call, using separate context names for independent trajectories such as +conditional and unconditional guidance. Pass a `raw_vision_callback` that returns the noisy vision latents, in addition +to the scheduler metadata callbacks shown below. Validate output quality and tune the cache parameters for each model +and scheduler; support and benchmark results do not transfer automatically from Cosmos 3. + +### Cosmos 3 + SeaCache is disabled by default. Enable it on the transformer and provide callbacks for the active scheduler step, sigma, and number of inference steps: diff --git a/src/diffusers/hooks/sea_cache.py b/src/diffusers/hooks/sea_cache.py index be12730990f8..c8d1efce0819 100644 --- a/src/diffusers/hooks/sea_cache.py +++ b/src/diffusers/hooks/sea_cache.py @@ -43,16 +43,19 @@ class SeaCacheConfig: `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. + change stays below `threshold`, the repeated transformer block stack is replaced with a cached residual. Cosmos 3 + is the primary optimized and benchmarked integration; its residual spans the decoder stack and final pathway + normalization, while input packing and modality prediction heads still execute. Wan T2V provides a built-in + eager-only example of the generic repeated-block integration. Other video transformers can integrate by registering + their block metadata and providing `raw_vision_callback`; their pipeline must enter a `cache_context` for each + transformer trajectory. 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. + Order used to predict the transformer 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`): @@ -69,8 +72,8 @@ class SeaCacheConfig: num_inference_steps_callback (`Callable[[], int]`): Callback returning the number of scheduler steps in the current pipeline call. raw_vision_callback (`Callable`, *optional*): - Advanced model adapter returning raw vision latents with shape `(C, T, H, W)`. Cosmos 3 uses its native - adapter when this is omitted. + Advanced model adapter returning raw vision latents with shape `(C, T, H, W)`. Cosmos 3 and compatible Wan + T2V transformers use built-in adapters when this is omitted. Example: ```python @@ -365,6 +368,17 @@ def _prepare_cosmos3_raw_vision_metadata( return raw_vision if raw_vision and has_noisy_vision else None +def _prepare_wan_t2v_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 + hidden_states = bound_arguments.get("hidden_states") + if not isinstance(hidden_states, torch.Tensor) or hidden_states.ndim != 5: + return None + return list(hidden_states.unbind(dim=0)) + + def _apply_sea_filter( hidden_states: torch.Tensor, sigma: float, @@ -952,10 +966,9 @@ 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 decoder-stack 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. + SeaCache reuses a cached residual for a transformer's repeated block stack. Model adapters may customize indicator + extraction and the residual boundary; the generic path uses registered block metadata and `raw_vision_callback`. + Compatibility and cache parameters must be validated for each model. Args: module (`torch.nn.Module`): @@ -964,12 +977,18 @@ def apply_sea_cache(module: torch.nn.Module, config: SeaCacheConfig) -> None: SeaCache configuration. """ from ..models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer + from ..models.transformers.transformer_wan import WanTransformer3DModel unwrapped_module = unwrap_module(module) is_cosmos3 = isinstance(unwrapped_module, Cosmos3OmniTransformer) + is_wan_t2v = isinstance(unwrapped_module, WanTransformer3DModel) and ( + unwrapped_module.config.in_channels == unwrapped_module.config.out_channels + ) raw_vision_callback = config.raw_vision_callback if raw_vision_callback is None and is_cosmos3: raw_vision_callback = _prepare_cosmos3_raw_vision_metadata + elif raw_vision_callback is None and is_wan_t2v: + raw_vision_callback = _prepare_wan_t2v_raw_vision_metadata post_norm_modules = None if is_cosmos3: diff --git a/tests/models/transformers/test_models_transformer_wan.py b/tests/models/transformers/test_models_transformer_wan.py index cbd848f086e4..3fa550f81c1e 100644 --- a/tests/models/transformers/test_models_transformer_wan.py +++ b/tests/models/transformers/test_models_transformer_wan.py @@ -27,6 +27,7 @@ LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, + SeaCacheTesterMixin, SingleFileTesterMixin, TorchAoTesterMixin, TorchCompileTesterMixin, @@ -108,6 +109,10 @@ class TestWanTransformer3D(WanTransformer3DTesterConfig, ModelTesterMixin): """Core model tests for Wan Transformer 3D.""" +class TestWanTransformer3DSeaCache(WanTransformer3DTesterConfig, SeaCacheTesterMixin): + cache_input_key = "hidden_states" + + class TestWanTransformer3DMemory(WanTransformer3DTesterConfig, MemoryTesterMixin): """Memory optimization tests for Wan Transformer 3D."""