Update app.py

app.py

@@ -68,29 +68,31 @@ def image_prompt(prompt, image1, image2, image3, image4):
             job_step_preview = job_data.get("job_step_preview")
             job_result = job_data.get("job_result")
 
-            # Update image and status in real-time
-            if job_stage == "RUNNING" and job_step_preview:
+            # 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)))
-                yield image, f"Job is running. Stage: {job_stage}"
-
-            elif job_stage == "SUCCESS":
+            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))
-                    yield image, "Job completed successfully."
-                    break
+                    return image, "Job completed successfully."
                 else:
-                    yield None, "Final image URL not found in the job data."
-                    break
+                    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":
-                yield None, "Job failed."
-                break
-
-            time.sleep(2)  # Wait 2 seconds before the next update
+                return None, "Job failed."
     else:
-        yield None, "Job ID not found."
+        return None, "Job ID not found."
 
 def gradio_app():
     with gr.Blocks() as demo:
@@ -104,7 +106,7 @@ def gradio_app():
         status = gr.Textbox(label="Status")
         
         generate_button = gr.Button("Generate Image")
-        generate_button.click(fn=image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status], stream=True)
+        generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status])
         
     demo.launch()