Update inference.py

inference.py

@@ -29,6 +29,14 @@ MODEL_URLS = {
     'resnet50_robust_face': 'https://huggingface.co/ttoosi/resnet50_robust_face/resolve/main/resnet50_imagenet_L2_eps_0.50_checkpoint150.pt'
 }
 
+# Per-model input size and normalization
+MODEL_PREPROC = {
+    "resnet50_robust":      {"size": 224, "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
+    "resnet50_standard":    {"size": 224, "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
+    # Typical for face models trained ArcFace/InsightFace-style
+    "resnet50_robust_face": {"size": 112, "mean": [0.5, 0.5, 0.5],      "std": [0.5, 0.5, 0.5]},
+}
+
 IMAGENET_MEAN = [0.485, 0.456, 0.406]
 IMAGENET_STD = [0.229, 0.224, 0.225]
 
@@ -282,7 +290,8 @@ def get_inference_configs(inference_type='IncreaseConfidence', eps=0.5, n_itr=50
 class GenerativeInferenceModel:
     def __init__(self):
         self.models = {}
-        self.normalizer = NormalizeByChannelMeanStd(IMAGENET_MEAN, IMAGENET_STD).to(device)
+        #self.normalizer = NormalizeByChannelMeanStd(IMAGENET_MEAN, IMAGENET_STD).to(device)
+        self.model_preproc = {}
         self.labels = get_imagenet_labels()
         
     def verify_model_integrity(self, model, model_type):
@@ -337,21 +346,20 @@ class GenerativeInferenceModel:
             return True
     
     def load_model(self, model_type):
-        """Load model from checkpoint or use pretrained model."""
         if model_type in self.models:
             print(f"Using cached {model_type} model")
             return self.models[model_type]
-        
-        # Record loading time for performance analysis
+    
         start_time = time.time()
         model_path = download_model(model_type)
-        
-        # Create a sequential model with normalizer and ResNet50
+    
+        # pick preproc for this model
+        pre = MODEL_PREPROC.get(model_type, {"size": 224, "mean": IMAGENET_MEAN, "std": IMAGENET_STD})
+        normalizer = NormalizeByChannelMeanStd(pre["mean"], pre["std"]).to(device)
+        self.model_preproc[model_type] = pre
+    
         resnet = models.resnet50()
-        model = nn.Sequential(
-            self.normalizer,  # Normalizer is part of the model sequence
-            resnet
-        )
+        model = nn.Sequential(normalizer, resnet)
         
         # Load the model checkpoint
         if model_path: