Update app.py

app.py

@@ -12,88 +12,76 @@ 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()
-    
+    session = requests.Session()
+    retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504])
+    session.mount('http://', HTTPAdapter(max_retries=retries))
+
+    # Read and encode images
+    image_paths = [image1, image2, image3, image4]
+    image_data = [
+        {
+            "cn_img": base64.b64encode(open(image_path, "rb").read()).decode('utf-8'),
+            "cn_stop": 1,
+            "cn_weight": 1,
+            "cn_type": "ImagePrompt"
+        } for image_path in image_paths if image_path
+    ]
+
     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"
-            }
-        ],
+        "image_prompts": image_data,
         "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  # Increase timeout as needed
+        timeout=10
     )
+
     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")
-
-            # Continuously print out or display the attributes every 2 seconds
-            print(f"Job Stage: {job_stage}")
-            if job_step_preview:
-                print("Step Preview Available")
-                image = Image.open(BytesIO(base64.b64decode(job_step_preview)))
-            if job_result:
-                print(f"Job Result: {job_result}")
-                
-            if 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))
-                    return image, "Job completed successfully."
-                else:
-                    return None, "Final image URL not found in the job data."
-            elif job_stage == "RUNNING":
-                if job_step_preview:
-                    image = Image.open(BytesIO(base64.b64decode(job_step_preview)))
-                time.sleep(2)  # Wait 2 seconds before the next update
-            elif job_stage == "FAILED":
-                return None, "Job failed."
-    else:
+
+    if not job_id:
         return None, "Job ID not found."
 
+    # Polling for job status
+    start_time = time.time()
+    max_wait_time = 300  # 5 minutes max wait time
+    while time.time() - start_time < max_wait_time:
+        query_url = f"{host}/v1/generation/query-job?job_id={job_id}&require_step_preview=true"
+        response = session.get(query_url, timeout=10)
+        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 there is a step preview, display it
+        if job_step_preview:
+            step_image = Image.open(BytesIO(base64.b64decode(job_step_preview)))
+            return step_image, "Processing..."  # Update the gr.Image widget with step preview
+
+        # If the job is completed successfully, display the final image
+        if 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)
+                final_image = Image.open(BytesIO(image_response.content))
+                return final_image, "Job completed successfully."
+            return None, "Final image URL not found in the job data."
+
+        # If the job failed
+        elif job_stage == "FAILED":
+            return None, "Job failed."
+
+        # If the job is still running, continue polling
+        time.sleep(2)
+
+    return None, "Job timed out."
+
 def gradio_app():
     with gr.Blocks() as demo:
         prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here")