Refactor progress tracking in generate_image function

tabs/images/events.py

@@ -413,9 +413,11 @@ def generate_image(
     resize_mode,
     scheduler, image_height, image_width, image_num_images_per_prompt, # type: ignore
     image_num_inference_steps, image_clip_skip, image_guidance_scale, image_seed, # type: ignore
-    refiner, vae
+    refiner, vae,
+    progress=gr.Progress(track_tqdm=True)
 ):
     try:
+        progress(0, "Configuring arguments...")
         base_args = {
             "model": model,
             "prompt": prompt,
@@ -507,7 +509,7 @@ def generate_image(
             base_args = BaseReq(**base_args.__dict__)
         
         return gr.update(
-            value=gen_img(base_args),
+            value=gen_img(base_args, progress),
             interactive=True
         )
     except Exception as e:

tabs/images/handlers.py

@@ -205,10 +205,12 @@ def cleanup(pipeline, loras = None, embeddings = None):
 
 
 # Gen Function
-def gen_img(request: BaseReq | BaseImg2ImgReq | BaseInpaintReq):
+def gen_img(request: BaseReq | BaseImg2ImgReq | BaseInpaintReq, progress=gr.Progress(track_tqdm=True)):
+    progress(0.1, "Loading Pipeline")
     pipeline_args = get_pipe(request)
     pipeline = pipeline_args["pipeline"]
     try:
+        progress(0.5, "Configuring Pipeline")
         positive_prompt_embeds, negative_prompt_embeds, positive_prompt_pooled, negative_prompt_pooled = get_prompt_attention(pipeline, request.prompt, request.negative_prompt)
 
         # Common Args
@@ -243,15 +245,16 @@ def gen_img(request: BaseReq | BaseImg2ImgReq | BaseInpaintReq):
             args['mask_image'] = resize_images([request.mask_image], request.height, request.width, request.resize_mode)[0]
 
         # Generate
+        progress(0.9, "Generating Images")
         images = pipeline(**args).images
 
         # Refiner
         if request.refiner:
             images = refiner(image=images, prompt=request.prompt, num_inference_steps=40, denoising_start=0.7).images
 
+        progress(1.0, "Cleaning Up")
+        cleanup(pipeline, request.loras, request.embeddings)
         return images
     except Exception as e:
         cleanup(pipeline, request.loras, request.embeddings)
         raise gr.Error(f"Error: {e}")
-    finally:
-        cleanup(pipeline, request.loras, request.embeddings)