fix(lora): read the pdd grid from the file header when the cache lacks it

The metadata cache keeps a failed read forever and --no-metadata returns
nothing, and in both cases the head loader returned None silently while
the backbone still applied. The grid is now read from the file's own
header when the cached metadata lacks it, and head-shaped tensors
without a grid log a warning.
This commit is contained in:
CalamitousFelicitousness
2026-09-16 08:37:13 +01:00
parent 909db4b797
commit c67b4f499f
2 changed files with 44 additions and 1 deletions
+20
View File
@@ -87,9 +87,29 @@ def load(name, metadata, state_dict):
return ParallelHeads(num_steps, block_size, heads)
def header_metadata(filename, name):
"""The metadata read from the file itself, for when the cached metadata lacks the grid; warns when head-shaped tensors have no grid."""
from safetensors import safe_open
try:
with safe_open(filename, framework='pt', device='cpu') as f:
metadata = f.metadata() or {}
if METADATA_STEPS in metadata:
log.debug(f'Network load: type=PDD name="{name}" grid read from the file header')
else:
heads = sum(1 for key in f.keys() if key.endswith('.weight') and len(f.get_slice(key).get_shape()) == 3)
if heads > 0:
log.warning(f'Network load: type=PDD name="{name}" heads={heads} no {METADATA_STEPS} metadata: heads ignored')
return metadata
except Exception as e:
log.warning(f'Network load: type=PDD name="{name}" header {e}')
return {}
def try_load(name, network_on_disk, lora_scale): # pylint: disable=unused-argument
"""Family loader for the native chain: a network carrying only the heads."""
metadata = getattr(network_on_disk, 'metadata', None) or {}
if METADATA_STEPS not in metadata:
metadata = header_metadata(network_on_disk.filename, name) # the metadata cache keeps a failed read forever and --no-metadata returns nothing
if METADATA_STEPS not in metadata:
return None
from modules.lora import native_adapter
+24 -1
View File
@@ -194,6 +194,29 @@ def test_load_without_heads_is_none():
assert network_pdd.load('x', METADATA, {'transformer_blocks.0.attn.to_q.lora_down': torch.zeros(64, HIDDEN)}) is None
def write_pdd_file(folder, metadata):
from safetensors.torch import save_file
path = os.path.join(folder, 'pdd.safetensors')
save_file({'proj_out.weight': torch.zeros(NUM_STEPS, VIDEO_OUT, HIDDEN), 'proj_out.bias': torch.zeros(NUM_STEPS, VIDEO_OUT)}, path, metadata=metadata)
return path
def test_try_load_reads_the_grid_from_the_file_header():
import tempfile
with tempfile.TemporaryDirectory() as folder:
disk = types.SimpleNamespace(name='pdd', filename=write_pdd_file(folder, METADATA), metadata={}) # the cached metadata lacks the grid
net = network_pdd.try_load('pdd', disk, 1.0)
assert net is not None and network_pdd.EXTRAS_KEY in net.extras, 'the grid in the file header must be enough'
assert net.extras[network_pdd.EXTRAS_KEY].nfe == NUM_STEPS // BLOCK
def test_try_load_without_grid_metadata_is_none():
import tempfile
with tempfile.TemporaryDirectory() as folder:
disk = types.SimpleNamespace(name='pdd', filename=write_pdd_file(folder, {'alpha': '1'}), metadata={})
assert network_pdd.try_load('pdd', disk, 1.0) is None
# ============================================================
# Tests: grid and fusion math
# ============================================================
@@ -352,7 +375,7 @@ def test_pin_overrides_steps_and_shift():
def main():
cat = category('detect')
for fn in (test_detect_grid, test_detect_rejects_bad_block, test_load_collects_heads, test_load_without_metadata_is_none, test_load_without_heads_is_none):
for fn in (test_detect_grid, test_detect_rejects_bad_block, test_load_collects_heads, test_load_without_metadata_is_none, test_load_without_heads_is_none, test_try_load_reads_the_grid_from_the_file_header, test_try_load_without_grid_metadata_is_none):
run_test(cat, fn)
cat = category('math')
for fn in (test_grid_intervals_match_reference, test_steps_for_counts_terminal_sigma, test_parallel_head_matches_reference_per_block, test_parallel_head_strength_blend, test_parallel_head_clamps_overflow, test_parallel_head_keeps_base_out_of_tree):