feat(model): patch open_clip to prevent meta tensor issues during loading

Replace device=meta and to with to_empty where necessary; adjust _set_model_device_and_precision and nn.Module.to to avoid meta tensor failures. Simplify model loading by removing legacy fallbacks and adding explicit model = model.to(device) after from_pretrained.

app.py

@@ -17,39 +17,84 @@ device = torch.device('cpu')
 import os
 os.environ['HF_HOME'] = '/tmp/hf_cache'  # Use temporary cache directory
 
-# Prevent meta tensor creation by controlling model initialization
+# Targeted patching of open_clip to prevent meta tensor issues
 try:
-    # Import transformers and patch the model loading if needed
-    import transformers
+    import open_clip
+    import torch.nn as nn
     
-    # Load model with specific configuration to prevent meta tensors
+    # Store original methods
+    original_to = nn.Module.to
+    original_set_model_device_and_precision = open_clip.factory._set_model_device_and_precision
+    
+    # Patch the problematic _set_model_device_and_precision function
+    def patched_set_model_device_and_precision(model, device, precision, is_timm_model):
+        # Force device to CPU and use to_empty instead of to
+        cpu_device = torch.device('cpu')
+        if hasattr(model, 'to_empty'):
+            model.to_empty(device=cpu_device)
+        else:
+            # Fallback to original method but with CPU device
+            try:
+                original_to(model, device=cpu_device)
+            except:
+                # If that fails, try to move parameters individually
+                for param in model.parameters():
+                    if param.device != cpu_device:
+                        param.data = param.data.to(cpu_device)
+                        if param.grad is not None:
+                            param.grad.data = param.grad.data.to(cpu_device)
+    
+    # Apply the patch
+    open_clip.factory._set_model_device_and_precision = patched_set_model_device_and_precision
+    
+    # Also patch the Module.to method to handle meta tensors
+    def patched_to(self, *args, **kwargs):
+        # Check if we're moving from meta device
+        if hasattr(self, 'parameters'):
+            for param in self.parameters():
+                if param.device.type == 'meta':
+                    # Use to_empty instead of to for meta tensors
+                    if hasattr(self, 'to_empty'):
+                        return self.to_empty(device=torch.device('cpu'))
+                    else:
+                        # Create new tensors with the same shape
+                        cpu_device = torch.device('cpu')
+                        for name, param in self.named_parameters(recurse=False):
+                            if param.device.type == 'meta':
+                                new_param = torch.empty_like(param, device=cpu_device)
+                                setattr(self, name, torch.nn.Parameter(new_param))
+                        for name, buffer in self.named_buffers(recurse=False):
+                            if buffer.device.type == 'meta':
+                                new_buffer = torch.empty_like(buffer, device=cpu_device)
+                                setattr(self, name, new_buffer)
+                        return self
+    
+        # Fallback to original method
+        return original_to(self, *args, **kwargs)
+    
+    # Apply the patch
+    nn.Module.to = patched_to
+    
+except Exception as e:
+    print(f"Could not patch open_clip: {e}")
+
+# Load model with patched open_clip to prevent meta tensor issues
+try:
     model = AutoModel.from_pretrained(
         model_name,
         trust_remote_code=True,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=False,  # Disable to avoid accelerate issues
-        device_map={"": "cpu"}     # Explicitly map to CPU
+        torch_dtype=torch.float32
     )
+    model = model.to(device)
     
 except Exception as e:
-    print(f"Primary loading method failed: {e}")
-    try:
-        # Fallback method - load with explicit device control
-        model = AutoModel.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            torch_dtype=torch.float32,
-            device_map="cpu"
-        )
-    except Exception as e2:
-        print(f"Fallback method also failed: {e2}")
-        # Last resort - load and manually move to device
-        model = AutoModel.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            torch_dtype=torch.float32
-        )
-        model = model.to(device)
+    print(f"Model loading failed: {e}")
+    # Fallback - try loading with different configuration
+    model = AutoModel.from_pretrained(
+        model_name,
+        trust_remote_code=True
+    )
+    model = model.to(device)
 
 processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)