Spaces:
Running
on
Zero
Running
on
Zero
#!/usr/bin/env python | |
from __future__ import annotations | |
import requests | |
import os | |
import random | |
import gradio as gr | |
import numpy as np | |
import spaces | |
import torch | |
import cv2 | |
import xformers | |
import triton | |
from PIL import Image | |
from io import BytesIO | |
from diffusers.utils import load_image | |
from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetInpaintPipeline, ControlNetModel, AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting, EulerDiscreteScheduler, DPMSolverMultistepScheduler | |
if not torch.cuda.is_available(): | |
DESCRIPTION += "\n<p>⚠️ This space is running on the CPU. This demo doesn't work on CPU 😞! Run on a GPU by duplicating this space or test our website for free and unlimited by <a href='https://squaadai.com'>clicking here</a>, which provides these and more options.</p>" | |
MAX_SEED = np.iinfo(np.int32).max | |
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1824")) | |
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1" | |
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1" | |
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
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 = "", | |
prompt_2: str = "", | |
negative_prompt_2: str = "", | |
use_negative_prompt: bool = False, | |
use_prompt_2: bool = False, | |
use_negative_prompt_2: bool = False, | |
seed: int = 0, | |
width: int = 1024, | |
height: int = 1024, | |
guidance_scale_base: float = 5.0, | |
num_inference_steps_base: int = 25, | |
controlnet_conditioning_scale: float = 1, | |
control_guidance_start: float = 0, | |
control_guidance_end: float = 1, | |
strength_img2img: float = 0.7, | |
use_vae: bool = False, | |
use_lora: bool = False, | |
use_lora2: bool = False, | |
model = 'stabilityai/stable-diffusion-xl-base-1.0', | |
vaecall = 'madebyollin/sdxl-vae-fp16-fix', | |
lora = '', | |
lora2 = '', | |
controlnet_model = 'diffusers/controlnet-canny-sdxl-1.0', | |
lora_scale: float = 0.7, | |
lora_scale2: float = 0.7, | |
use_img2img: bool = False, | |
use_controlnet: bool = False, | |
use_controlnetimg2img: bool = False, | |
url = '', | |
controlnet_img = '', | |
controlnet_img2img = '', | |
): | |
if torch.cuda.is_available(): | |
if not use_img2img: | |
scheduler = DPMSolverMultistepScheduler.from_pretrained(model, subfolder="scheduler") | |
pipe = DiffusionPipeline.from_pretrained(model, scheduler=scheduler, torch_dtype=torch.float16) | |
pipe.to(device) | |
if use_vae: | |
vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16) | |
pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16) | |
pipe.to(device) | |
if use_img2img: | |
pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16) | |
init_image = load_image(url) | |
if use_vae: | |
vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16) | |
pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16) | |
if use_controlnet: | |
controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16) | |
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, controlnet=controlnet, torch_dtype=torch.float16) | |
image = load_image(controlnet_img) | |
image = np.array(image) | |
image = cv2.Canny(image, 100, 200) | |
image = image[:, :, None] | |
image = np.concatenate([image, image, image], axis=2) | |
image = Image.fromarray(image) | |
if use_vae: | |
vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16) | |
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, controlnet=controlnet, vae=vae, torch_dtype=torch.float16) | |
if use_controlnetimg2img: | |
controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16) | |
pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, controlnet=controlnet, torch_dtype=torch.float16) | |
image_start = load_image(controlnet_img) | |
image = load_image(controlnet_img) | |
image_mask = load_image(controlnet_img2img) | |
image = np.array(image) | |
image = cv2.Canny(image, 100, 200) | |
image = image[:, :, None] | |
image = np.concatenate([image, image, image], axis=2) | |
image = Image.fromarray(image) | |
if use_vae: | |
vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16) | |
pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, controlnet=controlnet, vae=vae, torch_dtype=torch.float16) | |
if use_lora: | |
pipe.load_lora_weights(lora) | |
pipe.fuse_lora(lora_scale) | |
if use_lora2: | |
pipe.load_lora_weights(lora, adapter_name="1") | |
pipe.load_lora_weights(lora2, adapter_name="2") | |
pipe.set_adapters(["1", "2"], adapter_weights=[lora_scale, lora_scale2]) | |
generator = torch.Generator().manual_seed(seed) | |
if not use_negative_prompt: | |
negative_prompt = None # type: ignore | |
if not use_prompt_2: | |
prompt_2 = None # type: ignore | |
if not use_negative_prompt_2: | |
negative_prompt_2 = None # type: ignore | |
if use_controlnetimg2img: | |
image = pipe( | |
prompt=prompt, | |
strength=strength_img2img, | |
controlnet_conditioning_scale=controlnet_conditioning_scale, | |
eta=0.0, | |
mask_image=image_mask, | |
image=image_start, | |
control_image=image, | |
negative_prompt=negative_prompt, | |
width=width, | |
height=height, | |
guidance_scale=guidance_scale_base, | |
num_inference_steps=num_inference_steps_base, | |
generator=generator, | |
).images[0] | |
return image | |
if use_controlnet: | |
image = pipe( | |
prompt=prompt, | |
controlnet_conditioning_scale=controlnet_conditioning_scale, | |
control_guidance_start=control_guidance_start, | |
control_guidance_end=control_guidance_end, | |
image=image, | |
negative_prompt=negative_prompt, | |
prompt_2=prompt_2, | |
width=width, | |
height=height, | |
negative_prompt_2=negative_prompt_2, | |
guidance_scale=guidance_scale_base, | |
num_inference_steps=num_inference_steps_base, | |
generator=generator, | |
).images[0] | |
return image | |
elif use_img2img: | |
images = pipe( | |
prompt=prompt, | |
image=init_image, | |
strength=strength_img2img, | |
negative_prompt=negative_prompt, | |
prompt_2=prompt_2, | |
negative_prompt_2=negative_prompt_2, | |
width=width, | |
height=height, | |
guidance_scale=guidance_scale_base, | |
num_inference_steps=num_inference_steps_base, | |
generator=generator, | |
output_type="pil", | |
).images[0] | |
return images | |
else: | |
return pipe( | |
prompt=prompt, | |
negative_prompt=negative_prompt, | |
prompt_2=prompt_2, | |
negative_prompt_2=negative_prompt_2, | |
width=width, | |
height=height, | |
guidance_scale=guidance_scale_base, | |
num_inference_steps=num_inference_steps_base, | |
generator=generator, | |
output_type="pil", | |
).images[0] | |
theme = gr.themes.Monochrome( | |
text_size=gr.themes.Size(lg="18px", md="15px", sm="13px", xl="22px", xs="12px", xxl="24px", xxs="9px"), | |
font=[gr.themes.GoogleFont('Source Sans Pro'), 'ui-sans-serif', 'system-ui', 'sans-serif'], | |
) | |
with gr.Blocks(theme=theme, css="style.css") as demo: | |
gr.HTML() | |
gr.Markdown("""# Squaad AI 🧙 | |
### Run Stable Diffusion for free, in seconds | |
<small>Squaad AI was a free Artificial Intelligence website that brought together the best of the AI world in one place. Unfortunately, it had to cease operations after facing compatibility issues. This space is intended to make its legacy available as open source and allow anyone to run SD in a simple and free way.</small> | |
<div style="display: flex; flex-direction: column; align-items: flex-start;"> | |
<a target="_blank" href="https://discord.gg/CqXVwGDqFC"><img src="https://img.shields.io/badge/discord-join%20now-blue.svg?logo=discord" alt="discord.gg/CqXVwGDqFC"></a> | |
<a href="https://www.buymeacoffee.com/squaadai" target="_blank" style="margin-top: 10px;"><img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" alt="Buy Me A Coffee" style="height: auto !important;width: 130px !important;" ></a> | |
</div>""", elem_id="main_title") | |
with gr.Group(): | |
model = gr.Text(label='Model', value='stabilityai/stable-diffusion-xl-base-1.0', placeholder='e.g. stabilityai/stable-diffusion-xl-base-1.0', info='Enter the HUB path for this model! The model must be converted to diffusers to work. We recommend SG161222/RealVisXL_V4.0 for realistic generations, for example. By default, SDXL 1.0 (stabilityai/stable-diffusion-xl-base-1.0) will be used.') | |
vaecall = gr.Text(label='VAE', placeholder='e.g. madebyollin/sdxl-vae-fp16-fix', info='You can use the VAE that suits you best! Using an unfamiliar VAE or one that is not from the same base model mentioned above will result in errors or distorted images. You must enable "Use VAE" in advanced settings if you want to use a VAE.') | |
controlnet_model = gr.Text(label='Controlnet', placeholder='e.g diffusers/controlnet-canny-sdxl-1.0') | |
url = gr.Text(label='URL (Img2Img)', placeholder='e.g https://example.com/image.png') | |
controlnet_img = gr.Text(label='URL (Controlnet)', placeholder='e.g https://example.com/image.png') | |
controlnet_img2img = gr.Text(label='URL (Controlnet - IMG2IMG)', placeholder='e.g https://example.com/image.png') | |
with gr.Row(): | |
prompt = gr.Text( | |
placeholder="Input prompt", | |
label="Prompt", | |
show_label=False, | |
max_lines=1, | |
container=False, | |
) | |
run_button = gr.Button("Run", scale=0) | |
result = gr.Image(label="Result", show_label=False) | |
with gr.Accordion("Advanced options", open=False): | |
with gr.Row(): | |
use_controlnet = gr.Checkbox(label='Use Controlnet', value=False) | |
use_controlnetimg2img = gr.Checkbox(label='Use Controlnet Img2Img', value=False) | |
use_img2img = gr.Checkbox(label='Use Img2Img', value=False) | |
use_vae = gr.Checkbox(label='Use VAE', value=False) | |
use_lora = gr.Checkbox(label='Use Lora 1', value=False) | |
use_lora2 = gr.Checkbox(label='Use Lora 2', value=False) | |
use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False) | |
use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False) | |
use_negative_prompt_2 = gr.Checkbox(label="Use negative prompt 2", value=False) | |
with gr.Group(): | |
gr.Markdown("""### LoRAs | |
<small>To activate the use of LoRAs, you should check only the 'Use 1 LoRA' checkbox if you want to use just one LoRA, and only the 'Use 2 LoRAs' checkbox if you want to use 2 LoRAs (i.e., blend styles)! It's important that all LoRAs are from the same base version of the selected model. For example, if the model is SDXL 1.0, the LoRA must also be.</small>""" | |
) | |
lora = gr.Text(label='LoRA 1', placeholder='e.g. nerijs/pixel-art-xl', visible=False) | |
lora2 = gr.Text(label='LoRA 2', placeholder='e.g. nerijs/pixel-art-xl', visible=False) | |
lora_scale = gr.Slider( | |
info="The closer to 1, the more it will resemble LoRA, but errors may be visible.", | |
label="Lora Scale 1", | |
minimum=0.01, | |
maximum=1, | |
step=0.01, | |
value=0.7, | |
visible=False, | |
) | |
lora_scale2 = gr.Slider( | |
info="The closer to 1, the more it will resemble LoRA, but errors may be visible.", | |
label="Lora Scale 2", | |
minimum=0.01, | |
maximum=1, | |
step=0.01, | |
value=0.7, | |
visible=False, | |
) | |
with gr.Group(): | |
negative_prompt = gr.Text( | |
placeholder="Input Negative Prompt", | |
label="Negative prompt", | |
max_lines=1, | |
visible=False, | |
) | |
prompt_2 = gr.Text( | |
placeholder="Input Prompt 2", | |
label="Prompt 2", | |
max_lines=1, | |
visible=False, | |
) | |
negative_prompt_2 = gr.Text( | |
placeholder="Input Negative Prompt 2", | |
label="Negative prompt 2", | |
max_lines=1, | |
visible=False, | |
) | |
seed = gr.Slider( | |
label="Seed", | |
minimum=0, | |
maximum=MAX_SEED, | |
step=1, | |
value=0, | |
) | |
randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
with gr.Row(): | |
width = gr.Slider( | |
label="Width", | |
minimum=256, | |
maximum=MAX_IMAGE_SIZE, | |
step=32, | |
value=1024, | |
) | |
height = gr.Slider( | |
label="Height", | |
minimum=256, | |
maximum=MAX_IMAGE_SIZE, | |
step=32, | |
value=1024, | |
) | |
with gr.Row(): | |
guidance_scale_base = gr.Slider( | |
info="Scale for classifier-free guidance", | |
label="Guidance scale", | |
minimum=1, | |
maximum=20, | |
step=0.1, | |
value=5.0, | |
) | |
with gr.Row(): | |
num_inference_steps_base = gr.Slider( | |
info="Number of denoising steps", | |
label="Number of inference steps", | |
minimum=10, | |
maximum=100, | |
step=1, | |
value=25, | |
) | |
with gr.Row(): | |
controlnet_conditioning_scale = gr.Slider( | |
info="controlnet_conditioning_scale", | |
label="controlnet_conditioning_scale", | |
minimum=0.01, | |
maximum=2, | |
step=0.01, | |
value=1, | |
) | |
with gr.Row(): | |
control_guidance_start = gr.Slider( | |
info="control_guidance_start", | |
label="control_guidance_start", | |
minimum=0.01, | |
maximum=1, | |
step=0.01, | |
value=0, | |
) | |
with gr.Row(): | |
control_guidance_end = gr.Slider( | |
info="control_guidance_end", | |
label="control_guidance_end", | |
minimum=0.01, | |
maximum=1, | |
step=0.01, | |
value=1, | |
) | |
with gr.Row(): | |
strength_img2img = gr.Slider( | |
info="Strength for Img2Img", | |
label="Strength", | |
minimum=0, | |
maximum=1, | |
step=0.01, | |
value=0.7, | |
) | |
use_negative_prompt.change( | |
fn=lambda x: gr.update(visible=x), | |
inputs=use_negative_prompt, | |
outputs=negative_prompt, | |
queue=False, | |
api_name=False, | |
) | |
use_prompt_2.change( | |
fn=lambda x: gr.update(visible=x), | |
inputs=use_prompt_2, | |
outputs=prompt_2, | |
queue=False, | |
api_name=False, | |
) | |
use_negative_prompt_2.change( | |
fn=lambda x: gr.update(visible=x), | |
inputs=use_negative_prompt_2, | |
outputs=negative_prompt_2, | |
queue=False, | |
api_name=False, | |
) | |
use_vae.change( | |
fn=lambda x: gr.update(visible=x), | |
inputs=use_vae, | |
outputs=vaecall, | |
queue=False, | |
api_name=False, | |
) | |
use_lora.change( | |
fn=lambda x: [gr.update(visible=x), gr.update(visible=x)], | |
inputs=use_lora, | |
outputs=[lora, lora_scale], | |
queue=False, | |
api_name=False, | |
) | |
use_lora2.change( | |
fn=lambda x: [gr.update(visible=x), gr.update(visible=x), gr.update(visible=x), gr.update(visible=x)], | |
inputs=use_lora2, | |
outputs=[lora, lora2, lora_scale, lora_scale2], | |
queue=False, | |
api_name=False, | |
) | |
use_img2img.change( | |
fn=lambda x: gr.update(visible=x), | |
inputs=use_img2img, | |
outputs=url, | |
queue=False, | |
api_name=False, | |
) | |
use_controlnet.change( | |
fn=lambda x: gr.update(visible=x), | |
inputs=use_controlnet, | |
outputs=controlnet_img, | |
queue=False, | |
api_name=False, | |
) | |
use_controlnetimg2img.change( | |
fn=lambda x: gr.update(visible=x), | |
inputs=use_controlnetimg2img, | |
outputs=controlnet_img2img, | |
queue=False, | |
api_name=False, | |
) | |
gr.on( | |
triggers=[ | |
prompt.submit, | |
negative_prompt.submit, | |
prompt_2.submit, | |
negative_prompt_2.submit, | |
run_button.click, | |
], | |
fn=randomize_seed_fn, | |
inputs=[seed, randomize_seed], | |
outputs=seed, | |
queue=False, | |
api_name=False, | |
).then( | |
fn=generate, | |
inputs=[ | |
prompt, | |
negative_prompt, | |
prompt_2, | |
negative_prompt_2, | |
use_negative_prompt, | |
use_prompt_2, | |
use_negative_prompt_2, | |
seed, | |
width, | |
height, | |
guidance_scale_base, | |
num_inference_steps_base, | |
controlnet_conditioning_scale, | |
control_guidance_start, | |
control_guidance_end, | |
strength_img2img, | |
use_vae, | |
use_lora, | |
use_lora2, | |
model, | |
vaecall, | |
lora, | |
lora2, | |
controlnet_model, | |
lora_scale, | |
lora_scale2, | |
use_img2img, | |
use_controlnet, | |
use_controlnetimg2img, | |
url, | |
controlnet_img, | |
controlnet_img2img, | |
], | |
outputs=result, | |
api_name="run", | |
) | |
if __name__ == "__main__": | |
demo.queue(max_size=20).launch() |