Update app.py

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 # from datasets import load_dataset
 from PIL import Image
 import torch
@@ -18,7 +20,7 @@ base_dir = "./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"]
+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)
@@ -42,8 +44,8 @@ def infer(prompt, negative, scale):
 
     model_id = base_dir 
     # Use the Euler scheduler here instead
-    scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
-    pipe = DiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, use_safetensors=True)
+    scheduler = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_id)
+    pipe = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_id)
     # pipe = pipe.to("cuda")
     images = pipe(prompt).images
     return images