[WIP] Minimax h3 follow up (review & refactor)

docs/source/en/modular_diffusers/modular_pipeline.md

@@ -234,6 +234,12 @@ pipeline = ModularPipeline.from_pretrained(
 )
 ```
 
+When the pipeline blocks define workflows (check `pipeline.blocks.available_workflows`), pass `workflow=` to keep only that workflow's blocks — the same pruning as [`~ModularPipelineBlocks.get_workflow`]. The pipeline then only declares the components that workflow uses, and its docstring describes exactly that workflow's inputs.
+
+```py
+pipeline = ModularPipeline.from_pretrained("Qwen/Qwen-Image", workflow="inpainting")
+```
+
 ## Loading components
 
 A [`ModularPipeline`] doesn't automatically instantiate with components. It only loads the configuration and component specifications. You can load components with [`~ModularPipeline.load_components`].
@@ -252,6 +258,12 @@ You can also load specific components by name. The example below only loads the
 pipeline.load_components(names=["text_encoder"], dtype=torch.float16)
 ```
 
+On a pipeline whose blocks define workflows, `workflow=` loads only the components that workflow uses. The pipeline keeps all its blocks, so this is the convenient way to run one pipeline across workflows: each call adds just what the new workflow still misses.
+
+```py
+pipeline.load_components(workflow="inpainting", dtype=torch.float16)
+```
+
 After loading, printing the pipeline shows which components are loaded — the first two fields change from `null` to the component's library and class.
 
 ```py

src/diffusers/__init__.py

@@ -537,8 +537,6 @@
             "LTXModularPipeline",
             "MiniMaxH3Blocks",
             "MiniMaxH3ModularPipeline",
-            "MiniMaxH3Ref2VABlocks",
-            "MiniMaxH3Ref2VAModularPipeline",
             "QwenImageAutoBlocks",
             "QwenImageEditAutoBlocks",
             "QwenImageEditModularPipeline",
@@ -1365,8 +1363,6 @@
             LTXModularPipeline,
             MiniMaxH3Blocks,
             MiniMaxH3ModularPipeline,
-            MiniMaxH3Ref2VABlocks,
-            MiniMaxH3Ref2VAModularPipeline,
             QwenImageAutoBlocks,
             QwenImageEditAutoBlocks,
             QwenImageEditModularPipeline,

src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py

@@ -707,14 +707,8 @@ def _stitch_tiles(
             result_rows.append(torch.cat(result_row, dim=-1))
         return torch.cat(result_rows, dim=-2)
 
-    @apply_forward_hook
     def _encode_clip(self, x: torch.Tensor) -> torch.Tensor:
-        r"""
-        Encode one temporal clip, spatially tiled when tiling is enabled.
-
-        MiniMax-H3 encodes a keyframe or an image reference through this method rather than through [`~encode`],
-        because a single frame must not go through the temporal chunking, so it carries the offload hook too.
-        """
+        r"""Encode one temporal clip, spatially tiled when tiling is enabled."""
         if not self.use_tiling:
             return self.quant_conv(self.encoder(x))
 
@@ -768,17 +762,19 @@ def _decode_clip(self, z: torch.Tensor) -> torch.Tensor:
 
         return self._stitch_tiles(rows, y_overlaps, x_overlaps)
 
-    @apply_forward_hook
     def _encode(self, x: torch.Tensor) -> torch.Tensor:
         r"""
         Encode a video in `clip_length`-frame chunks and drop the `token_drop` trailing latent frames.
 
-        MiniMax-H3 encodes a video reference through this method rather than through [`~encode`], because the
-        posterior is sampled under a fixed generator rather than through the distribution object, so it carries the
-        offload hook too.
+        A single frame has no temporal extent to chunk, so it goes through the spatial encoder alone. Padding it up to
+        `clip_length` by repetition instead would run the temporal path over `clip_length` copies of the same image and
+        return `clip_length // temporal_compression_ratio - token_drop` latent frames rather than one — which is not
+        the conditioning MiniMax-H3 was trained with.
         """
         clip_length = self.config.clip_length
         num_frames = x.shape[2]
+        if num_frames == 1:
+            return self._encode_clip(x)
         if num_frames % clip_length != 0:
             pad_frames = x[:, :, -1:].repeat(1, 1, (-num_frames) % clip_length, 1, 1)
             x = torch.cat([x, pad_frames], dim=2)

src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py

@@ -93,10 +93,6 @@ class MiniMaxH3AudioEncoderOutput(BaseOutput):
     latent_dist: MiniMaxH3AudioDiagonalGaussianDistribution
 
 
-def _wn_conv1d(*args, **kwargs) -> nn.Module:
-    return weight_norm(nn.Conv1d(*args, **kwargs))
-
-
 def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
     r"""Kaiser-windowed sinc low-pass filter of shape `[1, 1, kernel_size]`.
 
@@ -234,9 +230,9 @@ def __init__(self, dim: int, dilation: int):
         super().__init__()
         self.block = nn.Sequential(
             MiniMaxH3AudioSnake1d(dim),
-            _wn_conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=((7 - 1) * dilation) // 2),
+            weight_norm(nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=((7 - 1) * dilation) // 2)),
             MiniMaxH3AudioSnake1d(dim),
-            _wn_conv1d(dim, dim, kernel_size=1),
+            weight_norm(nn.Conv1d(dim, dim, kernel_size=1)),
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -257,12 +253,14 @@ def __init__(self, dim: int, stride: int):
             MiniMaxH3AudioResidualUnit(dim // 2, dilation=3),
             MiniMaxH3AudioResidualUnit(dim // 2, dilation=9),
             MiniMaxH3AudioSnake1d(dim // 2),
-            _wn_conv1d(
-                dim // 2,
-                dim,
-                kernel_size=2 * stride,
-                stride=stride,
-                padding=math.ceil(stride / 2),
+            weight_norm(
+                nn.Conv1d(
+                    dim // 2,
+                    dim,
+                    kernel_size=2 * stride,
+                    stride=stride,
+                    padding=math.ceil(stride / 2),
+                )
             ),
         )
 
@@ -275,13 +273,13 @@ class MiniMaxH3AudioEncoder(nn.Module):
 
     def __init__(self, d_model: int, strides: tuple[int, ...], d_latent: int):
         super().__init__()
-        block: list[nn.Module] = [_wn_conv1d(1, d_model, kernel_size=7, padding=3)]
+        block: list[nn.Module] = [weight_norm(nn.Conv1d(1, d_model, kernel_size=7, padding=3))]
         for stride in strides:
             d_model *= 2
             block.append(MiniMaxH3AudioEncoderBlock(d_model, stride=stride))
         block += [
             MiniMaxH3AudioSnake1d(d_model),
-            _wn_conv1d(d_model, d_latent, kernel_size=3, padding=1),
+            weight_norm(nn.Conv1d(d_model, d_latent, kernel_size=3, padding=1)),
         ]
         self.block = nn.Sequential(*block)
 
@@ -403,12 +401,15 @@ def __init__(self, channels: int, kernel_size: int, dilation: tuple[int, ...]):
         super().__init__()
         self.convs1 = nn.ModuleList(
             [
-                _wn_conv1d(channels, channels, kernel_size, dilation=d, padding=(kernel_size * d - d) // 2)
+                weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=d, padding=(kernel_size * d - d) // 2))
                 for d in dilation
             ]
         )
         self.convs2 = nn.ModuleList(
-            [_wn_conv1d(channels, channels, kernel_size, dilation=1, padding=(kernel_size - 1) // 2) for _ in dilation]
+            [
+                weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=(kernel_size - 1) // 2))
+                for _ in dilation
+            ]
         )
         self.activations = nn.ModuleList(
             [
@@ -442,7 +443,7 @@ def __init__(
         self.num_kernels = len(resblock_kernel_sizes)
         self.num_upsamples = len(upsample_rates)
 
-        self.conv_pre = _wn_conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
+        self.conv_pre = weight_norm(nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3))
 
         # Each upsampler is wrapped in a one-element `ModuleList` in the original checkpoint
         # (`ups.<i>.0`); the extra nesting is kept so the state dict stays a passthrough.
@@ -471,7 +472,7 @@ def __init__(
                 self.resblocks.append(MiniMaxH3AudioAMPBlock(channels, kernel, tuple(dilation)))
 
         self.activation_post = MiniMaxH3AudioActivation1d(activation=MiniMaxH3AudioSnakeBeta(channels))
-        self.conv_post = _wn_conv1d(channels, 1, 7, 1, padding=3, bias=False)
+        self.conv_post = weight_norm(nn.Conv1d(channels, 1, 7, 1, padding=3, bias=False))
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = self.conv_pre(hidden_states)

src/diffusers/modular_pipelines/__init__.py

@@ -124,8 +124,6 @@
     _import_structure["minimax_h3"] = [
         "MiniMaxH3Blocks",
         "MiniMaxH3ModularPipeline",
-        "MiniMaxH3Ref2VABlocks",
-        "MiniMaxH3Ref2VAModularPipeline",
     ]
     _import_structure["z_image"] = [
         "ZImageAutoBlocks",
@@ -184,8 +182,6 @@
         from .minimax_h3 import (
             MiniMaxH3Blocks,
             MiniMaxH3ModularPipeline,
-            MiniMaxH3Ref2VABlocks,
-            MiniMaxH3Ref2VAModularPipeline,
         )
         from .modular_pipeline import (
             AutoPipelineBlocks,

src/diffusers/modular_pipelines/minimax_h3/__init__.py

@@ -21,9 +21,14 @@
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
-    _import_structure["modular_blocks_minimax_h3"] = ["MiniMaxH3Blocks", "MiniMaxH3Ref2VABlocks"]
-    _import_structure["modular_pipeline"] = ["MiniMaxH3ModularPipeline", "MiniMaxH3Ref2VAModularPipeline"]
-    _import_structure["packing_ref2va"] = ["MiniMaxH3Reference"]
+    _import_structure["modular_blocks_minimax_h3"] = ["MiniMaxH3Blocks"]
+    _import_structure["modular_pipeline"] = ["MiniMaxH3ModularPipeline"]
+    _import_structure["references"] = [
+        "MiniMaxH3AudioReference",
+        "MiniMaxH3ImageReference",
+        "MiniMaxH3Reference",
+        "MiniMaxH3VideoReference",
+    ]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
@@ -32,9 +37,14 @@
     except OptionalDependencyNotAvailable:
         from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
     else:
-        from .modular_blocks_minimax_h3 import MiniMaxH3Blocks, MiniMaxH3Ref2VABlocks
-        from .modular_pipeline import MiniMaxH3ModularPipeline, MiniMaxH3Ref2VAModularPipeline
-        from .packing_ref2va import MiniMaxH3Reference
+        from .modular_blocks_minimax_h3 import MiniMaxH3Blocks
+        from .modular_pipeline import MiniMaxH3ModularPipeline
+        from .references import (
+            MiniMaxH3AudioReference,
+            MiniMaxH3ImageReference,
+            MiniMaxH3Reference,
+            MiniMaxH3VideoReference,
+        )
 else:
     import sys