(SD) Fixes schedulers if recieving noise preds as numpy arrays

nod-ai · Dec 4, 2023 · 35cac51 · 35cac51
1 parent d72da38
commit 35cac51
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/apps/stable_diffusion/src/schedulers/shark_eulerancestraldiscrete.py b/apps/stable_diffusion/src/schedulers/shark_eulerancestraldiscrete.py
@@ -211,6 +211,11 @@ def step(
 
         sigma_from = self.sigmas[self.step_index]
         sigma_to = self.sigmas[self.step_index + 1]
+        noise_pred = (
+            torch.from_numpy(noise_pred)
+            if isinstance(noise_pred, np.ndarray)
+            else noise_pred
+        )
         noise = randn_tensor(
             torch.Size(noise_pred.shape),
             dtype=torch.float16,

diff --git a/apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py b/apps/stable_diffusion/src/schedulers/shark_eulerdiscrete.py
@@ -207,6 +207,11 @@ def step(
 
         sigma_hat = sigma * (gamma + 1)
 
+        noise_pred = (
+            torch.from_numpy(noise_pred)
+            if isinstance(noise_pred, np.ndarray)
+            else noise_pred
+        )
         noise = randn_tensor(
             noise_pred.shape,
             dtype=noise_pred.dtype,

diff --git a/apps/stable_diffusion/src/utils/stencils/stencil_utils.py b/apps/stable_diffusion/src/utils/stencils/stencil_utils.py
@@ -154,6 +154,7 @@ def controlnet_hint_conversion(
             controlnet_hint = hint_zoedepth(image)
         case _:
             return None
+    print(f"Controlnet hint is type {image.dtype}")
     controlnet_hint = controlnet_hint_shaping(
         controlnet_hint, height, width, dtype, num_images_per_prompt
     )