From f104668770a26819de8ff05b7a6e0d2cf57386ca Mon Sep 17 00:00:00 2001
From: Tung Nguyen <tung.nguyen@sipher.xyz>
Date: Wed, 30 Aug 2023 15:51:44 +0700
Subject: [PATCH] fix fooocus sampler error when sample batch size > 1

---
 modules/fooocus/patch.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/modules/fooocus/patch.py b/modules/fooocus/patch.py
index c6f9570..d68d216 100644
--- a/modules/fooocus/patch.py
+++ b/modules/fooocus/patch.py
@@ -445,6 +445,10 @@ def unet_forward_patched(
     alpha *= 0.001 * sharpness
     degraded_x0 = bilateral_blur(x0) * alpha + x0 * (1.0 - alpha)
 
+    # FIX: uc_mask is not always the same size as x0
+    if uc_mask.shape[0] < x0.shape[0]:
+        uc_mask = uc_mask.repeat(int(x0.shape[0] / uc_mask.shape[0]), 1, 1, 1)
+
     x0 = x0 * uc_mask + degraded_x0 * (1.0 - uc_mask)
 
     return x0