From daf84a5280661266ce28cc2276fa049b5fdb8717 Mon Sep 17 00:00:00 2001
From: hodanov <1031hoda@gmail.com>
Date: Sun, 21 May 2023 18:43:26 +0900
Subject: [PATCH] Modify sd_cli.py to use dotenv file. Add .env.example.

---
 .env.example |  3 +++
 .gitignore   |  1 +
 sd_cli.py    | 36 ++++++++++++++++++------------------
 3 files changed, 22 insertions(+), 18 deletions(-)
 create mode 100644 .env.example

diff --git a/.env.example b/.env.example
new file mode 100644
index 0000000..3f85f8c
--- /dev/null
+++ b/.env.example
@@ -0,0 +1,3 @@
+HUGGINGFACE_TOKEN=""
+MODEL_REPO_ID="stabilityai/stable-diffusion-2-1"
+MODEL_NAME="stable-diffusion-2-1"
diff --git a/.gitignore b/.gitignore
index 85e9ced..4ab8942 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
 .DS_Store
 __pycache__/
 outputs/
+.env
diff --git a/sd_cli.py b/sd_cli.py
index 0fd97c7..c4fd228 100644
--- a/sd_cli.py
+++ b/sd_cli.py
@@ -8,11 +8,7 @@ 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"])
+BASE_CACHE_PATH = "/vol/cache"
 
 
 def download_models():
@@ -24,22 +20,24 @@ def download_models():
     import torch
 
     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(
-        MODEL["repo_id"],
+        model_repo_id,
         subfolder="scheduler",
         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(
-        MODEL["repo_id"],
+        model_repo_id,
         use_auth_token=hugging_face_token,
         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 = (
@@ -58,13 +56,14 @@ stub_image = (
     .pip_install("xformers", pre=True)
     .run_function(
         download_models,
-        secrets=[Secret.from_name("my-huggingface-secret")],
+        secrets=[Secret.from_dotenv(__file__)],
     )
 )
 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:
     """
     A class that wraps the Stable Diffusion pipeline and scheduler.
@@ -74,16 +73,17 @@ class StableDiffusion:
         import diffusers
         import torch
 
-        if os.path.exists(CACHE_PATH):
-            print(f"The directory '{CACHE_PATH}' exists.")
+        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...")
+            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,
+            cache_path,
             subfolder="scheduler",
             solver_order=2,
             prediction_type="epsilon",
@@ -96,7 +96,7 @@ class StableDiffusion:
         )
 
         self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
-            CACHE_PATH,
+            cache_path,
             scheduler=scheduler,
             low_cpu_mem_usage=True,
             device_map="auto",