diff --git a/src/diffusers/hooks/_helpers.py b/src/diffusers/hooks/_helpers.py index 1c17d349562b..208fa810c61f 100644 --- a/src/diffusers/hooks/_helpers.py +++ b/src/diffusers/hooks/_helpers.py @@ -196,6 +196,7 @@ def _register_transformer_blocks_metadata(): MotifVideoTransformerBlock, ) from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock + from ..models.transformers.transformer_qwenimage21 import QwenImage21TransformerBlock from ..models.transformers.transformer_wan import WanTransformerBlock from ..models.transformers.transformer_z_image import ZImageTransformerBlock @@ -358,6 +359,15 @@ def _register_transformer_blocks_metadata(): ), ) + # QwenImage2.1 + TransformerBlockRegistry.register( + model_class=QwenImage21TransformerBlock, + metadata=TransformerBlockMetadata( + return_hidden_states_index=0, + return_encoder_hidden_states_index=None, + ), + ) + # HunyuanImage2.1 TransformerBlockRegistry.register( model_class=HunyuanImageTransformerBlock, diff --git a/src/diffusers/hooks/mag_cache.py b/src/diffusers/hooks/mag_cache.py index eb76e343b30e..e897915ff813 100644 --- a/src/diffusers/hooks/mag_cache.py +++ b/src/diffusers/hooks/mag_cache.py @@ -249,6 +249,10 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs): if diff > 0: output = output.clone() output[:, diff:, :] = output[:, diff:, :] + res + elif diff < 0: + # The residual was computed on a longer sequence (e.g. the QwenImage2.1 KV-cache prefill step + # covers prefix + target tokens, later steps only the target); the target tokens sit at the end. + output = output + res[:, -output.shape[1] :, :] else: logger.warning( f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. " @@ -371,6 +375,12 @@ def _perform_calibration_step(self, state: MagCacheState, current_residual: torc # norm(dim=-1) gives magnitude of each token vector curr_norm = torch.linalg.norm(current_residual.float(), dim=-1) prev_norm = torch.linalg.norm(state.previous_residual.float(), dim=-1) + if curr_norm.ndim == 2 and prev_norm.ndim == 2 and curr_norm.shape[1] != prev_norm.shape[1]: + # Sequence length changed between steps (e.g. KV-cache prefill vs. cached steps); compare the + # trailing tokens both steps share. + num_tokens = min(curr_norm.shape[1], prev_norm.shape[1]) + curr_norm = curr_norm[:, -num_tokens:] + prev_norm = prev_norm[:, -num_tokens:] # Avoid division by zero ratio = (curr_norm / (prev_norm + 1e-8)).mean().item() diff --git a/tests/models/transformers/test_models_transformer_qwenimage21.py b/tests/models/transformers/test_models_transformer_qwenimage21.py index f6d0a0004d0c..0308b39b6605 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage21.py +++ b/tests/models/transformers/test_models_transformer_qwenimage21.py @@ -25,6 +25,7 @@ from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, + MagCacheTesterMixin, MemoryTesterMixin, ModelTesterMixin, SingleFileTesterMixin, @@ -364,6 +365,10 @@ def test_taylorseer_cache_with_kv_cache(self, use_lite_mode): assert not torch.allclose(with_kv_cache[3], uncached[3], atol=1e-5) +class TestQwenImage21TransformerMagCache(QwenImage21TransformerTesterConfig, MagCacheTesterMixin): + """MagCache tests for QwenImage 2.1 Transformer.""" + + class TestQwenImage21TransformerSingleFile(QwenImage21TransformerTesterConfig, SingleFileTesterMixin): @property def ckpt_path(self): diff --git a/tests/pipelines/qwenimage21/test_qwenimage21.py b/tests/pipelines/qwenimage21/test_qwenimage21.py index 72b960c33936..21f07e2ece6c 100644 --- a/tests/pipelines/qwenimage21/test_qwenimage21.py +++ b/tests/pipelines/qwenimage21/test_qwenimage21.py @@ -38,6 +38,7 @@ BasePipelineTesterConfig, LoraMemoryTesterMixin, LoraTesterMixin, + MagCacheTesterMixin, MemoryTesterMixin, PipelineTesterMixin, ) @@ -347,3 +348,7 @@ def test_simple_inference_with_text_denoiser_lora_and_scale(self, base_pipe_outp class TestQwenImage21PipelineLoRAMemory(QwenImage21PipelineTesterConfig, LoraMemoryTesterMixin): """LoRA x memory-optimization tests for the Qwen-Image 2.1 pipeline.""" + + +class TestQwenImage21PipelineMagCache(QwenImage21PipelineTesterConfig, MagCacheTesterMixin): + """MagCache tests for the Qwen-Image 2.1 pipeline."""