diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index b670256cb..91b6e2d2b 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -33,7 +33,7 @@ from dataclasses import dataclass import torch -from modules import shared, sd_models +from modules import shared, sd_models, sd_models_utils from modules.logger import log from modules.lora import ( lora_convert, network, network_boft, network_full, network_glora, @@ -483,7 +483,7 @@ def try_load_lokr(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, - arch_name="generic"): + arch_name="generic"): # pylint: disable=unused-argument """Generic LoKR loader. Stores only the compact LoKR factors and dispatches to @@ -759,7 +759,7 @@ def try_load_norm(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, - arch_name="generic"): + arch_name="generic"): # pylint: disable=unused-argument """Generic Norm (LayerNorm / RMSNorm weight + bias delta) loader. Norm targets are never fused, so the chunk dispatch is dropped. @@ -873,6 +873,7 @@ def try_load_chain(name, network_on_disk, lora_scale, family_loaders): tuple of partial-applied generic loaders, each already bound to the arch's ``resolve_targets`` and prefix tuples. """ + sd_models_utils.state_dict_cache.enable() net = None for try_fn in family_loaders: sub = try_fn(name, network_on_disk, lora_scale) @@ -882,4 +883,5 @@ def try_load_chain(name, network_on_disk, lora_scale, family_loaders): net = sub else: net.modules.update(sub.modules) + sd_models_utils.state_dict_cache.disable() return net diff --git a/modules/sd_detect.py b/modules/sd_detect.py index 2f7ace9a3..d3cf26628 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -92,7 +92,7 @@ def guess_by_name(fn, current_guess): new_guess = 'ZetaChroma' elif 'chroma' in fn.lower() and 'xl' not in fn.lower(): new_guess = 'Chroma' - elif 'flux.2' in fn.lower() and 'klein' in fn.lower(): + elif ('flux.2' in fn.lower() or 'f2' in fn.lower()) and 'klein' in fn.lower(): new_guess = 'FLUX2Klein' elif 'flux.2' in fn.lower(): new_guess = 'FLUX2' diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py index 773d955b0..e2df10828 100644 --- a/modules/sd_models_utils.py +++ b/modules/sd_models_utils.py @@ -15,6 +15,30 @@ from modules.sd_checkpoint import CheckpointInfo # pylint: disable=unused-import debug = log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None +class StateDictCache: + _enabled: bool = True + _cache: dict[str, dict] = {} + + def get(self, key: str): + if not self._enabled: + return None + return self._cache.get(key, None) + + def set(self, key: str, value: dict): + if not self._enabled: + return + self._cache[key] = value + + def enable(self): + self._enabled = True + + def disable(self): + self._enabled = False + self._cache.clear() + +state_dict_cache = StateDictCache() + + class NoWatermark: def apply_watermark(self, img): return img @@ -97,12 +121,18 @@ def convert_to_faketensors(tensor): def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pylint: disable=unused-argument + cached = state_dict_cache.get(checkpoint_file) + if cached is not None: + return cached if not os.path.isfile(checkpoint_file): - log.error(f'Load dict: path="{checkpoint_file}" not a file') + log.error(f'Load dict: file="{checkpoint_file}" not a file') return None _, extension = os.path.splitext(checkpoint_file) if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt: - log.warning(f"Checkpoint loading disabled: {checkpoint_file}") + log.warning(f'Load dict: file="{checkpoint_file}" checkpoint loading disabled') + return None + if shared.state.interrupted: + log.warning(f'Load dict: file="{checkpoint_file}" interrupted before read') return None try: pl_sd = None @@ -123,6 +153,7 @@ def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pyl else: pl_sd = torch.load(f, map_location='cpu') sd = get_state_dict_from_checkpoint(pl_sd) + state_dict_cache.set(checkpoint_file, sd) del pl_sd except Exception as e: errors.display(e, f'Load model: {checkpoint_file}') diff --git a/pipelines/model_flux2_klein.py b/pipelines/model_flux2_klein.py index 9540c21e0..99fbca869 100644 --- a/pipelines/model_flux2_klein.py +++ b/pipelines/model_flux2_klein.py @@ -20,7 +20,12 @@ def load_flux2_klein(checkpoint_info, diffusers_load_config=None): if repo_id is None or repo_id.lower() == 'none': return None - pipe = diffusers.Flux2KleinPipeline.from_pretrained( + if '-kv' in repo_id: + cls = diffusers.Flux2KleinKVPipeline + else: + cls = diffusers.Flux2KleinPipeline + + pipe = cls.from_pretrained( repo_id, transformer=transformer, text_encoder=text_encoder, @@ -30,9 +35,9 @@ def load_flux2_klein(checkpoint_info, diffusers_load_config=None): pipe.task_args = { 'output_type': 'np', } - diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flux2klein"] = diffusers.Flux2KleinPipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux2klein"] = diffusers.Flux2KleinPipeline - diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux2klein"] = diffusers.Flux2KleinPipeline + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flux2klein"] = cls + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux2klein"] = cls + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux2klein"] = cls generic.load_vae_override(pipe, diffusers_load_config)