wrapper pipelines + make style

Author: miguelmartin75

scripts/convert_cosmos_to_diffusers.py

@@ -62,7 +62,7 @@
     EDMEulerScheduler,
     FlowMatchEulerDiscreteScheduler,
 )
-from diffusers.pipelines.cosmos.pipeline_cosmos2_5_predict import Cosmos_2_5_PredictBase
+from diffusers.pipelines.cosmos.pipeline_cosmos2_5_predict import Cosmos2_5_PredictBase
 
 
 def remove_keys_(key: str, state_dict: Dict[str, Any]):
@@ -527,7 +527,7 @@ def save_pipeline_cosmos_2_0(args, transformer, vae):
     pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
 
 
-def save_pipeline_cosmos_2_5(args, transformer, vae):
+def save_pipeline_cosmos2_5(args, transformer, vae):
     text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B"
     tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct"
 
@@ -538,7 +538,7 @@ def save_pipeline_cosmos_2_5(args, transformer, vae):
 
     scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True)
 
-    pipe = Cosmos_2_5_PredictBase(
+    pipe = Cosmos2_5_PredictBase(
         text_encoder=text_encoder,
         tokenizer=tokenizer,
         transformer=transformer,
@@ -613,6 +613,6 @@ def get_args():
             assert args.tokenizer_path is not None
             save_pipeline_cosmos_2_0(args, transformer, vae)
         elif "Cosmos-2.5" in args.transformer_type:
-            save_pipeline_cosmos_2_5(args, transformer, vae)
+            save_pipeline_cosmos2_5(args, transformer, vae)
         else:
             raise AssertionError(f"{args.transformer_type} not supported")

src/diffusers/__init__.py

@@ -464,7 +464,10 @@
             "CogView4ControlPipeline",
             "CogView4Pipeline",
             "ConsisIDPipeline",
-            "Cosmos_2_5_PredictBase",
+            "Cosmos2_5_PredictBase",
+            "Cosmos2_5_PredictImage2World",
+            "Cosmos2_5_PredictText2World",
+            "Cosmos2_5_PredictVideo2World",
             "Cosmos2TextToImagePipeline",
             "Cosmos2VideoToWorldPipeline",
             "CosmosTextToWorldPipeline",
@@ -1178,9 +1181,12 @@
             CogView4ControlPipeline,
             CogView4Pipeline,
             ConsisIDPipeline,
+            Cosmos2_5_PredictBase,
+            Cosmos2_5_PredictImage2World,
+            Cosmos2_5_PredictText2World,
+            Cosmos2_5_PredictVideo2World,
             Cosmos2TextToImagePipeline,
             Cosmos2VideoToWorldPipeline,
-            Cosmos_2_5_PredictBase,
             CosmosTextToWorldPipeline,
             CosmosVideoToWorldPipeline,
             CycleDiffusionPipeline,

src/diffusers/pipelines/__init__.py

@@ -165,7 +165,10 @@
     _import_structure["cogview4"] = ["CogView4Pipeline", "CogView4ControlPipeline"]
     _import_structure["consisid"] = ["ConsisIDPipeline"]
     _import_structure["cosmos"] = [
-        "Cosmos_2_5_PredictBase",
+        "Cosmos2_5_PredictBase",
+        "Cosmos2_5_PredictImage2World",
+        "Cosmos2_5_PredictText2World",
+        "Cosmos2_5_PredictVideo2World",
         "Cosmos2TextToImagePipeline",
         "CosmosTextToWorldPipeline",
         "CosmosVideoToWorldPipeline",
@@ -623,9 +626,12 @@
             StableDiffusionXLControlNetXSPipeline,
         )
         from .cosmos import (
+            Cosmos2_5_PredictBase,
+            Cosmos2_5_PredictImage2World,
+            Cosmos2_5_PredictText2World,
+            Cosmos2_5_PredictVideo2World,
             Cosmos2TextToImagePipeline,
             Cosmos2VideoToWorldPipeline,
-            Cosmos_2_5_PredictBase,
             CosmosTextToWorldPipeline,
             CosmosVideoToWorldPipeline,
         )

src/diffusers/pipelines/cosmos/__init__.py

@@ -22,7 +22,12 @@
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
-    _import_structure["pipeline_cosmos2_5_predict"] = ["Cosmos_2_5_PredictBase", "retrieve_latents"]
+    _import_structure["pipeline_cosmos2_5_predict"] = [
+        "Cosmos2_5_PredictBase",
+        "Cosmos2_5_PredictImage2World",
+        "Cosmos2_5_PredictText2World",
+        "Cosmos2_5_PredictVideo2World",
+    ]
     _import_structure["pipeline_cosmos2_text2image"] = ["Cosmos2TextToImagePipeline"]
     _import_structure["pipeline_cosmos2_video2world"] = ["Cosmos2VideoToWorldPipeline"]
     _import_structure["pipeline_cosmos_text2world"] = ["CosmosTextToWorldPipeline"]
@@ -36,7 +41,12 @@
     except OptionalDependencyNotAvailable:
         from ...utils.dummy_torch_and_transformers_objects import *
     else:
-        from .pipeline_cosmos2_5_predict import Cosmos_2_5_PredictBase, retrieve_latents
+        from .pipeline_cosmos2_5_predict import (
+            Cosmos2_5_PredictBase,
+            Cosmos2_5_PredictImage2World,
+            Cosmos2_5_PredictText2World,
+            Cosmos2_5_PredictVideo2World,
+        )
         from .pipeline_cosmos2_text2image import Cosmos2TextToImagePipeline
         from .pipeline_cosmos2_video2world import Cosmos2VideoToWorldPipeline
         from .pipeline_cosmos_text2world import CosmosTextToWorldPipeline