Update app_v4.py

app_v4.py

@@ -173,6 +173,53 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
         ).images[0]
         # print("Type: " + str(type(image)))
     return image
+progress = gr.Progress(track_tqdm=True)
+
+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, progress):
+    # Initialize with empty caption
+    final_prompt = user_prompt.strip()
+    # If no user prompt provided, generate a caption first
+    if not final_prompt:
+        # Generate a detailed caption
+        with progress:
+            progress(0.1, "Generating caption...")
+            mcaption = model.caption(control_image, length="normal")
+            detailed_caption = mcaption["caption"]
+            final_prompt = detailed_caption
+            yield f"Using caption: {final_prompt}", None, final_prompt
+    
+    # Show the final prompt being used
+    with progress:
+        progress(0.3, "Generating with prompt...")
+        yield f"Generating with: {final_prompt}", None, final_prompt
+    
+    # Generate the image
+    try:
+        with progress:
+            progress(0.5, "Generating image...")
+            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
+            )
+            
+            try:
+                debug_img = Image.open(image.save("/tmp/" + str(seed) + "output.png"))
+                save_image("/tmp/" + str(seed) + "output.png", debug_img)
+            except Exception as e:
+                print("Error 160: " + str(e))
+            log_params(final_prompt, scale, steps, controlnet_conditioning_scale, guidance_scale, seed, guidance_end, control_image, image)
+            yield f"Completed! Used prompt: {final_prompt}", image, final_prompt
+    except Exception as e:
+        print("Error: " + str(e))
+        yield f"Error: {str(e)}", None, None
 
 with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
     gr.Markdown("⚠️ WIP SPACE - UNFINISHED & BUGGY")
@@ -232,8 +279,8 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
             x_ip_token = request.headers['x-ip-token']
             client = Client("LPX55/zerogpu-experiments", hf_token=huggingface_token, headers={"x-ip-token": x_ip_token})
             cresult = client.predict(
-            		n=3,
-            		api_name="/predict"
+                    n=3,
+                    api_name="/predict"
             )
             print(f"X TOKEN: {x_ip_token}")
             print(cresult)
@@ -261,13 +308,12 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
     caption_state = gr.State()
     focus_state = gr.State()
     log_state = gr.State()
-
     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, max_tokens_slider, log_prompt
+            guidance_end, temperature_slider, top_p_slider, max_tokens_slider, log_prompt, progress
         ],
         outputs=[log_state, generated_image, prompt]
     )
@@ -297,12 +343,10 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
         inputs=[caption_state, focus_state],
         outputs=[prompt]
     )
-
     def hello(profile: gr.OAuthProfile | None) -> str:
         if profile is None:
             return "Hello guest! There is a bug with HF ZeroGPUs that are afffecting some usage on certain spaces. Testing out some possible solutions."
         return f"You are logged in as {profile.name}. If you run into incorrect messages about ZeroGPU runtime credits being out, PLEASE give me a heads up so I can investigate further."
 
-
     demo.load(hello, inputs=None, outputs=msg1)
-demo.queue().launch(show_error=True)
+demo.queue().launch(show_error=True)