Update app.py

app.py

@@ -97,16 +97,31 @@ def load_checkpoint_model(model_path, device):
         print("Failed to load checkpoint file:", e)
         logger.exception("Failed to load checkpoint file:")
         ckpt = {}
-    # unwrap common dict wrapper
-    if isinstance(ckpt, dict) and 'model_state_dict' in ckpt:
-        state_dict = ckpt['model_state_dict']
-    else:
-        # if checkpoint is a state dict directly
-        state_dict = ckpt if isinstance(ckpt, dict) else {}
+    # unwrap common dict wrapper (support both 'model_state_dict' and 'state_dict')
+    state_dict = {}
+    if isinstance(ckpt, dict):
+        if 'model_state_dict' in ckpt and isinstance(ckpt['model_state_dict'], dict):
+            state_dict = ckpt['model_state_dict']
+        elif 'state_dict' in ckpt and isinstance(ckpt['state_dict'], dict):
+            state_dict = ckpt['state_dict']
+        else:
+            # fallback: ckpt may already be a state dict
+            state_dict = ckpt
+
+    # If the state_dict is a single-key wrapper (e.g., {'state_dict': {...}} or {'model': {...}}), unwrap one more level
+    if isinstance(state_dict, dict) and len(state_dict) == 1:
+        sole_val = next(iter(state_dict.values()))
+        if isinstance(sole_val, dict):
+            # adopt inner dict as state_dict if it looks like parameters
+            inner_keys = list(sole_val.keys())[:8]
+            # Heuristic: keys with '.' and numeric shapes indicate a param dict
+            if any('.' in k for k in inner_keys):
+                logger.info(f"Unwrapping single-key checkpoint wrapper, inner keys sample: {inner_keys}")
+                state_dict = sole_val
 
     # Diagnostic: print a few state_dict keys so we can tell checkpoint format
     try:
-        sample_keys = list(state_dict.keys())[:8]
+        sample_keys = list(state_dict.keys())[:16]
         print("Checkpoint sample keys:", sample_keys)
         logger.info(f"Checkpoint sample keys: {sample_keys}")
     except Exception:
@@ -114,7 +129,13 @@ def load_checkpoint_model(model_path, device):
         logger.info("No state_dict keys to sample")
 
     # Heuristic: detect HF-style Swin checkpoint by looking for keys that start with 'swin.'
-    hf_like = any(k.startswith('swin.') or 'swin.embeddings' in k for k in state_dict.keys()) if state_dict else False
+    # Detect HF-like keys; strip common 'module.' prefix before checking
+    def key_is_hf_like(k: str) -> bool:
+        kk = k.replace('module.', '')
+        kk = kk.lower()
+        return kk.startswith('swin.') or 'swin.embeddings' in kk or 'swin.patch_embeddings' in kk
+
+    hf_like = any(key_is_hf_like(k) for k in state_dict.keys()) if state_dict else False
     hf_msg = f"hf_like_checkpoint_detected={hf_like} HF_AVAILABLE={HF_AVAILABLE} HF_MODEL_ID={'set' if HF_MODEL_ID else 'not-set'}"
     print(hf_msg)
     logger.info(hf_msg)