Address Codex review: FP16 fallback on pre-Ampere, forward HF token, hide edit models, download outside the load lock
This commit is contained in:
oobabooga
parent
0a2a423b08
commit
002b1216ee
4 changed files with 155 additions and 25 deletions
|
|
@ -94,6 +94,10 @@ class DiffusionBackend:
|
|||
# superseded (a new load) or cancelled (unload, incl. an arbiter eviction)
|
||||
# neither commits its pipeline nor stamps progress onto the current load.
|
||||
self._load_token = 0
|
||||
# Set by unload() to abort an in-flight download (which runs without the
|
||||
# lock, like the chat backend), so an eviction/unload can preempt a slow
|
||||
# load instead of blocking on the lock for the whole download.
|
||||
self._cancel_event = threading.Event()
|
||||
# The callback mutates this and generate_progress() reads it, both without
|
||||
# the lock (generate holds it for the whole call), so polling stays live.
|
||||
self._gen: Optional[_GenState] = None
|
||||
|
|
@ -106,7 +110,12 @@ class DiffusionBackend:
|
|||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
return "cuda", torch.bfloat16
|
||||
# BF16 needs Ampere+ (compute capability >= 8); pre-Ampere cards
|
||||
# (Turing/Volta/Pascal) only emulate it, so use FP16 there. (Checked by
|
||||
# capability, not torch.cuda.is_bf16_supported(), which returns True via
|
||||
# emulation on those cards and would still pick BF16.)
|
||||
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16
|
||||
return "cuda", dtype
|
||||
mps = getattr(torch.backends, "mps", None)
|
||||
if mps is not None and mps.is_available():
|
||||
return "mps", torch.float16
|
||||
|
|
@ -120,6 +129,24 @@ class DiffusionBackend:
|
|||
|
||||
return hf_hub_download(repo_id, gguf_filename, token = hf_token)
|
||||
|
||||
def _prefetch_files(
|
||||
self, repo_id: str, gguf_filename: Optional[str], base: str, base_files: list[str], hf_token: Optional[str]
|
||||
) -> None:
|
||||
"""Pre-download the GGUF + the given ``base_files`` into the HF cache,
|
||||
WITHOUT the lock and honoring ``_cancel_event``, so load_pipeline's
|
||||
from_single_file / from_pretrained hit the cache and the heavy download can
|
||||
be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``."""
|
||||
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
|
||||
|
||||
# GGUF transformer (hub repos only; a local path is already on disk).
|
||||
if gguf_filename and not Path(repo_id).expanduser().exists():
|
||||
hf_hub_download_with_xet_fallback(repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event)
|
||||
# Base repo (VAE / text-encoder / scheduler); list comes from the estimate.
|
||||
for rfilename in base_files:
|
||||
if self._cancel_event.is_set():
|
||||
raise RuntimeError("Cancelled")
|
||||
hf_hub_download_with_xet_fallback(base, rfilename, hf_token, cancel_event = self._cancel_event)
|
||||
|
||||
# ── Background load + progress ─────────────────────────────────────────
|
||||
|
||||
def begin_load(
|
||||
|
|
@ -149,6 +176,9 @@ class DiffusionBackend:
|
|||
raise RuntimeError("A diffusion load is already in progress.")
|
||||
self._load_token += 1
|
||||
token = self._load_token
|
||||
# Best-effort download preemption only; the token (not this event) is
|
||||
# the real guard that a superseded worker can't commit its pipeline.
|
||||
self._cancel_event.clear()
|
||||
# Seed with the family fallback; the worker resolves the real base
|
||||
# (a network lookup) and updates this, so begin_load never blocks.
|
||||
self._loading = _LoadingState(repo_id = repo_id, base_repo = fam.base_repo)
|
||||
|
|
@ -179,15 +209,18 @@ class DiffusionBackend:
|
|||
kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token")
|
||||
)
|
||||
kwargs["base_repo"] = base
|
||||
expected, base_files = self._estimate_download_bytes(
|
||||
kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token")
|
||||
)
|
||||
loading = self._loading
|
||||
if loading is not None:
|
||||
loading.base_repo = base
|
||||
loading.expected_bytes = self._estimate_download_bytes(
|
||||
kwargs["repo_id"],
|
||||
kwargs.get("gguf_filename"),
|
||||
base,
|
||||
kwargs.get("hf_token"),
|
||||
)
|
||||
loading.expected_bytes = expected
|
||||
# Download outside the lock so unload()/an eviction can preempt the
|
||||
# multi-GB pull; load_pipeline below then assembles from the cache.
|
||||
self._prefetch_files(
|
||||
kwargs["repo_id"], kwargs.get("gguf_filename"), base, base_files, kwargs.get("hf_token")
|
||||
)
|
||||
self.load_pipeline(**kwargs)
|
||||
with self._lock:
|
||||
# Only clear the marker if this load is still the current one; a
|
||||
|
|
@ -225,22 +258,26 @@ class DiffusionBackend:
|
|||
@staticmethod
|
||||
def _estimate_download_bytes(
|
||||
repo_id: str, gguf_filename: Optional[str], base_repo: str, hf_token: Optional[str]
|
||||
) -> int:
|
||||
) -> tuple[int, list[str]]:
|
||||
"""Total download size for the progress bar, plus the base-repo files to
|
||||
fetch (the prefetch reuses this list, so the base is listed only once)."""
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
api = HfApi()
|
||||
total = 0
|
||||
base_files: list[str] = []
|
||||
try:
|
||||
if gguf_filename:
|
||||
info = api.model_info(repo_id, files_metadata = True, token = hf_token)
|
||||
total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename)
|
||||
base_info = api.model_info(base_repo, files_metadata = True, token = hf_token)
|
||||
total += sum(
|
||||
s.size or 0 for s in base_info.siblings if _base_file_downloaded(s.rfilename)
|
||||
)
|
||||
for s in base_info.siblings:
|
||||
if _base_file_downloaded(s.rfilename):
|
||||
base_files.append(s.rfilename)
|
||||
total += s.size or 0
|
||||
except Exception as exc: # noqa: BLE001 — estimate is best-effort
|
||||
logger.warning("diffusion.size_estimate_failed: %s", exc)
|
||||
return total
|
||||
return total, base_files
|
||||
|
||||
@staticmethod
|
||||
def _cache_bytes(repo_id: str) -> int:
|
||||
|
|
@ -306,6 +343,9 @@ class DiffusionBackend:
|
|||
torch_dtype = dtype,
|
||||
config = base,
|
||||
subfolder = "transformer",
|
||||
# Forward the token: the config is fetched from the (possibly gated)
|
||||
# base repo before from_pretrained gets a chance to authenticate.
|
||||
token = hf_token,
|
||||
)
|
||||
|
||||
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer}
|
||||
|
|
@ -427,6 +467,9 @@ class DiffusionBackend:
|
|||
}
|
||||
|
||||
def unload(self) -> dict[str, Any]:
|
||||
# Abort an in-flight download (it runs without the lock and checks this),
|
||||
# so unload/an eviction returns promptly instead of waiting it out.
|
||||
self._cancel_event.set()
|
||||
with self._lock:
|
||||
self._unload_locked()
|
||||
# Cancel any in-flight load (its worker checks this token before
|
||||
|
|
|
|||
|
|
@ -216,6 +216,7 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path):
|
|||
gguf_filename = "model.gguf",
|
||||
base_repo = "base/repo",
|
||||
family_override = "z-image",
|
||||
hf_token = "hf_secret",
|
||||
)
|
||||
assert status["loaded"] is True
|
||||
assert status["family"] == "z-image"
|
||||
|
|
@ -226,6 +227,8 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path):
|
|||
# Transformer built from the local GGUF, pipeline assembled from the base repo.
|
||||
assert _FakeTransformer.last["path"] == str((tmp_path / "model.gguf").resolve())
|
||||
assert _FakeTransformer.last["subfolder"] == "transformer"
|
||||
# The token reaches the (possibly gated) base config fetch and the pipeline.
|
||||
assert _FakeTransformer.last["token"] == "hf_secret"
|
||||
assert _FakePipeline.last["base"] == "base/repo"
|
||||
assert "transformer" in _FakePipeline.last
|
||||
|
||||
|
|
@ -398,10 +401,12 @@ def test_generate_qwen_uses_true_cfg_scale(fake_runtime, tmp_path):
|
|||
|
||||
def test_begin_load_rejects_concurrent(monkeypatch):
|
||||
backend = DiffusionBackend()
|
||||
# The worker resolves the base via a network lookup; stub it so the test is offline.
|
||||
# The worker resolves the base + downloads, both over the network; stub them
|
||||
# so the test is offline.
|
||||
monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None)
|
||||
monkeypatch.setattr(DiffusionBackend, "_prefetch_files", lambda self, *a, **k: None)
|
||||
monkeypatch.setattr(
|
||||
DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: 0)
|
||||
DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: (0, []))
|
||||
)
|
||||
# Block the spawned worker so the load stays "in progress".
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -429,3 +434,60 @@ def test_unload_cancels_in_flight_load(fake_runtime):
|
|||
base_repo = fam.base_repo,
|
||||
_load_token = token,
|
||||
)
|
||||
|
||||
|
||||
def test_pick_dtype_bf16_only_on_ampere(fake_runtime, monkeypatch):
|
||||
# BF16 only on Ampere+ (cc >= 8); pre-Ampere cards must fall back to FP16.
|
||||
torch = sys.modules["torch"]
|
||||
backend = DiffusionBackend()
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True, raising = False)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 0), raising = False)
|
||||
assert backend._pick_device_and_dtype() == ("cuda", torch.bfloat16)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5), raising = False)
|
||||
assert backend._pick_device_and_dtype() == ("cuda", torch.float16)
|
||||
|
||||
|
||||
def test_unload_sets_cancel_event(fake_runtime):
|
||||
# unload signals an in-flight download (which runs without the lock) to abort.
|
||||
backend = DiffusionBackend()
|
||||
assert not backend._cancel_event.is_set()
|
||||
backend.unload()
|
||||
assert backend._cancel_event.is_set()
|
||||
|
||||
|
||||
def test_prefetch_aborts_when_cancelled(tmp_path):
|
||||
# A prefetch interrupted by unload (cancel event set) raises rather than
|
||||
# downloading the whole base, so the load can be preempted mid-download.
|
||||
backend = DiffusionBackend()
|
||||
backend._cancel_event.set()
|
||||
# Local gguf path so the transformer download is skipped; the base loop hits
|
||||
# the cancel check on its first file (no network).
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
with pytest.raises(RuntimeError, match = "Cancelled"):
|
||||
backend._prefetch_files(
|
||||
str(tmp_path), "model.gguf", "Tongyi-MAI/Z-Image-Turbo",
|
||||
["vae/diffusion_pytorch_model.safetensors"], None,
|
||||
)
|
||||
|
||||
|
||||
def test_prefetch_downloads_gguf_and_base(monkeypatch, tmp_path):
|
||||
backend = DiffusionBackend()
|
||||
calls: list = []
|
||||
monkeypatch.setattr(
|
||||
"utils.hf_xet_fallback.hf_hub_download_with_xet_fallback",
|
||||
lambda repo, fn, tok, **k: (calls.append((repo, fn)), f"/cache/{fn}")[1],
|
||||
)
|
||||
# Hub repo: the GGUF transformer and each base file are fetched.
|
||||
backend._prefetch_files(
|
||||
"unsloth/Z-Image-Turbo-GGUF", "model.gguf", "base/repo",
|
||||
["vae/x.safetensors", "text_encoder/y.safetensors"], "hf_tok",
|
||||
)
|
||||
assert ("unsloth/Z-Image-Turbo-GGUF", "model.gguf") in calls
|
||||
assert ("base/repo", "vae/x.safetensors") in calls
|
||||
assert ("base/repo", "text_encoder/y.safetensors") in calls
|
||||
# Local GGUF path: the transformer download is skipped, base still fetched.
|
||||
calls.clear()
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend._prefetch_files(str(tmp_path), "model.gguf", "base/repo", ["vae/x.safetensors"], None)
|
||||
assert all(repo != str(tmp_path) for repo, _ in calls)
|
||||
assert ("base/repo", "vae/x.safetensors") in calls
|
||||
|
|
|
|||
|
|
@ -985,11 +985,28 @@ export const IMAGE_GEN_TASKS = [
|
|||
"image-text-to-image",
|
||||
] as const;
|
||||
|
||||
// Editing/inpaint checkpoints are tagged image-to-image but need an input image,
|
||||
// which the text-to-image backend rejects (mirrors its _EDIT_KEYWORDS). Hidden by
|
||||
// id so they don't show in the Images picker only to 400 on load. Keeping the
|
||||
// image-to-image task itself is required: some supported models (FLUX.2-klein)
|
||||
// carry that tag too.
|
||||
const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint"] as const;
|
||||
function isImageEditModel(repoId: string | null | undefined): boolean {
|
||||
if (!repoId) return false;
|
||||
const id = repoId.toLowerCase();
|
||||
return IMAGE_EDIT_KEYWORDS.some((kw) => id.includes(kw));
|
||||
}
|
||||
|
||||
// Gate an on-device model by the picker's task scope. With a filter (the Images
|
||||
// page) keep only matching tasks; with no filter (chat) drop image-generation
|
||||
// models so a downloaded diffusion GGUF doesn't show up as a loadable chat model.
|
||||
function passesTaskGate(repoTask: string | null | undefined, filter: HfTaskFilter): boolean {
|
||||
if (filter) return taskMatchesFilter(repoTask, filter);
|
||||
// page) keep only matching, non-editing tasks; with no filter (chat) drop
|
||||
// image-generation models so a downloaded diffusion GGUF doesn't show up as a
|
||||
// loadable chat model.
|
||||
function passesTaskGate(
|
||||
repoTask: string | null | undefined,
|
||||
repoId: string | null | undefined,
|
||||
filter: HfTaskFilter,
|
||||
): boolean {
|
||||
if (filter) return taskMatchesFilter(repoTask, filter) && !isImageEditModel(repoId);
|
||||
return !(repoTask != null && (IMAGE_GEN_TASKS as readonly string[]).includes(repoTask));
|
||||
}
|
||||
|
||||
|
|
@ -1563,7 +1580,9 @@ export function HubModelPicker({
|
|||
// the chat classifier marks image tasks "unsupported".
|
||||
const isChatSupported = useCallback(
|
||||
(r: HfModelResult) => {
|
||||
if (task && taskMatchesFilter(r.pipelineTag, task)) return true;
|
||||
// Image tab: keep task-matching results, but drop editing checkpoints the
|
||||
// backend rejects (so they don't appear in Hub search only to 400 on load).
|
||||
if (task && taskMatchesFilter(r.pipelineTag, task)) return !isImageEditModel(r.id);
|
||||
return (
|
||||
classifyUnslothSupport({
|
||||
modelId: r.id,
|
||||
|
|
@ -1744,7 +1763,7 @@ export function HubModelPicker({
|
|||
const sortedCachedGguf = useMemo(
|
||||
() =>
|
||||
sortCachedRepos(
|
||||
cachedGguf.filter((c) => passesTaskGate(c.task, task)),
|
||||
cachedGguf.filter((c) => passesTaskGate(c.task, c.repo_id, task)),
|
||||
downloadedSort,
|
||||
loadTimes,
|
||||
),
|
||||
|
|
@ -1758,7 +1777,7 @@ export function HubModelPicker({
|
|||
task
|
||||
? []
|
||||
: sortCachedRepos(
|
||||
cachedModels.filter((c) => passesTaskGate(c.task, task)),
|
||||
cachedModels.filter((c) => passesTaskGate(c.task, c.repo_id, task)),
|
||||
downloadedSort,
|
||||
loadTimes,
|
||||
),
|
||||
|
|
@ -1779,7 +1798,7 @@ export function HubModelPicker({
|
|||
sortLocalModels(
|
||||
lmStudioModels.filter(
|
||||
(m) =>
|
||||
passesTaskGate(m.task, task) &&
|
||||
passesTaskGate(m.task, m.model_id ?? m.id, task) &&
|
||||
localModelMatchesFormat(m, formatFilter) &&
|
||||
matchesLocalQuery(m),
|
||||
),
|
||||
|
|
@ -1797,7 +1816,7 @@ export function HubModelPicker({
|
|||
sortLocalModels(
|
||||
localDirModels.filter(
|
||||
(m) =>
|
||||
passesTaskGate(m.task, task) &&
|
||||
passesTaskGate(m.task, m.model_id ?? m.id, task) &&
|
||||
(!chatOnly ||
|
||||
localModelIsGguf(m) ||
|
||||
(isMac && localModelIsMlx(m))) &&
|
||||
|
|
@ -1824,7 +1843,7 @@ export function HubModelPicker({
|
|||
sortLocalModels(
|
||||
customFolderModels.filter(
|
||||
(m) =>
|
||||
passesTaskGate(m.task, task) &&
|
||||
passesTaskGate(m.task, m.model_id ?? m.id, task) &&
|
||||
localModelMatchesFormat(m, formatFilter) &&
|
||||
matchesLocalQuery(m),
|
||||
),
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ import type {
|
|||
ModelSelectorChangeMeta,
|
||||
} from "@/components/assistant-ui/model-selector/types";
|
||||
import { ModelLoadDescription } from "@/features/chat/components/model-load-status";
|
||||
import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store";
|
||||
import { formatBytes, formatEta } from "@/features/hub/lib/format";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
|
@ -630,7 +631,12 @@ export function ImagesPage() {
|
|||
try {
|
||||
// Returns immediately — the load runs in the background; we poll for it.
|
||||
// The backend infers the family + base diffusers repo from the repo id.
|
||||
await loadDiffusionModel({ model_path: repoId, gguf_filename: ggufFilename });
|
||||
// Forward the saved HF token so gated bases (FLUX dev/klein) can download.
|
||||
await loadDiffusionModel({
|
||||
model_path: repoId,
|
||||
gguf_filename: ggufFilename,
|
||||
hf_token: hfApiToken(getHfToken()),
|
||||
});
|
||||
} catch (err) {
|
||||
dismissLoadToast();
|
||||
toast.error(err instanceof Error ? err.message : "Failed to start load");
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue