Update app.py

app.py

@@ -9,18 +9,11 @@ import time
 import base64
 import json
 
-# Local Dev
-import os
-from dotenv import load_dotenv
-
-load_dotenv()
-FAL_KEY = os.getenv("FAL_KEY")
-
 with open("examples/examples.json") as f:
     examples = json.load(f)
 
 # IC Light, Replace Background
-async def submit_ic_light_bria(image_data, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color):
+async def submit_ic_light_bria(api_key, image_data, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color):
     if not lightsource_start_color.startswith("#"):
         lightsource_start_color = f"#{lightsource_start_color}"
     if not lightsource_end_color.startswith("#"):
@@ -29,7 +22,6 @@ async def submit_ic_light_bria(image_data, positive_prompt, negative_prompt, lig
     retries = 3
     for attempt in range(retries):
         try:
-            fal_client.set_api_key(FAL_KEY)  # Set the API key before making the request
             handler = await fal_client.submit_async(
                 "comfy/martintmv-git/ic-light-bria",
                 arguments={
@@ -38,7 +30,8 @@ async def submit_ic_light_bria(image_data, positive_prompt, negative_prompt, lig
                     "Negative Prompt": negative_prompt,
                     "lightsource_start_color": lightsource_start_color,
                     "lightsource_end_color": lightsource_end_color
-                }
+                },
+                credentials={"api_key": api_key}  # Pass the user's API key dynamically
             )
 
             log_index = 0
@@ -70,18 +63,18 @@ async def submit_ic_light_bria(image_data, positive_prompt, negative_prompt, lig
                 return [f"Error: {str(e)}"], None
 
 # SDXL, Depth Anything, Replace Background
-async def submit_sdxl_rembg(image_data, positive_prompt, negative_prompt):
+async def submit_sdxl_rembg(api_key, image_data, positive_prompt, negative_prompt):
     retries = 3
     for attempt in range(retries):
         try:
-            fal_client.set_api_key(FAL_KEY)  # Set the API key before making the request
             handler = await fal_client.submit_async(
                 "comfy/martintmv-git/sdxl-depthanything-rembg",
                 arguments={
                     "loadimage_1": image_data,
                     "Positive prompt": positive_prompt,
                     "Negative prompt": negative_prompt
-                }
+                },
+                credentials={"api_key": api_key}  # Pass the user's API key dynamically
             )
 
             log_index = 0
@@ -108,16 +101,15 @@ async def submit_sdxl_rembg(image_data, positive_prompt, negative_prompt):
         except Exception as e:
             print(f"Attempt {attempt + 1} failed: {e}")
             if attempt < retries - 1:
-                time.sleep(2) # HTTP req retry mechanism
+                time.sleep(2)  # HTTP req retry mechanism
             else:
                 return [f"Error: {str(e)}"], None
 
 # SV3D, AnimateDiff
-async def submit_sv3d(image_data, fps, loop_frames_count, gif_loop):
+async def submit_sv3d(api_key, image_data, fps, loop_frames_count, gif_loop):
     retries = 3
     for attempt in range(retries):
         try:
-            fal_client.set_api_key(FAL_KEY)  # Set the API key before making the request
             handler = await fal_client.submit_async(
                 "comfy/martintmv-git/sv3d",
                 arguments={
@@ -126,7 +118,7 @@ async def submit_sv3d(image_data, fps, loop_frames_count, gif_loop):
                     "Loop Frames Count": loop_frames_count,
                     "GIF Loop": gif_loop
                 },
-                credentials={"fal_key": FAL_KEY}
+                credentials={"api_key": api_key}  # Pass the user's API key dynamically
             )
 
             log_index = 0
@@ -158,17 +150,17 @@ def convert_image_to_base64(image):
     image.save(buffered, format="PNG")
     return "data:image/png;base64," + base64.b64encode(buffered.getvalue()).decode()
 
-def submit_sync_ic_light_bria(image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color):
+def submit_sync_ic_light_bria(api_key, image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color):
     image_data = convert_image_to_base64(Image.open(image_upload))
-    return asyncio.run(submit_ic_light_bria(image_data, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color))
+    return asyncio.run(submit_ic_light_bria(api_key, image_data, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color))
 
-def submit_sync_sdxl_rembg(image_upload, positive_prompt, negative_prompt):
+def submit_sync_sdxl_rembg(api_key, image_upload, positive_prompt, negative_prompt):
     image_data = convert_image_to_base64(Image.open(image_upload))
-    return asyncio.run(submit_sdxl_rembg(image_data, positive_prompt, negative_prompt))
+    return asyncio.run(submit_sdxl_rembg(api_key, image_data, positive_prompt, negative_prompt))
 
-def submit_sync_sv3d(image_upload, fps, loop_frames_count, gif_loop):
+def submit_sync_sv3d(api_key, image_upload, fps, loop_frames_count, gif_loop):
     image_data = convert_image_to_base64(Image.open(image_upload))
-    return asyncio.run(submit_sv3d(image_data, fps, loop_frames_count, gif_loop))
+    return asyncio.run(submit_sv3d(api_key, image_data, fps, loop_frames_count, gif_loop))
 
 def run_gradio_app():
     with gr.Blocks() as demo:
@@ -203,55 +195,36 @@ def run_gradio_app():
                 output_result = gr.Image(label="Result")
 
         def validate_api_key(api_key):
-            global FAL_KEY
-            FAL_KEY = api_key
-            return gr.Row(visible=True)
+            return gr.Row(visible=True), api_key
 
-        api_key_submit.click(
-            fn=validate_api_key,
-            inputs=api_key_input,
-            outputs=main_content
-        )
+        api_key_submit.click(validate_api_key, [api_key_input], [main_content, api_key_input])
 
-        def update_ui(workflow):
+        def submit_handler(api_key, workflow, image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color, fps, loop_frames_count, gif_loop):
             if workflow == "IC Light, Replace Background":
-                return [gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)]
+                logs, result_image = submit_sync_ic_light_bria(api_key, image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color)
+                return logs, result_image
             elif workflow == "SDXL, Depth Anything, Replace Background":
-                return [gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)]
+                logs, result_image = submit_sync_sdxl_rembg(api_key, image_upload, positive_prompt, negative_prompt)
+                return logs, result_image
             elif workflow == "SV3D":
-                return [gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)]
+                logs, gif_url = submit_sync_sv3d(api_key, image_upload, fps, loop_frames_count, gif_loop)
+                response = requests.get(gif_url)
+                gif_bytes = BytesIO(response.content)
+                return logs, Image.open(gif_bytes)
 
-        workflow.change(fn=update_ui, inputs=workflow, outputs=[positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color, fps, loop_frames_count, gif_loop])
+        submit_btn.click(submit_handler, [api_key_input, workflow, image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color, fps, loop_frames_count, gif_loop], [output_logs, output_result])
 
-        def on_submit(image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color, fps, loop_frames_count, gif_loop, workflow):
+        def update_fields(workflow):
             if workflow == "IC Light, Replace Background":
-                logs, image = submit_sync_ic_light_bria(image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color)
-                return logs, image
+                return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
             elif workflow == "SDXL, Depth Anything, Replace Background":
-                logs, image = submit_sync_sdxl_rembg(image_upload, positive_prompt, negative_prompt)
-                return logs, image
+                return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
             elif workflow == "SV3D":
-                logs, gif_url = submit_sync_sv3d(image_upload, fps, loop_frames_count, gif_loop)
-                return logs, gif_url
-
-        submit_btn.click(
-            fn=on_submit,
-            inputs=[image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color, fps, loop_frames_count, gif_loop, workflow],
-            outputs=[output_logs, output_result]
-        )
-
-        gr.Examples(
-            examples=[
-                [example['input_image'], example['positive_prompt'], example['negative_prompt'], example.get('lightsource_start_color', "#FFFFFF"), example.get('lightsource_end_color', "#000000"), example.get('fps', 8), example.get('loop_frames_count', 30), example.get('gif_loop', True), example['workflow']]
-                for example in examples
-            ],
-            inputs=[image_upload, positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color, fps, loop_frames_count, gif_loop, workflow],
-            outputs=[output_logs, output_result],
-            fn=on_submit,
-            cache_examples=True
-        )
+                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+
+        workflow.change(update_fields, workflow, [positive_prompt, negative_prompt, lightsource_start_color, lightsource_end_color, fps, loop_frames_count, gif_loop])
 
     demo.launch()
 
 if __name__ == "__main__":
-    run_gradio_app()
+    run_gradio_app()