|
from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler |
|
import gradio as gr |
|
import torch |
|
from PIL import Image |
|
import utils |
|
import datetime |
|
import time |
|
import psutil |
|
import random |
|
|
|
|
|
start_time = time.time() |
|
is_colab = utils.is_google_colab() |
|
state = None |
|
current_steps = 25 |
|
|
|
class Model: |
|
def __init__(self, name, path="", prefix=""): |
|
self.name = name |
|
self.path = path |
|
self.prefix = prefix |
|
self.pipe_t2i = None |
|
self.pipe_i2i = None |
|
|
|
models = [ |
|
Model("Protogen x5.3 (Photorealism)", "darkstorm2150/Protogen_x5.3_Official_Release") |
|
] |
|
|
|
custom_model = None |
|
if is_colab: |
|
models.insert(0, Model("Custom model")) |
|
custom_model = models[0] |
|
|
|
last_mode = "txt2img" |
|
current_model = models[1] if is_colab else models[0] |
|
current_model_path = current_model.path |
|
|
|
if is_colab: |
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
current_model.path, |
|
torch_dtype=torch.float16, |
|
scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"), |
|
safety_checker=lambda images, clip_input: (images, False) |
|
) |
|
|
|
else: |
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
current_model.path, |
|
torch_dtype=torch.float16, |
|
scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler") |
|
) |
|
|
|
if torch.cuda.is_available(): |
|
pipe = pipe.to("cuda") |
|
pipe.enable_xformers_memory_efficient_attention() |
|
|
|
device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶" |
|
|
|
def error_str(error, title="Error"): |
|
return f"""#### {title} |
|
{error}""" if error else "" |
|
|
|
def update_state(new_state): |
|
global state |
|
state = new_state |
|
|
|
def update_state_info(old_state): |
|
if state and state != old_state: |
|
return gr.update(value=state) |
|
|
|
def custom_model_changed(path): |
|
models[0].path = path |
|
global current_model |
|
current_model = models[0] |
|
|
|
def on_model_change(model_name): |
|
|
|
prefix = "Enter prompt. \"" + next((m.prefix for m in models if m.name == model_name), None) + "\" is prefixed automatically" if model_name != models[0].name else "Don't forget to use the custom model prefix in the prompt!" |
|
|
|
return gr.update(visible = model_name == models[0].name), gr.update(placeholder=prefix) |
|
|
|
def on_steps_change(steps): |
|
global current_steps |
|
current_steps = steps |
|
|
|
def pipe_callback(step: int, timestep: int, latents: torch.FloatTensor): |
|
update_state(f"{step}/{current_steps} steps") |
|
|
|
def inference(model_name, prompt, guidance, steps, n_images=1, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""): |
|
|
|
update_state(" ") |
|
|
|
print(psutil.virtual_memory()) |
|
|
|
global current_model |
|
for model in models: |
|
if model.name == model_name: |
|
current_model = model |
|
model_path = current_model.path |
|
|
|
|
|
if seed == 0: |
|
seed = random.randint(0, 2147483647) |
|
|
|
generator = torch.Generator('cuda').manual_seed(seed) |
|
|
|
try: |
|
if img is not None: |
|
return img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed), f"Done. Seed: {seed}" |
|
else: |
|
return txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator, seed), f"Done. Seed: {seed}" |
|
except Exception as e: |
|
return None, error_str(e) |
|
|
|
def txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width, height, generator, seed): |
|
|
|
print(f"{datetime.datetime.now()} txt_to_img, model: {current_model.name}") |
|
|
|
global last_mode |
|
global pipe |
|
global current_model_path |
|
if model_path != current_model_path or last_mode != "txt2img": |
|
current_model_path = model_path |
|
|
|
update_state(f"Loading {current_model.name} text-to-image model...") |
|
|
|
if is_colab or current_model == custom_model: |
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
current_model_path, |
|
torch_dtype=torch.float16, |
|
scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"), |
|
safety_checker=lambda images, clip_input: (images, False) |
|
) |
|
else: |
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
current_model_path, |
|
torch_dtype=torch.float16, |
|
scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler") |
|
) |
|
|
|
|
|
|
|
if torch.cuda.is_available(): |
|
pipe = pipe.to("cuda") |
|
pipe.enable_xformers_memory_efficient_attention() |
|
last_mode = "txt2img" |
|
|
|
prompt = current_model.prefix + prompt |
|
result = pipe( |
|
prompt, |
|
negative_prompt = neg_prompt, |
|
num_images_per_prompt=n_images, |
|
num_inference_steps = int(steps), |
|
guidance_scale = guidance, |
|
width = width, |
|
height = height, |
|
generator = generator, |
|
callback=pipe_callback) |
|
|
|
|
|
|
|
return replace_nsfw_images(result) |
|
|
|
def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed): |
|
|
|
print(f"{datetime.datetime.now()} img_to_img, model: {model_path}") |
|
|
|
global last_mode |
|
global pipe |
|
global current_model_path |
|
if model_path != current_model_path or last_mode != "img2img": |
|
current_model_path = model_path |
|
|
|
update_state(f"Loading {current_model.name} image-to-image model...") |
|
|
|
if is_colab or current_model == custom_model: |
|
pipe = StableDiffusionImg2ImgPipeline.from_pretrained( |
|
current_model_path, |
|
torch_dtype=torch.float16, |
|
scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler"), |
|
safety_checker=lambda images, clip_input: (images, False) |
|
) |
|
else: |
|
pipe = StableDiffusionImg2ImgPipeline.from_pretrained( |
|
current_model_path, |
|
torch_dtype=torch.float16, |
|
scheduler=DPMSolverMultistepScheduler.from_pretrained(current_model.path, subfolder="scheduler") |
|
) |
|
|
|
|
|
|
|
if torch.cuda.is_available(): |
|
pipe = pipe.to("cuda") |
|
pipe.enable_xformers_memory_efficient_attention() |
|
last_mode = "img2img" |
|
|
|
prompt = current_model.prefix + prompt |
|
ratio = min(height / img.height, width / img.width) |
|
img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS) |
|
result = pipe( |
|
prompt, |
|
negative_prompt = neg_prompt, |
|
num_images_per_prompt=n_images, |
|
image = img, |
|
num_inference_steps = int(steps), |
|
strength = strength, |
|
guidance_scale = guidance, |
|
|
|
|
|
generator = generator, |
|
callback=pipe_callback) |
|
|
|
|
|
|
|
return replace_nsfw_images(result) |
|
|
|
def replace_nsfw_images(results): |
|
|
|
if is_colab: |
|
return results.images |
|
|
|
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="Protogen Web UI Colab"> |
|
<div> |
|
<h1>Protogen Web UI</h1> |
|
</div> |
|
<p> |
|
Demo for multiple fine-tuned Stable Diffusion models, trained on different styles: <br> |
|
<a href="https://huggingface.co/darkstorm2150/Protogen_x5.3_Official_Release">darkstorm2150 |
|
</a>, <a href="https://huggingface.co/darkstorm2150/Protogen_x5.8_Official_Release">darkstorm2150</a>, <a href="https://huggingface.co/darkstorm2150/Protogen_v2.2_Official_Release">darkstorm2150</a> + in colab notebook you can load any other Diffusers 🧨 SD model hosted on HuggingFace 🤗. |
|
</p> |
|
<p> |
|
Running on <b>{device}</b>{(" in a <b>Google Colab</b>." if is_colab else "")} |
|
</p> |
|
<p>You can also duplicate this space and upgrade to gpu by going to settings:<br> |
|
<a style="display:inline-block" href="https://huggingface.co/spaces/anzorq/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.Row(): |
|
|
|
with gr.Column(scale=55): |
|
with gr.Group(): |
|
model_name = gr.Dropdown(label="Model", choices=[m.name for m in models], value=current_model.name) |
|
with gr.Box(visible=False) as custom_model_group: |
|
custom_model_path = gr.Textbox(label="Custom model path", placeholder="Path to model, e.g. nitrosocke/Arcane-Diffusion", interactive=True) |
|
gr.HTML("<div><font size='2'>Custom models have to be downloaded first, so give it some time.</font></div>") |
|
|
|
with gr.Row(): |
|
prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,placeholder="Enter prompt. Style applied automatically").style(container=False) |
|
generate = gr.Button(value="Generate").style(rounded=(False, True, True, False)) |
|
|
|
|
|
|
|
gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto") |
|
|
|
state_info = gr.Textbox(label="State", show_label=False, max_lines=2).style(container=False) |
|
error_output = gr.Markdown() |
|
|
|
with gr.Column(scale=45): |
|
with gr.Tab("Options"): |
|
with gr.Group(): |
|
neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image") |
|
|
|
n_images = gr.Slider(label="Images", value=1, minimum=1, maximum=10, step=1) |
|
|
|
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) |
|
|
|
with gr.Row(): |
|
width = gr.Slider(label="Width", value=512, minimum=64, maximum=1024, step=8) |
|
height = gr.Slider(label="Height", value=512, minimum=64, maximum=1024, step=8) |
|
|
|
seed = gr.Slider(0, 2147483647, label='Seed (0 = random)', value=0, step=1) |
|
|
|
with gr.Tab("Image to image"): |
|
with gr.Group(): |
|
image = gr.Image(label="Image", height=256, tool="editor", type="pil") |
|
strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5) |
|
|
|
if is_colab: |
|
model_name.change(on_model_change, inputs=model_name, outputs=[custom_model_group, prompt], queue=False) |
|
custom_model_path.change(custom_model_changed, inputs=custom_model_path, outputs=None) |
|
|
|
steps.change(on_steps_change, inputs=[steps], outputs=[], queue=False) |
|
|
|
inputs = [model_name, prompt, guidance, steps, n_images, width, height, seed, image, strength, neg_prompt] |
|
outputs = [gallery, error_output] |
|
prompt.submit(inference, inputs=inputs, outputs=outputs) |
|
generate.click(inference, inputs=inputs, outputs=outputs) |
|
|
|
|
|
|
|
|
|
ex = gr.Examples( |
|
[ |
|
|
|
|
|
[models[1].name, "(extremely detailed CG unity 8k wallpaper), the most beautiful artwork in the world", 7.5, 50, "human, people, canvas frame, ((disfigured)), ((bad art)), ((deformed)),((extra limbs)),((close up)),((b&w)), weird colors, blurry, (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))), out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck))), Photoshop, video game, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, mutation, mutated, extra limbs, extra legs, extra arms, disfigured, deformed, cross-eye, body out of frame, blurry, bad art, bad anatomy"], |
|
|
|
|
|
|
|
|
|
|
|
], |
|
inputs=[model_name, prompt, guidance, steps, neg_prompt], |
|
outputs=outputs, |
|
fn=inference, |
|
cache_examples=False, |
|
) |
|
|
|
gr.HTML(""" |
|
<div style="border-top: 1px solid #303030;"> |
|
<br> |
|
<p>Models by <a href="https://huggingface.co/darkstorm2150/Protogen_x5.3_Official_Release">darkstorm2150, </a> and others. ❤️</p> |
|
|
|
</div> |
|
""") |
|
|
|
demo.load(update_state_info, inputs=state_info, outputs=state_info, every=0.5, show_progress=False) |
|
|
|
print(f"Space built in {time.time() - start_time:.2f} seconds") |
|
|
|
|
|
demo.queue(concurrency_count=1) |
|
demo.launch(debug=is_colab, share=is_colab) |
|
|
