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

Modify sd_cli.py to use dotenv file. Add .env.example.

Browse Source
This commit is contained in:
hodanov 2023-05-21 18:43:26 +09:00
parent 1a599cec7b
commit daf84a5280
3 changed files with 22 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.env.example Normal file
View File

@ -0,0 +1,3 @@
HUGGINGFACE_TOKEN=""
MODEL_REPO_ID="stabilityai/stable-diffusion-2-1"
MODEL_NAME="stable-diffusion-2-1"

1
.gitignore vendored
View File

@ -1,3 +1,4 @@
.DS_Store .DS_Store
__pycache__/ __pycache__/
outputs/ outputs/
.env

36
sd_cli.py
View File

@ -8,11 +8,7 @@ from modal import Image, Secret, Stub, method
stub = Stub("stable-diffusion-cli") stub = Stub("stable-diffusion-cli")
MODEL = { BASE_CACHE_PATH = "/vol/cache"
"repo_id": "runwayml/stable-diffusion-v1-5",
"name": "stable-diffusion-v1-5",
}
CACHE_PATH = os.path.join("/vol/cache", MODEL["name"])
def download_models(): def download_models():
@ -24,22 +20,24 @@ def download_models():
import torch import torch
hugging_face_token = os.environ["HUGGINGFACE_TOKEN"] hugging_face_token = os.environ["HUGGINGFACE_TOKEN"]
model_repo_id = os.environ["MODEL_REPO_ID"]
cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])
scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained( scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
MODEL["repo_id"], model_repo_id,
subfolder="scheduler", subfolder="scheduler",
use_auth_token=hugging_face_token, use_auth_token=hugging_face_token,
cache_dir=CACHE_PATH, cache_dir=cache_path,
) )
scheduler.save_pretrained(CACHE_PATH, safe_serialization=True) scheduler.save_pretrained(cache_path, safe_serialization=True)
pipe = diffusers.StableDiffusionPipeline.from_pretrained( pipe = diffusers.StableDiffusionPipeline.from_pretrained(
MODEL["repo_id"], model_repo_id,
use_auth_token=hugging_face_token, use_auth_token=hugging_face_token,
torch_dtype=torch.float16, torch_dtype=torch.float16,
cache_dir=CACHE_PATH, cache_dir=cache_path,
) )
pipe.save_pretrained(CACHE_PATH, safe_serialization=True) pipe.save_pretrained(cache_path, safe_serialization=True)
stub_image = ( stub_image = (
@ -58,13 +56,14 @@ stub_image = (
.pip_install("xformers", pre=True) .pip_install("xformers", pre=True)
.run_function( .run_function(
download_models, download_models,
secrets=[Secret.from_name("my-huggingface-secret")], secrets=[Secret.from_dotenv(__file__)],
) )
) )
stub.image = stub_image stub.image = stub_image
@stub.cls(gpu="A10G", secrets=[Secret.from_name("my-huggingface-secret")]) # @stub.cls(gpu="A10G", secrets=[Secret.from_name("my-huggingface-secret")])
@stub.cls(gpu="A10G", secrets=[Secret.from_dotenv(__file__)])
class StableDiffusion: class StableDiffusion:
""" """
A class that wraps the Stable Diffusion pipeline and scheduler. A class that wraps the Stable Diffusion pipeline and scheduler.
@ -74,16 +73,17 @@ class StableDiffusion:
import diffusers import diffusers
import torch import torch
if os.path.exists(CACHE_PATH): cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])
print(f"The directory '{CACHE_PATH}' exists.") if os.path.exists(cache_path):
print(f"The directory '{cache_path}' exists.")
else: else:
print(f"The directory '{CACHE_PATH}' does not exist. Download models...") print(f"The directory '{cache_path}' does not exist. Download models...")
download_models() download_models()
torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = True
scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained( scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
CACHE_PATH, cache_path,
subfolder="scheduler", subfolder="scheduler",
solver_order=2, solver_order=2,
prediction_type="epsilon", prediction_type="epsilon",
@ -96,7 +96,7 @@ class StableDiffusion:
) )
self.pipe = diffusers.StableDiffusionPipeline.from_pretrained( self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
CACHE_PATH, cache_path,
scheduler=scheduler, scheduler=scheduler,
low_cpu_mem_usage=True, low_cpu_mem_usage=True,
device_map="auto", device_map="auto",