Update app.py

app.py

@@ -1,83 +1,83 @@
 import gradio as gr
 import torch
 import torch.nn as nn
-from torchvision import transforms
+import torchvision.transforms as T
 from PIL import Image
-import requests
 from huggingface_hub import hf_hub_download
 
-########################
-# 1) Download & Load Model
-########################
-
-# Replace with your actual model repo on HF
-repo_id = "Abdu07/multitask-model"
-filename = "multitask_model.pth"
-
-# Download the model file from the Hub
+#####################################
+# 1) Define the same custom class
+#####################################
+class MultiTaskModel(nn.Module):
+    def __init__(self, backbone, feature_dim, num_obj_classes):
+        super(MultiTaskModel, self).__init__()
+        self.backbone = backbone
+        self.obj_head = nn.Linear(feature_dim, num_obj_classes)
+        self.bin_head = nn.Linear(feature_dim, 2)
+
+    def forward(self, x):
+        feats = self.backbone(x)
+        obj_logits = self.obj_head(feats)
+        bin_logits = self.bin_head(feats)
+        return obj_logits, bin_logits
+
+#####################################
+# 2) Allowlist the class
+#####################################
+import torch.serialization
+torch.serialization.add_safe_globals([MultiTaskModel])
+
+#####################################
+# 3) Download & Load the full model
+#####################################
+repo_id = "Abdu07/multitask-model"  # or your actual repo
+filename = "multitask_model.pth"    # the file you uploaded
 model_path = hf_hub_download(repo_id=repo_id, filename=filename)
-model = torch.load(model_path, map_location="cpu")  # or map_location="cuda" if you prefer
-model.eval()
 
-########################
-# 2) Define Label Mappings
-########################
+# Force PyTorch to load the full model object
+model = torch.load(model_path, map_location="cpu")  # default weights_only=True, but we added safe_globals
+model.eval()
 
-# For example, if your object labels are saved in code:
+#####################################
+# 4) Label Mappings
+#####################################
 idx_to_obj_label = {
+    # Fill in with your actual object label indices
     0: "cat",
     1: "dog",
     2: "car",
-    # ... fill in all your categories ...
+    # ...
 }
-
-bin_label_names = ["AI-Generated", "Real"]  # Adjust if 0=AI, 1=Real
-
-########################
-# 3) Define Transforms
-########################
-
-# Match the transforms you used during validation
-val_transforms = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225])
+bin_label_names = ["AI-Generated", "Real"]
+
+#####################################
+# 5) Validation Transforms
+#####################################
+val_transforms = T.Compose([
+    T.Resize(256),
+    T.CenterCrop(224),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225])
 ])
 
-########################
-# 4) Define the Inference Function
-########################
-
+#####################################
+# 6) Inference Function
+#####################################
 def predict_image(img: Image.Image) -> str:
-    """
-    Takes a PIL image, applies transforms, passes through the model,
-    and returns the combined prediction (object + AI/Real).
-    """
-    # Convert to RGB just in case
     img = img.convert("RGB")
-    
-    # Apply transforms
-    img_t = val_transforms(img)
-    # Add batch dimension
-    img_t = img_t.unsqueeze(0)
-
+    img_t = val_transforms(img).unsqueeze(0)
     with torch.no_grad():
         obj_logits, bin_logits = model(img_t)
         obj_pred = torch.argmax(obj_logits, dim=1).item()
         bin_pred = torch.argmax(bin_logits, dim=1).item()
-
-    # Map predictions to labels
     obj_name = idx_to_obj_label.get(obj_pred, "Unknown")
     bin_name = bin_label_names[bin_pred]
-
     return f"Object: {obj_name} | Authenticity: {bin_name}"
 
-########################
-# 5) Build Gradio UI
-########################
-
+#####################################
+# 7) Gradio UI
+#####################################
 demo = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
@@ -85,15 +85,10 @@ demo = gr.Interface(
     title="Multi-Task Image Classifier",
     description=(
         "Upload an image to get two predictions: "
-        "1) The primary object (from pseudo-labeling), "
-        "2) Whether the image is AI-generated or real."
+        "1) The primary object, 2) Whether the image is AI-generated or real."
     )
 )
 
-########################
-# 6) Launch the App
-########################
-
 def main():
     demo.launch(server_name="0.0.0.0", enable_queue=True)