Upload app.py

app.py

@@ -229,18 +229,6 @@ def list_available_checkpoints_from_hub(repo_id: str) -> Dict[str, str]:
             ckpts[disp] = f
     return ckpts
 
-def _shape_compatible_load(model: nn.Module, state: Dict):
-    """
-    Load only keys whose shapes match, warn about the rest.
-    This mirrors your flexible loader used elsewhere, but keeps strict as the default path.
-    """
-    model_k = model.state_dict()
-    filtered = {k: v for k, v in state.items() if k in model_k and model_k[k].shape == v.shape}
-    missing = sorted(set(model_k.keys()) - set(filtered.keys()))
-    if missing:
-        print(f"[WARN] Skipping {len(missing)} keys due to shape mismatch (likely classifier).")
-    model.load_state_dict(filtered, strict=False)
-
 _LOADED: Dict[str, nn.Module] = {}
 def load_model_from_hub(display_name: str, repo_filename: str) -> nn.Module:
     key = f"{MODEL_REPO}::{repo_filename}"
@@ -254,26 +242,15 @@ def load_model_from_hub(display_name: str, repo_filename: str) -> nn.Module:
         local_dir_use_symlinks=False
     )
 
-    # Resolve architecture (you already have this mapping logic earlier)
-    arch = DISPLAY_TO_ARCH.get(display_name) or infer_arch_from_filename(repo_filename)
+    arch = arch_from_filename(repo_filename)
     if arch not in MODEL_BUILDERS:
-        raise ValueError(f"Unknown or unsupported architecture for {display_name} / {repo_filename}")
+        raise RuntimeError(f"Unknown architecture inferred from filename: {arch}")
 
-    # Build fresh model and adapt head if checkpoint expects a Sequential head
     model = MODEL_BUILDERS[arch](NUM_CLASSES).to(DEVICE)
     state = torch.load(local_path, map_location="cpu")
     model = adapt_head_for_state_dict(model, arch, state, NUM_CLASSES)
-
-    # Strict-first, then fallback to shape-compatible load
-    try:
-        model.load_state_dict(state, strict=True)
-        print(f"[INFO] StrictLoad=True for {display_name} ({repo_filename})")
-    except Exception as e:
-        print(f"[ERROR] StrictLoad failed for {display_name} ({repo_filename}): {e}")
-        _shape_compatible_load(model, state)
-        print(f"[INFO] StrictLoad=False (shape-compatible fallback) for {display_name} ({repo_filename})")
-
-    model.to(DEVICE).eval()
+    model.load_state_dict(state, strict=True)
+    model.eval()
     _LOADED[key] = model
     return model
 
@@ -642,8 +619,11 @@ def build_app():
 
                         clear_btn = gr.Button("Clear files")
 
+                # Pre-define empty dataframe with column headers
+                empty_df = pd.DataFrame(columns=["Model 1", "Model 2", "Ensemble"])
+
                 # Predictions
-                results = gr.Dataframe(label="Predictions", wrap=True, interactive=False, type="pandas")
+                results = gr.Dataframe(label="Predictions", value=empty_df, wrap=True, interactive=False, type="pandas")
 
                 # Action
                 run_btn = gr.Button("🔍 Run inference", variant="primary")