fix(api): make pix2pix work with scheduler with float sigma

ssube · Apr 13, 2023 · 982dc10 · 982dc10
1 parent eb428e9
commit 982dc10
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/api/onnx_web/diffusers/pipelines/pix2pix.py b/api/onnx_web/diffusers/pipelines/pix2pix.py
@@ -323,7 +323,12 @@ def __call__(
             raise ValueError(
                 f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"
             )
-        latents = latents * self.scheduler.init_noise_sigma.numpy()
+
+        init_noise_sigma = self.scheduler.init_noise_sigma
+        if torch.is_tensor(init_noise_sigma):
+            init_noise_sigma = init_noise_sigma.numpy()
+
+        latents = latents * init_noise_sigma
 
         # 7. Check that shapes of latents and image match the UNet channels
         num_channels_image = image_latents.shape[1]