huggingface / huggingface/diffusers

Modular pipelines: variant blocksets sharing a `model_name` init the wrong pipeline class

Open
#14,451 0 comments 0 reactions 1 assignee Claimed by @yiyixuxu View on GitHub
bug modular-pipelines needs-code-example needs-env-info
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

`ModularPipelineBlocks.init_pipeline()` resolves the pipeline class through `MODULAR_PIPELINE_MAPPING[self.model_name]`, calling the map fn **without a config**. The config-branching map fns (`_flux2_klein_map_fn`, `_krea2_map_fn`, `_helios_pyramid_map_fn`, `_wan_map_fn`, `_wan_i2v_map_fn`) can only pick the variant class when given a repo config (the `model_index.json` fallback in `from_pretrained`), so from the blocks side they always return their `None` default. Any variant blockset that shares its base's `model_name` inits as the wrong pipeline class:

| variant blockset | `model_name` | `init_pipeline()` returns | should be |
|---|---|---|---|
| `Flux2KleinBaseAutoBlocks` | `flux2-klein` | `Flux2KleinModularPipeline` (the distilled default) | `Flux2KleinBaseModularPipeline` |
| `Krea2TurboAutoBlocks` | `krea2` | `Krea2ModularPipeline` | `Krea2TurboModularPipeline` |
| `HeliosPyramidDistilledAutoBlocks` | `helios-pyramid` | `HeliosPyramidModularPipeline` | `HeliosPyramidDistilledModularPipeline` |
| `Wan22Blocks` | `wan` | `WanModularPipeline` | `Wan22ModularPipeline` |
| `Wan22Image2VideoBlocks` | `wan-i2v` | `WanImage2VideoModularPipeline` | `Wan22Image2VideoModularPipeline` |
| `Cosmos3DistilledBlocks` | `cosmos3-omni` | `Cosmos3OmniModularPipeline` | `Cosmos3DistilledModularPipeline` |
| `LTX25AutoBlocks` | `ltx2` | `LTX2ModularPipeline` | `LTX25ModularPipeline` |

Consequences:

- **Krea2 turbo is functionally wrong**: `Krea2TurboModularPipeline` pins `requires_unconditional_embeds = False`; the base class consults the guider, so a turbo pipeline built from blocks can be asked to encode negative embeds.
- **The error is sticky**: `save_pretrained` writes `_class_name` from the live pipeline class, so a wrongly-inited variant round-trips the wrong class through its saved repo.
- **Most variants only override `default_blocks_name`**, so the damage is usually latent — but the fast test suites for krea2-turbo, flux2-klein-base, and cosmos3-distilled all go through `init_pipeline()` and currently exercise the base class.
- `Cosmos3DistilledBlocks` has no config branch at all, so even the `model_index.json` path can never return the distilled class; the `"ltx2.5"` mapping key is dead code (no blockset declares that name).

Proposed fix — the pattern qwenimage (4 names) and wan-animate-2 (`wan-animate-2` / `wan-animate-2-distilled`, #14413) already follow: give each variant blockset its own `model_name` with its own `_create_default_map_fn` entry. It's a small change now that every variant has its own blockset file, and it also fixes the save/reload round-trip. The config-branching map fns stay for the `model_index.json` fallback in `from_pretrained` (that path derives one model name per standard pipeline class, so it still needs the config to tell variants apart).

The same PR should add a guard test to `ModularPipelineTesterMixin` so this can't regress:

```python
def test_init_pipeline_class(self):
pipe = self.pipeline_blocks_class().init_pipeline()
assert type(pipe) is self.pipeline_class
```

Every tester already declares `pipeline_class`, so this is immediately red today for the krea2-turbo, flux2-klein-base, and cosmos3-distilled suites and goes green with the fix.

---

*Posted by Claude (Claude Code) on behalf of @yiyixuxu, as a reminder from the Wan-Animate-2 integration work (#14413).*

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.