Upload app.py with huggingface_hub

app.py

@@ -197,7 +197,7 @@ class ModelManager:
         except Exception as e:
             logger.warning(f"ROI annotations load failed: {e}")
         
-        # Load brain encoder
+        # Load brain encoder (optional - ridge is primary)
         try:
             model_path = hf_hub_download(repo_id=MODEL_REPO, filename="best_model.pt")
             checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
@@ -215,10 +215,13 @@ class ModelManager:
             if self.roi_annotations is not None:
                 self.brain_encoder.set_roi_assignments(self.roi_annotations)
             
+            # Free checkpoint memory
+            del checkpoint
+            
             logger.info("Brain encoder loaded successfully")
         except Exception as e:
-            logger.error(f"Brain encoder load failed: {e}")
-            raise
+            logger.warning(f"Brain encoder load failed (will use ridge only): {e}")
+            self.brain_encoder = None
         
         # Load ridge model
         try:
@@ -404,33 +407,46 @@ class ModelManager:
             pred_np = predictions.cpu().numpy().flatten()
         
         # ── Deep encoder for intermediates and uncertainty ──
-        with torch.no_grad():
-            deep_pred, intermediates = self.brain_encoder(input_features, return_intermediates=True)
-        
-        # Compute modality contributions using ridge (more reliable)
+        intermediates = {}
+        if self.brain_encoder is not None:
+            with torch.no_grad():
+                deep_pred, intermediates = self.brain_encoder(input_features, return_intermediates=True)
+            
+            # Compute uncertainty via dropout MC (deep encoder)
+            self.brain_encoder.train()
+            mc_predictions = []
+            for _ in range(10):
+                with torch.no_grad():
+                    mc_pred = self.brain_encoder(input_features)
+                    mc_predictions.append(mc_pred.cpu().numpy().flatten())
+            self.brain_encoder.eval()
+            mc_predictions = np.array(mc_predictions)
+            uncertainty = np.std(mc_predictions, axis=0)
+        else:
+            # Estimate uncertainty from ridge prediction variance across feature perturbation
+            mc_predictions = []
+            for _ in range(10):
+                noise = np.random.normal(0, 0.01, size=input_features_np.shape)
+                X_noisy = (input_features_np + noise - ridge['feat_mean']) / ridge['feat_std']
+                mp = ridge['model'].predict(X_noisy).flatten()
+                mc_predictions.append(mp)
+            mc_predictions = np.array(mc_predictions)
+            uncertainty = np.std(mc_predictions, axis=0)
+        
+        # Compute modality contributions using ridge
         modality_contributions = {}
         if self.ridge_model is not None:
+            ridge = self.ridge_model
             for key in ['image_multi_layer', 'text_multi_layer', 'audio_multi_layer']:
                 if key in features_dict:
                     modality_name = key.split('_')[0]
                     feat_np = features_dict[key].cpu().numpy()
                     X_n = (feat_np - ridge['feat_mean']) / ridge['feat_std']
                     mp = (ridge['model'].predict(X_n) * ridge['fmri_std'] + ridge['fmri_mean']).flatten()
-                    mp = np.clip(mp, -clip_val, clip_val)
+                    clip_val_mod = np.percentile(np.abs(mp), 99.5)
+                    mp = np.clip(mp, -clip_val_mod, clip_val_mod)
                     modality_contributions[modality_name] = mp
         
-        # Compute uncertainty via dropout MC (deep encoder)
-        self.brain_encoder.train()
-        mc_predictions = []
-        for _ in range(10):
-            with torch.no_grad():
-                mc_pred = self.brain_encoder(input_features)
-                mc_predictions.append(mc_pred.cpu().numpy().flatten())
-        self.brain_encoder.eval()
-        
-        mc_predictions = np.array(mc_predictions)
-        uncertainty = np.std(mc_predictions, axis=0)
-        
         # Compute ROI summaries using z-scored per-voxel predictions
         # This shows which regions are MORE or LESS activated compared to baseline
         if self.ridge_model is not None: