temp fix

app.py

@@ -1,153 +1,86 @@
+import os
 import gradio as gr
 from gradio_imageslider import ImageSlider
-import os
 from comfydeploy import ComfyDeploy
-import requests
 from PIL import Image
+import requests
 from io import BytesIO
-from dotenv import load_dotenv
 import base64
-from typing import Optional, Tuple, Union
 import glob
 import asyncio
+from dotenv import load_dotenv
 
 load_dotenv()
 
-# Initialize ComfyDeploy client
-client: ComfyDeploy = ComfyDeploy(bearer_auth=os.environ.get("COMFY_DEPLOY_API_KEY"))
-deployment_id: str = os.environ.get("COMFY_DEPLOYMENT_ID")
+API_KEY = os.environ.get("COMFY_DEPLOY_API_KEY")
+DEPLOYMENT_ID = os.environ.get("COMFY_DEPLOYMENT_ID")
+
+if not API_KEY or not DEPLOYMENT_ID:
+    raise ValueError(
+        "Please set COMFY_DEPLOY_API_KEY and COMFY_DEPLOYMENT_ID in your environment variables"
+    )
+
+client = ComfyDeploy(bearer_auth=API_KEY)
 
-# Add these global variables at the top of the file, after imports
-global_input_image = None
-global_image_slider = None
+
+def get_base64_from_image(image: Image.Image) -> str:
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 
 async def process_image(
-    image: Optional[Union[str, Image.Image]],
-    denoise: float,
-    steps: int,
-    tile_size: int,
-    downscale: float,
-    upscale: float,
-    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:
-        if isinstance(image, str):
-            with open(image, "rb") as img_file:
-                image_base64: str = base64.b64encode(img_file.read()).decode("utf-8")
-        else:
-            buffered: BytesIO = BytesIO()
-            image.save(buffered, format="PNG")
-            image_base64: str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    else:
-        return None, None
+    image: Image.Image, params: dict, progress: gr.Progress
+) -> tuple[Image.Image, Image.Image]:
+    progress(0, desc="Preparing inputs...")
+    image_base64 = get_base64_from_image(image)
 
-    # Prepare inputs
-    inputs: dict = {
+    inputs = {
         "image": f"data:image/png;base64,{image_base64}",
-        "denoise": str(denoise),
-        "steps": str(steps),
-        "tile_size": str(tile_size),
-        "downscale": str(downscale),
-        "upscale": str(upscale),
-        "color_match": str(color_match),
-        "controlnet_tile_end": str(controlnet_tile_end),
-        "controlnet_tile_strength": str(controlnet_tile_strength),
+        **{k: str(v) for k, v in params.items()},
     }
 
-    # Call ComfyDeploy API
     try:
         result = client.run.create(
-            request={"deployment_id": deployment_id, "inputs": inputs}
+            request={"deployment_id": DEPLOYMENT_ID, "inputs": inputs}
         )
-
-        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}")
-
-                if run_result.object.status == "success":
-                    for output in run_result.object.outputs:
-                        if output.data and output.data.images:
-                            image_url: str = output.data.images[0].url
-                            # Download and return both the original and processed images
-                            response: requests.Response = requests.get(image_url)
-                            processed_image: Image.Image = Image.open(
-                                BytesIO(response.content)
-                            )
-                            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
+        run_id = result.object.run_id
+        progress(0, desc="Starting processing...")
+
+        while True:
+            run_result = client.run.get(run_id=run_id)
+            progress_value = run_result.object.progress or 0
+            status = run_result.object.live_status or "Cold starting..."
+            progress(progress_value, desc=f"Status: {status}")
+
+            print(run_result)
+
+            if run_result.object.status == "success":
+                for output in run_result.object.outputs:
+                    if output.data and output.data.images:
+                        image_url: str = output.data.images[0].url
+
+                        response = requests.get(image_url)
+                        processed_image = Image.open(BytesIO(response.content))
+                        return image, processed_image
+            elif run_result.object.status == "failed":
+                print("Processing failed")
+                return None, None
+
+            await asyncio.sleep(2)
     except Exception as e:
         print(f"Error: {e}")
         return None, None
 
 
-async def run_async(
-    denoise,
-    steps,
-    tile_size,
-    downscale,
-    upscale,
-    color_match,
-    controlnet_tile_end,
-    controlnet_tile_strength,
-    progress: gr.Progress = gr.Progress(),
-):
-    global global_input_image
-    global global_image_slider
-
-    if not global_input_image:
-        return None
-
-    global_image_slider = None
-
-    original, processed = await process_image(
-        global_input_image,
-        denoise,
-        steps,
-        tile_size,
-        downscale,
-        upscale,
-        color_match,
-        controlnet_tile_end,
-        controlnet_tile_strength,
-        progress,
-    )
-
-    if original and processed:
-        global_image_slider = [original, processed]
-
-    return global_image_slider
+async def run_async(image: Image.Image, params: dict, progress: gr.Progress):
+    return await process_image(image, params, progress)
 
 
-# Wrapper function for gradio to use
-def run(*args):
-    return asyncio.run(run_async(*args))
+def run(image: Image.Image, params: dict):
+    return asyncio.run(run_async(image, params, gr.Progress()))
 
 
-# Function to load preset images
 def load_preset_images():
     image_files = glob.glob("images/inputs/*")
     return [
@@ -158,129 +91,75 @@ def load_preset_images():
     ]
 
 
-# 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 global_image_slider
-    global_input_image = image
-    global_image_slider = None
-    return None
+def build_example(input_image_path):
+    output_image_path = input_image_path.replace("inputs", "outputs")
+    return [
+        input_image_path,
+        0.4,
+        10,
+        1024,
+        1,
+        4,
+        0,
+        1,
+        0.7,
+        (input_image_path, output_image_path),
+    ]
 
 
-# Define Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# 🌊 Creative Image Upscaler")
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(
-                type="pil",
-                label="Input Image",
-                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)
-
-            with gr.Accordion("Advanced Parameters", open=False):
-                denoise: gr.Slider = gr.Slider(0, 1, value=0.4, label="Denoise")
-                steps: gr.Slider = gr.Slider(1, 25, value=10, step=1, label="Steps")
-                tile_size: gr.Slider = gr.Slider(
-                    256, 2048, value=1024, step=256, label="Tile Size"
-                )
-                downscale: gr.Slider = gr.Slider(
-                    1, 4, value=1, step=1, label="Downscale"
-                )
-                upscale: gr.Slider = gr.Slider(1, 4, value=4, step=1, label="Upscale")
-                color_match: gr.Slider = gr.Slider(0, 1, value=0, label="Color Match")
-                controlnet_tile_end: gr.Slider = gr.Slider(
+            input_image = gr.Image(type="pil", label="Input Image", interactive=True)
+            params = {
+                "denoise": gr.Slider(0, 1, value=0.4, label="Denoise"),
+                "steps": gr.Slider(1, 25, value=10, label="Steps"),
+                "tile_size": gr.Slider(256, 2048, value=1024, label="Tile Size"),
+                "downscale": gr.Slider(1, 4, value=1, label="Downscale"),
+                "upscale": gr.Slider(1, 4, value=4, label="Upscale"),
+                "color_match": gr.Slider(0, 1, value=0, label="Color Match"),
+                "controlnet_tile_end": gr.Slider(
                     0, 1, value=1, label="ControlNet Tile End"
-                )
-                controlnet_tile_strength: gr.Slider = gr.Slider(
+                ),
+                "controlnet_tile_strength": gr.Slider(
                     0, 1, value=0.7, label="ControlNet Tile Strength"
-                )
+                ),
+            }
 
-        with gr.Column():
+            process_btn = gr.Button("Run", variant="primary", size="lg")
             image_slider = ImageSlider(
-                label="Compare Original and Processed",
-                type="pil",
-                value=lambda: global_image_slider,
-                interactive=True,
+                label="Compare Original and Processed", type="pil", interactive=True
             )
 
-            process_btn: gr.Button = gr.Button("Run", variant="primary", size="lg")
             process_btn.click(
-                fn=run,
-                inputs=[
-                    denoise,
-                    steps,
-                    tile_size,
-                    downscale,
-                    upscale,
-                    color_match,
-                    controlnet_tile_end,
-                    controlnet_tile_strength,
-                ],
+                fn=lambda img,
+                denoise,
+                steps,
+                tile_size,
+                downscale,
+                upscale,
+                color_match,
+                controlnet_tile_end,
+                controlnet_tile_strength: run(
+                    img,
+                    {
+                        "denoise": denoise,
+                        "steps": steps,
+                        "tile_size": tile_size,
+                        "downscale": downscale,
+                        "upscale": upscale,
+                        "color_match": color_match,
+                        "controlnet_tile_end": controlnet_tile_end,
+                        "controlnet_tile_strength": controlnet_tile_strength,
+                    },
+                ),
+                inputs=[input_image, *params.values()],
                 outputs=[image_slider],
             )
 
-    def build_example(input_image_path):
-        output_image_path = input_image_path.replace("inputs", "outputs")
-        return [
-            input_image_path,
-            0.4,  # denoise
-            10,  # steps
-            1024,  # tile_size
-            1,  # downscale
-            4,  # upscale
-            0,  # color_match
-            1,  # controlnet_tile_end
-            0.7,  # controlnet_tile_strength
-            (input_image_path, output_image_path),
-        ]
-
-    # Build examples
-    input_images = glob.glob("images/inputs/*")
-    examples = [build_example(img) for img in input_images]
-
-    # Update the gr.Examples call
-    gr.Examples(
-        examples=examples,
-        inputs=[
-            input_image,
-            denoise,
-            steps,
-            tile_size,
-            downscale,
-            upscale,
-            color_match,
-            controlnet_tile_end,
-            controlnet_tile_strength,
-            image_slider,
-        ],
-    )
+    examples = [build_example(img) for img in glob.glob("images/inputs/*")]
+    gr.Examples(examples=examples, inputs=[input_image, *params.values(), image_slider])
 
 if __name__ == "__main__":
     demo.queue().launch(debug=True, share=True)