From c059bab3e05a885ac2499b4901593f62ccc6f8d2 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Sun, 9 Aug 2026 10:37:07 -0700 Subject: [PATCH] add modular blockset for wan 2.2 vace --- src/diffusers/__init__.py | 4 + src/diffusers/modular_pipelines/__init__.py | 4 + .../modular_pipelines/modular_pipeline.py | 1 + .../modular_pipelines/wan/__init__.py | 4 + .../modular_pipelines/wan/before_denoise.py | 45 +++ .../modular_pipelines/wan/decoders.py | 34 ++ .../modular_pipelines/wan/denoise.py | 178 +++++++++- .../modular_pipelines/wan/encoders.py | 310 +++++++++++++++++- .../wan/modular_blocks_wan_vace.py | 208 ++++++++++++ .../modular_pipelines/wan/modular_pipeline.py | 17 + .../dummy_torch_and_transformers_objects.py | 30 ++ .../wan/test_modular_pipeline_wan_vace.py | 74 +++++ 12 files changed, 907 insertions(+), 2 deletions(-) create mode 100644 src/diffusers/modular_pipelines/wan/modular_blocks_wan_vace.py create mode 100644 tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 2deff4b71cb2..ed82277c2908 100644 --- a/src/diffusers/__init__.py +++ b/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, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 008a654c3fa3..d858997e4566 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/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, diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 9bf1ddca3b98..344a0e9e0ad8 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/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")), diff --git a/src/diffusers/modular_pipelines/wan/__init__.py b/src/diffusers/modular_pipelines/wan/__init__.py index 284b6c9fa436..b7ce7cde2dee 100644 --- a/src/diffusers/modular_pipelines/wan/__init__.py +++ b/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, ) diff --git a/src/diffusers/modular_pipelines/wan/before_denoise.py b/src/diffusers/modular_pipelines/wan/before_denoise.py index 6b0874037b0d..509838f31716 100644 --- a/src/diffusers/modular_pipelines/wan/before_denoise.py +++ b/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" diff --git a/src/diffusers/modular_pipelines/wan/decoders.py b/src/diffusers/modular_pipelines/wan/decoders.py index bc26df234039..529c9291c250 100644 --- a/src/diffusers/modular_pipelines/wan/decoders.py +++ b/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" diff --git a/src/diffusers/modular_pipelines/wan/denoise.py b/src/diffusers/modular_pipelines/wan/denoise.py index 63093c37f7ec..240745a1ec8e 100644 --- a/src/diffusers/modular_pipelines/wan/denoise.py +++ b/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, diff --git a/src/diffusers/modular_pipelines/wan/encoders.py b/src/diffusers/modular_pipelines/wan/encoders.py index c7d5df48e7be..519dd384becf 100644 --- a/src/diffusers/modular_pipelines/wan/encoders.py +++ b/src/diffusers/modular_pipelines/wan/encoders.py @@ -23,7 +23,7 @@ from ...configuration_utils import FrozenDict from ...guiders import ClassifierFreeGuidance from ...image_processor import PipelineImageInput -from ...models import AutoencoderKLWan +from ...models import AutoencoderKLWan, WanVACETransformer3DModel from ...utils import is_ftfy_available, is_torchvision_available, logging from ...video_processor import VideoProcessor from ..modular_pipeline import ModularPipelineBlocks, PipelineState @@ -563,6 +563,314 @@ def __call__(self, components: WanModularPipeline, state: PipelineState) -> Pipe return components, state +class WanVaceEncoderStep(ModularPipelineBlocks): + model_name = "wan-vace" + + @property + def description(self) -> str: + return ( + "Vace Encoder step that preprocesses the control video, mask and reference images and encodes them " + "into the conditioning latents used by the VACE control branch of the transformer" + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", WanVACETransformer3DModel), + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "video", + type_hint=list[PIL.Image.Image], + description="The control video to condition the generation on. If not provided, an empty video is used.", + ), + InputParam( + "mask", + type_hint=list[PIL.Image.Image], + description="The mask that defines which video regions to condition on (black) and which to generate (white). Can only be passed if `video` is passed as well.", + ), + InputParam( + "reference_images", + type_hint=PIL.Image.Image | list[PIL.Image.Image], + description="One or more reference images as extra conditioning for the generation.", + ), + InputParam( + "conditioning_scale", + type_hint=float | list[float] | torch.Tensor, + default=1.0, + description="The conditioning scale applied in each control layer of the model. If a float, it is applied uniformly to all layers; a list or tensor must have the same length as the number of control layers.", + ), + InputParam("height"), + InputParam("width"), + InputParam("num_frames", type_hint=int, default=81), + InputParam("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "vace_conditioning_latents", + type_hint=torch.Tensor, + description="The concatenated video and mask conditioning latents fed into the VACE control branch of the transformer", + ), + OutputParam( + "num_reference_images", + type_hint=int, + description="Number of reference images prepended on the frame dimension of the conditioning latents", + ), + ] + + @staticmethod + def check_inputs(components, block_state): + base = components.vae_scale_factor_spatial * components.patch_size_spatial + if (block_state.height is not None and block_state.height % base != 0) or ( + block_state.width is not None and block_state.width % base != 0 + ): + raise ValueError( + f"`height` and `width` have to be divisible by {base} but are {block_state.height} and {block_state.width}." + ) + if block_state.num_frames is not None and ( + block_state.num_frames < 1 or (block_state.num_frames - 1) % components.vae_scale_factor_temporal != 0 + ): + raise ValueError( + f"`num_frames` has to be greater than 0, and (num_frames - 1) must be divisible by {components.vae_scale_factor_temporal}, but got {block_state.num_frames}." + ) + if isinstance(block_state.generator, list): + raise ValueError("Passing a list of generators is not yet supported. This may be supported in the future.") + if block_state.video is not None: + if block_state.mask is not None and len(block_state.video) != len(block_state.mask): + raise ValueError( + f"Length of `video` {len(block_state.video)} and `mask` {len(block_state.mask)} do not match. Please make sure that" + " they have the same length." + ) + elif block_state.mask is not None: + raise ValueError("`mask` can only be passed if `video` is passed as well.") + + @staticmethod + def preprocess_conditions( + components, + video, + mask, + reference_images, + height, + width, + num_frames, + dtype, + device, + ): + if video is not None: + base = components.vae_scale_factor_spatial * components.patch_size_spatial + video_height, video_width = components.video_processor.get_default_height_width(video[0]) + + if video_height * video_width > height * width: + scale = min(width / video_width, height / video_height) + video_height, video_width = int(video_height * scale), int(video_width * scale) + + if video_height % base != 0 or video_width % base != 0: + logger.warning( + f"Video height and width should be divisible by {base}, but got {video_height} and {video_width}. " + ) + video_height = (video_height // base) * base + video_width = (video_width // base) * base + + video = components.video_processor.preprocess_video(video, video_height, video_width) + image_size = (video_height, video_width) # Use the height/width of video (with possible rescaling) + else: + video = torch.zeros(1, 3, num_frames, height, width, dtype=dtype, device=device) + image_size = (height, width) # Use the height/width provider by user + + if mask is not None: + mask = components.video_processor.preprocess_video(mask, image_size[0], image_size[1]) + mask = torch.clamp((mask + 1) / 2, min=0, max=1) + else: + mask = torch.ones_like(video) + + video = video.to(dtype=dtype, device=device) + mask = mask.to(dtype=dtype, device=device) + + # Make a list of list of images where the outer list corresponds to video batch size and the inner list + # corresponds to list of conditioning images per video + if reference_images is None or isinstance(reference_images, PIL.Image.Image): + reference_images = [[reference_images] for _ in range(video.shape[0])] + elif isinstance(reference_images, (list, tuple)) and isinstance(next(iter(reference_images)), PIL.Image.Image): + reference_images = [reference_images] + elif ( + isinstance(reference_images, (list, tuple)) + and isinstance(next(iter(reference_images)), list) + and isinstance(next(iter(reference_images[0])), PIL.Image.Image) + ): + reference_images = reference_images + else: + raise ValueError( + "`reference_images` has to be of type `PIL.Image.Image` or `list` of `PIL.Image.Image`, or " + f"`list` of `list` of `PIL.Image.Image`, but is {type(reference_images)}" + ) + + if video.shape[0] != len(reference_images): + raise ValueError( + f"Batch size of `video` {video.shape[0]} and length of `reference_images` {len(reference_images)} does not match." + ) + + reference_images_preprocessed = [] + for i, reference_images_batch in enumerate(reference_images): + preprocessed_images = [] + for j, image in enumerate(reference_images_batch): + if image is None: + continue + image = components.video_processor.preprocess(image, None, None) + img_height, img_width = image.shape[-2:] + scale = min(image_size[0] / img_height, image_size[1] / img_width) + new_height, new_width = int(img_height * scale), int(img_width * scale) + resized_image = torch.nn.functional.interpolate( + image, size=(new_height, new_width), mode="bilinear", align_corners=False + ).squeeze(0) # [C, H, W] + top = (image_size[0] - new_height) // 2 + left = (image_size[1] - new_width) // 2 + canvas = torch.ones(3, *image_size, device=device, dtype=dtype) + canvas[:, top : top + new_height, left : left + new_width] = resized_image + preprocessed_images.append(canvas) + reference_images_preprocessed.append(preprocessed_images) + + return video, mask, reference_images_preprocessed + + @staticmethod + def prepare_video_latents(components, video, mask, reference_images, generator, device): + vae_dtype = components.vae.dtype + video = video.to(dtype=vae_dtype) + + latents_mean = torch.tensor(components.vae.config.latents_mean, device=device, dtype=torch.float32).view( + 1, components.vae.config.z_dim, 1, 1, 1 + ) + latents_std = 1.0 / torch.tensor(components.vae.config.latents_std, device=device, dtype=torch.float32).view( + 1, components.vae.config.z_dim, 1, 1, 1 + ) + + mask = torch.where(mask > 0.5, 1.0, 0.0).to(dtype=vae_dtype) + inactive = video * (1 - mask) + reactive = video * mask + inactive = retrieve_latents(components.vae.encode(inactive), generator, sample_mode="argmax") + reactive = retrieve_latents(components.vae.encode(reactive), generator, sample_mode="argmax") + inactive = ((inactive.float() - latents_mean) * latents_std).to(vae_dtype) + reactive = ((reactive.float() - latents_mean) * latents_std).to(vae_dtype) + latents = torch.cat([inactive, reactive], dim=1) + + latent_list = [] + for latent, reference_images_batch in zip(latents, reference_images): + for reference_image in reference_images_batch: + reference_image = reference_image.to(dtype=vae_dtype) + reference_image = reference_image[None, :, None, :, :] # [1, C, 1, H, W] + reference_latent = retrieve_latents( + components.vae.encode(reference_image), generator, sample_mode="argmax" + ) + reference_latent = ((reference_latent.float() - latents_mean) * latents_std).to(vae_dtype) + reference_latent = reference_latent.squeeze(0) # [C, 1, H, W] + reference_latent = torch.cat([reference_latent, torch.zeros_like(reference_latent)], dim=0) + latent = torch.cat([reference_latent.squeeze(0), latent], dim=1) + latent_list.append(latent) + return torch.stack(latent_list) + + @staticmethod + def prepare_masks(components, mask, reference_images): + transformer_patch_size = components.patch_size_spatial + + mask_list = [] + for mask_, reference_images_batch in zip(mask, reference_images): + num_channels, num_frames, height, width = mask_.shape + new_num_frames = ( + num_frames + components.vae_scale_factor_temporal - 1 + ) // components.vae_scale_factor_temporal + new_height = ( + height // (components.vae_scale_factor_spatial * transformer_patch_size) * transformer_patch_size + ) + new_width = ( + width // (components.vae_scale_factor_spatial * transformer_patch_size) * transformer_patch_size + ) + mask_ = mask_[0, :, :, :] + mask_ = mask_.view( + num_frames, + new_height, + components.vae_scale_factor_spatial, + new_width, + components.vae_scale_factor_spatial, + ) + mask_ = mask_.permute(2, 4, 0, 1, 3).flatten(0, 1) # [8x8, num_frames, new_height, new_width] + mask_ = torch.nn.functional.interpolate( + mask_.unsqueeze(0), size=(new_num_frames, new_height, new_width), mode="nearest-exact" + ).squeeze(0) + num_ref_images = len(reference_images_batch) + if num_ref_images > 0: + mask_padding = torch.zeros_like(mask_[:, :num_ref_images, :, :]) + mask_ = torch.cat([mask_padding, mask_], dim=1) + mask_list.append(mask_) + return torch.stack(mask_list) + + @torch.no_grad() + def __call__(self, components: WanModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(components, block_state) + + device = components._execution_device + dtype = torch.float32 + + height = block_state.height or components.default_height + width = block_state.width or components.default_width + num_frames = block_state.num_frames or components.default_num_frames + + video, mask, reference_images = self.preprocess_conditions( + components, + block_state.video, + block_state.mask, + block_state.reference_images, + height, + width, + num_frames, + dtype, + device, + ) + if video.shape[0] != 1: + raise ValueError( + "Generating with more than one video is not yet supported. This may be supported in the future." + ) + block_state.num_reference_images = len(reference_images[0]) + + conditioning_latents = self.prepare_video_latents( + components, video, mask, reference_images, block_state.generator, device + ) + mask = self.prepare_masks(components, mask, reference_images) + block_state.vace_conditioning_latents = torch.cat([conditioning_latents, mask], dim=1) + + conditioning_scale = block_state.conditioning_scale + if isinstance(conditioning_scale, (int, float)): + conditioning_scale = [conditioning_scale] * components.num_vace_layers + if isinstance(conditioning_scale, list): + if len(conditioning_scale) != components.num_vace_layers: + raise ValueError( + f"Length of `conditioning_scale` {len(conditioning_scale)} does not match number of layers {components.num_vace_layers}." + ) + conditioning_scale = torch.tensor(conditioning_scale) + if isinstance(conditioning_scale, torch.Tensor): + if conditioning_scale.size(0) != components.num_vace_layers: + raise ValueError( + f"Length of `conditioning_scale` {conditioning_scale.size(0)} does not match number of layers {components.num_vace_layers}." + ) + conditioning_scale = conditioning_scale.to(device=device) + block_state.conditioning_scale = conditioning_scale + + self.set_block_state(state, block_state) + return components, state + + class WanPrepareFirstFrameLatentsStep(ModularPipelineBlocks): model_name = "wan" diff --git a/src/diffusers/modular_pipelines/wan/modular_blocks_wan_vace.py b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_vace.py new file mode 100644 index 000000000000..c4a61f0e335c --- /dev/null +++ b/src/diffusers/modular_pipelines/wan/modular_blocks_wan_vace.py @@ -0,0 +1,208 @@ +# 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. + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + WanPrepareLatentsStep, + WanSetTimestepsStep, + WanTextInputStep, + WanVaceAdditionalInputsStep, +) +from .decoders import WanVaceTrimReferenceLatentsStep, WanVaeDecoderStep +from .denoise import ( + Wan22VaceDenoiseStep, +) +from .encoders import ( + WanTextEncoderStep, + WanVaceEncoderStep, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# ==================== +# 1. DENOISE +# ==================== + +# inputs (text + vace_conditioning_latents) -> additional_inputs -> set_timesteps -> prepare_latents -> denoise + + +# auto_docstring +class Wan22VaceCoreDenoiseStep(SequentialPipelineBlocks): + """ + denoise block that takes encoded text and vace conditioning latents and runs the denoising process. + + Components: + transformer (`WanTransformer3DModel`) scheduler (`UniPCMultistepScheduler`) guider (`ClassifierFreeGuidance`) + guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanVACETransformer3DModel`) + + Configs: + boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low + noise stages. + + Inputs: + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + prompt_embeds (`Tensor`): + Pre-generated text embeddings. Can be generated from text_encoder step. + negative_prompt_embeds (`Tensor`, *optional*): + Pre-generated negative text embeddings. Can be generated from text_encoder step. + num_frames (`int`, *optional*): + TODO: Add description. + num_reference_images (`int`, *optional*, defaults to 0): + Number of reference images prepended on the frame dimension of the conditioning latents. Can be generated + in vace_encoder step. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + height (`int`, *optional*): + TODO: Add description. + width (`int`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + vace_conditioning_latents (`Tensor`): + The conditioning latents fed into the VACE control branch of the transformer. Can be generated in + vace_encoder step. + conditioning_scale (`Tensor`): + The per-layer conditioning scale tensor applied to the VACE control branch. Can be generated in + vace_encoder step. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "wan-vace" + block_classes = [ + WanTextInputStep, + WanVaceAdditionalInputsStep, + WanSetTimestepsStep, + WanPrepareLatentsStep, + Wan22VaceDenoiseStep, + ] + block_names = [ + "input", + "additional_inputs", + "set_timesteps", + "prepare_latents", + "denoise", + ] + + @property + def description(self): + return "denoise block that takes encoded text and vace conditioning latents and runs the denoising process." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# ==================== +# 2. BLOCKS (Wan2.2 VACE) +# ==================== + + +# auto_docstring +class Wan22VaceBlocks(SequentialPipelineBlocks): + """ + Modular pipeline for controllable video generation using Wan2.2 VACE. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) guider (`ClassifierFreeGuidance`) transformer + (`WanVACETransformer3DModel`) vae (`AutoencoderKLWan`) video_processor (`VideoProcessor`) scheduler + (`UniPCMultistepScheduler`) guider_2 (`ClassifierFreeGuidance`) transformer_2 (`WanVACETransformer3DModel`) + + Configs: + boundary_ratio (default: 0.875): The boundary ratio to divide the denoising loop into high noise and low + noise stages. + + Inputs: + prompt (`None`, *optional*): + TODO: Add description. + negative_prompt (`None`, *optional*): + TODO: Add description. + max_sequence_length (`None`, *optional*, defaults to 512): + TODO: Add description. + video (`list`, *optional*): + The control video to condition the generation on. If not provided, an empty video is used. + mask (`list`, *optional*): + The mask that defines which video regions to condition on (black) and which to generate (white). Can only + be passed if `video` is passed as well. + reference_images (`Image | list`, *optional*): + One or more reference images as extra conditioning for the generation. + conditioning_scale (`float | list | Tensor`, *optional*, defaults to 1.0): + The conditioning scale applied in each control layer of the model. If a float, it is applied uniformly to + all layers; a list or tensor must have the same length as the number of control layers. + height (`None`, *optional*): + TODO: Add description. + width (`None`, *optional*): + TODO: Add description. + num_frames (`int`, *optional*, defaults to 81): + TODO: Add description. + generator (`None`, *optional*): + TODO: Add description. + num_videos_per_prompt (`None`, *optional*, defaults to 1): + TODO: Add description. + num_inference_steps (`None`, *optional*, defaults to 50): + TODO: Add description. + timesteps (`None`, *optional*): + TODO: Add description. + sigmas (`None`, *optional*): + TODO: Add description. + latents (`Tensor | NoneType`, *optional*): + TODO: Add description. + attention_kwargs (`None`, *optional*): + TODO: Add description. + output_type (`str`, *optional*, defaults to np): + The output type of the decoded videos + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "wan-vace" + block_classes = [ + WanTextEncoderStep, + WanVaceEncoderStep, + Wan22VaceCoreDenoiseStep, + WanVaceTrimReferenceLatentsStep, + WanVaeDecoderStep, + ] + block_names = [ + "text_encoder", + "vace_encoder", + "denoise", + "trim_latents", + "decode", + ] + + @property + def description(self): + return "Modular pipeline for controllable video generation using Wan2.2 VACE." + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/wan/modular_pipeline.py b/src/diffusers/modular_pipelines/wan/modular_pipeline.py index a360440c9251..94727d70d3e6 100644 --- a/src/diffusers/modular_pipelines/wan/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/wan/modular_pipeline.py @@ -139,3 +139,20 @@ class Wan22Image2VideoModularPipeline(Wan22ModularPipeline): """ default_blocks_name = "Wan22Image2VideoBlocks" + + +class Wan22VaceModularPipeline(Wan22ModularPipeline): + """ + A ModularPipeline for Wan2.2 VACE. + + > [!WARNING] > This is an experimental feature and is likely to change in the future. + """ + + default_blocks_name = "Wan22VaceBlocks" + + @property + def num_vace_layers(self): + num_vace_layers = 8 + if hasattr(self, "transformer") and self.transformer is not None: + num_vace_layers = len(self.transformer.config.vace_layers) + return num_vace_layers diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 8058937aa2a8..c956ecaf1095 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -782,6 +782,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Wan22VaceBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class Wan22VaceModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class WanBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py b/tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py new file mode 100644 index 000000000000..f4b00f7e0812 --- /dev/null +++ b/tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py @@ -0,0 +1,74 @@ +# coding=utf-8 +# 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 PIL import Image + +from diffusers.modular_pipelines import Wan22VaceBlocks, Wan22VaceModularPipeline + +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +class TestWan22VaceModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = Wan22VaceModularPipeline + pipeline_blocks_class = Wan22VaceBlocks + pretrained_model_name_or_path = "akshan-main/tiny-wan22-vace-modular-pipe" + + params = frozenset(["prompt", "height", "width", "num_frames", "video", "mask", "reference_images"]) + batch_params = frozenset() + optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) + output_name = "videos" + + def get_dummy_inputs(self, seed=0): + generator = self.get_generator(seed) + num_frames = 9 + video = [Image.new("RGB", (16, 16))] * num_frames + mask = [Image.new("L", (16, 16), 0)] * num_frames + inputs = { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "video": video, + "mask": mask, + "num_inference_steps": 2, + "height": 16, + "width": 16, + "num_frames": num_frames, + "max_sequence_length": 16, + "output_type": "pt", + } + return inputs + + def test_inference_with_reference_image(self): + pipe = self.get_pipeline().to("cpu") + + inputs = self.get_dummy_inputs() + inputs["reference_images"] = Image.new("RGB", (16, 16)) + videos = pipe(**inputs, output=self.output_name) + assert videos.shape == (1, 9, 3, 16, 16) + assert torch.isnan(videos).sum() == 0 + + @pytest.mark.skip(reason="Batching is not yet supported with this pipeline") + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip(reason="Batching is not yet supported with this pipeline") + def test_inference_batch_single_identical(self): + pass + + @pytest.mark.skip(reason="num_videos_per_prompt") + def test_num_images_per_prompt(self): + pass