Merge remote-tracking branch 'origin/diffusion-lora-ux' into diffusion-lora-training
This commit is contained in:
Daniel Han
commit
b4b66907fe
12 changed files with 260 additions and 160 deletions
|
|
@ -29,8 +29,10 @@ from loggers import get_logger
|
|||
from utils.hardware import clear_gpu_cache
|
||||
|
||||
from .diffusion_families import (
|
||||
DIFFUSION_CANCELLED_MSG,
|
||||
DIFFUSION_NOT_LOADED_MSG,
|
||||
DiffusionFamily,
|
||||
detect_family,
|
||||
detect_family_for_pick,
|
||||
resolve_base_repo,
|
||||
resolve_local_gguf_child,
|
||||
supported_family_names,
|
||||
|
|
@ -43,11 +45,10 @@ from .diffusion_device import (
|
|||
from .diffusion_memory import (
|
||||
OFFLOAD_NONE,
|
||||
apply_memory_plan,
|
||||
estimate_gguf_dense_mib,
|
||||
estimate_gguf_resident_mib,
|
||||
estimate_image_runtime_mib,
|
||||
estimate_safetensors_dense_mib,
|
||||
file_size_mib,
|
||||
infer_gguf_quant_label,
|
||||
plan_diffusion_memory,
|
||||
snapshot_device_memory,
|
||||
)
|
||||
|
|
@ -410,22 +411,6 @@ class DiffusionBackend:
|
|||
base, rfilename, hf_token, cancel_event = self._cancel_event
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _detect_family_for_pick(
|
||||
repo_id: str, gguf_filename: Optional[str], family_override: Optional[str]
|
||||
) -> Optional[DiffusionFamily]:
|
||||
"""Detect the family from the repo id, falling back to the combined
|
||||
path/filename for a direct local .gguf pick. The frontend splits such a
|
||||
pick into (parent dir, basename), so the family keyword can live only in
|
||||
the filename (e.g. /models/z-image-turbo-Q4_K_M.gguf) while the parent
|
||||
directory carries none; scan it too when the directory alone is
|
||||
undetectable. Only used as a fallback, so remote 'org/name' picks and
|
||||
explicit overrides behave exactly as before."""
|
||||
fam = detect_family(repo_id, family_override)
|
||||
if fam is None and gguf_filename and not family_override:
|
||||
fam = detect_family(f"{repo_id}/{gguf_filename}", family_override)
|
||||
return fam
|
||||
|
||||
def validate_load_request(
|
||||
self,
|
||||
repo_id: str,
|
||||
|
|
@ -441,7 +426,7 @@ class DiffusionBackend:
|
|||
undetectable family, and ValueError/FileNotFoundError for a bad local path.
|
||||
Touches no GPU, network, or state."""
|
||||
kind = resolve_model_kind(gguf_filename, model_kind)
|
||||
fam = self._detect_family_for_pick(repo_id, gguf_filename, family_override)
|
||||
fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
|
||||
if fam is None:
|
||||
raise ValueError(
|
||||
f"'{repo_id}' is not a supported diffusion image model. Supported families: "
|
||||
|
|
@ -564,7 +549,7 @@ class DiffusionBackend:
|
|||
# Resolve the base repo and estimate sizes on this thread (both network
|
||||
# calls) so begin_load returns instantly; the bar shows raw bytes until
|
||||
# the total lands. This is the only writer of _loading's fields here.
|
||||
fam = self._detect_family_for_pick(
|
||||
fam = detect_family_for_pick(
|
||||
kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override")
|
||||
)
|
||||
kind = resolve_model_kind(kwargs.get("gguf_filename"), kwargs.get("model_kind"))
|
||||
|
|
@ -780,7 +765,6 @@ class DiffusionBackend:
|
|||
plan = self._plan_memory(
|
||||
target,
|
||||
single_file_path,
|
||||
gguf_filename,
|
||||
base,
|
||||
fam,
|
||||
memory_mode,
|
||||
|
|
@ -973,7 +957,11 @@ class DiffusionBackend:
|
|||
quant = transformer_quant_engaged,
|
||||
attention_backend = attention_engaged,
|
||||
compile_kwargs = {
|
||||
"fullgraph": cache_engaged is None,
|
||||
# Mirrors apply_speed_optims' fullgraph decision: an active
|
||||
# step cache OR a planned offload graph-breaks, so the cached
|
||||
# bundle must be keyed on the same fullgraph setting.
|
||||
"fullgraph": cache_engaged is None
|
||||
and plan.offload_policy == OFFLOAD_NONE,
|
||||
"dynamic": effective_speed != SPEED_MAX,
|
||||
"mode": "max-autotune-no-cudagraphs"
|
||||
if effective_speed == SPEED_MAX
|
||||
|
|
@ -989,8 +977,21 @@ class DiffusionBackend:
|
|||
family = fam,
|
||||
speed_mode = effective_speed,
|
||||
cache_active = cache_engaged is not None,
|
||||
# The planned offload policy: group/model/sequential offload installs
|
||||
# compiler-disabled onload hooks, so compile must drop fullgraph.
|
||||
offload_active = plan.offload_policy != OFFLOAD_NONE,
|
||||
logger = logger,
|
||||
)
|
||||
if transformer_quant_engaged is not None and not speed_applied.get("compiled"):
|
||||
# Promotion above could not engage compile (e.g. the family is not
|
||||
# compile-friendly, or compile_repeated_blocks failed): the quantized
|
||||
# transformer is now running eager, which is far slower than the GGUF
|
||||
# path it replaced. Surface it loudly rather than hiding the regression.
|
||||
logger.warning(
|
||||
"diffusion.transformer_quant: %s engaged but the transformer is NOT "
|
||||
"compiled; eager torchao quant is ~30x slower than GGUF here",
|
||||
transformer_quant_engaged,
|
||||
)
|
||||
# Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
|
||||
# also before placement so the offload hooks move the smaller weights.
|
||||
te_quant = quantize_text_encoders(
|
||||
|
|
@ -1046,6 +1047,9 @@ class DiffusionBackend:
|
|||
if eager_patched:
|
||||
uninstall_patches()
|
||||
uninstall_arch_patches()
|
||||
# Also free the half-built pipe's VRAM: the failed load never
|
||||
# commits _state, so nothing else reclaims it until the next unload.
|
||||
clear_gpu_cache()
|
||||
|
||||
logger.info(
|
||||
"diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s",
|
||||
|
|
@ -1144,7 +1148,6 @@ class DiffusionBackend:
|
|||
self,
|
||||
target: DiffusionDeviceTarget,
|
||||
single_file_path: Optional[str],
|
||||
gguf_filename: Optional[str],
|
||||
base: str,
|
||||
fam: DiffusionFamily,
|
||||
memory_mode: Optional[str],
|
||||
|
|
@ -1158,9 +1161,10 @@ class DiffusionBackend:
|
|||
offload policy + VAE memory savers. Kept on the backend so the cached base
|
||||
repo (companion text-encoder / VAE) feeds the size estimate.
|
||||
|
||||
The size estimate is per-kind: a GGUF dequantises (a 4-bit file ~4x), a
|
||||
safetensors single-file loads near its on-disk size, and a full pipeline is
|
||||
one cached download (transformer + companions) that is already compressed."""
|
||||
The size estimate is per-kind: diffusers keeps GGUF weights packed (per-matmul
|
||||
transient dequant), so a GGUF loads near its on-disk size; a safetensors
|
||||
single-file loads near its on-disk size (it carries its dtype); and a full
|
||||
pipeline is one cached download (transformer + companions), already compressed."""
|
||||
device_memory = snapshot_device_memory(target)
|
||||
if kind == "pipeline":
|
||||
# The whole repo (transformer + companions) is one cached download; the
|
||||
|
|
@ -1172,18 +1176,18 @@ class DiffusionBackend:
|
|||
else:
|
||||
if kind == "single_file":
|
||||
# Safetensors single-file: no dequant expansion (it carries its dtype).
|
||||
transformer_dense = estimate_safetensors_dense_mib(file_size_mib(single_file_path))
|
||||
else:
|
||||
transformer_dense = estimate_gguf_dense_mib(
|
||||
file_size_mib(single_file_path), infer_gguf_quant_label(gguf_filename)
|
||||
transformer_resident = estimate_safetensors_dense_mib(
|
||||
file_size_mib(single_file_path)
|
||||
)
|
||||
else:
|
||||
transformer_resident = estimate_gguf_resident_mib(file_size_mib(single_file_path))
|
||||
# The companion components (VAE + text encoders) load near their on-disk
|
||||
# size; sum whatever the prefetch already placed in the base-repo cache.
|
||||
companion = self._cache_bytes(base)
|
||||
companion_mib = int(companion // (1024 * 1024)) if companion else None
|
||||
model_dense_mib = None
|
||||
if transformer_dense is not None:
|
||||
model_dense_mib = transformer_dense + (companion_mib or 0)
|
||||
if transformer_resident is not None:
|
||||
model_dense_mib = transformer_resident + (companion_mib or 0)
|
||||
runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = fam.name)
|
||||
return plan_diffusion_memory(
|
||||
target = target,
|
||||
|
|
@ -1392,7 +1396,7 @@ class DiffusionBackend:
|
|||
with self._lock:
|
||||
state = self._state
|
||||
if state is None:
|
||||
raise RuntimeError("No diffusion model is loaded.")
|
||||
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
|
||||
# Register under _lock so unload()/a load can signal THIS generation.
|
||||
# A cancel that arrived before now either nulled _state (we raised
|
||||
# above) or targets an older generation, so nothing is lost.
|
||||
|
|
@ -1633,7 +1637,7 @@ class DiffusionBackend:
|
|||
# A cancelled denoise returns early with a partial/garbage image;
|
||||
# don't hand it back to be persisted.
|
||||
if cancel.is_set():
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
# The first compiled generation just paid the compile cost; persist the
|
||||
# warm torch.compile cache bundle when saving is enabled (distributor /
|
||||
# first-run warm). Idempotent + best-effort -- never fails a generation.
|
||||
|
|
|
|||
|
|
@ -20,6 +20,14 @@ from pathlib import Path, PurePosixPath
|
|||
from typing import Optional
|
||||
|
||||
|
||||
# Runtime->route contract: the RuntimeError messages a backend raises for
|
||||
# client-recoverable generate states. The /images/generate route matches these
|
||||
# EXACTLY to return 409 (vs a sanitized 500 for real failures), so both engines
|
||||
# must raise them verbatim -- keep them named here, not as scattered literals.
|
||||
DIFFUSION_NOT_LOADED_MSG = "No diffusion model is loaded."
|
||||
DIFFUSION_CANCELLED_MSG = "Diffusion generation was cancelled."
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class DiffusionFamily:
|
||||
name: str
|
||||
|
|
@ -338,6 +346,24 @@ def supported_family_names() -> tuple[str, ...]:
|
|||
return tuple(fam.name for fam in _FAMILIES)
|
||||
|
||||
|
||||
def detect_family_for_pick(
|
||||
repo_id: str,
|
||||
gguf_filename: Optional[str] = None,
|
||||
override: Optional[str] = None,
|
||||
) -> Optional[DiffusionFamily]:
|
||||
"""``detect_family``, falling back to the combined path/filename for a direct
|
||||
local ``.gguf`` pick. The frontend splits such a pick into (parent dir, basename),
|
||||
so the family keyword can live only in the filename (e.g.
|
||||
``/models/z-image-turbo-Q4_K_M.gguf``) while the parent directory carries none;
|
||||
scan the combined string too when the directory alone is undetectable. Only a
|
||||
fallback, so remote ``org/name`` picks and explicit overrides behave exactly as
|
||||
``detect_family``. Shared by both engines so validation and load can't diverge."""
|
||||
fam = detect_family(repo_id, override)
|
||||
if fam is None and gguf_filename and not override:
|
||||
fam = detect_family(f"{repo_id}/{gguf_filename}", override)
|
||||
return fam
|
||||
|
||||
|
||||
def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
|
||||
"""The companion diffusers repo: caller-supplied if given, else the family fallback."""
|
||||
base = (base_repo or "").strip()
|
||||
|
|
|
|||
|
|
@ -218,51 +218,22 @@ def file_size_mib(path: Any) -> Optional[int]:
|
|||
return None
|
||||
|
||||
|
||||
def infer_gguf_quant_label(filename: Optional[str]) -> Optional[str]:
|
||||
"""Pull a quant tag (Q4_K_M, Q8_0, BF16, ...) out of a GGUF filename."""
|
||||
if not filename:
|
||||
return None
|
||||
from pathlib import Path
|
||||
def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]:
|
||||
"""Approximate the RESIDENT device size of a GGUF transformer loaded through
|
||||
diffusers' ``GGUFQuantizationConfig``.
|
||||
|
||||
stem = Path(filename).name
|
||||
if stem.lower().endswith(".gguf"):
|
||||
stem = stem[:-5]
|
||||
parts = [p.upper() for p in stem.replace("-", "_").split("_") if p]
|
||||
for index, part in enumerate(parts):
|
||||
if part in ("BF16", "F16", "FP16", "FP8", "Q8", "Q6", "Q5", "Q4", "Q3", "Q2"):
|
||||
suffix = parts[index + 1 :]
|
||||
# Quant names carry either a K-family suffix (Q4_K_M) or a legacy
|
||||
# numeric one (Q8_0, Q5_1); keep up to two suffix tokens.
|
||||
if suffix and suffix[0] in ("K", "M", "S", "L", "XS", "XXS", "0", "1"):
|
||||
return "_".join([part] + suffix[:2])
|
||||
return part
|
||||
if part.startswith("IQ") or part.startswith("UD"):
|
||||
return "_".join(parts[index : index + 3])
|
||||
return None
|
||||
The weights stay PACKED on the device as quantised bytes (``GGUFParameter`` /
|
||||
uint8); ``GGUFLinear.forward`` dequantises each weight to the bf16 compute dtype
|
||||
transiently for its matmul and frees it immediately, so the persistent footprint
|
||||
is ~= the on-disk tensor size, NOT the unpacked bf16 size. Measured on
|
||||
Z-Image-Turbo: Q2_K 3.64 GiB -> 3.68 GiB, Q8_0 7.22 GiB -> 7.25 GiB resident.
|
||||
The transient per-op dequant is covered by the separate runtime headroom.
|
||||
|
||||
|
||||
def estimate_gguf_dense_mib(storage_mib: Optional[int], quant: Optional[str]) -> Optional[int]:
|
||||
"""Approximate the dequantised (device) size of a GGUF from its on-disk size
|
||||
and quant label. The compute dtype is bf16/fp16, so a 4-bit file roughly
|
||||
quadruples once unpacked; higher-bit quants expand less."""
|
||||
(The prior per-quant expansion assumed a full unpack that never happens on this
|
||||
path; it over-estimated e.g. Q2 ~7.6x, forcing needless offload.)"""
|
||||
if storage_mib is None:
|
||||
return None
|
||||
q = (quant or "").upper()
|
||||
if any(t in q for t in ("BF16", "F16", "FP16")):
|
||||
return storage_mib
|
||||
if "FP8" in q or "Q8" in q:
|
||||
return int(storage_mib * 2.0)
|
||||
if "Q6" in q:
|
||||
return int(storage_mib * 2.8)
|
||||
if "Q5" in q:
|
||||
return int(storage_mib * 3.3)
|
||||
if "Q4" in q or "IQ4" in q or "UD" in q:
|
||||
return int(storage_mib * 4.0)
|
||||
if "Q3" in q or "IQ3" in q:
|
||||
return int(storage_mib * 5.3)
|
||||
if "Q2" in q or "Q1" in q or "IQ2" in q or "IQ1" in q:
|
||||
return int(storage_mib * 8.0)
|
||||
return int(storage_mib * 4.0) # unknown: assume 4-bit-ish
|
||||
return int(storage_mib * 1.05) # small margin for allocator + bf16 norms/biases
|
||||
|
||||
|
||||
def estimate_safetensors_dense_mib(storage_mib: Optional[int]) -> Optional[int]:
|
||||
|
|
|
|||
|
|
@ -164,11 +164,17 @@ def apply_speed_optims(
|
|||
family: Any,
|
||||
speed_mode: str = SPEED_OFF,
|
||||
cache_active: bool = False,
|
||||
offload_active: bool = False,
|
||||
logger: Any = None,
|
||||
) -> dict[str, bool]:
|
||||
"""Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline,
|
||||
BEFORE placement / offload. Returns which optimisations actually engaged. Every
|
||||
step is best-effort: a pipeline that doesn't support one is simply skipped."""
|
||||
step is best-effort: a pipeline that doesn't support one is simply skipped.
|
||||
|
||||
``offload_active`` is the planned offload policy != none: group/model/sequential
|
||||
offloading installs ``@torch.compiler.disable``d onload hooks, so the compile must
|
||||
drop ``fullgraph`` (same reason as an active step cache) or it crashes at the first
|
||||
denoise step."""
|
||||
applied = {
|
||||
"channels_last": False,
|
||||
"cudnn_benchmark": False,
|
||||
|
|
@ -178,12 +184,11 @@ def apply_speed_optims(
|
|||
"compiled_dequant": False,
|
||||
}
|
||||
mode = normalize_speed_mode(speed_mode)
|
||||
# TF32 is the one PROCESS-GLOBAL flag we flip (on max). Restore it whenever this
|
||||
# load isn't max, so a later default/off diffusion load -- or chat inference in the
|
||||
# same long-lived process -- doesn't silently inherit a prior max load's TF32 and
|
||||
# lose the bit-identical default the regression harness checks.
|
||||
if mode != SPEED_MAX:
|
||||
_restore_tf32(logger)
|
||||
# TF32 and cudnn.benchmark are the process-global flags this may flip (TF32 on max,
|
||||
# cudnn.benchmark on any non-off CUDA load). The caller snapshots them before this
|
||||
# call and restores on unload / failed load via snapshot_backend_flags /
|
||||
# restore_backend_flags, so a later `off` load -- or chat inference in the same
|
||||
# process -- never inherits them. We keep no separate bookkeeping here.
|
||||
if mode == SPEED_OFF:
|
||||
return applied
|
||||
|
||||
|
|
@ -212,11 +217,19 @@ def apply_speed_optims(
|
|||
applied["compiled_dequant"] = gguf_compile.install_compiled_dequant(logger)
|
||||
elif compile_eligible(target, is_gguf = is_gguf, family = family):
|
||||
applied["compiled"] = _compile_repeated_blocks(
|
||||
pipe, logger, max_autotune = False, cache_active = cache_active
|
||||
pipe,
|
||||
logger,
|
||||
max_autotune = False,
|
||||
cache_active = cache_active,
|
||||
offload_active = offload_active,
|
||||
)
|
||||
elif mode == SPEED_MAX and compile_eligible(target, is_gguf = is_gguf, family = family):
|
||||
applied["compiled"] = _compile_repeated_blocks(
|
||||
pipe, logger, max_autotune = True, cache_active = cache_active
|
||||
pipe,
|
||||
logger,
|
||||
max_autotune = True,
|
||||
cache_active = cache_active,
|
||||
offload_active = offload_active,
|
||||
)
|
||||
|
||||
if mode == SPEED_MAX:
|
||||
|
|
@ -247,6 +260,7 @@ def _compile_repeated_blocks(
|
|||
*,
|
||||
max_autotune: bool = False,
|
||||
cache_active: bool = False,
|
||||
offload_active: bool = False,
|
||||
) -> bool:
|
||||
transformer = getattr(pipe, "transformer", None)
|
||||
fn = getattr(transformer, "compile_repeated_blocks", None)
|
||||
|
|
@ -259,14 +273,34 @@ def _compile_repeated_blocks(
|
|||
# / max-autotune) are deliberately NOT used: they crash on the regionally-compiled
|
||||
# block because its static output buffer is overwritten across denoise steps.
|
||||
#
|
||||
# fullgraph drops to False when a step cache is engaged: FBCache's per-step decision is
|
||||
# ``@torch.compiler.disable``d, i.e. a graph break, which fullgraph=True rejects ("Skip
|
||||
# inlining torch.compiler.disable()d function"). The break is cheap and the rest of the
|
||||
# block still compiles.
|
||||
kwargs: dict[str, Any] = {"fullgraph": not cache_active, "dynamic": not max_autotune}
|
||||
# fullgraph drops to False when a step cache OR CPU offloading is engaged: both insert
|
||||
# an ``@torch.compiler.disable``d function into the forward -- FBCache's per-step
|
||||
# decision, and group/model/sequential offload's ``ModuleGroup.onload_`` streaming hook
|
||||
# -- i.e. a graph break, which fullgraph=True rejects ("Skip inlining
|
||||
# torch.compiler.disable()d function"). The break is cheap and the rest of the block
|
||||
# still compiles.
|
||||
kwargs: dict[str, Any] = {
|
||||
"fullgraph": not (cache_active or offload_active),
|
||||
"dynamic": not max_autotune,
|
||||
}
|
||||
if max_autotune:
|
||||
kwargs["mode"] = "max-autotune-no-cudagraphs"
|
||||
try:
|
||||
import torch
|
||||
|
||||
# Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block
|
||||
# shape through compile_repeated_blocks; Z-Image needs ~11, above dynamo's default
|
||||
# recompile_limit of 8. Once the limit is hit a resident load hard-errors under
|
||||
# fullgraph (and an offload/cache load silently drops the overflow blocks to eager),
|
||||
# so raise it well past that (64) for headroom on larger heterogeneous DiTs. This is
|
||||
# diffusers' own documented fix for regional-compile recompilation (their guide bumps
|
||||
# cache_size_limit). Deliberately NOT force_parameter_static_shapes=False: it doesn't
|
||||
# cut the variant count here and makes each compile ~6x slower (24s -> 143s cold).
|
||||
dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None)
|
||||
if dynamo_cfg is not None:
|
||||
for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver
|
||||
if hasattr(dynamo_cfg, _limit_attr):
|
||||
setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
|
||||
fn(**kwargs)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 — optimisation only
|
||||
|
|
@ -284,22 +318,10 @@ def _enable_cudnn_benchmark(logger: Any) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
# The TF32 flag values from before the first max load flipped them, so a later
|
||||
# non-max load / unload can put the process back exactly as it found it (rather than
|
||||
# forcing a hardcoded default that might clobber another component's choice).
|
||||
_tf32_prev: Optional[tuple[bool, bool]] = None
|
||||
|
||||
|
||||
def _enable_tf32(logger: Any) -> bool:
|
||||
global _tf32_prev
|
||||
try:
|
||||
import torch
|
||||
|
||||
if _tf32_prev is None:
|
||||
_tf32_prev = (
|
||||
torch.backends.cuda.matmul.allow_tf32,
|
||||
torch.backends.cudnn.allow_tf32,
|
||||
)
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
return True
|
||||
|
|
@ -308,26 +330,6 @@ def _enable_tf32(logger: Any) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def restore_tf32(logger: Any = None) -> None:
|
||||
"""Put the process-global TF32 flags back to their pre-max-load values. No-op if
|
||||
a max load never set them. Called on a non-max load and on unload."""
|
||||
_restore_tf32(logger)
|
||||
|
||||
|
||||
def _restore_tf32(logger: Any) -> None:
|
||||
global _tf32_prev
|
||||
if _tf32_prev is None:
|
||||
return
|
||||
try:
|
||||
import torch
|
||||
torch.backends.cuda.matmul.allow_tf32 = _tf32_prev[0]
|
||||
torch.backends.cudnn.allow_tf32 = _tf32_prev[1]
|
||||
except Exception as exc: # noqa: BLE001 — best-effort restore
|
||||
_warn(logger, "tf32_restore", exc)
|
||||
finally:
|
||||
_tf32_prev = None
|
||||
|
||||
|
||||
def _fuse_qkv(pipe: Any, logger: Any) -> bool:
|
||||
for owner in (pipe, getattr(pipe, "transformer", None)):
|
||||
fn = getattr(owner, "fuse_qkv_projections", None)
|
||||
|
|
|
|||
|
|
@ -35,8 +35,10 @@ from typing import Any, Optional
|
|||
|
||||
from core.inference.diffusion_device import resolve_diffusion_device_target
|
||||
from core.inference.diffusion_families import (
|
||||
DIFFUSION_CANCELLED_MSG,
|
||||
DIFFUSION_NOT_LOADED_MSG,
|
||||
DiffusionFamily,
|
||||
detect_family,
|
||||
detect_family_for_pick,
|
||||
family_sd_cpp_supported,
|
||||
resolve_base_repo,
|
||||
resolve_local_gguf_child,
|
||||
|
|
@ -248,7 +250,10 @@ class SdCppDiffusionBackend:
|
|||
raise ValueError(
|
||||
"gguf_filename is required: the native engine loads single-file GGUF checkpoints only."
|
||||
)
|
||||
fam = detect_family(repo_id, family_override)
|
||||
# Use the filename-fallback detector the route validated with, so a local
|
||||
# .gguf pick whose family keyword lives only in the basename doesn't pass
|
||||
# validation and then dead-end here on a no-GPU (native-routed) host.
|
||||
fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
|
||||
if fam is None:
|
||||
raise ValueError(
|
||||
f"'{repo_id}' is not a supported diffusion image model. Supported families: "
|
||||
|
|
@ -505,7 +510,7 @@ class SdCppDiffusionBackend:
|
|||
with self._lock:
|
||||
state = self._state
|
||||
if state is None:
|
||||
raise RuntimeError("No diffusion model is loaded.")
|
||||
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
|
||||
self._active_generate_cancel = cancel
|
||||
engine = self._resolve_engine()
|
||||
try:
|
||||
|
|
@ -554,7 +559,7 @@ class SdCppDiffusionBackend:
|
|||
lora_dir = str(Path(tmpdir) / "loras")
|
||||
for index in range(max(1, int(batch_size))):
|
||||
if cancel.is_set():
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
# Distinct seed per batch image (sd-cli is one image/run here),
|
||||
# so a batch is reproducible image-by-image from the base seed.
|
||||
# Mask to sd-cli's int64 range, NOT 53 bits: the request model and
|
||||
|
|
@ -593,7 +598,7 @@ class SdCppDiffusionBackend:
|
|||
images.append(im.copy())
|
||||
seeds.append(seed_i)
|
||||
if cancel.is_set():
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
# ``seeds`` is the per-image seed (each sd-cli run used seed+index), so
|
||||
# the route can persist the real seed for every image in the batch.
|
||||
return {
|
||||
|
|
@ -603,7 +608,7 @@ class SdCppDiffusionBackend:
|
|||
"repo_id": state.repo_id,
|
||||
}
|
||||
except SdCppCancelled as exc:
|
||||
raise RuntimeError("Diffusion generation was cancelled.") from exc
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc
|
||||
finally:
|
||||
self._gen = None
|
||||
with self._lock:
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ import time
|
|||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
from utils.process_lifetime import child_popen_kwargs
|
||||
from core.inference.sd_cpp_args import (
|
||||
SdCppGenParams,
|
||||
SdCppModelFiles,
|
||||
|
|
@ -352,6 +353,10 @@ class SdCppEngine:
|
|||
# Own session/process group so cancellation/timeout can kill the whole
|
||||
# tree, not just the parent (POSIX only; harmless flag elsewhere).
|
||||
start_new_session = (os.name == "posix"),
|
||||
# Bind the child to the parent's lifetime (Linux PR_SET_PDEATHSIG), so a
|
||||
# hard parent crash mid-generation can't orphan sd-cli holding VRAM/RAM --
|
||||
# matching every llama.cpp Popen site. Composes with start_new_session.
|
||||
**child_popen_kwargs(),
|
||||
)
|
||||
# Drain stdout on a reader thread so the timeout is enforced even when the
|
||||
# child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain
|
||||
|
|
|
|||
|
|
@ -10377,6 +10377,10 @@ async def generate_diffusion_image(
|
|||
):
|
||||
from core.inference import image_gallery
|
||||
from core.inference.diffusion_engine_router import get_active_diffusion_engine
|
||||
from core.inference.diffusion_families import (
|
||||
DIFFUSION_CANCELLED_MSG,
|
||||
DIFFUSION_NOT_LOADED_MSG,
|
||||
)
|
||||
|
||||
backend = get_active_diffusion_engine()
|
||||
try:
|
||||
|
|
@ -10414,11 +10418,15 @@ async def generate_diffusion_image(
|
|||
# doesn't support) — a 400 with the reason, not a generic 500.
|
||||
raise HTTPException(status_code = 400, detail = str(exc))
|
||||
except RuntimeError as exc:
|
||||
# Only "no model loaded" / cancelled are client-state (409). The native
|
||||
# sd.cpp engine also raises RuntimeError for execution failures (nonzero
|
||||
# exit, timeout, missing output), which are server errors (500).
|
||||
# Only "no model loaded" / user-cancelled are client-state (409); both engines
|
||||
# raise these two EXACT messages. The native sd.cpp engine also raises
|
||||
# RuntimeError for execution failures (nonzero exit, timeout, missing output)
|
||||
# whose text can embed the raw sd-cli tail (local paths / argv) -- those are
|
||||
# server errors (500) returned as a fixed literal, never echoed. Match the
|
||||
# sentinels exactly, not as a substring, so an sd-cli failure that merely
|
||||
# contains "cancelled" can't misroute to 409 and leak that output.
|
||||
msg = str(exc)
|
||||
if "No diffusion model is loaded" in msg or "cancelled" in msg.lower():
|
||||
if msg in (DIFFUSION_NOT_LOADED_MSG, DIFFUSION_CANCELLED_MSG):
|
||||
raise HTTPException(status_code = 409, detail = msg)
|
||||
logger.error("diffusion.generate_failed: %s", exc, exc_info = True)
|
||||
raise HTTPException(status_code = 500, detail = "Image generation failed.")
|
||||
|
|
|
|||
|
|
@ -472,7 +472,8 @@ def test_generate_img2img_unsupported_family_raises(fake_runtime, tmp_path, monk
|
|||
base_repo = "base/repo",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion.detect_family", lambda repo_id, override = None: plain
|
||||
"core.inference.diffusion.detect_family_for_pick",
|
||||
lambda repo_id, gguf_filename = None, override = None: plain,
|
||||
)
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
|
|
@ -987,6 +988,38 @@ def test_generate_without_load_raises(fake_runtime):
|
|||
backend.generate(prompt = "x")
|
||||
|
||||
|
||||
def test_failed_load_restores_backend_flags(fake_runtime, tmp_path, monkeypatch):
|
||||
# A failure AFTER apply_speed_optims (here an OOM in apply_memory_plan) must go
|
||||
# through the load's try/finally and restore the process-global TF32 / cudnn flags,
|
||||
# so a later `off` load is still bit-identical, and must not commit a partial state.
|
||||
# Regression: a refactor dropped this guard, leaking the flags on a failed load.
|
||||
(tmp_path / "model.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
|
||||
restored: list = []
|
||||
cleared: list = []
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion.restore_backend_flags", lambda snap: restored.append(snap)
|
||||
)
|
||||
monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: cleared.append(True))
|
||||
monkeypatch.setattr(
|
||||
"core.inference.diffusion.apply_memory_plan",
|
||||
lambda *a, **k: (_ for _ in ()).throw(RuntimeError("CUDA out of memory")),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match = "out of memory"):
|
||||
backend.load_pipeline(
|
||||
str(tmp_path),
|
||||
gguf_filename = "model.gguf",
|
||||
family_override = "z-image",
|
||||
base_repo = "base/repo",
|
||||
speed_mode = "max",
|
||||
)
|
||||
assert restored, "restore_backend_flags was not called on the failed-load path"
|
||||
assert cleared, "clear_gpu_cache was not called on the failed-load path (VRAM leak)"
|
||||
assert backend._state is None and backend.is_loaded is False
|
||||
|
||||
|
||||
def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch):
|
||||
from core.inference import diffusion
|
||||
from core.inference.diffusion_families import detect_family
|
||||
|
|
|
|||
|
|
@ -27,9 +27,8 @@ from core.inference.diffusion_memory import (
|
|||
DeviceMemory,
|
||||
MemoryPlan,
|
||||
apply_memory_plan,
|
||||
estimate_gguf_dense_mib,
|
||||
estimate_gguf_resident_mib,
|
||||
estimate_image_runtime_mib,
|
||||
infer_gguf_quant_label,
|
||||
normalize_memory_mode,
|
||||
plan_diffusion_memory,
|
||||
snapshot_device_memory,
|
||||
|
|
@ -70,29 +69,15 @@ def test_normalize_memory_mode_accepts_and_rejects():
|
|||
# ── filename / size estimates ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename,expected",
|
||||
[
|
||||
("z-image-turbo-Q4_K_M.gguf", "Q4_K_M"),
|
||||
("flux1-dev-Q8_0.gguf", "Q8_0"),
|
||||
("model-BF16.gguf", "BF16"),
|
||||
("qwen-image-IQ4_XS.gguf", "IQ4_XS"),
|
||||
("no-quant-here.gguf", None),
|
||||
(None, None),
|
||||
],
|
||||
)
|
||||
def test_infer_gguf_quant_label(filename, expected):
|
||||
assert infer_gguf_quant_label(filename) == expected
|
||||
|
||||
|
||||
def test_estimate_gguf_dense_mib_expansion():
|
||||
# 4-bit roughly quadruples once dequantised to bf16; F16 is already dense.
|
||||
assert estimate_gguf_dense_mib(1000, "Q4_K_M") == 4000
|
||||
assert estimate_gguf_dense_mib(1000, "Q8_0") == 2000
|
||||
assert estimate_gguf_dense_mib(1000, "BF16") == 1000
|
||||
assert estimate_gguf_dense_mib(None, "Q4_K_M") is None
|
||||
# Unknown quant falls back to the conservative 4-bit-ish factor.
|
||||
assert estimate_gguf_dense_mib(1000, None) == 4000
|
||||
def test_estimate_gguf_resident_mib_matches_packed_size():
|
||||
# GGUF weights stay packed (uint8) on-device; diffusers dequantises per-matmul
|
||||
# transiently, so the resident footprint ~= the on-disk size regardless of quant
|
||||
# level (measured on Z-Image-Turbo: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). A
|
||||
# small margin covers allocator overhead. The prior per-quant expansion over-
|
||||
# estimated (Q2 ~7.6x) and forced needless offload on a roomy card.
|
||||
assert estimate_gguf_resident_mib(1000) == 1050
|
||||
assert estimate_gguf_resident_mib(7220) == 7581
|
||||
assert estimate_gguf_resident_mib(None) is None
|
||||
|
||||
|
||||
def test_estimate_image_runtime_scales_with_pixels_and_family():
|
||||
|
|
|
|||
|
|
@ -303,6 +303,37 @@ def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch):
|
|||
assert "CUDA" not in resp.json()["detail"]
|
||||
|
||||
|
||||
def test_generate_execution_error_with_cancelled_substring_is_sanitized_500(client, monkeypatch):
|
||||
# A native sd-cli execution failure whose raw tail merely CONTAINS "cancelled"
|
||||
# must stay a sanitized 500, not misroute to 409 and echo that output (path/arg
|
||||
# leak). Regression: the handler matched "cancelled" as a substring.
|
||||
backend = diffusion_module.get_diffusion_backend()
|
||||
backend.loaded = True
|
||||
|
||||
def _fail(**kwargs):
|
||||
raise RuntimeError("sd-cli exited 1. Last output:\nop cancelled at /home/u/models/x.gguf")
|
||||
|
||||
monkeypatch.setattr(backend, "generate", _fail)
|
||||
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
||||
assert resp.status_code == 500
|
||||
assert resp.json()["detail"] == "Image generation failed."
|
||||
assert "cancelled" not in resp.json()["detail"] and "models" not in resp.json()["detail"]
|
||||
|
||||
|
||||
def test_generate_user_cancellation_returns_409(client, monkeypatch):
|
||||
# The exact cancellation sentinel both engines raise is client-state (409).
|
||||
backend = diffusion_module.get_diffusion_backend()
|
||||
backend.loaded = True
|
||||
|
||||
def _cancel(**kwargs):
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
|
||||
monkeypatch.setattr(backend, "generate", _cancel)
|
||||
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
||||
assert resp.status_code == 409
|
||||
assert resp.json()["detail"] == "Diffusion generation was cancelled."
|
||||
|
||||
|
||||
def test_load_unknown_family_returns_400(client, monkeypatch):
|
||||
def _raise(*a, **k):
|
||||
raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.")
|
||||
|
|
|
|||
|
|
@ -239,6 +239,25 @@ def test_speed_default_dense_falls_back_to_regional_compile(monkeypatch):
|
|||
assert called == {"compiled_dequant": 0}
|
||||
|
||||
|
||||
def test_offload_active_drops_fullgraph(monkeypatch):
|
||||
# Group/model/sequential offload installs a torch.compiler.disable'd onload hook;
|
||||
# compiling with fullgraph=True then crashes at the first denoise step. Same reason
|
||||
# as an active step cache -> fullgraph must drop to False when offload is planned.
|
||||
# (Dense model: on this branch GGUF `default` takes the compiled-dequant path.)
|
||||
_stub_torch(monkeypatch)
|
||||
pipe = _Pipe(with_compile = True)
|
||||
applied = apply_speed_optims(
|
||||
pipe,
|
||||
_target(),
|
||||
is_gguf = False,
|
||||
family = _family(),
|
||||
speed_mode = SPEED_DEFAULT,
|
||||
offload_active = True,
|
||||
)
|
||||
assert applied["compiled"] is True
|
||||
assert pipe.compile_kwargs["fullgraph"] is False
|
||||
|
||||
|
||||
def test_speed_default_gguf_compiles_only_dequant(monkeypatch):
|
||||
# GGUF `default` is the LIGHT path: compile ONLY the dequant op chain, NOT the
|
||||
# regional block compile.
|
||||
|
|
|
|||
|
|
@ -219,6 +219,17 @@ def test_begin_load_requires_gguf_filename():
|
|||
b.begin_load("unsloth/Z-Image-Turbo-GGUF")
|
||||
|
||||
|
||||
def test_begin_load_resolves_family_from_filename_only(monkeypatch):
|
||||
# A local .gguf pick whose family keyword lives only in the basename (parent dir
|
||||
# carries none) must resolve via the same filename fallback the route validated
|
||||
# with -- not dead-end with "Could not infer" on a native (no-GPU) host.
|
||||
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
||||
monkeypatch.setattr(b, "_run_load", lambda **kwargs: None) # skip the download thread
|
||||
b.begin_load("/models/gguf-store", gguf_filename = "Z-Image-Turbo-Q4_K_M.gguf")
|
||||
# Validation passed (no ValueError) and the family was inferred from the filename.
|
||||
assert b._loading is not None and b._loading.repo_id == "/models/gguf-store"
|
||||
|
||||
|
||||
def test_ensure_binary_returns_found(monkeypatch):
|
||||
monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli")
|
||||
assert ensure_sd_cpp_binary() == "/usr/bin/sd-cli"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue