Allow batch of different sigmas when noise scaling.

comfyanonymous · Jan 30, 2025 · 8d8dc9a · 8d8dc9a
1 parent 2f98c24
commit 8d8dc9a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py
@@ -31,6 +31,7 @@ def calculate_denoised(self, sigma, model_output, model_input):
         return model_input - model_output * sigma
 
     def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
         if max_denoise:
             noise = noise * torch.sqrt(1.0 + sigma ** 2.0)
         else:
@@ -61,9 +62,11 @@ def calculate_denoised(self, sigma, model_output, model_input):
         return model_input - model_output * sigma
 
     def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
         return sigma * noise + (1.0 - sigma) * latent_image
 
     def inverse_noise_scaling(self, sigma, latent):
+        sigma = sigma.view(sigma.shape[:1] + (1,) * (latent.ndim - 1))
         return latent / (1.0 - sigma)
 
 class ModelSamplingDiscrete(torch.nn.Module):