bugfix for schedule (#64)

replicate · Dec 21, 2024 · 2610fcc · 2610fcc
1 parent 7807dd3
commit 2610fcc
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/predict.py b/predict.py
@@ -514,7 +514,8 @@ def bf16_predict(
         )
         timesteps = get_schedule(
             num_inference_steps,
-            x.shape[1],
+            # equivalent to inp["img"].shape[1], needs to be here for prompt strength in img2img
+            (x.shape[-1] * x.shape[-2]) // 4,
             shift=self.shift,
         )