Resolve class labels via id2label comma-separated synonyms

app.py

@@ -29,7 +29,13 @@ from model_catalog import (
     get_profile_by_label,
     parse_model_label,
 )
-from model_loader import PIPELINE_MANAGER, _to_float, _to_int, run_inference
+from model_loader import (
+    PIPELINE_MANAGER,
+    _to_float,
+    _to_int,
+    default_class_label_for_pipe,
+    run_inference,
+)
 
 
 DEFAULT_MODEL = MODEL_LABELS[0]
@@ -162,24 +168,33 @@ def _gpu_duration(
     return int(min(300, max(base, step_budget * 0.6 + 30)))
 
 
-def _load_model_core(model_label: str) -> str:
+def _load_model_core(model_label: str) -> tuple[str, str | None]:
     collection, variant = parse_model_label(model_label)
     message, _ = PIPELINE_MANAGER.load(collection, variant)
     PIPELINE_MANAGER.move_to_cuda()
-    return message
+    pipe = PIPELINE_MANAGER.pipe
+    profile = get_profile_by_label(model_label)
+    suggested_label = default_class_label_for_pipe(pipe, profile) if pipe is not None else None
+    return message, suggested_label
 
 
 @spaces.GPU(size="xlarge", duration=120)
-def _load_on_gpu(model_label: str) -> str:
+def _load_on_gpu(model_label: str) -> tuple[str, str | None]:
     return _load_model_core(model_label)
 
 
 def load_model(model_label: str):
     try:
-        message = _load_on_gpu(model_label)
+        message, suggested_label = _load_on_gpu(model_label)
     except Exception as exc:
         raise gr.Error(f"Failed to load `{model_label}`: {exc}") from exc
-    return (message, *_config_from_profile(get_profile_by_label(model_label)))
+    profile = get_profile_by_label(model_label)
+    config = _config_from_profile(profile)
+    if suggested_label:
+        config = list(config)
+        config[1] = gr.update(value=suggested_label)
+        config = tuple(config)
+    return (message, *config)
 
 
 @spaces.GPU(size="xlarge", duration=_gpu_duration)
@@ -318,7 +333,7 @@ def build_demo() -> gr.Blocks:
                     class_label = gr.Textbox(
                         label="class_labels",
                         value=DEFAULT_PROFILE.default_class_label,
-                        info="ImageNet class name (e.g. golden retriever) or id (e.g. 207)",
+                        info="ImageNet class id (e.g. 207) or any synonym from id2label (e.g. golden retriever, tabby)",
                     )
                     with gr.Row():
                         seed = gr.Number(label="seed", value=DEFAULT_PROFILE.default_seed, precision=0)
@@ -421,7 +436,8 @@ def build_demo() -> gr.Blocks:
                     examples=[
                         [DEFAULT_MODEL, "golden retriever", 42],
                         [DEFAULT_MODEL, "207", 0],
-                        [DEFAULT_MODEL, "tabby, tabby cat", 123],
+                        [DEFAULT_MODEL, "tabby", 123],
+                        [DEFAULT_MODEL, "tabby, tabby cat", 456],
                     ],
                     inputs=[model, class_label, seed],
                     label="Quick examples",

model_loader.py

@@ -202,6 +202,110 @@ def build_inference_steps(profile: ModelProfile, steps: Any) -> int | list[int]:
     return steps
 
 
+def _split_synonyms(label_text: str) -> list[str]:
+    return [part.strip() for part in label_text.split(",") if part.strip()]
+
+
+def _get_pipe_id2label(pipe: DiffusionPipeline) -> dict[int, str]:
+    id2label: dict[int, str] | None = None
+    if hasattr(pipe, "id2label"):
+        raw = pipe.id2label
+        if isinstance(raw, dict) and raw:
+            id2label = {int(key): str(value) for key, value in raw.items()}
+    if id2label is None:
+        config = getattr(pipe, "config", None)
+        raw = getattr(config, "id2label", None) if config is not None else None
+        if isinstance(raw, dict) and raw:
+            id2label = {int(key): str(value) for key, value in raw.items()}
+    return id2label or {}
+
+
+def _build_label2id(id2label: dict[int, str]) -> dict[str, int]:
+    label2id: dict[str, int] = {}
+    for class_id, value in id2label.items():
+        synonyms = _split_synonyms(value)
+        if not synonyms:
+            continue
+        for synonym in synonyms:
+            label2id[synonym] = int(class_id)
+        label2id[value.strip()] = int(class_id)
+    return label2id
+
+
+def resolve_class_labels(
+    pipe: DiffusionPipeline,
+    class_label: str,
+    *,
+    default: str,
+) -> int | str:
+    """Resolve a class name or id using the loaded model's id2label synonyms."""
+    label = str(class_label or "").strip() or default
+    if label.isdigit():
+        return int(label)
+    if label.replace(".", "", 1).isdigit():
+        return int(float(label))
+
+    id2label = _get_pipe_id2label(pipe)
+    if not id2label:
+        return label
+
+    label2id = _build_label2id(id2label)
+    if label in label2id:
+        return label2id[label]
+
+    for part in _split_synonyms(label):
+        if part in label2id:
+            return label2id[part]
+
+    normalized = label.casefold()
+    for class_id, value in id2label.items():
+        if value.strip().casefold() == normalized:
+            return int(class_id)
+        for part in _split_synonyms(value):
+            if part.casefold() == normalized:
+                return int(class_id)
+
+    if hasattr(pipe, "get_label_ids"):
+        for candidate in _split_synonyms(label):
+            try:
+                return pipe.get_label_ids(candidate)[0]
+            except (ValueError, TypeError):
+                continue
+
+    return label
+
+
+def primary_label_for_id(pipe: DiffusionPipeline, class_id: int, *, fallback: str) -> str:
+    """Return the first synonym from id2label for a class id."""
+    id2label = _get_pipe_id2label(pipe)
+    value = id2label.get(int(class_id))
+    if not value:
+        return fallback
+    synonyms = _split_synonyms(value)
+    return synonyms[0] if synonyms else fallback
+
+
+def default_class_label_for_pipe(pipe: DiffusionPipeline, profile: ModelProfile) -> str:
+    """Pick a sensible default label using id2label synonyms when available."""
+    id2label = _get_pipe_id2label(pipe)
+    if not id2label:
+        return profile.default_class_label
+
+    preferred_ids = (207, 285, 281)
+    for class_id in preferred_ids:
+        if class_id in id2label:
+            return primary_label_for_id(pipe, class_id, fallback=profile.default_class_label)
+
+    for value in id2label.values():
+        for synonym in _split_synonyms(value):
+            if synonym.casefold() == profile.default_class_label.casefold():
+                return synonym
+
+    first_value = next(iter(id2label.values()), profile.default_class_label)
+    synonyms = _split_synonyms(first_value)
+    return synonyms[0] if synonyms else profile.default_class_label
+
+
 def run_inference(
     profile: ModelProfile,
     pipe: DiffusionPipeline,
@@ -228,13 +332,11 @@ def run_inference(
     }
     call_kwargs.update(extra_kwargs if extra_kwargs is not None else profile.extra_call_kwargs)
 
-    label = str(class_label or "").strip() or profile.default_class_label
-    if label.isdigit():
-        call_kwargs["class_labels"] = int(label)
-    elif label.replace(".", "", 1).isdigit():
-        call_kwargs["class_labels"] = int(float(label))
-    else:
-        call_kwargs["class_labels"] = label
+    call_kwargs["class_labels"] = resolve_class_labels(
+        pipe,
+        class_label,
+        default=profile.default_class_label,
+    )
 
     native = profile.infer_resolution()
     if height > 0 and width > 0: