Update app.py

app.py

@@ -1,3 +1,35 @@
+from unittest.mock import MagicMock
+
+# Fix 1: Mock device_mesh for torch compatibility
+class MockDeviceMesh:
+    def init_device_mesh(self, *args, **kwargs):
+        return None
+
+if 'torch.distributed.device_mesh' not in sys.modules:
+    sys.modules['torch.distributed.device_mesh'] = MockDeviceMesh()
+
+# Fix 2: Mock functional_tensor for torchvision compatibility
+try:
+    from torchvision.transforms import functional_tensor
+except ImportError:
+    mock_functional_tensor = MagicMock()
+    
+    try:
+        from torchvision.transforms.functional import rgb_to_grayscale
+        mock_functional_tensor.rgb_to_grayscale = rgb_to_grayscale
+    except ImportError:
+        def rgb_to_grayscale(img, num_output_channels=1):
+            import torch
+            if len(img.shape) == 4:
+                return img.mean(dim=1, keepdim=True)
+            elif len(img.shape) == 3:
+                return img.mean(dim=0, keepdim=True)
+            return img
+        mock_functional_tensor.rgb_to_grayscale = rgb_to_grayscale
+    
+    sys.modules['torchvision.transforms.functional_tensor'] = mock_functional_tensor
+
+# Now continue with original imports
 import spaces
 import time
 import os
@@ -14,6 +46,7 @@ from pulid.pipeline_flux import PuLIDPipeline
 from pulid.utils import resize_numpy_image_long
 
 
+
 def get_models(name: str, device: torch.device, offload: bool):
     t5 = load_t5(device, max_length=128)
     clip = load_clip(device)