Update app.py

app.py

@@ -1,17 +1,35 @@
+"""
+{
+    "inputs": {
+        "prompt": "A text prompt to generate the image from.",
+        "image1": "Filepath for the first image prompt.",
+        "image2": "Filepath for the second image prompt.",
+        "image3": "Filepath for the third image prompt.",
+        "image4": "Filepath for the fourth image prompt."
+    },
+    "outputs": {
+        "output_image": "The final generated image or step preview.",
+        "status": "Current status of the image generation process (e.g., Running, Success, Failed)."
+    },
+    "description": "This Gradio app takes a text prompt and up to four image prompts to generate an AI-powered image. It provides real-time updates on the job's progress and displays the final image once complete."
+}
+"""
+
 import requests
 from requests.adapters import HTTPAdapter
 from requests.packages.urllib3.util.retry import Retry
 import json
 import base64
-import time
 import gradio as gr
 from PIL import Image
 from io import BytesIO
 import os
 
+# Host server configuration
 host = "http://18.119.36.46:8888"
 
 def image_prompt(prompt, image1, image2, image3, image4):
+    """Handles the image generation process by sending the text and image prompts to the server."""
     try:
         with open(image1, "rb") as img1, open(image2, "rb") as img2, open(image3, "rb") as img3, open(image4, "rb") as img4:
             source1 = img1.read()
@@ -55,9 +73,9 @@ def image_prompt(prompt, image1, image2, image3, image4):
             url=f"{host}/v2/generation/text-to-image-with-ip",
             data=json.dumps(params),
             headers={"Content-Type": "application/json"},
-            timeout=15  # Slightly increased timeout
+            timeout=15
         )
-        response.raise_for_status()  # Raise exception for HTTP errors
+        response.raise_for_status()  # Ensure we raise errors for bad HTTP status codes
         result = response.json()
         
         job_id = result.get('job_id')
@@ -74,6 +92,7 @@ def image_prompt(prompt, image1, image2, image3, image4):
         return None, f"An unexpected error occurred: {str(e)}"
 
 def query_status(job_id):
+    """Polls the server to get the current status of the image generation job."""
     try:
         session = requests.Session()
         retries = Retry(total=5, backoff_factor=2, status_forcelist=[502, 503, 504])
@@ -81,7 +100,7 @@ def query_status(job_id):
         
         query_url = f"http://18.119.36.46:8888/v1/generation/query-job?job_id={job_id}&require_step_preview=true"
         response = session.get(query_url, timeout=15)
-        response.raise_for_status()  # Raise exception for HTTP errors
+        response.raise_for_status()
         job_data = response.json()
         
         job_stage = job_data.get("job_stage")
@@ -91,7 +110,7 @@ def query_status(job_id):
             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=15)
-                image_response.raise_for_status()  # Raise exception for HTTP errors
+                image_response.raise_for_status()
                 image = Image.open(BytesIO(image_response.content))
                 return image, "Job completed successfully."
             else:
@@ -115,20 +134,23 @@ def query_status(job_id):
         return None, f"An unexpected error occurred: {str(e)}"
 
 def gradio_app():
+    """Defines the Gradio app layout and functionality."""
     with gr.Blocks() as demo:
-        prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here")
+        gr.Markdown("# AI Image Generator\nEnter a text prompt and upload up to 4 images to generate a unique AI-powered image.")
+        
+        prompt = gr.Textbox(label="Text Prompt", placeholder="Describe the image you want to generate")
         with gr.Row():
             image1 = gr.Image(label="Image Prompt 1", type="filepath")
             image2 = gr.Image(label="Image Prompt 2", type="filepath")
             image3 = gr.Image(label="Image Prompt 3", type="filepath")
             image4 = gr.Image(label="Image Prompt 4", type="filepath")
-        output_image = gr.Image(label="Generated Image")
-        status = gr.Textbox(label="Status")
-        
-        job_id = gr.State()
         
         generate_button = gr.Button("Generate Image")
-        generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=job_id)
+        
+        output_image = gr.Image(label="Generated Image", visible=False)
+        status = gr.Textbox(label="Status", value="Awaiting input...", interactive=False)
+        
+        job_id = gr.State()
         
         def update_image(job_id):
             if job_id:
@@ -136,6 +158,7 @@ def gradio_app():
                 return image, status_text
             return None, "No job ID found."
         
+        generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[job_id])
         demo.load(update_image, inputs=job_id, outputs=[output_image, status], every=2)
     
     demo.launch()