From c67b4f499f383d09cf85961abebb95b397c43cec Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Wed, 16 Sep 2026 08:37:13 +0100 Subject: [PATCH] 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. --- modules/lora/network_pdd.py | 20 ++++++++++++++++++++ test/test-pdd.py | 25 ++++++++++++++++++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/modules/lora/network_pdd.py b/modules/lora/network_pdd.py index 3ae55c708..86b1744b2 100644 --- a/modules/lora/network_pdd.py +++ b/modules/lora/network_pdd.py @@ -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 diff --git a/test/test-pdd.py b/test/test-pdd.py index 52a29677b..de5e3d04b 100644 --- a/test/test-pdd.py +++ b/test/test-pdd.py @@ -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):