Spaces:
Running
on
Zero
Running
on
Zero
import os | |
import random | |
import uuid | |
import json | |
import gradio as gr | |
import numpy as np | |
from PIL import Image | |
import spaces | |
import torch | |
from diffusers import DiffusionPipeline | |
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler | |
from typing import Tuple | |
bad_words = json.loads(os.getenv('BAD_WORDS', "[]")) | |
bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]")) | |
default_negative = os.getenv("default_negative","") | |
def check_text(prompt, negative=""): | |
for i in bad_words: | |
if i in prompt: | |
return True | |
for i in bad_words_negative: | |
if i in negative: | |
return True | |
return False | |
style_list = [ | |
{ | |
"name": "Photo", | |
"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": "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": "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": "3D Model", | |
"prompt": "professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting", | |
"negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting", | |
}, | |
{ | |
"name": "(No style)", | |
"prompt": "{prompt}", | |
"negative_prompt": "", | |
}, | |
] | |
styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list} | |
STYLE_NAMES = list(styles.keys()) | |
DEFAULT_STYLE_NAME = "Photo" | |
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 | |
DESCRIPTION = """## Text to Image | |
""" | |
if not torch.cuda.is_available(): | |
DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>" | |
MAX_SEED = np.iinfo(np.int32).max | |
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1" | |
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 = 1 | |
if torch.cuda.is_available(): | |
pipe = StableDiffusionXLPipeline.from_pretrained( | |
"SG161222/RealVisXL_V5.0_Lightning", | |
torch_dtype=torch.float16, | |
use_safetensors=True, | |
add_watermarker=False, | |
variant="fp16" | |
) | |
pipe2 = StableDiffusionXLPipeline.from_pretrained( | |
"SG161222/RealVisXL_V4.0_Lightning", | |
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_NAME, | |
seed: int = 0, | |
width: int = 1024, | |
height: int = 1024, | |
guidance_scale: float = 3, | |
randomize_seed: bool = False, | |
use_resolution_binning: bool = True, | |
progress=gr.Progress(track_tqdm=True), | |
): | |
if check_text(prompt, negative_prompt): | |
raise ValueError("Prompt contains restricted words.") | |
prompt, negative_prompt = apply_style(style, prompt, negative_prompt) | |
seed = int(randomize_seed_fn(seed, randomize_seed)) | |
generator = torch.Generator().manual_seed(seed) | |
if not use_negative_prompt: | |
negative_prompt = "" # type: ignore | |
negative_prompt += default_negative | |
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 = [ | |
"3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)", | |
"A glass cup of cold coffee placed on a rustic wooden table, surrounded by soft morning light. The coffee is rich, dark, and topped with a light layer of creamy froth, droplets of condensation sliding down the glass.", | |
"Vector illustration of a horse, vector graphic design with flat colors on an brown background in the style of vector art, using simple shapes and graphics with simple details, professionally designed as a tshirt logo ready for print on a white background. --ar 89:82 --v 6.0 --style raw", | |
"Man in brown leather jacket posing for camera, in the style of sleek and stylized, clockpunk, subtle shades, exacting precision, ferrania p30 --ar 67:101 --v 5", | |
"Commercial photography, giant burger, white lighting, studio light, 8k octane rendering, high resolution photography, insanely detailed, fine details, on white isolated plain, 8k, commercial photography, stock photo, professional color grading, --v 4 --ar 9:16 " | |
] | |
css = ''' | |
.gradio-container { | |
max-width: 590px !important; | |
margin: 0 auto !important; | |
} | |
h1 { | |
text-align: center; | |
} | |
footer { | |
visibility: hidden; | |
} | |
''' | |
with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo: | |
gr.Markdown(DESCRIPTION) | |
with gr.Group(): | |
with gr.Row(): | |
prompt = gr.Text( | |
show_label=False, | |
max_lines=1, | |
placeholder="Enter your prompt", | |
container=False, | |
) | |
run_button = gr.Button("Run", scale=0, variant="primary") | |
result = gr.Gallery(label="Result", columns=1, preview=True) | |
with gr.Accordion("Advanced options", open=False): | |
use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True) | |
negative_prompt = gr.Text( | |
label="Negative prompt", | |
max_lines=1, | |
placeholder="Enter a negative prompt", | |
value="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck", | |
visible=True, | |
) | |
with gr.Row(): | |
num_inference_steps = gr.Slider( | |
label="Steps", | |
minimum=10, | |
maximum=60, | |
step=1, | |
value=20, | |
) | |
with gr.Row(): | |
num_images_per_prompt = gr.Slider( | |
label="Images", | |
minimum=1, | |
maximum=4, | |
step=1, | |
value=2, | |
) | |
seed = gr.Slider( | |
label="Seed", | |
minimum=0, | |
maximum=MAX_SEED, | |
step=1, | |
value=0, | |
visible=True | |
) | |
randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
with gr.Row(visible=True): | |
width = gr.Slider( | |
label="Width", | |
minimum=512, | |
maximum=2048, | |
step=8, | |
value=1024, | |
) | |
height = gr.Slider( | |
label="Height", | |
minimum=512, | |
maximum=2048, | |
step=8, | |
value=1024, | |
) | |
with gr.Row(): | |
guidance_scale = gr.Slider( | |
label="Guidance Scale", | |
minimum=0.1, | |
maximum=20.0, | |
step=0.1, | |
value=3.0, | |
) | |
with gr.Row(visible=True): | |
style_selection = gr.Radio( | |
show_label=True, | |
container=True, | |
interactive=True, | |
choices=STYLE_NAMES, | |
value=DEFAULT_STYLE_NAME, | |
label="Image 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(ssr_mode=True, show_error=True, share=True) |