Upload folder using huggingface_hub

test/check_pearson.py

@@ -16,7 +16,7 @@ sys.path.append(str(src_dir))
 
 from training import make_data_loaders, evaluate_stage1, evaluate_stage2, SUBJECTS
 from medarc_architecture import MultiSubjectConvLinearEncoder
-from matcha_architecture import CFM
+from stage2.CFM import CFM
 
 def plot_heatmap(acc_map, title, save_path):
     """
@@ -108,20 +108,19 @@ def main():
 
     print("\n--- Stage 2 Checkout ---")
     cfm_params = cfg.stage2.cfm
-    decoder_params = cfg.stage2.decoder
-    latent_dim = cfg.stage2.get("latent_dim", 128)
+    velocity_net_params = cfg.stage2.velocity_net
+    source_ve_params = cfg.stage2.source_ve
+    transport_params = cfg.stage2.transport
     stage2_models = nn.ModuleDict()
 
     for sub in subjects_list:
         sub_key = str(sub)
         cfm_model = CFM(
-            in_channels=2 * target_dim,
-            out_channel=target_dim,
+            feat_dim=target_dim,
             cfm_params=cfm_params,
-            decoder_params=decoder_params,
-            n_spks=1,
-            voxel_dim=target_dim,
-            latent_dim=latent_dim,
+            velocity_net_params=velocity_net_params,
+            source_ve_params=source_ve_params,
+            transport_params=transport_params,
         ).to(device)
         stage2_models[sub_key] = cfm_model
 

test/load_model.py

@@ -54,17 +54,16 @@ def export_stage2(cfg, target_dim=1000):
     print('\nExporting Stage 2...')
     
     cfm_params = cfg.stage2.cfm
-    decoder_params = cfg.stage2.decoder
-    latent_dim = cfg.stage2.get('latent_dim', 128)
+    velocity_net_params = cfg.stage2.velocity_net
+    source_ve_params = cfg.stage2.source_ve
+    transport_params = cfg.stage2.transport
     
     model = CFM(
-        in_channels=2 * target_dim,
-        out_channel=target_dim,
+        feat_dim=target_dim,
         cfm_params=cfm_params,
-        decoder_params=decoder_params,
-        n_spks=1,
-        voxel_dim=target_dim,
-        latent_dim=latent_dim,
+        velocity_net_params=velocity_net_params,
+        source_ve_params=source_ve_params,
+        transport_params=transport_params,
     )
     
     weights_path = ROOT / 'output' / 'debug_run' / 'stage2_epoch_0.pt'
@@ -74,22 +73,31 @@ def export_stage2(cfg, target_dim=1000):
         print('Loaded Stage 2 weights.')
         
     model.eval()
+
+    class Stage2ExportWrapper(nn.Module):
+        def __init__(self, cfm_model: nn.Module, steps: int):
+            super().__init__()
+            self.cfm_model = cfm_model
+            self.steps = steps
+
+        def forward(self, mu):
+            return self.cfm_model(mu, n_timesteps=self.steps)
+
+    export_wrapper = Stage2ExportWrapper(model, steps=10)
+    export_wrapper.eval()
     
     # Dummy input: (B, C, T) 
     B, C, T = 1, target_dim, 10
     mu = torch.randn(B, C, T)
-        
-    mask = torch.ones(B, 1, T)
-    n_timesteps = torch.tensor(10)
     
     save_path = 'stage2_model.onnx'
     try:
         torch.onnx.export(
-            model,
-            (mu, mask, n_timesteps),
+            export_wrapper,
+            mu,
             save_path,
             opset_version=14,
-            input_names=['mu', 'mask', 'n_timesteps'],
+            input_names=['mu'],
             output_names=['output']
         )
         print(f'Saved {save_path}')
@@ -102,7 +110,17 @@ if __name__ == '__main__':
         cfg = OmegaConf.load(config_path)
     else:
         # Dummy config if config not found
-        cfg = OmegaConf.create({'stage1': {'model': {}}, 'stage2': {'cfm': {}, 'decoder': {'channels': [32, 32], 'dropout': 0.0, 'attention_head_dim': 16, 'n_blocks': 1, 'num_mid_blocks': 1, 'num_heads': 2, 'act_fn': 'snakebeta', 'down_block_type': 'transformer', 'mid_block_type': 'transformer', 'up_block_type': 'transformer'}}})
+        cfg = OmegaConf.create(
+            {
+                'stage1': {'model': {}},
+                'stage2': {
+                    'cfm': {'solver': 'euler', 'kld_weight': 3.0, 'kld_target_std': 1.0, 'detach_ut': False},
+                    'velocity_net': {'hidden_dim': 256, 'n_blocks': 2, 'n_heads': 4, 'dropout': 0.1, 'max_seq_len': 2048, 'temporal_attn_layers': 1},
+                    'source_ve': {'depth': 2, 'num_heads': 4, 'num_queries': 8, 'dropout': 0.1, 'use_variational': True},
+                    'transport': {'path_type': 'Linear', 'prediction': 'velocity', 'time_dist_type': 'uniform', 'time_dist_shift': 1.0},
+                },
+            }
+        )
 
     export_stage1(cfg)
     export_stage2(cfg, target_dim=1000)

test/overfit_test.py

@@ -68,7 +68,6 @@ def main():
     i = args.subject_idx
     x1 = fmri[:, i].transpose(1, 2)        # (B, V, T)
     mu = mu_anchor[:, i].transpose(1, 2)    # (B, V, T)
-    mask = torch.ones(x1.shape[0], 1, x1.shape[2], device=device)
 
     # --- Stage 2: single-subject CFM ---
     cfm = CFM(