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

Refactor to deploy modal app.

Browse Source
This commit is contained in:
hodanov 2023-06-27 22:25:32 +09:00
parent 643e0e2ea6
commit 364cb78992
4 changed files with 89 additions and 118 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
Makefile
View File

@ -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

3
README.md
View File

@ -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.

56
entrypoint.py Normal file
View File

@ -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)

135
sd_cli.py → sdcli.py
View File

@ -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)