Update app_v2.py

app_v2.py

@@ -106,13 +106,14 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
     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):
+                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 final_prompt:
+        # Generate caption
         caption_gen = caption(
             input_image=control_image,
             prompt=system_prompt,
@@ -121,17 +122,18 @@ def process_image(control_image, user_prompt, system_prompt, scale, steps,
             max_new_tokens=max_new_tokens,
             log_prompt=log_prompt
         )
+        
         # Get the full caption by exhausting the generator
         generated_caption = ""
         for chunk in caption_gen:
             generated_caption += chunk
-            yield {"__type__": "update_caption", "caption": generated_caption}, None
+            yield generated_caption, None  # Update caption in real-time
         
         final_prompt = generated_caption
-        yield {"__type__": "update_caption", "caption": f"Using caption: {final_prompt}"}, None
+        yield f"Using caption: {final_prompt}", None
     
     # Show the final prompt being used
-    yield {"__type__": "update_caption", "caption": f"Generating with: {final_prompt}"}, None
+    yield f"Generating with: {final_prompt}", None
     
     # Generate the image
     try:
@@ -145,9 +147,9 @@ def process_image(control_image, user_prompt, system_prompt, scale, steps,
             seed=seed,
             guidance_end=guidance_end
         )
-        yield {"__type__": "update_caption", "caption": f"Completed! Used prompt: {final_prompt}"}, image
+        yield f"Completed! Used prompt: {final_prompt}", image
     except Exception as e:
-        yield {"__type__": "update_caption", "caption": f"Error: {str(e)}"}, None
+        yield f"Error: {str(e)}", None
         raise
         
 def handle_outputs(outputs):
@@ -207,13 +209,9 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as iface:
         inputs=[
             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
+            guidance_end, temperature, top_p, max_new_tokens, log_prompt
         ],
         outputs=[output_caption, generated_image]
-    ).then(
-        handle_outputs,
-        outputs=[output_caption, generated_image]
     )
     
     caption_button.click(