Update app.py

app.py

@@ -2,22 +2,38 @@ import os, json, re
 import gradio as gr
 from PIL import Image
 import torch
-from transformers import AutoProcessor, AutoModelForCausalLM
+from transformers import AutoProcessor, AutoTokenizer, AutoModelForCausalLM
 
 MODEL_ID = os.environ.get("MODEL_ID", "GrassData/cliptagger-12b")
+BASE_PROCESSOR_ID = os.environ.get("BASE_PROCESSOR_ID", "google/gemma-3-12b-it")
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-# Load processor & model
-processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN, trust_remote_code=True)
+# ---- Load processor (from base) + model (from your FT) ----
+try:
+    # Processor comes from base VLM repo (has preprocessor_config.json)
+    processor = AutoProcessor.from_pretrained(
+        BASE_PROCESSOR_ID, token=HF_TOKEN, trust_remote_code=True
+    )
+except Exception as e:
+    raise RuntimeError(f"Failed to load processor from {BASE_PROCESSOR_ID}: {e}")
+
+# Optional: get a fast tokenizer if processor doesn't expose one
+tokenizer = getattr(processor, "tokenizer", None)
+if tokenizer is None:
+    tokenizer = AutoTokenizer.from_pretrained(
+        BASE_PROCESSOR_ID, token=HF_TOKEN, trust_remote_code=True, use_fast=True
+    )
+
+# Your fine-tuned weights
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     token=HF_TOKEN,
     torch_dtype=DTYPE,
     device_map="auto",
-    trust_remote_code=True
+    trust_remote_code=True,
 )
 
 # Prompts (system + user, as given)