Merge pull request #4876 from vladmandic/feat/native-transformer-loader

Feat/native transformer loader

modules/lora/lora_load.py

@@ -19,7 +19,7 @@ exclude_errors = [
 
 # shared.sd_model_type -> dotted module path of a pipeline native loader
 # exposing ``try_load(name, network_on_disk, lora_scale)``. New archs add an
-# entry here and ship a per-arch ``try_load`` (either binding native_loader's
+# entry here and ship a per-arch ``try_load`` (either binding native_adapter's
 # generic helpers via try_load_chain, or rolling their own).
 _NATIVE_DISPATCH = {
     'zimage':     'pipelines.z_image.zimage_lora',

modules/lora/native_adapter.py

@@ -1,7 +1,6 @@
 """Shared scaffolding for native adapter loaders.
 
-The four native adapter loaders (z-image, chroma, ernie, flux2) all implement
-the same algorithm:
+Each per-arch native adapter loader implements the same algorithm:
 
 1. Read the safetensors state dict
 2. Test for family-specific markers; bail out if absent
@@ -25,7 +24,7 @@ diffusers paths plus optional chunk descriptors).
 
 Per-arch loader modules import this module and pass their own ``prefixes``,
 ``bare_prefixes``, ``bare_diffusers_prefixes``, and ``resolve_targets`` to the
-generic helpers. Loader business logic itself lands in subsequent commits.
+generic helpers.
 """
 
 import os
@@ -45,7 +44,8 @@ from modules.lora import lora_common as l
 
 
 # Universal prefix list shared by every native arch loader. Per-arch loaders
-# extend this with arch-specific entries (e.g. flux2 adds ``"lycoris_"``).
+# extend this with arch-specific entries when their files use additional
+# vendor-specific naming conventions.
 KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_")
 
 
@@ -56,9 +56,10 @@ KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_")
 BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers"
 
 
-# Default network-key prefix. Single-component arches (flux2, zimage, chroma,
-# ernie) keep this default; multi-component arches (anima: transformer plus
-# llm_adapter plus text_encoder) pass a callable that picks per ``prefix_used``.
+# Default network-key prefix. Single-component arches keep this default;
+# multi-component arches (those with separate text-encoder or adapter
+# components alongside the transformer) pass a callable that picks per
+# ``prefix_used``.
 NETWORK_PREFIX_DEFAULT = "lora_transformer_"
 
 

modules/sd_models.py

@@ -1412,6 +1412,19 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi
     jobid = shared.state.begin('Load model')
     if sd_model is None:
         sd_model = model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner
+    loaded_ckpt = getattr(sd_model, 'sd_checkpoint_info', None) if sd_model is not None else None
+    changed_checkpoint = loaded_ckpt is None or checkpoint_info is None or loaded_ckpt.filename != checkpoint_info.filename
+    if op == 'model' and sd_model is not None and changed_checkpoint and shared.opts.sd_unet not in (None, 'Default', 'None'):
+        old_class = type(sd_model).__name__
+        try:
+            new_pipeline, _ = sd_detect.detect_pipeline(checkpoint_info.path, op)
+        except Exception:
+            new_pipeline = None
+        new_class = getattr(new_pipeline, '__name__', None)
+        if new_class is not None and new_class != old_class:
+            log.info(f'Load model: pipeline cls={old_class} changed={new_class} unet="{shared.opts.sd_unet}" set to default')
+            shared.opts.data["sd_unet"] = 'Default'
+            sd_unet.loaded_unet = None
     if sd_model is None:  # previous model load failed
         current_checkpoint_info = None
     else:

modules/sd_models_utils.py

@@ -100,26 +100,29 @@ def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pyl
     if not os.path.isfile(checkpoint_file):
         log.error(f'Load dict: path="{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}")
+        return None
     try:
         pl_sd = None
-        with progress.open(checkpoint_file, 'rb', description=f'[cyan]Load {what}: [yellow]{checkpoint_file}', auto_refresh=True, console=console) as f:
-            _, extension = os.path.splitext(checkpoint_file)
-            if extension.lower() == ".ckpt" and shared.opts.sd_disable_ckpt:
-                log.warning(f"Checkpoint loading disabled: {checkpoint_file}")
-                return None
-            if shared.opts.stream_load:
-                if extension.lower() == ".safetensors":
-                    buffer = f.read()
-                    pl_sd = safetensors.torch.load(buffer)
-                else:
-                    buffer = io.BytesIO(f.read())
-                    pl_sd = torch.load(buffer, map_location='cpu')
-            else:
-                if extension.lower() == ".safetensors":
-                    pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu')
+        # safetensors.torch.load_file opens its own handle by path, so wrapping
+        # with progress.open leaves the bar stuck at 0/total. Skip the wrapper
+        # on that path; other paths actually read through f and update.
+        if extension.lower() == ".safetensors" and not shared.opts.stream_load:
+            pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu')
+        else:
+            with progress.open(checkpoint_file, 'rb', description=f'[cyan]Load {what}: [yellow]{checkpoint_file}', auto_refresh=True, console=console) as f:
+                if shared.opts.stream_load:
+                    if extension.lower() == ".safetensors":
+                        buffer = f.read()
+                        pl_sd = safetensors.torch.load(buffer)
+                    else:
+                        buffer = io.BytesIO(f.read())
+                        pl_sd = torch.load(buffer, map_location='cpu')
                 else:
                     pl_sd = torch.load(f, map_location='cpu')
-            sd = get_state_dict_from_checkpoint(pl_sd)
+        sd = get_state_dict_from_checkpoint(pl_sd)
         del pl_sd
     except Exception as e:
         errors.display(e, f'Load model: {checkpoint_file}')

modules/sd_unet.py

@@ -46,6 +46,12 @@ def load_unet(model, repo_id: str | None = None):
                 return
 
     if shared.opts.sd_unet == 'Default' or shared.opts.sd_unet == 'None':
+        # Switching back to Default reverts a previously-loaded custom transformer.
+        if loaded_unet in (None, 'Default', 'None'):
+            return
+        log.info(f'Load module: type=UNet name="Default" (was="{loaded_unet}") reverting to base transformer')
+        loaded_unet = shared.opts.sd_unet
+        sd_models.reload_model_weights(force=True)
         return
 
     if shared.opts.sd_unet not in list(unet_dict):
@@ -74,7 +80,7 @@ def load_unet(model, repo_id: str | None = None):
                 model.prior_pipe.text_encoder = prior_text_encoder.to(devices.device, dtype=devices.dtype)
         elif any([m in model.__class__.__name__ for m in dit_models]) or hasattr(model, 'transformer'): # noqa: C419 # pylint: disable=use-a-generator
             loaded_unet = shared.opts.sd_unet
-            sd_models.load_diffuser() # TODO model load: force-reloading entire model as loading transformers only leads to massive memory usage
+            sd_models.reload_model_weights(force=True) # full reload: in-place transformer swap leaks memory
         else:
             if not hasattr(model, 'unet') or model.unet is None:
                 log.error('Load module: type=UNET not found in current model')

modules/ui_settings.py

@@ -391,6 +391,10 @@ def create_quicksettings(interfaces):
         if shared.opts.notification_audio_enable and os.path.exists(os.path.join(paths.script_path, shared.opts.notification_audio_path)):
             gr.Audio(interactive=False, value=os.path.join(paths.script_path, shared.opts.notification_audio_path), elem_id="audio_notification", visible=False)
 
+        def sync_checkpoint_unet(value, progress=False, force=False):
+            checkpoint_update, settings_text = run_settings_single(value, key='sd_model_checkpoint', progress=progress, force=force)
+            return checkpoint_update, get_value_for_setting('sd_unet'), settings_text
+
         for k, _item in quicksettings_list:
             component = shared.settings_components[k]
             info = shared.opts.data_labels[k]
@@ -405,20 +409,32 @@ def create_quicksettings(interfaces):
                 change_handlers = [component.blur]
             else:
                 change_handlers = [component.release if hasattr(component, 'release') else component.change]
+            progress_flag = info.refresh is not None
+            if k == 'sd_model_checkpoint':
+                def fn(value, progress=progress_flag):
+                    return sync_checkpoint_unet(value, progress=progress)
+                outputs = [component, shared.settings_components['sd_unet'], text_settings]
+            else:
+                def fn(value, k=k, progress=progress_flag):
+                    return run_settings_single(value, key=k, progress=progress)
+                outputs = [component, text_settings]
             for change_handler in change_handlers:
                 change_handler(
-                    fn=lambda value, k=k, progress=info.refresh is not None: run_settings_single(value, key=k, progress=progress),
+                    fn=fn,
                     inputs=[component],
-                    outputs=[component, text_settings],
+                    outputs=outputs,
                     show_progress='full' if info.refresh is not None else 'hidden',
                 )
 
+        def sync_checkpoint_unet_forced(value, _dummy):
+            return sync_checkpoint_unet(value, force=True)
+
         button_set_checkpoint = gr.Button('Change model', elem_id='change_checkpoint', visible=False)
         button_set_checkpoint.click(
-            fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint', force=True),
+            fn=sync_checkpoint_unet_forced,
             _js="consumeDesiredCheckpointName",
             inputs=[shared.settings_components['sd_model_checkpoint'], dummy_component],
-            outputs=[shared.settings_components['sd_model_checkpoint'], text_settings],
+            outputs=[shared.settings_components['sd_model_checkpoint'], shared.settings_components['sd_unet'], text_settings],
         )
         button_set_refiner = gr.Button('Change refiner', elem_id='change_refiner', visible=False)
         button_set_refiner.click(