Upload 2 files

model.py

@@ -0,0 +1,62 @@
+import os
+import torch
+import torch.nn as nn
+import torchvision.models as models
+from transformers import PreTrainedModel, AutoConfig
+
+# Define the model architecture based on EfficientNetV2-S
+class AIDetectorModel(nn.Module):
+    def __init__(self):
+        super(AIDetectorModel, self).__init__()
+        # Load EfficientNetV2-S as base model
+        self.base_model = models.efficientnet_v2_s(weights=None)
+        
+        # Replace classifier with custom layers
+        self.base_model.classifier = nn.Sequential(
+            nn.Linear(self.base_model.classifier[1].in_features, 1024),
+            nn.ReLU(),
+            nn.Dropout(p=0.3),
+            nn.Linear(1024, 512),
+            nn.ReLU(),
+            nn.Dropout(p=0.3),
+            nn.Linear(512, 2)  # 2 classes: real or AI-generated
+        )
+    
+    def forward(self, x):
+        return self.base_model(x)
+
+# Wrapper class to make the model compatible with Hugging Face
+class AIDetectorForImageClassification(PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = AIDetectorModel()
+        
+        # Load the trained weights
+        model_path = os.path.join(os.getcwd(), "best_model_improved.pth")
+        try:
+            # Try to load with strict=True first
+            self.model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
+            print(f"Model loaded successfully from {model_path}")
+        except Exception as e:
+            print(f"Error with strict loading: {e}")
+            print("Trying with strict=False...")
+            # If that fails, try with strict=False
+            self.model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")), strict=False)
+            print("Model loaded with strict=False")
+    
+    def forward(self, pixel_values, labels=None, **kwargs):
+        logits = self.model(pixel_values)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+        
+        return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}
+
+# Function to create and load the model
+def get_model():
+    config = AutoConfig.from_pretrained("./")
+    model = AIDetectorForImageClassification(config)
+    return model

preprocessor_config.json

@@ -0,0 +1,10 @@
+{
+  "do_normalize": true,
+  "do_resize": true,
+  "image_mean": [0.485, 0.456, 0.406],
+  "image_std": [0.229, 0.224, 0.225],
+  "size": {
+    "height": 224,
+    "width": 224
+  }
+}