diffusers/matrix-game-2-modular
021
1# Copyright 2025 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15from typing import Any, List, Tuple16 17import torch18 19from diffusers.configuration_utils import FrozenDict20from diffusers.guiders import ClassifierFreeGuidance21from diffusers.models import AutoModel, WanTransformer3DModel22from diffusers.schedulers import UniPCMultistepScheduler23from diffusers.utils import logging24from diffusers.utils.torch_utils import randn_tensor25from diffusers.modular_pipelines import (26 BlockState,27 LoopSequentialPipelineBlocks,28 ModularPipelineBlocks,29 PipelineState,30 ModularPipeline31)32from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam33 34 35logger = logging.get_logger(__name__) # pylint: disable=invalid-name36 37 38class MatrixGameWanLoopDenoiser(ModularPipelineBlocks):39 model_name = "MatrixGameWan"40 frame_seq_length = 88041 42 @property43 def expected_components(self) -> List[ComponentSpec]:44 return [45 ComponentSpec(46 "guider",47 ClassifierFreeGuidance,48 config=FrozenDict({"guidance_scale": 5.0}),49 default_creation_method="from_config",50 ),51 ComponentSpec("transformer", AutoModel),52 ]53 54 @property55 def description(self) -> str:56 return (57 "Step within the denoising loop that denoise the latents with guidance. "58 "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "59 "object (e.g. `MatrixGameWanDenoiseLoopWrapper`)"60 )61 62 @property63 def inputs(self) -> List[Tuple[str, Any]]:64 return [65 InputParam("attention_kwargs"),66 InputParam(67 "latents",68 required=True,69 type_hint=torch.Tensor,70 description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",71 ),72 InputParam(73 "image_mask_latents",74 required=True,75 type_hint=torch.Tensor,76 ),77 InputParam(78 "image_embeds",79 required=True,80 type_hint=torch.Tensor,81 ),82 InputParam(83 "keyboard_conditions",84 required=True,85 type_hint=torch.Tensor,86 ),87 InputParam(88 "mouse_conditions",89 required=True,90 type_hint=torch.Tensor,91 ),92 InputParam(93 "num_inference_steps",94 required=True,95 type_hint=int,96 default=4,97 description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",98 ),99 InputParam(100 kwargs_type="guider_input_fields",101 description=(102 "All conditional model inputs that need to be prepared with guider. "103 "It should contain prompt_embeds/negative_prompt_embeds. "104 "Please add `kwargs_type=guider_input_fields` to their parameter spec (`OutputParam`) when they are created and added to the pipeline state"105 ),106 ),107 ]108 109 @torch.no_grad()110 def __call__(111 self, components: ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor112 ) -> PipelineState:113 cond_concat = block_state.image_mask_latents114 keyboard_conditions = block_state.keyboard_conditions115 mouse_conditions = block_state.mouse_conditions116 visual_context = block_state.image_embeds117 118 transformer_dtype = components.transformer.dtype119 components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)120 121 # Prepare mini‐batches according to guidance method and `guider_input_fields`122 # Each guider_state_batch will have .prompt_embeds, .time_ids, text_embeds, image_embeds.123 # e.g. for CFG, we prepare two batches: one for uncond, one for cond124 # for first batch, guider_state_batch.prompt_embeds correspond to block_state.prompt_embeds125 # for second batch, guider_state_batch.prompt_embeds correspond to block_state.negative_prompt_embeds126 guider_state = components.guider.prepare_inputs(block_state, {})127 128 # run the denoiser for each guidance batch129 for guider_state_batch in guider_state:130 components.guider.prepare_models(components.transformer)131 cond_kwargs = guider_state_batch.as_dict()132 133 # Predict the noise residual134 # store the noise_pred in guider_state_batch so that we can apply guidance across all batches135 guider_state_batch.noise_pred = components.transformer(136 x=block_state.latents.to(transformer_dtype),137 t=t.expand(block_state.latents.shape[0], block_state.num_frames_per_block),138 visual_context=visual_context.to(transformer_dtype),139 cond_concat=cond_concat.to(transformer_dtype),140 keyboard_cond=keyboard_conditions,141 mouse_cond=mouse_conditions,142 kv_cache=block_state.kv_cache,143 kv_cache_mouse=block_state.kv_cache_mouse,144 kv_cache_keyboard=block_state.kv_cache_keyboard,145 crossattn_cache=block_state.kv_cache_cross_attn,146 current_start=block_state.current_frame_idx * self.frame_seq_length,147 num_frames_per_block=block_state.num_frames_per_block,148 )[0]149 components.guider.cleanup_models(components.transformer)150 151 # Perform guidance152 block_state.noise_pred = components.guider(guider_state)[0]153 154 return components, block_state155 156 157class MatrixGameWanLoopAfterDenoiser(ModularPipelineBlocks):158 model_name = "MatrixGameWan"159 160 @property161 def expected_components(self) -> List[ComponentSpec]:162 return [163 ComponentSpec("scheduler", UniPCMultistepScheduler),164 ]165 166 @property167 def description(self) -> str:168 return (169 "step within the denoising loop that update the latents. "170 "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "171 "object (e.g. `MatrixGameWanDenoiseLoopWrapper`)"172 )173 174 @property175 def inputs(self) -> List[Tuple[str, Any]]:176 return []177 178 @property179 def intermediate_inputs(self) -> List[str]:180 return [181 InputParam("generator"),182 ]183 184 @property185 def intermediate_outputs(self) -> List[OutputParam]:186 return [OutputParam("latents", type_hint=torch.Tensor, description="The denoised latents")]187 188 @torch.no_grad()189 def __call__(self, components: ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):190 # Perform scheduler step using the predicted output191 latents_dtype = block_state.latents.dtype192 193 step_index = components.scheduler.index_for_timestep(t)194 sigma_t = components.scheduler.sigmas[step_index]195 196 latents = block_state.latents.double() - sigma_t.double() * block_state.noise_pred.double()197 block_state.latents = latents198 199 if block_state.latents.dtype != latents_dtype:200 block_state.latents = block_state.latents.to(latents_dtype)201 202 return components, block_state203 204 205class MatrixGameWanDenoiseLoopWrapper(LoopSequentialPipelineBlocks):206 model_name = "MatrixGameWan"207 frame_seq_length = 880208 local_attn_size = 6209 num_transformer_blocks = 30210 211 def _initialize_kv_cache(self, batch_size, dtype, device):212 """213 Initialize a Per-GPU KV cache for the Wan model.214 """215 cache = []216 if self.local_attn_size != -1:217 # Use the local attention size to compute the KV cache size218 kv_cache_size = self.local_attn_size * self.frame_seq_length219 else:220 # Use the default KV cache size221 kv_cache_size = 15 * 1 * self.frame_seq_length # 32760222 223 for _ in range(self.num_transformer_blocks):224 cache.append({225 "k": torch.zeros((batch_size, kv_cache_size, 12, 128), dtype=dtype, device=device),226 "v": torch.zeros((batch_size, kv_cache_size, 12, 128), dtype=dtype, device=device),227 "global_end_index": torch.tensor([0], dtype=torch.long, device=device),228 "local_end_index": torch.tensor([0], dtype=torch.long, device=device)229 })230 231 return cache # always store the clean cache232 233 def _initialize_kv_cache_mouse_and_keyboard(self, batch_size, dtype, device):234 """235 Initialize a Per-GPU KV cache for the Wan model.236 """237 kv_cache_mouse = []238 kv_cache_keyboard = []239 if self.local_attn_size != -1:240 kv_cache_size = self.local_attn_size241 else:242 kv_cache_size = 15 * 1243 for _ in range(self.num_transformer_blocks):244 kv_cache_keyboard.append({245 "k": torch.zeros([batch_size, kv_cache_size, 16, 64], dtype=dtype, device=device),246 "v": torch.zeros([batch_size, kv_cache_size, 16, 64], dtype=dtype, device=device),247 "global_end_index": torch.tensor([0], dtype=torch.long, device=device),248 "local_end_index": torch.tensor([0], dtype=torch.long, device=device)249 })250 kv_cache_mouse.append({251 "k": torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, 16, 64], dtype=dtype, device=device),252 "v": torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, 16, 64], dtype=dtype, device=device),253 "global_end_index": torch.tensor([0], dtype=torch.long, device=device),254 "local_end_index": torch.tensor([0], dtype=torch.long, device=device)255 })256 return kv_cache_mouse, kv_cache_keyboard # always store the clean cache257 258 def _initialize_crossattn_cache(self, batch_size, dtype, device):259 """260 Initialize a Per-GPU cross-attention cache for the Wan model.261 """262 crossattn_cache = []263 264 for _ in range(self.num_transformer_blocks):265 crossattn_cache.append({266 "k": torch.zeros([batch_size, 257, 12, 128], dtype=dtype, device=device),267 "v": torch.zeros([batch_size, 257, 12, 128], dtype=dtype, device=device),268 "is_init": False269 })270 271 return crossattn_cache272 273 @property274 def description(self) -> str:275 return (276 "Pipeline block that iteratively denoise the latents over `timesteps`. "277 "The specific steps with each iteration can be customized with `sub_blocks` attributes"278 )279 280 @property281 def loop_expected_components(self) -> List[ComponentSpec]:282 return [283 ComponentSpec(284 "guider",285 ClassifierFreeGuidance,286 config=FrozenDict({"guidance_scale": 5.0}),287 default_creation_method="from_config",288 ),289 ComponentSpec("scheduler", UniPCMultistepScheduler),290 ComponentSpec("transformer", AutoModel),291 ]292 293 @property294 def loop_inputs(self) -> List[InputParam]:295 return [296 InputParam(297 "timesteps",298 required=True,299 type_hint=torch.Tensor,300 description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",301 ),302 InputParam(303 "num_inference_steps",304 required=True,305 type_hint=int,306 description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",307 ),308 InputParam(309 "num_frames_per_block",310 required=True,311 type_hint=int,312 default=3,313 ),314 ]315 316 @torch.no_grad()317 def __call__(318 self, components: ModularPipeline, state: PipelineState319 ) -> PipelineState:320 block_state = self.get_block_state(state)321 transformer_dtype = components.transformer.dtype322 323 num_frames_per_block = block_state.num_frames_per_block324 latents = block_state.latents.to(transformer_dtype)325 image_mask_latents = block_state.image_mask_latents.to(transformer_dtype)326 mouse_conditions = block_state.mouse_conditions.unsqueeze(0).to(transformer_dtype)327 keyboard_conditions = block_state.keyboard_conditions.unsqueeze(0).to(transformer_dtype)328 visual_context = block_state.image_embeds329 330 batch_size, num_channels, num_frames, height, width = latents.shape331 output = torch.zeros(332 (batch_size, num_channels, num_frames, height, width),333 device=latents.device,334 dtype=latents.dtype,335 )336 337 current_frame_idx = 0338 num_blocks = num_frames // num_frames_per_block339 340 kv_cache = self._initialize_kv_cache(batch_size, latents.dtype, latents.device)341 kv_cache_mouse, kv_cache_keyboard = self._initialize_kv_cache_mouse_and_keyboard(batch_size, latents.dtype, latents.device)342 kv_cache_cross_attn = self._initialize_crossattn_cache(batch_size, latents.dtype, latents.device)343 344 block_state.kv_cache = kv_cache345 block_state.kv_cache_mouse = kv_cache_mouse346 block_state.kv_cache_keyboard = kv_cache_keyboard347 block_state.kv_cache_cross_attn = kv_cache_cross_attn348 349 for _ in range(num_blocks):350 block_state.current_frame_idx = current_frame_idx351 block_state.image_mask_latents = image_mask_latents[352 :, :, current_frame_idx : current_frame_idx + num_frames_per_block353 ]354 cond_idx = 1 + 4 * (current_frame_idx + num_frames_per_block - 1)355 block_state.mouse_conditions = mouse_conditions[:, :cond_idx]356 block_state.keyboard_conditions = keyboard_conditions[:, :cond_idx]357 358 block_state.latents = latents[359 :, :, current_frame_idx : current_frame_idx + num_frames_per_block360 ]361 for i, t in enumerate(block_state.timesteps):362 components, block_state = self.loop_step(363 components, block_state, i=i, t=t364 )365 366 if i < (block_state.num_inference_steps - 1):367 t1 = components.scheduler.timesteps[i+1]368 block_state.latents = components.scheduler.add_noise(369 block_state.latents,370 randn_tensor(371 block_state.latents.shape,372 device=block_state.latents.device,373 dtype=block_state.latents.dtype374 ),375 t1.expand(block_state.latents.shape[0])376 )377 378 output[379 :, :, current_frame_idx : current_frame_idx + num_frames_per_block380 ] = block_state.latents381 382 components.transformer(383 x=block_state.latents,384 t=t.expand(block_state.latents.shape[0], block_state.num_frames_per_block) * 0.0,385 visual_context=visual_context,386 cond_concat=block_state.image_mask_latents,387 keyboard_cond=block_state.keyboard_conditions,388 mouse_cond=block_state.mouse_conditions,389 kv_cache=block_state.kv_cache,390 kv_cache_mouse=block_state.kv_cache_mouse,391 kv_cache_keyboard=block_state.kv_cache_keyboard,392 crossattn_cache=block_state.kv_cache_cross_attn,393 current_start=block_state.current_frame_idx * self.frame_seq_length,394 num_frames_per_block=block_state.num_frames_per_block,395 )[0]396 current_frame_idx += num_frames_per_block397 398 block_state.latents = output399 self.set_block_state(state, block_state)400 401 return components, state402 403 404class MatrixGameWanDenoiseStep(MatrixGameWanDenoiseLoopWrapper):405 block_classes = [406 MatrixGameWanLoopDenoiser,407 MatrixGameWanLoopAfterDenoiser,408 ]409 block_names = ["denoiser", "after_denoiser"]410 411 @property412 def description(self) -> str:413 return (414 "Denoise step that iteratively denoise the latents. \n"415 "Its loop logic is defined in `MatrixGameWanDenoiseLoopWrapper.__call__` method \n"416 "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"417 " - `MatrixGameWanLoopDenoiser`\n"418 " - `MatrixGameWanLoopAfterDenoiser`\n"419 "This block supports both text2vid tasks."420 )421