load selected variant inside gpu calls, fix get_cmap removal

app.py

@@ -4,7 +4,6 @@ import colorsys
 import os
 
 import gradio as gr
-import matplotlib.cm as cm
 import matplotlib.pyplot as plt
 import numpy as np
 import spaces
@@ -347,9 +346,14 @@ def _ensure_ade20k_embs():
     print("Pascal Context text embeddings computed.")
 
 
-def _init_model():
-    """Load model + move to GPU + compute text embeddings."""
-    load_variant(_model["name"] or DEFAULT_VARIANT)
+def _init_model(name=None):
+    """Load model + move to GPU + compute text embeddings.
+
+    Must be called with the requested variant inside each @spaces.GPU function:
+    on ZeroGPU those run in a forked worker, so global state mutated in one call
+    (e.g. by a dropdown handler) does not survive into the next.
+    """
+    load_variant(name or _model["name"] or DEFAULT_VARIANT)
     _move_models_to_device()
     _ensure_ade20k_embs()
 
@@ -464,7 +468,7 @@ def vis_depth(spatial):
     h, w = spatial.shape[0], spatial.shape[1]
     depth = PCA(n_components=1).fit_transform(feat).reshape(h, w)
     depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-8)
-    colored = cm.get_cmap("inferno")(depth)[:, :, :3].astype(np.float32)
+    colored = plt.get_cmap("inferno")(depth)[:, :, :3].astype(np.float32)
     return to_uint8(colored)
 
 
@@ -573,7 +577,7 @@ def vis_depth_dpt(depth_map, h, w):
     """Colour a depth map with the turbo colormap → PIL Image."""
     d = depth_map.squeeze()
     d = (d - d.min()) / (d.max() - d.min() + 1e-8)
-    colored = cm.get_cmap("turbo")(d)[:, :, :3].astype(np.float32)
+    colored = plt.get_cmap("turbo")(d)[:, :, :3].astype(np.float32)
     return to_uint8(upsample(colored, h, w))
 
 
@@ -643,11 +647,10 @@ def vis_segmentation_dpt(seg_map, orig_image):
 # ── Gradio callbacks ────────────────────────────────────────────────────────
 
 
-@spaces.GPU
-def on_variant_change(variant_name):
-    load_variant(variant_name)
-    _move_models_to_device()
-    _ensure_ade20k_embs()
+def on_variant_change(variant_name):  # noqa: ARG001
+    # Only clears stale outputs. The variant is loaded inside the inference
+    # handlers themselves — loading it here would happen in a ZeroGPU worker
+    # whose state is thrown away when the call returns.
     return (
         None,
         None,
@@ -661,10 +664,10 @@ def on_variant_change(variant_name):
 
 
 @spaces.GPU
-def on_pca_extract(image, resolution, _pca_state):
+def on_pca_extract(image, variant, resolution, _pca_state):
     if image is None:
         return None, None, None, None
-    _init_model()
+    _init_model(variant)
     resolution = int(resolution)
     spatial = extract_features(image, resolution)
     h, w = image.shape[:2]
@@ -681,11 +684,11 @@ def on_pca_extract(image, resolution, _pca_state):
 
 
 @spaces.GPU
-def on_recluster(image, resolution, n_clusters, pca_state):
+def on_recluster(image, variant, resolution, n_clusters, pca_state):
     if image is None:
         gr.Warning("Upload an image first.")
         return None, pca_state
-    _init_model()
+    _init_model(variant)
     resolution = int(resolution)
     if (
         pca_state is not None
@@ -706,11 +709,11 @@ def on_recluster(image, resolution, n_clusters, pca_state):
 
 
 @spaces.GPU
-def on_zeroseg_custom(image, resolution, class_names_str):
+def on_zeroseg_custom(image, variant, resolution, class_names_str):
     if image is None or not class_names_str or not class_names_str.strip():
         gr.Warning("Upload an image and enter at least one class name.")
         return None, None, "", ""
-    _init_model()
+    _init_model(variant)
     resolution = int(resolution)
     classes = [c.strip() for c in class_names_str.split(",") if c.strip()]
     if not classes:
@@ -740,11 +743,11 @@ def on_zeroseg_custom(image, resolution, class_names_str):
 
 
 @spaces.GPU
-def on_depth_normals_predict(image, dpt_variant, resolution):  # noqa: ARG001
+def on_depth_normals_predict(image, dpt_variant, resolution):
     """Run DPT depth and normals prediction."""
     if image is None:
         return None, None
-    _init_model()
+    _init_model(dpt_variant)
     dev = _device()
 
     h, w = image.shape[:2]
@@ -761,11 +764,11 @@ def on_depth_normals_predict(image, dpt_variant, resolution):  # noqa: ARG001
 
 
 @spaces.GPU
-def on_segmentation_predict(image, dpt_variant, resolution):  # noqa: ARG001
+def on_segmentation_predict(image, dpt_variant, resolution):
     """Run DPT segmentation prediction."""
     if image is None:
         return None
-    _init_model()
+    _init_model(dpt_variant)
     dev = _device()
 
     img = Image.fromarray(image).convert("RGB")
@@ -998,12 +1001,12 @@ with gr.Blocks(head=head, title="TIPSv2 Feature Explorer", css=custom_css) as de
 
     pca_btn.click(
         fn=on_pca_extract,
-        inputs=[pca_input, resolution_dd, pca_state],
+        inputs=[pca_input, variant_dd, resolution_dd, pca_state],
         outputs=[pca_out, depth_out, kmeans_out, pca_state],
     )
     recluster_btn.click(
         fn=on_recluster,
-        inputs=[pca_input, resolution_dd, n_clusters, pca_state],
+        inputs=[pca_input, variant_dd, resolution_dd, n_clusters, pca_state],
         outputs=[kmeans_out, pca_state],
     )
 
@@ -1021,7 +1024,7 @@ with gr.Blocks(head=head, title="TIPSv2 Feature Explorer", css=custom_css) as de
 
     custom_btn.click(
         fn=on_zeroseg_custom,
-        inputs=[custom_input, resolution_dd, custom_classes],
+        inputs=[custom_input, variant_dd, resolution_dd, custom_classes],
         outputs=[custom_overlay, custom_mask, custom_detected, custom_undetected],
     )