Use safe_progress helper

lib/__init__.py

@@ -11,6 +11,7 @@ from .utils import (
     enable_progress_bars,
     load_json,
     read_file,
+    safe_progress,
     timer,
 )
 
@@ -27,5 +28,6 @@ __all__ = [
     "generate",
     "load_json",
     "read_file",
+    "safe_progress",
     "timer",
 ]

lib/inference.py

@@ -16,7 +16,7 @@ from spaces import GPU
 from .config import Config
 from .loader import Loader
 from .logger import Logger
-from .utils import load_json, timer
+from .utils import load_json, safe_progress, timer
 
 
 def parse_prompt_with_arrays(prompt: str) -> list[str]:
@@ -128,8 +128,8 @@ def generate(
     log = Logger("generate")
     log.info(f"Generating {num_images} image{'s' if num_images > 1 else ''}")
 
-    if Config.ZERO_GPU and progress is not None:
-        progress((100, 100), desc="ZeroGPU init")
+    if Config.ZERO_GPU:
+        safe_progress(progress, 100, 100, "ZeroGPU init")
 
     if not torch.cuda.is_available():
         raise Error("CUDA not available")
@@ -197,7 +197,7 @@ def generate(
     desc_loras = "Loading LoRAs"
     if total_loras > 0:
         with timer(f"Loading {total_loras} LoRA{'s' if total_loras > 1 else ''}"):
-            progress((0, total_loras), desc=desc_loras)
+            safe_progress(progress, 0, total_loras, desc_loras)
             for i, (lora, weight) in enumerate(loras_and_weights):
                 if lora and lora.lower() != "none" and lora not in loras:
                     config = Config.CIVIT_LORAS.get(lora)
@@ -210,7 +210,7 @@ def generate(
                             )
                             weights.append(weight)
                             loras.append(lora)
-                            progress((i + 1, total_loras), desc=desc_loras)
+                            safe_progress(progress, i + 1, total_loras, desc_loras)
                         except Exception:
                             raise Error(f"Error loading {config['name']} LoRA")
 
@@ -247,6 +247,7 @@ def generate(
 
     images = []
     current_seed = seed
+    safe_progress(progress, 0, num_images, f"Generating image 0/{num_images}")
     for i in range(num_images):
         try:
             generator = torch.Generator(device=pipe.device).manual_seed(current_seed)
@@ -299,16 +300,8 @@ def generate(
 
         try:
             image = pipe(**kwargs).images[0]
-            if scale > 1:
-                msg = f"Upscaling {scale}x"
-                with timer(msg, logger=log.info):
-                    progress((0, 100), desc=msg)
-                    image = upscaler.predict(image)
-                    progress((100, 100), desc=msg)
             images.append((image, str(current_seed)))
             current_seed += 1
-        except Exception as e:
-            raise Error(f"{e}")
         finally:
             if embeddings:
                 pipe.unload_textual_inversion()
@@ -317,6 +310,15 @@ def generate(
             CURRENT_STEP = 0
             CURRENT_IMAGE += 1
 
+    if scale > 1:
+        msg = f"Upscaling {scale}x"
+        with timer(msg, logger=log.info):
+            safe_progress(progress, 0, num_images, desc=msg)
+            for i, image in enumerate(images):
+                image = upscaler.predict(image[0])
+                images[i] = image
+                safe_progress(progress, i + 1, num_images, desc=msg)
+
     # cleanup
     loader.collect()
     gc.collect()

lib/loader.py

@@ -9,7 +9,7 @@ from diffusers.models.attention_processor import AttnProcessor2_0, IPAdapterAttn
 from .config import Config
 from .logger import Logger
 from .upscaler import RealESRGAN
-from .utils import timer
+from .utils import safe_progress, timer
 
 
 class Loader:
@@ -62,6 +62,11 @@ class Loader:
             return True
         return False
 
+    def _should_unload_freeu(self, freeu=False):
+        if self._has_freeu and not freeu:
+            return True
+        return False
+
     def _should_unload_ip_adapter(self, model="", ip_adapter=""):
         # unload if model changed
         if self.model and self.model.lower() != model.lower():
@@ -128,7 +133,7 @@ class Loader:
         if self._should_unload_deepcache(deepcache):  # remove deepcache first
             self._unload_deepcache()
 
-        if self._has_freeu and not freeu:
+        if self._should_unload_freeu(freeu):
             self._unload_freeu()
 
         if self._should_unload_upscaler(scale):
@@ -154,6 +159,11 @@ class Loader:
             return True
         return False
 
+    def _should_load_freeu(self, freeu=False):
+        if not self._has_freeu and freeu:
+            return True
+        return False
+
     def _should_load_deepcache(self, interval=1):
         has_deepcache = hasattr(self.pipe, "deepcache")
         if not has_deepcache and interval != 1:
@@ -176,11 +186,9 @@ class Loader:
         if self._should_load_upscaler(scale):
             try:
                 msg = f"Loading {scale}x upscaler"
-                # fmt: off
                 with timer(msg, logger=self.log.info):
                     self.upscaler = RealESRGAN(scale, device=self.pipe.device)
                     self.upscaler.load_weights()
-                # fmt: on
             except Exception as e:
                 self.log.error(f"Error loading {scale}x upscaler: {e}")
                 self.upscaler = None
@@ -194,7 +202,7 @@ class Loader:
 
     # https://github.com/ChenyangSi/FreeU
     def _load_freeu(self, freeu=False):
-        if not self._has_freeu and freeu:
+        if self._should_load_freeu(freeu):
             self.log.info("Enabling FreeU")
             self.pipe.enable_freeu(b1=1.5, b2=1.6, s1=0.9, s2=0.2)
 
@@ -355,30 +363,34 @@ class Loader:
             [
                 self._is_kl_vae and taesd,
                 self._is_tiny_vae and not taesd,
-                not self._has_freeu and freeu,
+                self._should_load_freeu(freeu),
                 self._should_load_deepcache(deepcache),
                 self._should_load_ip_adapter(ip_adapter),
                 self._should_load_upscaler(scale),
             ]
         )
 
-        msg = "Loading additional features"
-        if self._is_kl_vae and taesd or self._is_tiny_vae and not taesd:
-            self._load_vae(taesd, model)
-            progress((CURRENT_STEP, TOTAL_STEPS), desc=msg)
-            CURRENT_STEP += 1
+        desc = "Configuring pipeline"
         if not self._has_freeu and freeu:
             self._load_freeu(freeu)
-            progress((CURRENT_STEP, TOTAL_STEPS), desc=msg)
+            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)
             CURRENT_STEP += 1
+
         if self._should_load_deepcache(deepcache):
             self._load_deepcache(deepcache)
-            progress((CURRENT_STEP, TOTAL_STEPS), desc=msg)
+            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)
             CURRENT_STEP += 1
+
         if self._should_load_ip_adapter(ip_adapter):
             self._load_ip_adapter(ip_adapter)
-            progress((CURRENT_STEP, TOTAL_STEPS), desc=msg)
+            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)
             CURRENT_STEP += 1
+
         if self._should_load_upscaler(scale):
             self._load_upscaler(scale)
-            progress((CURRENT_STEP, TOTAL_STEPS), desc=msg)
+            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)
+            CURRENT_STEP += 1
+
+        if self._is_kl_vae and taesd or self._is_tiny_vae and not taesd:
+            self._load_vae(taesd, model)
+            safe_progress(progress, CURRENT_STEP, TOTAL_STEPS, desc)

lib/utils.py

@@ -58,6 +58,11 @@ def enable_progress_bars():
     diffusers_logging.enable_progress_bar()
 
 
+def safe_progress(progress, current=0, total=0, desc=""):
+    if progress is not None:
+        progress((current, total), desc=desc)
+
+
 def download_repo_files(repo_id, allow_patterns, token=None):
     was_disabled = are_progress_bars_disabled()
     enable_progress_bars()