347 lines
11 KiB
Python
347 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import os
|
|
import time
|
|
from urllib.request import Request, urlopen
|
|
|
|
from modal import Image, Mount, Secret, Stub, method
|
|
|
|
BASE_CACHE_PATH = "/vol/cache"
|
|
BASE_CACHE_PATH_LORA = "/vol/cache/lora"
|
|
BASE_CACHE_PATH_TEXTUAL_INVERSION = "/vol/cache/textual_inversion"
|
|
|
|
|
|
def download_files(urls, file_names, file_path):
|
|
"""
|
|
Download files.
|
|
"""
|
|
file_names = file_names.split(",")
|
|
urls = urls.split(",")
|
|
|
|
for file_name, url in zip(file_names, urls):
|
|
req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
|
|
downloaded = urlopen(req).read()
|
|
|
|
dir_names = os.path.join(file_path, file_name)
|
|
os.makedirs(os.path.dirname(dir_names), exist_ok=True)
|
|
with open(dir_names, mode="wb") as f:
|
|
f.write(downloaded)
|
|
|
|
|
|
def download_models():
|
|
"""
|
|
Downloads the model from Hugging Face and saves it to the cache path using
|
|
diffusers.StableDiffusionPipeline.from_pretrained().
|
|
"""
|
|
import diffusers
|
|
|
|
hugging_face_token = os.environ["HUGGING_FACE_TOKEN"]
|
|
model_repo_id = os.environ["MODEL_REPO_ID"]
|
|
cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])
|
|
|
|
vae = diffusers.AutoencoderKL.from_pretrained(
|
|
"stabilityai/sd-vae-ft-mse",
|
|
use_auth_token=hugging_face_token,
|
|
cache_dir=cache_path,
|
|
)
|
|
vae.save_pretrained(cache_path, safe_serialization=True)
|
|
|
|
pipe = diffusers.StableDiffusionPipeline.from_pretrained(
|
|
model_repo_id,
|
|
use_auth_token=hugging_face_token,
|
|
cache_dir=cache_path,
|
|
)
|
|
pipe.save_pretrained(cache_path, safe_serialization=True)
|
|
|
|
|
|
def build_image():
|
|
"""
|
|
Build the Docker image.
|
|
"""
|
|
download_models()
|
|
|
|
if os.environ["LORA_NAMES"] != "":
|
|
download_files(
|
|
os.getenv("LORA_DOWNLOAD_URLS"),
|
|
os.getenv("LORA_NAMES"),
|
|
BASE_CACHE_PATH_LORA,
|
|
)
|
|
|
|
if os.environ["TEXTUAL_INVERSION_NAMES"] != "":
|
|
download_files(
|
|
os.getenv("TEXTUAL_INVERSION_DOWNLOAD_URLS"),
|
|
os.getenv("TEXTUAL_INVERSION_NAMES"),
|
|
BASE_CACHE_PATH_TEXTUAL_INVERSION,
|
|
)
|
|
|
|
|
|
stub_image = Image.from_dockerfile(
|
|
path="./Dockerfile",
|
|
context_mount=Mount.from_local_file("./requirements.txt"),
|
|
).run_function(
|
|
build_image,
|
|
secrets=[Secret.from_dotenv(__file__)],
|
|
)
|
|
stub = Stub("stable-diffusion-cli")
|
|
stub.image = stub_image
|
|
|
|
|
|
@stub.cls(gpu="A10G", secrets=[Secret.from_dotenv(__file__)])
|
|
class StableDiffusion:
|
|
"""
|
|
A class that wraps the Stable Diffusion pipeline and scheduler.
|
|
"""
|
|
|
|
def __enter__(self):
|
|
import diffusers
|
|
import torch
|
|
|
|
use_vae = os.environ["USE_VAE"] == "true"
|
|
self.upscaler = os.environ["UPSCALER"]
|
|
self.use_face_enhancer = os.environ["USE_FACE_ENHANCER"] == "true"
|
|
|
|
cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])
|
|
if os.path.exists(cache_path):
|
|
print(f"The directory '{cache_path}' exists.")
|
|
else:
|
|
print(f"The directory '{cache_path}' does not exist. Download models...")
|
|
download_models()
|
|
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
|
|
self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
|
|
cache_path,
|
|
custom_pipeline="lpw_stable_diffusion",
|
|
torch_dtype=torch.float16,
|
|
)
|
|
|
|
# TODO: Add support for other schedulers.
|
|
# self.pipe.scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
|
|
self.pipe.scheduler = diffusers.DPMSolverMultistepScheduler.from_pretrained(
|
|
cache_path,
|
|
subfolder="scheduler",
|
|
)
|
|
|
|
if use_vae:
|
|
self.pipe.vae = diffusers.AutoencoderKL.from_pretrained(
|
|
cache_path,
|
|
subfolder="vae",
|
|
)
|
|
|
|
self.pipe.to("cuda")
|
|
|
|
if os.environ["LORA_NAMES"] != "":
|
|
names = os.getenv("LORA_NAMES").split(",")
|
|
urls = os.getenv("LORA_DOWNLOAD_URLS").split(",")
|
|
for name, url in zip(names, urls):
|
|
path = os.path.join(BASE_CACHE_PATH_LORA, name)
|
|
if os.path.exists(path):
|
|
print(f"The directory '{path}' exists.")
|
|
else:
|
|
print(f"The directory '{path}' does not exist. Download it...")
|
|
download_files(url, name, BASE_CACHE_PATH_LORA)
|
|
self.pipe.load_lora_weights(".", weight_name=path)
|
|
|
|
if os.environ["TEXTUAL_INVERSION_NAMES"] != "":
|
|
names = os.getenv("TEXTUAL_INVERSION_NAMES").split(",")
|
|
urls = os.getenv("TEXTUAL_INVERSION_DOWNLOAD_URLS").split(",")
|
|
for name, url in zip(names, urls):
|
|
path = os.path.join(BASE_CACHE_PATH_TEXTUAL_INVERSION, name)
|
|
if os.path.exists(path):
|
|
print(f"The directory '{path}' exists.")
|
|
else:
|
|
print(f"The directory '{path}' does not exist. Download it...")
|
|
download_files(url, name, BASE_CACHE_PATH_TEXTUAL_INVERSION)
|
|
self.pipe.load_textual_inversion(path)
|
|
|
|
self.pipe.enable_xformers_memory_efficient_attention()
|
|
|
|
@method()
|
|
def count_token(self, p: str, n: str) -> int:
|
|
"""
|
|
Count the number of tokens in the prompt and negative prompt.
|
|
"""
|
|
from transformers import CLIPTokenizer
|
|
|
|
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
|
|
token_size_p = len(tokenizer.tokenize(p))
|
|
token_size_n = len(tokenizer.tokenize(n))
|
|
token_size = token_size_p
|
|
if token_size_p <= token_size_n:
|
|
token_size = token_size_n
|
|
|
|
max_embeddings_multiples = 1
|
|
max_length = tokenizer.model_max_length - 2
|
|
if token_size > max_length:
|
|
max_embeddings_multiples = token_size // max_length + 1
|
|
|
|
print(f"token_size: {token_size}, max_embeddings_multiples: {max_embeddings_multiples}")
|
|
|
|
return max_embeddings_multiples
|
|
|
|
@method()
|
|
def run_inference(self, inputs: dict[str, int | str]) -> list[bytes]:
|
|
"""
|
|
Runs the Stable Diffusion pipeline on the given prompt and outputs images.
|
|
"""
|
|
import torch
|
|
|
|
generator = torch.Generator("cuda").manual_seed(inputs["seed"])
|
|
with torch.inference_mode():
|
|
with torch.autocast("cuda"):
|
|
base_images = self.pipe.text2img(
|
|
[inputs["prompt"]] * int(inputs["batch_size"]),
|
|
negative_prompt=[inputs["n_prompt"]] * int(inputs["batch_size"]),
|
|
height=inputs["height"],
|
|
width=inputs["width"],
|
|
num_inference_steps=inputs["steps"],
|
|
guidance_scale=7.5,
|
|
max_embeddings_multiples=inputs["max_embeddings_multiples"],
|
|
generator=generator,
|
|
).images
|
|
|
|
if self.upscaler != "":
|
|
uplcaled_images = self.upscale(
|
|
base_images=base_images,
|
|
scale_factor=4,
|
|
half_precision=False,
|
|
tile=700,
|
|
)
|
|
base_images.extend(uplcaled_images)
|
|
|
|
image_output = []
|
|
for image in base_images:
|
|
with io.BytesIO() as buf:
|
|
image.save(buf, format="PNG")
|
|
image_output.append(buf.getvalue())
|
|
|
|
return image_output
|
|
|
|
@method()
|
|
def upscale(
|
|
self,
|
|
base_images: list[Image.Image],
|
|
scale_factor: float = 4,
|
|
half_precision: bool = False,
|
|
tile: int = 0,
|
|
tile_pad: int = 10,
|
|
pre_pad: int = 0,
|
|
) -> list[Image.Image]:
|
|
"""
|
|
Upscales the given images using the given model.
|
|
https://github.com/xinntao/Real-ESRGAN
|
|
"""
|
|
import numpy
|
|
import torch
|
|
from basicsr.archs.rrdbnet_arch import RRDBNet
|
|
from PIL import Image
|
|
from realesrgan import RealESRGANer
|
|
from tqdm import tqdm
|
|
|
|
model_name = self.upscaler
|
|
if model_name == "RealESRGAN_x4plus":
|
|
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
|
|
netscale = 4
|
|
elif model_name == "RealESRNet_x4plus":
|
|
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
|
|
netscale = 4
|
|
elif model_name == "RealESRGAN_x4plus_anime_6B":
|
|
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
|
|
netscale = 4
|
|
elif model_name == "RealESRGAN_x2plus":
|
|
upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
|
|
netscale = 2
|
|
else:
|
|
raise NotImplementedError("Model name not supported")
|
|
|
|
upsampler = RealESRGANer(
|
|
scale=netscale,
|
|
model_path=os.path.join(BASE_CACHE_PATH, "esrgan", f"{model_name}.pth"),
|
|
dni_weight=None,
|
|
model=upscale_model,
|
|
tile=tile,
|
|
tile_pad=tile_pad,
|
|
pre_pad=pre_pad,
|
|
half=half_precision,
|
|
gpu_id=None,
|
|
)
|
|
|
|
from gfpgan import GFPGANer
|
|
|
|
if self.use_face_enhancer:
|
|
face_enhancer = GFPGANer(
|
|
model_path=os.path.join(BASE_CACHE_PATH, "esrgan", "GFPGANv1.3.pth"),
|
|
upscale=netscale,
|
|
arch="clean",
|
|
channel_multiplier=2,
|
|
bg_upsampler=upsampler,
|
|
)
|
|
|
|
torch.cuda.empty_cache()
|
|
upscaled_imgs = []
|
|
with tqdm(total=len(base_images)) as progress_bar:
|
|
for i, img in enumerate(base_images):
|
|
img = numpy.array(img)
|
|
if self.use_face_enhancer:
|
|
_, _, enhance_result = face_enhancer.enhance(
|
|
img,
|
|
has_aligned=False,
|
|
only_center_face=False,
|
|
paste_back=True,
|
|
)
|
|
else:
|
|
enhance_result, _ = upsampler.enhance(img)
|
|
|
|
upscaled_imgs.append(Image.fromarray(enhance_result))
|
|
progress_bar.update(1)
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
return upscaled_imgs
|
|
|
|
|
|
@stub.local_entrypoint()
|
|
def entrypoint(
|
|
prompt: str,
|
|
n_prompt: str,
|
|
height: int = 512,
|
|
width: int = 512,
|
|
samples: int = 5,
|
|
batch_size: int = 1,
|
|
steps: int = 20,
|
|
seed: int = -1,
|
|
):
|
|
"""
|
|
This function is the entrypoint for the Runway CLI.
|
|
The function pass the given prompt to StableDiffusion on Modal,
|
|
gets back a list of images and outputs images to local.
|
|
"""
|
|
import util
|
|
|
|
inputs: dict[str, int | str] = {
|
|
"prompt": prompt,
|
|
"n_prompt": n_prompt,
|
|
"height": height,
|
|
"width": width,
|
|
"samples": samples,
|
|
"batch_size": batch_size,
|
|
"steps": steps,
|
|
"seed": seed,
|
|
}
|
|
|
|
directory = util.make_directory()
|
|
|
|
sd = StableDiffusion()
|
|
inputs["max_embeddings_multiples"] = sd.count_token(p=prompt, n=n_prompt)
|
|
for i in range(samples):
|
|
if seed == -1:
|
|
inputs["seed"] = util.generate_seed()
|
|
start_time = time.time()
|
|
images = sd.run_inference.call(inputs)
|
|
util.save_images(directory, images, int(inputs["seed"]), i)
|
|
total_time = time.time() - start_time
|
|
print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).")
|
|
|
|
util.save_prompts(inputs)
|