Latte: Latent Diffusion Transformer for Video Generation

.gitignore

@@ -175,4 +175,4 @@ tags
 .ruff_cache
 
 # wandb
-wandb
+wandb

docs/source/en/api/models/latte_transformer3d.md

@@ -0,0 +1,20 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+## LatteTransformer3DModel
+
+A Diffusion Transformer model for 3D data from [Latte](https://github.com/Vchitect/Latte).
+
+## LatteTransformer3DModel
+
+[[autodoc]] LatteTransformer3DModel
+

src/diffusers/__init__.py

@@ -88,6 +88,7 @@
             "HunyuanDiT2DMultiControlNetModel",
             "I2VGenXLUNet",
             "Kandinsky3UNet",
+            "LatteTransformer3DModel",
             "LuminaNextDiT2DModel",
             "ModelMixin",
             "MotionAdapter",
@@ -272,6 +273,7 @@
             "LDMTextToImagePipeline",
             "LEditsPPPipelineStableDiffusion",
             "LEditsPPPipelineStableDiffusionXL",
+            "LattePipeline",
             "LuminaText2ImgPipeline",
             "MarigoldDepthPipeline",
             "MarigoldNormalsPipeline",
@@ -513,6 +515,7 @@
             HunyuanDiT2DMultiControlNetModel,
             I2VGenXLUNet,
             Kandinsky3UNet,
+            LatteTransformer3DModel,
             LuminaNextDiT2DModel,
             ModelMixin,
             MotionAdapter,
@@ -675,6 +678,7 @@
             LDMTextToImagePipeline,
             LEditsPPPipelineStableDiffusion,
             LEditsPPPipelineStableDiffusionXL,
+            LattePipeline,
             LuminaText2ImgPipeline,
             MarigoldDepthPipeline,
             MarigoldNormalsPipeline,

src/diffusers/models/__init__.py

@@ -41,6 +41,7 @@
     _import_structure["transformers.dit_transformer_2d"] = ["DiTTransformer2DModel"]
     _import_structure["transformers.dual_transformer_2d"] = ["DualTransformer2DModel"]
     _import_structure["transformers.hunyuan_transformer_2d"] = ["HunyuanDiT2DModel"]
+    _import_structure["transformers.latte_transformer_3d"] = ["LatteTransformer3DModel"]
     _import_structure["transformers.lumina_nextdit2d"] = ["LuminaNextDiT2DModel"]
     _import_structure["transformers.pixart_transformer_2d"] = ["PixArtTransformer2DModel"]
     _import_structure["transformers.prior_transformer"] = ["PriorTransformer"]
@@ -86,6 +87,7 @@
             DiTTransformer2DModel,
             DualTransformer2DModel,
             HunyuanDiT2DModel,
+            LatteTransformer3DModel,
             LuminaNextDiT2DModel,
             PixArtTransformer2DModel,
             PriorTransformer,

src/diffusers/models/attention.py

@@ -359,7 +359,10 @@ def __init__(
                 out_bias=attention_out_bias,
             )  # is self-attn if encoder_hidden_states is none
         else:
-            self.norm2 = None
+            if norm_type == "ada_norm_single": # For Latte
+                self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+            else: 
+                self.norm2 = None
             self.attn2 = None
 
         # 3. Feed-forward
@@ -456,6 +459,7 @@ def forward(
             attention_mask=attention_mask,
             **cross_attention_kwargs,
         )
+
         if self.norm_type == "ada_norm_zero":
             attn_output = gate_msa.unsqueeze(1) * attn_output
         elif self.norm_type == "ada_norm_single":

src/diffusers/models/modeling_outputs.py

@@ -28,4 +28,4 @@ class Transformer2DModelOutput(BaseOutput):
             distributions for the unnoised latent pixels.
     """
 
-    sample: "torch.Tensor"  # noqa: F821
+    sample: "torch.Tensor"  # noqa: F821

src/diffusers/models/transformers/__init__.py

@@ -5,6 +5,7 @@
     from .dit_transformer_2d import DiTTransformer2DModel
     from .dual_transformer_2d import DualTransformer2DModel
     from .hunyuan_transformer_2d import HunyuanDiT2DModel
+    from .latte_transformer_3d import LatteTransformer3DModel
     from .lumina_nextdit2d import LuminaNextDiT2DModel
     from .pixart_transformer_2d import PixArtTransformer2DModel
     from .prior_transformer import PriorTransformer