From 364cb789928180ada2f2342880d1bb2160f301ae Mon Sep 17 00:00:00 2001 From: hodanov <1031hoda@gmail.com> Date: Tue, 27 Jun 2023 22:25:32 +0900 Subject: [PATCH] Refactor to deploy modal app. --- Makefile | 13 ++-- README.md | 3 +- entrypoint.py | 56 ++++++++++++++++++ sd_cli.py => sdcli.py | 135 +++++++----------------------------------- 4 files changed, 89 insertions(+), 118 deletions(-) create mode 100644 entrypoint.py rename sd_cli.py => sdcli.py (75%) diff --git a/Makefile b/Makefile index 5268315..dee827e 100644 --- a/Makefile +++ b/Makefile @@ -1,8 +1,11 @@ +deploy: + modal deploy sdcli.py + run: - modal run sd_cli.py \ - --prompt "A woman with bob hair" \ + modal run entrypoint.py \ + --prompt "a photograph of an astronaut riding a horse" \ --n-prompt "" \ - --height 768 \ + --height 512 \ --width 512 \ - --samples 5 \ - --steps 30 + --samples 1 \ + --steps 50 diff --git a/README.md b/README.md index c5badc1..c0706c8 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,8 @@ To use the script, execute the below. 1. git clone the repository. 2. Create the `.env` file and set a huggingface API token and a model with reference to `.env.example`. 3. Open the Makefile and set prompts. -4. Execute `make run` command. +4. Execute `make deploy` command. An application will be deployed to Modal by the command. +5. Execute `make run` command. Images are generated and output to the `outputs/` directory. diff --git a/entrypoint.py b/entrypoint.py new file mode 100644 index 0000000..4a17657 --- /dev/null +++ b/entrypoint.py @@ -0,0 +1,56 @@ +import time + +import modal + +stub = modal.Stub("run-stable-diffusion-cli") +stub.run_inference = modal.Function.from_name("stable-diffusion-cli", "StableDiffusion.run_inference") + + +@stub.local_entrypoint() +def main( + 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 + + 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 = sd.run_inference(seed=seed_generated) + images = stub.app.run_inference.call( + prompt=prompt, + n_prompt=n_prompt, + height=height, + width=width, + batch_size=batch_size, + steps=steps, + seed=seed_generated, + ) + util.save_images(directory, images, seed_generated, i) + 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, + "n_prompt": n_prompt, + "height": height, + "width": width, + "samples": samples, + "batch_size": batch_size, + "steps": steps, + } + util.save_prompts(prompts) diff --git a/sd_cli.py b/sdcli.py similarity index 75% rename from sd_cli.py rename to sdcli.py index ef328da..406d2c3 100644 --- a/sd_cli.py +++ b/sdcli.py @@ -2,7 +2,6 @@ from __future__ import annotations import io import os -import time from urllib.request import Request, urlopen from modal import Image, Mount, Secret, Stub, method @@ -94,38 +93,20 @@ class StableDiffusion(ClsMixin): A class that wraps the Stable Diffusion pipeline and scheduler. """ - def __init__( - self, - prompt: str, - n_prompt: str, - height: int = 512, - width: int = 512, - samples: int = 1, - batch_size: int = 1, - steps: int = 30, - ): + def __enter__(self): import diffusers import torch - self.prompt = prompt - self.n_prompt = n_prompt - self.height = height - self.width = width - self.samples = samples - self.batch_size = batch_size - self.steps = steps self.use_vae = os.environ["USE_VAE"] == "true" self.upscaler = os.environ["UPSCALER"] self.use_face_enhancer = os.environ["USE_FACE_ENHANCER"] == "true" self.use_hires_fix = os.environ["USE_HIRES_FIX"] == "true" - self.cache_path = os.path.join(BASE_CACHE_PATH, os.environ["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. Download models...") download_models() - self.max_embeddings_multiples = self.count_token(p=prompt, n=n_prompt) torch.backends.cuda.matmul.allow_tf32 = True @@ -203,23 +184,34 @@ class StableDiffusion(ClsMixin): return max_embeddings_multiples @method() - def run_inference(self, seed: int) -> list[bytes]: + 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, + ) -> 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( - self.prompt * self.batch_size, - negative_prompt=self.n_prompt * self.batch_size, - height=self.height, - width=self.width, - num_inference_steps=self.steps, + prompt * batch_size, + negative_prompt=n_prompt * batch_size, + height=height, + width=width, + num_inference_steps=steps, guidance_scale=7.5, - max_embeddings_multiples=self.max_embeddings_multiples, + max_embeddings_multiples=max_embeddings_multiples, generator=generator, ).images @@ -236,12 +228,12 @@ class StableDiffusion(ClsMixin): with torch.inference_mode(): with torch.autocast("cuda"): hires_fixed = self.pipe.img2img( - prompt=self.prompt * self.batch_size, - negative_prompt=self.n_prompt * self.batch_size, - num_inference_steps=self.steps, + prompt=prompt * batch_size, + negative_prompt=n_prompt * batch_size, + num_inference_steps=steps, strength=0.3, guidance_scale=7.5, - max_embeddings_multiples=self.max_embeddings_multiples, + max_embeddings_multiples=max_embeddings_multiples, generator=generator, image=img, ).images @@ -336,84 +328,3 @@ class StableDiffusion(ClsMixin): torch.cuda.empty_cache() return upscaled_imgs - - # TODO: Implement this - # @method() - # def img2img( - # self, - # prompt: str, - # n_prompt: str, - # batch_size: int = 1, - # steps: int = 20, - # strength: float = 0.3, - # max_embeddings_multiples: int = 1, - # # image: Image.Image = None, - # base_images: list[Image.Image], - # ): - # import torch - - # torch.cuda.empty_cache() - # for img in base_images: - # 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=strength, - # guidance_scale=7.5, - # max_embeddings_multiples=max_embeddings_multiples, - # generator=generator, - # image=img, - # ).images - # base_images.extend(hires_fixed) - # torch.cuda.empty_cache() - - -@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 - - directory = util.make_directory() - - sd = StableDiffusion.remote( - prompt=prompt, - n_prompt=n_prompt, - height=height, - width=width, - batch_size=batch_size, - steps=steps, - ) - for i in range(samples): - if seed == -1: - seed_generated = util.generate_seed() - start_time = time.time() - images = sd.run_inference(seed=seed_generated) - util.save_images(directory, images, seed_generated, i) - 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, - "n_prompt": n_prompt, - "height": height, - "width": width, - "samples": samples, - "batch_size": batch_size, - "steps": steps, - } - util.save_prompts(prompts)