Update app.py

app.py

@@ -7,11 +7,14 @@ import os
 import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import TorchAoConfig # not for Zero GPU
+from torchao.quantization import Int8DynamicActivationInt8WeightConfig, Float8DynamicActivationFloat8WeightConfig # not for Zero GPU
 import spaces # 1. Import the spaces library
 
 IS_CUDA = torch.cuda.is_available()
 IS_ZEROGPU = True if os.getenv("SPACES_ZERO_GPU", None) else False
 if IS_ZEROGPU: torch.set_float32_matmul_precision("high")
+IS_QUANT = True
 
 # ── Configuration ────────────────────────────────────────────────────────────
 MODEL_ID = "Reubencf/gemma3-konkani"
@@ -28,8 +31,14 @@ print(f"[Init] Loading model pipeline for the first time: {MODEL_ID}...")
 def load_model():
     try:
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16 if IS_CUDA else torch.float32,
-                                                     device_map="auto", token=HF_TOKEN)
+        if IS_QUANT: # not for Zero GPU
+            quant_config = Float8DynamicActivationFloat8WeightConfig() if IS_CUDA else Int8DynamicActivationInt8WeightConfig()
+            quantization_config = TorchAoConfig(quant_type=quant_config)
+            model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16 if IS_CUDA else torch.float32,
+                                                         device_map="auto", quantization_config=quantization_config, token=HF_TOKEN)
+        else:
+            model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16 if IS_CUDA else torch.float32,
+                                                         device_map="auto", token=HF_TOKEN)
         print("[Init] Model loaded successfully.")
         return model, tokenizer
     except Exception as e:
@@ -79,7 +88,7 @@ def generate_response(message, history=[], system_message="", max_tokens=DEF_TOK
         # Generate the response
         gen_kwargs = dict(
             input_ids=inputs["input_ids"],
-            attention_mask=inputs["attention_mask"],
+            #attention_mask=inputs["attention_mask"],
             max_new_tokens=max_tokens,
             do_sample=True,
             temperature=temperature,
@@ -129,4 +138,4 @@ demo = gr.ChatInterface(
 # ── Launch ────────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
     print("🚀 Starting Gradio app for ZeroGPU...")
-    demo.queue().launch()
+    demo.queue().launch()