From a515fbfed38dfc6a3c495019b0acfc1aa0f79346 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 10:51:01 +0000 Subject: [PATCH] Version the conditioning cache key and reject non-finite flow_shift Two correctness fixes: - The cache keyed the checkpoint and its companion base by name only, so a Hub repo advancing to a new commit, or a local directory updated in place, kept returning embeddings from the previous text encoder. Pair both with a revision marker: the locally resolved commit sha for a Hub repo, config plus text-encoder file stats for a directory. Neither loads the encoders, so a warm run still keeps them off the GPU. - flow_shift only checked positivity, but JSON accepts 1e309, which floats to inf, and inf <= 0 is False while NaN fails every comparison. The sigma table then evaluates s * u / (1 + (s - 1) * u) as NaN, which poisons every sampled sigma and saves a corrupted adapter while progress looks normal. Require a finite value. --- .../core/inference/diffusion_cond_cache.py | 63 ++++++++++++++++++- .../core/training/diffusion_train_common.py | 11 +++- .../tests/test_diffusion_cond_cache.py | 30 +++++++++ .../test_diffusion_dit_trainer_flow_shift.py | 8 +++ 4 files changed, 109 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion_cond_cache.py b/studio/backend/core/inference/diffusion_cond_cache.py index 608169f4a5..86b0c6038d 100644 --- a/studio/backend/core/inference/diffusion_cond_cache.py +++ b/studio/backend/core/inference/diffusion_cond_cache.py @@ -149,9 +149,15 @@ def install( # base identity must key the cache too: the same checkpoint reloaded against a different # base would otherwise hit entries encoded by the previous base's encoder. Defaults to # the checkpoint itself (a full pipeline is its own base). + # The identifiers alone are not versions: both are paired with a revision marker so a + # Hub repo advancing to a new commit, or a local directory updated in place, misses + # instead of returning embeddings from the previous text encoder. + base_ref = base_repo if base_repo else repo_id load_fp = { "repo": str(repo_id), - "base": str(base_repo) if base_repo else str(repo_id), + "repo_rev": _source_revision(repo_id), + "base": str(base_ref), + "base_rev": _source_revision(base_ref), "dtype": str(dtype), "te_quant": str(te_quant) if te_quant is not None else "none", "diffusers": _diffusers_version(), @@ -224,3 +230,58 @@ def _diffusers_version() -> Optional[str]: return str(getattr(diffusers, "__version__", None)) except Exception: # noqa: BLE001 return None + + +def _source_revision(ref: Any) -> str: + """Revision/content marker for a checkpoint reference, resolved WITHOUT loading it. + + A bare repo id or directory path is not a version: a Hub repo that advances to a + new revision, or a local directory edited in place, keeps the same string while its + text encoders change, so cached embeddings from the old encoder would be reused + silently. This must stay cheap and must NOT touch the encoders themselves -- the + whole point of the cache is that a warm run never loads them -- so it reads the + local commit sha (no network) for a Hub repo and file stats for a local directory. + """ + try: + name = str(ref or "").strip() + if not name: + return "none" + if os.path.isdir(name): + # Local dir: stat the config + text-encoder files, which is what an in-place + # update touches. Bounded to the top level and text_encoder* subdirs. + parts: list[str] = [] + roots = [name] + with os.scandir(name) as it: + roots += [ + e.path for e in it + if e.is_dir() and e.name.startswith(("text_encoder", "tokenizer")) + ] + for root in roots: + with os.scandir(root) as it: + for e in it: + if not e.is_file(): + continue + st = e.stat() + parts.append(f"{os.path.relpath(e.path, name)}:{st.st_size}:{st.st_mtime_ns}") + digest = hashlib.sha256("|".join(sorted(parts)).encode()).hexdigest() + return f"dir-{digest[:16]}" + if "/" in name: + # Hub repo: the cache's refs/ file holds the resolved commit sha. + from huggingface_hub import constants # noqa: PLC0415 + + org, _, repo = name.partition("/") + base = os.path.join(constants.HF_HUB_CACHE, f"models--{org}--{repo}") + ref_file = os.path.join(base, "refs", "main") + if os.path.isfile(ref_file): + with open(ref_file, encoding = "utf-8") as fh: + sha = fh.read().strip() + if sha: + return f"rev-{sha[:16]}" + snaps = os.path.join(base, "snapshots") + if os.path.isdir(snaps): + names = sorted(os.listdir(snaps)) + if len(names) == 1: + return f"rev-{names[0][:16]}" + return "unresolved" + except Exception: # noqa: BLE001 — best-effort, never block a load + return "unresolved" diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 1ffe7c6c64..3280b24346 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -639,8 +639,15 @@ class DiffusionLoraConfig: ) from exc if not isinstance(flow_shift, str): flow_shift = float(flow_shift) - if flow_shift <= 0: - raise ValueError("flow_shift must be > 0 (1.0 disables the shift), or 'auto'") + # isfinite as well as positive: JSON accepts 1e309, which floats to inf, and + # inf <= 0 is False (NaN fails every comparison), so a positivity-only guard + # let it through to the sigma table, where s * u / (1 + (s - 1) * u) is NaN. + # That poisons every sampled sigma and the run saves a corrupted adapter while + # reporting normal progress. + if not math.isfinite(flow_shift) or flow_shift <= 0: + raise ValueError( + "flow_shift must be a finite number > 0 (1.0 disables the shift), or 'auto'" + ) try: cfg_dropout = float(self.cfg_dropout or 0.0) except (TypeError, ValueError) as exc: diff --git a/studio/backend/tests/test_diffusion_cond_cache.py b/studio/backend/tests/test_diffusion_cond_cache.py index a786939196..62e69041ac 100644 --- a/studio/backend/tests/test_diffusion_cond_cache.py +++ b/studio/backend/tests/test_diffusion_cond_cache.py @@ -140,6 +140,36 @@ def test_companion_base_keys_apart(cache_env): assert third.calls == 0 +def test_a_local_base_updated_in_place_keys_apart(cache_env, tmp_path): + # A directory path is not a version: editing the text encoder in place must MISS, or the + # run silently conditions on embeddings from the encoder that was there before. + base = tmp_path / "base" + (base / "text_encoder").mkdir(parents = True) + weights = base / "text_encoder" / "model.safetensors" + weights.write_bytes(b"v1") + first = _EncodePipe() + _install(first, repo_id = "org/model-GGUF", base_repo = str(base)) + first.encode_prompt("a sloth") + # Unchanged base -> warm hit (the cache still has to work). + warm = _EncodePipe() + _install(warm, repo_id = "org/model-GGUF", base_repo = str(base)) + warm.encode_prompt("a sloth") + assert (first.calls, warm.calls) == (1, 0) + # Same path, new contents -> re-encode. + weights.write_bytes(b"v2-different-length") + updated = _EncodePipe() + _install(updated, repo_id = "org/model-GGUF", base_repo = str(base)) + updated.encode_prompt("a sloth") + assert updated.calls == 1 + + +def test_source_revision_never_raises(): + # Best-effort by contract: a missing path, a bare name and junk all resolve to a marker + # instead of blocking the load. + for ref in (None, "", "no/such/repo-xyz", "/does/not/exist", 1234): + assert isinstance(cond_cache._source_revision(ref), str) + + def test_lora_attached_bypasses_the_cache(cache_env): pipe = _EncodePipe() _install(pipe) diff --git a/studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py b/studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py index a0339d3e5f..d21a0eed62 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py +++ b/studio/backend/tests/test_diffusion_dit_trainer_flow_shift.py @@ -91,6 +91,14 @@ def test_flow_shift_explicit_values_and_validation(): DiffusionLoraConfig( base_model = "b", data_dir = "d", output_dir = "o", flow_shift = "bogus" ).normalized() + # Non-finite must be rejected too: JSON accepts 1e309, which floats to inf, and + # inf <= 0 is False (NaN fails every comparison), so a positivity-only guard passed + # them through to the sigma table as NaN and the run saved a corrupted adapter. + for bad in (float("inf"), float("-inf"), float("nan"), 1e309): + with pytest.raises(ValueError, match = "flow_shift"): + DiffusionLoraConfig( + base_model = "b", data_dir = "d", output_dir = "o", flow_shift = bad + ).normalized() def test_cfg_dropout_and_weighting_scheme_validation():