add conversion script

Author: a-r-r-o-w

src/diffusers/__init__.py

@@ -106,6 +106,7 @@
             "ControlNetModel",
             "ControlNetUnionModel",
             "ControlNetXSAdapter",
+            "CosmosTransformer3DModel",
             "DiTTransformer2DModel",
             "FluxControlNetModel",
             "FluxMultiControlNetModel",
@@ -620,6 +621,7 @@
             ControlNetModel,
             ControlNetUnionModel,
             ControlNetXSAdapter,
+            CosmosTransformer3DModel,
             DiTTransformer2DModel,
             FluxControlNetModel,
             FluxMultiControlNetModel,

src/diffusers/models/__init__.py

@@ -69,6 +69,7 @@
     _import_structure["transformers.transformer_2d"] = ["Transformer2DModel"]
     _import_structure["transformers.transformer_allegro"] = ["AllegroTransformer3DModel"]
     _import_structure["transformers.transformer_cogview3plus"] = ["CogView3PlusTransformer2DModel"]
+    _import_structure["transformers.transformer_cosmos"] = ["CosmosTransformer3DModel"]
     _import_structure["transformers.transformer_flux"] = ["FluxTransformer2DModel"]
     _import_structure["transformers.transformer_hunyuan_video"] = ["HunyuanVideoTransformer3DModel"]
     _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"]
@@ -133,6 +134,7 @@
             CogVideoXTransformer3DModel,
             CogView3PlusTransformer2DModel,
             ConsisIDTransformer3DModel,
+            CosmosTransformer3DModel,
             DiTTransformer2DModel,
             DualTransformer2DModel,
             FluxTransformer2DModel,

src/diffusers/models/transformers/__init__.py

@@ -18,6 +18,7 @@
     from .transformer_2d import Transformer2DModel
     from .transformer_allegro import AllegroTransformer3DModel
     from .transformer_cogview3plus import CogView3PlusTransformer2DModel
+    from .transformer_cosmos import CosmosTransformer3DModel
     from .transformer_flux import FluxTransformer2DModel
     from .transformer_hunyuan_video import HunyuanVideoTransformer3DModel
     from .transformer_ltx import LTXVideoTransformer3DModel

src/diffusers/models/transformers/transformer_cosmos.py

@@ -356,7 +356,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return (emb / norm).type_as(hidden_states)
 
 
-class CosmosTransformer(ModelMixin, ConfigMixin):
+class CosmosTransformer3DModel(ModelMixin, ConfigMixin):
     r"""
     A Transformer model for video-like data used in [Cosmos](https://github.com/NVIDIA/Cosmos).
 
@@ -423,9 +423,9 @@ def __init__(
             hidden_size=attention_head_dim, max_size=max_size, patch_size=patch_size, rope_scale=rope_scale
         )
 
-        self.learnable_pos_embedder = None
+        self.learnable_pos_embed = None
         if extra_pos_embed_type == "learnable":
-            self.learnable_pos_embedder = CosmosLearnablePositionalEmbed(
+            self.learnable_pos_embed = CosmosLearnablePositionalEmbed(
                 hidden_size=hidden_size,
                 max_size=max_size,
                 patch_size=patch_size,
@@ -477,7 +477,7 @@ def forward(
 
         # 2. Generate positional embeddings
         image_rotary_emb = self.rope(hidden_states, fps=fps)
-        extra_pos_emb = self.learnable_pos_embedder(hidden_states) if self.config.extra_pos_embed_type else None
+        extra_pos_emb = self.learnable_pos_embed(hidden_states) if self.config.extra_pos_embed_type else None
 
         # 3. Patchify input
         batch_size, num_channels, num_frames, height, width = hidden_states.shape

src/diffusers/utils/dummy_pt_objects.py

@@ -351,6 +351,21 @@ def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["torch"])
 
 
+class CosmosTransformer3DModel(metaclass=DummyObject):
+    _backends = ["torch"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["torch"])
+
+    @classmethod
+    def from_config(cls, *args, **kwargs):
+        requires_backends(cls, ["torch"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["torch"])
+
+
 class DiTTransformer2DModel(metaclass=DummyObject):
     _backends = ["torch"]