fix tests

Author: a-r-r-o-w

src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py

@@ -546,9 +546,8 @@ def __call__(
         sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
         sigmas = torch.linspace(0, 1, num_inference_steps, dtype=sigmas_dtype)
         timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, device=device, sigmas=sigmas)
-        if self.scheduler.config.final_sigmas_type == "sigma_min":
+        if self.scheduler.config.get("final_sigmas_type", "zero") == "sigma_min":
             # Replace the last sigma (which is zero) with the minimum sigma value
-            timesteps[-1] = timesteps[-2]
             self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2]
 
         # 5. Prepare latent variables

src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py

@@ -633,7 +633,6 @@ def __call__(
         timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, device=device, sigmas=sigmas)
         if self.scheduler.config.final_sigmas_type == "sigma_min":
             # Replace the last sigma (which is zero) with the minimum sigma value
-            timesteps[-1] = timesteps[-2]
             self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2]
 
         # 5. Prepare latent variables
@@ -687,7 +686,9 @@ def __call__(
                 c_in = 1 - current_t
                 c_skip = 1 - current_t
                 c_out = -current_t
-                timestep = current_t.expand(latents.shape[0]).to(transformer_dtype)  # [B, 1, T, 1, 1]
+                timestep = current_t.view(1, 1, 1, 1, 1).expand(
+                    latents.size(0), -1, latents.size(2), -1, -1
+                )  # [B, 1, T, 1, 1]
 
                 cond_latent = latents * c_in
                 cond_latent = cond_indicator * conditioning_latents + (1 - cond_indicator) * cond_latent

tests/pipelines/cosmos/test_cosmos2_text2image.py

@@ -22,7 +22,12 @@
 import torch
 from transformers import AutoTokenizer, T5EncoderModel
 
-from diffusers import AutoencoderKLWan, Cosmos2TextToImagePipeline, CosmosTransformer3DModel, FlowMatchEulerDiscreteScheduler
+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

tests/pipelines/cosmos/test_cosmos2_video2world.py

@@ -23,7 +23,12 @@
 import torch
 from transformers import AutoTokenizer, T5EncoderModel
 
-from diffusers import AutoencoderKLWan, Cosmos2VideoToWorldPipeline, CosmosTransformer3DModel, EDMEulerScheduler
+from diffusers import (
+    AutoencoderKLWan,
+    Cosmos2VideoToWorldPipeline,
+    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