Initial commit

models/loader.py

@@ -49,9 +49,7 @@ class ModelLoader:
             self.model_registry[ModelType.VQVAE] = VQVAE
         except ImportError:
             print("Warning: VQVAE model implementation not found")
-
-        # Temporary removal of CNN and GAN models
-        """
+            
         try:
             from .cnn import CNN
             self.model_registry[ModelType.SRCNN] = CNN
@@ -60,11 +58,9 @@ class ModelLoader:
             
         try:
             from .gan import UncertainESRGAN
-            self.model_registry[ModelType.GAN] = UncertainESRGAN        
+            self.model_registry[ModelType.GAN] = UncertainESRGAN
         except ImportError:
-            print("Warning: Model implementation not found")
-        """
-        
+            print("Warning: GAN model implementation not found")
 
     def load_model(self, model_type: str, checkpoint_path: str, config_overrides: Optional[Dict] = None):
         """
@@ -94,7 +90,12 @@ class ModelLoader:
             # Initialize model with potentially modified config
             model_class = self.model_registry[model_type]
             model = model_class(**model_config)
-            model = torch.compile(model)
+            
+            # Apply torch compilation for performance
+            try:
+                model = torch.compile(model)
+            except Exception as e:
+                print(f"Warning: Could not compile model: {e}")
             
             # Move model to device
             model = model.to(self.device)
@@ -102,35 +103,31 @@ class ModelLoader:
             # Load checkpoint
             try:
                 checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=True)
-                print(f"Successfully loaded {model_type.value} model from {checkpoint_path}")
-
-                model.load_state_dict(checkpoint['state_dict'], strict = False)
-
-                print(checkpoint['state_dict'].keys())
-                print(checkpoint.keys())                    
-
-                model.uncertainty_tracker.ema_errors = checkpoint['state_dict'].get(
-                    '_orig_mod.uncertainty_tracker.ema_errors',
-                    torch.zeros(model.uncertainty_tracker.block_size**2).to(self.device)
-                )
-                model.uncertainty_tracker.ema_quantile = checkpoint['state_dict'].get(
-                    '_orig_mod.uncertainty_tracker.ema_quantile',
-                    torch.zeros(model.uncertainty_tracker.block_size**2).to(self.device)
-                )
-                model.uncertainty_tracker._initialize_buffers(64, 64, self.device)
-
-                # Store model
-
-                model.uncertainty_tracker.ema_errors = checkpoint['state_dict']['_orig_mod.uncertainty_tracker.ema_errors']
-                model.uncertainty_tracker.ema_quantile = checkpoint['state_dict']['_orig_mod.uncertainty_tracker.ema_quantile']
-                
-                model.uncertainty_tracker.calibrated = checkpoint.get('calibrated', True)
+                print(f"Checkpoint keys: {checkpoint.keys()}")
+                model.load_state_dict(checkpoint['state_dict'], strict=False)
                 
-                if model.uncertainty_tracker.calibrated:
-                    if 'block_scale_means' in checkpoint:
-                        model.uncertainty_tracker.block_scale_means = checkpoint['block_scale_means']
-                    if 'block_scale_stds' in checkpoint:
-                        model.uncertainty_tracker.block_scale_stds = checkpoint['block_scale_stds']
+                # Explicitly load uncertainty tracker buffers
+                if hasattr(model, 'uncertainty_tracker'):
+                    # Load EMA buffers
+                    model.uncertainty_tracker.ema_errors = checkpoint['state_dict'].get(
+                        '_orig_mod.uncertainty_tracker.ema_errors',
+                        torch.zeros(model.uncertainty_tracker.block_size**2).to(self.device)
+                    )
+                    model.uncertainty_tracker.ema_quantile = checkpoint['state_dict'].get(
+                        '_orig_mod.uncertainty_tracker.ema_quantile',
+                        torch.zeros(model.uncertainty_tracker.block_size**2).to(self.device)
+                    )
+                    
+                    # Initialize buffers with proper dimensions (64x64 for super-resolution output)
+                    model.uncertainty_tracker._initialize_buffers(64, 64, self.device)
+                    
+                    # Load calibration state if available
+                    model.uncertainty_tracker.calibrated = checkpoint.get('calibrated', False)
+                    if model.uncertainty_tracker.calibrated:
+                        if 'block_scale_means' in checkpoint:
+                            model.uncertainty_tracker.block_scale_means = checkpoint['block_scale_means']
+                        if 'block_scale_stds' in checkpoint:
+                            model.uncertainty_tracker.block_scale_stds = checkpoint['block_scale_stds']
                 
                 print(f"Successfully loaded {model_type.value} model from {checkpoint_path}")
             except Exception as e: