mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
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:
@@ -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
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user