minor

app_v3.py

@@ -1,5 +1,5 @@
 import subprocess
-subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+# subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
 
 import torch
 import spaces
@@ -113,6 +113,17 @@ def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_
         ).images[0]
     return image
 
+def generate_caption(control_image):
+    if control_image is None:
+        return None, None
+    
+    # Generate a detailed caption
+    mcaption = model.caption(control_image, length="long")
+    detailed_caption = mcaption["caption"]
+    print(f"Detailed caption: {detailed_caption}")
+    
+    return detailed_caption
+
 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):
@@ -125,22 +136,7 @@ def process_image(control_image, user_prompt, system_prompt, scale, steps,
         mcaption = model.caption(control_image, length="long")
         detailed_caption = mcaption["caption"]
         print(f"Detailed caption: {detailed_caption}")
-        # Generate caption
-        # caption_gen = caption(
-        #     input_image=control_image,
-        #     prompt=system_prompt,
-        #     temperature=temperature,
-        #     top_p=top_p,
-        #     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 generated_caption, None  # Update caption in real-time
-        
+
         final_prompt = detailed_caption
         yield f"Using caption: {final_prompt}", None, final_prompt
     
@@ -225,11 +221,15 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
         ],
         outputs=[output_caption, generated_image, prompt]
     )
-    
-    # caption_button.click(
-    #     fn=caption,
-    #     inputs=[control_image, system_prompt, temperature_slider, top_p_slider, max_tokens_slider, log_prompt],
-    #     outputs=output_caption,
-    # )
+    control_image.change(
+        generate_caption,
+        inputs=[control_image],
+        outputs=[prompt]
+    )
+    caption_button.click(
+        fn=generate_caption,
+        inputs=[control_image],
+        outputs=[prompt]
+    )
 
-demo.launch()
+demo.launch(show_error=True)