|
import spaces |
|
from diffusers import ( |
|
StableDiffusionPipeline, |
|
DPMSolverMultistepScheduler, |
|
DiffusionPipeline, |
|
) |
|
import gradio as gr |
|
import torch |
|
from PIL import Image |
|
import time |
|
import psutil |
|
import random |
|
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker |
|
|
|
start_time = time.time() |
|
current_steps = 25 |
|
|
|
SAFETY_CHECKER = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=torch.float16) |
|
|
|
UPSCALER = DiffusionPipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16) |
|
UPSCALER.to("cuda") |
|
|
|
|
|
|
|
class Model: |
|
def __init__(self, name, path=""): |
|
self.name = name |
|
self.path = path |
|
|
|
if path != "": |
|
self.pipe_t2i = StableDiffusionPipeline.from_pretrained( |
|
path, torch_dtype=torch.float16, safety_checker=SAFETY_CHECKER |
|
) |
|
self.pipe_t2i.scheduler = DPMSolverMultistepScheduler.from_config( |
|
self.pipe_t2i.scheduler.config |
|
) |
|
else: |
|
self.pipe_t2i = None |
|
|
|
|
|
models = [ |
|
|
|
|
|
Model("anything-v4.0", "xyn-ai/anything-v4.0"), |
|
] |
|
|
|
MODELS = {m.name: m for m in models} |
|
|
|
device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶" |
|
|
|
|
|
def error_str(error, title="Error"): |
|
return ( |
|
f"""#### {title} |
|
{error}""" |
|
if error |
|
else "" |
|
) |
|
|
|
@spaces.GPU |
|
def inference( |
|
prompt, |
|
neg_prompt, |
|
guidance, |
|
steps, |
|
seed, |
|
model_name, |
|
): |
|
|
|
print(psutil.virtual_memory()) |
|
|
|
if seed == 0: |
|
seed = random.randint(0, 2147483647) |
|
|
|
generator = torch.Generator("cuda").manual_seed(seed) |
|
|
|
try: |
|
low_res_image, up_res_image = txt_to_img( |
|
model_name, |
|
prompt, |
|
neg_prompt, |
|
guidance, |
|
steps, |
|
generator, |
|
) |
|
return low_res_image, up_res_image, f"Done. Seed: {seed}", |
|
except Exception as e: |
|
return None, None, error_str(e) |
|
|
|
|
|
def txt_to_img( |
|
model_name, |
|
prompt, |
|
neg_prompt, |
|
guidance, |
|
steps, |
|
generator, |
|
): |
|
pipe = MODELS[model_name].pipe_t2i |
|
|
|
if torch.cuda.is_available(): |
|
pipe = pipe.to("cuda") |
|
pipe.enable_xformers_memory_efficient_attention() |
|
|
|
low_res_latents = pipe( |
|
prompt, |
|
negative_prompt=neg_prompt, |
|
num_inference_steps=int(steps), |
|
guidance_scale=guidance, |
|
generator=generator, |
|
output_type="latent", |
|
).images |
|
|
|
with torch.no_grad(): |
|
low_res_image = pipe.decode_latents(low_res_latents) |
|
low_res_image = pipe.numpy_to_pil(low_res_image) |
|
|
|
up_res_image = UPSCALER( |
|
prompt=prompt, |
|
negative_prompt=neg_prompt, |
|
image=low_res_latents, |
|
num_inference_steps=20, |
|
guidance_scale=0, |
|
generator=generator, |
|
).images |
|
|
|
pipe.to("cpu") |
|
torch.cuda.empty_cache() |
|
|
|
return low_res_image[0], up_res_image[0] |
|
|
|
|
|
def replace_nsfw_images(results): |
|
for i in range(len(results.images)): |
|
if results.nsfw_content_detected[i]: |
|
results.images[i] = Image.open("nsfw.png") |
|
return results.images |
|
|
|
|
|
with gr.Blocks(css="style.css") as demo: |
|
gr.HTML( |
|
f""" |
|
<div class="finetuned-diffusion-div"> |
|
<div style="text-align: center"> |
|
<h1>Anything v4 model + <a href="https://huggingface.co/stabilityai/sd-x2-latent-upscaler">Stable Diffusion Latent Upscaler</a></h1> |
|
<p> |
|
Demo for the <a href="https://huggingface.co/andite/anything-v4.0">Anything v4</a> model hooked with the ultra-fast <a href="https://huggingface.co/stabilityai/sd-x2-latent-upscaler">Latent Upscaler</a> |
|
</p> |
|
</div> |
|
<!-- |
|
<p>To skip the queue, you can duplicate this Space<br> |
|
<a style="display:inline-block" href="https://huggingface.co/spaces/patrickvonplaten/finetuned_diffusion?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></p> |
|
--> |
|
</div> |
|
""" |
|
) |
|
|
|
with gr.Column(scale=100): |
|
with gr.Group(visible=False): |
|
model_name = gr.Dropdown( |
|
label="Model", |
|
choices=[m.name for m in models], |
|
value=models[0].name, |
|
visible=False |
|
) |
|
|
|
with gr.Row(elem_id="prompt-container"): |
|
with gr.Column(): |
|
prompt = gr.Textbox( |
|
label="Enter your prompt", |
|
show_label=False, |
|
max_lines=1, |
|
placeholder="Enter your prompt", |
|
elem_id="prompt-text-input", |
|
container=False, |
|
) |
|
neg_prompt = gr.Textbox( |
|
label="Enter your negative prompt", |
|
show_label=False, |
|
max_lines=1, |
|
placeholder="Enter a negative prompt", |
|
elem_id="negative-prompt-text-input", |
|
container=False, |
|
) |
|
generate = gr.Button("Generate image", scale=0) |
|
|
|
with gr.Accordion("Advanced Options", open=False): |
|
with gr.Group(): |
|
with gr.Row(): |
|
guidance = gr.Slider( |
|
label="Guidance scale", value=7.5, maximum=15 |
|
) |
|
steps = gr.Slider( |
|
label="Steps", |
|
value=current_steps, |
|
minimum=2, |
|
maximum=75, |
|
step=1, |
|
) |
|
|
|
seed = gr.Slider( |
|
0, 2147483647, label="Seed (0 = random)", value=0, step=1 |
|
) |
|
|
|
|
|
with gr.Column(scale=100): |
|
with gr.Row(): |
|
with gr.Column(scale=75): |
|
up_res_image = gr.Image(label="Upscaled 1024px Image", width=1024, height=1024) |
|
with gr.Column(scale=25): |
|
low_res_image = gr.Image(label="Original 512px Image", width=512, height=512) |
|
error_output = gr.Markdown() |
|
|
|
inputs = [ |
|
prompt, |
|
neg_prompt, |
|
guidance, |
|
steps, |
|
seed, |
|
model_name, |
|
] |
|
outputs = [low_res_image, up_res_image, error_output] |
|
prompt.submit(inference, inputs=inputs, outputs=outputs) |
|
generate.click(inference, inputs=inputs, outputs=outputs) |
|
|
|
ex = gr.Examples( |
|
[ |
|
["a mecha robot in a favela", "low quality", 7.5, 25, 33, models[0].name], |
|
["the spirit of a tamagotchi wandering in the city of Paris", "low quality, bad render", 7.5, 50, 85, models[0].name], |
|
], |
|
inputs=[prompt, neg_prompt, guidance, steps, seed, model_name], |
|
outputs=outputs, |
|
fn=inference, |
|
cache_examples=True, |
|
) |
|
ex.dataset.headers = [""] |
|
|
|
gr.HTML( |
|
""" |
|
<div style="border-top: 1px solid #303030;"> |
|
<br> |
|
<p>Space by 🤗 Hugging Face, models by Stability AI, andite, linaqruf and others ❤️</p> |
|
<p>This space uses the <a href="https://github.com/LuChengTHU/dpm-solver">DPM-Solver++</a> sampler by <a href="https://arxiv.org/abs/2206.00927">Cheng Lu, et al.</a>.</p> |
|
<p>This is a Demo Space For:<br> |
|
<a href="https://huggingface.co/stabilityai/sd-x2-latent-upscaler">Stability AI's Latent Upscaler</a> |
|
</div> |
|
""" |
|
) |
|
|
|
print(f"Space built in {time.time() - start_time:.2f} seconds") |
|
|
|
demo.queue(api_open=False) |
|
demo.launch(show_api=False) |
|
|
