From 8d8dc9a262bfee9f7eb3a2fc381b79a69274f225 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 30 Jan 2025 06:49:52 -0500
Subject: [PATCH] Allow batch of different sigmas when noise scaling.

---
 comfy/model_sampling.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py
index 4370516b..ff27b09a 100644
--- a/comfy/model_sampling.py
+++ b/comfy/model_sampling.py
@@ -31,6 +31,7 @@ class EPS:
         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 @@ class CONST:
         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):