Update app.py

app.py

@@ -35,47 +35,33 @@ def load_model_from_pickle(pickle_path="best_model.pkl"):
         if not os.path.exists(pickle_path):
             return f"❌ Model file not found: {pickle_path}\n\nPlease ensure best_model.pkl is uploaded to the HuggingFace Space."
         
-        # Strategy: Use a custom unpickler that forces all CUDA tensors to CPU during deserialization
-        # This prevents the "Attempting to deserialize object on a CUDA device" error entirely
-        class CPUUnpickler(pickle.Unpickler):
-            def find_class(self, module, name):
-                # Remap torch CUDA classes to CPU equivalents
-                if module == 'torch.storage' and name.endswith('Storage'):
-                    # Force all storage types to CPU (e.g., cuda:0 -> cpu)
-                    return super().find_class('torch', name.replace('Cuda', '').replace('cuda', ''))
-                return super().find_class(module, name)
-            
-            def persistent_load(self, pid):
-                # Handle torch storage tuples: ('storage', storage_type, key, location, size)
-                if isinstance(pid, tuple) and len(pid) > 0 and pid[0] == 'storage':
-                    storage_type, key, location, size = pid[1], pid[2], pid[3], pid[4]
-                    # Force location to CPU regardless of original device
-                    if 'cuda' in str(location).lower():
-                        location = 'cpu'
-                    # Rebuild the storage descriptor with CPU location
-                    return ('storage', storage_type, key, location, size)
-                return pid
+        # ULTIMATE FIX: Monkey-patch torch.cuda.is_available during unpickling
+        # This tricks PyTorch into thinking CUDA is available, then we immediately move to CPU
+        import torch.cuda
+        import torch.serialization
+        
+        original_is_available = torch.cuda.is_available
+        original_device_count = torch.cuda.device_count
+        
+        def fake_is_available():
+            return True
+        
+        def fake_device_count():
+            return 1
         
         try:
-            # Attempt 1: Custom CPU unpickler (most robust for CUDA->CPU)
+            # Temporarily make PyTorch think CUDA is available
+            torch.cuda.is_available = fake_is_available
+            torch.cuda.device_count = fake_device_count
+            
+            # Load the pickle (PyTorch will accept CUDA device references now)
             with open(pickle_path, 'rb') as f:
-                model_package = CPUUnpickler(f).load()
-        except Exception as e1:
-            try:
-                # Attempt 2: torch.load with map_location (standard approach)
-                model_package = torch.load(pickle_path, map_location=torch.device('cpu'))
-            except Exception as e2:
-                try:
-                    # Attempt 3: torch.load with weights_only=False (for older torch versions)
-                    model_package = torch.load(pickle_path, map_location='cpu', weights_only=False)
-                except Exception as e3:
-                    # All strategies failed - return combined error info
-                    return (f"❌ All loading strategies failed:\n\n"
-                            f"1. Custom CPU unpickler: {str(e1)[:100]}\n"
-                            f"2. torch.load(map_location): {str(e2)[:100]}\n"
-                            f"3. torch.load(weights_only=False): {str(e3)[:100]}\n\n"
-                            f"The pickle may be corrupted or incompatible with this PyTorch version.\n"
-                            f"Try re-saving with: model.to('cpu'); torch.save({{'model': model, ...}}, 'file.pkl')")
+                model_package = torch.load(f, map_location=lambda storage, loc: storage.cpu())
+        
+        finally:
+            # Always restore original functions
+            torch.cuda.is_available = original_is_available
+            torch.cuda.device_count = original_device_count
         
         # Success! Model loaded with one of the strategies above
         # Handle a few common package shapes.