Update app.py

app.py

@@ -7,23 +7,41 @@ import torch
 import numpy as np
 
 def preprocess_image(img):
+    print(f"Input image type: {type(img)}")
+    print(f"Input image shape: {img.shape if hasattr(img, 'shape') else 'No shape attribute'}")
+    
     # Convert numpy array to Tensor and ensure correct shape
     if isinstance(img, np.ndarray):
         img = K.image_to_tensor(img, keepdim=False).float() / 255.0
+    elif isinstance(img, torch.Tensor):
+        img = img.float()
+        if img.max() > 1.0:
+            img = img / 255.0
+    else:
+        raise ValueError(f"Unsupported image type: {type(img)}")
+    
+    print(f"After conversion to tensor - shape: {img.shape}")
     
     # Ensure 3D tensor (C, H, W)
     if img.ndim == 2:
         img = img.unsqueeze(0)
+    elif img.ndim == 3 and img.shape[0] not in [1, 3]:
+        img = img.permute(2, 0, 1)
+    
+    print(f"After ensuring 3D - shape: {img.shape}")
     
     # Ensure 3 channel image
     if img.shape[0] == 1:
-        img = img.repeat(3, 1, 1)
+        img = img.expand(3, -1, -1)
     elif img.shape[0] > 3:
         img = img[:3]  # Take only the first 3 channels if more than 3
     
+    print(f"After ensuring 3 channels - shape: {img.shape}")
+    
     # Add batch dimension
     img = img.unsqueeze(0)
     
+    print(f"Final tensor shape: {img.shape}")
     return img
 
 def inference(img_1, img_2):