Update app.py

app.py

@@ -7,14 +7,18 @@ import time
 import gradio as gr
 from PIL import Image
 from io import BytesIO
+import os
 
+# Your host address - this will need to scale in the future
 host = "http://18.119.36.46:8888"
 
 def image_prompt(prompt, image1, image2, image3, image4):
     try:
+        # Reading image files and encoding them
         image_sources = [open(image, "rb").read() for image in [image1, image2, image3, image4]]
         encoded_images = [base64.b64encode(img).decode('utf-8') for img in image_sources]
 
+        # Prepare the payload with all the image prompts
         params = {
             "prompt": prompt,
             "image_prompts": [
@@ -28,17 +32,19 @@ def image_prompt(prompt, image1, image2, image3, image4):
             "async_process": True
         }
         
+        # Setup retry strategy for robust request handling
         session = requests.Session()
         retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504])
         session.mount('http://', HTTPAdapter(max_retries=retries))
 
+        # Initiating the job
         response = session.post(
             url=f"{host}/v2/generation/text-to-image-with-ip",
             data=json.dumps(params),
             headers={"Content-Type": "application/json"},
-            timeout=10
+            timeout=10  # Timeout can be adjusted as needed
         )
-        response.raise_for_status()
+        response.raise_for_status()  # Ensure we catch any HTTP errors
         result = response.json()
         job_id = result.get('job_id')
 
@@ -46,13 +52,14 @@ def image_prompt(prompt, image1, image2, image3, image4):
             while True:
                 query_url = f"{host}/v1/generation/query-job?job_id={job_id}&require_step_preview=true"
                 response = session.get(query_url, timeout=10)
-                response.raise_for_status()
+                response.raise_for_status()  # Catch any issues with querying the job
                 job_data = response.json()
 
                 job_stage = job_data.get("job_stage")
                 job_step_preview = job_data.get("job_step_preview")
                 job_result = job_data.get("job_result")
 
+                # Real-time update to image and status
                 if job_stage == "RUNNING" and job_step_preview:
                     image = Image.open(BytesIO(base64.b64decode(job_step_preview)))
                     yield image, f"Job is running smoothly. Current stage: {job_stage}. Hang tight!"
@@ -73,7 +80,7 @@ def image_prompt(prompt, image1, image2, image3, image4):
                     yield None, "Job failed. Let's check the parameters and try again."
                     break
 
-                time.sleep(2)
+                time.sleep(2)  # Pause for 2 seconds before checking again
 
         else:
             yield None, "Job ID not found. Did we miss something in the setup?"