init commit

app.py

@@ -1,42 +1,149 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import DiffusionPipeline
 import torch
+import requests
+from PIL import Image
+from io import BytesIO
+import time
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# TODO: PNGinfo追加, 使用データの展開(コピーボタン)
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1344
 
-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)
+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"}
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+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 infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+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 = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
-    
-    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] 
-    
-    return image
+        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",
@@ -46,46 +153,110 @@ examples = [
 css="""
 #col-container {
     margin: 0 auto;
-    max-width: 520px;
+    max-width: 50vw;
 }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
+
+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"""
-        # Text-to-Image Gradio Template
-        Currently running on {power_device}.
+        # 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():
-            
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                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")
+
+        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)
-        
-        result = gr.Image(label="Result", show_label=False)
+
+        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",
-                visible=False,
             )
-            
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -93,54 +264,39 @@ 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)
             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,
-                )
-        
+                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]
+            examples=examples,
+            inputs=[prompt]
         )
 
-    run_button.click(
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
+        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]
     )
 
-demo.queue().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])
+
+demo.launch()

requirements.txt

@@ -3,4 +3,7 @@ diffusers
 invisible_watermark
 torch
 transformers
-xformers
+xformers
+numpy==2.0.0
+Pillow==10.3.0
+Requests==2.32.3