Rename everything

Author: miguelmartin75

scripts/convert_cosmos_to_diffusers.py

@@ -62,7 +62,7 @@
     EDMEulerScheduler,
     FlowMatchEulerDiscreteScheduler,
 )
-from diffusers.pipelines.cosmos.pipeline_cosmos25_predict import Cosmos25PredictBase
+from diffusers.pipelines.cosmos.pipeline_cosmos2_5_predict import Cosmos_2_5_PredictBase
 
 
 def remove_keys_(key: str, state_dict: Dict[str, Any]):
@@ -538,7 +538,7 @@ def save_pipeline_cosmos_2_5(args, transformer, vae):
 
     scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True)
 
-    pipe = Cosmos25PredictBase(
+    pipe = Cosmos_2_5_PredictBase(
         text_encoder=text_encoder,
         tokenizer=tokenizer,
         transformer=transformer,

scripts/plot_flow_unipc_sigmas.ipynb

@@ -463,6 +463,7 @@
             "CogView4ControlPipeline",
             "CogView4Pipeline",
             "ConsisIDPipeline",
+            "Cosmos_2_5_PredictBase",
             "Cosmos2TextToImagePipeline",
             "Cosmos2VideoToWorldPipeline",
             "CosmosTextToWorldPipeline",
@@ -1172,6 +1173,7 @@
             CogView4ControlPipeline,
             CogView4Pipeline,
             ConsisIDPipeline,
+            Cosmos_2_5_PredictBase,
             Cosmos2TextToImagePipeline,
             Cosmos2VideoToWorldPipeline,
             CosmosTextToWorldPipeline,

src/diffusers/__init__.py

@@ -165,6 +165,7 @@
     _import_structure["cogview4"] = ["CogView4Pipeline", "CogView4ControlPipeline"]
     _import_structure["consisid"] = ["ConsisIDPipeline"]
     _import_structure["cosmos"] = [
+        "Cosmos_2_5_PredictBase",
         "Cosmos2TextToImagePipeline",
         "CosmosTextToWorldPipeline",
         "CosmosVideoToWorldPipeline",
@@ -616,6 +617,7 @@
             StableDiffusionXLControlNetXSPipeline,
         )
         from .cosmos import (
+            Cosmos_2_5_PredictBase,
             Cosmos2TextToImagePipeline,
             Cosmos2VideoToWorldPipeline,
             CosmosTextToWorldPipeline,

src/diffusers/pipelines/__init__.py

@@ -22,6 +22,7 @@
 
     _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_text2image"] = ["Cosmos2TextToImagePipeline"]
     _import_structure["pipeline_cosmos2_video2world"] = ["Cosmos2VideoToWorldPipeline"]
     _import_structure["pipeline_cosmos_text2world"] = ["CosmosTextToWorldPipeline"]
@@ -35,6 +36,7 @@
     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_text2image import Cosmos2TextToImagePipeline
         from .pipeline_cosmos2_video2world import Cosmos2VideoToWorldPipeline
         from .pipeline_cosmos_text2world import CosmosTextToWorldPipeline

src/diffusers/pipelines/cosmos/__init__.py

@@ -71,11 +71,11 @@ def retrieve_latents(
     Examples:
         ```python
         >>> import torch
-        >>> from diffusers import Cosmos25PredictBase
+        >>> from diffusers import Cosmos_2_5_PredictBase
         >>> from diffusers.utils import export_to_video, load_image, load_video
 
         >>> model_id = "nvidia/Cosmos-Predict2.5-Base-2B"
-        >>> pipe = Cosmos25PredictBase.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+        >>> pipe = Cosmos_2_5_PredictBase.from_pretrained(model_id, torch_dtype=torch.bfloat16)
         >>> pipe.to("cuda")
 
         >>> prompt = "A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess."
@@ -124,7 +124,7 @@ def retrieve_latents(
 """
 
 
-class Cosmos25PredictBase(DiffusionPipeline):
+class Cosmos_2_5_PredictBase(DiffusionPipeline):
     r"""
     Pipeline for [Cosmos Predict2.5](https://github.com/nvidia-cosmos/cosmos-predict2.5) base model.
 

…ines/cosmos/pipeline_cosmos25_predict.py …nes/cosmos/pipeline_cosmos2_5_predict.pysrc/diffusers/pipelines/cosmos/pipeline_cosmos25_predict.py renamed to src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py src/diffusers/pipelines/cosmos/pipeline_cosmos25_predict.py renamed to src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py

@@ -23,7 +23,7 @@
 import numpy as np
 import torch
 
-from diffusers import AutoencoderKLWan, Cosmos25PredictBase, CosmosTransformer3DModel, FlowUniPCMultistepScheduler
+from diffusers import AutoencoderKLWan, Cosmos_2_5_PredictBase, CosmosTransformer3DModel, FlowUniPCMultistepScheduler
 
 from ...testing_utils import enable_full_determinism, torch_device
 from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
@@ -96,15 +96,15 @@ def forward(self, input_ids: torch.LongTensor, output_hidden_states: bool = Fals
         return SimpleNamespace(hidden_states=hidden_states)
 
 
-class Cosmos25PredictBaseWrapper(Cosmos25PredictBase):
+class Cosmos_2_5_PredictBaseWrapper(Cosmos_2_5_PredictBase):
     @staticmethod
     def from_pretrained(*args, **kwargs):
         kwargs["safety_checker"] = DummyCosmosSafetyChecker()
-        return Cosmos25PredictBase.from_pretrained(*args, **kwargs)
+        return Cosmos_2_5_PredictBase.from_pretrained(*args, **kwargs)
 
 
-class Cosmos25PredictPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
-    pipeline_class = Cosmos25PredictBaseWrapper
+class Cosmos_2_5_PredictPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+    pipeline_class = Cosmos_2_5_PredictBaseWrapper
     params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
     batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
     image_params = TEXT_TO_IMAGE_IMAGE_PARAMS