app.py

from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, UNet2DConditionModel
from diffusers.utils import load_image
from diffusers import (
    DDIMScheduler,
    PNDMScheduler,
    LMSDiscreteScheduler,
    EulerDiscreteScheduler,
    EulerAncestralDiscreteScheduler,
    DPMSolverMultistepScheduler,
)
import torch
import os
import random
import numpy as np
from PIL import Image
from typing import Tuple
import gradio as gr
import spaces
DESCRIPTION =  """
    # CosmicMan
    - CosmicMan: A Text-to-Image Foundation Model for Humans (CVPR 2024 (Highlight))
    """

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

schedule_map = {
    "ddim"   : DDIMScheduler,
    "pndm"   : PNDMScheduler,
    "lms"    : LMSDiscreteScheduler,
    "euler"  : EulerDiscreteScheduler,
    "euler_a": EulerAncestralDiscreteScheduler,
    "dpm"    : DPMSolverMultistepScheduler,
}

examples = [
    "A fit Caucasian elderly woman, her wavy white hair above shoulders, wears a pink floral cotton long-sleeve shirt and a cotton hat against a natural landscape in an upper body shot",
    "A closeup of a doll with a purple ribbon around her neck, best quality, extremely detailed",
    "A closeup of a girl with a butterfly painted on her face",
    "A headshot, an asian elderly male, a blue wall, bald above eyes gray hair",
    "A closeup portrait shot against a white wall, a fit Caucasian adult female with wavy blonde hair falling above her chest wears a short sleeve silk floral dress and a floral silk normal short sleeve white blouse",
    "A headshot, an adult caucasian male, fit, a white wall, red crew cut curly hair, short sleeve normal blue t-shirt, best quality, extremely detailed",
    "A closeup of a man wearing a red shirt with a flower design on it",
    "There is a man wearing a mask and holding a cell phone",
    "Two boys playing in the yard",
]

style_list = [
    {
        "name": "(No style)",
        "prompt": "{prompt}",
        "negative_prompt": "",
    },
    {
        "name": "Cinematic",
        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
    },
    {
        "name": "Photographic",
        "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
    },
    {
        "name": "Anime",
        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
    },
    {
        "name": "Fantasy art",
        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
        "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
    },
    {
        "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",
        "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured",
    }
]

styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "(No style)"
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
MAX_SEED = np.iinfo(np.int32).max
NUM_IMAGES_PER_PROMPT = 1 
 
def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    if not negative:
        negative = ""
    return p.replace("{prompt}", positive), n + negative

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

print("Loading Model!")
schedule: str = "euler_a"
base_model_path: str = "stabilityai/stable-diffusion-xl-base-1.0"
refiner_model_path: str = "stabilityai/stable-diffusion-xl-refiner-1.0" 
unet_path: str = "cosmicman/CosmicMan-SDXL"
SCHEDULER = schedule_map[schedule]
scheduler = SCHEDULER.from_pretrained(base_model_path, subfolder="scheduler", torch_dtype=torch.float16)
unet = UNet2DConditionModel.from_pretrained(unet_path, torch_dtype=torch.float16)

pipe = StableDiffusionXLPipeline.from_pretrained(
    base_model_path,
    unet=unet,
    scheduler=scheduler,
    torch_dtype=torch.float16,
    use_safetensors=True
).to("cuda")
pipe.watermark = NoWatermark()
refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    base_model_path, # we found use base_model_path instead of refiner_model_path may get a better performance
    scheduler=scheduler,
    torch_dtype=torch.float16,
    use_safetensors=True
).to("cuda")
refiner.watermark = NoWatermark()
print("Finish Loading Model!")

@spaces.GPU(duration=120)
def generate_image(prompt, 
                   n_prompt="", 
                   style: str = DEFAULT_STYLE_NAME,
                   steps: int = 50, 
                   height: int = 1024,
                   width: int = 1024,
                   scale: float = 7.5, 
                   img_num: int = 4, 
                   seeds: int = 42,
                   random_seed: bool = False,
):
    print("Beign to generate")
    image_list = []
    for i in range(img_num):
        seed = int(randomize_seed_fn(seeds, random_seed))
        generator = torch.Generator(device="cuda").manual_seed(seed)
        positive_prompt, negative_prompt = apply_style(style, prompt, n_prompt)
        image = pipe(positive_prompt, num_inference_steps=steps, 
                guidance_scale=scale, height=height, 
                width=width, negative_prompt=negative_prompt, 
                generator=generator, output_type="latent").images[0]
        image = refiner(positive_prompt, negative_prompt=negative_prompt, image=image[None, :]).images[0]
        image_list.append((image,f"Seed {seed}"))
    return image_list

with gr.Blocks(theme=gr.themes.Soft(),css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Group():
        with gr.Row():
            with gr.Column():
                input_prompt = gr.Textbox(label="Input prompt", lines=3, max_lines=5)
                negative_prompt = gr.Textbox(label="Negative prompt",value="")
                run_button = gr.Button("Run", scale=0)
        result = gr.Gallery(label="Result", show_label=False, elem_id="gallery", columns=[2], rows=[2], object_fit="contain", height="auto")
    with gr.Accordion("Advanced options", open=False):
        with gr.Row():
            style_selection = gr.Radio(
                    show_label=True,
                    container=True,
                    interactive=True,
                    choices=STYLE_NAMES,
                    value=DEFAULT_STYLE_NAME,
                    label="Image Style",
                )
        with gr.Row():       
            height = gr.Slider(minimum=512, maximum=1536, value=1024, label="Height", step=64)
            width = gr.Slider(minimum=512, maximum=1536, value=1024, label="Witdh", step=64) 
        with gr.Row():
            steps =  gr.Slider(minimum=1, maximum=50, value=30, label="Number of diffusion steps", step=1)
            scale = gr.Number(minimum=1, maximum=12, value=7.5, label="Number of scale")        
        with gr.Row():     
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            random_seed = gr.Checkbox(label="Randomize seed", value=True)
            img_num = gr.Slider(minimum=1, maximum=4, value=4, label="Number of images", step=1)   

    gr.Examples(
        examples=examples,
        inputs=input_prompt,
        outputs=result,
        fn=generate_image,
        cache_examples=0,
    )

    gr.on(
        triggers=[
            input_prompt.submit,
            negative_prompt.submit,
            run_button.click,
        ],
        fn=generate_image, 
        inputs = [input_prompt, negative_prompt, style_selection, steps, height, width, scale, img_num, seed, random_seed],
        outputs= result,
        api_name="run")


if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_api=False, debug=False)