Spaces:
Sleeping
Sleeping
| 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 | |
| # Your host address - this will need to scale in the future | |
| host = "http://18.119.36.46:8888" | |
| def image_prompt(prompt, image1, image2, image3, image4): | |
| try: | |
| # Reading image files and encoding them | |
| image_sources = [open(image, "rb").read() for image in [image1, image2, image3, image4]] | |
| encoded_images = [base64.b64encode(img).decode('utf-8') for img in image_sources] | |
| # Prepare the payload with all the image prompts | |
| params = { | |
| "prompt": prompt, | |
| "image_prompts": [ | |
| { | |
| "cn_img": encoded_img, | |
| "cn_stop": 1, | |
| "cn_weight": 1, | |
| "cn_type": "ImagePrompt" | |
| } for encoded_img in encoded_images | |
| ], | |
| "async_process": True | |
| } | |
| # Setup retry strategy for robust request handling | |
| session = requests.Session() | |
| retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504]) | |
| session.mount('http://', HTTPAdapter(max_retries=retries)) | |
| # Initiating the job | |
| response = session.post( | |
| url=f"{host}/v2/generation/text-to-image-with-ip", | |
| data=json.dumps(params), | |
| headers={"Content-Type": "application/json"}, | |
| timeout=10 # Timeout can be adjusted as needed | |
| ) | |
| response.raise_for_status() # Ensure we catch any HTTP errors | |
| 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) | |
| response.raise_for_status() # Catch any issues with querying the job | |
| 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") | |
| # Real-time update to image and status | |
| if job_stage == "RUNNING" and job_step_preview: | |
| image = Image.open(BytesIO(base64.b64decode(job_step_preview))) | |
| yield image, f"Job is running smoothly. Current stage: {job_stage}. Hang tight!" | |
| elif 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) | |
| image = Image.open(BytesIO(image_response.content)) | |
| yield image, "Job completed successfully. Enjoy your masterpiece!" | |
| break | |
| else: | |
| yield None, "Final image URL not found. Something went amiss." | |
| break | |
| elif job_stage == "FAILED": | |
| yield None, "Job failed. Let's check the parameters and try again." | |
| break | |
| time.sleep(2) # Pause for 2 seconds before checking again | |
| else: | |
| yield None, "Job ID not found. Did we miss something in the setup?" | |
| except Exception as e: | |
| yield None, f"An error occurred: {str(e)}. We'll need to debug this." | |
| 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(fn=image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status], stream=True) | |
| demo.launch() | |
| if __name__ == "__main__": | |
| gradio_app() |