Merge branch 'nightly' into feature/support-for-audio-models
This commit is contained in:
Roland Tannous
GitHub
commit
87f2b2a9db
9 changed files with 238 additions and 67 deletions
|
|
@ -41,6 +41,8 @@ class LlamaCppBackend:
|
|||
self._is_vision: bool = False
|
||||
self._healthy = False
|
||||
self._lock = threading.Lock()
|
||||
self._stdout_lines: list[str] = []
|
||||
self._stdout_thread: Optional[threading.Thread] = None
|
||||
|
||||
atexit.register(self._cleanup)
|
||||
|
||||
|
|
@ -115,6 +117,26 @@ class LlamaCppBackend:
|
|||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
# ── Stdout drain (prevents pipe deadlock on Windows) ─────────
|
||||
|
||||
def _drain_stdout(self):
|
||||
"""
|
||||
Read lines from the subprocess stdout in a background thread.
|
||||
|
||||
This prevents a pipe-buffer deadlock on Windows where the default
|
||||
pipe buffer is only ~4 KB. Without draining, llama-server blocks
|
||||
on writes and never becomes healthy.
|
||||
"""
|
||||
try:
|
||||
for line in self._process.stdout:
|
||||
line = line.rstrip()
|
||||
if line:
|
||||
self._stdout_lines.append(line)
|
||||
logger.info(f"[llama-server] {line}")
|
||||
except (ValueError, OSError):
|
||||
# Pipe closed — process is terminating
|
||||
pass
|
||||
|
||||
# ── Lifecycle ─────────────────────────────────────────────────
|
||||
|
||||
def load_model(
|
||||
|
|
@ -122,6 +144,8 @@ class LlamaCppBackend:
|
|||
*,
|
||||
# Local mode: pass a path to a .gguf file
|
||||
gguf_path: Optional[str] = None,
|
||||
# Vision projection (mmproj) for local vision models
|
||||
mmproj_path: Optional[str] = None,
|
||||
# HF mode: let llama-server download via -hf "repo:quant"
|
||||
hf_repo: Optional[str] = None,
|
||||
hf_variant: Optional[str] = None,
|
||||
|
|
@ -186,6 +210,14 @@ class LlamaCppBackend:
|
|||
if n_threads is not None:
|
||||
cmd.extend(["--threads", str(n_threads)])
|
||||
|
||||
# Append mmproj for local vision models
|
||||
if mmproj_path:
|
||||
if not Path(mmproj_path).is_file():
|
||||
logger.warning(f"mmproj file not found: {mmproj_path}")
|
||||
else:
|
||||
cmd.extend(["--mmproj", mmproj_path])
|
||||
logger.info(f"Using mmproj for vision: {mmproj_path}")
|
||||
|
||||
logger.info(f"Starting llama-server: {' '.join(cmd)}")
|
||||
|
||||
# Set LD_LIBRARY_PATH so llama-server can find its shared libs
|
||||
|
|
@ -196,6 +228,7 @@ class LlamaCppBackend:
|
|||
existing_ld = env.get("LD_LIBRARY_PATH", "")
|
||||
env["LD_LIBRARY_PATH"] = f"{binary_dir}:{existing_ld}" if existing_ld else binary_dir
|
||||
|
||||
self._stdout_lines = []
|
||||
self._process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
|
|
@ -204,6 +237,12 @@ class LlamaCppBackend:
|
|||
env=env,
|
||||
)
|
||||
|
||||
# Start background thread to drain stdout and prevent pipe deadlock
|
||||
self._stdout_thread = threading.Thread(
|
||||
target=self._drain_stdout, daemon=True, name="llama-stdout"
|
||||
)
|
||||
self._stdout_thread.start()
|
||||
|
||||
self._gguf_path = gguf_path
|
||||
self._hf_repo = hf_repo
|
||||
self._hf_variant = hf_variant
|
||||
|
|
@ -256,6 +295,9 @@ class LlamaCppBackend:
|
|||
logger.warning(f"Error killing llama-server process: {e}")
|
||||
finally:
|
||||
self._process = None
|
||||
if self._stdout_thread is not None:
|
||||
self._stdout_thread.join(timeout=2)
|
||||
self._stdout_thread = None
|
||||
|
||||
def _cleanup(self):
|
||||
"""atexit handler to ensure llama-server is terminated."""
|
||||
|
|
@ -273,8 +315,10 @@ class LlamaCppBackend:
|
|||
while time.monotonic() < deadline:
|
||||
# Check if process crashed
|
||||
if self._process.poll() is not None:
|
||||
# Read remaining output for error info
|
||||
output = self._process.stdout.read() if self._process.stdout else ""
|
||||
# Give the drain thread a moment to collect final output
|
||||
if self._stdout_thread is not None:
|
||||
self._stdout_thread.join(timeout=2)
|
||||
output = "\n".join(self._stdout_lines[-50:])
|
||||
logger.error(
|
||||
f"llama-server exited with code {self._process.returncode}. "
|
||||
f"Output: {output[:2000]}"
|
||||
|
|
|
|||
|
|
@ -130,6 +130,7 @@ async def load_model(
|
|||
# Local mode: llama-server loads via -m <path>
|
||||
success = llama_backend.load_model(
|
||||
gguf_path=config.gguf_file,
|
||||
mmproj_path=config.gguf_mmproj_file,
|
||||
model_identifier=config.identifier,
|
||||
is_vision=config.is_vision,
|
||||
n_ctx=request.max_seq_length,
|
||||
|
|
|
|||
|
|
@ -464,6 +464,30 @@ def has_audio_input_model(model_name: str) -> bool:
|
|||
except Exception as e:
|
||||
logger.debug(f"Could not determine if {model_name} has audio input: {e}")
|
||||
return False
|
||||
def _is_mmproj(filename: str) -> bool:
|
||||
"""Check if a GGUF filename is a vision projection (mmproj) file."""
|
||||
return "mmproj" in filename.lower()
|
||||
|
||||
|
||||
def detect_mmproj_file(path: str) -> Optional[str]:
|
||||
"""
|
||||
Find the mmproj (vision projection) GGUF file in a directory.
|
||||
|
||||
Args:
|
||||
path: Directory to search — or a .gguf file (uses its parent dir).
|
||||
|
||||
Returns:
|
||||
Full path to the mmproj .gguf file, or None if not found.
|
||||
"""
|
||||
p = Path(path)
|
||||
search_dir = p.parent if p.is_file() else p
|
||||
if not search_dir.is_dir():
|
||||
return None
|
||||
|
||||
for f in search_dir.glob("*.gguf"):
|
||||
if _is_mmproj(f.name):
|
||||
return str(f.resolve())
|
||||
return None
|
||||
|
||||
|
||||
def detect_gguf_model(path: str) -> Optional[str]:
|
||||
|
|
@ -474,6 +498,9 @@ def detect_gguf_model(path: str) -> Optional[str]:
|
|||
1. path is a direct .gguf file path
|
||||
2. path is a directory containing .gguf files
|
||||
|
||||
Skips mmproj (vision projection) files — those must be passed via
|
||||
``--mmproj``, not ``-m``. Use :func:`detect_mmproj_file` instead.
|
||||
|
||||
Returns the full path to the .gguf file if found, None otherwise.
|
||||
For HuggingFace repo detection, use detect_gguf_model_remote() instead.
|
||||
"""
|
||||
|
|
@ -481,11 +508,16 @@ def detect_gguf_model(path: str) -> Optional[str]:
|
|||
|
||||
# Case 1: direct .gguf file
|
||||
if p.suffix == ".gguf" and p.is_file():
|
||||
if _is_mmproj(p.name):
|
||||
return None
|
||||
return str(p.resolve())
|
||||
|
||||
# Case 2: directory containing .gguf files
|
||||
# Case 2: directory containing .gguf files (skip mmproj)
|
||||
if p.is_dir():
|
||||
gguf_files = sorted(p.glob("*.gguf"), key=lambda f: f.stat().st_size, reverse=True)
|
||||
gguf_files = sorted(
|
||||
(f for f in p.glob("*.gguf") if not _is_mmproj(f.name)),
|
||||
key=lambda f: f.stat().st_size, reverse=True,
|
||||
)
|
||||
if gguf_files:
|
||||
return str(gguf_files[0].resolve())
|
||||
|
||||
|
|
@ -721,7 +753,8 @@ def scan_exported_models(exports_dir: str = "./exports") -> List[Tuple[str, str,
|
|||
continue
|
||||
|
||||
# Check for flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/)
|
||||
gguf_files = list(run_dir.glob("*.gguf"))
|
||||
# Filter out mmproj (vision projection) files — they aren't loadable as main models
|
||||
gguf_files = [f for f in run_dir.glob("*.gguf") if not _is_mmproj(f.name)]
|
||||
if gguf_files:
|
||||
base_model = None
|
||||
export_meta = run_dir / "export_metadata.json"
|
||||
|
|
@ -971,6 +1004,7 @@ class ModelConfig:
|
|||
audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
|
||||
has_audio_input: bool = False # Accepts audio input (ASR/speech understanding)
|
||||
gguf_file: Optional[str] = None # Full path to the .gguf file (local mode)
|
||||
gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection)
|
||||
gguf_hf_repo: Optional[str] = None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
|
||||
gguf_variant: Optional[str] = None # Quantization variant (e.g. "Q4_K_M")
|
||||
base_model: Optional[str] = None # Base model (for LoRAs)
|
||||
|
|
@ -1074,16 +1108,44 @@ class ModelConfig:
|
|||
if gguf_file:
|
||||
display_name = Path(gguf_file).stem
|
||||
logger.info(f"Detected local GGUF model: {gguf_file}")
|
||||
|
||||
# Detect vision: check if base model is vision, then look for mmproj
|
||||
mmproj_file = None
|
||||
gguf_is_vision = False
|
||||
gguf_dir = Path(gguf_file).parent
|
||||
|
||||
# Determine if this is a vision model from export metadata
|
||||
base_is_vision = False
|
||||
meta_path = gguf_dir / "export_metadata.json"
|
||||
if meta_path.exists():
|
||||
try:
|
||||
meta = json.loads(meta_path.read_text())
|
||||
base = meta.get("base_model")
|
||||
if base and is_vision_model(base, hf_token=hf_token):
|
||||
base_is_vision = True
|
||||
logger.info(f"GGUF base model '{base}' is a vision model")
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not read export metadata: {e}")
|
||||
|
||||
# If vision (or mmproj happens to exist), find the mmproj file
|
||||
mmproj_file = detect_mmproj_file(gguf_file)
|
||||
if mmproj_file:
|
||||
gguf_is_vision = True
|
||||
logger.info(f"Detected mmproj for vision: {mmproj_file}")
|
||||
elif base_is_vision:
|
||||
logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}")
|
||||
|
||||
return cls(
|
||||
identifier=identifier,
|
||||
display_name=display_name,
|
||||
path=path,
|
||||
is_local=True,
|
||||
is_cached=True,
|
||||
is_vision=False,
|
||||
is_vision=gguf_is_vision,
|
||||
is_lora=False,
|
||||
is_gguf=True,
|
||||
gguf_file=gguf_file,
|
||||
gguf_mmproj_file=mmproj_file,
|
||||
)
|
||||
else:
|
||||
# Check if the HF repo contains GGUF files
|
||||
|
|
|
|||
|
|
@ -103,9 +103,6 @@ function ModelRow({
|
|||
{vramStatus === "tight" && (
|
||||
<span className="text-[9px] font-medium text-amber-400">TIGHT</span>
|
||||
)}
|
||||
{vramStatus === "fits" && (
|
||||
<span className="text-[9px] font-medium text-emerald-500/90">FIT</span>
|
||||
)}
|
||||
{meta ? (
|
||||
<span className="text-[10px] text-muted-foreground">{meta}</span>
|
||||
) : null}
|
||||
|
|
@ -253,9 +250,6 @@ function GgufVariantExpander({
|
|||
{fitStatus === "tight" && (
|
||||
<span className="text-[9px] font-medium text-amber-400">TIGHT</span>
|
||||
)}
|
||||
{fitStatus === "fits" && (
|
||||
<span className="text-[9px] font-medium text-emerald-500/90">FIT</span>
|
||||
)}
|
||||
<span className="text-[10px] text-muted-foreground">
|
||||
{formatBytes(v.size_bytes)}
|
||||
</span>
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ import {
|
|||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { useTrainingRuntimeStore } from "@/features/training";
|
||||
import { Link, useRouterState } from "@tanstack/react-router";
|
||||
import { AnimatePresence, motion } from "motion/react";
|
||||
import { motion } from "motion/react";
|
||||
import { useState } from "react";
|
||||
import { TOUR_OPEN_EVENT } from "@/features/tour";
|
||||
|
||||
|
|
@ -58,9 +58,9 @@ export function Navbar() {
|
|||
|
||||
return (
|
||||
<header className="relative top-0 z-40 h-16 w-full">
|
||||
<div className="mx-auto flex h-full max-w-7xl items-center justify-between px-4 sm:px-6">
|
||||
<div className="mx-auto grid h-full max-w-7xl grid-cols-[1fr_auto_1fr] items-center px-4 sm:px-6">
|
||||
{/* Left: logo */}
|
||||
<Link to="/studio" className="flex items-center select-none">
|
||||
<Link to="/studio" className="flex items-center justify-self-start select-none">
|
||||
<img
|
||||
src="/blacklogo.png"
|
||||
alt="Unsloth"
|
||||
|
|
@ -117,23 +117,21 @@ export function Navbar() {
|
|||
/>
|
||||
)}
|
||||
<span className="relative z-10 flex items-center gap-1.5">
|
||||
<AnimatePresence mode="popLayout">
|
||||
{active && item.icon && (
|
||||
<motion.span
|
||||
key={item.href}
|
||||
initial={{ width: 0, opacity: 0 }}
|
||||
animate={{ width: "auto", opacity: 1 }}
|
||||
exit={{ width: 0, opacity: 0 }}
|
||||
transition={{ duration: 0.2, ease: [0.165, 0.84, 0.44, 1] }}
|
||||
className="overflow-hidden"
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={item.icon}
|
||||
className="size-3.5 -mt-px"
|
||||
/>
|
||||
</motion.span>
|
||||
)}
|
||||
</AnimatePresence>
|
||||
<span className="inline-flex size-3.5 items-center justify-center overflow-hidden">
|
||||
<motion.span
|
||||
initial={false}
|
||||
animate={{
|
||||
opacity: active ? 1 : 0,
|
||||
scale: active ? 1 : 0.9,
|
||||
}}
|
||||
transition={{ duration: 0.2, ease: [0.165, 0.84, 0.44, 1] }}
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={item.icon}
|
||||
className="size-3.5 -mt-px"
|
||||
/>
|
||||
</motion.span>
|
||||
</span>
|
||||
{item.label}
|
||||
</span>
|
||||
</Link>
|
||||
|
|
@ -142,7 +140,7 @@ export function Navbar() {
|
|||
</nav>
|
||||
|
||||
{/* Right: docs/tour desktop */}
|
||||
<div className="hidden items-center gap-2 md:flex">
|
||||
<div className="hidden items-center justify-self-end gap-2 md:flex">
|
||||
<AnimatedThemeToggler
|
||||
className="flex h-9 w-9 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-accent hover:text-foreground [&_svg]:size-4"
|
||||
title="Toggle theme"
|
||||
|
|
@ -182,17 +180,20 @@ export function Navbar() {
|
|||
</HoverCardContent>
|
||||
</HoverCard>
|
||||
|
||||
{tourId ? (
|
||||
<button
|
||||
type="button"
|
||||
onClick={openTour}
|
||||
className="flex h-9 items-center gap-1.5 rounded-md px-3 text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
|
||||
title="Tour"
|
||||
>
|
||||
<HugeiconsIcon icon={CursorInfo02Icon} className="size-4" />
|
||||
<span className="text-sm font-medium">Tour</span>
|
||||
</button>
|
||||
) : null}
|
||||
<button
|
||||
type="button"
|
||||
onClick={tourId ? openTour : undefined}
|
||||
className={cn(
|
||||
"flex h-9 items-center gap-1.5 rounded-md px-3 text-muted-foreground transition-colors hover:bg-accent hover:text-foreground",
|
||||
!tourId && "invisible pointer-events-none",
|
||||
)}
|
||||
title="Tour"
|
||||
aria-hidden={!tourId}
|
||||
tabIndex={tourId ? 0 : -1}
|
||||
>
|
||||
<HugeiconsIcon icon={CursorInfo02Icon} className="size-4" />
|
||||
<span className="text-sm font-medium">Tour</span>
|
||||
</button>
|
||||
</div>
|
||||
|
||||
{/* Right: mobile */}
|
||||
|
|
|
|||
|
|
@ -33,11 +33,17 @@ import {
|
|||
import { MODEL_TYPE_TO_HF_TASK } from "@/config/training";
|
||||
import {
|
||||
useDebouncedValue,
|
||||
useGpuInfo,
|
||||
useHfModelSearch,
|
||||
useHfTokenValidation,
|
||||
useInfiniteScroll,
|
||||
} from "@/hooks";
|
||||
import { formatCompact } from "@/lib/utils";
|
||||
import {
|
||||
type TrainingMethod as VramTrainingMethod,
|
||||
type VramFitStatus,
|
||||
buildModelVramMap,
|
||||
} from "@/lib/vram";
|
||||
import { useTrainingConfigStore } from "@/features/training";
|
||||
import type { TrainingMethod } from "@/types/training";
|
||||
import {
|
||||
|
|
@ -50,6 +56,7 @@ import { useEffect, useMemo, useRef, useState } from "react";
|
|||
import { useShallow } from "zustand/react/shallow";
|
||||
|
||||
export function ModelSelectionStep() {
|
||||
const gpu = useGpuInfo();
|
||||
const {
|
||||
modelType,
|
||||
selectedModel,
|
||||
|
|
@ -93,6 +100,24 @@ export function ModelSelectionStep() {
|
|||
|
||||
const resultIds = useMemo(() => hfResults.map((r) => r.id), [hfResults]);
|
||||
|
||||
// Match Studio behavior: only show exception signals (OOM/TIGHT) in training flows.
|
||||
const vramMap = useMemo(() => {
|
||||
const fitMap = buildModelVramMap(
|
||||
hfResults,
|
||||
trainingMethod as VramTrainingMethod,
|
||||
gpu,
|
||||
);
|
||||
const map = new Map<string, { status: VramFitStatus | null; detail: string | null }>();
|
||||
for (const r of hfResults) {
|
||||
const fit = fitMap.get(r.id);
|
||||
map.set(r.id, {
|
||||
status: fit?.status ?? null,
|
||||
detail: r.totalParams ? formatCompact(r.totalParams) : null,
|
||||
});
|
||||
}
|
||||
return map;
|
||||
}, [hfResults, gpu, trainingMethod]);
|
||||
|
||||
const comboboxAnchorRef = useRef<HTMLDivElement>(null);
|
||||
const { scrollRef, sentinelRef } = useInfiniteScroll(
|
||||
fetchMore,
|
||||
|
|
@ -218,19 +243,21 @@ export function ModelSelectionStep() {
|
|||
>
|
||||
<ComboboxList className="p-1 !max-h-none !overflow-visible">
|
||||
{(id: string) => {
|
||||
const r = hfResults.find((r) => r.id === id);
|
||||
const sizeLabel = r?.totalParams
|
||||
? formatCompact(r.totalParams)
|
||||
: null;
|
||||
const entry = vramMap.get(id);
|
||||
const sizeLabel = entry?.detail ?? null;
|
||||
const fitStatus = entry?.status ?? null;
|
||||
const exceeds = fitStatus === "exceeds";
|
||||
return (
|
||||
<ComboboxItem
|
||||
key={id}
|
||||
value={id}
|
||||
className="justify-between"
|
||||
className={`justify-between ${exceeds ? "opacity-50" : ""}`}
|
||||
>
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<span className="min-w-0 flex-1 truncate">
|
||||
<span
|
||||
className={`min-w-0 flex-1 truncate ${exceeds ? "line-through decoration-muted-foreground/50" : ""}`}
|
||||
>
|
||||
{id}
|
||||
</span>
|
||||
</TooltipTrigger>
|
||||
|
|
@ -241,11 +268,23 @@ export function ModelSelectionStep() {
|
|||
{id}
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
{sizeLabel ? (
|
||||
<span className="text-xs text-muted-foreground shrink-0">
|
||||
{sizeLabel}
|
||||
</span>
|
||||
) : null}
|
||||
<span className="flex items-center gap-1.5 shrink-0">
|
||||
{fitStatus === "exceeds" && (
|
||||
<span className="text-[9px] font-medium text-red-400">
|
||||
OOM
|
||||
</span>
|
||||
)}
|
||||
{fitStatus === "tight" && (
|
||||
<span className="text-[9px] font-medium text-amber-400">
|
||||
TIGHT
|
||||
</span>
|
||||
)}
|
||||
{sizeLabel ? (
|
||||
<span className="text-xs text-muted-foreground">
|
||||
{sizeLabel}
|
||||
</span>
|
||||
) : null}
|
||||
</span>
|
||||
</ComboboxItem>
|
||||
);
|
||||
}}
|
||||
|
|
|
|||
|
|
@ -37,8 +37,7 @@ import { formatCompact } from "@/lib/utils";
|
|||
import {
|
||||
type TrainingMethod as VramTrainingMethod,
|
||||
type VramFitStatus,
|
||||
checkVramFit,
|
||||
estimateLoadingVram,
|
||||
buildModelVramMap,
|
||||
} from "@/lib/vram";
|
||||
import {
|
||||
listLocalModels,
|
||||
|
|
@ -218,22 +217,23 @@ export function ModelSection() {
|
|||
// Keyed by model id so the render callback is a simple O(1) lookup.
|
||||
// Re-computes when the training method changes (QLoRA=4-bit vs LoRA/Full=fp16).
|
||||
const vramMap = useMemo(() => {
|
||||
const method = trainingMethod as VramTrainingMethod;
|
||||
const fitMap = buildModelVramMap(
|
||||
hfResults,
|
||||
trainingMethod as VramTrainingMethod,
|
||||
gpu,
|
||||
);
|
||||
const map = new Map<
|
||||
string,
|
||||
{ est: number; status: VramFitStatus | null; detail: string | null }
|
||||
>();
|
||||
for (const r of hfResults) {
|
||||
const detail = r.totalParams ? formatCompact(r.totalParams) : null;
|
||||
if (r.totalParams) {
|
||||
const est = estimateLoadingVram(r.totalParams, method);
|
||||
const status = gpu.available
|
||||
? checkVramFit(est, gpu.memoryTotalGb)
|
||||
: null;
|
||||
map.set(r.id, { est, status, detail });
|
||||
} else {
|
||||
map.set(r.id, { est: 0, status: null, detail });
|
||||
}
|
||||
const fit = fitMap.get(r.id);
|
||||
map.set(r.id, {
|
||||
est: fit?.est ?? 0,
|
||||
status: fit?.status ?? null,
|
||||
detail,
|
||||
});
|
||||
}
|
||||
return map;
|
||||
}, [hfResults, gpu, trainingMethod]);
|
||||
|
|
|
|||
|
|
@ -268,6 +268,7 @@
|
|||
}
|
||||
html {
|
||||
@apply font-sans;
|
||||
scrollbar-gutter: stable;
|
||||
}
|
||||
h1,
|
||||
h2,
|
||||
|
|
|
|||
|
|
@ -93,3 +93,32 @@ export function checkVramFit(
|
|||
if (ratio <= 1.0) return "tight";
|
||||
return "exceeds";
|
||||
}
|
||||
|
||||
export interface ModelVramMapInput {
|
||||
id: string;
|
||||
totalParams?: number;
|
||||
}
|
||||
|
||||
export interface ModelVramMapEntry {
|
||||
est: number;
|
||||
status: VramFitStatus | null;
|
||||
}
|
||||
|
||||
export function buildModelVramMap(
|
||||
models: ModelVramMapInput[],
|
||||
method: TrainingMethod,
|
||||
gpu: { available: boolean; memoryTotalGb: number },
|
||||
): Map<string, ModelVramMapEntry> {
|
||||
const map = new Map<string, ModelVramMapEntry>();
|
||||
for (const model of models) {
|
||||
if (!model.totalParams) {
|
||||
map.set(model.id, { est: 0, status: null });
|
||||
continue;
|
||||
}
|
||||
|
||||
const est = estimateLoadingVram(model.totalParams, method);
|
||||
const status = gpu.available ? checkVramFit(est, gpu.memoryTotalGb) : null;
|
||||
map.set(model.id, { est, status });
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue