Fix progress bars

app.py

@@ -87,7 +87,13 @@ async def generate_fn(*args):
     if prompt is None or prompt.strip() == "":
         raise gr.Error("You must enter a prompt")
     try:
-        images = await async_call(generate, *args, Info=gr.Info, Error=gr.Error)
+        images = await async_call(
+            generate,
+            *args,
+            Info=gr.Info,
+            Error=gr.Error,
+            progress=gr.Progress(),
+        )
     except RuntimeError:
         raise gr.Error("RuntimeError: Please try again")
     return images

lib/inference.py

@@ -9,7 +9,6 @@ from itertools import product
 from typing import Callable, TypeVar
 
 import anyio
-import gradio as gr
 import numpy as np
 import spaces
 import torch
@@ -123,7 +122,7 @@ def generate(
     clip_skip=False,
     Info: Callable[[str], None] = None,
     Error=Exception,
-    progress=gr.Progress(),
+    progress=None,
 ):
     if not torch.cuda.is_available():
         raise Error("CUDA not available")
@@ -150,10 +149,15 @@ def generate(
         IP_ADAPTER = ""
 
     if progress is not None:
+        TQDM = False
         progress((0, inference_steps), desc=f"Generating image {CURRENT_IMAGE}/{num_images}")
+    else:
+        TQDM = True
 
     def callback_on_step_end(pipeline, step, timestep, latents):
         nonlocal CURRENT_IMAGE
+        if progress is None:
+            return latents
         strength = denoising_strength if KIND == "img2img" else 1
         total_steps = min(int(inference_steps * strength), inference_steps)
         current_step = step + 1
@@ -177,6 +181,7 @@ def generate(
         freeu,
         deepcache,
         scale,
+        TQDM,
         DEVICE,
     )
 

lib/loader.py

@@ -110,7 +110,7 @@ class Loader:
             self.upscaler = RealESRGAN(device=device, scale=scale)
             self.upscaler.load_weights()
 
-    def _load_pipeline(self, kind, model, device, **kwargs):
+    def _load_pipeline(self, kind, model, tqdm, device, **kwargs):
         pipeline = Config.PIPELINES[kind]
         if self.pipe is None:
             print(f"Loading {model}...")
@@ -131,7 +131,9 @@ class Loader:
 
         if not isinstance(self.pipe, pipeline):
             self.pipe = pipeline.from_pipe(self.pipe).to(device)
-        self.pipe.set_progress_bar_config(disable=True)
+
+        if not tqdm:
+            self.pipe.set_progress_bar_config(disable=True)
 
     def _load_vae(self, taesd=False, model=""):
         vae_type = type(self.pipe.vae)
@@ -202,6 +204,7 @@ class Loader:
         freeu,
         deepcache,
         scale,
+        tqdm,
         device,
     ):
         scheduler_kwargs = {
@@ -242,7 +245,7 @@ class Loader:
             pipe_kwargs["torch_dtype"] = torch.float16
 
         self._unload(kind, model, ip_adapter, scale)
-        self._load_pipeline(kind, model, device, **pipe_kwargs)
+        self._load_pipeline(kind, model, tqdm, device, **pipe_kwargs)
 
         # error loading model
         if self.pipe is None: