File size: 8,151 Bytes
90f10d2 fcff745 90f10d2 271bd15 90f10d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | """Fetching tensors from sources into local bytes."""
from pathlib import Path
import json
import os
from vlib import ctx
from vlib.ui import _fail, _now_iso, _step, _warn
from vlib.net import _cleanup_tmp, _safe_id, download_range
from vlib.tensors import _is_ggml_quant, _is_quant_fetched, _write_safetensors_streaming, _write_single_tensor_gguf, encode_from_f32, write_safetensors
from vlib.sources import _hf_config, _public_source, detect_output_tensor, largest_2d
def _source_arch(source):
if source.kind == "safetensors":
cfg = _hf_config(source.repo) if source.repo else None
if cfg:
arch = (cfg.get("architectures") or [None])[0] or cfg.get("model_type")
else:
arch = None
return arch, cfg
source._load()
if getattr(source, "_reader", None) is not None:
for f in source._reader.fields.values():
if f.name == "general.architecture":
try:
v = f.parts[f.data[0]]
return (v.decode() if isinstance(v, bytes) else str(v)), None
except Exception:
return None, None
return (source._remote[0] if source._remote else None), None
def _resolve_output_name(source, config, names):
if source.kind == "gguf":
for n in ("output.weight", "token_embd.weight"):
if n in names:
return n
n = detect_output_tensor(names, config)
if n:
return n
else:
n = detect_output_tensor(names, config)
if n:
return n
shapes = {n: source.ref(n).shape for n in names}
return largest_2d(names, shapes)
def _stream_safetensors_tensor(source, name, dest_path):
"""Stream a safetensors tensor's raw bytes to dest_path. Returns byte count."""
if source.path:
hs, tensors = source._local()
begin, end = tensors[name]["data_offsets"]
with open(source.path, "rb") as src, open(dest_path, "wb") as out:
src.seek(8 + hs + begin)
remaining = end - begin
while remaining:
b = src.read(min(8 << 20, remaining))
if not b:
break
out.write(b)
remaining -= len(b)
return end - begin
url, hs, tensors = source._shard_info(name)
begin, end = tensors[name]["data_offsets"]
import shutil
combined = download_range(url, 8 + hs + begin, 8 + hs + end - 1, dest_path.parent, None, label=name)
shutil.move(str(combined), str(dest_path))
return end - begin
def _fetch_tensor(source, name, tmp_dir):
"""Return (kind, raw_path, data, dtype, shape) for one tensor.
kind == 'file' → raw bytes at raw_path; kind == 'bytes' → data holds bytes.
GGML-quantized tensors keep their exact raw blocks (no requant)."""
ref = source.ref(name)
tmp_dir.mkdir(parents=True, exist_ok=True)
if source.kind == "safetensors":
raw_path = tmp_dir / "raw.bin"
n = _stream_safetensors_tensor(source, name, raw_path)
return ("file", raw_path, n, ref.dtype, ref.shape)
if _is_ggml_quant(ref.dtype):
# passthrough: exact quant blocks, stored in a single-tensor GGUF artifact
raw_path = tmp_dir / "raw.bin"
n = _stream_gguf_tensor(source, name, raw_path)
return ("file", raw_path, n, ref.dtype, ref.shape)
f32 = source.read_f32(name)
# Preserve original precision family: BF16 stays BF16 (not F16 downconvert),
# F32 stays F32, F16 stays F16. BF16->F32->BF16 round-trip is exact.
_rd = ref.dtype.upper()
if _rd in ("F32", "FP32"):
dtype = "F32"
elif _rd in ("BF16", "BF16E8M"):
dtype = "BF16"
else:
dtype = "F16"
return ("bytes", None, encode_from_f32(f32, dtype), dtype, ref.shape)
def _stream_gguf_tensor(source, name, dest_path):
"""Stream a GGUF tensor's exact raw blocks to dest_path. Returns byte count."""
import shutil
import numpy as np
if source.path:
source._load()
for t in source._reader.tensors:
if t.name == name:
with open(dest_path, "wb") as out:
raw = t.data.tobytes() if t.data.dtype == np.uint8 else memoryview(t.data).cast("B").tobytes()
out.write(raw)
return t.n_bytes
raise KeyError(name)
url, parsed = source._owner(name)
dims, gtype, data_offset, n_bytes = parsed[2][name]
start = parsed[1] + data_offset
combined = download_range(url, start, start + n_bytes - 1, dest_path.parent, None, label=name)
shutil.move(str(combined), str(dest_path))
return n_bytes
def _write_fetched(fetched, tensor_name, out_path):
"""Write a fetched tensor to a standalone artifact (gguf when quantized)."""
kind, raw_path, data, dtype, shape = fetched
if _is_quant_fetched(fetched):
if not str(out_path).endswith(".gguf"):
_fixed = str(out_path)[:-len(".safetensors")] + ".gguf" if str(out_path).endswith(".safetensors") else str(out_path) + ".gguf"
_warn(f" Quant artifact needs .gguf — writing {_fixed} instead of the requested extension.")
out_path = _fixed
_write_single_tensor_gguf(tensor_name, raw_path, raw_path.stat().st_size, dtype, shape, str(out_path))
return Path(out_path)
tmp = str(out_path) + ".tmp"
if kind == "file":
_write_safetensors_streaming(tensor_name, str(raw_path), raw_path.stat().st_size, dtype, shape, tmp)
else:
write_safetensors({tensor_name: (data, dtype, shape)}, tmp)
os.replace(tmp, str(out_path))
return Path(out_path)
def _write_voicepack_json(out_path, source_id, pack):
"""Provenance sidecar for a voicepack: source + per-tensor names/shapes/dtypes.
Written next to the pack file (same stem, .json suffix)."""
meta = {
"source": _public_source(source_id),
"tensors": [{"name": n, "shape": [int(x) for x in v[2]], "dtype": v[1]}
for n, v in pack.items()],
"downloaded_at": _now_iso(),
}
jpath = Path(str(out_path)).with_suffix(".json")
jtmp = jpath.parent / (jpath.name + ".tmp")
jtmp.write_text(json.dumps(meta, indent=2) + "\n")
jtmp.replace(jpath)
return jpath
def _fetch_pack(source, source_id, resolved):
"""Fetch each resolved tensor's raw bytes. Returns pack {name: (bytes, dtype, shape)}.
Shared by get --multi and pack: always safetensors (GGUF quants dequant once with a warning)."""
pack = {}
for tname in resolved:
ref = source.ref(tname)
_step(f"Extracting {tname} ({ref.dtype} {tuple(ref.shape)})…")
tmp_dir = ctx.VOICES_DIR / ".parts" / _safe_id(source_id, tname)
try:
fetched = _fetch_tensor(source, tname, tmp_dir)
except Exception as e:
_fail(f" ✗ Could not extract tensor {tname}: {e}")
kind, raw_path, data, dtype, shape = fetched
try:
if kind == "file":
# safetensors raw file: read bytes directly (preserves BF16/F16/F32).
with open(str(raw_path), "rb") as f:
b = f.read()
pack[tname] = (b, dtype, tuple(shape))
else:
# bytes already encoded (F32/F16/BF16 preserved).
pack[tname] = (data, dtype, tuple(shape))
# If fetched was GGUF quant passthrough, _fetch_tensor returns
# ("file", raw, n, qtype, shape) with quant dtype — dequant once for safetensors pack.
if _is_ggml_quant(dtype):
_warn(f" {tname} is {dtype} quant — dequanting once to F32 for voicepack.safetensors (unavoidable, safetensors has no quant).")
try:
f32 = source.read_f32(tname)
pack[tname] = (f32.astype("float32").tobytes(), "F32", tuple(shape))
except Exception as e:
_fail(f" ✗ Could not dequant {tname} for pack: {e}")
finally:
_cleanup_tmp(tmp_dir)
if not pack:
_fail(" ✗ Nothing to pack.")
return pack
|