Update app.py

app.py

@@ -1,10 +1,17 @@
 from fastapi import FastAPI
 import uvicorn
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 model_name = "TheBloke/Wizard-Vicuna-13B-Uncensored-HF"
 
-model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",            # Auto-distribute layers across CPU/GPU
+    low_cpu_mem_usage=True,       # Reduces CPU RAM during loading
+    torch_dtype="auto",           # Automatically select dtype (float16/32)
+    offload_folder="offload",
+    trust_remote_code=True
+)
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token