from diffusers import (
StableDiffusionPipeline,
StableDiffusionImg2ImgPipeline,
DPMSolverMultistepScheduler,
)
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)
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
)
self.pipe_i2i = StableDiffusionImg2ImgPipeline(**self.pipe_t2i.components, safety_checker=SAFETY_CHECKER)
else:
self.pipe_t2i = None
self.pipe_i2i = None
models = [
Model("2.2", "darkstorm2150/Protogen_v2.2_Official_Release"),
Model("3.4", "darkstorm2150/Protogen_x3.4_Official_Release"),
Model("5.3", "darkstorm2150/Protogen_v5.3_Official_Release"),
Model("5.8", "darkstorm2150/Protogen_x5.8_Official_Release"),
Model("Dragon", "darkstorm2150/Protogen_Dragon_Official_Release"),
]
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 ""
)
def inference(
model_name,
prompt,
guidance,
steps,
n_images=1,
width=512,
height=512,
seed=0,
img=None,
strength=0.5,
neg_prompt="",
):
print(psutil.virtual_memory()) # print memory usage
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_name,
prompt,
n_images,
neg_prompt,
img,
strength,
guidance,
steps,
width,
height,
generator,
seed,
),
f"Done. Seed: {seed}",
)
else:
return (
txt_to_img(
model_name,
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_name,
prompt,
n_images,
neg_prompt,
guidance,
steps,
width,
height,
generator,
seed,
):
pipe = MODELS[model_name].pipe_t2i
if torch.cuda.is_available():
pipe = pipe.to("cuda")
pipe.enable_xformers_memory_efficient_attention()
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,
)
pipe.to("cpu")
return replace_nsfw_images(result)
def img_to_img(
model_name,
prompt,
n_images,
neg_prompt,
img,
strength,
guidance,
steps,
width,
height,
generator,
seed,
):
pipe = MODELS[model_name].pipe_i2i
if torch.cuda.is_available():
pipe = pipe.to("cuda")
pipe.enable_xformers_memory_efficient_attention()
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,
)
pipe.to("cpu")
return replace_nsfw_images(result)
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(
"""
Protogen Diffusion
Demo for multiple fine-tuned Protogen Stable Diffusion models.
You can also duplicate this space and upgrade to gpu by going to settings:
"""
)
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=models[0].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. darkstorm2150/Protogen_x3.4_Official_Release",
interactive=True,
)
gr.HTML(
"Custom models have to be downloaded first, so give it some time.
"
)
with gr.Row():
prompt = gr.Textbox(
label="Prompt",
show_label=False,
max_lines=2,
placeholder="Enter prompt.",
).style(container=False)
generate = gr.Button(value="Generate").style(
rounded=(False, True, True, False)
)
# image_out = gr.Image(height=512)
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=4, 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,
)
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[0].name, "portrait of a beautiful alyx vance half life", 10, 25],
[models[1].name, "Brad Pitt with sunglasses, highly realistic", 7.5, 25],
],
inputs=[model_name, prompt, guidance, steps],
outputs=outputs,
fn=inference,
cache_examples=False,
)
gr.HTML(
"""
"""
)
print(f"Space built in {time.time() - start_time:.2f} seconds")
demo.queue(concurrency_count=1)
demo.launch()