feat(diffusers/models): update models to v0.29

mindspore-lab · Sep 2, 2024 · 46a4ed6 · 46a4ed6
1 parent b14ef56
commit 46a4ed6
Show file tree

Hide file tree

Showing 46 changed files with 5,671 additions and 833 deletions.
diff --git a/mindone/diffusers/models/__init__.py b/mindone/diffusers/models/__init__.py
@@ -23,15 +23,22 @@
     "autoencoders.autoencoder_kl_temporal_decoder": ["AutoencoderKLTemporalDecoder"],
     "autoencoders.autoencoder_tiny": ["AutoencoderTiny"],
     "autoencoders.consistency_decoder_vae": ["ConsistencyDecoderVAE"],
+    "autoencoders.vq_model": ["VQModel"],
     "controlnet": ["ControlNetModel"],
+    "controlnet_sd3": ["SD3ControlNetModel", "SD3MultiControlNetModel"],
+    "controlnet_xs": ["ControlNetXSAdapter", "UNetControlNetXSModel"],
     "dual_transformer_2d": ["DualTransformer2DModel"],
     "embeddings": ["ImageProjection"],
     "modeling_utils": ["ModelMixin"],
+    "transformers.dit_transformer_2d": ["DiTTransformer2DModel"],
+    "transformers.dual_transformer_2d": ["DualTransformer2DModel"],
+    "transformers.hunyuan_transformer_2d": ["HunyuanDiT2DModel"],
+    "transformers.pixart_transformer_2d": ["PixArtTransformer2DModel"],
     "transformers.prior_transformer": ["PriorTransformer"],
     "transformers.t5_film_transformer": ["T5FilmDecoder"],
     "transformers.transformer_2d": ["Transformer2DModel"],
-    "transformers.transformer_temporal": ["TransformerTemporalModel"],
     "transformers.transformer_sd3": ["SD3Transformer2DModel"],
+    "transformers.transformer_temporal": ["TransformerTemporalModel"],
     "unets.unet_1d": ["UNet1DModel"],
     "unets.unet_2d": ["UNet2DModel"],
     "unets.unet_2d_condition": ["UNet2DConditionModel"],
@@ -42,7 +49,6 @@
     "unets.unet_stable_cascade": ["StableCascadeUNet"],
     "unets.unet_spatio_temporal_condition": ["UNetSpatioTemporalConditionModel"],
     "unets.uvit_2d": ["UVit2DModel"],
-    "vq_model": ["VQModel"],
 }
 
 if TYPE_CHECKING:
@@ -53,12 +59,18 @@
         AutoencoderKLTemporalDecoder,
         AutoencoderTiny,
         ConsistencyDecoderVAE,
+        VQModel,
     )
     from .controlnet import ControlNetModel
+    from .controlnet_sd3 import SD3ControlNetModel, SD3MultiControlNetModel
+    from .controlnet_xs import ControlNetXSAdapter, UNetControlNetXSModel
     from .embeddings import ImageProjection
     from .modeling_utils import ModelMixin
     from .transformers import (
+        DiTTransformer2DModel,
         DualTransformer2DModel,
+        HunyuanDiT2DModel,
+        PixArtTransformer2DModel,
         PriorTransformer,
         SD3Transformer2DModel,
         T5FilmDecoder,
@@ -78,7 +90,6 @@
         UNetSpatioTemporalConditionModel,
         UVit2DModel,
     )
-    from .vq_model import VQModel
 
 else:
     import sys

diff --git a/mindone/diffusers/models/attention.py b/mindone/diffusers/models/attention.py
@@ -765,7 +765,6 @@ def __init__(
         if inner_dim is None:
             inner_dim = int(dim * mult)
         dim_out = dim_out if dim_out is not None else dim
-        linear_cls = nn.Dense
 
         if activation_fn == "gelu":
             act_fn = GELU(dim, inner_dim, bias=bias)
@@ -782,7 +781,7 @@ def __init__(
         # project dropout
         net.append(nn.Dropout(p=dropout))
         # project out
-        net.append(linear_cls(inner_dim, dim_out, has_bias=bias))
+        net.append(nn.Dense(inner_dim, dim_out, has_bias=bias))
         # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
         if final_dropout:
             net.append(nn.Dropout(p=dropout))