add use-cpu

app.py

@@ -1,8 +1,12 @@
 from transformers import AutoModel, AutoTokenizer
 import gradio as gr
 
+use_cpu = os.environ.get("USE_CPU", "all")
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
+if not use_cpu:
+    model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
+else:
+    model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).float()
 model = model.eval()
 
 def predict(input, history=None):