Split fetch/apply; drop Xet CAS store after loading PinkCherry

pinkcherry_loader.py

@@ -4,21 +4,25 @@ The PinkCherry checkpoint is a ComfyUI-layout state dict: fused `qkv_proj` atten
 swiglu, `blocks.N` / `token_refiner.blocks.N` naming. The diffusers `MiniMaxH3Transformer3DModel`
 expects SD3-style names (`to_q/to_k/to_v`, `ff.net.0.proj` / `ff.net.2`, `transformer_blocks.N`,
 `refiner_blocks`). This module remaps one to the other and verifies full coverage.
+
+Storage strategy: the checkpoint is downloaded first (Xet, fast), loaded into RAM, and its on-disk
+files — the materialized safetensors plus the Xet CAS blob — are deleted before the base model is
+fetched, so peak disk usage never approaches the 150 GB ZeroGPU limit.
 """
 
 from __future__ import annotations
 
 import os
 import re
-import tempfile
-import urllib.request
+import shutil
+from pathlib import Path
 
 import safetensors.torch
 import torch
+from huggingface_hub import hf_hub_download
 
 PINKCHERRY_REPO = "SexGod1979/PinkCherry_MiniMax-H3"
 PINKCHERRY_FILE = "beta-0.6-fl2va/PinkCherry_fl2va_MiniMax_H3_bf16_beta-0.6.safetensors"
-PINKCHERRY_URL = f"https://huggingface.co/{PINKCHERRY_REPO}/resolve/main/{PINKCHERRY_FILE}"
 
 # ComfyUI -> diffusers, per attention/block module. `qkv_proj` is fused Q,K,V in that order and is
 # handled separately by the splitter.
@@ -82,89 +86,84 @@ def _remap_key(key: str) -> str | None:
     return None
 
 
-def _download_pinkcherry() -> str:
-    """Fetch the checkpoint over plain HTTP into a temp file.
-
-    Deliberately bypasses the HF cache: `hf_hub_download` on a Xet-backed repo leaves a second compressed CAS
-    copy on disk, and base model (77.3 GB) + PinkCherry (61.7 GB) + CAS blob (~35 GB) + venv busts the 150 GB
-    ZeroGPU storage limit mid-boot. A plain fetch materializes the file exactly once; the caller deletes it after
-    loading.
-    """
-    fd, path = tempfile.mkstemp(suffix=".safetensors", prefix="pinkcherry-")
-    os.close(fd)
-    print(f"[pinkcherry] downloading {PINKCHERRY_URL} -> {path}", flush=True)
-    req = urllib.request.Request(PINKCHERRY_URL, headers={"User-Agent": "pinkcherry-loader/1.0"})
-    with urllib.request.urlopen(req, timeout=3600) as resp, open(path, "wb") as out:
-        total = int(resp.headers.get("Content-Length") or 0)
-        done = 0
-        while True:
-            chunk = resp.read(8 * 1024 * 1024)
-            if not chunk:
-                break
-            out.write(chunk)
-            done += len(chunk)
-            if total:
-                print(f"[pinkcherry] {done / 1024**3:.1f} / {total / 1024**3:.1f} GiB", flush=True)
-    print(f"[pinkcherry] downloaded {done / 1024**3:.1f} GiB", flush=True)
-    return path
-
-
-def load_pinkcherry_state_dict(transformer) -> None:
-    """Download, remap, and load the PinkCherry weights into `transformer`; the temp file is removed on exit."""
-    path = _download_pinkcherry()
+def _remap(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    remapped: dict[str, torch.Tensor] = {}
+    for key, tensor in state.items():
+        if key == "__metadata__":
+            continue
+        if key == "rope.inv_freq":
+            # ComfyUI stores the rotary table; the diffusers model computes it internally.
+            continue
+        if key.endswith("attn.qkv_proj.weight"):
+            inner = tensor.shape[0] // 3
+            q, k, v = tensor[:inner], tensor[inner : 2 * inner], tensor[2 * inner :]
+            base = key[: -len("attn.qkv_proj.weight")]
+            base = _BLOCK_RE.sub(r"transformer_blocks.\1.", base)
+            base = _REFINER_RE.sub(r"token_refiner.refiner_blocks.\1.", base)
+            remapped[base + "attn.to_q.weight"] = q
+            remapped[base + "attn.to_k.weight"] = k
+            remapped[base + "attn.to_v.weight"] = v
+            continue
+        if key.endswith("mlp.fc1.weight"):
+            # ComfyUI's swiglu stores `[gate; value]`; diffusers' SwiGLU reads `[value; gate]`,
+            # so the two halves swap places (same transform as the official conversion script).
+            gate, value = tensor.chunk(2, dim=0)
+            tensor = torch.cat([value, gate], dim=0).contiguous()
+        new = _remap_key(key)
+        if new is None:
+            raise KeyError(f"no mapping for PinkCherry key {key!r}")
+        remapped[new] = tensor
+    return remapped
+
+
+def _drop_xet_cas() -> None:
+    """Remove the Xet content-addressed blob store so its compressed copy stops counting against storage."""
+    from huggingface_hub.constants import HF_HUB_CACHE
+
+    for candidate in (Path(os.environ.get("HF_XET_CACHE", "")), Path(HF_HUB_CACHE) / "xet"):
+        if candidate.is_dir():
+            try:
+                shutil.rmtree(candidate)
+                print(f"[pinkcherry] removed Xet CAS store {candidate}", flush=True)
+            except OSError as error:
+                print(f"[pinkcherry] could not remove {candidate}: {error}", flush=True)
+
+
+def fetch_pinkcherry_state_dict() -> dict[str, torch.Tensor]:
+    """Download the checkpoint (Xet, fast), load it into RAM, free its disk footprint, return the remapped dict."""
+    path = hf_hub_download(PINKCHERRY_REPO, PINKCHERRY_FILE)
+    print(f"[pinkcherry] loading {path}", flush=True)
+    state = safetensors.torch.load_file(path)
+    remapped = _remap(state)
+    del state
     try:
-        print(f"[pinkcherry] loading {path}", flush=True)
-        state = safetensors.torch.load_file(path)
-
-        remapped: dict[str, torch.Tensor] = {}
-        for key, tensor in state.items():
-            if key == "__metadata__":
-                continue
-            if key == "rope.inv_freq":
-                # ComfyUI stores the rotary table; the diffusers model computes it internally.
-                continue
-            if key.endswith("attn.qkv_proj.weight"):
-                inner = tensor.shape[0] // 3
-                q, k, v = tensor[:inner], tensor[inner : 2 * inner], tensor[2 * inner :]
-                base = key[: -len("attn.qkv_proj.weight")]
-                base = _BLOCK_RE.sub(r"transformer_blocks.\1.", base)
-                base = _REFINER_RE.sub(r"token_refiner.refiner_blocks.\1.", base)
-                remapped[base + "attn.to_q.weight"] = q
-                remapped[base + "attn.to_k.weight"] = k
-                remapped[base + "attn.to_v.weight"] = v
-                continue
-            if key.endswith("mlp.fc1.weight"):
-                # ComfyUI's swiglu stores `[gate; value]`; diffusers' SwiGLU reads `[value; gate]`,
-                # so the two halves swap places (same transform as the official conversion script).
-                gate, value = tensor.chunk(2, dim=0)
-                tensor = torch.cat([value, gate], dim=0).contiguous()
-            new = _remap_key(key)
-            if new is None:
-                raise KeyError(f"no mapping for PinkCherry key {key!r}")
-            remapped[new] = tensor
-
-        expected = set(transformer.state_dict().keys())
-        missing = expected - set(remapped.keys())
-        unexpected = set(remapped.keys()) - expected
-        if missing:
-            raise KeyError(f"remap left {len(missing)} model keys uncovered: {sorted(missing)[:10]} ...")
-        if unexpected:
-            raise KeyError(f"remap produced {len(unexpected)} keys the model does not have: {sorted(unexpected)[:10]} ...")
-        for key in expected:
-            if remapped[key].shape != transformer.state_dict()[key].shape:
-                raise ValueError(
-                    f"shape mismatch on {key}: PinkCherry {tuple(remapped[key].shape)} vs model {tuple(transformer.state_dict()[key].shape)}"
-                )
-
-        missing, unexpected = transformer.load_state_dict(remapped, strict=False)
-        if missing:
-            raise KeyError(f"load left {len(missing)} keys missing: {sorted(missing)[:10]} ...")
-        if unexpected:
-            raise KeyError(f"load got {len(unexpected)} unexpected keys: {sorted(unexpected)[:10]} ...")
-        print(f"[pinkcherry] loaded {len(remapped)} tensors into the transformer", flush=True)
-    finally:
-        try:
-            os.unlink(path)
-            print(f"[pinkcherry] removed temp file {path}", flush=True)
-        except OSError:
-            pass
+        Path(path).unlink()
+        print(f"[pinkcherry] removed materialized checkpoint {path}", flush=True)
+    except OSError as error:
+        print(f"[pinkcherry] could not unlink {path}: {error}", flush=True)
+    _drop_xet_cas()
+    print(f"[pinkcherry] remapped {len(remapped)} tensors, disk footprint freed", flush=True)
+    return remapped
+
+
+def apply_pinkcherry_state_dict(transformer, remapped: dict[str, torch.Tensor]) -> None:
+    """Validate and load the remapped PinkCherry weights into the (base-loaded) transformer."""
+    expected = set(transformer.state_dict().keys())
+    missing = expected - set(remapped.keys())
+    unexpected = set(remapped.keys()) - expected
+    if missing:
+        raise KeyError(f"remap left {len(missing)} model keys uncovered: {sorted(missing)[:10]} ...")
+    if unexpected:
+        raise KeyError(f"remap produced {len(unexpected)} keys the model does not have: {sorted(unexpected)[:10]} ...")
+    for key in expected:
+        if remapped[key].shape != transformer.state_dict()[key].shape:
+            raise ValueError(
+                f"shape mismatch on {key}: PinkCherry {tuple(remapped[key].shape)} vs model {tuple(transformer.state_dict()[key].shape)}"
+            )
+
+    missing, unexpected = transformer.load_state_dict(remapped, strict=False)
+    if missing:
+        raise KeyError(f"load left {len(missing)} keys missing: {sorted(missing)[:10]} ...")
+    if unexpected:
+        raise KeyError(f"load got {len(unexpected)} unexpected keys: {sorted(unexpected)[:10]} ...")
+    print(f"[pinkcherry] loaded {len(remapped)} tensors into the transformer", flush=True)