remove central registry based on review

Author: a-r-r-o-w

src/diffusers/hooks/first_block_cache.py

@@ -17,7 +17,6 @@
 
 import torch
 
-from ..models.metadata import TransformerBlockRegistry
 from ..utils import get_logger
 from ..utils.torch_utils import unwrap_module
 from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS
@@ -72,7 +71,13 @@ def __init__(self, state_manager: StateManager, threshold: float):
         self._metadata = None
 
     def initialize_hook(self, module):
-        self._metadata = TransformerBlockRegistry.get(unwrap_module(module).__class__)
+        unwrapped_module = unwrap_module(module)
+        if not hasattr(unwrapped_module, "_diffusers_transformer_block_metadata"):
+            raise ValueError(
+                f"Module {unwrapped_module} does not have any registered metadata. "
+                "Make sure to register the metadata using `diffusers.models.metadata.register_transformer_block`."
+            )
+        self._metadata = unwrapped_module._diffusers_transformer_block_metadata
         return module
 
     def new_forward(self, module: torch.nn.Module, *args, **kwargs):
@@ -150,7 +155,13 @@ def __init__(self, state_manager: StateManager, is_tail: bool = False):
         self._metadata = None
 
     def initialize_hook(self, module):
-        self._metadata = TransformerBlockRegistry.get(unwrap_module(module).__class__)
+        unwrapped_module = unwrap_module(module)
+        if not hasattr(unwrapped_module, "_diffusers_transformer_block_metadata"):
+            raise ValueError(
+                f"Module {unwrapped_module} does not have any registered metadata. "
+                "Make sure to register the metadata using `diffusers.models.metadata.register_transformer_block`."
+            )
+        self._metadata = unwrapped_module._diffusers_transformer_block_metadata
         return module
 
     def new_forward(self, module: torch.nn.Module, *args, **kwargs):

src/diffusers/models/attention.py

@@ -22,7 +22,7 @@
 from .activations import GEGLU, GELU, ApproximateGELU, FP32SiLU, LinearActivation, SwiGLU
 from .attention_processor import Attention, JointAttnProcessor2_0
 from .embeddings import SinusoidalPositionalEmbedding
-from .metadata import TransformerBlockMetadata, TransformerBlockRegistry
+from .metadata import TransformerBlockMetadata, register_transformer_block
 from .normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm, SD35AdaLayerNormZeroX
 
 
@@ -259,7 +259,7 @@ def forward(
 
 
 @maybe_allow_in_graph
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=None,

src/diffusers/models/metadata.py

@@ -44,20 +44,10 @@ def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None)
         return args[index]
 
 
-class TransformerBlockRegistry:
-    _registry = {}
-
-    @classmethod
-    def register(cls, metadata: TransformerBlockMetadata):
-        def inner(model_class: Type):
-            metadata._cls = model_class
-            cls._registry[model_class] = metadata
-            return model_class
-
-        return inner
-
-    @classmethod
-    def get(cls, model_class: Type) -> TransformerBlockMetadata:
-        if model_class not in cls._registry:
-            raise ValueError(f"Model class {model_class} not registered.")
-        return cls._registry[model_class]
+def register_transformer_block(metadata: TransformerBlockMetadata):
+    def inner(model_class: Type):
+        metadata._cls = model_class
+        model_class._diffusers_transformer_block_metadata = metadata
+        return model_class
+
+    return inner

src/diffusers/models/transformers/cogvideox_transformer_3d.py

@@ -26,7 +26,7 @@
 from ..attention_processor import AttentionProcessor, CogVideoXAttnProcessor2_0, FusedCogVideoXAttnProcessor2_0
 from ..cache_utils import CacheMixin
 from ..embeddings import CogVideoXPatchEmbed, TimestepEmbedding, Timesteps
-from ..metadata import TransformerBlockMetadata, TransformerBlockRegistry
+from ..metadata import TransformerBlockMetadata, register_transformer_block
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNorm, CogVideoXLayerNormZero
@@ -36,7 +36,7 @@
 
 
 @maybe_allow_in_graph
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=1,

src/diffusers/models/transformers/transformer_cogview4.py

@@ -26,7 +26,7 @@
 from ..attention_processor import Attention
 from ..cache_utils import CacheMixin
 from ..embeddings import CogView3CombinedTimestepSizeEmbeddings
-from ..metadata import TransformerBlockMetadata, TransformerBlockRegistry
+from ..metadata import TransformerBlockMetadata, register_transformer_block
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormContinuous
@@ -456,7 +456,7 @@ def __call__(
 
 
 @maybe_allow_in_graph
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=1,

src/diffusers/models/transformers/transformer_flux.py

@@ -34,7 +34,7 @@
 )
 from ..cache_utils import CacheMixin
 from ..embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
-from ..metadata import TransformerBlockMetadata, TransformerBlockRegistry
+from ..metadata import TransformerBlockMetadata, register_transformer_block
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
@@ -44,7 +44,7 @@
 
 
 @maybe_allow_in_graph
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=1,
         return_encoder_hidden_states_index=0,
@@ -116,7 +116,7 @@ def forward(
 
 
 @maybe_allow_in_graph
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=1,
         return_encoder_hidden_states_index=0,

src/diffusers/models/transformers/transformer_hunyuan_video.py

@@ -33,7 +33,7 @@
     Timesteps,
     get_1d_rotary_pos_embed,
 )
-from ..metadata import TransformerBlockMetadata, TransformerBlockRegistry
+from ..metadata import TransformerBlockMetadata, register_transformer_block
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, FP32LayerNorm
@@ -311,7 +311,7 @@ def forward(
         return conditioning, token_replace_emb
 
 
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=None,
@@ -496,7 +496,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return freqs_cos, freqs_sin
 
 
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=1,
@@ -578,7 +578,7 @@ def forward(
         return hidden_states, encoder_hidden_states
 
 
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=1,
@@ -663,7 +663,7 @@ def forward(
         return hidden_states, encoder_hidden_states
 
 
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=1,
@@ -749,7 +749,7 @@ def forward(
         return hidden_states, encoder_hidden_states
 
 
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=1,

src/diffusers/models/transformers/transformer_ltx.py

@@ -28,7 +28,7 @@
 from ..attention_processor import Attention
 from ..cache_utils import CacheMixin
 from ..embeddings import PixArtAlphaTextProjection
-from ..metadata import TransformerBlockMetadata, TransformerBlockRegistry
+from ..metadata import TransformerBlockMetadata, register_transformer_block
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormSingle, RMSNorm
@@ -197,7 +197,7 @@ def forward(
 
 
 @maybe_allow_in_graph
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=None,

src/diffusers/models/transformers/transformer_mochi.py

@@ -27,7 +27,7 @@
 from ..attention_processor import MochiAttention, MochiAttnProcessor2_0
 from ..cache_utils import CacheMixin
 from ..embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed
-from ..metadata import TransformerBlockMetadata, TransformerBlockRegistry
+from ..metadata import TransformerBlockMetadata, register_transformer_block
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import AdaLayerNormContinuous, RMSNorm
@@ -117,7 +117,7 @@ def forward(
 
 
 @maybe_allow_in_graph
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     metadata=TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=1,

src/diffusers/models/transformers/transformer_wan.py

@@ -27,7 +27,7 @@
 from ..attention_processor import Attention
 from ..cache_utils import CacheMixin
 from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed
-from ..metadata import TransformerBlockMetadata, TransformerBlockRegistry
+from ..metadata import TransformerBlockMetadata, register_transformer_block
 from ..modeling_outputs import Transformer2DModelOutput
 from ..modeling_utils import ModelMixin
 from ..normalization import FP32LayerNorm
@@ -222,7 +222,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
 
 @maybe_allow_in_graph
-@TransformerBlockRegistry.register(
+@register_transformer_block(
     TransformerBlockMetadata(
         return_hidden_states_index=0,
         return_encoder_hidden_states_index=None,