Infer model architecture dimensions dynamically from HF config.json and clean up YAML configs

Author: amepas

src/maxdiffusion/configs/base_flux2klein_9B.yml

@@ -270,9 +270,6 @@ width: 1024
 batch_size: 4
 interactive: False
 
-# 9B Architecture Dimensions
-depth: 24 # num_single_layers
-num_double_layers: 8
-hidden_size: 4096
-num_attention_heads: 32
+# Note: Architecture dimensions (depth, num_double_layers, num_attention_heads) are
+# automatically inferred from pretrained_model_name_or_path (transformer/config.json).
 

src/maxdiffusion/generate_flux2klein.py

@@ -168,10 +168,10 @@ def main(argv):
     pyconfig._config.keys["logical_axis_rules"] = tuple(new_rules)
 
   # 3. Resolve weights repository snapshots
-  is_9B = config.depth == 24
   repo_id = getattr(config, "pretrained_model_name_or_path", None)
   if not repo_id:
-    repo_id = "black-forest-labs/FLUX.2-klein-9B" if is_9B else "black-forest-labs/FLUX.2-klein-4B"
+    depth_val = getattr(config, "depth", None)
+    repo_id = "black-forest-labs/FLUX.2-klein-9B" if depth_val == 24 else "black-forest-labs/FLUX.2-klein-4B"
   max_logging.log(f"Target model detected: {repo_id}")
 
   from huggingface_hub import snapshot_download