update

Author: a-r-r-o-w

src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py

@@ -633,9 +633,13 @@ def __call__(
                 current_sigma = self.scheduler.sigmas[i]
                 is_augment_sigma_greater = augment_sigma >= current_sigma
 
+                c_in_augment = self.scheduler._get_conditioning_c_in(augment_sigma)
+                c_in_original = self.scheduler._get_conditioning_c_in(current_sigma)
+
                 current_cond_indicator = cond_indicator * 0 if is_augment_sigma_greater else cond_indicator
                 cond_noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=torch.float32)
                 cond_latent = conditioning_latents + cond_noise * augment_sigma[:, None, None, None, None]
+                cond_latent = cond_latent * c_in_augment / c_in_original
                 cond_latent = current_cond_indicator * cond_latent + (1 - current_cond_indicator) * latents
                 cond_latent = self.scheduler.scale_model_input(cond_latent, t)
                 cond_latent = cond_latent.to(transformer_dtype)
@@ -654,6 +658,7 @@ def __call__(
                     current_uncond_indicator = uncond_indicator * 0 if is_augment_sigma_greater else uncond_indicator
                     uncond_noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=torch.float32)
                     uncond_latent = conditioning_latents + uncond_noise * augment_sigma[:, None, None, None, None]
+                    uncond_latent = uncond_latent * c_in_augment / c_in_original
                     uncond_latent = current_uncond_indicator * uncond_latent + (1 - current_uncond_indicator) * latents
                     uncond_latent = self.scheduler.scale_model_input(uncond_latent, t)
                     uncond_latent = uncond_latent.to(transformer_dtype)

src/diffusers/schedulers/scheduling_edm_euler.py

@@ -161,7 +161,7 @@ def set_begin_index(self, begin_index: int = 0):
         self._begin_index = begin_index
 
     def precondition_inputs(self, sample, sigma):
-        c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
+        c_in = self._get_conditioning_c_in(sigma)
         scaled_sample = sample * c_in
         return scaled_sample
 
@@ -440,5 +440,9 @@ def add_noise(
         noisy_samples = original_samples + noise * sigma
         return noisy_samples
 
+    def _get_conditioning_c_in(self, sigma):
+        c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
+        return c_in
+
     def __len__(self):
         return self.config.num_train_timesteps