update

Author: a-r-r-o-w

src/diffusers/models/transformers/latte_transformer_3d.py

@@ -27,7 +27,6 @@
 
 class LatteTransformer3DModel(ModelMixin, ConfigMixin):
     _supports_gradient_checkpointing = True
-    _no_split_modules = ["BasicTransformerBlock"]
 
     """
     A 3D Transformer model for video-like data, paper: https://arxiv.org/abs/2401.03048, offical code:

src/diffusers/models/transformers/transformer_allegro.py

@@ -222,7 +222,6 @@ class AllegroTransformer3DModel(ModelMixin, ConfigMixin):
     """
 
     _supports_gradient_checkpointing = True
-    _no_split_modules = ["AllegroTransformerBlock"]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/transformer_ltx.py

@@ -295,7 +295,6 @@ class LTXVideoTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin
     """
 
     _supports_gradient_checkpointing = True
-    _no_split_modules = ["LTXTransformerBlock"]
 
     @register_to_config
     def __init__(

src/diffusers/models/transformers/transformer_sd3.py

@@ -125,7 +125,6 @@ class SD3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrigi
     """
 
     _supports_gradient_checkpointing = True
-    _no_split_modules = ["JointTransformerBlock", "SD3SingleTransformerBlock"]
 
     @register_to_config
     def __init__(

tests/models/transformers/test_models_transformer_cogview3plus.py

@@ -71,7 +71,7 @@ def prepare_init_args_and_inputs_for_common(self):
         init_dict = {
             "patch_size": 2,
             "in_channels": 4,
-            "num_layers": 1,
+            "num_layers": 2,
             "attention_head_dim": 4,
             "num_attention_heads": 2,
             "out_channels": 4,