mirror of
https://github.com/vladmandic/automatic
synced 2026-09-06 13:00:44 +02:00
feat(te): reset text-encoder override when base model architecture changes
When a checkpoint change switches the model type, a custom sd_text_encoder no longer fits, so reset it to Default and clear loaded_te, mirroring the sd_unet reset. The type is resolved with detect_pipeline on both the loaded and incoming checkpoints, so same-arch switches (Krea2 Base and Turbo share one pipeline class) do not reset. The checkpoint handler also returns sd_text_encoder alongside sd_unet so the dropdown reflects it.
This commit is contained in:
+19
-9
@@ -1443,6 +1443,7 @@ def reload_text_encoder(initial=False):
|
||||
|
||||
|
||||
def reload_model_weights(sd_model=None, info: CheckpointInfo | None = None, op='model', force=False, revision=None):
|
||||
global loaded_te # pylint: disable=global-statement
|
||||
checkpoint_info = info or select_checkpoint(op=op) # are we selecting model or dictionary
|
||||
if checkpoint_info is None:
|
||||
unload_model_weights(op=op)
|
||||
@@ -1452,17 +1453,26 @@ def reload_model_weights(sd_model=None, info: CheckpointInfo | None = None, op='
|
||||
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__
|
||||
reset_unet = shared.opts.sd_unet not in (None, 'Default', 'None')
|
||||
reset_te = shared.opts.sd_text_encoder not in (None, 'Default', 'None')
|
||||
if op == 'model' and sd_model is not None and changed_checkpoint and (reset_unet or reset_te):
|
||||
# compare detected model type, not pipeline class: custom-loader arches (e.g. Krea2) load as a
|
||||
# concrete class but detect as generic DiffusionPipeline, so a class compare would falsely reset
|
||||
# across same-arch checkpoints (Base vs Turbo). detect both sides so the comparison is symmetric.
|
||||
try:
|
||||
new_pipeline, _ = sd_detect.detect_pipeline(checkpoint_info.path, op)
|
||||
_, new_type = sd_detect.detect_pipeline(checkpoint_info.path, op)
|
||||
_, old_type = sd_detect.detect_pipeline(loaded_ckpt.path, op) if loaded_ckpt is not None else (None, None)
|
||||
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
|
||||
new_type = old_type = None
|
||||
if new_type is not None and old_type is not None and new_type != old_type: # architecture changed: custom components no longer fit
|
||||
if reset_unet:
|
||||
log.info(f'Load model: type="{old_type}" changed="{new_type}" unet="{shared.opts.sd_unet}" set to default')
|
||||
shared.opts.data["sd_unet"] = 'Default'
|
||||
sd_unet.loaded_unet = None
|
||||
if reset_te:
|
||||
log.info(f'Load model: type="{old_type}" changed="{new_type}" te="{shared.opts.sd_text_encoder}" set to default')
|
||||
shared.opts.data["sd_text_encoder"] = 'Default'
|
||||
loaded_te = None
|
||||
if sd_model is None: # previous model load failed
|
||||
current_checkpoint_info = None
|
||||
else:
|
||||
|
||||
+10
-8
@@ -393,9 +393,11 @@ 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):
|
||||
def sync_checkpoint_components(value, progress=False, force=False):
|
||||
# a checkpoint change can reset sd_unet / sd_text_encoder to Default (arch changed);
|
||||
# push both back so the dropdowns reflect it, not just the stored option
|
||||
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
|
||||
return checkpoint_update, get_value_for_setting('sd_unet'), get_value_for_setting('sd_text_encoder'), settings_text
|
||||
|
||||
for k, _item in quicksettings_list:
|
||||
component = shared.settings_components[k]
|
||||
@@ -414,8 +416,8 @@ def create_quicksettings(interfaces):
|
||||
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]
|
||||
return sync_checkpoint_components(value, progress=progress)
|
||||
outputs = [component, shared.settings_components['sd_unet'], shared.settings_components['sd_text_encoder'], text_settings]
|
||||
else:
|
||||
def fn(value, k=k, progress=progress_flag):
|
||||
return run_settings_single(value, key=k, progress=progress)
|
||||
@@ -428,15 +430,15 @@ def create_quicksettings(interfaces):
|
||||
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)
|
||||
def sync_checkpoint_components_forced(value, _dummy):
|
||||
return sync_checkpoint_components(value, force=True)
|
||||
|
||||
button_set_checkpoint = gr.Button('Change model', elem_id='change_checkpoint', visible=False)
|
||||
button_set_checkpoint.click(
|
||||
fn=sync_checkpoint_unet_forced,
|
||||
fn=sync_checkpoint_components_forced,
|
||||
_js="consumeDesiredCheckpointName",
|
||||
inputs=[shared.settings_components['sd_model_checkpoint'], dummy_component],
|
||||
outputs=[shared.settings_components['sd_model_checkpoint'], shared.settings_components['sd_unet'], text_settings],
|
||||
outputs=[shared.settings_components['sd_model_checkpoint'], shared.settings_components['sd_unet'], shared.settings_components['sd_text_encoder'], text_settings],
|
||||
)
|
||||
button_set_refiner = gr.Button('Change refiner', elem_id='change_refiner', visible=False)
|
||||
button_set_refiner.click(
|
||||
|
||||
Reference in New Issue
Block a user