Update app.py

app.py

@@ -1,146 +1,162 @@
-import gradio as gr
-import numpy as np
-import random
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
+"""
+{
+    "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."
+    }
+}
+"""
 
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
+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
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+host = "http://18.119.36.46:8888"
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+# 📂 Get the directory where the script is located
+script_dir = os.path.dirname(os.path.abspath(__file__))
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
+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()
     
-    image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
+    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
+    }
     
-    return image
+    session = requests.Session()
+    retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504])
+    session.mount('http://', HTTPAdapter(max_retries=retries))
 
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 520px;
-}
-"""
-
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
-with gr.Blocks(css=css) as demo:
+    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()
     
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""
-        # Text-to-Image Gradio Template
-        Currently running on {power_device}.
-        """)
-        
-        with gr.Row():
+    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()
             
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
+            job_stage = job_data.get("job_stage")
             
-            run_button = gr.Button("Run", scale=0)
-        
-        result = gr.Image(label="Result", show_label=False)
+            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."
 
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-            
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
-            with gr.Row():
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-            
-            with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,
-                )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=12,
-                    step=1,
-                    value=2,
-                )
-        
-        gr.Examples(
-            examples = examples,
-            inputs = [prompt]
-        )
+def create_status_image():
+    if os.path.exists("output.png"):
+        return "output.png"
+    else:
+        return None
 
-    run_button.click(
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
-    )
+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()
 
-demo.queue().launch()
+if __name__ == "__main__":
+    gradio_app()