Spaces:
Runtime error
Runtime error
#!/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 | |
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() |