Spaces:
Running
Running
| import gradio as gr | |
| import numpy as np | |
| import random | |
| import torch | |
| import requests | |
| from PIL import Image | |
| from io import BytesIO | |
| import time | |
| # TODO: PNGinfo追加, 使用データの展開(コピーボタン) | |
| MAX_SEED = np.iinfo(np.int32).max | |
| MAX_IMAGE_SIZE = 1344 | |
| model_url = { "ImageUltra":"https://api.stability.ai/v2beta/stable-image/generate/ultra", | |
| "ImageCore":"https://api.stability.ai/v2beta/stable-image/generate/core", | |
| "StableDiffusion3": "https://api.stability.ai/v2beta/stable-image/generate/sd3"} | |
| service_url = { "Conservative_Upscale":"https://api.stability.ai/v2beta/stable-image/upscale/conservative", | |
| "Creative_Upscale":"https://api.stability.ai/v2beta/stable-image/upscale/creative", | |
| "Erase":"https://api.stability.ai/v2beta/stable-image/edit/erase", | |
| "Inpaint":"https://api.stability.ai/v2beta/stable-image/edit/inpaint", | |
| "Outpaint":"https://api.stability.ai/v2beta/stable-image/edit/outpaint", | |
| "SR":"https://api.stability.ai/v2beta/stable-image/edit/search-and-replace", | |
| "RMBG":"https://api.stability.ai/v2beta/stable-image/edit/remove-background", | |
| "Sketch":"https://api.stability.ai/v2beta/stable-image/control/sketch", | |
| "Structure":"https://api.stability.ai/v2beta/stable-image/control/structure" | |
| } | |
| def bytes_to_image(image): | |
| image = BytesIO(image) | |
| image = Image.open(image).convert('RGB') | |
| return image | |
| def image_to_bytes(image): | |
| byte_io = BytesIO() | |
| image.save(byte_io, format='PNG') | |
| byte_data = byte_io.getvalue() | |
| return byte_data | |
| def send_request(url, api_key, file,data): | |
| response = requests.post( | |
| url, | |
| headers={ | |
| "Authorization": f"Bearer {api_key}", | |
| "Accept": "image/*" | |
| }, | |
| files=file, | |
| data=data, | |
| ) | |
| return response | |
| def generate(prompt, negative_prompt, seed, mode, submode, input_image, mask, CNstrength, search_prompt, op_left, op_right, op_up, op_down, randomize_seed, aspect, model, sd3_model, preset, api_key): | |
| if randomize_seed: | |
| seed = 0 | |
| file = {} | |
| data = { | |
| "prompt": prompt, | |
| "negative_prompt": negative_prompt, | |
| "output_format": "png", | |
| "seed": seed, | |
| "aspect_ratio": aspect | |
| } | |
| if input_image is not None: | |
| file["image"] = image_to_bytes(input_image) | |
| if mask is not None: | |
| file["mask"] = image_to_bytes(mask) | |
| if mode == "Generate": | |
| file["none"] = "" | |
| if model == "ImageUltra": | |
| url = model_url[model] | |
| elif model == "ImageCore": | |
| url = model_url[model] | |
| data["style_preset"] = preset | |
| elif model == "StableDiffusion3": | |
| url = model_url[model] | |
| data["model"] = sd3_model | |
| else: | |
| raise ValueError("Invalid model type") | |
| elif mode == "Upscale": | |
| if submode == "Conservative": | |
| url = service_url["Conservative_Upscale"] | |
| elif submode == "Creative": | |
| url = service_url["Creative_Upscale"] | |
| elif mode == "Edit": | |
| if submode == "Erase": | |
| url = service_url["Erase"] | |
| elif submode == "Inpaint": | |
| url = service_url["Inpaint"] | |
| elif submode == "Outpaint": | |
| url = service_url["Outpaint"] | |
| data["left"] = op_left | |
| data["right"] = op_right | |
| data["up"] = op_up | |
| data["down"] = op_down | |
| elif submode == "Search and Replace": | |
| url = service_url["SR"] | |
| data["search_prompt"] = search_prompt | |
| elif submode == "Remove Background": | |
| url = service_url["RMBG"] | |
| elif mode == "Control": | |
| data["control_strength"] = CNstrength | |
| if submode == "Sketch": | |
| url = service_url["Sketch"] | |
| elif submode == "Structure": | |
| url = service_url["Structure"] | |
| response = send_request(url, api_key, file, data) | |
| if response.status_code == 200: | |
| if mode == "Upscale" and submode == "Creative": | |
| generation_id = response.json().get("id") | |
| if not generation_id: | |
| raise Exception("No generation ID returned for creative upscale") | |
| # Polling for the result | |
| result_url = f"https://api.stability.ai/v2beta/stable-image/upscale/creative/result/{generation_id}" | |
| while True: | |
| result_response = requests.get( | |
| result_url, | |
| headers={ | |
| 'accept': "image/*", | |
| 'authorization': f"Bearer {api_key}" | |
| } | |
| ) | |
| if result_response.status_code == 202: | |
| print("Generation in-progress, try again in 10 seconds.") | |
| time.sleep(10) | |
| elif result_response.status_code == 200: | |
| print("Generation complete!") | |
| image = result_response.content | |
| image = bytes_to_image(image) | |
| copy_filed_value = f"prompt:{prompt}, negative:{negative_prompt}, mode:{mode}, submode:{submode}" | |
| return image, seed, copy_filed_value | |
| else: | |
| raise Exception(str(result_response.json())) | |
| else: | |
| image = response.content | |
| image = bytes_to_image(image) | |
| copy_filed_value = f"prompt:{prompt}, negative:{negative_prompt}, mode:{mode}, submode:{submode}" | |
| return image, seed, copy_filed_value | |
| else: | |
| raise Exception(str(response.json())) | |
| 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: 50vw; | |
| } | |
| """ | |
| def update_mode(mode): | |
| submode_update = gr.update(choices=["None"], visible=False) | |
| image_label_update = gr.update(visible=False) | |
| img_input_update = gr.update(visible=False) | |
| mask_update = gr.update(visible=False) | |
| if mode == "Generate": | |
| submode_update = gr.update(visible=False) | |
| elif mode == "Upscale": | |
| submode_update = gr.update(choices=["Conservative", "Creative"],value="Conservative" ,visible=True) | |
| img_input_update = gr.update(visible=True) | |
| image_label_update = gr.update(visible=True) | |
| elif mode == "Edit": | |
| submode_update = gr.update(choices=["Erase", "Inpaint", "Outpaint", "Search and Replace", "Remove Background"],value="Erase", visible=True) | |
| img_input_update = gr.update(visible=True) | |
| image_label_update = gr.update(visible=True) | |
| elif mode == "Control": | |
| submode_update = gr.update(choices=["Sketch", "Structure"],value="Sketch", visible=True) | |
| img_input_update = gr.update(visible=True) | |
| image_label_update = gr.update(visible=True) | |
| return submode_update, img_input_update, mask_update, image_label_update | |
| def update_submode(submode): | |
| mask = gr.update(visible=False) | |
| outpaint = gr.update(visible=False) | |
| cn = gr.update(visible=False) | |
| search_prompt = gr.update(visible=False) | |
| if submode in ["Erase", "Inpaint"]: | |
| mask = gr.update(visible=True) | |
| else: | |
| if submode == "Outpaint": | |
| outpaint = gr.update(visible=True) | |
| elif submode == "Control": | |
| cn = gr.update(visible=True) | |
| elif submode == "Search and Replace": | |
| search_prompt = gr.update(visible=True) | |
| return mask, outpaint, cn, search_prompt | |
| with gr.Blocks(css=css) as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown(f""" | |
| # Demo Stable Image API | |
| Learn more about the [Stable Diffusion 3 series](https://stability.ai/news/stable-diffusion-3). Try on [Stability AI API](https://platform.stability.ai/docs/api-reference#tag/Generate/paths/~1v2beta~1stable-image~1generate~1sd3/post), [Stable Assistant](https://stability.ai/stable-assistant), or on Discord via [Stable Artisan](https://stability.ai/stable-artisan). Run locally with [ComfyUI](https://github.com/comfyanonymous/ComfyUI) or [diffusers](https://github.com/huggingface/diffusers) | |
| """) | |
| with gr.Row(): | |
| api_key = gr.Text(label="API Key", type="password", placeholder="Enter your API key", max_lines=1, container=False) | |
| with gr.Row(): | |
| model = gr.Dropdown(label="Model", choices=["ImageUltra", "ImageCore", "StableDiffusion3"], value="ImageUltra") | |
| mode = gr.Dropdown(label="Mode", choices=["Generate", "Upscale", "Edit", "Control"], value="Generate") | |
| submode = gr.Dropdown(label="Submode", choices=["None"], visible=False, value="None") | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Text( | |
| label="Prompt", | |
| show_label=False, | |
| max_lines=1, | |
| placeholder="Enter your prompt", | |
| container=False, | |
| ) | |
| search_prompt = gr.Text( | |
| label="search prompt", | |
| visible=False, | |
| show_label=False, | |
| max_lines=1, | |
| placeholder="Enter a search prompt", | |
| ) | |
| run_button = gr.Button("Run", scale=0) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_label = gr.Markdown(value = "input image",visible=False) | |
| image = gr.Image(type='pil',label="img input", width="20vw", height="20vw",show_label=True,visible=False, interactive=True, container=False) | |
| with gr.Column(visible=False) as mask: | |
| mask_label = gr.Markdown(value="input mask") | |
| mask_input = gr.Image(type='pil',label="mask", width="20vw", height="20vw", show_label=True, interactive=True, container=False) | |
| with gr.Row(): | |
| result = gr.Image(label="Result", width="20vw", height="20%") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| negative_prompt = gr.Text( | |
| label="Negative prompt", | |
| max_lines=1, | |
| placeholder="Enter a negative prompt", | |
| ) | |
| seed = gr.Slider( | |
| label="Seed", | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=0, | |
| ) | |
| CN_strength = gr.Slider(label="Control Strength", minimum=0, maximum=1, step=0.01, value=0.5, visible=False) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| with gr.Row(): | |
| aspect = gr.Radio(choices=["1:1", "16:9", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"], label="Aspect raito", value="1:1") | |
| style_preset = gr.Radio(choices=["3d-model", "analog-film", "anime", "cinematic", "comic-book", "digital-art", "enhance", "fantasy-art", "isometric", "line-art", "low-poly", "modeling-compound", "neon-punk", "origami", "photographic", "pixel-art", "tile-texture"], label="Style_preset", value="anime", info="This parameter is only available for ImageCore model.") | |
| sd3_model = gr.Dropdown(label="SD3 Model Size", choices=["sd3-medium", "sd3-large", "sd3-large-turbo"], value="sd3-medium") | |
| with gr.Row(visible=False) as outpaint_scale: | |
| paint = gr.Markdown(value = "Outpain Scale") | |
| op_left = gr.Slider(label="left", minimum=0, maximum=2000, step=4, value=200) | |
| op_right = gr.Slider(label="right", minimum=0, maximum=2000, step=4, value=200) | |
| op_up = gr.Slider(label="up", minimum=0, maximum=2000, step=4, value=200) | |
| op_down = gr.Slider(label="down", minimum=0, maximum=2000, step=4, value=200) | |
| gr.Examples( | |
| examples=examples, | |
| inputs=[prompt] | |
| ) | |
| copy_filed = gr.TextArea( | |
| value="", | |
| label="Copy Field", | |
| max_lines=1, | |
| placeholder="Copy the field", | |
| show_copy_button=True, | |
| container=False) | |
| gr.on( | |
| triggers=[run_button.click, prompt.submit, negative_prompt.submit], | |
| fn=generate, | |
| inputs=[prompt, negative_prompt, seed, mode, submode, image, mask_input, CN_strength, search_prompt, op_left, op_right, op_up, op_down, randomize_seed, aspect, model, sd3_model, style_preset, api_key], | |
| outputs=[result, seed, copy_filed] | |
| ) | |
| mode.change(fn=update_mode, inputs=mode, outputs=[submode, image, mask, image_label]) | |
| submode.change(fn=update_submode, inputs=submode, outputs=[mask,outpaint_scale,CN_strength,search_prompt]) | |
| demo.launch() |