app.py

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