diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index 9cbe5bc8108f..2a9ae6b5a383 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -27,6 +27,7 @@ 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" _cls: Type = None _cached_parameter_indices: dict[str, int] = None @@ -174,6 +175,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 +232,20 @@ def _register_transformer_blocks_metadata(): ), ) + # Cosmos3 (Omni) + # MoT dual-stream decoder layer: forward(und_seq, gen_seq, rotary_emb) -> (und_seq, gen_seq). + # The gen stream carries the denoised tokens, so it is mapped to hidden_states for cache + # decisions; the und (text/understanding) stream is quasi-static across denoising steps. + 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", + ), + ) + # Flux TransformerBlockRegistry.register( model_class=FluxTransformerBlock, diff --git a/src/diffusers/hooks/first_block_cache.py b/src/diffusers/hooks/first_block_cache.py index 685ccd383674..77a628e1c9c5 100644 --- a/src/diffusers/hooks/first_block_cache.py +++ b/src/diffusers/hooks/first_block_cache.py @@ -76,7 +76,9 @@ def initialize_hook(self, module): return module def new_forward(self, module: torch.nn.Module, *args, **kwargs): - original_hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs) + original_hidden_states = self._metadata._get_parameter_from_args_kwargs( + self._metadata.hidden_states_argument_name, args, kwargs + ) output = self.fn_ref.original_forward(*args, **kwargs) is_output_tuple = isinstance(output, tuple) @@ -155,11 +157,13 @@ def initialize_hook(self, module): return module def new_forward(self, module: torch.nn.Module, *args, **kwargs): - original_hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs) + original_hidden_states = self._metadata._get_parameter_from_args_kwargs( + self._metadata.hidden_states_argument_name, args, kwargs + ) original_encoder_hidden_states = None if self._metadata.return_encoder_hidden_states_index is not None: original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs( - "encoder_hidden_states", args, kwargs + self._metadata.encoder_hidden_states_argument_name, args, kwargs ) shared_state = self.state_manager.get_state() diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py index 3e331ba74628..fc0fd1a27907 100644 --- a/src/diffusers/models/transformers/transformer_cosmos3.py +++ b/src/diffusers/models/transformers/transformer_cosmos3.py @@ -21,6 +21,7 @@ from ...loaders import PeftAdapterMixin 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 @@ -294,7 +295,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/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 12405755309f..25ecc4439905 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -1641,32 +1641,33 @@ 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, - ) + with self.transformer.cache_context("cond"): + preds_vision, preds_sound, preds_action = self.transformer( + input_ids=cond_packed_static["input_ids"], + text_indexes=cond_packed_static["text_indexes"], + position_ids=cond_packed_static["position_ids"], + und_len=cond_packed_static["und_len"], + sequence_length=cond_packed_static["sequence_length"], + vision_tokens=[vision_tokens], + vision_token_shapes=cond_packed_static["vision_token_shapes"], + vision_sequence_indexes=cond_packed_static["vision_sequence_indexes"], + vision_mse_loss_indexes=cond_packed_static["vision_mse_loss_indexes"], + vision_timesteps=vision_timesteps, + vision_noisy_frame_indexes=cond_packed_static["vision_noisy_frame_indexes"], + sound_tokens=[sound_tokens] if sound_tokens is not None else None, + sound_token_shapes=cond_packed_static.get("sound_token_shapes"), + sound_sequence_indexes=cond_packed_static.get("sound_sequence_indexes"), + sound_mse_loss_indexes=cond_packed_static.get("sound_mse_loss_indexes"), + sound_timesteps=sound_timesteps, + sound_noisy_frame_indexes=cond_packed_static.get("sound_noisy_frame_indexes"), + action_tokens=[action_tokens] if action_tokens is not None else None, + action_token_shapes=cond_packed_static.get("action_token_shapes"), + action_sequence_indexes=cond_packed_static.get("action_sequence_indexes"), + action_mse_loss_indexes=cond_packed_static.get("action_mse_loss_indexes"), + action_timesteps=action_timesteps, + action_noisy_frame_indexes=cond_packed_static.get("action_noisy_frame_indexes"), + action_domain_ids=[action_domain_id] if action_domain_id is not None else None, + ) cond_v_vision, cond_v_sound, cond_v_action = self._mask_velocity_predictions( preds_vision, preds_sound, @@ -1680,32 +1681,33 @@ 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, - ) + with self.transformer.cache_context("uncond"): + preds_vision, preds_sound, preds_action = self.transformer( + input_ids=uncond_packed_static["input_ids"], + text_indexes=uncond_packed_static["text_indexes"], + position_ids=uncond_packed_static["position_ids"], + und_len=uncond_packed_static["und_len"], + sequence_length=uncond_packed_static["sequence_length"], + vision_tokens=[vision_tokens], + vision_token_shapes=uncond_packed_static["vision_token_shapes"], + vision_sequence_indexes=uncond_packed_static["vision_sequence_indexes"], + vision_mse_loss_indexes=uncond_packed_static["vision_mse_loss_indexes"], + vision_timesteps=vision_timesteps, + vision_noisy_frame_indexes=uncond_packed_static["vision_noisy_frame_indexes"], + sound_tokens=[sound_tokens] if sound_tokens is not None else None, + sound_token_shapes=uncond_packed_static.get("sound_token_shapes"), + sound_sequence_indexes=uncond_packed_static.get("sound_sequence_indexes"), + sound_mse_loss_indexes=uncond_packed_static.get("sound_mse_loss_indexes"), + sound_timesteps=sound_timesteps, + sound_noisy_frame_indexes=uncond_packed_static.get("sound_noisy_frame_indexes"), + action_tokens=[action_tokens] if action_tokens is not None else None, + action_token_shapes=uncond_packed_static.get("action_token_shapes"), + action_sequence_indexes=uncond_packed_static.get("action_sequence_indexes"), + action_mse_loss_indexes=uncond_packed_static.get("action_mse_loss_indexes"), + action_timesteps=action_timesteps, + action_noisy_frame_indexes=uncond_packed_static.get("action_noisy_frame_indexes"), + action_domain_ids=[action_domain_id] if action_domain_id is not None else None, + ) uncond_v_vision, uncond_v_sound, uncond_v_action = self._mask_velocity_predictions( preds_vision, preds_sound,