Update app.py

app.py

@@ -4,7 +4,12 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 import os
 from threading import Thread
+from accelerate import init_empty_weights
 
+max_memory = {
+    0: "40GiB",  # بسته به VRAM کارت گرافیکت، مثلاً اگه 8 گیگ داری، بذار 6 یا 5
+    "cpu": "32GiB",  # بسته به RAM سیستمت
+}
 MODEL_LIST = ["THUDM/GLM-4-Z1-32B-0414"]
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -46,7 +51,8 @@ def stream_chat(message, history: list, temperature: float, max_length: int, top
         low_cpu_mem_usage=True,
         trust_remote_code=True,
         quantization_config=quantization_config,
-        device_map="auto"
+        device_map="auto"و
+        max_memory=max_memory,
     )
         
     print(f'message is - {message}')