Update app.py

app.py

@@ -7,10 +7,12 @@ 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()
@@ -58,33 +60,38 @@ def image_prompt(prompt, image1, image2, image3, image4):
     result = response.json()
     
     job_id = result.get('job_id')
-    if job_id:
-        while True:
-            query_url = f"http://18.119.36.46:8888/v1/generation/query-job?job_id={job_id}&require_step_preview=true"
-            response = session.get(query_url, timeout=10)  # Increase timeout as needed
-            job_data = response.json()
-            
-            job_stage = job_data.get("job_stage")
-            
-            if job_stage == "SUCCESS":
-                final_image_url = job_data.get("job_result")[0].get("url")
-                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
-                    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:
-                    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."
+    return job_id
+
+def query_status(job_id):
+    session = requests.Session()
+    retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504])
+    session.mount('http://', HTTPAdapter(max_retries=retries))
+    
+    query_url = f"http://18.119.36.46:8888/v1/generation/query-job?job_id={job_id}&require_step_preview=true"
+    response = session.get(query_url, timeout=10)  # Increase timeout as needed
+    job_data = response.json()
+    
+    job_stage = job_data.get("job_stage")
+    
+    if job_stage == "SUCCESS":
+        final_image_url = job_data.get("job_result")[0].get("url")
+        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
+            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:
+            image = Image.open(BytesIO(base64.b64decode(step_preview_base64)))
+            return image, "Job is running, step preview available."
+        return None, "Job is running, no step preview available."
+    elif job_stage == "FAILED":
+        return None, "Job failed."
     else:
-        return None, "Job ID not found."
+        return None, "Unknown job stage."
 
 def gradio_app():
     with gr.Blocks() as demo:
@@ -97,9 +104,20 @@ def gradio_app():
         output_image = gr.Image(label="Generated Image")
         status = gr.Textbox(label="Status")
         
+        job_id = gr.State()
+        
         generate_button = gr.Button("Generate Image")
-        generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status])
+        generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=job_id)
         
+        def update_image(job_id):
+            if job_id:
+                image, status_text = query_status(job_id)
+                return image, status_text
+            return None, "No job ID found."
+        
+        # ⏲️ Update the image every 2 seconds
+        demo.load(update_image, inputs=job_id, outputs=[output_image, status], every=2)
+    
     demo.launch()
 
 if __name__ == "__main__":