Merge branch 'main' into remove-test-pipelines-common

Author: sayakpaul

docs/source/en/modular_diffusers/modular_pipeline.md

@@ -165,7 +165,7 @@ ModularPipeline {
 }
 ```
 
-If you pass a repository to [`~ModularPipelineBlocks.init_pipeline`], it overrides the loading path by matching your block's components against the pipeline config in that repository (`model_index.json` or `modular_model_index.json`).
+If you pass a repository to [`~ModularPipelineBlocks.init_pipeline`], it overrides the loading path by matching your block's components against the pipeline config in that repository (`model_index.json` or `modular_model_index.json`). See [Modular repository](#modular-repository) for how loading specs are recorded and saved.
 
 In the example below, the `pretrained_model_name_or_path` will be updated to `"stabilityai/stable-diffusion-xl-base-1.0"`.
 
@@ -415,6 +415,34 @@ pipeline = ModularPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base
 pipeline.save_pretrained("local/path", repo_id="my-username/sdxl-modular", push_to_hub=True)
 ```
 
+By default, [`~ModularPipeline.save_pretrained`] writes each currently loaded component that Diffusers can serialize. Components that are not loaded, or cannot be serialized, are not written and keep their existing loading specifications. This gives you two ways to save, depending on what you want.
+
+### Save a self-contained copy
+
+Load all the components, then save. Every spec points at the result, so it reloads entirely from one place, including offline.
+
+```py
+pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3")
+pipe.load_components()
+pipe.save_pretrained("path/to/local-copy")
+```
+
+The destination recorded in `modular_model_index.json` is wherever [`~ModularPipeline.save_pretrained`] writes: the `save_directory` for a local save, or the `repo_id` when you pass `push_to_hub=True`.
+
+### Keep references to existing components
+
+Load only what's new (or nothing at all). Only loaded components are saved; everything else stays a pointer to its original repository. Use this mode when you want to replace one component while continuing to load the others from their original repository. For example, save a custom transformer while the remaining components continue to load from the base repository.
+
+```py
+pipe = ModularPipeline.from_pretrained("black-forest-labs/FLUX.2-dev")
+pipe.update_components(transformer=my_custom_transformer)  # the only component in memory
+pipe.save_pretrained("local/path", repo_id="my-username/flux2-custom-transformer", push_to_hub=True)
+```
+
+Pass `overwrite_modular_index=False` to keep the loading specs in `modular_model_index.json` as they are. A saved component whose loading spec is empty is still filled in with the destination, since there is nothing to preserve.
+
+Note that moving the files any other way (uploading with `hf upload`, downloading a repository with `hf download --local-dir`) doesn't rewrite the index, so the copy still points to the old location; update the index manually in that case.
+
 A modular repository can also include custom pipeline blocks as Python code. This allows you to share specialized blocks that aren't native to Diffusers. For example, [diffusers/Florence2-image-Annotator](https://huggingface.co/diffusers/Florence2-image-Annotator) contains custom blocks alongside the loading configuration:
 
 ```

examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py

@@ -684,17 +684,22 @@ def collate_fn(examples):
     )
 
     # Scheduler and math around the number of training steps.
-    overrode_max_train_steps = False
-    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
+    num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes
     if args.max_train_steps is None:
-        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
-        overrode_max_train_steps = True
+        len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
+        num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
+        num_training_steps_for_scheduler = (
+            args.num_train_epochs * accelerator.num_processes * num_update_steps_per_epoch
+        )
+    else:
+        num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes
 
     lr_scheduler = get_scheduler(
         args.lr_scheduler,
         optimizer=optimizer,
-        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
-        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+        num_warmup_steps=num_warmup_steps_for_scheduler,
+        num_training_steps=num_training_steps_for_scheduler,
     )
     unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
         unet, optimizer, train_dataloader, lr_scheduler
@@ -704,8 +709,14 @@ def collate_fn(examples):
     vae.to(accelerator.device, dtype=weight_dtype)
     # We need to recalculate our total training steps as the size of the training dataloader may have changed.
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
-    if overrode_max_train_steps:
+    if args.max_train_steps is None:
         args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        if num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes:
+            logger.warning(
+                f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
+                f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
+                f"This inconsistency may result in the learning rate scheduler not functioning properly."
+            )
     # Afterwards we recalculate our number of training epochs
     args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 

src/diffusers/__init__.py

@@ -569,6 +569,8 @@
             "Wan22Image2VideoBlocks",
             "Wan22Image2VideoModularPipeline",
             "Wan22ModularPipeline",
+            "Wan22VaceBlocks",
+            "Wan22VaceModularPipeline",
             "WanAnimate2Blocks",
             "WanAnimate2DistilledBlocks",
             "WanAnimate2DistilledModularPipeline",
@@ -1423,6 +1425,8 @@
             Wan22Image2VideoBlocks,
             Wan22Image2VideoModularPipeline,
             Wan22ModularPipeline,
+            Wan22VaceBlocks,
+            Wan22VaceModularPipeline,
             WanAnimate2Blocks,
             WanAnimate2DistilledBlocks,
             WanAnimate2DistilledModularPipeline,

src/diffusers/hooks/hooks.py

@@ -276,6 +276,20 @@ def _set_context(self, name: str | None = None) -> None:
         for registry in self._get_child_registries():
             registry._set_context(name)
 
+    def invalidate_child_registries_cache(self) -> None:
+        """Invalidate the cached child-registry list across this module's tree.
+
+        `_get_child_registries` caches the registries it finds by walking `named_modules()`, keyed on the registry that
+        built it. Registering or removing hooks on descendant modules (e.g. block hooks added by `enable_cache`)
+        changes which modules carry a `_diffusers_hook`, which stales that cache. Call this after any operation that
+        adds or removes hooks in the subtree so the list is rebuilt on next use. Clears the cache on every registry in
+        the tree, since the same registries appear in ancestor caches.
+        """
+        for _, module in unwrap_module(self._module_ref).named_modules():
+            module = unwrap_module(module)
+            if hasattr(module, "_diffusers_hook"):
+                module._diffusers_hook._child_registries_cache = None
+
     def _get_child_registries(self) -> list["HookRegistry"]:
         """Return registries of child modules, using a cached list when available.
 

src/diffusers/models/cache_utils.py

@@ -69,6 +69,7 @@ def enable_cache(self, config) -> None:
         from ..hooks import (
             FasterCacheConfig,
             FirstBlockCacheConfig,
+            HookRegistry,
             MagCacheConfig,
             PyramidAttentionBroadcastConfig,
             TaylorSeerCacheConfig,
@@ -101,6 +102,11 @@ def enable_cache(self, config) -> None:
         else:
             raise ValueError(f"Cache config {type(config)} is not supported.")
 
+        # Applying a cache technique registers hooks on child blocks, which stales any
+        # `_child_registries_cache` built earlier (e.g. by a prior `cache_context`). Invalidate it
+        # so `_set_context` reaches the freshly-registered block hooks.
+        HookRegistry.check_if_exists_or_initialize(self).invalidate_child_registries_cache()
+
         self._cache_config = config
 
     def disable_cache(self) -> None:
@@ -144,6 +150,9 @@ def disable_cache(self) -> None:
         else:
             raise ValueError(f"Cache config {type(self._cache_config)} is not supported.")
 
+        # Removing the cache hooks stales any `_child_registries_cache` that included them.
+        registry.invalidate_child_registries_cache()
+
         self._cache_config = None
 
     def _reset_stateful_cache(self, recurse: bool = True) -> None:

src/diffusers/models/transformers/prior_transformer.py

@@ -318,5 +318,8 @@ def forward(
         return PriorTransformerOutput(predicted_image_embedding=predicted_image_embedding)
 
     def post_process_latents(self, prior_latents):
-        prior_latents = (prior_latents * self.clip_std) + self.clip_mean
+        # `clip_std` / `clip_mean` are parameters of this model, not of a submodule, so group offloading onloads
+        # them only for the duration of `forward`. This runs after the denoising loop, hence the explicit move.
+        device = prior_latents.device
+        prior_latents = (prior_latents * self.clip_std.to(device)) + self.clip_mean.to(device)
         return prior_latents

src/diffusers/modular_pipelines/__init__.py

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

src/diffusers/modular_pipelines/modular_pipeline.py

@@ -36,7 +36,7 @@
     _unwrap_model,
     simple_get_class_obj,
 )
-from ..utils import PushToHubMixin, is_accelerate_available, logging
+from ..utils import PushToHubMixin, deprecate, is_accelerate_available, logging
 from ..utils.dynamic_modules_utils import get_class_from_dynamic_module, resolve_trust_remote_code
 from ..utils.hub_utils import _resolve_revision, load_or_create_model_card, populate_model_card
 from ..utils.torch_utils import empty_device_cache, is_compiled_module
@@ -135,6 +135,7 @@ def _helios_pyramid_map_fn(config_dict=None):
         ("wan-animate-2", _create_default_map_fn("WanAnimate2ModularPipeline")),
         ("wan-animate-2-distilled", _create_default_map_fn("WanAnimate2DistilledModularPipeline")),
         ("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")),
@@ -1974,10 +1975,13 @@ def save_pretrained(
             push_to_hub (`bool`, *optional*, defaults to `False`):
                 Whether to push the pipeline to the Hugging Face model hub after saving it.
             **kwargs: Additional keyword arguments:
-                - `overwrite_modular_index` (`bool`, *optional*, defaults to `False`):
-                    When saving a Modular Pipeline, its components in `modular_model_index.json` may reference repos
-                    different from the destination repo. Setting this to `True` updates all component references in
-                    `modular_model_index.json` so they point to the repo specified by `repo_id`.
+                - `overwrite_modular_index` (`bool`, *optional*, defaults to `True`):
+                    Whether to update `modular_model_index.json` so each saved component's loading spec points to the
+                    destination: `repo_id` when pushing to the Hub, otherwise `save_directory`. Components that are not
+                    loaded are not saved and always keep their recorded loading specs. Pass `False` to also preserve
+                    the recorded specs of the components being saved (e.g. for an index that deliberately references
+                    other repositories); components without a load id (such as custom models added with
+                    `update_components`) are still rewritten since they have no recorded source.
                 - `repo_id` (`str`, *optional*):
                     The repository ID to push the pipeline to. Defaults to the last component of `save_directory`.
                 - `commit_message` (`str`, *optional*):
@@ -1989,7 +1993,17 @@ def save_pretrained(
                 - `token` (`str`, *optional*):
                     The Hugging Face token to use for authentication.
         """
-        overwrite_modular_index = kwargs.pop("overwrite_modular_index", False)
+        if "overwrite_modular_index" not in kwargs:
+            deprecate(
+                "overwrite_modular_index",
+                "0.43.0",
+                "The default of `overwrite_modular_index` in `ModularPipeline.save_pretrained` changed from `False`"
+                " to `True`: the saved `modular_model_index.json` now points each saved component at the destination"
+                " (the save directory, or `repo_id` when pushing to the Hub). Pass `overwrite_modular_index=False`"
+                " to keep the previously recorded loading specs, or pass `True` explicitly to silence this warning.",
+                standard_warn=False,
+            )
+        overwrite_modular_index = kwargs.pop("overwrite_modular_index", True)
         repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
 
         if push_to_hub:
@@ -2060,6 +2074,9 @@ def save_pretrained(
                 library, class_name, component_spec_dict = self.config[component_name]
                 component_spec_dict["pretrained_model_name_or_path"] = repo_id if push_to_hub else save_directory
                 component_spec_dict["subfolder"] = component_name
+                component_spec_dict["variant"] = variant if save_method_accept_variant else None
+                # a revision pinned for the original source doesn't exist at the destination
+                component_spec_dict["revision"] = None
                 self.register_to_config(**{component_name: (library, class_name, component_spec_dict)})
 
         self.save_config(save_directory=save_directory)

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

@@ -402,6 +402,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"