Update app.py

app.py

@@ -22,24 +22,25 @@ def preprocess_image(img):
     
     print(f"After conversion to tensor - shape: {img.shape}")
     
-    # Ensure 3D tensor (C, H, W)
+    # Ensure 4D tensor (B, 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)
+        img = img.unsqueeze(0).unsqueeze(0)
+    elif img.ndim == 3:
+        if img.shape[0] in [1, 3]:
+            img = img.unsqueeze(0)
+        else:
+            img = img.unsqueeze(1)
+    elif img.ndim == 4:
+        if img.shape[1] not in [1, 3]:
+            img = img.permute(0, 3, 1, 2)
     
-    print(f"After ensuring 3D - shape: {img.shape}")
+    print(f"After ensuring 4D - shape: {img.shape}")
     
     # Ensure 3 channel image
-    if img.shape[0] == 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)
+    if img.shape[1] == 1:
+        img = img.repeat(1, 3, 1, 1)
+    elif img.shape[1] > 3:
+        img = img[:, :3]  # Take only the first 3 channels if more than 3
     
     print(f"Final tensor shape: {img.shape}")
     return img
@@ -70,4 +71,4 @@ with gr.Blocks(theme='huggingface') as demo_app:
     gr.Examples(examples=examples, inputs=[input_image1, input_image2])
 
 if __name__ == "__main__":
-    demo_app.launch()
+    demo_app.launch(share=True)