app.py

"""
{
    "inputs": {
        "prompt": "A text prompt to generate the image from.",
        "image_prompts": [
            {
                "cn_img": "Base64 encoded image data for the first image prompt.",
                "cn_stop": "ControlNet stop value for the first image prompt.",
                "cn_weight": "ControlNet weight value for the first image prompt.",
                "cn_type": "Type of the first image prompt."
            },
            {
                "cn_img": "Base64 encoded image data for the second image prompt.",
                "cn_stop": "ControlNet stop value for the second image prompt.",
                "cn_weight": "ControlNet weight value for the second image prompt.",
                "cn_type": "Type of the second image prompt."
            },
            {
                "cn_img": "Base64 encoded image data for the third image prompt.",
                "cn_stop": "ControlNet stop value for the third image prompt.",
                "cn_weight": "ControlNet weight value for the third image prompt.",
                "cn_type": "Type of the third image prompt."
            },
            {
                "cn_img": "Base64 encoded image data for the fourth image prompt.",
                "cn_stop": "ControlNet stop value for the fourth image prompt.",
                "cn_weight": "ControlNet weight value for the fourth image prompt.",
                "cn_type": "Type of the fourth image prompt."
            }
        ],
        "async_process": "Boolean to indicate if the process should be asynchronous."
    },
    "outputs": {
        "job_id": "The ID of the job submitted for image generation.",
        "job_stage": "The current stage of the job (e.g., RUNNING, SUCCESS, FAILED).",
        "final_image_url": "The URL of the final generated image.",
        "step_preview": "Base64 encoded step preview image data."
    }
}
"""

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
import os

host = "http://18.119.36.46:8888"

# 📂 Get the directory where the script is located
script_dir = os.path.dirname(os.path.abspath(__file__))

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))

    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"http://18.119.36.46:8888/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
                    with open("output.png", "wb") as f:
                        f.write(image_response.content)
                    return "output.png", "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:
                    with open("output.png", "wb") as f:
                        f.write(base64.b64decode(step_preview_base64))
                time.sleep(5)
            elif job_stage == "FAILED":
                return None, "Job failed."
    else:
        return None, "Job ID not found."

def create_status_image():
    if os.path.exists("output.png"):
        return "output.png"
    else:
        return None

def gradio_app():
    with gr.Blocks() as demo:
        prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here")
        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")
        
        generate_button = gr.Button("Generate Image")
        generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status])
        
        # 🖼️ Display the status image
        status_image = gr.Image(label="Queue Status", interactive=False)
        
        # ⏲️ Update the image every 5 seconds
        demo.load(create_status_image, every=5, outputs=status_image)
    
    demo.launch()

if __name__ == "__main__":
    gradio_app()