_always_upcast_modules -> _precision_sensitive_module_patterns

Author: a-r-r-o-w

src/diffusers/models/autoencoders/autoencoder_asym_kl.py

@@ -60,7 +60,7 @@ class AsymmetricAutoencoderKL(ModelMixin, ConfigMixin):
             Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper.
     """
 
-    _always_upcast_modules = ["decoder"]
+    _precision_sensitive_module_patterns = ["decoder"]
 
     @register_to_config
     def __init__(

src/diffusers/models/autoencoders/vq_model.py

@@ -71,7 +71,7 @@ class VQModel(ModelMixin, ConfigMixin):
             Type of normalization layer to use. Can be one of `"group"` or `"spatial"`.
     """
 
-    _always_upcast_modules = ["quantize"]
+    _precision_sensitive_module_patterns = ["quantize"]
 
     @register_to_config
     def __init__(

src/diffusers/models/modeling_utils.py

@@ -163,7 +163,7 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
     _keys_to_ignore_on_load_unexpected = None
     _no_split_modules = None
     _keep_in_fp32_modules = None
-    _always_upcast_modules = None
+    _precision_sensitive_module_patterns = None
 
     def __init__(self):
         super().__init__()
@@ -344,10 +344,11 @@ def enable_layerwise_upcasting(
         memory footprint from model weights, but may lead to some quality degradation in the outputs. Most degradations
         are negligible, mostly stemming from weight casting in normalization and modulation layers.
 
-        By default, most models in diffusers set the `_always_upcast_modules` attribute to ignore patch embedding,
-        positional embedding and normalization layers. This is because these layers are most likely precision-critical
-        for quality. If you wish to change this behavior, you can set the `_always_upcast_modules` attribute to `None`,
-        or call [`~hooks.layerwise_upcasting.apply_layerwise_upcasting`] with custom arguments.
+        By default, most models in diffusers set the `_precision_sensitive_module_patterns` attribute to ignore patch
+        embedding, positional embedding and normalization layers. This is because these layers are most likely
+        precision-critical for quality. If you wish to change this behavior, you can set the
+        `_precision_sensitive_module_patterns` attribute to `None`, or call
+        [`~hooks.layerwise_upcasting.apply_layerwise_upcasting`] with custom arguments.
 
         Example:
             Using [`~models.ModelMixin.enable_layerwise_upcasting`]:
@@ -386,8 +387,8 @@ def enable_layerwise_upcasting(
             user_provided_patterns = False
         if self._keep_in_fp32_modules is not None:
             skip_modules_pattern += tuple(self._keep_in_fp32_modules)
-        if self._always_upcast_modules is not None:
-            skip_modules_pattern += tuple(self._always_upcast_modules)
+        if self._precision_sensitive_module_patterns is not None:
+            skip_modules_pattern += tuple(self._precision_sensitive_module_patterns)
         skip_modules_pattern = tuple(set(skip_modules_pattern))
 
         if is_peft_available() and not user_provided_patterns:

src/diffusers/models/transformers/auraflow_transformer_2d.py

@@ -276,7 +276,7 @@ class AuraFlowTransformer2DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin
     """
 
     _no_split_modules = ["AuraFlowJointTransformerBlock", "AuraFlowSingleTransformerBlock", "AuraFlowPatchEmbed"]
-    _always_upcast_modules = ["pos_embed", "norm"]
+    _precision_sensitive_module_patterns = ["pos_embed", "norm"]
     _supports_gradient_checkpointing = True
 
     @register_to_config

src/diffusers/models/transformers/cogvideox_transformer_3d.py

@@ -212,7 +212,7 @@ class CogVideoXTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
             Scaling factor to apply in 3D positional embeddings across temporal dimensions.
     """
 
-    _always_upcast_modules = ["patch_embed", "norm"]
+    _precision_sensitive_module_patterns = ["patch_embed", "norm"]
     _supports_gradient_checkpointing = True
     _no_split_modules = ["CogVideoXBlock", "CogVideoXPatchEmbed"]
 

src/diffusers/models/transformers/dit_transformer_2d.py

@@ -64,7 +64,7 @@ class DiTTransformer2DModel(ModelMixin, ConfigMixin):
             A small constant added to the denominator in normalization layers to prevent division by zero.
     """
 
-    _always_upcast_modules = ["pos_embed", "norm"]
+    _precision_sensitive_module_patterns = ["pos_embed", "norm"]
     _supports_gradient_checkpointing = True
 
     @register_to_config

src/diffusers/models/transformers/hunyuan_transformer_2d.py

@@ -244,7 +244,7 @@ class HunyuanDiT2DModel(ModelMixin, ConfigMixin):
             Whether or not to use style condition and image meta size. True for version <=1.1, False for version >= 1.2
     """
 
-    _always_upcast_modules = ["pos_embed", "norm", "pooler"]
+    _precision_sensitive_module_patterns = ["pos_embed", "norm", "pooler"]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/latte_transformer_3d.py

@@ -65,7 +65,7 @@ class LatteTransformer3DModel(ModelMixin, ConfigMixin):
             The number of frames in the video-like data.
     """
 
-    _always_upcast_modules = ["pos_embed", "norm"]
+    _precision_sensitive_module_patterns = ["pos_embed", "norm"]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/lumina_nextdit2d.py

@@ -221,7 +221,7 @@ class LuminaNextDiT2DModel(ModelMixin, ConfigMixin):
             overall scale of the model's operations.
     """
 
-    _always_upcast_modules = ["patch_embedder", "norm", "ffn_norm"]
+    _precision_sensitive_module_patterns = ["patch_embedder", "norm", "ffn_norm"]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/pixart_transformer_2d.py

@@ -79,7 +79,7 @@ class PixArtTransformer2DModel(ModelMixin, ConfigMixin):
 
     _supports_gradient_checkpointing = True
     _no_split_modules = ["BasicTransformerBlock", "PatchEmbed"]
-    _always_upcast_modules = ["pos_embed", "norm", "adaln_single"]
+    _precision_sensitive_module_patterns = ["pos_embed", "norm", "adaln_single"]
 
     @register_to_config
     def __init__(