fix:config

apps/protein_folding/helixfold3/helixfold/config/helixfold.yaml

@@ -65,7 +65,7 @@ CONFIG_DIFFS:
   model:
     global_config:
       subbatch_size: 96 # model.global_config.subbatch_size
-  num_recycle: 3 # model.num_recycle
-  heads:
-    confidence_head:
-      weight: 0.0 # model.heads.confidence_head.weight
+    num_recycle: 3 # model.num_recycle
+    heads:
+      confidence_head:
+        weight: 0.0 # model.heads.confidence_head.weight

apps/protein_folding/helixfold3/helixfold/inference.py

@@ -490,8 +490,7 @@ def main(cfg: DictConfig):
     msa_templ_data_pipeline_dict = get_msa_templates_pipeline(cfg=cfg)
 
     ### Create model
-    model_config = config.model_config(cfg.job_id)
-    #print(f'>>> model_config:\n{model_config}')
+    model_config = config.model_config(cfg.CONFIG_DIFFS)
 
     model = RunModel(model_config)
 

apps/protein_folding/helixfold3/helixfold/model/config.py

@@ -41,18 +41,17 @@ def model_config(config_diffs: Union[str, DictConfig, dict[str, dict[str, Any]]]
       print(f'Updated config from `CONFIG_DIFFS.{preset_name}`: {updated_config}')
 
       return cfg
-    if 'model' in config_diffs and (updated_config := config_diffs['model']) is not None:
-      cfg.update(dict(updated_config))
-      print(f'Updated config from `CONFIG_DIFFS`: {updated_config}')
-
+
+    # update from detailed configuration
+    if any(root_kw in config_diffs for root_kw in CONFIG_ALLATOM):
+      for root_kw in CONFIG_ALLATOM:
+        if root_kw in config_diffs and (updated_config := config_diffs[root_kw]) is not None:
+          cfg.update(dict(updated_config))
+          print(f'Updated config from `CONFIG_DIFFS`: {updated_config}')
       return cfg
 
   raise ValueError(f'Invalid config_diffs ({type(config_diffs)}): {config_diffs}')
 
-
-
-
-
 
 # preset for runs
 CONFIG_DIFFS = {