Update app.py

app.py

@@ -28,26 +28,33 @@ os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
 os.environ["HF_HOME"] = "/tmp/hf_home"
 
 from transformers import (
-    Qwen2_5_VLForConditionalGeneration,
-    Qwen2VLForConditionalGeneration,
     AutoProcessor,
+    AutoModelForImageTextToText,   # Universal VLM loader — Qwen2VL + Qwen3VL dono
     TextIteratorStreamer,
     BitsAndBytesConfig,
 )
 
-try:
-    from peft import PeftModel, PeftConfig
-    PEFT_AVAILABLE = True
-except:
-    PEFT_AVAILABLE = False
-    print("⚠️ PEFT not available")
-
+# Specific class imports — graceful fallback
 try:
     from transformers import Qwen3VLForConditionalGeneration
     QWEN3_AVAILABLE = True
-except:
+    print("✅ Qwen3VLForConditionalGeneration available")
+except ImportError:
     QWEN3_AVAILABLE = False
-    print("⚠️ Qwen3VL not available in current transformers version")
+    print("⚠️ Qwen3VL direct import not available — using AutoModel fallback")
+
+try:
+    from transformers import Qwen2VLForConditionalGeneration
+    QWEN2_AVAILABLE = True
+except ImportError:
+    QWEN2_AVAILABLE = False
+
+try:
+    from transformers import Qwen2_5_VLForConditionalGeneration
+    QWEN25_AVAILABLE = True
+except ImportError:
+    QWEN25_AVAILABLE = False
+
 
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
@@ -223,6 +230,42 @@ RULES:
 
 ---"""
 
+def load_vl_model(model_id: str, quantization_config=None, pre_quantized: bool = False):
+    """
+    Universal VLM loader — Qwen2VL / Qwen3VL / any VLM
+    pre_quantized=True  → model already has weights quantized, no extra config needed
+    pre_quantized=False → apply quantization_config during load
+    """
+    load_kwargs = {
+        "torch_dtype": "auto",
+        "device_map": "auto",
+        "trust_remote_code": True,
+    }
+    if quantization_config is not None and not pre_quantized:
+        load_kwargs["quantization_config"] = quantization_config
+
+    # Try 1: Qwen3VL (newest)
+    if QWEN3_AVAILABLE:
+        try:
+            return Qwen3VLForConditionalGeneration.from_pretrained(
+                model_id, **load_kwargs).eval()
+        except Exception as e:
+            print(f"      Qwen3VL failed: {e}, trying AutoModel...")
+
+    # Try 2: AutoModelForImageTextToText (universal fallback)
+    try:
+        return AutoModelForImageTextToText.from_pretrained(
+            model_id, **load_kwargs).eval()
+    except Exception as e:
+        print(f"      AutoModel failed: {e}, trying Qwen2VL...")
+
+    # Try 3: Qwen2VL last resort
+    if QWEN2_AVAILABLE:
+        return Qwen2VLForConditionalGeneration.from_pretrained(
+            model_id, **load_kwargs).eval()
+
+    raise RuntimeError(f"No compatible loader found for {model_id}")
+
 
 # ╔══════════════════════════════════════════╗
 # ║         MODEL LOADING                   ║
@@ -240,7 +283,7 @@ bnb_4bit_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=True,
 )
 
-# ── Model 1: Chhagan_ML-VL-OCR-v1 (LoRA, keep) ──
+# ── Model 1: Chhagan_ML-VL-OCR-v1 (LoRA on Qwen2VL base) ──
 print("\n1️⃣  Chhagan_ML-VL-OCR-v1 (LoRA Refined)...")
 MODEL_ID_C1 = "Chhagan005/Chhagan_ML-VL-OCR-v1"
 CHHAGAN_V1_AVAILABLE = False
@@ -248,15 +291,11 @@ processor_c1 = model_c1 = None
 
 if PEFT_AVAILABLE:
     try:
-        try:
-            config = PeftConfig.from_pretrained(MODEL_ID_C1)
-            base_id = config.base_model_name_or_path
-        except:
-            base_id = "Qwen/Qwen2.5-VL-2B-Instruct"
+        config       = PeftConfig.from_pretrained(MODEL_ID_C1)
+        base_id      = config.base_model_name_or_path
         processor_c1 = AutoProcessor.from_pretrained(base_id, trust_remote_code=True)
-        base_c1 = Qwen2VLForConditionalGeneration.from_pretrained(
-            base_id, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
-        model_c1 = PeftModel.from_pretrained(base_c1, MODEL_ID_C1).to(device).eval()
+        base_c1      = load_vl_model(base_id)
+        model_c1     = PeftModel.from_pretrained(base_c1, MODEL_ID_C1).to(device).eval()
         print("   ✅ Loaded!")
         CHHAGAN_V1_AVAILABLE = True
     except Exception as e:
@@ -264,92 +303,41 @@ if PEFT_AVAILABLE:
 else:
     print("   ⚠️ PEFT not available")
 
-# ── Model 2: Chhagan-DocVL-Qwen3 (Qwen3 fine-tuned, keep) ──
+# ── Model 2: Chhagan-DocVL-Qwen3 (LoRA on Qwen3VL base) ──
 print("\n2️⃣  Chhagan-DocVL-Qwen3 (Qwen3-VL Refined)...")
 MODEL_ID_C2 = "Chhagan005/Chhagan-DocVL-Qwen3"
 CHHAGAN_QWEN3_AVAILABLE = False
 processor_c2 = model_c2 = None
 
-if QWEN3_AVAILABLE:
+if PEFT_AVAILABLE:
     try:
-        try:
-            if PEFT_AVAILABLE:
-                config = PeftConfig.from_pretrained(MODEL_ID_C2)
-                base_id = config.base_model_name_or_path
-                processor_c2 = AutoProcessor.from_pretrained(base_id, trust_remote_code=True)
-                base_c2 = Qwen3VLForConditionalGeneration.from_pretrained(
-                    base_id, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
-                model_c2 = PeftModel.from_pretrained(base_c2, MODEL_ID_C2).to(device).eval()
-            else:
-                raise Exception("No PEFT")
-        except:
-            print("   Loading as full fine-tuned...")
-            processor_c2 = AutoProcessor.from_pretrained(MODEL_ID_C2, trust_remote_code=True)
-            model_c2 = Qwen3VLForConditionalGeneration.from_pretrained(
-                MODEL_ID_C2, attn_implementation="flash_attention_2",
-                torch_dtype=torch.float16, device_map="auto", trust_remote_code=True
-            ).to(device).eval()
+        config       = PeftConfig.from_pretrained(MODEL_ID_C2)
+        base_id      = config.base_model_name_or_path
+        processor_c2 = AutoProcessor.from_pretrained(base_id, trust_remote_code=True)
+        base_c2      = load_vl_model(base_id)
+        model_c2     = PeftModel.from_pretrained(base_c2, MODEL_ID_C2).to(device).eval()
         print("   ✅ Loaded!")
         CHHAGAN_QWEN3_AVAILABLE = True
     except Exception as e:
         print(f"   ❌ Failed: {e}")
 else:
-    print("   ⚠️ Qwen3VL not in transformers version")
+    print("   ⚠️ PEFT not available")
 
-# ── Model 3: CSM-DocExtract-VL-Q4KM (NEW — Qwen3VL, pre-quantized Q4KM) ──
-print("\n3️⃣  CSM-DocExtract-VL-Q4KM (8B Q4KM — Qwen3VL architecture)...")
+# ── Model 3: CSM-DocExtract-VL-Q4KM (Qwen3VL, PRE-QUANTIZED Q4KM) ──
+print("\n3️⃣  CSM-DocExtract-VL-Q4KM (8B Q4KM — pre-quantized)...")
 MODEL_ID_Q4KM = "Chhagan005/CSM-DocExtract-VL-Q4KM"
 CSM_Q4KM_AVAILABLE = False
 processor_q4km = model_q4km = None
 
 try:
-    # Processor load
-    processor_q4km = AutoProcessor.from_pretrained(
-        MODEL_ID_Q4KM, trust_remote_code=True
-    )
-
-    # Model is qwen3_vl type + ALREADY pre-quantized Q4KM
-    # So: use Qwen3VL class + NO extra quantization_config
-    if QWEN3_AVAILABLE:
-        model_q4km = Qwen3VLForConditionalGeneration.from_pretrained(
-            MODEL_ID_Q4KM,
-            torch_dtype="auto",       # model already has Q4KM weights
-            device_map="auto",
-            trust_remote_code=True,
-        ).eval()
-        print("   ✅ Loaded! (Qwen3VL Q4KM pre-quantized)")
-        CSM_Q4KM_AVAILABLE = True
-    else:
-        # Qwen3VL not in transformers → use AutoModel fallback
-        from transformers import AutoModelForCausalLM
-        try:
-            from transformers import AutoModelForVisualQuestionAnswering
-            model_q4km = AutoModelForVisualQuestionAnswering.from_pretrained(
-                MODEL_ID_Q4KM,
-                torch_dtype="auto",
-                device_map="auto",
-                trust_remote_code=True,
-            ).eval()
-        except:
-            # Last fallback: force load with Qwen2_5 but ignore arch warning
-            import warnings
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore")
-                model_q4km = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                    MODEL_ID_Q4KM,
-                    torch_dtype="auto",
-                    device_map="auto",
-                    trust_remote_code=True,
-                    ignore_mismatched_sizes=True,
-                ).eval()
-        print("   ✅ Loaded! (fallback loader)")
-        CSM_Q4KM_AVAILABLE = True
-
+    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)")
+    CSM_Q4KM_AVAILABLE = True
 except Exception as e:
     print(f"   ❌ Failed: {e}")
 
-
-# ── Model 4: CSM-DocExtract-VL 4BNB (NEW, replaces Nanonets) ──
+# ── Model 4: CSM-DocExtract-VL 4BNB (Qwen3VL, BitsAndBytes 4-bit) ──
 print("\n4️⃣  CSM-DocExtract-VL 4BNB (BitsAndBytes 4-bit)...")
 MODEL_ID_4BNB = "Chhagan005/CSM-DocExtract-VL"
 CSM_4BNB_AVAILABLE = False
@@ -357,25 +345,10 @@ processor_4bnb = model_4bnb = None
 
 try:
     processor_4bnb = AutoProcessor.from_pretrained(MODEL_ID_4BNB, trust_remote_code=True)
-    try:
-        model_4bnb = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            MODEL_ID_4BNB,
-            quantization_config=bnb_4bit_config,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            trust_remote_code=True,
-        ).eval()
-    except:
-        if QWEN3_AVAILABLE:
-            model_4bnb = Qwen3VLForConditionalGeneration.from_pretrained(
-                MODEL_ID_4BNB,
-                quantization_config=bnb_4bit_config,
-                torch_dtype=torch.float16,
-                device_map="auto",
-                trust_remote_code=True,
-            ).eval()
-        else:
-            raise Exception("Architecture detection failed")
+    model_4bnb     = load_vl_model(
+                         MODEL_ID_4BNB,
+                         quantization_config=bnb_4bit_config,
+                         pre_quantized=False)
     print("   ✅ Loaded! (~6-7GB VRAM)")
     CSM_4BNB_AVAILABLE = True
 except Exception as e:
@@ -385,10 +358,10 @@ print("\n" + "="*70)
 print("📊 MODEL STATUS")
 print("="*70)
 status = [
-    ("Chhagan_ML-VL-OCR-v1",     CHHAGAN_V1_AVAILABLE,    "LoRA Fine-tuned"),
-    ("Chhagan-DocVL-Qwen3",      CHHAGAN_QWEN3_AVAILABLE, "Qwen3-VL Fine-tuned"),
-    ("CSM-DocExtract-VL-Q4KM",   CSM_Q4KM_AVAILABLE,      "8B Q4KM ~6-7GB"),
-    ("CSM-DocExtract-VL 4BNB",   CSM_4BNB_AVAILABLE,      "BitsAndBytes 4-bit ~6-7GB"),
+    ("Chhagan_ML-VL-OCR-v1",   CHHAGAN_V1_AVAILABLE,    "LoRA Fine-tuned"),
+    ("Chhagan-DocVL-Qwen3",    CHHAGAN_QWEN3_AVAILABLE, "Qwen3-VL Fine-tuned"),
+    ("CSM-DocExtract-Q4KM",    CSM_Q4KM_AVAILABLE,      "Qwen3VL Q4KM pre-quantized"),
+    ("CSM-DocExtract-4BNB",    CSM_4BNB_AVAILABLE,      "Qwen3VL BitsAndBytes 4-bit"),
 ]
 for name, ok, note in status:
     print(f"  {'✅' if ok else '❌'} {name:<35} {note}")
@@ -397,6 +370,7 @@ loaded = sum(x[1] for x in status)
 print(f"  Total loaded: {loaded}/4\n")
 
 
+
 # ╔══════════════════════════════════════════╗
 # ║         PYTHON PIPELINE FUNCTIONS       ║
 # ╚══════════════════════════════════════════╝