diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 03a1c39b78..0ededa6614 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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 diff --git a/studio/backend/core/inference/diffusion_controlnet.py b/studio/backend/core/inference/diffusion_controlnet.py index 6b7cdec9fb..b318b040c8 100644 --- a/studio/backend/core/inference/diffusion_controlnet.py +++ b/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( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5f1a5103a4..daf895072b 100644 --- a/studio/backend/routes/inference.py +++ b/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, diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index e91abad053..5708c3a3ae 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/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 diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index e33da62113..ad38261705 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/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; diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 9d7107a927..9471882067 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -224,6 +224,16 @@ const ASPECT_RATIOS: Record = { }; 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 = { + 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(""); const [controlImage, setControlImage] = useState(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("passthrough"); const [controlStrength, setControlStrength] = useState(0.7); const [availableControlNets, setAvailableControlNets] = useState([]); // 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 }) {
Control type - - Passthrough (already a map) - Canny (trace edges) + {controlTypeOptions.map((t) => ( + + {CONTROL_TYPE_LABELS[t] ?? + `${t.charAt(0).toUpperCase()}${t.slice(1)} (map)`} + + ))}