Update app.py

app.py

@@ -19,6 +19,7 @@ import torch, gradio as gr
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForVision2Seq
 from peft import PeftModel
+from transformers import BitsAndBytesConfig
 
 # ---------------------------------------------------------------------
 # 1  Load base + LoRA  (≈ 12 GB VRAM in bf16; falls back to CPU if needed)
@@ -31,8 +32,16 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype  = torch.float16 if device == "cuda" else torch.float32
 
 processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
+
+quant_cfg = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
+
 base      = AutoModelForVision2Seq.from_pretrained(
-                BASE_MODEL, torch_dtype=dtype, trust_remote_code=True)
+                BASE_MODEL, device_map="auto", torch_dtype=torch.float16, quantization_config=quant_cfg, trust_remote_code=True)
 model     = PeftModel.from_pretrained(base, ADAPTER_REPO, revision=ADAPTER_REV)
 model.to(device).eval()