add modular blockset for wan 2.2 vace

src/diffusers/__init__.py

@@ -553,6 +553,8 @@
             "Wan22Image2VideoBlocks",
             "Wan22Image2VideoModularPipeline",
             "Wan22ModularPipeline",
+            "Wan22VaceBlocks",
+            "Wan22VaceModularPipeline",
             "WanBlocks",
             "WanImage2VideoAutoBlocks",
             "WanImage2VideoModularPipeline",
@@ -1379,6 +1381,8 @@
             Wan22Image2VideoBlocks,
             Wan22Image2VideoModularPipeline,
             Wan22ModularPipeline,
+            Wan22VaceBlocks,
+            Wan22VaceModularPipeline,
             WanBlocks,
             WanImage2VideoAutoBlocks,
             WanImage2VideoModularPipeline,

src/diffusers/modular_pipelines/__init__.py

@@ -52,10 +52,12 @@
         "Wan22Blocks",
         "WanImage2VideoAutoBlocks",
         "Wan22Image2VideoBlocks",
+        "Wan22VaceBlocks",
         "WanModularPipeline",
         "Wan22ModularPipeline",
         "WanImage2VideoModularPipeline",
         "Wan22Image2VideoModularPipeline",
+        "Wan22VaceModularPipeline",
     ]
     _import_structure["helios"] = [
         "HeliosAutoBlocks",
@@ -211,6 +213,8 @@
             Wan22Image2VideoBlocks,
             Wan22Image2VideoModularPipeline,
             Wan22ModularPipeline,
+            Wan22VaceBlocks,
+            Wan22VaceModularPipeline,
             WanBlocks,
             WanImage2VideoAutoBlocks,
             WanImage2VideoModularPipeline,

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -133,6 +133,7 @@ def _helios_pyramid_map_fn(config_dict=None):
         ("stable-diffusion-3", _create_default_map_fn("StableDiffusion3ModularPipeline")),
         ("wan", _wan_map_fn),
         ("wan-i2v", _wan_i2v_map_fn),
+        ("wan-vace", _create_default_map_fn("Wan22VaceModularPipeline")),
         ("flux", _create_default_map_fn("FluxModularPipeline")),
         ("flux-kontext", _create_default_map_fn("FluxKontextModularPipeline")),
         ("flux2", _create_default_map_fn("Flux2ModularPipeline")),

src/diffusers/modular_pipelines/wan/__init__.py

@@ -25,9 +25,11 @@
     _import_structure["modular_blocks_wan22"] = ["Wan22Blocks"]
     _import_structure["modular_blocks_wan22_i2v"] = ["Wan22Image2VideoBlocks"]
     _import_structure["modular_blocks_wan_i2v"] = ["WanImage2VideoAutoBlocks"]
+    _import_structure["modular_blocks_wan_vace"] = ["Wan22VaceBlocks"]
     _import_structure["modular_pipeline"] = [
         "Wan22Image2VideoModularPipeline",
         "Wan22ModularPipeline",
+        "Wan22VaceModularPipeline",
         "WanImage2VideoModularPipeline",
         "WanModularPipeline",
     ]
@@ -43,9 +45,11 @@
         from .modular_blocks_wan22 import Wan22Blocks
         from .modular_blocks_wan22_i2v import Wan22Image2VideoBlocks
         from .modular_blocks_wan_i2v import WanImage2VideoAutoBlocks
+        from .modular_blocks_wan_vace import Wan22VaceBlocks
         from .modular_pipeline import (
             Wan22Image2VideoModularPipeline,
             Wan22ModularPipeline,
+            Wan22VaceModularPipeline,
             WanImage2VideoModularPipeline,
             WanModularPipeline,
         )

src/diffusers/modular_pipelines/wan/before_denoise.py

@@ -404,6 +404,51 @@ def __call__(self, components: WanModularPipeline, state: PipelineState) -> Pipe
         return components, state
 
 
+class WanVaceAdditionalInputsStep(ModularPipelineBlocks):
+    model_name = "wan-vace"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Input processing step that extends `num_frames` with the reference image frames so that the initial "
+            "noise latents match the frame dimension of the vace conditioning latents.\n\n"
+            "This block should be placed after the encoder steps and the text input step."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(name="num_videos_per_prompt", default=1),
+            InputParam(name="batch_size", required=True),
+            InputParam(name="num_frames", type_hint=int),
+            InputParam(
+                name="num_reference_images",
+                type_hint=int,
+                default=0,
+                description="Number of reference images prepended on the frame dimension of the conditioning latents. Can be generated in vace_encoder step.",
+            ),
+        ]
+
+    @staticmethod
+    def check_inputs(block_state):
+        if block_state.batch_size != 1:
+            raise ValueError("Passing a list of prompts is not yet supported. This may be supported in the future.")
+        if block_state.num_videos_per_prompt != 1:
+            raise ValueError(
+                "Generating multiple videos per prompt is not yet supported. This may be supported in the future."
+            )
+
+    def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        self.check_inputs(block_state)
+
+        num_frames = block_state.num_frames or components.default_num_frames
+        block_state.num_frames = num_frames + block_state.num_reference_images * components.vae_scale_factor_temporal
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
 class WanSetTimestepsStep(ModularPipelineBlocks):
     model_name = "wan"
 

src/diffusers/modular_pipelines/wan/decoders.py

@@ -29,6 +29,40 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
+class WanVaceTrimReferenceLatentsStep(ModularPipelineBlocks):
+    model_name = "wan-vace"
+
+    @property
+    def description(self) -> str:
+        return "Step that removes the prepended reference image frames from the denoised latents before decoding"
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                "latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The denoised latents from the denoising step",
+            ),
+            InputParam(
+                "num_reference_images",
+                type_hint=int,
+                default=0,
+                description="Number of reference image frames to remove from the front of the latents. Can be generated in vace_encoder step.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        block_state.latents = block_state.latents[:, :, block_state.num_reference_images :]
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
 class WanVaeDecoderStep(ModularPipelineBlocks):
     model_name = "wan"
 

src/diffusers/modular_pipelines/wan/denoise.py

@@ -18,7 +18,7 @@
 
 from ...configuration_utils import FrozenDict
 from ...guiders import ClassifierFreeGuidance
-from ...models import WanTransformer3DModel
+from ...models import WanTransformer3DModel, WanVACETransformer3DModel
 from ...schedulers import UniPCMultistepScheduler
 from ...utils import logging
 from ..modular_pipeline import (
@@ -355,6 +355,156 @@ def __call__(
         return components, block_state
 
 
+class Wan22VaceLoopDenoiser(ModularPipelineBlocks):
+    model_name = "wan-vace"
+
+    def __init__(
+        self,
+        guider_input_fields: dict[str, Any] = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")},
+    ):
+        """Initialize a denoiser block that calls the denoiser model. This block is used in Wan2.2 VACE.
+
+        Args:
+            guider_input_fields: A dictionary that maps each argument expected by the denoiser model
+                (for example, "encoder_hidden_states") to data stored on `block_state`. The value can be either:
+
+                - A tuple of strings. For instance, `{"encoder_hidden_states": ("prompt_embeds",
+                  "negative_prompt_embeds")}` tells the guider to read `block_state.prompt_embeds` and
+                  `block_state.negative_prompt_embeds` and pass them as the conditional and unconditional batches of
+                  `encoder_hidden_states`.
+                - A string. For example, `{"encoder_hidden_image": "image_embeds"}` makes the guider forward
+                  `block_state.image_embeds` for both conditional and unconditional batches.
+        """
+        if not isinstance(guider_input_fields, dict):
+            raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}")
+        self._guider_input_fields = guider_input_fields
+        super().__init__()
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec(
+                "guider",
+                ClassifierFreeGuidance,
+                config=FrozenDict({"guidance_scale": 4.0}),
+                default_creation_method="from_config",
+            ),
+            ComponentSpec(
+                "guider_2",
+                ClassifierFreeGuidance,
+                config=FrozenDict({"guidance_scale": 3.0}),
+                default_creation_method="from_config",
+            ),
+            ComponentSpec("transformer", WanVACETransformer3DModel),
+            ComponentSpec("transformer_2", WanVACETransformer3DModel),
+        ]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Step within the denoising loop that denoise the latents with guidance. "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "object (e.g. `WanDenoiseLoopWrapper`)"
+        )
+
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [
+            ConfigSpec(
+                name="boundary_ratio",
+                default=0.875,
+                description="The boundary ratio to divide the denoising loop into high noise and low noise stages.",
+            ),
+        ]
+
+    @property
+    def inputs(self) -> list[tuple[str, Any]]:
+        inputs = [
+            InputParam("attention_kwargs"),
+            InputParam(
+                "num_inference_steps",
+                required=True,
+                type_hint=int,
+                description="The number of inference steps to use for the denoising process. Can be generated in set_timesteps step.",
+            ),
+            InputParam(
+                "vace_conditioning_latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The conditioning latents fed into the VACE control branch of the transformer. Can be generated in vace_encoder step.",
+            ),
+            InputParam(
+                "conditioning_scale",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The per-layer conditioning scale tensor applied to the VACE control branch. Can be generated in vace_encoder step.",
+            ),
+        ]
+        guider_input_names = []
+        for value in self._guider_input_fields.values():
+            if isinstance(value, tuple):
+                guider_input_names.extend(value)
+            else:
+                guider_input_names.append(value)
+
+        for name in guider_input_names:
+            inputs.append(InputParam(name=name, required=True, type_hint=torch.Tensor))
+        return inputs
+
+    @torch.no_grad()
+    def __call__(
+        self, components: WanModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+    ) -> PipelineState:
+        boundary_timestep = components.config.boundary_ratio * components.num_train_timesteps
+        if t >= boundary_timestep:
+            block_state.current_model = components.transformer
+            block_state.guider = components.guider
+        else:
+            block_state.current_model = components.transformer_2
+            block_state.guider = components.guider_2
+
+        block_state.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
+
+        # The guider splits model inputs into separate batches for conditional/unconditional predictions.
+        # For CFG with guider_inputs = {"encoder_hidden_states": (prompt_embeds, negative_prompt_embeds)}:
+        # you will get a guider_state with two batches:
+        #   guider_state = [
+        #       {"encoder_hidden_states": prompt_embeds, "__guidance_identifier__": "pred_cond"},      # conditional batch
+        #       {"encoder_hidden_states": negative_prompt_embeds, "__guidance_identifier__": "pred_uncond"},  # unconditional batch
+        #   ]
+        # Other guidance methods may return 1 batch (no guidance) or 3+ batches (e.g., PAG, APG).
+        guider_state = block_state.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields)
+
+        # run the denoiser for each guidance batch
+        for guider_state_batch in guider_state:
+            block_state.guider.prepare_models(block_state.current_model)
+            cond_kwargs = guider_state_batch.as_dict()
+            cond_kwargs = {
+                k: v.to(block_state.dtype) if isinstance(v, torch.Tensor) else v
+                for k, v in cond_kwargs.items()
+                if k in self._guider_input_fields.keys()
+            }
+
+            # Predict the noise residual
+            # store the noise_pred in guider_state_batch so that we can apply guidance across all batches
+            # the vace conditioning latents and scale are shared across the conditional/unconditional batches
+            guider_state_batch.noise_pred = block_state.current_model(
+                hidden_states=block_state.latent_model_input.to(block_state.dtype),
+                timestep=t.expand(block_state.latent_model_input.shape[0]).to(block_state.dtype),
+                control_hidden_states=block_state.vace_conditioning_latents.to(block_state.dtype),
+                control_hidden_states_scale=block_state.conditioning_scale.to(block_state.dtype),
+                attention_kwargs=block_state.attention_kwargs,
+                return_dict=False,
+                **cond_kwargs,
+            )[0]
+            block_state.guider.cleanup_models(block_state.current_model)
+
+        # Perform guidance
+        block_state.noise_pred = block_state.guider(guider_state)[0]
+
+        return components, block_state
+
+
 class WanLoopAfterDenoiser(ModularPipelineBlocks):
     model_name = "wan"
 
@@ -519,6 +669,32 @@ def description(self) -> str:
         )
 
 
+class Wan22VaceDenoiseStep(WanDenoiseLoopWrapper):
+    model_name = "wan-vace"
+    block_classes = [
+        WanLoopBeforeDenoiser,
+        Wan22VaceLoopDenoiser(
+            guider_input_fields={
+                "encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds"),
+            }
+        ),
+        WanLoopAfterDenoiser,
+    ]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoise step that iteratively denoise the latents. \n"
+            "Its loop logic is defined in `WanDenoiseLoopWrapper.__call__` method \n"
+            "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+            " - `WanLoopBeforeDenoiser`\n"
+            " - `Wan22VaceLoopDenoiser`\n"
+            " - `WanLoopAfterDenoiser`\n"
+            "This block supports controllable video generation tasks for Wan2.2 VACE."
+        )
+
+
 class Wan22Image2VideoDenoiseStep(WanDenoiseLoopWrapper):
     block_classes = [
         WanImage2VideoLoopBeforeDenoiser,