refactor(flux2): expose single try_load entry point for dispatcher

Move the family-iteration loop into pipelines/flux/flux2_lora.try_load
so the f2 branch in lora_load.load_safetensors becomes a one-line call.

Per-family try_load_* entry points stay public for direct testing.

Addresses review feedback on PR #4841.
This commit is contained in:
CalamitousFelicitousness
2026-05-10 16:33:18 +01:00
parent 69506551a7
commit 0c44b93b25
2 changed files with 27 additions and 19 deletions
+1 -19
View File
@@ -103,25 +103,7 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
return chroma_net
if shared.sd_model_type == 'f2':
from pipelines.flux import flux2_lora
lora_scale = shared.opts.extra_networks_default_multiplier
f2_net = None
for try_fn in (
flux2_lora.try_load_lora,
flux2_lora.try_load_lokr,
flux2_lora.try_load_loha,
flux2_lora.try_load_oft,
flux2_lora.try_load_ia3,
flux2_lora.try_load_glora,
flux2_lora.try_load_norm,
flux2_lora.try_load_full,
):
sub = try_fn(name, network_on_disk, lora_scale)
if sub is None:
continue
if f2_net is None:
f2_net = sub
else:
f2_net.modules.update(sub.modules)
f2_net = flux2_lora.try_load(name, network_on_disk, shared.opts.extra_networks_default_multiplier)
if f2_net is not None:
lora_cache[name] = f2_net
return f2_net
+26
View File
@@ -423,6 +423,32 @@ def _bfl_to_diffusers_targets(base):
# === Native loaders ===
def try_load(name, network_on_disk, lora_scale):
"""Run every Flux2 family loader in dispatch order, merge any that match.
Per-family ``try_load_*`` entry points stay public; this is the single
umbrella the dispatcher in ``modules.lora.lora_load.load_safetensors``
calls. Order matters only for marker-cost: LoRA / LoKR are most common
so their fast bail-out runs first; the rare families come last.
Returns a ``Network`` with the union of modules from every matching
family loader, or ``None`` if no loader recognized the file.
"""
net = None
for try_fn in (
try_load_lora, try_load_lokr, try_load_loha, try_load_oft,
try_load_ia3, try_load_glora, try_load_norm, try_load_full,
):
sub = try_fn(name, network_on_disk, lora_scale)
if sub is None:
continue
if net is None:
net = sub
else:
net.modules.update(sub.modules)
return net
def try_load_lora(name, network_on_disk, lora_scale):
"""Load a Flux2/Klein LoRA (plus DoRA via the universal ``finalize_updown`` hook) as native modules.