diff --git a/modules/lora/lora_nunchaku.py b/modules/lora/lora_nunchaku.py index 1fc8ff9b1..b7499387e 100644 --- a/modules/lora/lora_nunchaku.py +++ b/modules/lora/lora_nunchaku.py @@ -1,12 +1,19 @@ import time from modules import shared, errors from modules.logger import log -from modules.lora import lora_load, lora_common +from modules.lora import lora_load, lora_common, network previously_loaded = [] # we maintain private state here +def wrap_network(network_on_disk): + net = network.Network(network_on_disk.name, network_on_disk) + net.mentioned_name = network_on_disk.name + network_on_disk.read_hash() # nothing else on this path fills the hash infotext reads + return net + + def load_nunchaku(names, strengths): global previously_loaded # pylint: disable=global-statement strengths = [s[0] if isinstance(s, list) else s for s in strengths] @@ -26,7 +33,7 @@ def load_nunchaku(names, strengths): from nunchaku.lora.flux.compose import compose_lora composed_lora = compose_lora(loras) shared.sd_model.transformer.update_lora_params(composed_lora) - lora_common.loaded_networks = [n[0] for n in networks] # used by infotext + lora_common.loaded_networks[:] = [wrap_network(n[0]) for n in networks] # read by infotext and the trigger tags t1 = time.time() lora_common.timer.load = t1 - t0 log.debug(f"Network load: type=LoRA method=nunchaku loras={names} strength={strengths} time={t1-t0:.3f}") diff --git a/test/test-sdnq-lora-factors.py b/test/test-sdnq-lora-factors.py index af44cb1b0..cb864d213 100644 --- a/test/test-sdnq-lora-factors.py +++ b/test/test-sdnq-lora-factors.py @@ -154,6 +154,9 @@ class MockNOD: self.shorthash = '' self.sd_version = 'unknown' + def read_hash(self): + pass + def make_net(name, layer, A, B, te_mult=1.0, alpha=None, dora=False): net = network.Network(name, MockNOD(name)) @@ -2411,6 +2414,16 @@ def test_remove_factors_after_device_move(): return True +def test_nunchaku_entries_carry_the_network_interface(): + from modules.lora import lora_nunchaku + nod = MockNOD('composed') + net = lora_nunchaku.wrap_network(nod) + assert len(net.modules) == 0, 'a composed set owns no modules: the reported method probes this to tell native from nunchaku' + assert net.network_on_disk is nod, 'infotext reads the hash through network_on_disk' + assert net.name == nod.name + return True + + def test_stacked_shape_mismatch_falls_back(): from types import SimpleNamespace layer = build_layer('uint4') @@ -2880,7 +2893,7 @@ def run_tests(): for fn in [test_factor_add_inside_compiled_graph, test_rank_bucket_graph_reuse, test_recompile_wall_resets_on_unload]: run_test(CAT_COMPILE, fn) log.warning('=== Robustness ===') - for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back]: + for fn in [test_remove_factors_after_device_move, test_stacked_shape_mismatch_falls_back, test_nunchaku_entries_carry_the_network_interface]: run_test(CAT_ROBUST, fn) log.warning('=== Block weights ===') for fn in [test_block_index_sd_unet_layout, test_block_index_sdxl_unet_layout, test_block_index_flux_chains_concatenate,