Update app.py

app.py

@@ -1,44 +1,3 @@
-"""
-{
-    "inputs": {
-        "prompt": "A text prompt to generate the image from.",
-        "image_prompts": [
-            {
-                "cn_img": "Base64 encoded image data for the first image prompt.",
-                "cn_stop": "ControlNet stop value for the first image prompt.",
-                "cn_weight": "ControlNet weight value for the first image prompt.",
-                "cn_type": "Type of the first image prompt."
-            },
-            {
-                "cn_img": "Base64 encoded image data for the second image prompt.",
-                "cn_stop": "ControlNet stop value for the second image prompt.",
-                "cn_weight": "ControlNet weight value for the second image prompt.",
-                "cn_type": "Type of the second image prompt."
-            },
-            {
-                "cn_img": "Base64 encoded image data for the third image prompt.",
-                "cn_stop": "ControlNet stop value for the third image prompt.",
-                "cn_weight": "ControlNet weight value for the third image prompt.",
-                "cn_type": "Type of the third image prompt."
-            },
-            {
-                "cn_img": "Base64 encoded image data for the fourth image prompt.",
-                "cn_stop": "ControlNet stop value for the fourth image prompt.",
-                "cn_weight": "ControlNet weight value for the fourth image prompt.",
-                "cn_type": "Type of the fourth image prompt."
-            }
-        ],
-        "async_process": "Boolean to indicate if the process should be asynchronous."
-    },
-    "outputs": {
-        "job_id": "The ID of the job submitted for image generation.",
-        "job_stage": "The current stage of the job (e.g., RUNNING, SUCCESS, FAILED).",
-        "final_image_url": "The URL of the final generated image.",
-        "step_preview": "Base64 encoded step preview image data."
-    }
-}
-"""
-
 import requests
 from requests.adapters import HTTPAdapter
 from requests.packages.urllib3.util.retry import Retry
@@ -47,13 +6,11 @@ import base64
 import time
 import gradio as gr
 from PIL import Image
+from io import BytesIO
 import os
 
 host = "http://18.119.36.46:8888"
 
-# 📂 Get the directory where the script is located
-script_dir = os.path.dirname(os.path.abspath(__file__))
-
 def image_prompt(prompt, image1, image2, image3, image4):
     source1 = open(image1, "rb").read()
     source2 = open(image2, "rb").read()
@@ -114,28 +71,21 @@ def image_prompt(prompt, image1, image2, image3, image4):
                 if final_image_url:
                     final_image_url = final_image_url.replace("127.0.0.1", "18.119.36.46")
                     image_response = session.get(final_image_url, timeout=10)  # Increase timeout as needed
-                    with open("output.png", "wb") as f:
-                        f.write(image_response.content)
-                    return "output.png", "Job completed successfully."
+                    image = Image.open(BytesIO(image_response.content))
+                    return image, "Job completed successfully."
                 else:
                     return None, "Final image URL not found in the job data."
             elif job_stage == "RUNNING":
                 step_preview_base64 = job_data.get("job_step_preview")
                 if step_preview_base64:
-                    with open("output.png", "wb") as f:
-                        f.write(base64.b64decode(step_preview_base64))
+                    image = Image.open(BytesIO(base64.b64decode(step_preview_base64)))
+                    return image, "Job is running, step preview available."
                 time.sleep(5)
             elif job_stage == "FAILED":
                 return None, "Job failed."
     else:
         return None, "Job ID not found."
 
-def create_status_image():
-    if os.path.exists("output.png"):
-        return "output.png"
-    else:
-        return None
-
 def gradio_app():
     with gr.Blocks() as demo:
         prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here")
@@ -150,12 +100,6 @@ def gradio_app():
         generate_button = gr.Button("Generate Image")
         generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status])
         
-        # 🖼️ Display the status image
-        status_image = gr.Image(label="Queue Status", interactive=False)
-        
-        # ⏲️ Update the image every 5 seconds
-        demo.load(create_status_image, every=5, outputs=status_image)
-    
     demo.launch()
 
 if __name__ == "__main__":