Update inference_pytorch.py

inference_pytorch.py

@@ -31,6 +31,10 @@ class PyTorchInference:
         self.model_arch = model_arch
         self.device = device
 
+        # Hardcoded input size mapping - based on actual model definitions
+        self.input_size = self._get_input_size(model_arch)
+        print(f"Using input size: {self.input_size} for model {model_arch}")
+
         # Load checkpoint
         state_dict = self.load_checkpoint_state(checkpoint_path)
         state_dict = self.normalize_state_dict_keys(state_dict)
@@ -41,7 +45,7 @@ class PyTorchInference:
         if feature_dim is None:
             feature_dim = self.resolve_feature_dim(state_dict)
 
-        # Create model
+        # Create model - don't pass img_size, let the model use its default
         self.model = create_model(
             model_arch,
             pretrained=False,
@@ -54,9 +58,19 @@ class PyTorchInference:
         self.model.to(device)
         self.model.eval()
 
-        # Get transform
-        config = resolve_data_config({}, model=self.model)
+        # Get transform - override with our correct input size
+        # We manually set the input_size instead of relying on the model's config
+        config = resolve_data_config({'input_size': (3, self.input_size, self.input_size)}, model=self.model)
         self.transform = create_transform(**config)
+        print(f"Created transform with input size: {self.input_size}")
+
+    def _get_input_size(self, model_arch):
+        """Get input size based on model architecture - hardcoded to match actual model definitions"""
+        if model_arch == 'lsnet_xl_artist_448':
+            return 448
+        else:
+            # All other artist models use 224
+            return 224
 
     @staticmethod
     def load_checkpoint_state(checkpoint_path: str):
@@ -116,6 +130,7 @@ class PyTorchInference:
         """
         image = image.convert("RGB")
         tensor = self.transform(image)
+        print(f"Preprocessed image to tensor shape: {tensor.shape}")
         return tensor.unsqueeze(0)
 
     def predict(self, image, top_k=5, threshold=0.0):
@@ -134,7 +149,8 @@ class PyTorchInference:
 
         with torch.no_grad():
             input_tensor = input_tensor.to(self.device)
+            print(f"Running inference on tensor shape: {input_tensor.shape}")
             # Use return_features=False to get classification logits
             logits = self.model(input_tensor, return_features=False)
 
-        return logits.cpu().numpy()[0]
+        return logits.cpu().numpy()[0]