Update app.py

app.py

@@ -9,84 +9,88 @@ 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]
+    source1 = open(image1, "rb").read()
+    source2 = open(image2, "rb").read()
+    source3 = open(image3, "rb").read()
+    source4 = open(image4, "rb").read()
+    
+    params = {
+        "prompt": prompt,
+        "image_prompts": [
+            {
+                "cn_img": base64.b64encode(source1).decode('utf-8'),
+                "cn_stop": 1,
+                "cn_weight": 1,
+                "cn_type": "ImagePrompt"
+            },{
+                "cn_img": base64.b64encode(source2).decode('utf-8'),
+                "cn_stop": 1,
+                "cn_weight": 1,
+                "cn_type": "ImagePrompt"
+            },{
+                "cn_img": base64.b64encode(source3).decode('utf-8'),
+                "cn_stop": 1,
+                "cn_weight": 1,
+                "cn_type": "ImagePrompt"
+            },{
+                "cn_img": base64.b64encode(source4).decode('utf-8'),
+                "cn_stop": 1,
+                "cn_weight": 1,
+                "cn_type": "ImagePrompt"
+            }
+        ],
+        "async_process": True
+    }
+    
+    session = requests.Session()
+    retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504])
+    session.mount('http://', HTTPAdapter(max_retries=retries))
 
-        # Prepare the payload with all the image prompts
-        params = {
-            "prompt": prompt,
-            "image_prompts": [
-                {
-                    "cn_img": encoded_img,
-                    "cn_stop": 1,
-                    "cn_weight": 1,
-                    "cn_type": "ImagePrompt"
-                } for encoded_img in encoded_images
-            ],
-            "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 can be adjusted as needed
-        )
-        response.raise_for_status()  # Ensure we catch any HTTP errors
-        result = response.json()
-        job_id = result.get('job_id')
-
-        if job_id:
-            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()  # 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")
+    response = session.post(
+        url=f"{host}/v2/generation/text-to-image-with-ip",
+        data=json.dumps(params),
+        headers={"Content-Type": "application/json"},
+        timeout=10  # Increase timeout as needed
+    )
+    result = response.json()
+    
+    job_id = result.get('job_id')
+    if job_id:
+        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)  # Increase timeout as needed
+            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!"
+            # Update image and status in real-time
+            if job_stage == "RUNNING" and job_step_preview:
+                image = Image.open(BytesIO(base64.b64decode(job_step_preview)))
+                yield image, f"Job is running. Stage: {job_stage}"
 
-                elif job_stage == "SUCCESS":
-                    final_image_url = 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)
-                        image = Image.open(BytesIO(image_response.content))
-                        yield image, "Job completed successfully. Enjoy your masterpiece!"
-                        break
-                    else:
-                        yield None, "Final image URL not found. Something went amiss."
-                        break
-
-                elif job_stage == "FAILED":
-                    yield None, "Job failed. Let's check the parameters and try again."
+            elif job_stage == "SUCCESS":
+                final_image_url = 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))
+                    yield image, "Job completed successfully."
                     break
+                else:
+                    yield None, "Final image URL not found in the job data."
+                    break
+            elif job_stage == "FAILED":
+                yield None, "Job failed."
+                break
 
-                time.sleep(2)  # Pause for 2 seconds before checking again
-
-        else:
-            yield None, "Job ID not found. Did we miss something in the setup?"
-
-    except Exception as e:
-        yield None, f"An error occurred: {str(e)}. We'll need to debug this."
+            time.sleep(2)  # Wait 2 seconds before the next update
+    else:
+        yield None, "Job ID not found."
 
 def gradio_app():
     with gr.Blocks() as demo: