Update app.py

app.py

@@ -7,84 +7,79 @@ import time
 import gradio as gr
 from PIL import Image
 from io import BytesIO
-import os
 
 host = "http://18.119.36.46:8888"
 
 def image_prompt(prompt, image1, image2, image3, image4):
-    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))
+    try:
+        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]
 
-    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")
-            
-            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)))
-                    time.sleep(5)
-                    continue
-            elif job_stage == "FAILED":
-                return None, "Job failed."
-    else:
-        return None, "Job ID not found."
+        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
+        }
+        
+        session = requests.Session()
+        retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504])
+        session.mount('http://', HTTPAdapter(max_retries=retries))
+
+        response = session.post(
+            url=f"{host}/v2/generation/text-to-image-with-ip",
+            data=json.dumps(params),
+            headers={"Content-Type": "application/json"},
+            timeout=10
+        )
+        response.raise_for_status()
+        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()
+                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")
+
+                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!"
+
+                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."
+                    break
+
+                time.sleep(2)
+
+        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."
 
 def gradio_app():
     with gr.Blocks() as demo:
@@ -98,7 +93,7 @@ def gradio_app():
         status = gr.Textbox(label="Status")
         
         generate_button = gr.Button("Generate Image")
-        generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status])
+        generate_button.click(fn=image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status], stream=True)
         
     demo.launch()