chunzhi/stable-diffusion-cli-on-modal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactoring.

Browse Source
This commit is contained in:
hodanov
2023-07-02 22:44:32 +09:00
parent 11e2a6b790
commit d810577f3b
7 changed files with 202 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
setup_files/.env.example
+3
View File
@@ -0,0 +1,3 @@
# `HUGGING_FACE_TOKEN` is the token for the Hugging Face API.
# The token can be found at https://huggingface.co/settings/token.
HUGGING_FACE_TOKEN=""
setup_files/Dockerfile
+13
View File
@@ -0,0 +1,13 @@
FROM python:3.11.3-slim-bullseye
COPY ./requirements.txt /
RUN apt update \
&& apt install -y wget git libgl1-mesa-glx libglib2.0-0 \
&& apt autoremove -y \
&& apt clean -y \
&& pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117 \
&& mkdir -p /vol/cache/esrgan \
&& wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /vol/cache/esrgan \
&& wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth -P /vol/cache/esrgan \
&& wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P /vol/cache/esrgan \
&& wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth -P /vol/cache/esrgan \
&& wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P /vol/cache/esrgan
setup_files/config.sample.yml
+38
View File
@@ -0,0 +1,38 @@
##########
# This is the config file to set a base model, vae and some tools.
# Rename the file to `config.yml` before running the script.
# Execute `modal deploy ./setup_files/setup.py` every time modify this file.
##########
##########
# You can use a diffusers model and VAE on hugging face.
model:
name: stable-diffusion-2-1
repo_id: stabilityai/stable-diffusion-2-1
vae:
name: sd-vae-ft-mse
repo_id: stabilityai/sd-vae-ft-mse
##########
# Add LoRA if you want to use one. You can use a download url such as the below.
# ex)
# loras:
# - name: hogehoge.safetensors
# download_url: https://hogehoge/xxxx
# - name: fugafuga.safetensors
# download_url: https://fugafuga/xxxx
##########
# You can use Textual Inversion and ControlNet also. Usage is the same as `loras`.
# ex)
# textual_inversions:
# - name: hogehoge
# download_url: https://hogehoge/xxxx
# - name: fugafuga
# download_url: https://fugafuga/xxxx
# cotrolnets:
# - name: control_v11f1e_sd15_tile
# repo_id: lllyasviel/control_v11f1e_sd15_tile
# upscaler:
# name: RealESRGAN_x2plus
# use_face_enhancer: false
# use_hires_fix: false
setup_files/requirements.txt
+20
View File
@@ -0,0 +1,20 @@
accelerate
diffusers[torch]==0.17.1
onnxruntime==1.15.1
safetensors==0.3.1
torch==2.0.1+cu117
transformers==4.30.2
xformers==0.0.20
realesrgan==0.3.0
basicsr>=1.4.2
facexlib>=0.3.0
gfpgan>=1.3.8
numpy
opencv-python
Pillow
torchvision
tqdm
controlnet_aux
pyyaml
setup_files/setup.py
+422
View File
@@ -0,0 +1,422 @@
from __future__ import annotations
import io
import os
from urllib.request import Request, urlopen
import diffusers
import yaml
from modal import Image, Mount, Secret, Stub, method
from modal.cls import ClsMixin
BASE_CACHE_PATH = "/vol/cache"
BASE_CACHE_PATH_LORA = "/vol/cache/lora"
BASE_CACHE_PATH_TEXTUAL_INVERSION = "/vol/cache/textual_inversion"
BASE_CACHE_PATH_CONTROLNET = "/vol/cache/controlnet"
def download_file(url, file_name, file_path):
"""
Download files.
"""
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_controlnet(name: str, repo_id: str, token: str):
"""
Download a controlnet.
"""
cache_path = os.path.join(BASE_CACHE_PATH_CONTROLNET, name)
controlnet = diffusers.ControlNetModel.from_pretrained(
repo_id,
use_auth_token=token,
cache_dir=cache_path,
)
controlnet.save_pretrained(cache_path, safe_serialization=True)
def download_vae(name: str, repo_id: str, token: str):
"""
Download a vae.
"""
cache_path = os.path.join(BASE_CACHE_PATH, name)
vae = diffusers.AutoencoderKL.from_pretrained(
repo_id,
use_auth_token=token,
cache_dir=cache_path,
)
vae.save_pretrained(cache_path, safe_serialization=True)
def download_model(name: str, repo_id: str, token: str):
"""
Download a model.
"""
cache_path = os.path.join(BASE_CACHE_PATH, name)
pipe = diffusers.StableDiffusionPipeline.from_pretrained(
repo_id,
use_auth_token=token,
cache_dir=cache_path,
)
pipe.save_pretrained(cache_path, safe_serialization=True)
def build_image():
"""
Build the Docker image.
"""
token = os.environ["HUGGING_FACE_TOKEN"]
config = {}
with open("/config.yml", "r") as file:
config = yaml.safe_load(file)
model = config.get("model")
if model is not None:
download_model(name=model["name"], repo_id=model["repo_id"], token=token)
vae = config.get("vae")
if vae is not None:
download_vae(name=model["name"], repo_id=vae["repo_id"], token=token)
controlnets = config.get("controlnets")
if controlnets is not None:
for controlnet in controlnets:
download_controlnet(name=controlnet["name"], repo_id=controlnet["repo_id"], token=token)
loras = config.get("loras")
if loras is not None:
for lora in loras:
download_file(
url=lora["download_url"],
file_name=lora["name"],
file_path=BASE_CACHE_PATH_LORA,
)
textual_inversions = config.get("textual_inversions")
if textual_inversions is not None:
for textual_inversion in textual_inversions:
download_file(
url=textual_inversion["download_url"],
file_name=textual_inversion["name"],
file_path=BASE_CACHE_PATH_TEXTUAL_INVERSION,
)
stub = Stub("stable-diffusion-cli")
base_stub = Image.from_dockerfile(
path="./setup_files/Dockerfile",
context_mount=Mount.from_local_file("./setup_files/requirements.txt"),
)
stub.image = base_stub.extend(
dockerfile_commands=[
"FROM base",
"COPY ./config.yml /",
],
context_mount=Mount.from_local_file("./setup_files/config.yml"),
).run_function(
build_image,
secrets=[Secret.from_dotenv(__file__)],
)
@stub.cls(
gpu="A10G",
secrets=[Secret.from_dotenv(__file__)],
)
class StableDiffusion(ClsMixin):
"""
A class that wraps the Stable Diffusion pipeline and scheduler.
"""
def __enter__(self):
import torch
config = {}
with open("/config.yml", "r") as file:
config = yaml.safe_load(file)
self.cache_path = os.path.join(BASE_CACHE_PATH, config["model"]["name"])
if os.path.exists(self.cache_path):
print(f"The directory '{self.cache_path}' exists.")
else:
print(f"The directory '{self.cache_path}' does not exist.")
torch.cuda.memory._set_allocator_settings("max_split_size_mb:256")
self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
self.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(
self.cache_path,
subfolder="scheduler",
)
vae = config.get("vae")
if vae is not None:
self.pipe.vae = diffusers.AutoencoderKL.from_pretrained(
self.cache_path,
subfolder="vae",
)
self.pipe.to("cuda")
loras = config.get("loras")
if loras is not None:
for lora in loras:
path = os.path.join(BASE_CACHE_PATH_LORA, 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_file(lora["download_url"], lora["name"], BASE_CACHE_PATH_LORA)
self.pipe.load_lora_weights(".", weight_name=path)
textual_inversions = config.get("textual_inversions")
if textual_inversions is not None:
for textual_inversion in textual_inversions:
path = os.path.join(BASE_CACHE_PATH_TEXTUAL_INVERSION, 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_file(
textual_inversion["download_url"],
textual_inversion["name"],
BASE_CACHE_PATH_TEXTUAL_INVERSION,
)
self.pipe.load_textual_inversion(path)
self.pipe.enable_xformers_memory_efficient_attention()
# TODO: Add support for controlnets.
# controlnet = diffusers.ControlNetModel.from_pretrained(
# "lllyasviel/control_v11f1e_sd15_tile",
# # "lllyasviel/sd-controlnet-canny",
# # self.cache_path,
# # subfolder="controlnet",
# torch_dtype=torch.float16,
# )
# self.controlnet_pipe = diffusers.StableDiffusionControlNetPipeline.from_pretrained(
# self.cache_path,
# controlnet=controlnet,
# custom_pipeline="lpw_stable_diffusion",
# # custom_pipeline="stable_diffusion_controlnet_img2img",
# scheduler=self.pipe.scheduler,
# vae=self.pipe.vae,
# torch_dtype=torch.float16,
# )
# self.controlnet_pipe.to("cuda")
# self.controlnet_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(
self.cache_path,
subfolder="tokenizer",
)
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,
prompt: str,
n_prompt: str,
height: int = 512,
width: int = 512,
samples: int = 1,
batch_size: int = 1,
steps: int = 30,
seed: int = 1,
upscaler: str = "",
use_face_enhancer: bool = False,
use_hires_fix: bool = False,
) -> list[bytes]:
"""
Runs the Stable Diffusion pipeline on the given prompt and outputs images.
"""
import torch
max_embeddings_multiples = self.count_token(p=prompt, n=n_prompt)
generator = torch.Generator("cuda").manual_seed(seed)
with torch.inference_mode():
with torch.autocast("cuda"):
base_images = self.pipe.text2img(
prompt * batch_size,
negative_prompt=n_prompt * batch_size,
height=height,
width=width,
num_inference_steps=steps,
guidance_scale=7.5,
max_embeddings_multiples=max_embeddings_multiples,
generator=generator,
).images
# for image in base_images:
# image = self.resize_image(image=image, scale_factor=2)
# with torch.inference_mode():
# with torch.autocast("cuda"):
# generatedWithControlnet = self.controlnet_pipe(
# prompt=prompt * batch_size,
# negative_prompt=n_prompt * batch_size,
# num_inference_steps=steps,
# strength=0.3,
# guidance_scale=7.5,
# max_embeddings_multiples=max_embeddings_multiples,
# generator=generator,
# image=image,
# ).images
# base_images.extend(generatedWithControlnet)
if upscaler != "":
upscaled = self.upscale(
base_images=base_images,
half_precision=False,
tile=700,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer,
use_hires_fix=use_hires_fix,
)
base_images.extend(upscaled)
if use_hires_fix:
for img in upscaled:
with torch.inference_mode():
with torch.autocast("cuda"):
hires_fixed = self.pipe.img2img(
prompt=prompt * batch_size,
negative_prompt=n_prompt * batch_size,
num_inference_steps=steps,
strength=0.3,
guidance_scale=7.5,
max_embeddings_multiples=max_embeddings_multiples,
generator=generator,
image=img,
).images
base_images.extend(hires_fixed)
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 resize_image(self, image: Image.Image, scale_factor: int) -> Image.Image:
from PIL import Image
image = image.convert("RGB")
width, height = image.size
img = image.resize((width * scale_factor, height * scale_factor), resample=Image.LANCZOS)
return img
@method()
def upscale(
self,
base_images: list[Image.Image],
half_precision: bool = False,
tile: int = 0,
tile_pad: int = 10,
pre_pad: int = 0,
upscaler: str = "",
use_face_enhancer: bool = False,
use_hires_fix: bool = False,
) -> list[Image.Image]:
"""
Upscales the given images using a upscaler.
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 = 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 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,
)
upscaled_imgs = []
with tqdm(total=len(base_images)) as progress_bar:
for img in base_images:
img = numpy.array(img)
if 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)
return upscaled_imgs