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

Implement SDXLTxt2Img.

Browse Source
This commit is contained in:
hodanov 2023-12-10 16:46:03 +09:00
parent 8a0a28e999
commit 6b522e20eb
7 changed files with 277 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
Makefile
View File

@ -1,4 +1,4 @@
deploy:
app:
cd ./setup_files && modal deploy __main__.py
# `--upscaler` is a name of upscaler you want to use.
@ -7,8 +7,8 @@ deploy:
# - `RealESRNet_x4plus`
# - `RealESRGAN_x4plus_anime_6B`
# - `RealESRGAN_x2plus`
run:
cd ./sdcli && modal run txt2img.py \
img_by_sd15_txt2img:
cd ./sdcli && modal run sd15_txt2img.py \
--prompt "a photograph of an astronaut riding a horse" \
--n-prompt "" \
--height 512 \
@ -17,4 +17,15 @@ run:
--steps 30 \
--upscaler "RealESRGAN_x2plus" \
--use-face-enhancer "False" \
--fix-by-controlnet-tile "True"
--fix-by-controlnet-tile "True" \
--output-format "avif"
img_by_sdxl_txt2img:
cd ./sdcli && modal run sdxl_txt2img.py \
--prompt "A dog is running on the grass" \
--height 1024 \
--width 1024 \
--samples 1 \
--upscaler "RealESRGAN_x2plus" \
--output-format "avif"

2
sdcli/txt2img.py → sdcli/sd15_txt2img.py
View File

@ -4,7 +4,7 @@ import modal
import util
stub = modal.Stub("run-stable-diffusion-cli")
stub.run_inference = modal.Function.from_name("stable-diffusion-cli", "Txt2Img.run_inference")
stub.run_inference = modal.Function.from_name("stable-diffusion-cli", "SD15Txt2Img.run_inference")
@stub.local_entrypoint()

51
sdcli/sdxl_txt2img.py Normal file
View File

@ -0,0 +1,51 @@
import time
import modal
import util
stub = modal.Stub("run-stable-diffusion-cli")
stub.run_inference = modal.Function.from_name("stable-diffusion-cli", "SDXLTxt2Img.run_inference")
@stub.local_entrypoint()
def main(
prompt: str,
height: int = 1024,
width: int = 1024,
samples: int = 5,
seed: int = -1,
upscaler: str = "",
use_face_enhancer: str = "False",
output_format: str = "png",
):
"""
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.
"""
directory = util.make_directory()
seed_generated = seed
for i in range(samples):
if seed == -1:
seed_generated = util.generate_seed()
start_time = time.time()
images = stub.run_inference.remote(
prompt=prompt,
height=height,
width=width,
seed=seed_generated,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer == "True",
output_format=output_format,
)
util.save_images(directory, images, seed_generated, i, output_format)
total_time = time.time() - start_time
print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).")
prompts: dict[str, int | str] = {
"prompt": prompt,
"height": height,
"width": width,
"samples": samples,
}
util.save_prompts(prompts)

6
setup_files/__main__.py
View File

@ -1,12 +1,14 @@
from __future__ import annotations
import stable_diffusion_1_5
import stable_diffusion_xl
from setup import stub
from stable_diffusion_1_5 import Txt2Img
@stub.function(gpu="A10G")
def main():
Txt2Img
stable_diffusion_1_5.SD15Txt2Img
stable_diffusion_xl.SDXLTxt2Img
if __name__ == "__main__":

26
setup_files/setup.py
View File

@ -64,6 +64,26 @@ def download_model(name: str, model_url: str, token: str):
pipe.save_pretrained(cache_path, safe_serialization=True)
def download_model_sdxl(name: str, model_url: str, token: str):
"""
Download a sdxl model.
"""
cache_path = os.path.join(BASE_CACHE_PATH, name)
pipe = diffusers.StableDiffusionXLPipeline.from_single_file(
pretrained_model_link_or_path=model_url,
use_auth_token=token,
cache_dir=cache_path,
)
pipe.save_pretrained(cache_path, safe_serialization=True)
refiner_cache_path = cache_path + "-refiner"
refiner = diffusers.StableDiffusionXLImg2ImgPipeline.from_single_file(
"https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/main/sd_xl_refiner_1.0.safetensors",
cache_dir=refiner_cache_path,
)
refiner.save_pretrained(refiner_cache_path, safe_serialization=True)
def build_image():
"""
Build the Docker image.
@ -76,8 +96,12 @@ def build_image():
config = yaml.safe_load(file)
model = config.get("model")
use_xl = config.get("use_xl")
if model is not None:
download_model(name=model["name"], model_url=model["url"], token=token)
if use_xl is not None and use_xl:
download_model_sdxl(name=model["name"], model_url=model["url"], token=token)
else:
download_model(name=model["name"], model_url=model["url"], token=token)
vae = config.get("vae")
if vae is not None:

2
setup_files/stable_diffusion_1_5.py
View File

@ -18,7 +18,7 @@ from setup import (
gpu="A10G",
secrets=[Secret.from_dotenv(__file__)],
)
class Txt2Img:
class SD15Txt2Img:
"""
A class that wraps the Stable Diffusion pipeline and scheduler.
"""

180
setup_files/stable_diffusion_xl.py Normal file
View File

@ -0,0 +1,180 @@
from __future__ import annotations
import io
import os
import PIL.Image
from modal import Secret, method
from setup import BASE_CACHE_PATH, stub
@stub.cls(
gpu="A10G",
secrets=[Secret.from_dotenv(__file__)],
)
class SDXLTxt2Img:
"""
A class that wraps the Stable Diffusion pipeline and scheduler.
"""
def __enter__(self):
import diffusers
import torch
import yaml
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.")
self.pipe = diffusers.AutoPipelineForText2Image.from_pretrained(
self.cache_path,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
)
self.refiner_cache_path = self.cache_path + "-refiner"
self.refiner = diffusers.StableDiffusionXLImg2ImgPipeline.from_pretrained(
self.refiner_cache_path,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
)
@method()
def run_inference(
self,
prompt: str,
height: int = 1024,
width: int = 1024,
seed: int = 1,
upscaler: str = "",
use_face_enhancer: bool = False,
output_format: str = "png",
) -> list[bytes]:
"""
Runs the Stable Diffusion pipeline on the given prompt and outputs images.
"""
import pillow_avif # noqa
import torch
generator = torch.Generator("cuda").manual_seed(seed)
self.pipe.to("cuda")
generated_images = self.pipe(
prompt=prompt,
height=height,
width=width,
generator=generator,
).images
base_images = generated_images
for image in base_images:
self.refiner.to("cuda")
refined_images = self.refiner(
prompt=prompt,
image=image,
).images
generated_images.extend(refined_images)
base_images = refined_images
if upscaler != "":
upscaled = self._upscale(
base_images=base_images,
half_precision=False,
tile=700,
upscaler=upscaler,
use_face_enhancer=use_face_enhancer,
)
generated_images.extend(upscaled)
image_output = []
for image in generated_images:
with io.BytesIO() as buf:
image.save(buf, format=output_format)
image_output.append(buf.getvalue())
return image_output
def _upscale(
self,
base_images: list[PIL.Image],
half_precision: bool = False,
tile: int = 0,
tile_pad: int = 10,
pre_pad: int = 0,
upscaler: str = "",
use_face_enhancer: bool = False,
) -> list[PIL.Image]:
"""
Upscale the generated images by the upscaler when `upscaler` is selected.
The upscaler can be selected from the following list:
- `RealESRGAN_x4plus`
- `RealESRNet_x4plus`
- `RealESRGAN_x4plus_anime_6B`
- `RealESRGAN_x2plus`
https://github.com/xinntao/Real-ESRGAN
"""
import numpy
from basicsr.archs.rrdbnet_arch import RRDBNet
from gfpgan import GFPGANer
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,
)
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 = []
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(PIL.Image.fromarray(enhance_result))
return upscaled_imgs