app.py

import spaces
import random
import numpy as np
from PIL import Image

import torch
import torchvision.transforms.functional as F
from diffusers import ControlNetModel, StableDiffusionControlNetPipeline, AutoencoderTiny, AutoencoderKL
import gradio as gr

device = "cuda"
weight_type = torch.float16

controlnet = ControlNetModel.from_pretrained(
    "IDKiro/sdxs-512-dreamshaper-sketch", torch_dtype=weight_type
).to(device)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "IDKiro/sdxs-512-dreamshaper", controlnet=controlnet, torch_dtype=weight_type
)
pipe.to(device)

vae_tiny = AutoencoderTiny.from_pretrained(
    "IDKiro/sdxs-512-dreamshaper", subfolder="vae"
)
vae_tiny.to(device, dtype=weight_type)

vae_large = AutoencoderKL.from_pretrained(
    "IDKiro/sdxs-512-dreamshaper", subfolder="vae_large"
)
vae_tiny.to(device, dtype=weight_type)

style_list = [
    {
        "name": "No Style",
        "prompt": "{prompt}",
    },
    {
        "name": "Cinematic",
        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
    },
    {
        "name": "3D Model",
        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
    },
    {
        "name": "Anime",
        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
    },
    {
        "name": "Digital Art",
        "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
    },
    {
        "name": "Photographic",
        "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
    },
    {
        "name": "Pixel art",
        "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
    },
    {
        "name": "Fantasy art",
        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
    },
    {
        "name": "Neonpunk",
        "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
    },
    {
        "name": "Manga",
        "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
    },
]

styles = {k["name"]: k["prompt"] for k in style_list}
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "No Style"
MAX_SEED = np.iinfo(np.int32).max


def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed


@spaces.GPU
def run(
    image,
    prompt,
    prompt_template,
    style_name,
    controlnet_conditioning_scale,
    vae_type="tiny vae",
    device_type="GPU",
    param_dtype="torch.float16",
):
    if vae_type == "tiny vae":
        pipe.vae = vae_tiny
    elif vae_type == "large vae":
        pipe.vae = vae_large

    if device_type == "CPU":
        device = "cpu"
        param_dtype = "torch.float32"
    else:
        device = "cuda"

    pipe.to(
        torch_device=device,
        torch_dtype=torch.float16 if param_dtype == "torch.float16" else torch.float32,
    )

    print(f"prompt: {prompt}")
    print("sketch updated")
    if image is None:
        ones = Image.new("L", (512, 512), 255)
        return ones
    prompt = prompt_template.replace("{prompt}", prompt)
    control_image = Image.fromarray(255 - np.array(image["composite"])[:, :, -1])

    output_pil = pipe(
        prompt=prompt,
        image=control_image,
        width=512,
        height=512,
        guidance_scale=0.0,
        num_inference_steps=1,
        num_images_per_prompt=1,
        output_type="pil",
        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
    ).images[0]

    return output_pil


with gr.Blocks(theme="monochrome") as demo:
    gr.Markdown("# SDXS-512-DreamShaper-Sketch")
    gr.Markdown(
        "[SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions](https://arxiv.org/abs/2403.16627) | [GitHub](https://github.com/IDKiro/sdxs)"
    )
    with gr.Row():
        with gr.Column():
            gr.Markdown("## INPUT")
            image = gr.Sketchpad(
                type="pil",
                image_mode="RGBA",
                brush=gr.Brush(colors=["#000000"], color_mode="fixed", default_size=8),
                crop_size="1:1",
            )

            prompt = gr.Textbox(label="Prompt", value="", show_label=True)
            with gr.Row():
                style = gr.Dropdown(
                    label="Style",
                    choices=STYLE_NAMES,
                    value=DEFAULT_STYLE_NAME,
                    scale=1,
                )
                prompt_temp = gr.Textbox(
                    label="Prompt Style Template",
                    value=styles[DEFAULT_STYLE_NAME],
                    scale=2,
                    max_lines=1,
                )

            controlnet_conditioning_scale = gr.Slider(
                label="Control Strength", minimum=0, maximum=1, step=0.01, value=0.8
            )

            vae_choices = ["tiny vae", "large vae"]
            vae_type = gr.Radio(
                vae_choices,
                label="Image Decoder Type",
                value=vae_choices[0],
                interactive=True,
                info="To save GPU memory, use tiny vae. For better quality, use large vae.",
            )

            device_choices = ["GPU", "CPU"]
            device_type = gr.Radio(
                device_choices,
                label="Device",
                value=device_choices[0],
                interactive=True,
                info="Many thanks to the community for the GPU!",
            )

            dtype_choices = ["torch.float16", "torch.float32"]
            param_dtype = gr.Radio(
                dtype_choices,
                label="torch.weight_type",
                value=dtype_choices[0],
                interactive=True,
                info="To save GPU memory, use torch.float16. For better quality, use torch.float32.",
            )

        with gr.Column():
            gr.Markdown("## OUTPUT")
            result = gr.Image(
                label="Result",
                show_label=False,
                show_download_button=True,
            )
            run_button = gr.Button("Run")
            gr.Markdown("### Instructions")
            gr.Markdown("**1**. Enter a text prompt (e.g. cat)")
            gr.Markdown("**2**. Start sketching")
            gr.Markdown("**3**. Change the image style using a style template")
            gr.Markdown("**4**. Adjust the effect of sketch guidance using the slider")

    inputs = [
        image,
        prompt,
        prompt_temp,
        style,
        controlnet_conditioning_scale,
        vae_type,
        device_type,
        param_dtype,
    ]
    outputs = [result]

    prompt.submit(fn=run, inputs=inputs, outputs=outputs)
    style.change(lambda x: styles[x], inputs=[style], outputs=[prompt_temp]).then(
        fn=run,
        inputs=inputs,
        outputs=outputs,
    )
    image.change(
        run,
        inputs=inputs,
        outputs=outputs,
    )
    run_button.click(
        run,
        inputs=inputs,
        outputs=outputs,
    )

if __name__ == "__main__":
    demo.queue().launch()