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.

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/<branch> 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"

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:

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)

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():