Update app.py

app.py

@@ -331,36 +331,83 @@ if PEFT_AVAILABLE:
 else:
     print("   ⚠️ PEFT not available")
 
-# ── Model 3: CSM-DocExtract-VL-Q4KM (Qwen3VL, PRE-QUANTIZED Q4KM) ──
-print("\n3️⃣  CSM-DocExtract-VL-Q4KM (8B Q4KM — pre-quantized)...")
+# ── Model 3: CSM-DocExtract-VL-Q4KM (Full Qwen3VL, pre-quantized) ──
+print("\n3️⃣  CSM-DocExtract-VL-Q4KM (Full Qwen3VL, pre-quantized BNB)...")
 MODEL_ID_Q4KM = "Chhagan005/CSM-DocExtract-VL-Q4KM"
 CSM_Q4KM_AVAILABLE = False
 processor_q4km = model_q4km = None
 
 try:
-    processor_q4km = AutoProcessor.from_pretrained(MODEL_ID_Q4KM, trust_remote_code=True)
-    model_q4km     = load_vl_model(MODEL_ID_Q4KM, pre_quantized=True)
-    print("   ✅ Loaded! (pre-quantized Q4KM ~6-7GB)")
+    processor_q4km = AutoProcessor.from_pretrained(
+        MODEL_ID_Q4KM, trust_remote_code=True
+    )
+    # Pre-quantized safetensors → torch_dtype=auto, NO extra quantization_config
+    model_q4km = Qwen3VLForConditionalGeneration.from_pretrained(
+        MODEL_ID_Q4KM,
+        torch_dtype="auto",
+        device_map="auto",
+        trust_remote_code=True,
+    ).eval()
+    print("   ✅ Loaded! (Qwen3VL pre-quantized BNB ~6.4GB)")
     CSM_Q4KM_AVAILABLE = True
 except Exception as e:
-    print(f"   ❌ Failed: {e}")
-
-# ── Model 4: CSM-DocExtract-VL 4BNB (Qwen3VL, BitsAndBytes 4-bit) ──
-print("\n4️⃣  CSM-DocExtract-VL 4BNB (BitsAndBytes 4-bit)...")
+    try:
+        model_q4km = AutoModelForImageTextToText.from_pretrained(
+            MODEL_ID_Q4KM,
+            torch_dtype="auto",
+            device_map="auto",
+            trust_remote_code=True,
+        ).eval()
+        print("   ✅ Loaded! (AutoModel fallback)")
+        CSM_Q4KM_AVAILABLE = True
+    except Exception as e2:
+        print(f"   ❌ Failed: {e2}")
+
+# ── Model 4: CSM-DocExtract-VL (Full Qwen3VL, BNB INT4 trained) ──
+print("\n4️⃣  CSM-DocExtract-VL 4BNB (Full Qwen3VL, BNB INT4 trained)...")
 MODEL_ID_4BNB = "Chhagan005/CSM-DocExtract-VL"
 CSM_4BNB_AVAILABLE = False
 processor_4bnb = model_4bnb = None
+system_prompt_4bnb = "You are a helpful assistant."   # default
 
 try:
-    processor_4bnb = AutoProcessor.from_pretrained(MODEL_ID_4BNB, trust_remote_code=True)
-    model_4bnb     = load_vl_model(
-                         MODEL_ID_4BNB,
-                         quantization_config=bnb_4bit_config,
-                         pre_quantized=False)
-    print("   ✅ Loaded! (~6-7GB VRAM)")
+    # Read custom system_prompt.txt — this model was trained with it
+    try:
+        from huggingface_hub import hf_hub_download
+        sp_path = hf_hub_download(repo_id=MODEL_ID_4BNB, filename="system_prompt.txt")
+        with open(sp_path, "r", encoding="utf-8") as f:
+            system_prompt_4bnb = f.read().strip()
+        print(f"   📋 system_prompt.txt loaded: {system_prompt_4bnb[:80]}...")
+    except Exception as sp_err:
+        print(f"   ⚠️ system_prompt.txt not loaded: {sp_err} — using default")
+
+    processor_4bnb = AutoProcessor.from_pretrained(
+        MODEL_ID_4BNB, trust_remote_code=True
+    )
+    # BNB INT4 trained safetensors → torch_dtype=auto, NO extra quantization_config
+    # (ignore .gguf files — those are for llama.cpp, not transformers)
+    model_4bnb = Qwen3VLForConditionalGeneration.from_pretrained(
+        MODEL_ID_4BNB,
+        torch_dtype="auto",
+        device_map="auto",
+        trust_remote_code=True,
+        ignore_mismatched_sizes=True,   # GGUF files present — ignore safely
+    ).eval()
+    print("   ✅ Loaded! (Qwen3VL BNB INT4 trained ~6.4GB)")
     CSM_4BNB_AVAILABLE = True
 except Exception as e:
-    print(f"   ❌ Failed: {e}")
+    try:
+        model_4bnb = AutoModelForImageTextToText.from_pretrained(
+            MODEL_ID_4BNB,
+            torch_dtype="auto",
+            device_map="auto",
+            trust_remote_code=True,
+        ).eval()
+        print("   ✅ Loaded! (AutoModel fallback)")
+        CSM_4BNB_AVAILABLE = True
+    except Exception as e2:
+        print(f"   ❌ Failed: {e2}")
+
 
 print("\n" + "="*70)
 print("📊 MODEL STATUS")
@@ -714,7 +761,8 @@ def build_unified_summary(front_result: str, back_result: str, mrz_data: dict) -
 # ║         STEP PIPELINE FUNCTIONS         ║
 # ╚══════════════════════════════════════════╝
 
-def run_step1_extraction(model, processor, image, device, temperature, top_p, top_k, repetition_penalty):
+def run_step1_extraction(model, processor, image, device, temperature, top_p, top_k, repetition_penalty, system_prompt=None):
+
     """Step 1: LLM → Raw OCR, original script, NO translation, NO coordinates"""
 
     def _generate(prompt_text):
@@ -724,10 +772,15 @@ def run_step1_extraction(model, processor, image, device, temperature, top_p, to
         except ImportError:
             HAS_QWEN_VL_UTILS = False
 
-        messages = [{"role": "user", "content": [
-            {"type": "image", "image": image},
-            {"type": "text",  "text": prompt_text},
-        ]}]
+        sys_msg = system_prompt or "You are a helpful assistant."
+        messages = [
+            {"role": "system", "content": sys_msg},
+            {"role": "user", "content": [
+                {"type": "image", "image": image},
+                {"type": "text",  "text": prompt_text},
+            ]}
+        ]
+
 
         # Step A: Build prompt string
         try:
@@ -1081,8 +1134,13 @@ def generate_dual_card_ocr(model_name: str, text: str,
         full_output += "⏳ **Step 1/2 — Raw OCR (original script, no translation)...**\n\n"
         yield full_output, full_output
 
+        # Model 4 ke liye system prompt pass karo
+        sys_p = system_prompt_4bnb if model_name == "CSM-DocExtract-4BNB 💎" else None
+
         step1_raw = run_step1_extraction(model, processor, image_front, device,
-                                         temperature, top_p, top_k, repetition_penalty)
+                                          temperature, top_p, top_k, repetition_penalty,
+                                          system_prompt=sys_p)
+
         front_meta = parse_step1_output(step1_raw)
         front_meta_saved = front_meta