github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Stability improvements to model offloading. Added D-OPSD bleed loss as well.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-29 11:23:47 -06:00
parent 64a20f51a6
commit 683fe8afc0
11 changed files with 390 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
toolkit/models/base_model.py
View File

@@ -185,6 +185,8 @@ class BaseModel:
self.supports_video_control_images = False
# D-OPSD: cache per-item teacher text embeds (item's own media as reference 1)
self.dopsd_self_ref = False
# weight of the normal-target loss added alongside the D-OPSD teacher loss
self.dopsd_bleed_strength = 1.0
# forces cache_tensors_to_disk on latent-caching datasets (BaseSDTrainProcess)
self.require_pixel_tensor_cache = False
# control images will come in as a list for encoding some things if true

65
toolkit/models/v2/_mixin.py
View File

@@ -218,6 +218,11 @@ class OstrisModelMixin:
config=config,
subfolder=subfolder,
use_comfy_weights=use_comfy_weights,
# comfy candidate ranking prefers the file whose shipped
# quantization matches the request (strip any ARA suffix);
# quantization itself happens in aitk_post_load below
qtype=(qtype or "").split("|", 1)[0] or None,
quantize_on_load=False,
**kwargs,
)
return model.aitk_post_load(
@@ -254,6 +259,53 @@ class OstrisModelMixin:
if qtype is not None and "|" in qtype:
qtype, ara_path = qtype.split("|", 1)
# pre-quantized checkpoint: keep the shipped quantization IFF it
# exactly matches the request. Anything else — no quantization
# requested (full finetuning), a different backend, or an ARA —
# restores/requantizes to what was asked for.
if getattr(self, "aitk_is_quantized", False):
from toolkit.util.ostris_quant import OstrisLinear
from toolkit.util.quantize import (
dequantize_ostris_to_linear,
get_qtype,
ostristype,
)
shipped = sorted(
{
getattr(m.ostris_quantizer, "qtype", None)
for m in self.modules()
if isinstance(m, OstrisLinear)
}
- {None}
)
matches = (
qtype is not None
and ara_path is None
and shipped == [qtype]
)
if matches:
status_fn(f"Checkpoint is pre-quantized ({qtype}); keeping it")
else:
target_is_ostris = qtype is not None and isinstance(
get_qtype(qtype), ostristype
)
if qtype is None or ara_path is not None or not target_is_ostris:
# full precision requested, an ARA (quantizes fresh from
# full weights), or a quanto/torchao backend that cannot
# re-quantize an OstrisLinear: dequantize first
status_fn(
f"Dequantizing shipped {'/'.join(shipped)} weights "
f"-> {qtype or 'full precision'}"
)
dequantize_ostris_to_linear(self)
else:
# ostris -> ostris: quantize_module re-quantizes each
# layer in place (dequant -> requant, one layer at a time)
status_fn(f"Requantizing {'/'.join(shipped)} -> {qtype}")
self.aitk_is_quantized = False
self.aitk_qtype = None
if qtype and not getattr(self, "aitk_is_quantized", False):
from toolkit.util.quantize import (
attach_ara_and_quantize,
@@ -308,8 +360,6 @@ class OstrisModelMixin:
)
self.aitk_is_quantized = True
self.aitk_qtype = qtype
elif qtype and getattr(self, "aitk_is_quantized", False):
status_fn("Checkpoint is pre-quantized; skipping quantization")
if offload and offload > 0:
from toolkit.memory_management import MemoryManager
@@ -337,6 +387,7 @@ class OstrisModelMixin:
config=None,
subfolder: Optional[str] = None,
use_comfy_weights: bool = True,
quantize_on_load: bool = True,
**kwargs,
):
"""Load a model universally from a given name or path.
@@ -370,7 +421,9 @@ class OstrisModelMixin:
and not name_or_path.endswith(".safetensors")
and not os.path.exists(name_or_path)
):
comfy_path = cls.resolve_comfy_weights(name_or_path, subfolder=subfolder)
comfy_path = cls.resolve_comfy_weights(
name_or_path, subfolder=subfolder, qtype=qtype
)
if comfy_path is not None:
if config_path is None:
# the standard repo supplies the config for the comfy file
@@ -399,7 +452,9 @@ class OstrisModelMixin:
name_or_path, subfolder=subfolder, dtype=dtype, **kwargs
)
if qtype is not None:
# quantize_on_load=False: qtype was only a comfy-candidate ranking
# hint (the .load() path quantizes in aitk_post_load instead)
if qtype is not None and quantize_on_load:
model.quantize_(
qtype, device=quantize_device, exclude=exclude_quant_modules
)
@@ -589,6 +644,7 @@ class OstrisModelMixin:
local_only: bool = False,
hf_token: Optional[str] = None,
status_fn: Optional[callable] = None,
qtype: Optional[str] = None,
) -> Optional[str]:
"""The comfy-format weight file replacing a standard ``name_or_path``,
or None when this class has none registered for it. Best-ranked local
@@ -618,6 +674,7 @@ class OstrisModelMixin:
hf_token=hf_token,
status_fn=status_fn,
local_only=local_only,
qtype=qtype,
)
# ------------------------------------------------------------------

61
toolkit/models/v2/resolver.py
View File

@@ -16,22 +16,41 @@ from typing import Callable, Iterable, Optional
from toolkit.paths import MODELS_PATH
def comfy_precision_rank(filename: str) -> int:
"""Load-preference rank for a comfy weight filename:
convrot8 (0) > float8 mixed (1) > float8 (2) > bf16 (3) > fp16 (4) >
anything else, e.g. nvfp4 or unmarked (5)."""
def comfy_precision_rank(filename: str, qtype: Optional[str] = None) -> int:
"""Load-preference rank for a comfy weight filename, given the REQUESTED
quantization. A file whose shipped quantization matches the request loads
with no work; anything else costs a dequantize/requantize pass, so:
- convrot* requested: convrot (0) > fp8 mixed (1) > fp8 (2) > bf16 (3) >
fp16 (4) > other (5)
- float8/qfloat8 requested: fp8 mixed (0) > fp8 (1) > bf16 (2) > fp16 (3)
> convrot (4) > other (5)
- nvfp4 requested: nvfp4 (0) > bf16 (1) > fp16 (2) > convrot (3) > fp8
(4) > other (5)
- no quantization requested (full precision) or any other fresh-quant
backend: bf16 (0) > fp16 (1) > convrot (2) > fp8 mixed (3) > fp8 (4) >
other (5) — clean weights beat paying a dequantize
"""
name = os.path.basename(filename).lower()
if "convrot" in name:
return 0
is_fp8 = "fp8" in name or "float8" in name or "e4m3" in name
if is_fp8 and "mixed" in name:
return 1
if is_fp8:
return 2
if "bf16" in name:
return 3
if "fp16" in name:
return 4
is_convrot = "convrot" in name
is_nvfp4 = "nvfp4" in name
is_fp8 = ("fp8" in name or "float8" in name or "e4m3" in name) and not is_nvfp4
is_fp8_mixed = is_fp8 and "mixed" in name
is_bf16 = "bf16" in name
is_fp16 = "fp16" in name and not is_fp8
qt = (qtype or "").lower()
if qt.startswith("convrot"):
order = [is_convrot, is_fp8_mixed, is_fp8, is_bf16, is_fp16]
elif "float8" in qt or qt == "qfloat8":
order = [is_fp8_mixed, is_fp8, is_bf16, is_fp16, is_convrot]
elif "nvfp4" in qt:
order = [is_nvfp4, is_bf16, is_fp16, is_convrot, is_fp8]
else:
order = [is_bf16, is_fp16, is_convrot, is_fp8_mixed, is_fp8]
for rank, flag in enumerate(order):
if flag:
return rank
return 5
@@ -51,15 +70,17 @@ def resolve_comfy_candidates(
hf_token: Optional[str] = None,
status_fn: Optional[Callable[[str], None]] = None,
local_only: bool = False,
qtype: Optional[str] = None,
) -> Optional[str]:
"""Pick the best comfy weight file among precision variants of one
component (repo-relative paths, ranked by comfy_precision_rank then list
order). The best-ranked LOCAL candidate wins; only when no candidate is
local is the best-ranked one downloaded to its comfy-layout location
under MODELS_PATH."""
component (repo-relative paths, ranked by comfy_precision_rank for the
requested qtype, then list order). The best-ranked LOCAL candidate wins;
only when no candidate is local is the best-ranked one downloaded to its
comfy-layout location under MODELS_PATH."""
candidates = list(candidates)
ordered = sorted(
candidates, key=lambda c: (comfy_precision_rank(c), candidates.index(c))
candidates,
key=lambda c: (comfy_precision_rank(c, qtype=qtype), candidates.index(c)),
)
for repo_rel in ordered:
found = resolve_comfy_file(