[WIP] Refactor Model Design by DN6 · Pull Request #13794 · huggingface/diffusers

DN6 · 2026-05-22T19:59:03Z

What does this PR do?

This refactor turns models into self-contained modules that declare their capabilities in one place. Per-model conversion code moves next to the model and a unified metadata() API makes feature attributes inspectable from any model class.

Motivation

Today, features are added to models through a mix of class attributes and mixins. Mixins define their own class attributes as well, so when examining a model class it isn't immediately clear which attributes and features are relevant or available.

Models are defined in a single file, so we end up using centralized utility files for things like model-specific weight and LoRA conversions. These files have grown enormous as they accumulate code to handle per-model variants and their idiosyncrasies.

The new design makes the mixins model-agnostic and has each mixin reach for the per-model metadata it needs through small handler objects attached to the model class.

Proposed Structure

Using Flux as a reference:

models/transformers/flux/
├── __init__.py
├── _ip_adapter.py        # FluxIPAdapterMixin + converters (internal)
├── _lora.py              # FLUX_LORA handler + per-format converters (internal)
├── _weight_mapping.py    # FLUX_WEIGHT_MAPPING handler + key tables (internal)
└── model.py              # FluxTransformer2DModel class declaration

Two patterns live next to model.py, picked per subsystem based on whether the behavior actually generalizes across models:

handler + shared mixin — for features where the steps are the same across models and only the data/conversion function varies. The model opts in by inheriting the shared mixin in loaders/ and assigning its handler as a class attribute. LoRA and single-file weight mapping fit here:
```
class FluxTransformer2DModel(ModelMixin, LoRAModelMixin, ...):
    _lora = LoRAHandler("...")               # handler instance, consumed by LoRAModelMixin
```
Per-model mixin — for features that vary too much across models for a single shared mixin to be useful. Each model gets its own mixin declared right next to the model and inherited directly. IP-Adapter is the showcase:
```
class FluxTransformer2DModel(..., FluxIPAdapterMixin):
    ...
```

This should simplify developing on top of these models — modifications or enhancements stay within one folder. If a model has a very specific feature that doesn't generalize across others, it can be kept isolated there too (e.g. FreeNoise for the AnimateDiff UNet). Additionally, if a custom model is modifying an existing diffusers model (Self-Forcing Wan), the folder method of organizing the model lends itself well to custom code loading with AutoModel.

Features Introduced

Model capability introspection via `Model.metadata()`

Each model exposes a metadata() classmethod that returns a metadata object, keyed by the class attribute that controls each feature. The displayed row tells you exactly what to set or inherit to change the behavior.

>>> print(FluxTransformer2DModel.metadata())
FluxTransformer2DModel feature attributes
──────────────────────────────────────────────────────────────────────────────────
  _supports_gradient_checkpointing  True
  _supports_group_offloading        True
  _no_split_modules                 FluxTransformerBlock, FluxSingleTransformerBlock
  _skip_layerwise_casting_patterns  pos_embed, norm
  _repeated_blocks                  FluxTransformerBlock, FluxSingleTransformerBlock
  _cp_plan                          True
  _weight_mapping                   flux-depth, flux-dev, flux-fill, flux-schnell
  _lora                             bfl, kohya, kontext, xlabs
  _supports_cache                   True
  _supports_ip_adapter              True

The returned ModelMetadata exposes each feature value as an attribute (meta._supports_ip_adapter, meta._lora, ...), supports keys() / values() / items() for mapping-style iteration, and in for presence checks. meta.describe(verbose=True) adds an indented description and docs link under each row. Which can be useful for agents.

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

yiyixuxu

Thanks so much for working on this!
I really like the direction this is taking, very exciting!
It's a big PR with a lot of changes, so I'll go through it step by step - left some initial questions:)

yiyixuxu · 2026-05-29T02:30:22Z

    _default_processor_cls = None
    _available_processors = []
    _supports_qkv_fusion = True
+    _parallel_config = None


why do we need it here?

follow up on this
any reason we need it here? it should live on the attention processor, no?

Sorry, this proposal should also probably be included in a follow up.

IMO it is better to set attention backend/ CP at the attention module level, rather than at the processor level like we have currently. e.g. If you enable CP and then switch out the processor, you would need to re-enable CP again with a new processor. This feels a bit weird.

model.enable_parallelism(config) # writes to current processor model.set_processor(NewProcessor()) # NewProcessor._parallel_config defaults back to None # CP is silently disabled. The next forward pass runs unsharded.

But we can discuss in a follow up PR. I'll revert for now.

@DN6 sorry not to digress but we change between attention backends through set_attention_backend() or with the context manager:

diffusers/src/diffusers/models/attention_dispatch.py

Line 373 in 21ba394

def attention_backend(backend: str | AttentionBackendName = AttentionBackendName.NATIVE):

So, assuming we're doing:

model.enable_parallelism(config) model.set_attention_backend(new_backend)

where new_backend is different from the one used during enable_parallelism(), we actually throw an error (if the backend doesn't support CP), see:

diffusers/src/diffusers/models/modeling_utils.py

Lines 628 to 637 in 21ba394

if parallel_config_set and not _AttentionBackendRegistry._is_context_parallel_available(backend):

compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel)

raise ValueError(

f"Context parallelism is enabled but current attention backend '{backend.value}' "

f"does not support context parallelism. "

f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()`."

)

_check_attention_backend_requirements(backend)

_maybe_download_kernel_for_backend(backend)

Even the other opposite ordering (set_attention_backend() and enable_parallelism()) checks for compatibility:

diffusers/src/diffusers/models/modeling_utils.py

Lines 1660 to 1665 in 21ba394

raise ValueError(

f"Context parallelism is enabled but the attention processor '{processor.__class__.__name__}' "

f"is using backend '{attention_backend.value}' which does not support context parallelism. "

f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()` before "

f"calling `model.enable_parallelism()`."

)

IMO, it's better to separate parallelism from attention. On the surface, it may seem like our current parallelism support (CP) is largely operating at the attention level; that may not hold up for long, especially once we ship support for TP.

My points are based on the assumption that we want to encourage users to use set_attention_backend() rather than the attention processor method that you cited above.

@sayakpaul I'm not referring to using processors to set attention backends.

Parallelism and attention backends are currently coupled to processors because the _parallel_config and _attention_backend attributes are defined and set on processors.

Processors are swappable and enable different types of attention operations (independent of backend) e.g. for Flux we have FluxIPAdapterAttnProcessor and Flux2 has Flux2KVAttnProcessor . If you were to enable parallelism on Flux and then load an IP Adapter, you would have to enable parallelism again for it to work.

pipeline.transformer.enable_parallelism() pipeline.load_ip_adapter() # parallelism is now disabled.

The _attention_backend attribute is also set at the processor level. We have a fallback that uses the backend set in the global AttentionRegistry if the processor attribute is set to None . But this also a bug that needs fixing and can lead to weird situations such as.

model_A.set_attention_backend("sage_hub") # e.g. set a new processor to experiment with KV cacheing model_A.set_processors(...) model_B.set_attention_backend("flash_hub") # the attention registry global backend is now flash_hub # this runs with the flash_hub backend because the _attention_backend attribute on the processor is set to None when a new processor is set. # It will attempt to use the global backend which is based on the last model that set the backend. model_A(**kwargs)

Moving these attributes into the attention module addresses this issue. But again, I will remove this part from the scope of this PR and introduce a dedicated PR for it. We can discuss where parallelism lives in that PR as well.

Wow that is a good point. model_B and model_A shouldn't have this kind of interference!

yiyixuxu · 2026-05-29T04:22:49Z

+                f"{DOCS_BASE}/optimization/memory#gradient-checkpointing",
+            )
+        if cls._supports_group_offloading:
+            rows["_supports_group_offloading"] = (


I really like the metadata + mixin direction but I'd like to understand a bit more: is there a fundamental reason why some features get their own self-contained mixin (Cache, Lora) and many others stay as methods on ModelMixin?

yiyixuxu · 2026-05-29T05:47:36Z

+        return "\n".join(lines)
+
+
+def register_metadata(metadata):


Is there a reason we have to keep two systems to attach metadata to a class? the register_metadata vs the class attribute like _cp_plan?

I think that is because for handling with stuff like _cp_plan we don't have a mixin like CPMixin. However for others, we have dedicated mixins.

But do we need this registration, though? All the available feature set can be queried through the main model class, no?

yiyixuxu · 2026-05-29T05:48:59Z

+
+
+@dataclass
+class TransformerBlockOutput(TransformerModuleOutput):


these are not used yet no?

sayakpaul · 2026-05-29T09:06:01Z

    _cached_parameter_indices: dict[str, int] = None

+    def _register(self, cls):
+        """Attach this metadata to ``cls`` and register it in :class:`TransformerBlockRegistry`.


Single "" is what we use across diffusers. Also, :class:` isn't something we do.

Suggested change

"""Attach this metadata to ``cls`` and register it in :class:`TransformerBlockRegistry`.

"""Attach this metadata to `cls` and register it in `TransformerBlockRegistry`.

sayakpaul

Did a pass on LoRA. Will now do a pass on modeling_utils.py.

I think it might be better to do the LoRA-related separation in another PR because it's difficult to truly assess if it's in a working state.

We are likely missing andling Flux Control LoRA and as such, handling text-encoder LoRA modules. It's also not clear how this would affect lora_pipeline.py and lora_conversion_utils.py.
One-off utilities that are not shared across multiple different functions / classes.
It's not clear how the existing https://github.com/huggingface/diffusers/blob/main/src/diffusers/loaders/peft.py gets affected by this.

sayakpaul · 2026-05-29T09:06:37Z

    _cls: Type = None
    _cached_parameter_indices: dict[str, int] = None

+    def _register(self, cls):


I am a little unclear about the docstring. How is this method metadata (which is what the docstring reads)?

sayakpaul · 2026-05-29T09:07:13Z

+    def _register(self, cls):
+        """Attach this metadata to ``cls`` and register it in :class:`TransformerBlockRegistry`.
+
+        Lets ``@register_metadata(TransformerBlockMetadata(...))`` work for block classes that opt into the decorator


sayakpaul · 2026-05-29T09:08:51Z

+        self._cls = cls
+        cls._block_metadata = self
+        TransformerBlockRegistry._registry[cls] = self


Is it only applicable to classes with _repeated_blocks set?

sayakpaul · 2026-05-29T09:28:28Z



 @maybe_allow_in_graph
+@register_metadata(TransformerBlockMetadata(return_hidden_states_index=1, return_encoder_hidden_states_index=0))


What is the purpose of this?

sayakpaul · 2026-05-29T09:31:01Z

+            "_supports_cache": (
+                True,
+                "True",
+                "Supports caching techniques (PAB / FasterCache / FirstBlockCache) via `enable_cache`.",


Suggested change

"Supports caching techniques (PAB / FasterCache / FirstBlockCache) via `enable_cache`.",

"Supports caching techniques (e.g. FasterCache) via `enable_cache`.",

sayakpaul · 2026-05-29T10:00:54Z

+        r"""
+        Add an adapter to the underlying model.
+
+        ``source`` can be either:


There is nothing called source here.

sayakpaul · 2026-05-29T10:02:32Z

+
+        _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name)
+
+    def _inject_adapter(self, state_dict, lora_config, adapter_name, peft_kwargs):


Will prefer it inline as it's not shared.

sayakpaul · 2026-05-29T10:02:42Z

+            self._rollback_adapter(adapter_name, e)
+            raise
+
+    def _maybe_apply_deferred_hotswap_prep(self, lora_config):


sayakpaul · 2026-05-29T10:02:50Z

+        prepare_model_for_compiled_hotswap(self, config=lora_config, **self._lora_hotswap_kwargs)
+        self._lora_hotswap_kwargs = None
+
+    def _hotswap_adapter(self, state_dict, lora_config, adapter_name):


sayakpaul · 2026-05-29T10:02:56Z

+            self._rollback_adapter(adapter_name, e)
+            raise
+
+    def _rollback_adapter(self, adapter_name, error):


sayakpaul · 2026-05-29T10:07:55Z

+
+
+@dataclass
+class AttnProcessorOutput(TransformerModuleOutput):


Why does AttnProcessorOutput has to live in src/diffusers/models/transformers/utils.py? Could it not be used by VAEs or other components under src/diffusers/models/?

sayakpaul · 2026-05-29T10:12:55Z

+class ModelMetadata:
+    """Snapshot of a model class's feature attributes.
+
+    Constructed by :meth:`ModelMixin.metadata` — walks ``cls.__mro__`` collecting rows from each mixin's ``_metadata``


:meth: can be dangerous abbreviation 🤪

sayakpaul · 2026-05-29T10:14:24Z

+        return "\n".join(lines)
+
+
+def register_metadata(metadata):


I think that is because for handling with stuff like _cp_plan we don't have a mixin like CPMixin. However for others, we have dedicated mixins.

sayakpaul · 2026-05-29T10:15:54Z

+        return "\n".join(lines)
+
+
+def register_metadata(metadata):


But do we need this registration, though? All the available feature set can be queried through the main model class, no?

DN6 · 2026-06-01T05:24:57Z

@yiyixuxu @sayakpaul Realised that the scope of the PR is probably too wide and is leading to some difficulty with the reviews. I've updated the PR to reduce scope to only the essentials. I will open PRs for LoRA and TransformerRegistry changes in follow ups.

@yiyixu Yes, ideally we would want anything that's not Model Loading or Hub related to live in a dedicated Mixin. The current challenge with the refactor is not introducing backward breaking changes as we migrate models. e.g If we were to add a CPMixin for enable_parallelism and remove the method from ModelMixin (practically speaking this change should be okay, but it is technically backwards breaking). Which is why some methods still remain on ModelMixin.

For from_single_file, I would like to actually move this into core ModelMixin because apart from the auto config fetching and conversion, it should reuse all the same logic as from_pretrained. IMO it makes sense to have it live in ModelMixin. For now I've introduced a new WeightMappingMixin as a temporary measure because adding the new from_single_file method to ModelMixin would break single file loading for anything that's not Flux.

yiyixuxu

thanks @DN6
I left some questions on the weight mapper

about this

ideally we would want anything that's not Model Loading or Hub related to live in a dedicated Mixin. ...

got it! the dedicated-mixin direction sounds good to me, but I think 1.0 is a good place to do this though (to deprecate the methods on ModelMixin and move them properly). How about we merge this PR and then do this togther in a follow-up? (@dg845 and I can help as well)

yiyixuxu · 2026-06-24T02:54:19Z

+
+
+# Assigned to ``FluxTransformer2DModel`` as the ``_weight_mapping`` class attribute in ``flux/model.py``.
+FLUX_WEIGHT_MAPPING = WeightMappingHandler(


umm the WeightMappingHandler already has detect_confnig/map_to_diffusers etc, as methods, why can't we just subclass it here and override these diretly, instead of passing them in as callables? feel like it'd give clearer signature to document and we could even enforce them

Yeah this works for me. You mean something like this?

@dataclass class FluxWeightMappingHandler(WeightMappingHandler): def detect_config(self, state_dict): ... def map_to_diffusers(self, state_dict, **kwargs): ... FLUX_WEIGHT_MAPPING = FluxWeightMappingHandler(available_configs=..., default_config="flux-dev", ...).

yiyixuxu · 2026-06-24T02:57:51Z

    _default_processor_cls = None
    _available_processors = []
    _supports_qkv_fusion = True
+    _parallel_config = None


follow up on this
any reason we need it here? it should live on the attention processor, no?

…ightMappingMixin

HuggingFaceDocBuilderDev · 2026-07-01T17:40:06Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

yiyixuxu

thanks, i left some comments
feel free to address as you see fit & merge

yiyixuxu · 2026-07-02T03:04:59Z

+
+    # ---- conversion ----
+
+    def map_to_diffusers(self, state_dict: dict, **kwargs) -> dict:


I think it is better if we raise a NotImplementedError here too to enforce this method

yiyixuxu · 2026-07-02T03:11:02Z

+            if source in key:
+                tensors = forward_fn(value, **ctx)
+                for target, tensor in zip(targets, tensors):
+                    new_key = WeightMappingHandler.rename_key(key.replace(source, target), rename_patterns)


ohh maybe just make rename_key a regular function then if you need to use it this way

yiyixuxu · 2026-07-02T03:15:24Z

+        Default implementation returns ``None`` so the caller falls back to ``_default_config``. Subclasses override to
+        inspect the state_dict for architecture-distinguishing keys or shapes.
+        """
+        return None


Should we throw a warning if there are multiple configs and this is not implemented?

# Conflicts: # src/diffusers/hooks/_helpers.py # src/diffusers/models/modeling_utils.py # src/diffusers/models/transformers/flux/model.py # src/diffusers/models/transformers/transformer_bria_fibo.py

# Conflicts: # src/diffusers/hooks/_helpers.py # src/diffusers/loaders/single_file_model.py

github-actions Bot added size/L PR with diff > 200 LOC models utils single-file hooks labels May 22, 2026

DN6 marked this pull request as ready for review May 27, 2026 16:35

DN6 requested review from asomoza, dg845, sayakpaul and yiyixuxu May 27, 2026 16:35

yiyixuxu reviewed May 29, 2026

View reviewed changes

sayakpaul reviewed May 29, 2026

View reviewed changes

github-actions Bot added the pipelines label Jun 1, 2026

yiyixuxu reviewed Jun 24, 2026

View reviewed changes

github-actions Bot added documentation Improvements or additions to documentation lora quantization tests modular-pipelines examples CI guiders schedulers experimental loaders ip-adapter labels Jul 1, 2026

DN6 marked this pull request as draft July 1, 2026 17:11

Refactor model design: metadata system + Flux port to package with We…

4c90ebc

…ightMappingMixin

DN6 force-pushed the refactor-model-metadata branch from e368850 to 4c90ebc Compare July 1, 2026 17:26

github-actions Bot removed documentation Improvements or additions to documentation lora quantization tests modular-pipelines pipelines examples CI guiders schedulers experimental loaders ip-adapter labels Jul 1, 2026

DN6 added 2 commits July 1, 2026 22:59

clean up

c90e57f

clean up

14d0722

DN6 added 2 commits July 1, 2026 23:50

clean up

1953e6b

fix

698e7f3

github-actions Bot added the pipelines label Jul 1, 2026

DN6 added 2 commits July 1, 2026 23:59

clean up

f514d7a

clean up

4763076

DN6 marked this pull request as ready for review July 1, 2026 18:54

yiyixuxu approved these changes Jul 2, 2026

View reviewed changes

DN6 added 2 commits October 29, 2026 19:32

Merge remote-tracking branch 'origin/main' into refactor-model-metadata

7f58967

# Conflicts: # src/diffusers/hooks/_helpers.py # src/diffusers/models/modeling_utils.py # src/diffusers/models/transformers/flux/model.py # src/diffusers/models/transformers/transformer_bria_fibo.py

Merge remote-tracking branch 'origin/main' into refactor-model-metadata

7c88feb

# Conflicts: # src/diffusers/hooks/_helpers.py # src/diffusers/loaders/single_file_model.py

	if parallel_config_set and not _AttentionBackendRegistry._is_context_parallel_available(backend):
	compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel)
	raise ValueError(
	f"Context parallelism is enabled but current attention backend '{backend.value}' "
	f"does not support context parallelism. "
	f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()`."
	)

	_check_attention_backend_requirements(backend)
	_maybe_download_kernel_for_backend(backend)

	raise ValueError(
	f"Context parallelism is enabled but the attention processor '{processor.__class__.__name__}' "
	f"is using backend '{attention_backend.value}' which does not support context parallelism. "
	f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()` before "
	f"calling `model.enable_parallelism()`."
	)



		@dataclass
		class TransformerBlockOutput(TransformerModuleOutput):

	"""Attach this metadata to ``cls`` and register it in :class:`TransformerBlockRegistry`.
	"""Attach this metadata to `cls` and register it in `TransformerBlockRegistry`.



		@maybe_allow_in_graph
		@register_metadata(TransformerBlockMetadata(return_hidden_states_index=1, return_encoder_hidden_states_index=0))

	"Supports caching techniques (PAB / FasterCache / FirstBlockCache) via `enable_cache`.",
	"Supports caching techniques (e.g. FasterCache) via `enable_cache`.",


		_maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name)

		def _inject_adapter(self, state_dict, lora_config, adapter_name, peft_kwargs):



		@dataclass
		class AttnProcessorOutput(TransformerModuleOutput):



		# Assigned to ``FluxTransformer2DModel`` as the ``_weight_mapping`` class attribute in ``flux/model.py``.
		FLUX_WEIGHT_MAPPING = WeightMappingHandler(


		# ---- conversion ----

		def map_to_diffusers(self, state_dict: dict, **kwargs) -> dict:

Conversation

DN6 commented May 22, 2026

What does this PR do?

Motivation

Proposed Structure

Features Introduced

Model capability introspection via Model.metadata()

Before submitting

Who can review?

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DN6 commented Jun 1, 2026

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Model capability introspection via `Model.metadata()`

sayakpaul Jun 25, 2026 •

edited

Loading