From a5adb6b9fb242d4f39522b5d232598525df19811 Mon Sep 17 00:00:00 2001
From: hodanov <1031hoda@gmail.com>
Date: Sat, 17 Jun 2023 17:41:05 +0900
Subject: [PATCH] Modify to use LoRA.

---
 .env.example | 14 ++++++++++++++
 sd_cli.py    | 42 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/.env.example b/.env.example
index 90a3abb..e19e5d7 100644
--- a/.env.example
+++ b/.env.example
@@ -1,4 +1,18 @@
 HUGGING_FACE_TOKEN=""
 MODEL_REPO_ID="stabilityai/stable-diffusion-2-1"
 MODEL_NAME="stable-diffusion-2-1"
+
+# Modify `USE_VAE` to `true` if you want to use VAE.
 USE_VAE="false"
+
+# Add LoRA if you want to use one. You can use a download link of civitai.
+# ex)
+#  - `LORA_NAMES="hogehoge.safetensors"`
+#  - `LORA_DOWNLOAD_URLS="https://civitai.com/api/download/models/xxxxxx"`
+#
+# If you have multiple LoRAs you want to use, separate by commas like the below:
+# ex)
+#  - `LORA_NAMES="hogehoge.safetensors,mogumogu.safetensors"`
+#  - `LORA_DOWNLOAD_URLS="https://civitai.com/api/download/models/xxxxxx,https://civitai.com/api/download/models/xxxxxx"`
+LORA_NAMES=""
+LORA_DOWNLOAD_URLS=""
diff --git a/sd_cli.py b/sd_cli.py
index 01b7a01..e1d0dd6 100644
--- a/sd_cli.py
+++ b/sd_cli.py
@@ -3,12 +3,31 @@ from __future__ import annotations
 import io
 import os
 import time
+from urllib.request import Request, urlopen
 
 from modal import Image, Mount, Secret, Stub, method
 
 import util
 
 BASE_CACHE_PATH = "/vol/cache"
+BASE_CACHE_PATH_LORA = "/vol/cache/lora"
+
+
+def download_loras():
+    """
+    Download LoRA.
+    """
+    lora_names = os.getenv("LORA_NAMES").split(",")
+    lora_download_urls = os.getenv("LORA_DOWNLOAD_URLS").split(",")
+
+    for name, url in zip(lora_names, lora_download_urls):
+        req = Request(url, headers={"User-Agent": "Mozilla/5.0"})
+        downloaded = urlopen(req).read()
+
+        dir_names = os.path.join(BASE_CACHE_PATH_LORA, name)
+        os.makedirs(os.path.dirname(dir_names), exist_ok=True)
+        with open(dir_names, mode="wb") as f:
+            f.write(downloaded)
 
 
 def download_models():
@@ -45,11 +64,21 @@ def download_models():
     pipe.save_pretrained(cache_path, safe_serialization=True)
 
 
+def build_image():
+    """
+    Build the Docker image.
+    """
+    download_models()
+
+    if os.environ["LORA_NAMES"] != "":
+        download_loras()
+
+
 stub_image = Image.from_dockerfile(
     path="./Dockerfile",
     context_mount=Mount.from_local_file("./requirements.txt"),
 ).run_function(
-    download_models,
+    build_image,
     secrets=[Secret.from_dotenv(__file__)],
 )
 stub = Stub("stable-diffusion-cli")
@@ -100,6 +129,17 @@ class StableDiffusion:
                 torch_dtype=torch.float16,
             ).to("cuda")
 
+        if os.environ["LORA_NAMES"] != "":
+            lora_names = os.getenv("LORA_NAMES").split(",")
+            for lora_name in lora_names:
+                path_to_lora = os.path.join(BASE_CACHE_PATH_LORA, lora_name)
+                if os.path.exists(path_to_lora):
+                    print(f"The directory '{path_to_lora}' exists.")
+                else:
+                    print(f"The directory '{path_to_lora}' does not exist. Download loras...")
+                    download_loras()
+                self.pipe.load_lora_weights(".", weight_name=path_to_lora)
+
         self.pipe.enable_xformers_memory_efficient_attention()
 
     @method()