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)