Update app.py

app.py

@@ -1,10 +1,17 @@
+import importlib.util
 import os
+import sys
 import cv2
 import gradio as gr
 import numpy as np
 import torch
+from huggingface_hub import hf_hub_download
 from PIL import Image
-from transformers import AutoConfig, AutoModelForImageClassification, ViTImageProcessor
+from transformers import (
+    AutoConfig,
+    AutoModelForImageClassification,
+    ViTImageProcessor,
+)
 
 MODEL_REPO = "SARVM/Refined_ViT"
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -28,20 +35,50 @@ processor = ViTImageProcessor(
     image_std=[0.229, 0.224, 0.225],
 )
 
-print("Loading custom model via trust_remote_code...")
+print("Loading custom model code...")
 
-# 2. Load Config and Model automatically using remote code
-config = AutoConfig.from_pretrained(
-    MODEL_REPO,
+# 2. Download custom config and modeling files directly from HF Hub
+config_path = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename="configuration_vit_refiner.py",
     token=HF_TOKEN,
-    trust_remote_code=True
 )
+model_path = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename="modeling_vit_refiner.py",
+    token=HF_TOKEN,
+)
+
+# 3. Dynamically import configuration module
+config_spec = importlib.util.spec_from_file_location(
+    "configuration_vit_refiner", config_path
+)
+config_module = importlib.util.module_from_spec(config_spec)
+sys.modules["configuration_vit_refiner"] = config_module
+config_spec.loader.exec_module(config_module)
+
+# 4. Dynamically import modeling module
+model_spec = importlib.util.spec_from_file_location(
+    "modeling_vit_refiner", model_path
+)
+model_module = importlib.util.module_from_spec(model_spec)
+sys.modules["modeling_vit_refiner"] = model_module
+model_spec.loader.exec_module(model_module)
+
+# 5. Get the actual custom classes from imported modules
+ViTRefinerConfig = config_module.ViTRefinerConfig
+ViTRefinerModel = model_module.ViTRefinerModel
+
+# 6. Explicitly register 'vit_refiner' with Transformers AutoClasses
+AutoConfig.register("vit_refiner", ViTRefinerConfig)
+AutoModelForImageClassification.register(ViTRefinerConfig, ViTRefinerModel)
 
+# 7. Now safely load via standard AutoClasses
+config = AutoConfig.from_pretrained(MODEL_REPO, token=HF_TOKEN)
 model = AutoModelForImageClassification.from_pretrained(
     MODEL_REPO,
     config=config,
     token=HF_TOKEN,
-    trust_remote_code=True
 )
 
 model.to(DEVICE)
@@ -146,7 +183,7 @@ def predict(image):
     with torch.no_grad():
         outputs = model(**inputs)
 
-    # Extract logits defensively (handles both dict and Object outputs)
+    # Extract logits and attentions safely
     if isinstance(outputs, dict):
         logits = outputs.get("logits")
         attentions = outputs.get("attentions", None)
@@ -165,7 +202,7 @@ def predict(image):
     prediction = model.config.id2label.get(predicted_index, str(predicted_index))
     confidence_text = f"{confidence * 100:.2f}%"
 
-    # Generate attention visualization (safely handles None attentions)
+    # Generate attention visualization
     overlay = create_attention_overlay(image, attentions)
 
     return prediction, confidence_text, overlay