Skip to content

C
coderai
Commit eeb3bba1 authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

video: param-weighted VRAM estimate, smarter auto offload, runtime accel LoRA 

VRAM estimation (manager.py):
- Weight the effective quant multiplier by REAL per-component parameter
  shares (new _component_param_shares scans safetensors by component folder)
  instead of a blind 70/30 split. Wan2.2 is 99.6% quantizable (two 14B
  experts + text encoder 4-bit, only the 0.13B VAE dense), so the old 0.475x
  multiplier inflated ~25.8 GB -> 42.7 GB and forced needless offload. Now
  ~0.28x -> ~25.8 GB. VAE forced dense (conv-only, bnb can't quantize).

Auto offload decision (video.py):
- 'auto': when peak footprint exceeds free VRAM, go straight to `model` CPU
  offload (active component on GPU, near full-GPU speed) — no full-GPU gamble,
  no slow balanced+disk path.
- 'auto-borderline' (new mode): same, except a marginal overshoot (<=3 GB)
  tries full-GPU first to keep both experts resident and use free VRAM,
  falling back to model offload on OOM.

Acceleration LoRA (acceleration.py + video.py):
- Keep the distill/Lightning LoRA as an ACTIVE RUNTIME ADAPTER instead of
  fusing. Fusing into CPU-offloaded bitsandbytes 4-bit weights triggers a
  dequant->merge->requant per Linear on the CPU — minutes/hours per expert,
  appearing to hang (high CPU, empty VRAM). Runtime adapters apply at forward
  time on the GPU at negligible cost and natively cover transformer_2.
- _sync_video_loras preserves the accel adapters across per-request LoRA swaps
  and re-includes them in every set_adapters; _unload_video_loras deletes only
  per-request adapters, keeping accel.

UI (models.html):
- Add "Auto borderline-aware" offload strategy option + updated hint.
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
parent f55f6578
Hide whitespace changes
Inline Side-by-side
Showing with 305 additions and 84 deletions
codai/admin/templates/models.html
View file @ eeb3bba1
......@@ -671,15 +671,17 @@ window.__DEFAULT_WHISPER_SERVER_PATH__ = {{ default_whisper_server_path|tojson }
<div class="form-row" style="margin:0">
<label class="form-label">Strategy</label>
<select id="cfg-offload-strategy" class="form-input">
<option value="auto">Auto (pick from free VRAM)</option>
<option value="auto">Auto (over-VRAM → straight to model offload)</option>
<option value="auto-borderline">Auto borderline-aware (try full-GPU if marginally over)</option>
<option value="none">None (GPU only)</option>
<option value="model">CPU offload (model — active component on GPU)</option>
<option value="group">Group offload + stream (block-level, prefetched)</option>
<option value="balanced">Balanced (fill GPU, spill to CPU → disk)</option>
<option value="model">CPU offload (model — module-by-module)</option>
<option value="sequential">CPU offload (sequential — most aggressive)</option>
<option value="cpu">CPU RAM (legacy)</option>
<option value="disk">Disk</option>
</select>
<span class="form-hint">Auto picks full-GPU when the weights fit, else Balanced. Pick <b>Balanced</b> + lower the GPU % below for a model that's just over VRAM; <b>model</b>/<b>sequential</b> keep less on GPU (slower, but the safest fit).</span>
<span class="form-hint"><b>Auto</b> goes straight to <b>model</b> offload when the model's peak exceeds free VRAM (no full-GPU gamble). <b>Auto borderline-aware</b> is the same, except when the model only marginally overshoots (≤3 GB) it tries full-GPU first to keep both experts resident and use the free VRAM, falling back to model offload on OOM. <b>model</b> keeps only the active component on GPU (best for two-expert Wan2.2 — near full-GPU speed). <b>group</b> streams blocks with prefetch (lowest VRAM, fast; incompatible with bitsandbytes 4-bit → falls back to sequential). <b>balanced</b> fills GPU to the % below then spills.</span>
</div>
<div class="form-row" style="margin:0">
<label class="form-label">Offload directory</label>
......
codai/api/video.py
View file @ eeb3bba1
......@@ -1134,6 +1134,45 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
_clear_mem()
return None
def _load_group_offload():
"""Block-level group offloading with CUDA-stream prefetch: keeps only a
few transformer blocks resident and overlaps the next block's CPU→GPU
copy with compute, so it's much faster than sequential at similar VRAM.
Applied per heavy component; the small VAE stays resident for decode.
Assigns the outer `pipe`. Raises on OOM / if nothing could be offloaded."""
nonlocal pipe
from diffusers.hooks import apply_group_offloading
_mem_snapshot("before group offload load")
print(" Video load strategy: group offload + stream "
"(block-level GPU↔CPU with prefetch)")
pipe = PClass.from_pretrained(**_with_quant(dict(
pretrained_model_name_or_path=model_name,
torch_dtype=torch_dtype, low_cpu_mem_usage=True)))
_onload, _offl = torch.device('cuda'), torch.device('cpu')
_applied = []
for _cn in ('transformer', 'transformer_2', 'text_encoder', 'text_encoder_2'):
_comp = getattr(pipe, _cn, None)
if _comp is None or not hasattr(_comp, 'parameters'):
continue
try:
apply_group_offloading(
_comp, onload_device=_onload, offload_device=_offl,
offload_type='block_level', num_blocks_per_group=1,
use_stream=True, low_cpu_mem_usage=True)
_applied.append(_cn)
except Exception as _ge:
print(f" [group-offload] {_cn} not offloaded ({_ge})")
if not _applied:
raise RuntimeError("group offloading applied to no components "
"(incompatible — e.g. bitsandbytes-quantized weights)")
try:
if getattr(pipe, 'vae', None) is not None:
pipe.vae.to(_onload)
except Exception:
pass
_report_loaded(pipe, f"group offload + stream ({','.join(_applied)})")
return pipe
def _load_sequential():
"""Most aggressive fit: stream each submodule GPU↔CPU during the
forward pass (slowest, lowest VRAM). Assigns the outer `pipe`. Raises
......@@ -1149,6 +1188,26 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
_report_loaded(pipe, "sequential CPU offload")
return pipe
def _try_group_then_sequential():
"""Group offload (stream), then sequential CPU offload on OOM/incompat.
Returns the pipe, or None (caller falls through to disk offload)."""
try:
return _load_group_offload()
except (RuntimeError, MemoryError) as _e:
if not _is_oom(_e) and 'no components' not in str(_e):
raise
print(f" Video: group offload unavailable/OOM ({str(_e).splitlines()[0]}) "
f"— trying sequential CPU offload…")
_clear_mem()
try:
return _load_sequential()
except (RuntimeError, MemoryError) as _e:
if not _is_oom(_e):
raise
print(f" Video: sequential CPU offload OOM ({_e}) — trying disk offload…")
_clear_mem()
return None
def _try_balanced_then_sequential():
"""Balanced chain (configured% → 60 → 40), then sequential CPU offload
if all balanced steps OOM. Returns the pipe, or None if even sequential
......@@ -1174,7 +1233,7 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
# Even sequential OOM'd → continue to the disk-offload attempts.
# ── Attempt 0: full GPU ──────────────────────────────────────────────
if offload not in ('model', 'sequential', 'disk', 'balanced'):
if offload not in ('model', 'group', 'sequential', 'disk', 'balanced'):
_mem_snapshot("before full-GPU load")
_q = " + quantized" if _quant_config is not None else ""
print(f" Video load strategy: full GPU ({torch_dtype}{_q})")
......@@ -1198,19 +1257,15 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
except (RuntimeError, MemoryError) as e:
if not _is_oom(e):
raise
print(f" Video: full-GPU OOM ({e}) — falling back to balanced "
f"GPU+CPU (starting at {_gpu_pct:.0f}% GPU)…")
_clear_mem()
# Graceful degrade: balanced at the configured %, then 60%, then
# 40%, then sequential CPU offload, before the slower disk paths.
pipe = _try_balanced_then_sequential()
if pipe is not None:
return pipe
print(" Video: balanced + sequential all OOM — trying disk offload…")
# Auto degrade: prefer model CPU offload (only the active expert
# resident — near full-GPU speed and fits this two-expert model),
# then group offload + stream, then sequential, then disk. Balanced
# is reserved for an explicit offload_strategy=balanced.
print(f" Video: full-GPU OOM ({e}) — trying model CPU offload…")
_clear_mem()
# ── Attempt 1: model CPU offload ─────────────────────────────────────
if offload not in ('sequential', 'disk'):
if offload not in ('group', 'sequential', 'disk'):
_mem_snapshot("before model-CPU-offload load")
print(f" Video load strategy: model CPU offload"
f" (each module GPU↔CPU during forward pass)")
......@@ -1226,9 +1281,28 @@ def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str =
if not _is_oom(e):
raise
print(f" Video: model CPU offload OOM ({e})"
f" — trying GPU+CPU+disk offload…")
f" — trying group offload + stream…")
_clear_mem()
# ── Attempt 1.5: group offload + stream → sequential ─────────────────
# Block-level offloading with CUDA-stream prefetch: lowest VRAM after
# sequential but hides the transfer behind compute. Runs for auto (after
# model offload) and for an explicit offload_strategy=group; on OOM or
# incompatibility (e.g. bitsandbytes weights) it drops to sequential.
if offload == 'sequential':
try:
return _load_sequential()
except (RuntimeError, MemoryError) as e:
if not _is_oom(e):
raise
print(f" Video: sequential CPU offload OOM ({e}) — trying disk offload…")
_clear_mem()
elif offload != 'disk':
pipe = _try_group_then_sequential()
if pipe is not None:
return pipe
_clear_mem()
# ── Attempt 2: GPU + CPU + disk offload via device_map='auto' ──────────
os.makedirs(_offload_dir, exist_ok=True)
_cpu_gb = min(32, max(2, int(_psutil.virtual_memory().available * 0.50 / 1e9)))
......@@ -1417,9 +1491,26 @@ def _apply_character_refs(kw: dict, character_references: List[str], strength: f
def _unload_video_loras(pipe):
"""Remove any LoRA adapters so a cached pipeline is clean for the next request."""
"""Remove per-request LoRA adapters so a cached pipeline is clean for the next
request — but PRESERVE the acceleration/distill adapters (Lightning), which are
kept as permanent runtime adapters (never fused) and must survive every swap.
"""
accel = set(getattr(pipe, '_coderai_accel_adapters', []) or [])
try:
if hasattr(pipe, 'unload_lora_weights'):
if accel:
# Delete only the per-request adapters; keep the accel ones, then
# re-activate accel alone so the pipe is in a clean accel-only state.
present = _present_adapters(pipe)
to_delete = [a for a in present if a not in accel]
if to_delete and hasattr(pipe, 'delete_adapters'):
pipe.delete_adapters(to_delete)
try:
pipe.set_adapters(
list(accel),
[getattr(pipe, '_coderai_accel_weight', 1.0)] * len(accel))
except Exception:
pass
elif hasattr(pipe, 'unload_lora_weights'):
pipe.unload_lora_weights()
except Exception as e:
print(f" [video][lora] unload failed: {e}")
......@@ -1530,10 +1621,19 @@ def _sync_video_loras(pipe, loras) -> None:
_unload_video_loras(pipe)
pipe._coderai_active_loras = desired # _unload reset it; restore for dedup
return
# Always re-include the acceleration/distill adapters (kept as runtime
# adapters, never fused) alongside the per-request LoRAs — otherwise this
# set_adapters would deactivate the distill LoRA and the 4-step preset would
# collapse the clip to a solid colour.
_accel = list(getattr(pipe, '_coderai_accel_adapters', []) or [])
_accel_w = getattr(pipe, '_coderai_accel_weight', 1.0)
_names = _accel + [n for n, _ in loaded]
_weights = [_accel_w] * len(_accel) + [w for _, w in loaded]
try:
pipe.set_adapters([n for n, _ in loaded], [w for _, w in loaded])
pipe.set_adapters(_names, _weights)
print(f" [video][lora] applied: {[n for n, _ in loaded]} "
f"weights={[w for _, w in loaded]}")
f"weights={[w for _, w in loaded]}"
+ (f" (+ accel {_accel})" if _accel else ""))
except Exception as e:
print(f" [video][lora] could not activate LoRA weights: {e}")
_unload_video_loras(pipe)
......@@ -2300,19 +2400,27 @@ async def video_generations(request: VideoGenerationRequest,
if pipe is None:
_offload = _model_cfg.get('offload_strategy') or None
# 'auto' (the default) means "let coderai pick from available VRAM" — it is
# NOT a diffusers strategy, and passing it through lands on the full-GPU
# path that then disk-thrashes. Normalise it to None so the VRAM check
# below decides between full-GPU and balanced GPU+CPU.
if _offload == 'auto':
# Two auto modes (neither is a diffusers strategy — both are normalised to
# None so the VRAM check below picks the concrete strategy):
# * 'auto' — when the peak footprint exceeds free VRAM, go STRAIGHT to
# `model` CPU offload (no full-GPU gamble).
# * 'auto-borderline' — same, EXCEPT when the model only marginally
# overshoots (within _BORDERLINE_GB), try full-GPU first to keep both
# experts resident and actually use the free VRAM; it falls back to
# model offload on OOM. Best when the estimate is conservative and the
# model very likely fits.
_auto_borderline = (_offload == 'auto-borderline')
if _offload in ('auto', 'auto-borderline'):
_offload = None
# Auto-select "balanced" strategy when the model (including runtime
# reserve: KV/activation spike, VAE decode) exceeds available VRAM even
# after eviction. Going straight to "balanced" (GPU-first + CPU spill)
# avoids the expensive OOM → free → reload cycle that wastes ~1 hr of
# shard reloading only to end up at the same place. The GPU cap is 80%
# of free VRAM (or the per-model balanced_gpu_percent if configured) so
# we leave breathing room for activations and the decode spike.
# Auto-select the strategy from how the model's PEAK footprint compares
# to free VRAM after eviction:
# * peak fits → full GPU (fastest).
# * peak doesn't fit → `model` CPU offload directly (active component on
# GPU, inactive ones on CPU — near full-GPU speed). We do NOT gamble on
# a full-GPU load that OOMs at the decode/activation peak only to fall
# back anyway, and we do NOT jump to the slow balanced+disk path.
# `_load_video_pipeline`'s ladder still escalates model → group →
# sequential → disk if even model offload can't fit (truly huge model).
if _offload is None:
try:
import torch as _t
......@@ -2322,27 +2430,29 @@ async def video_generations(request: VideoGenerationRequest,
# `_get_model_used_vram_gb` is the *measured total* footprint —
# it already includes the runtime/activation reserve AND the
# fused acceleration LoRA (it's measured after fusion). So do
# NOT re-add those (that over-counts and wrongly forces the
# slow balanced+disk path). Per-request LoRAs are extra.
# NOT re-add those. Per-request LoRAs are extra.
_base_gb = multi_model_manager._get_model_used_vram_gb(
model_key, model_name)
# full-GPU only needs the WEIGHTS to fit at load time (the
# bundled ~runtime reserve is a gen-time allowance, and the
# full-GPU path has its own OOM→offload fallback). Keep a
# headroom margin so a model that *marginally* fits uses the
# much faster full-GPU strategy rather than balanced+disk.
_need_gb = _base_gb + _lora_extra_gb
_margin = 2.5 # ≈ the bundled runtime reserve
if _base_gb > 0 and _free_gb < (_need_gb - _margin):
_gpu_pct = float(_model_cfg.get('balanced_gpu_percent') or 80)
print(f" VRAM well short for full-GPU load "
f"({_need_gb:.1f} GB measured need + LoRA; "
f"{_free_gb:.1f} GB free) — auto-selecting balanced "
f"strategy ({_gpu_pct:.0f}% GPU + CPU spill)")
_offload = 'balanced'
else:
_BORDERLINE_GB = 3.0 # how far over free VRAM still counts as "borderline"
_over_gb = _need_gb - _free_gb
if _base_gb > 0 and _free_gb < _need_gb:
if _auto_borderline and _over_gb <= _BORDERLINE_GB:
# Marginal overshoot + the estimate is conservative → try
# full-GPU first (keeps both experts resident, uses the
# free VRAM). Leaving _offload=None routes to the loader's
# full-GPU attempt, which falls back to model offload on OOM.
print(f" Peak VRAM need {_need_gb:.1f} GB marginally over "
f"{_free_gb:.1f} GB free (+{_over_gb:.1f} GB, borderline) "
f"— trying full GPU first (falls back to model offload on OOM)")
else:
print(f" Peak VRAM need {_need_gb:.1f} GB > {_free_gb:.1f} GB "
f"free — auto-selecting `model` CPU offload "
f"(active component on GPU, near full-GPU speed)")
_offload = 'model'
elif _base_gb > 0:
print(f" Full-GPU load looks viable "
f"({_need_gb:.1f} GB measured need, {_free_gb:.1f} GB "
f"({_need_gb:.1f} GB peak need, {_free_gb:.1f} GB "
f"free) — using full GPU (it falls back to offload on OOM)")
except Exception:
pass
......
codai/models/acceleration.py
View file @ eeb3bba1
......@@ -325,32 +325,31 @@ def apply_accel_to_pipeline(pipe, accel: Optional[dict]) -> None:
raise RuntimeError("no distill adapter registered on the pipeline")
try:
pipe.set_adapters(loaded_adapters, [weight] * len(loaded_adapters))
except Exception:
pass
# Bake them in, then drop the adapter handles so per-request LoRAs are clean.
# CRITICAL: diffusers' Wan fuse_lora defaults to components=["transformer"],
# so without naming transformer_2 the low-noise expert's distill adapter is
# never fused — and the subsequent unload strips it off, leaving that expert
# undistilled. At 4 steps that collapses the clip to a solid colour. Fuse
# BOTH experts explicitly.
_fuse_components = ["transformer"]
if has_t2:
_fuse_components.append("transformer_2")
try:
pipe.fuse_lora(components=_fuse_components, lora_scale=weight)
except TypeError:
# Older diffusers without the `components` kwarg — best effort.
pipe.fuse_lora(lora_scale=weight)
try:
pipe.unload_lora_weights()
except Exception:
pass
except Exception as e:
log.warning("[accel] could not activate distill adapters %s: %s",
loaded_adapters, e)
# Keep the distill LoRA(s) as ACTIVE RUNTIME ADAPTERS rather than fusing.
#
# Fusing merges the adapter into the base weights. For a bitsandbytes 4-bit
# model under CPU/disk offload the weights live on the CPU, so fuse_lora has
# to dequantize → merge → requantize every Linear on the CPU — minutes-to-
# hours per 14B expert, which looks exactly like a hang (high CPU, empty
# VRAM, no progress). Runtime adapters instead apply at forward time on
# whichever device the module is currently on (the GPU during the pass), at
# negligible cost, and `set_adapters` natively covers transformer_2 — so the
# low-noise expert is distilled too (the old fuse path needed an explicit
# components=[...,"transformer_2"] or it collapsed to a solid colour).
#
# `_sync_video_loras` preserves these accel adapters across per-request LoRA
# swaps and re-includes them in every set_adapters call.
pipe._coderai_accel_adapters = list(loaded_adapters)
pipe._coderai_accel_weight = weight
try:
pipe._coderai_accel_fused = True
pipe._coderai_accel_fused = True # accel is effective (distilled preset applies)
except Exception:
pass
log.info("[accel] fused distillation LoRA(s) %s (weight=%s) into %s%s",
loaded_adapters, weight, type(pipe).__name__,
log.info("[accel] activated distillation LoRA(s) %s (weight=%s, runtime adapters) "
"on %s%s", loaded_adapters, weight, type(pipe).__name__,
" (both experts)" if len(loaded_adapters) > 1 else "")
except Exception as e:
log.warning("[accel] failed to fuse acceleration LoRA (high=%s low=%s): %s "
......
codai/models/manager.py
View file @ eeb3bba1
......@@ -1892,26 +1892,62 @@ class MultiModelManager:
return 1.0
def _effective_quant_multiplier(self, cfg: dict,
load_in_4bit: bool, load_in_8bit: bool) -> float:
load_in_4bit: bool, load_in_8bit: bool,
model_key: str = None,
resolved_name: str = None) -> float:
"""Fraction of full-precision weight size that stays resident.
Per-component quantization wins when present: the big weight-bearing
components (transformer(s), unet, text encoders) dominate, so we assume
the quantized components carry ~80 % of the weight and the rest (VAE,
etc.) stays dense. Falls back to the global 4/8-bit flags.
When per-component quantization is configured AND the model's real
per-component parameter shares can be scanned from disk, weight each
component's quant divisor by its ACTUAL share of the parameters — so a
pipeline that is e.g. 99 % 4-bit (Wan2.2: two 14B experts + text encoder
quantized, only a 0.13B VAE dense) is estimated at ~0.28× rather than the
blunt 0.475× a fixed 70/30 split would give (which inflated 25.8 GB →
42.7 GB and forced needless offload). Falls back to a fixed-share
heuristic, then to the global 4/8-bit flags.
"""
comp_q = cfg.get("component_quantization") or {}
def _comp_divisor(name: str) -> float:
# VAEs are conv-only → bitsandbytes/quanto can't quantize them; they
# stay dense regardless of what the config asks for.
low = str(name).lower()
if low == 'vae' or low.endswith('_vae') or low.startswith('vae'):
return 1.0
if name in comp_q:
return self._quant_divisor(comp_q[name])
if load_in_4bit:
return 4.0
if load_in_8bit:
return 2.0
return 1.0
# --- Preferred: weight by real per-component parameter shares ---
shares = {}
if comp_q and model_key:
try:
shares = self._component_param_shares(model_key, resolved_name, cfg)
except Exception:
shares = {}
if shares:
total = sum(shares.values())
if total > 0:
mult = 0.0
for comp, numel in shares.items():
div = _comp_divisor(comp)
eff = (1.0 / div) if div > 1.0 else 1.0
mult += (numel / total) * eff
# bitsandbytes keeps fp32 absmax + scale tensors and fp16 compute
# buffers alongside the packed 4-bit weights (~+12 %).
return mult * 1.12
divisors = [self._quant_divisor(v) for v in comp_q.values()]
divisors = [d for d in divisors if d > 1.0]
if divisors:
avg_div = sum(divisors) / len(divisors)
# Lean slightly HIGH: assume ~70 % of the footprint is shrunk by
# quantization and ~30 % (VAE, embeddings, fp16 compute buffers,
# bitsandbytes scale/absmax tensors) stays dense. Better to slightly
# over-estimate and evict enough than to under-estimate and OOM —
# but not so high it exceeds the card and forces needless offload.
# (e.g. all-4bit → ~0.475× raw, ×1.15 overhead ≈ 0.55× → ~+10% over
# the measured resident size.)
# No per-component scan available: assume ~70 % of the footprint is
# shrunk by quantization and ~30 % (VAE, embeddings, fp16 compute
# buffers, bitsandbytes scale/absmax tensors) stays dense.
quantized_share = 0.7
return quantized_share / avg_div + (1.0 - quantized_share)
if load_in_4bit:
......@@ -2042,6 +2078,80 @@ class MultiModelManager:
self._storage_dtype_cache[ck] = result
return result
# component (top-level subfolder) -> total params, per model. Cached.
_component_shares_cache: Dict[str, Dict[str, int]] = {}
def _component_param_shares(self, model_key: str, resolved_name: str,
cfg: dict) -> Dict[str, int]:
"""Map each pipeline component → its parameter count (numel).
Groups every ``.safetensors`` file by its TOP-LEVEL subfolder name
(``transformer``, ``transformer_2``, ``text_encoder``, ``vae`` …) — i.e.
the diffusers component layout — and sums the parameters per component.
Used to weight the effective quantization multiplier by each component's
real share of the model, so a pipeline that is 99 % quantized isn't
treated as if 30 % stays dense. Returns {} when it can't be determined.
"""
ck = resolved_name or model_key
cached = self._component_shares_cache.get(ck)
if cached is not None:
return cached
import os
shares: Dict[str, int] = {}
candidates = []
for v in (resolved_name, cfg.get('path'), cfg.get('model_path'), cfg.get('model')):
if v and isinstance(v, str):
candidates.append(v)
def _record(root_dir: str, file_path: str):
rel = os.path.relpath(file_path, root_dir)
comp = rel.split(os.sep)[0]
# Files at the pipeline root (no subfolder) → bucket as 'root'.
if comp.endswith('.safetensors'):
comp = 'root'
numel, _ = self._safetensors_numel_bytes(os.path.realpath(file_path))
if numel > 0:
shares[comp] = shares.get(comp, 0) + numel
try:
scanned = False
for c in candidates:
if os.path.isdir(c):
for r, _, files in os.walk(c):
for fn in files:
if fn.endswith('.safetensors'):
_record(c, os.path.join(r, fn))
scanned = True
break
if not scanned:
from huggingface_hub import scan_cache_dir
from codai.models.cache import get_all_cache_dirs, is_huggingface_model_id
hf_dir = get_all_cache_dirs().get("huggingface")
for c in candidates:
repo_id = c.split(":", 1)[1] if ":" in c else c
if not (hf_dir and is_huggingface_model_id(repo_id)):
continue
info = scan_cache_dir(hf_dir)
for repo in info.repos:
if repo.repo_id != repo_id:
continue
revs = sorted(repo.revisions,
key=lambda rv: rv.last_modified, reverse=True)
if revs:
snap = str(revs[0].snapshot_path)
for fobj in revs[0].files:
fp = str(fobj.file_path)
if fp.endswith('.safetensors'):
_record(snap, fp)
if shares:
break
except Exception:
shares = {}
self._component_shares_cache[ck] = shares
return shares
@staticmethod
def _load_bytes_per_elem(cfg: dict) -> float:
"""Bytes/elem for the configured LOAD precision (default fp16 = 2)."""
......@@ -2189,7 +2299,7 @@ class MultiModelManager:
# diffusion pipeline whose transformer/text_encoder are 4-bit isn't
# wildly over-estimated at full precision.
quant_mult = self._effective_quant_multiplier(
cfg, load_in_4bit, load_in_8bit)
cfg, load_in_4bit, load_in_8bit, model_key, resolved_name)
# Precision normalization: used_vram_gb / disk-scan baselines are STORAGE
# sizes at the on-disk dtype (e.g. Wan2.2 ships fp32 → 4 bytes/elem), but
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment