Spaces:
Running on Zero
Running on Zero
File size: 4,628 Bytes
fed6c68 | 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 | # Copyright (c) 2026 Bytedance Ltd. and/or its affiliate
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Load Bernini Renderer transformer weights from a safetensors checkpoint.
Accepts a local directory, a ``*.safetensors.index.json`` file, or a Hugging
Face repo id. For each transformer a list of candidate key prefixes is tried
in order (``diff_dec.transformer.`` -> ``transformer.`` -> bare keys); keys may
carry a leading ``ema.`` prefix, and when both an EMA and a non-EMA copy of a
tensor exist the EMA copy is used.
"""
import json
import logging
import os
from safetensors import safe_open
logger = logging.getLogger("bernini.weights")
HIGH_NOISE_PREFIXES = ["diff_dec.transformer.", "transformer.", ""]
LOW_NOISE_PREFIXES = ["diff_dec.transformer_2.", "transformer_2.", ""]
def _resolve_dir(path: str) -> str:
if os.path.isfile(path) and path.endswith("safetensors.index.json"):
return os.path.dirname(path)
if os.path.isdir(path):
return path
from huggingface_hub import snapshot_download # treat as a HF repo id
return snapshot_download(path, allow_patterns=["*.safetensors", "*.json"])
def _build_weight_map(ckpt_dir: str) -> dict:
"""Map every tensor key to the shard file that holds it."""
index_files = sorted(f for f in os.listdir(ckpt_dir) if f.endswith("safetensors.index.json"))
if index_files:
with open(os.path.join(ckpt_dir, index_files[0])) as f:
return json.load(f)["weight_map"]
shard_files = sorted(f for f in os.listdir(ckpt_dir) if f.endswith(".safetensors"))
if not shard_files:
raise FileNotFoundError(f"no .safetensors weights found in {ckpt_dir}")
weight_map = {}
for fn in shard_files:
with safe_open(os.path.join(ckpt_dir, fn), framework="pt", device="cpu") as f:
for key in f.keys():
weight_map[key] = fn
return weight_map
def _select_keys(weight_map: dict, prefix: str) -> dict:
"""Map target tensor name -> source key for `prefix`, preferring EMA copies."""
chosen = {}
for key in weight_map:
is_ema = key.startswith("ema.")
base = key[len("ema.") :] if is_ema else key
if not base.startswith(prefix):
continue
target = base[len(prefix) :]
if target not in chosen or (is_ema and not chosen[target].startswith("ema.")):
chosen[target] = key
return chosen
def load_transformer_state_dict(ckpt_path: str, prefixes: list):
"""Return ``(state_dict, prefix_used)`` for one transformer."""
ckpt_dir = _resolve_dir(ckpt_path)
weight_map = _build_weight_map(ckpt_dir)
for prefix in prefixes:
chosen = _select_keys(weight_map, prefix)
if not chosen:
continue
shards = {}
for target, source in chosen.items():
shards.setdefault(weight_map[source], []).append((source, target))
state_dict = {}
for shard_file, items in shards.items():
with safe_open(os.path.join(ckpt_dir, shard_file), framework="pt", device="cpu") as f:
for source, target in items:
state_dict[target] = f.get_tensor(source)
return state_dict, prefix
raise ValueError(f"no weights matching prefixes {prefixes} found in {ckpt_path}")
def load_weights(model, high_noise_ckpt: str, low_noise_ckpt: str):
"""Load the high-noise and low-noise transformer weights into a BerniniRendererModel."""
high, prefix = load_transformer_state_dict(high_noise_ckpt, HIGH_NOISE_PREFIXES)
miss, unexpected = model.diff_dec.transformer.load_state_dict(high, strict=False, assign=False)
logger.info("high-noise: %d tensors via prefix '%s', missing=%d unexpected=%d",
len(high), prefix, len(miss), len(unexpected))
low, prefix = load_transformer_state_dict(low_noise_ckpt, LOW_NOISE_PREFIXES)
miss, unexpected = model.diff_dec.transformer_2.load_state_dict(low, strict=False, assign=False)
logger.info("low-noise: %d tensors via prefix '%s', missing=%d unexpected=%d",
len(low), prefix, len(miss), len(unexpected))
|