support all variants with added vairant detection

app.py

@@ -3,6 +3,7 @@ import spaces
 import os, sys, importlib.util, re
 import gradio as gr
 from PIL import Image
+import torch
 
 # ——— Monkey-patch mmdet to remove its mmcv-version assertion ———
 spec = importlib.util.find_spec('mmdet')
@@ -23,8 +24,33 @@ def load_inferencer(checkpoint_path=None, device=None):
     kwargs = {'pose2d': 'rtmo', 'scope': 'mmpose', 'device': device, 'det_cat_ids': [0]}
     if checkpoint_path:
         kwargs['pose2d_weights'] = checkpoint_path
+        # detect model variant
+        variant = detect_rtmo_variant(checkpoint_path)
+        kwargs['pose2d'] = variant
     return MMPoseInferencer(**kwargs)
 
+def detect_rtmo_variant(checkpoint_path: str) -> str:
+    """
+    Inspect an RTMO .pth checkpoint and return its variant alias:
+    one of 'rtmo-l', 'rtmo-m', 'rtmo-s', 'rtmo-t', or 'unknown'.
+    """
+    ckpt = torch.load(checkpoint_path, map_location='cpu')
+    state_dict = ckpt.get('state_dict', ckpt)
+
+    key = 'backbone.stem.conv.conv.weight'
+    if key not in state_dict:
+        raise KeyError(f"Cannot find '{key}' in checkpoint.")
+
+    out_ch = state_dict[key].shape[0]
+
+    mapping = {
+    24: "rtmo-t_8xb32-600e_body7-416x416",
+    32: "rtmo-s_8xb32-600e_body7-640x640",
+    48: "rtmo-m_16xb16-600e_body7-640x640",
+    64: "rtmo-l_16xb16-600e_body7-640x640",
+    }
+    return mapping.get(out_ch, f'unknown (stem out_channels={out_ch})')
+
 # ——— Gradio prediction function ———
 @spaces.GPU()
 def predict(image: Image.Image, checkpoint):