Update app.py

app.py

@@ -1,16 +1,21 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from spaces import GPU  # Import ZeroGPU decorator
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from spaces import GPU
 
-# Load model & tokenizer (runs on CPU at startup)
+# Load model & tokenizer
 MODEL_NAME = "ubiodee/plutus_llm"
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    bnb_8bit_compute_dtype=torch.float16,
+    bnb_8bit_use_double_quant=True
+)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float16,
+    quantization_config=quantization_config,
     device_map="auto",
-    load_in_8bit=True
+    torch_dtype=torch.float16
 )
 
 # Set padding token