prithivMLmods's picture
Update app.py
d671d5d verified
raw
history blame
10.6 kB
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
@spaces.GPU(duration=30)
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)