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):