Load the model config from the single root config.yaml instead of weights/model.yaml

infer.py

@@ -24,10 +24,10 @@ MASK_ALPHA = 0.65
 
 
 def load_model(weights_dir, device):
-    """Build the segmenter from model.yaml and load model.safetensors.
+    """Build the segmenter from config.yaml (repo root) and load model.safetensors.
 
     Loaded in bfloat16: that's the precision the model was trained in."""
-    model = Swinv2Segmenter(weights_dir / "model.yaml")
+    model = Swinv2Segmenter(weights_dir.parent / "config.yaml")
     load_safetensors(model, weights_dir / "model.safetensors")
     return model.to(device).to(torch.bfloat16).eval()