log progress

app.py

@@ -9,6 +9,7 @@ from dotenv import load_dotenv
 import base64
 from typing import Optional, Tuple, Union
 import glob
+import asyncio
 
 load_dotenv()
 
@@ -21,11 +22,7 @@ global_input_image = None
 global_image_slider = None
 
 
-def clear_output():
-    return None
-
-
-def process_image(
+async def process_image(
     image: Optional[Union[str, Image.Image]],
     denoise: float,
     steps: int,
@@ -35,6 +32,7 @@ def process_image(
     color_match: float,
     controlnet_tile_end: float,
     controlnet_tile_strength: float,
+    progress: gr.Progress = gr.Progress(),
 ) -> Tuple[Optional[Image.Image], Optional[Image.Image]]:
     # Convert image to base64
     if image is not None:
@@ -69,9 +67,23 @@ def process_image(
 
         if result and result.object:
             run_id: str = result.object.run_id
+            progress(0, desc="Starting processing...")
             # Wait for the result
             while True:
                 run_result = client.run.get(run_id=run_id)
+                progress_value = (
+                    run_result.object.progress
+                    if run_result.object.progress is not None
+                    else 0
+                )
+                status = (
+                    run_result.object.live_status
+                    if run_result.object.live_status is not None
+                    else "Cold starting..."
+                )
+                progress(progress_value, desc=f"Status: {status}")
+                print(f"Progress: {progress_value:.2f}%, Status: {status}")
+
                 if run_result.object.status == "success":
                     for output in run_result.object.outputs:
                         if output.data and output.data.images:
@@ -84,13 +96,16 @@ def process_image(
                             return image, processed_image
                     return None, None
                 elif run_result.object.status == "failed":
+                    print("Processing failed")
                     return None, None
+
+                await asyncio.sleep(2)  # Wait for 2 seconds before checking again
     except Exception as e:
         print(f"Error: {e}")
         return None, None
 
 
-def run(
+async def run_async(
     denoise,
     steps,
     tile_size,
@@ -99,6 +114,7 @@ def run(
     color_match,
     controlnet_tile_end,
     controlnet_tile_strength,
+    progress: gr.Progress = gr.Progress(),
 ):
     global global_input_image
     global global_image_slider
@@ -106,11 +122,9 @@ def run(
     if not global_input_image:
         return None
 
-    # Set image_slider to None before processing
     global_image_slider = None
 
-    # Process the image
-    original, processed = process_image(
+    original, processed = await process_image(
         global_input_image,
         denoise,
         steps,
@@ -120,6 +134,7 @@ def run(
         color_match,
         controlnet_tile_end,
         controlnet_tile_strength,
+        progress,
     )
 
     if original and processed:
@@ -128,6 +143,11 @@ def run(
     return global_image_slider
 
 
+# Wrapper function for gradio to use
+def run(*args):
+    return asyncio.run(run_async(*args))
+
+
 # Function to load preset images
 def load_preset_images():
     image_files = glob.glob("images/inputs/*")
@@ -139,10 +159,22 @@ def load_preset_images():
     ]
 
 
-def set_input_image(images, evt: gr.SelectData):
+# def set_input_image(images, evt: gr.SelectData):
+#     global global_input_image
+#     global global_image_slider
+
+#     global_input_image = images[evt.index][0]
+#     global_image_slider = None
+
+#     return global_input_image
+
+
+def update_global_input_image(image):
     global global_input_image
-    global_input_image = images[evt.index][0]
-    return global_input_image
+    global global_image_slider
+    global_input_image = image
+    global_image_slider = None
+    return None
 
 
 # Define Gradio interface
@@ -156,18 +188,22 @@ with gr.Blocks() as demo:
                 value=lambda: global_input_image,
                 interactive=True,
             )
+            input_image.change(
+                fn=update_global_input_image,
+                inputs=[input_image],
+            )
 
             # Add preset images
-            gr.Markdown("### Preset Images")
-            preset_images = load_preset_images()
-            gallery = gr.Gallery(
-                [img["image"] for img in preset_images],
-                label="Preset Images",
-                columns=5,
-                height=130,
-                allow_preview=False,
-            )
-            gallery.select(set_input_image, gallery, input_image)
+            # gr.Markdown("### Preset Images")
+            # preset_images = load_preset_images()
+            # gallery = gr.Gallery(
+            #     [img["image"] for img in preset_images],
+            #     label="Preset Images",
+            #     columns=5,
+            #     height=130,
+            #     allow_preview=False,
+            # )
+            # gallery.select(set_input_image, gallery, input_image)
 
             with gr.Accordion("Advanced Parameters", open=False):
                 denoise: gr.Slider = gr.Slider(0, 1, value=0.4, label="Denoise")
@@ -195,7 +231,7 @@ with gr.Blocks() as demo:
                 interactive=True,
             )
 
-            process_btn: gr.Button = gr.Button("Run")
+            process_btn: gr.Button = gr.Button("Run", variant="primary", size="lg")
             process_btn.click(
                 fn=run,
                 inputs=[
@@ -248,4 +284,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(debug=True, share=True)
+    demo.queue().launch(debug=True, share=True)