Harden ControlNet resolve, gallery metadata, and the control-type picker

Check cancellation immediately after a ControlNet from_pretrained and before
any device placement, so an unload/eviction that raced the download does not
allocate several GB onto the GPU after the load was already cleared.

Require a loadable weight or shard index (not just config.json) before a local
ControlNet folder is advertised, so an interrupted copy is hidden instead of
failing deep in from_pretrained as a generic 500.

Do not record a strength-0 ControlNet in the gallery recipe: it is treated as
disabled and skipped, so the image is unconditioned and the metadata must not
claim a ControlNet was applied.

Build the control-type picker from the selected ControlNet's advertised
control_types instead of a hardcoded passthrough/canny pair, so a union model
with a precomputed depth or pose map sends the correct control_mode.

studio/backend/core/inference/diffusion.py

@@ -1360,6 +1360,13 @@ class DiffusionBackend:
                 # raise on a blank credential instead of falling back, so coerce to None.
                 token = state.hf_token or None,
             )
+            if cancel.is_set():
+                # An unload/eviction raced the blocking download above and may have already
+                # cleared the load. Bail BEFORE any device placement so we don't allocate
+                # several GB onto the GPU after _unload_locked() freed it (which would OOM
+                # or make the unload appear to free memory only to repopulate it).
+                del cn_model
+                raise RuntimeError(DIFFUSION_CANCELLED_MSG)
             # Placement must follow the base model's offload policy. A resident base moves
             # the ControlNet resident too; an offloaded (low-VRAM) base streams it through
             # the device with group offloading instead of forcing the whole module onto the

studio/backend/core/inference/diffusion_controlnet.py

@@ -97,6 +97,26 @@ def sanitize_id(raw: str) -> str:
     return stem or "controlnet"
 
 
+def _has_controlnet_weights(p: Path) -> bool:
+    """True when ``p`` holds a loadable diffusers ControlNet weight (or shard index).
+
+    A config-only folder (interrupted copy/download) would otherwise be advertised and
+    then fail deep inside ``from_pretrained`` as a generic 500. Accept the standard
+    single-file weights, a sharded weight index, or any ``.safetensors`` shard."""
+    names = (
+        "diffusion_pytorch_model.safetensors",
+        "diffusion_pytorch_model.bin",
+        "diffusion_pytorch_model.safetensors.index.json",
+        "diffusion_pytorch_model.bin.index.json",
+    )
+    if any((p / n).exists() for n in names):
+        return True
+    try:
+        return any(child.suffix == ".safetensors" for child in p.iterdir())
+    except OSError:
+        return False
+
+
 def _scan_local() -> list[ControlNetCatalogEntry]:
     """A local ControlNet is a directory containing a diffusers config + weights."""
     entries: list[ControlNetCatalogEntry] = []
@@ -108,7 +128,9 @@ def _scan_local() -> list[ControlNetCatalogEntry]:
     for p in children:
         if not p.is_dir():
             continue
-        if not (p / "config.json").exists():
+        # Require BOTH the config and a loadable weight/index: a config-only folder is an
+        # incomplete copy/download, and advertising it would fail later in from_pretrained.
+        if not (p / "config.json").exists() or not _has_controlnet_weights(p):
             continue
         entries.append(
             ControlNetCatalogEntry(

studio/backend/routes/inference.py

@@ -11230,7 +11230,10 @@ async def generate_diffusion_image(
                         "controlnet": (
                             f"{request.controlnet.id}:{request.controlnet.control_type}:"
                             f"{request.controlnet.strength:g}"
-                            if request.controlnet
+                            # strength 0 is treated as disabled and skipped before loading /
+                            # conditioning, so the image is unconditioned; don't claim a
+                            # ControlNet was applied in the recipe/metadata.
+                            if request.controlnet and request.controlnet.strength > 0
                             else None
                         ),
                         "created_at": created_at,

studio/backend/tests/test_diffusion_controlnet.py

@@ -71,6 +71,7 @@ def test_resolve_controlnet_local(tmp_path, monkeypatch):
     cn = d / "my-cn"
     cn.mkdir()
     (cn / "config.json").write_text("{}")
+    (cn / "diffusion_pytorch_model.safetensors").write_bytes(b"x")  # a loadable weight
     monkeypatch.setattr(dc, "controlnets_dir", lambda: d)
     entries = {e.id for e in dc.list_controlnets()}
     assert "my-cn" in entries
@@ -78,6 +79,21 @@ def test_resolve_controlnet_local(tmp_path, monkeypatch):
     assert r.is_local and r.path == str(cn)
 
 
+def test_scan_local_skips_config_only_folder(tmp_path, monkeypatch):
+    # A folder with config.json but no weight/index (interrupted copy) must NOT be
+    # advertised: it would otherwise fail deep in from_pretrained as a generic 500.
+    d = tmp_path / "controlnets"
+    d.mkdir()
+    incomplete = d / "incomplete-cn"
+    incomplete.mkdir()
+    (incomplete / "config.json").write_text("{}")
+    monkeypatch.setattr(dc, "controlnets_dir", lambda: d)
+    assert "incomplete-cn" not in {e.id for e in dc.list_controlnets()}
+    # A sharded weight index counts as a loadable weight.
+    (incomplete / "diffusion_pytorch_model.safetensors.index.json").write_text("{}")
+    assert "incomplete-cn" in {e.id for e in dc.list_controlnets()}
+
+
 def test_preprocess_control_passthrough_and_canny():
     from PIL import Image
 

studio/frontend/src/features/images/api.ts

@@ -111,7 +111,9 @@ export interface ControlNetSpecInput {
   id: string;
   // Base64/data-URL control image (a source image or an already-made control map).
   image: string;
-  control_type: "passthrough" | "canny";
+  // "canny" preprocesses edges from a source image; any other type (passthrough, or a
+  // union type like depth/pose) is an already-made map the backend maps to a control mode.
+  control_type: string;
   strength: number;
   guidance_start?: number;
   guidance_end?: number;

studio/frontend/src/features/images/images-page.tsx

@@ -224,6 +224,16 @@ const ASPECT_RATIOS: Record<string, [number, number]> = {
 };
 const ASPECT_OPTIONS = ["custom", ...Object.keys(ASPECT_RATIOS)];
 
+// Friendly labels for ControlNet control types. "canny" traces edges from a source image;
+// every other type is an already-made map (passthrough/depth/pose/...). Unknown types fall
+// back to a capitalized "(map)" label so a new backend type still renders.
+const CONTROL_TYPE_LABELS: Record<string, string> = {
+  passthrough: "Passthrough (already a map)",
+  canny: "Canny (trace edges)",
+  depth: "Depth (map)",
+  pose: "Pose (map)",
+};
+
 // Z-Image accepts 256–2048, in multiples of 16. Snap any value into range.
 const MIN_DIM = 256;
 const MAX_DIM = 2048;
@@ -900,7 +910,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
   // per loaded family; applied at generate time only when a model + control image are set.
   const [controlnetId, setControlnetId] = useState<string>("");
   const [controlImage, setControlImage] = useState<string | null>(null);
-  const [controlType, setControlType] = useState<"passthrough" | "canny">("passthrough");
+  // Free-form: a union ControlNet advertises depth/pose/etc alongside the preprocessing
+  // "canny", and the backend maps the exact control_type to the union control_mode. The
+  // picker is built from the selected model's control_types, so it isn't limited to two.
+  const [controlType, setControlType] = useState<string>("passthrough");
   const [controlStrength, setControlStrength] = useState(0.7);
   const [availableControlNets, setAvailableControlNets] = useState<DiffusionControlNetInfo[]>([]);
   // Advanced options live in a right-docked panel (like Chat's settings panel). Closed by
@@ -1027,6 +1040,25 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
     };
   }, [controlnetCapable, status?.family]);
 
+  // The control types offered for the selected ControlNet. A union model advertises
+  // several (canny/depth/pose/passthrough); a plain model advertises its own. Fall back
+  // to the preprocessing pair when nothing is selected.
+  const controlTypeOptions = useMemo(() => {
+    const cn = availableControlNets.find((c) => c.id === controlnetId);
+    const types = cn?.control_types?.length ? cn.control_types : ["passthrough", "canny"];
+    return types;
+  }, [availableControlNets, controlnetId]);
+
+  // Keep controlType valid for the selected model: if the current choice isn't among the
+  // model's advertised types, snap to the first (prefer passthrough when offered).
+  useEffect(() => {
+    if (!controlTypeOptions.includes(controlType)) {
+      setControlType(
+        controlTypeOptions.includes("passthrough") ? "passthrough" : controlTypeOptions[0],
+      );
+    }
+  }, [controlTypeOptions, controlType]);
+
   const selected = useMemo(
     () => images.find((i) => i.id === selectedId) ?? images[0] ?? null,
     [images, selectedId],
@@ -2163,16 +2195,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
                     <ImageDropzone value={controlImage} onChange={setControlImage} />
                     <div className="flex items-center gap-2">
                       <span className="shrink-0 text-xs text-muted-foreground">Control type</span>
-                      <Select
-                        value={controlType}
-                        onValueChange={(v) => setControlType(v as "passthrough" | "canny")}
-                      >
+                      <Select value={controlType} onValueChange={setControlType}>
                         <SelectTrigger className="h-8 flex-1 text-xs">
                           <SelectValue />
                         </SelectTrigger>
                         <SelectContent>
-                          <SelectItem value="passthrough">Passthrough (already a map)</SelectItem>
-                          <SelectItem value="canny">Canny (trace edges)</SelectItem>
+                          {controlTypeOptions.map((t) => (
+                            <SelectItem key={t} value={t}>
+                              {CONTROL_TYPE_LABELS[t] ??
+                                `${t.charAt(0).toUpperCase()}${t.slice(1)} (map)`}
+                            </SelectItem>
+                          ))}
                         </SelectContent>
                       </Select>
                     </div>