update tests

Author: a-r-r-o-w

tests/pipelines/cosmos/test_cosmos2_text2_image.py renamed to tests/pipelines/cosmos/test_cosmos2_text2image.py

@@ -22,7 +22,7 @@
 import torch
 from transformers import AutoTokenizer, T5EncoderModel
 
-from diffusers import AutoencoderKLWan, Cosmos2TextToImagePipeline, CosmosTransformer3DModel, EDMEulerScheduler
+from diffusers import AutoencoderKLWan, Cosmos2TextToImagePipeline, CosmosTransformer3DModel, FlowMatchEulerDiscreteScheduler
 from diffusers.utils.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
@@ -89,17 +89,7 @@ def get_dummy_components(self):
         )
 
         torch.manual_seed(0)
-        scheduler = EDMEulerScheduler(
-            sigma_min=0.002,
-            sigma_max=80,
-            sigma_data=0.5,
-            sigma_schedule="karras",
-            num_train_timesteps=1000,
-            prediction_type="epsilon",
-            rho=7.0,
-            final_sigmas_type="sigma_min",
-            use_flow_sigmas=True,
-        )
+        scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True)
         text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
         tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
 

tests/pipelines/cosmos/test_cosmos2_video2world.py

@@ -90,17 +90,7 @@ def get_dummy_components(self):
         )
 
         torch.manual_seed(0)
-        scheduler = EDMEulerScheduler(
-            sigma_min=0.002,
-            sigma_max=80,
-            sigma_data=0.5,
-            sigma_schedule="karras",
-            num_train_timesteps=1000,
-            prediction_type="epsilon",
-            rho=7.0,
-            final_sigmas_type="sigma_min",
-            use_flow_sigmas=True,
-        )
+        scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True)
         text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
         tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")