Update svae_cadence.py

svae_cadence.py

@@ -43,13 +43,14 @@ from geolip_svae.model_v2 import PatchSVAEv2
 # ═══════════════════════════════════════════════════════════════
 
 HF_REPO = 'AbstractPhil/geolip-SVAE'
-VERSION = 'v51_conduit_proto'
+VERSION = 'version2_v2_conduit_proto_2'
 LOCAL_DIR = f'/content/{VERSION}_checkpoints'
 LOG_PATH = os.path.join(LOCAL_DIR, 'training_log.json')
 
 CFG = dict(
     # Architecture (inherited from Fresnel v50)
-    V=48, D=4, ps=4, hidden=384, depth=4, n_cross=2,
+    V=16, D=4, ps=4, hidden=384, depth=4, n_cross=2,
+    stage_hidden=128, stage_V=64,
 
     # Training
     img_size=64,
@@ -60,7 +61,7 @@ CFG = dict(
     val_size=10000,
 
     # CV soft hand
-    target_cv=0.125,
+    target_cv=0.2915,
     cv_weight=0.3,
     boost=0.5,
     sigma=0.15,
@@ -273,6 +274,8 @@ def save_checkpoint(model, opt, sched, epoch, val_mse, log,
             'V': CFG['V'], 'D': CFG['D'], 'patch_size': CFG['ps'],
             'hidden': CFG['hidden'], 'depth': CFG['depth'],
             'n_cross_layers': CFG['n_cross'],
+            'stage_hidden': CFG.get('stage_hidden', 128),
+            'stage_V': CFG.get('stage_V', 16),
             'img_size': CFG['img_size'],
             'model_type': 'v2',
         },
@@ -329,6 +332,8 @@ def train():
         V=CFG['V'], D=D, ps=CFG['ps'],
         hidden=CFG['hidden'], depth=CFG['depth'],
         n_cross=CFG['n_cross'],
+        stage_hidden=CFG.get('stage_hidden', 128),
+        stage_V=CFG.get('stage_V', 16),
     ).to(device)
 
     n_params = sum(p.numel() for p in model.parameters())