update

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import numpy as np
-import random
-import torch
+
 import requests
 from PIL import Image
 from io import BytesIO
@@ -11,44 +10,68 @@ import time
 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"
+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')
+    image = Image.open(image).convert("RGB")
     return image
 
+
 def image_to_bytes(image):
     byte_io = BytesIO()
-    image.save(byte_io, format='PNG')
+    image.save(byte_io, format="PNG")
     byte_data = byte_io.getvalue()
     return byte_data
 
-def send_request(url, api_key, file,data):
+
+def send_request(url, api_key, file, data):
     response = requests.post(
         url,
-        headers={
-            "Authorization": f"Bearer {api_key}",
-            "Accept": "image/*"
-        },
+        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):
+
+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,
+    preset,
+    api_key,
+):
     if randomize_seed:
         seed = 0
 
@@ -58,7 +81,7 @@ def generate(prompt, negative_prompt, seed, mode, submode, input_image, mask, CN
         "negative_prompt": negative_prompt,
         "output_format": "png",
         "seed": seed,
-        "aspect_ratio": aspect
+        "aspect_ratio": aspect,
     }
     if input_image is not None:
         file["image"] = image_to_bytes(input_image)
@@ -67,14 +90,22 @@ def generate(prompt, negative_prompt, seed, mode, submode, input_image, mask, CN
 
     if mode == "Generate":
         file["none"] = ""
-        if model == "ImageUltra":
-            url = model_url[model]
-        elif model == "ImageCore":
-            url = model_url[model]
+        if model == "Ultra":
+            url = model_url["ImageUltra"]
+        elif model == "Core":
+            url = model_url["ImageCore"]
             data["style_preset"] = preset
-        elif model == "StableDiffusion3":
-            url = model_url[model]
-            data["model"] = sd3_model
+
+        elif model == "sd3-medium":
+            url = model_url["StableDiffusion3"]
+            data["model"] = model
+        elif model == "sd3-large":
+            url = model_url["StableDiffusion3"]
+            data["model"] = model
+        elif model == "sd3-large-turbo":
+            url = model_url["StableDiffusion3"]
+            data["model"] = model
+
         else:
             raise ValueError("Invalid model type")
 
@@ -121,10 +152,7 @@ def generate(prompt, negative_prompt, seed, mode, submode, input_image, mask, CN
             while True:
                 result_response = requests.get(
                     result_url,
-                    headers={
-                        'accept': "image/*",
-                        'authorization': f"Bearer {api_key}"
-                    }
+                    headers={"accept": "image/*", "authorization": f"Bearer {api_key}"},
                 )
                 if result_response.status_code == 202:
                     print("Generation in-progress, try again in 10 seconds.")
@@ -144,13 +172,15 @@ def generate(prompt, negative_prompt, seed, mode, submode, input_image, mask, CN
             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="""
+css = """
 #col-container {
     margin: 0 auto;
     max-width: 50vw;
@@ -158,6 +188,13 @@ css="""
 """
 
 
+def update_style_visibility(model):
+    if model == "Core":
+        return gr.update(visible=True)
+    else:
+        return gr.update(visible=False)
+
+
 def update_mode(mode):
     submode_update = gr.update(choices=["None"], visible=False)
     image_label_update = gr.update(visible=False)
@@ -167,20 +204,35 @@ def update_mode(mode):
     if mode == "Generate":
         submode_update = gr.update(visible=False)
     elif mode == "Upscale":
-        submode_update = gr.update(choices=["Conservative", "Creative"],value="Conservative" ,visible=True)
+        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)
+        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)
+        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)
@@ -194,7 +246,6 @@ def update_submode(submode):
         if submode == "Outpaint":
             outpaint = gr.update(visible=True)
 
-
         elif submode == "Control":
             cn = gr.update(visible=True)
 
@@ -204,22 +255,39 @@ def update_submode(submode):
     return mask, outpaint, cn, search_prompt
 
 
-
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(css=css, theme="NoCrypt/miku") as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""
+        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)
+            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")
+            model = gr.Dropdown(
+                label="Model",
+                choices=["Ultra", "Core", "sd3-medium", "sd3-large", "sd3-large-turbo"],
+                value="Ultra",
+            )
+            mode = gr.Dropdown(
+                label="Mode",
+                choices=["Generate", "Upscale", "Edit", "Control"],
+                value="Generate",
+            )
 
-        submode = gr.Dropdown(label="Submode", choices=["None"], visible=False, value="None")
+        submode = gr.Dropdown(
+            label="Submode", choices=["None"], visible=False, value="None"
+        )
 
         with gr.Row():
             with gr.Column():
@@ -242,11 +310,28 @@ with gr.Blocks(css=css) as demo:
 
         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)
+                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)
+                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%")
@@ -264,39 +349,113 @@ with gr.Blocks(css=css) as demo:
                 step=1,
                 value=0,
             )
-            CN_strength = gr.Slider(label="Control Strength", minimum=0, maximum=1, step=0.01, value=0.5, visible=False)
+            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")
+            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",
+            )
+            with gr.Row(visible=False) as style:
+                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.",
+                )
             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]
-        )
+                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)
+            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]
+        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,
+            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()
+    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],
+    )
+    model.change(fn=update_style_visibility, inputs=model, outputs=style)
+demo.launch()