app.py

#!/usr/bin/env python

import json
import os
import random
from typing import Tuple
import uuid

from diffusers import DiffusionPipeline
import gradio as gr
import numpy as np
from PIL import Image
import spaces
import torch

from gradio_imagefeed import ImageFeed


DEFAULT_STYLE = "Photograph"
DEFAULT_NEGATIVE = (
    "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon,"
    " drawing, anime, asian, bad anatomy:1.4), text, cropped, out of frame,"
    " worst quality, low quality, morbid, mutilated,"
    " extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation,"
    " deformed, blurry, bad anatomy, bad proportions, extra limbs,"
    " plastic, fake, missing arms, missing legs, fat, ugly, huge breasts,"
    " extra arms, extra legs, fused fingers, too many fingers"
)
STYLES = {
    "Photograph": (
        (
            "realistic photograph of {positive}, ultra fine detail, lifelike,"
            " high-resolution, sharp, realistic colors, photorealistic, Nikon, 35mm"
        ),
        DEFAULT_NEGATIVE,
    ),
    "Cinematic": (
        (
            "cinematic photograph of {positive}, 35mm photograph, film, bokeh,"
            " professional, 4k, highly detailed"
        ),
        DEFAULT_NEGATIVE,
    ),
    "Still Photo": (
        (
            "cinematic still photograph of {positive}, emotional, harmonious, vignette,"
            " highly detailed, bokeh, cinemascope, moody, epic, gorgeous, film grain,"
            " grainy, high resolution"
        ),
        DEFAULT_NEGATIVE,
    ),
    'No Style': (
        '{positive}',
        '')
}


def apply_style(name: str, pos: str, neg: str) -> Tuple[str, str]:
    try:
        def_pos, def_neg = STYLES[name]
    except KeyError:
        def_pos, def_neg = "{positive}", ""
    finally:
        pos = def_pos.replace("{positive}", pos).strip().strip(",")
        neg = def_neg + (", " + neg).strip().strip(",")
        return (pos, neg)


DESCRIPTION = ""
MAX_SEED = np.iinfo(np.int32).max
CACHE_EXAMPLES = False  # "lazy"
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

NUM_IMAGES_PER_PROMPT = 3

if torch.cuda.is_available():
    pipe = DiffusionPipeline.from_pretrained(
        "SG161222/RealVisXL_V4.0",
        torch_dtype=torch.float16,
        use_safetensors=True,
        add_watermarker=False,
        variant="fp16",
    )
    pipe2 = DiffusionPipeline.from_pretrained(
        "SG161222/RealVisXL_V3.0_Turbo",
        torch_dtype=torch.float16,
        use_safetensors=True,
        add_watermarker=False,
        variant="fp16",
    )
    if ENABLE_CPU_OFFLOAD:
        pipe.enable_model_cpu_offload()
        pipe2.enable_model_cpu_offload()
    else:
        pipe.to(device)
        pipe2.to(device)
        print("Loaded on Device!")

    if USE_TORCH_COMPILE:
        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
        pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
        print("Model Compiled!")


def save_image(img):
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name


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


@spaces.GPU(enable_queue=True)
def generate(
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    style: str = DEFAULT_STYLE,
    seed: int = 0,
    width: int = 896,
    height: int = 1152,
    guidance_scale: float = 3,
    randomize_seed: bool = False,
    use_resolution_binning: bool = True,
    progress=gr.Progress(track_tqdm=True),
):

    prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
    seed = int(randomize_seed_fn(seed, randomize_seed))
    generator = torch.Generator().manual_seed(seed)

    options = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "width": width,
        "height": height,
        "guidance_scale": guidance_scale,
        "num_inference_steps": 25,
        "generator": generator,
        "num_images_per_prompt": NUM_IMAGES_PER_PROMPT,
        "use_resolution_binning": use_resolution_binning,
        "output_type": "pil",
    }

    images = pipe(**options).images + pipe2(**options).images
    image_paths = [save_image(img) for img in images]

    return image_paths, seed


examples = [
    (
        "college life of 21 year old, depth of field, bokeh, shallow"
        " focus, minimalism, fujifilm xh2s with Canon EF lens, cinematic --ar 85:128"
        " --v 6.0 --style raw"
    ),
]

css = ""

with gr.Blocks(css=css, theme="rawrsor1/Everforest") as demo:
    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=4,
                placeholder="Enter a Prompt",
                container=False,
            )
            run_button = gr.Button("Run")
        #result = ImageFeed(label="Result")
        result = gr.Gallery(label="Result", preview=True)

    with gr.Accordion("Advanced", open=False):
        use_negative_prompt = gr.Checkbox(
            label="Use Negative", value=True, visible=True
        )
        negative_prompt = gr.Text(
            label="Negative Prompt",
            max_lines=4,
            placeholder="",
            value="",
            visible=True,
        )
        with gr.Row():
            num_inference_steps = gr.Slider(
                label="Steps",
                minimum=10,
                maximum=60,
                step=1,
                value=30,
                interactive=True
            )
        with gr.Row():
            num_images_per_prompt = gr.Slider(
                label="Image Count",
                minimum=1,
                maximum=5,
                step=1,
                value=2,
                interactive=True
            )
        seed = gr.Slider(
            label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, visible=True
        )
        randomize_seed = gr.Checkbox(label="New Seed", value=True)
        with gr.Row(visible=True):
            width = gr.Slider(
                label="Width",
                minimum=512,
                maximum=2048,
                step=16,
                value=896,
                interactive=True
            )
            height = gr.Slider(
                label="Height",
                minimum=512,
                maximum=2048,
                step=16,
                value=1152,
                interactive=True
            )
        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance",
                minimum=0.1,
                maximum=20.0,
                step=0.1,
                value=6,
                interactive=True
            )
    with gr.Row(visible=True):
        style_selection = gr.Radio(
            show_label=True,
            container=True,
            interactive=True,
            choices=list(STYLES.keys()),
            value=DEFAULT_STYLE,
            label="Style",
        )
    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=[result, seed],
        fn=generate,
        cache_examples=CACHE_EXAMPLES,
    )

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        api_name=False,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            run_button.click,
        ],
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            use_negative_prompt,
            style_selection,
            seed,
            width,
            height,
            guidance_scale,
            randomize_seed,
        ],
        outputs=[result, seed],
        api_name="run",
    )

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