Update app_v2.py

app_v2.py

@@ -105,9 +105,14 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
         ).images[0]
     return image
 
-def process_image(control_image, user_prompt, system_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, temperature, top_p, max_new_tokens, log_prompt):
+def process_image(control_image, user_prompt, system_prompt, scale, steps, 
+                 controlnet_conditioning_scale, guidance_scale, seed, 
+                 guidance_end, temperature, top_p, max_new_tokens, log_prompt):
+    # Initialize with empty caption
+    final_prompt = user_prompt.strip()
+    
     # If no user prompt provided, generate a caption first
-    if not user_prompt.strip():
+    if not final_prompt:
         caption_gen = caption(
             input_image=control_image,
             prompt=system_prompt,
@@ -117,19 +122,38 @@ def process_image(control_image, user_prompt, system_prompt, scale, steps, contr
             log_prompt=log_prompt
         )
         # Get the full caption by exhausting the generator
-        user_prompt = "".join([chunk for chunk in caption_gen])
+        generated_caption = ""
+        for chunk in caption_gen:
+            generated_caption += chunk
+            yield {"__type__": "update_caption", "caption": generated_caption}, None
+        
+        final_prompt = generated_caption
+        yield {"__type__": "update_caption", "caption": f"Using caption: {final_prompt}"}, None
     
-    # Generate the image using the prompt (either user-provided or generated)
-    return generate_image(
-        prompt=user_prompt,
-        scale=scale,
-        steps=steps,
-        control_image=control_image,
-        controlnet_conditioning_scale=controlnet_conditioning_scale,
-        guidance_scale=guidance_scale,
-        seed=seed,
-        guidance_end=guidance_end
-    )
+    # Show the final prompt being used
+    yield {"__type__": "update_caption", "caption": f"Generating with: {final_prompt}"}, None
+    
+    # Generate the image
+    try:
+        image = generate_image(
+            prompt=final_prompt,
+            scale=scale,
+            steps=steps,
+            control_image=control_image,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            guidance_scale=guidance_scale,
+            seed=seed,
+            guidance_end=guidance_end
+        )
+        yield {"__type__": "update_caption", "caption": f"Completed! Used prompt: {final_prompt}"}, image
+    except Exception as e:
+        yield {"__type__": "update_caption", "caption": f"Error: {str(e)}"}, None
+        raise
+        
+def handle_outputs(outputs):
+    if isinstance(outputs, dict) and outputs.get("__type__") == "update_caption":
+        return outputs["caption"], None
+    return outputs
 
 with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as iface:
     gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")
@@ -181,11 +205,15 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as iface:
     generate_button.click(
         fn=process_image,
         inputs=[
-            control_image, prompt, system_prompt, scale, steps, controlnet_conditioning_scale, 
-            guidance_scale, seed, guidance_end, temperature_slider, top_p_slider, 
+            control_image, prompt, system_prompt, scale, steps, 
+            controlnet_conditioning_scale, guidance_scale, seed, 
+            guidance_end, temperature_slider, top_p_slider, 
             max_tokens_slider, log_prompt
         ],
-        outputs=[generated_image]
+        outputs=[output_caption, generated_image]
+    ).then(
+        handle_outputs,
+        outputs=[output_caption, generated_image]
     )
     
     caption_button.click(