From 8218c9bf42d1759e4c0ed5c8e1e579971f1dab66 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 10 Jul 2026 04:32:31 +0000
Subject: [PATCH] studio: use largest single GPU for the diffusion catalog fit
 budget

The catalog fit budget used gpu.memoryTotalGb, which sums VRAM across
every GPU. That sum is right for the chat/llama.cpp path (tensor-split
shards across cards) but wrong for the diffusion/video catalog: those
backends place the whole pipeline on a single device (pipe.to or cpu
offload, never device_map), so on a multi-GPU host the fit toggle and
bare-group-click routing credited VRAM no single card has. On a 4x24 GB
plus 128 GB RAM host the 114 GB Wan A14B bf16 group passed the toggle
(0.7*96 + 0.7*128 budget) and a click would OOM, the exact load the
toggle exists to prevent. Expose maxDeviceMemoryGb (largest single
device) from use-gpu-info and use it for deviceBudget; the chat path
keeps the sum. Single-GPU hosts are unchanged.
---
 .../components/assistant-ui/model-selector/pickers.tsx |  7 ++++++-
 studio/frontend/src/hooks/use-gpu-info.ts              | 10 ++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index 767e586a07..8ae2111d1e 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -2074,8 +2074,13 @@ export function HubModelPicker({
     [downloadedSet],
   );
   const deviceBudget = useMemo(
+    // Largest single device, NOT the multi-GPU sum: the diffusion/video
+    // backends place the whole pipeline on one device (pipe.to / cpu-offload,
+    // never device_map), so summed VRAM would pass groups no single card can
+    // hold (e.g. a 114 GB group "fits" a 4x24 GB host) and a bare group click
+    // would OOM -- the exact load the fit toggle exists to prevent.
     () => ({
-      gpuGb: gpu.available ? gpu.memoryTotalGb : 0,
+      gpuGb: gpu.available ? gpu.maxDeviceMemoryGb : 0,
       systemRamGb: gpu.systemRamAvailableGb || 0,
     }),
     [gpu],
diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts
index 1e313acdf3..9f4cc48eb9 100644
--- a/studio/frontend/src/hooks/use-gpu-info.ts
+++ b/studio/frontend/src/hooks/use-gpu-info.ts
@@ -8,7 +8,13 @@ import type { SystemInfoResponse } from "./use-system";
 export interface GpuInfo {
   available: boolean;
   name: string;
+  /** Sum across every GPU. Right for the chat/llama.cpp path (tensor-split
+   * shards across cards); wrong as a single-placement budget. */
   memoryTotalGb: number;
+  /** Largest single device. The diffusion/video backends place the whole
+   * pipeline on one device (pipe.to / cpu-offload, never device_map), so
+   * their fit budget must not credit VRAM from the other cards. */
+  maxDeviceMemoryGb: number;
   cpuCore: number;
   cpuThread: number;
   systemRamAvailableGb: number;
@@ -19,6 +25,7 @@ const DEFAULT_GPU: GpuInfo = {
   available: false,
   name: "Unknown",
   memoryTotalGb: 0,
+  maxDeviceMemoryGb: 0,
   cpuCore: 0,
   cpuThread: 0,
   systemRamAvailableGb: 0,
@@ -59,6 +66,9 @@ async function fetchGpuOnce(): Promise<GpuInfo> {
               available: true,
               name: devices[0]?.name ?? "Unknown",
               memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0),
+              maxDeviceMemoryGb: devices.reduce(
+                (max, d) => Math.max(max, d.memory_total_gb ?? 0), 0,
+              ),
             }
           : { ...DEFAULT_GPU, ...base };
       cachedGpu = info;