app.py

#!/usr/bin/env python

from __future__ import annotations

import os
import random

import gradio as gr
import numpy as np
import PIL.Image
import torch
from diffusers import AutoencoderKL, DiffusionPipeline

DESCRIPTION = "# Segmind Stable Diffusion"
if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"

MAX_SEED = np.iinfo(np.int32).max
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
ENABLE_REFINER = os.getenv("ENABLE_REFINER", "0") == "1"

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
    pipe = DiffusionPipeline.from_pretrained(
        "segmind/SSD-1B",
        vae=vae,
        torch_dtype=torch.float16,
        use_safetensors=True,
        variant="fp16",
    )
    if ENABLE_REFINER:
        refiner = DiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-refiner-1.0",
            vae=vae,
            torch_dtype=torch.float16,
            use_safetensors=True,
            variant="fp16",
        )

    if ENABLE_CPU_OFFLOAD:
        pipe.enable_model_cpu_offload()
        if ENABLE_REFINER:
            refiner.enable_model_cpu_offload()
    else:
        pipe.to(device)
        if ENABLE_REFINER:
            refiner.to(device)

    if USE_TORCH_COMPILE:
        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
        if ENABLE_REFINER:
            refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)


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


def generate(
    prompt: str,
    negative_prompt: str = "",
    prompt_2: str = "",
    negative_prompt_2: str = "",
    use_negative_prompt: bool = False,
    use_prompt_2: bool = False,
    use_negative_prompt_2: bool = False,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale_base: float = 5.0,
    guidance_scale_refiner: float = 5.0,
    num_inference_steps_base: int = 25,
    num_inference_steps_refiner: int = 25,
    apply_refiner: bool = False,
) -> PIL.Image.Image:
    generator = torch.Generator().manual_seed(seed)

    if not use_negative_prompt:
        negative_prompt = None  # type: ignore
    if not use_prompt_2:
        prompt_2 = None  # type: ignore
    if not use_negative_prompt_2:
        negative_prompt_2 = None  # type: ignore

    if not apply_refiner:
        return pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            prompt_2=prompt_2,
            negative_prompt_2=negative_prompt_2,
            width=width,
            height=height,
            guidance_scale=guidance_scale_base,
            num_inference_steps=num_inference_steps_base,
            generator=generator,
            output_type="pil",
        ).images[0]
    else:
        latents = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            prompt_2=prompt_2,
            negative_prompt_2=negative_prompt_2,
            width=width,
            height=height,
            guidance_scale=guidance_scale_base,
            num_inference_steps=num_inference_steps_base,
            generator=generator,
            output_type="latent",
        ).images
        image = refiner(
            prompt=prompt,
            negative_prompt=negative_prompt,
            prompt_2=prompt_2,
            negative_prompt_2=negative_prompt_2,
            guidance_scale=guidance_scale_refiner,
            num_inference_steps=num_inference_steps_refiner,
            image=latents,
            generator=generator,
        ).images[0]
        return image


examples = ['art by Steve McCurry, Afghan Girl, green eyes, worn out clothes, Peshawar, Pakistan, 1984, 8k, uhd, dslr, film grain, insane details, detailed skin, skin imperfections, award-winning photo', '(masterpiece, high quality), concept art of a beautiful ornate mechanical donmf41dr4g0n steampunk owl, looking at viewer, anime owl, sitting on a tree branch, extremely detailed eyes, dark nebula bioluminescence sky, (3d style:0.5), bokeh, ultra realistic, sharp focus, intricate, sharp details, highly detailed, rich color, 8K,', 'front shot, portrait photo of a 25 y.o (afghan woman), (red headscarf), blue eyes, looks away, natural skin, skin moles, cozy interior, (cinematic, film grain:1.1)', 'cinematic photo (art by Mathias Goeritz:0.9) , photograph, Lush Girlfriend, looking at the camera smiling, Rich ginger hair, Winter, tilt shift, Horror, specular lighting, film grain, Samsung Galaxy, F/5, (cinematic still:1.2), freckles . 35mm photograph, film, bokeh, professional, 4k, highly detailed', 'beautiful lady, freckles, big smile, blue eyes, short ginger hair, dark makeup, wearing a floral blue vest top, hyperdetailed photography, sharp focus on face, soft light, dark background, head and shoulders portrait, cover', 'professional portrait photo by Alberto Burri of an anthropomorphic cat wearing fancy hat and jacket walking in autumn forest. cinematic focus on the cat, dynamic pose, dynamic background, dynamic composition, dynamic lighting, realistic proportions, hdr, raytracing, 8k resolution, ultra realistic, photorealistic, extreme detailed, ultra detailed, intricate details, highly detailed atmosphere, highly detailed textures.', 'cowboy, riding a Tiger, badlands with clouds in the background, style of Howard Terpning, rich color palette adds depth to the scene, highly detailed textures and elements capture the essence of a wild west adventure. hyper-detailed, high quality visuals, dim Lighting, ultra-realistic, sharply focused, octane render, 8k UHD,', 'abstract beauty, centered, looking at the camera, approaching perfection, dynamic, moonlight, highly detailed, digital painting, artstation, concept art, smooth, sharp focus, illustration, art by Carne Griffiths and Wadim Kashin', 'a bird in perfect geometrical harmony, intricate and hyperdetailed by Android Jones, Emily Kell, Rossdraws, Andree Wallin, Sung Choi, fluid ink and oil splash painting artistic fantasy art, album cover art, amazing depth, perfect ratio, hdr, 8k resolution, perfect balanced styles, intense light, muted colors,hyperdetailed masterwork by head of prompt engineering', 'symmetrical composition BREAK tangerine tango and ultramarine green color blocking digital painting, torso clay sculpture of an anthropomorphic manticore goddess, low lit darkness, astral, space, rainbows explosion, paint smudges, paint drops, splashes, impasto, maximum details, art by Charlie Bowater and Anna Dittmann and Henry Asencio', 'cinematic film still of Futuristic hero with golden dark armour with machine gun,  muscular body, athletic, shallow depth of field, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy', 'cinematic still, film grain, Concept art of an ancient Gnome male mage, clad in dirty robes that whisper secrets of forgotten spells. His face is etched with wrinkles, gray hair stringy and greasy, hinting at centuries spent honing his arcane craft. He wields a magical staff adorned with runes and orbs of power, exuding an otherworldly aura that chills the blood. The overall atmosphere is grim and foreboding, style of Apterus, haunting illustrations, extremely high-resolution details, photographic, realism pushed to extreme, fine texture, 4k,  ultra-detailed, high quality, high contrast,', 'surreal double exposure portrait of a hyperdetailed woman under the bright Savannah sunlight, by Dave White,Conrad Roset,', "A wolf looking at the camera, with an irridescent coat, neon realism, ocean waves at night, muted colors, the melting beauty of the Universe in surrealism in the mirror reflection of nebula fragments. Alcohol ink. Works by Gustav Klimt, Escher, Jean Baptiste Monge, Alberto Seveso, Camilla D'Errico, Epic scale, super detail. Detailed. , Mysterious", 'ONCE UPON A TIME, THERE WAS A BEAUTIFUL, CUTE, FLUFFY BABY LITTLE BIRD NAMED PEDRO WHO LOVED GOD VERY MUCH and god love him too. tender, glowing eyes, strybk, the Image should capture the heart of the adult viewer', 'STICKER, A detailed illustration a print of vivid winking racoon with glasses,in the form of Harry Potter,fantasy red splash, 2023 t-shirt design,pastel tetradic colors, 3D vector art, cute and quirky, fantasy art, watercolor effect, bokeh, Adobe Illustrator, hand-drawn, digital painting, low-poly, soft lighting, isometric style, retro aesthetic, focused on the character, 4K resolution, photorealistic rendering, using Cinema 4D, sticker, 2d cute, fantasy, dreamy, vector illustration, 2d flat, centered, by Tim Burton, professional, sleek, modern, minimalist, graphic, line art, vector graphics\n \n (Highest Quality, 4k, masterpiece, Amazing Details:1.1), a colorful exotic dragon_bird, feathers, film grain, Fujifilm XT3, photography,\n \n text "silvia " name under cute catrina doll with cute smile sitting with flowers, skulls and jack lanterns and hearts, anime, illustration, fashion, painting, 3d render\n \n with smoke, ice and fire and ultra realistic in detail., typography, dark fantasy, wildlife photography, vibrant, cinematic and the words “common coding” on a black background\n \n Fantasy Art. In an ethereal and grandiose composition, depict a regal mage locked in an epic battle with an immense and ferocious dragon atop a foreboding, shadowy mountain precipice. The brooding sky sets a menacing tone as the mage unleashes a torrent of awe-inspiring elemental magic, while the dragon\'s fiery breath clashes with the mage\'s powerful spellcasting. Sharp focus. Masterpiece.']

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(
        value="Duplicate Space for private use",
        elem_id="duplicate-button",
        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
    )
    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0)
        result = gr.Image(label="Result", show_label=False)
    with gr.Accordion("Advanced options", open=False):
        with gr.Row():
            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
            use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False)
            use_negative_prompt_2 = gr.Checkbox(label="Use negative prompt 2", value=False)
        negative_prompt = gr.Text(
            label="Negative prompt",
            max_lines=1,
            placeholder="Enter a negative prompt",
            visible=False,
        )
        prompt_2 = gr.Text(
            label="Prompt 2",
            max_lines=1,
            placeholder="Enter your prompt",
            visible=False,
        )
        negative_prompt_2 = gr.Text(
            label="Negative prompt 2",
            max_lines=1,
            placeholder="Enter a negative prompt",
            visible=False,
        )

        seed = gr.Slider(
            label="Seed",
            minimum=0,
            maximum=MAX_SEED,
            step=1,
            value=0,
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Row(visible=False):
            width = gr.Slider(
                label="Width",
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=1024,
            )
            height = gr.Slider(
                label="Height",
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=1024,
            )
        apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER)
        with gr.Row():
            guidance_scale_base = gr.Slider(
                label="Guidance scale for base",
                minimum=1,
                maximum=20,
                step=0.1,
                value=5.0,
            )
            num_inference_steps_base = gr.Slider(
                label="Number of inference steps for base",
                minimum=10,
                maximum=100,
                step=1,
                value=25,
            )
        with gr.Row(visible=False) as refiner_params:
            guidance_scale_refiner = gr.Slider(
                label="Guidance scale for refiner",
                minimum=1,
                maximum=20,
                step=0.1,
                value=5.0,
            )
            num_inference_steps_refiner = gr.Slider(
                label="Number of inference steps for refiner",
                minimum=10,
                maximum=100,
                step=1,
                value=25,
            )

    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=result,
        fn=generate,
        cache_examples=CACHE_EXAMPLES,
    )

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        queue=False,
        api_name=False,
    )
    use_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_prompt_2,
        outputs=prompt_2,
        queue=False,
        api_name=False,
    )
    use_negative_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt_2,
        outputs=negative_prompt_2,
        queue=False,
        api_name=False,
    )
    apply_refiner.change(
        fn=lambda x: gr.update(visible=x),
        inputs=apply_refiner,
        outputs=refiner_params,
        queue=False,
        api_name=False,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            prompt_2.submit,
            negative_prompt_2.submit,
            run_button.click,
        ],
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            prompt_2,
            negative_prompt_2,
            use_negative_prompt,
            use_prompt_2,
            use_negative_prompt_2,
            seed,
            # 1024,
            # 1024,
            width,
            height,
            guidance_scale_base,
            guidance_scale_refiner,
            num_inference_steps_base,
            num_inference_steps_refiner,
            apply_refiner,
        ],
        outputs=result,
        api_name="run",
    )

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