Update app.py

app.py

@@ -13,6 +13,18 @@ from share_btn import community_icon_html, loading_icon_html, share_js
 # word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt")
 # word_list = word_list_dataset["train"]['text']
 
+# 下载模型
+base_dir = "/root/.cache/huggingface/hub"
+if not os.path.isdir(base_dir):
+    os.makedirs(base_dir)
+
+cmd_list = ["cd", base_dir, "&&", "git lfs install", "&&", "git clone", "https://gitee.com/modelee/stable-diffusion-2.git", "models"]
+cmd_str = " ".join(cmd_list)
+print("cmd_str:", cmd_str)
+ret, out = subprocess.getstatusoutput(cmd_str)
+print("ret:", ret)
+print("out:", out)
+
 is_gpu_busy = True
 def infer(prompt, negative, scale):
     global is_gpu_busy
@@ -28,7 +40,7 @@ def infer(prompt, negative, scale):
     #     image_b64 = (f"data:image/jpeg;base64,{image}")
     #     images.append(image_b64)
 
-    model_id = "stabilityai/stable-diffusion-2"
+    model_id = base_dir +"models"
     # Use the Euler scheduler here instead
     scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
     pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float32)