From 89f84c6a009232ff8416dd0df73bab43edcc4f01 Mon Sep 17 00:00:00 2001 From: hodanov <1031hoda@gmail.com> Date: Sun, 21 May 2023 17:52:48 +0900 Subject: [PATCH] Add sd_cli.py. --- sd_cli.py | 192 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 192 insertions(+) create mode 100644 sd_cli.py diff --git a/sd_cli.py b/sd_cli.py new file mode 100644 index 0000000..0fd97c7 --- /dev/null +++ b/sd_cli.py @@ -0,0 +1,192 @@ +from __future__ import annotations +import io +import os +import time +from datetime import date +from pathlib import Path +from modal import Image, Secret, Stub, method + +stub = Stub("stable-diffusion-cli") + +MODEL = { + "repo_id": "runwayml/stable-diffusion-v1-5", + "name": "stable-diffusion-v1-5", +} +CACHE_PATH = os.path.join("/vol/cache", MODEL["name"]) + + +def download_models(): + """ + Downloads the model from Hugging Face and saves it to the cache path using + diffusers.StableDiffusionPipeline.from_pretrained(). + """ + import diffusers + import torch + + hugging_face_token = os.environ["HUGGINGFACE_TOKEN"] + + scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained( + MODEL["repo_id"], + subfolder="scheduler", + use_auth_token=hugging_face_token, + cache_dir=CACHE_PATH, + ) + scheduler.save_pretrained(CACHE_PATH, safe_serialization=True) + + pipe = diffusers.StableDiffusionPipeline.from_pretrained( + MODEL["repo_id"], + use_auth_token=hugging_face_token, + torch_dtype=torch.float16, + cache_dir=CACHE_PATH, + ) + pipe.save_pretrained(CACHE_PATH, safe_serialization=True) + + +stub_image = ( + Image.debian_slim(python_version="3.10") + .pip_install( + "accelerate", + "diffusers[torch]>=0.15.1", + "ftfy", + "torch", + "torchvision", + "transformers~=4.25.1", + "triton", + "safetensors", + "torch>=2.0", + ) + .pip_install("xformers", pre=True) + .run_function( + download_models, + secrets=[Secret.from_name("my-huggingface-secret")], + ) +) +stub.image = stub_image + + +@stub.cls(gpu="A10G", secrets=[Secret.from_name("my-huggingface-secret")]) +class StableDiffusion: + """ + A class that wraps the Stable Diffusion pipeline and scheduler. + """ + + def __enter__(self): + import diffusers + import torch + + 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 + + scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained( + CACHE_PATH, + subfolder="scheduler", + solver_order=2, + prediction_type="epsilon", + thresholding=False, + algorithm_type="dpmsolver++", + solver_type="midpoint", + denoise_final=True, # important if steps are <= 10 + low_cpu_mem_usage=True, + device_map="auto", + ) + + self.pipe = diffusers.StableDiffusionPipeline.from_pretrained( + CACHE_PATH, + scheduler=scheduler, + low_cpu_mem_usage=True, + device_map="auto", + ).to("cuda") + if self.pipe.safety_checker is not None: + self.pipe.safety_checker = lambda images, **kwargs: (images, False) + self.pipe.enable_xformers_memory_efficient_attention() + + @method() + def run_inference( + self, + prompt: str, + n_prompt: str, + steps: int = 30, + batch_size: int = 1, + height: int = 512, + width: int = 512, + ) -> list[bytes]: + """ + Runs the Stable Diffusion pipeline on the given prompt and outputs images. + """ + import torch + + with torch.inference_mode(): + with torch.autocast("cuda"): + images = self.pipe( + [prompt] * batch_size, + negative_prompt=[n_prompt] * batch_size, + height=height, + width=width, + num_inference_steps=steps, + guidance_scale=7.5, + ).images + + # Convert to PNG bytes + image_output = [] + for image in images: + with io.BytesIO() as buf: + image.save(buf, format="PNG") + image_output.append(buf.getvalue()) + return image_output + + +@stub.local_entrypoint() +def entrypoint( + prompt: str, + n_prompt: str, + samples: int = 5, + steps: int = 30, + batch_size: int = 1, + height: int = 512, + width: int = 512, +): + """ + 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. + + The function is called with the following arguments: + - prompt: the prompt to run inference on + - n_prompt: the negative prompt to run inference on + - samples: the number of samples to generate + - steps: the number of steps to run inference for + - batch_size: the batch size to use + - height: the height of the output image + - width: the width of the output image + """ + print(f"steps => {steps}, sapmles => {samples}, batch_size => {batch_size}") + + directory = Path(f"./outputs/{date.today().strftime('%Y-%m-%d')}") + if not directory.exists(): + directory.mkdir(exist_ok=True, parents=True) + + stable_diffusion = StableDiffusion() + for i in range(samples): + start_time = time.time() + images = stable_diffusion.run_inference.call( + prompt, + n_prompt, + steps, + batch_size, + height, + width, + ) + total_time = time.time() - start_time + print( + f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image)." + ) + for j, image_bytes in enumerate(images): + output_path = directory / f"output_{j}_{i}.png" + print(f"Saving it to {output_path}") + with open(output_path, "wb") as file: + file.write(image_bytes)