Upload 3 files

vae_model.py

@@ -318,22 +318,35 @@ class DemoVAE(BaseEstimator):
 
     def load(self, path):
         try:
-            # First try loading with weights_only=True (PyTorch 2.6+ default)
+            # Try different loading methods based on PyTorch version
+            print(f"Attempting to load model from {path}")
             try:
-                print(f"Attempting to load model with weights_only=True from {path}")
-                checkpoint = torch.load(path, weights_only=True)
-                print("Successfully loaded checkpoint with weights_only=True")
-            except (TypeError, ValueError):
-                # If weights_only parameter is not supported (older PyTorch) or fails
-                print("Falling back to load with default parameters")
-                # Add necessary global variables to the safe list if using PyTorch 2.6+
-                if hasattr(torch.serialization, 'add_safe_globals'):
+                # For PyTorch 2.6+, explicitly set weights_only=False for backward compatibility
+                if hasattr(torch, '__version__') and torch.__version__.startswith('2.6'):
                     import numpy as np
-                    # Add numpy scalar to safe globals
-                    torch.serialization.add_safe_globals(['numpy._core.multiarray.scalar', 
-                                                      'numpy.core.multiarray.scalar'])
-                checkpoint = torch.load(path)
-                print("Successfully loaded checkpoint with default parameters")
+                    # Add all necessary numpy types to safe globals list
+                    if hasattr(torch.serialization, 'add_safe_globals'):
+                        torch.serialization.add_safe_globals([
+                            'numpy._core.multiarray.scalar', 
+                            'numpy.core.multiarray.scalar',
+                            'numpy.ndarray',
+                            'numpy._globals._NoValue'
+                        ])
+                    with torch.serialization.safe_globals(['numpy._core.multiarray.scalar']):
+                        checkpoint = torch.load(path, weights_only=False)
+                else:
+                    # For older PyTorch versions
+                    checkpoint = torch.load(path)
+            except Exception as e:
+                print(f"Primary loading method failed: {str(e)}")
+                # Last resort - try with context manager if available
+                if hasattr(torch.serialization, 'safe_globals'):
+                    with torch.serialization.safe_globals(['numpy._core.multiarray.scalar', 'numpy.core.multiarray.scalar']):
+                        checkpoint = torch.load(path, weights_only=False)
+                else:
+                    # Fall back to default with no safety
+                    checkpoint = torch.load(path)
+            print("Successfully loaded checkpoint")
             
             # Initialize from checkpoint
             self.set_params(**checkpoint['params'])