diff --git a/modules/sd_models.py b/modules/sd_models.py index ce6e37f6a..b08a70b72 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1475,8 +1475,9 @@ def reload_model_weights(sd_model=None, info: CheckpointInfo | None = None, op=' 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 reset_unet = shared.opts.sd_unet not in (None, 'Default', 'None') + reset_unet_secondary = shared.opts.sd_unet_secondary 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): + if op == 'model' and sd_model is not None and changed_checkpoint and (reset_unet or reset_unet_secondary 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. @@ -1490,6 +1491,10 @@ def reload_model_weights(sd_model=None, info: CheckpointInfo | None = None, op=' 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_unet_secondary: + log.info(f'Load model: type="{old_type}" changed="{new_type}" unet_secondary="{shared.opts.sd_unet_secondary}" set to default') + shared.opts.data["sd_unet_secondary"] = 'Default' + sd_unet.loaded_unet_secondary = 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' diff --git a/modules/sd_unet.py b/modules/sd_unet.py index 7072bd019..890662be0 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -5,11 +5,14 @@ from modules.logger import log unet_dict = {} loaded_unet = None +loaded_unet_secondary = None failed_unet = [] debug = os.environ.get('SD_LOAD_DEBUG', None) is not None dit_models = ['Flux', 'StableDiffusion3', 'HiDream', 'Lumina2', 'Chroma', 'Wan', 'Qwen', 'Anima'] +# model types (shared.sd_model_type keyspace) the secondary UNET override applies to +DUAL_TRANSFORMER_TYPES = ('ideogram4', 'wanai') def load_unet_sdxl_nunchaku(repo_id): @@ -101,6 +104,35 @@ def load_unet(model, repo_id: str | None = None): devices.torch_gc() +def load_unet_secondary(model): # pylint: disable=unused-argument + """Onchange handler for the secondary UNET override: a change means a + full reload; for single-transformer models the selection is stored and + applies on the next dual-transformer load. + """ + global loaded_unet_secondary # pylint: disable=global-statement + selected = shared.opts.sd_unet_secondary + + if selected is None or selected in ('Default', 'None'): + if loaded_unet_secondary in (None, 'Default', 'None'): + return + log.info(f'Load module: type=UNet slot=secondary name="Default" (was="{loaded_unet_secondary}") reverting to base transformer') + loaded_unet_secondary = selected + sd_models.reload_model_weights(force=True) + return + + if selected not in list(unet_dict): + log.error(f'Load module: type=UNet slot=secondary not found: {selected}') + return + if selected == loaded_unet_secondary or selected in failed_unet: + return + if shared.sd_model_type not in DUAL_TRANSFORMER_TYPES: + log.warning(f'Load module: type=UNet slot=secondary name="{selected}" stored: model type={shared.sd_model_type} has a single transformer, applies on next dual-transformer load') + return + loaded_unet_secondary = selected + sd_models.reload_model_weights(force=True) + devices.torch_gc() + + def refresh_unet_list(): unet_dict.clear() for file in files_cache.list_files(shared.opts.unet_dir, ext_filter=[".safetensors", ".gguf", ".pth"]): diff --git a/modules/shared_items.py b/modules/shared_items.py index a8e8e55ec..4bef8c855 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -116,6 +116,12 @@ def refresh_unet_list(): modules.sd_unet.refresh_unet_list() +def sd_unet_secondary_visible(): + import modules.sd_unet + from modules import shared + return shared.sd_model_type in modules.sd_unet.DUAL_TRANSFORMER_TYPES + + def sd_te_items(): import modules.model_te predefined = ['Default'] diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 7d0ac016e..2075256c9 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -94,6 +94,7 @@ def create_settings(cmd_opts): "sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", DropdownEditable, lambda: {"choices": list_checkpoint_titles()}, refresh=refresh_checkpoints), "sd_model_refiner": OptionInfo('None', "Refiner model", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_titles()}, refresh=refresh_checkpoints), "sd_unet": OptionInfo("Default", "UNET model", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list), + "sd_unet_secondary": OptionInfo("Default", "UNET model secondary", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items(), "visible": shared_items.sd_unet_secondary_visible()}, refresh=shared_items.refresh_unet_list), "latent_history": OptionInfo(20, "Latent history size", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}), "advanced_sep": OptionInfo("

Advanced Options

", "", gr.HTML), diff --git a/modules/ui_settings.py b/modules/ui_settings.py index b89b0d335..a1048d520 100644 --- a/modules/ui_settings.py +++ b/modules/ui_settings.py @@ -10,9 +10,20 @@ ui_system_tabs = None # required for system-info dummy_component = gr.Textbox(visible=False, value='dummy') loadsave = ui_loadsave.UiLoadsave(shared.cmd_opts.ui_config) quicksettings_names = {x: i for i, x in enumerate(shared.opts.quicksettings_list) if x != 'quicksettings'} +# companion settings render beside their parent wherever it lives: +# quicksettings the parent and the companion follows, keeping its own visibility +companion_settings = (('sd_unet', 'sd_unet_secondary'),) +for parent_key, companion_key in companion_settings: + if parent_key in quicksettings_names and companion_key not in quicksettings_names: + quicksettings_names[companion_key] = quicksettings_names[parent_key] + 0.5 quicksettings_list = [] hidden_list = [] components = [] +# settings with model-dependent visibility: their wrapper group (settings page) or +# refresh button (quicksettings) is registered here so visibility pushes toggle the +# whole control, not just the inner dropdown +dynamic_visibility_keys = ('sd_unet_secondary',) +dynamic_visibility: dict = {} def apply_setting(key, value): @@ -76,11 +87,15 @@ def create_setting_component(key, is_quicksettings=False): dirtyable_setting = gr.Group(elem_classes="dirtyable", visible=args.get("visible", True)) dirtyable_setting.__enter__() dirty_indicator = gr.Button("", elem_classes="modification-indicator", elem_id=f"modification_indicator_{key}") + if key in dynamic_visibility_keys: + dynamic_visibility.setdefault(key, []).append(dirtyable_setting) if info.refresh is not None: if is_quicksettings: res = comp(label=info.label, value=fun(), elem_id=elem_id, **args) - ui_common.create_refresh_button(res, info.refresh, info.component_args, f"settings_{key}_refresh") + refresh_button = ui_common.create_refresh_button(res, info.refresh, info.component_args, f"settings_{key}_refresh", visible=args.get("visible", True)) + if key in dynamic_visibility_keys: + dynamic_visibility.setdefault(key, []).append(refresh_button) else: with gr.Row(): res = comp(label=info.label, value=fun(), elem_id=elem_id, **args) @@ -202,7 +217,7 @@ def run_settings_single(value, key, progress=False, force=False): shared.opts.save(silent=True) if key == 'sd_text_encoder': sd_models.reload_text_encoder() # apply the change now; reloads the model for encoders with no in-place swap - if key not in ['sd_model_checkpoint', 'sd_model_refiner', 'sd_vae', 'sd_te', 'sd_unet'] or force: + if key not in ['sd_model_checkpoint', 'sd_model_refiner', 'sd_vae', 'sd_te', 'sd_unet', 'sd_unet_secondary'] or force: log.debug(f'Setting changed: {key}="{value}" progress={progress} force={force}') return get_value_for_setting(key), shared.opts.dumpjson() @@ -267,6 +282,7 @@ def create_ui(disabled_tabs=None): with gr.Tabs(elem_id="settings"): quicksettings_list.clear() + dynamic_visibility.clear() for (section_id, section_text) in sections: items = [item for item in shared.opts.data_labels.items() if item[1].section[0] == section_id] # find all items in this section hidden = section_id is None or 'hidden' in section_id.lower() or 'hidden' in section_text.lower() @@ -364,7 +380,7 @@ def create_quicksettings(interfaces): with gr.Row(elem_id="quicksettings", variant="compact"): quicksetting_components = [] quicksetting_keys = [] - for k, _item in sorted(quicksettings_list, key=lambda x: quicksettings_names.get(x[1], x[0])): + for k, _item in sorted(quicksettings_list, key=lambda x: quicksettings_names.get(x[0], 0)): component = create_setting_component(k, is_quicksettings=True) quicksetting_components.append(component) quicksetting_keys.append(k) @@ -394,10 +410,20 @@ def create_quicksettings(interfaces): 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_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 + # a checkpoint change can reset sd_unet / sd_unet_secondary / sd_text_encoder + # (arch changed); push the current values back to the dropdowns. The secondary + # dropdown and its dynamic_visibility companions also toggle visibility here, + # since get_value_for_setting strips 'visible'. + from modules import sd_unet, shared_items 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'), get_value_for_setting('sd_text_encoder'), settings_text + secondary_visible = shared.sd_model_type in sd_unet.DUAL_TRANSFORMER_TYPES + secondary_update = gr.update( + value=shared.opts.sd_unet_secondary, + choices=shared_items.sd_unet_items(), + visible=secondary_visible, + ) + companion_updates = [gr.update(visible=secondary_visible) for _ in dynamic_visibility.get('sd_unet_secondary', [])] + return checkpoint_update, get_value_for_setting('sd_unet'), secondary_update, *companion_updates, get_value_for_setting('sd_text_encoder'), settings_text for k, _item in quicksettings_list: component = shared.settings_components[k] @@ -417,7 +443,7 @@ def create_quicksettings(interfaces): if k == 'sd_model_checkpoint': def fn(value, progress=progress_flag): return sync_checkpoint_components(value, progress=progress) - outputs = [component, shared.settings_components['sd_unet'], shared.settings_components['sd_text_encoder'], text_settings] + outputs = [component, shared.settings_components['sd_unet'], shared.settings_components['sd_unet_secondary'], *dynamic_visibility.get('sd_unet_secondary', []), 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) @@ -438,7 +464,7 @@ def create_quicksettings(interfaces): 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'], shared.settings_components['sd_text_encoder'], text_settings], + outputs=[shared.settings_components['sd_model_checkpoint'], shared.settings_components['sd_unet'], shared.settings_components['sd_unet_secondary'], *dynamic_visibility.get('sd_unet_secondary', []), 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( diff --git a/pipelines/generic_transformer.py b/pipelines/generic_transformer.py index 99b77a716..9de674781 100644 --- a/pipelines/generic_transformer.py +++ b/pipelines/generic_transformer.py @@ -21,10 +21,16 @@ def load_transformer( modules_dtype_dict=None, use_safetensors=True, native_spec=None, + override_slot='primary', **kwargs): """Load a DiT transformer from the base repo, or from a user-selected - single file when the UNET dropdown (``shared.opts.sd_unet``) is set. + single file when the slot's UNET override dropdown is set. + + ``override_slot`` selects which dropdown this call consumes: ``'primary'`` + reads ``shared.opts.sd_unet``, ``'secondary'`` reads + ``shared.opts.sd_unet_secondary`` (dual-transformer arches give each + transformer its own slot). With ``native_spec`` set and a .safetensors override selected, dispatches to :func:`pipelines.native_transformer.load`. Without a spec, a single-file @@ -73,12 +79,19 @@ def load_transformer( fallback = True from modules import sd_unet - if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - if shared.opts.sd_unet not in list(sd_unet.unet_dict): - log.error(f'Load module: type=transformer file="{shared.opts.sd_unet}" not found') - elif os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]): - local_file = sd_unet.unet_dict[shared.opts.sd_unet] - override_name = shared.opts.sd_unet + if override_slot == 'primary': + override_opt, tracker_attr = 'sd_unet', 'loaded_unet' + elif override_slot == 'secondary': + override_opt, tracker_attr = 'sd_unet_secondary', 'loaded_unet_secondary' + else: + raise ValueError(f'load_transformer: unknown override_slot={override_slot}') + selected = getattr(shared.opts, override_opt, None) + if selected is not None and selected != 'Default': + if selected not in list(sd_unet.unet_dict): + log.error(f'Load module: type=transformer slot={override_slot} file="{selected}" not found') + elif os.path.exists(sd_unet.unet_dict[selected]): + local_file = sd_unet.unet_dict[selected] + override_name = selected if repo_id.startswith(shared.opts.ckpt_dir) and os.path.exists(repo_id): log.error(f'Load model: transformer="{repo_id}" is incorrectly placed in the checkpoints folder') @@ -120,8 +133,8 @@ def load_transformer( log.warning(f'Load model: transformer="{local_file}" override incompatible with cls={cls_name.__name__} ({e})') if fallback: log.warning(f'Load model: transformer="{local_file}" ignoring override and loading base transformer') - shared.opts.data['sd_unet'] = 'Default' - sd_unet.loaded_unet = None + shared.opts.data[override_opt] = 'Default' + setattr(sd_unet, tracker_attr, None) transformer = load_from_repo() # 3. load safetensors with diffusers loader @@ -144,10 +157,10 @@ def load_transformer( else: transformer = load_from_repo() - # mark the dropdown selection as loaded so the sd_unet onchange callback + # mark the dropdown selection as loaded so the slot's onchange callback # does not force a redundant full reload for an already-consumed override - if transformer is not None and override_name is not None and shared.opts.sd_unet == override_name: - sd_unet.loaded_unet = override_name + if transformer is not None and override_name is not None and getattr(shared.opts, override_opt, None) == override_name: + setattr(sd_unet, tracker_attr, override_name) sd_models.allow_post_quant = False # we already handled it if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: @@ -173,12 +186,15 @@ def load_transformer( log.debug(f'Load model: transformer="{repo_id}" quant="{quant_type}" size={module_size:.3f} params={param_num:.3f} memory={module_memory}') try: - actual_dtype = transformer.dtype - if isinstance(actual_dtype, torch.dtype) and isinstance(dtype, torch.dtype) and actual_dtype != dtype: - force = shared.opts.force_dtype - log.warning(f'Load model: transformer="{repo_id}" dtype desired={dtype} actual={actual_dtype} force={force}') - if force: - transformer = transformer.to(dtype) + # quantized models legitimately report the storage dtype (e.g. fp8 comfy_quant + # adopted via SDNQ); the compute dtype lives in the dequantizers, not the params + if getattr(transformer, 'quantization_config', None) is None: + actual_dtype = transformer.dtype + if isinstance(actual_dtype, torch.dtype) and isinstance(dtype, torch.dtype) and actual_dtype != dtype: + force = shared.opts.force_dtype + log.warning(f'Load model: transformer="{repo_id}" dtype desired={dtype} actual={actual_dtype} force={force}') + if force: + transformer = transformer.to(dtype) except Exception: pass diff --git a/pipelines/ideogram/__init__.py b/pipelines/ideogram/__init__.py new file mode 100644 index 000000000..f0b6585f7 --- /dev/null +++ b/pipelines/ideogram/__init__.py @@ -0,0 +1,62 @@ +"""Ideogram 4 native loader spec. + +Kept import-light: only the converter and spec live here so the native +loader can import them without pulling the pipeline or the qwen patch. +""" + +import re + +import torch +import diffusers + +from pipelines.native_transformer import TransformerSpec + + +QKV_TARGETS = ("to_q", "to_k", "to_v") +QKV_RE = re.compile(r"^(layers\.\d+\.attention)\.qkv\.(.+)$") +OUT_RE = re.compile(r"^(layers\.\d+\.attention)\.o\.(.+)$") + + +def convert_ideogram4_transformer_checkpoint(state_dict: dict, **kwargs) -> dict: # pylint: disable=unused-argument + """Fused community layout to ``Ideogram4Transformer2DModel`` layout. + + ``layers.N.attention.qkv.`` splits into ``to_q/to_k/to_v.``: + row-stacked tensors (weight, row-wise scale, bias) are dim-0 sliced into + thirds, non-row-stacked sidecars (markers, scalar scales) copied verbatim + to all three. ``attention.o.`` renames to ``attention.to_out.0.``, + everything else passes through. Returns a new dict; the input (which may + be the cached state dict) is not mutated. + """ + converted: dict = {} + for key, value in state_dict.items(): + fused_match = QKV_RE.match(key) + if fused_match is not None: + prefix, suffix = fused_match.group(1), fused_match.group(2) + fused_weight = state_dict.get(f"{prefix}.qkv.weight") + fused_rows = fused_weight.shape[0] if torch.is_tensor(fused_weight) and fused_weight.ndim >= 1 else None + row_stacked = ( + suffix != "comfy_quant" + and torch.is_tensor(value) and value.ndim >= 1 + and fused_rows is not None and value.shape[0] == fused_rows and fused_rows % 3 == 0 + ) + if row_stacked: + third = fused_rows // 3 + for i, target in enumerate(QKV_TARGETS): + converted[f"{prefix}.{target}.{suffix}"] = value[i * third:(i + 1) * third] + else: + for target in QKV_TARGETS: + converted[f"{prefix}.{target}.{suffix}"] = value + continue + out_match = OUT_RE.match(key) + if out_match is not None: + converted[f"{out_match.group(1)}.to_out.0.{out_match.group(2)}"] = value + continue + converted[key] = value + return converted + + +IDEOGRAM4_SPEC = TransformerSpec( + cls=diffusers.Ideogram4Transformer2DModel, + converter=convert_ideogram4_transformer_checkpoint, + converter_handles_quant=True, +) diff --git a/pipelines/model_ideogram4.py b/pipelines/model_ideogram4.py index cbb09e5d3..3dcdac0c1 100644 --- a/pipelines/model_ideogram4.py +++ b/pipelines/model_ideogram4.py @@ -1,3 +1,5 @@ +import dataclasses + import diffusers import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae @@ -41,11 +43,16 @@ def load_ideogram4(checkpoint_info, diffusers_load_config=None): if repo_id is None or repo_id.lower() == 'none': return None + from pipelines.ideogram import IDEOGRAM4_SPEC transformer_cls = diffusers.Ideogram4Transformer2DModel - transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="transformer", load_config=diffusers_load_config) + transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="transformer", load_config=diffusers_load_config, native_spec=IDEOGRAM4_SPEC) if shared.opts.model_ideogram4_enable_cg: - unconditional_transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, subfolder="unconditional_transformer", load_config=diffusers_load_config) + # the spec subfolder swap makes an unconditional override fetch its config from the matching subfolder + unconditional_transformer = generic.load_transformer( + repo_id, cls_name=transformer_cls, subfolder="unconditional_transformer", load_config=diffusers_load_config, + native_spec=dataclasses.replace(IDEOGRAM4_SPEC, subfolder="unconditional_transformer"), override_slot='secondary', + ) else: unconditional_transformer = None text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLModel, load_config=diffusers_load_config) diff --git a/pipelines/model_wanai.py b/pipelines/model_wanai.py index 2a7a421ea..9a9a1bdda 100644 --- a/pipelines/model_wanai.py +++ b/pipelines/model_wanai.py @@ -42,7 +42,7 @@ def load_wan(checkpoint_info, diffusers_load_config=None): boundary_ratio = 1000.0 elif shared.opts.model_wan_stage == 'combined' or shared.opts.model_wan_stage == 'both': transformer = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer') - transformer_2 = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer_2') + transformer_2 = generic.load_transformer(repo_id, cls_name=transformer_cls, load_config=diffusers_load_config, subfolder='transformer_2', override_slot='secondary') # load with the checkpoint's boundary; the slider override is applied at runtime in set_pipeline_args boundary_ratio = None else: diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index ab3a0e42b..2fbf3f797 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -128,6 +128,10 @@ class TransformerSpec: VAE) that all-in-one exports bundle alongside the transformer; their keys are dropped before prefix detection rather than treated as a malformed file. + + ``converter_handles_quant`` runs the converter before comfy_quant + detection; such converters must translate marker/scale sidecar keys along + with the weights. Float-oriented converters keep the default. """ cls: type @@ -135,6 +139,7 @@ class TransformerSpec: prefixes: tuple[str, ...] = DEFAULT_PREFIXES ignored_prefixes: tuple[str, ...] = DEFAULT_IGNORED_PREFIXES converter: Callable[[dict], dict] | None = None + converter_handles_quant: bool = False siblings: dict[str, SiblingSpec] = field(default_factory=dict) acceptable_missing: tuple[str, ...] = DEFAULT_ACCEPTABLE_MISSING zero_init_missing: tuple[str, ...] = () @@ -298,6 +303,7 @@ def load( dtype=effective_dtype, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict, + converter_handles_quant=spec.converter_handles_quant, **kwargs, ) del transformer_sd @@ -422,9 +428,9 @@ def check_forbidden_markers( def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str] | None: - """Detect ComfyUI ``comfy_quant`` pre-quantized layers in a state dict. + """Detect ``comfy_quant`` pre-quantized layers in a state dict. - ComfyUI's quantized checkpoints mark each quantized layer with a + ComfyUI-format quantized checkpoints mark each quantized layer with a ``.comfy_quant`` uint8 tensor whose bytes are a JSON object naming the storage format, alongside a ``.weight_scale`` tensor. Returns the set of marked module names and the format string, or ``None`` when no @@ -646,28 +652,17 @@ def build_component_prequantized( **kwargs, ) -> object: """Build a component from a comfy_quant pre-quantized state dict, mapping - the marked layers onto SDNQ quantized layers without dequantizing. + the marked layers onto SDNQ layers without dequantizing. - The supported comfy formats (``int8_tensorwise``, ``float8_e4m3fn``) are - strict subsets of SDNQ's symmetric quantization for the corresponding - weights dtype: same storage layout (unpacked 8-bit ``[out, in]``), same - dequant math (``weight * scale``, no zero point), so the file's tensors - are adopted bit-exact. The model is built under ``init_empty_weights`` - and each marked Linear is swapped for an SDNQ wrapper with per-tensor - dequant geometry (``group_size=-1``, scalar scale); the file dictates - which layers are quantized, independent of the user's quantization - settings. Weights load through ``SDNQQuantizer``'s pre-quantized path, - which preserves the quantized weight dtype and the fp32 scale (unlike - :func:`build_component_quantized`, floating-point SDNQ params are - deliberately not cast to the target dtype). - - Layers are assembled in canonical dequant layout first; - ``apply_sdnq_options_to_model`` then enables quantized matmul per the - user's settings (transposing eligible layers), matching the order used by + The supported formats are subsets of SDNQ's symmetric quantization: same + storage layout (unpacked 8-bit ``[out, in]``), same dequant math + (``weight * scale``, no zero point), so tensors are adopted bit-exact. + The file dictates which layers are quantized, independent of the user's + quantization settings, and floating-point SDNQ params are not cast to the + target dtype (the fp32 scales must survive). Layers are assembled in + canonical dequant layout; ``apply_sdnq_options_to_model`` then applies + the user's quantized-matmul settings, matching ``modules.sdnq.loader.load_sdnq_model``. - - Caller is responsible for prefix stripping; the state dict arrives with - its comfy marker keys intact and is remapped here. """ import rich.progress as rp from accelerate import init_empty_weights @@ -803,6 +798,22 @@ def build_component_prequantized( return component +def apply_converter(converter: Callable[[dict], dict], state_dict: dict, cls: type, component_name: str) -> dict: + """Run a spec converter, wrapping any failure as + :class:`OverrideArchMismatch` (with the original chained) so a wrong-arch + file degrades to the base-repo fallback instead of a raw converter crash. + """ + log.debug(f'Load model: transformer=native {component_name} converter={converter.__name__} keys={len(state_dict)}') + try: + return converter(state_dict) + except Exception as e: + raise OverrideArchMismatch( + f"Load model: type={cls.__name__} native_transformer converter " + f"{converter.__name__} rejected the override ({type(e).__name__}: {e}); " + f"file does not look like a {cls.__name__} checkpoint" + ) from e + + def build_component( *, component_name: str, @@ -817,21 +828,21 @@ def build_component( dtype=None, modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None, + converter_handles_quant: bool = False, **kwargs, ) -> object: """Convert (if needed), instantiate, load weights, dtype-cast, quantize, and offload-place a single component. Raises on any hard failure. - Transformer state dicts carrying ComfyUI ``comfy_quant`` markers are - dispatched to :func:`build_component_prequantized`, which adopts the - file's int8 tensors as SDNQ layers regardless of quantization settings - (the converter and ``quant_args`` are bypassed: the file is already - quantized). For the transformer component under SDNQ, the per-tensor - pre-mode path in :func:`build_component_quantized` is used so - quantization is applied in flight (one layer's worth of bf16 in memory - at a time). All other cases (siblings, non-quantized loads, - NVIDIAModelOptConfig, layerwise quant) go through the standard - load_state_dict + post-quantize path. + Transformer state dicts carrying ``comfy_quant`` markers dispatch + to :func:`build_component_prequantized` (``quant_args`` are bypassed: the + file is already quantized); a ``converter_handles_quant`` converter runs + before that detection, float-oriented converters after it. Under SDNQ the + transformer uses the per-tensor pre-mode path in + :func:`build_component_quantized` so quantization is applied in flight. + All other cases (siblings, non-quantized loads, NVIDIAModelOptConfig, + layerwise quant) go through the standard load_state_dict + post-quantize + path. ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global default is used. ``modules_to_not_convert`` and ``modules_dtype_dict`` @@ -840,6 +851,10 @@ def build_component( reach ``cls.from_config`` for both construction paths. """ try: + if converter is not None and converter_handles_quant and component_name == "transformer": + state_dict = apply_converter(converter, state_dict, cls, component_name) + converter = None # consumed; must not run again on the non-comfy path below + comfy_quant = detect_comfy_quant(state_dict, cls.__name__) if component_name == "transformer" else None if comfy_quant is not None: marked_names, comfy_format = comfy_quant @@ -863,15 +878,7 @@ def build_component( return component if converter is not None: - log.debug(f'Load model: transformer=native {component_name} converter={converter.__name__} keys={len(state_dict)}') - try: - sd = converter(state_dict) - except Exception as e: - raise OverrideArchMismatch( - f"Load model: type={cls.__name__} native_transformer converter " - f"{converter.__name__} rejected the override ({type(e).__name__}: {e}); " - f"file does not look like a {cls.__name__} checkpoint" - ) from e + sd = apply_converter(converter, state_dict, cls, component_name) else: sd = state_dict diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 616fc8c5a..a8cd14687 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -8,7 +8,7 @@ Covers the pure helpers that own per-arch knob handling: - ``strip_prefix`` for single/multi prefix detection and mixed-prefix rejection - ``partition_siblings`` for inline-sibling key partitioning - ``check_forbidden_markers`` for structural-mismatch rejection -- ``detect_comfy_quant`` for ComfyUI comfy_quant marker detection and format gating +- ``detect_comfy_quant`` for comfy_quant marker detection and format gating - ``remap_comfy_quant`` for comfy_quant -> SDNQ key translation - ``is_noop_converter`` for diffusers no-op lambda detection - ``validate_state_dict_load`` for unexpected / missing key handling @@ -954,6 +954,104 @@ def test_load_converter_crash_raises_mismatch(): os.unlink(path) +# ============================================================ +# Ideogram 4 converter (fused community layout -> diffusers layout) +# ============================================================ + +def ideogram_converter(): + from pipelines.ideogram import convert_ideogram4_transformer_checkpoint + return convert_ideogram4_transformer_checkpoint + + +def test_ideogram4_converter_splits_fused_qkv_weight(): + convert = ideogram_converter() + fused = torch.arange(48, dtype=torch.float32).reshape(12, 4) + out = convert({'layers.0.attention.qkv.weight': fused}) + assert set(out.keys()) == {f'layers.0.attention.{n}.weight' for n in ('to_q', 'to_k', 'to_v')} + assert torch.equal(out['layers.0.attention.to_q.weight'], fused[0:4]) + assert torch.equal(out['layers.0.attention.to_k.weight'], fused[4:8]) + assert torch.equal(out['layers.0.attention.to_v.weight'], fused[8:12]) + + +def test_ideogram4_converter_splits_rowwise_scale(): + convert = ideogram_converter() + fused = torch.zeros((12, 4), dtype=torch.int8) + scale = torch.arange(12, dtype=torch.float32).reshape(12, 1) + out = convert({'layers.3.attention.qkv.weight': fused, 'layers.3.attention.qkv.weight_scale': scale}) + assert torch.equal(out['layers.3.attention.to_q.weight_scale'], scale[0:4]) + assert torch.equal(out['layers.3.attention.to_v.weight_scale'], scale[8:12]) + assert out['layers.3.attention.to_k.weight_scale'].shape == (4, 1) + + +def test_ideogram4_converter_scalar_scale_copied_not_sliced(): + convert = ideogram_converter() + fused = torch.zeros((12, 4), dtype=torch.int8) + scale = torch.tensor(0.5) + out = convert({'layers.0.attention.qkv.weight': fused, 'layers.0.attention.qkv.weight_scale': scale}) + for name in ('to_q', 'to_k', 'to_v'): + assert out[f'layers.0.attention.{name}.weight_scale'] is scale + + +def test_ideogram4_converter_duplicates_marker(): + convert = ideogram_converter() + fused = torch.zeros((12, 4), dtype=torch.int8) + marker = comfy_marker('int8_tensorwise') + out = convert({'layers.0.attention.qkv.weight': fused, 'layers.0.attention.qkv.comfy_quant': marker}) + for name in ('to_q', 'to_k', 'to_v'): + assert out[f'layers.0.attention.{name}.comfy_quant'] is marker + + +def test_ideogram4_converter_renames_o_with_sidecars(): + convert = ideogram_converter() + sd = { + 'layers.5.attention.o.weight': torch.zeros((4, 4), dtype=torch.int8), + 'layers.5.attention.o.weight_scale': torch.tensor(1.0), + 'layers.5.attention.o.comfy_quant': comfy_marker('int8_tensorwise'), + } + out = convert(sd) + assert set(out.keys()) == { + 'layers.5.attention.to_out.0.weight', + 'layers.5.attention.to_out.0.weight_scale', + 'layers.5.attention.to_out.0.comfy_quant', + } + + +def test_ideogram4_converter_passthrough(): + convert = ideogram_converter() + sd = { + 'input_proj.weight': torch.zeros(4), + 'llm_cond_proj.weight': torch.zeros(4), + 'layers.0.feed_forward.w1.weight': torch.zeros(4), + 'layers.0.attention.norm_q.weight': torch.zeros(4), + 'layers.0.adaln_modulation.bias': torch.zeros(4), + 'final_layer.linear.weight': torch.zeros(4), + } + out = convert(sd) + assert set(out.keys()) == set(sd.keys()) + for k, v in sd.items(): + assert out[k] is v + + +def test_ideogram4_converter_defensive_bias_split(): + convert = ideogram_converter() + fused = torch.zeros((12, 4), dtype=torch.float32) + bias = torch.arange(12, dtype=torch.float32) + out = convert({'layers.0.attention.qkv.weight': fused, 'layers.0.attention.qkv.bias': bias}) + assert torch.equal(out['layers.0.attention.to_k.bias'], bias[4:8]) + + +def test_ideogram4_converter_does_not_mutate_input(): + convert = ideogram_converter() + sd = { + 'layers.0.attention.qkv.weight': torch.zeros((12, 4)), + 'layers.0.attention.o.weight': torch.zeros((4, 4)), + 'input_proj.weight': torch.zeros(4), + } + keys_before = set(sd.keys()) + convert(sd) + assert set(sd.keys()) == keys_before + + # ============================================================ # Integration: comfy_quant pre-quantized load # ============================================================ @@ -1238,6 +1336,204 @@ def test_load_transformer_syncs_loaded_unet(): os.unlink(path) +class MockFusedAttention(torch.nn.Module): + """Split-attention module tree matching Ideogram4Transformer2DModel's + naming (to_q/to_k/to_v + to_out ModuleList), fed by fused checkpoints.""" + + def __init__(self, dim: int): + super().__init__() + self.to_q = torch.nn.Linear(dim, dim, bias=False) + self.to_k = torch.nn.Linear(dim, dim, bias=False) + self.to_v = torch.nn.Linear(dim, dim, bias=False) + self.to_out = torch.nn.ModuleList([torch.nn.Linear(dim, dim, bias=False), torch.nn.Dropout(0.0)]) + + +class MockFusedBlock(torch.nn.Module): + def __init__(self, dim: int): + super().__init__() + self.attention = MockFusedAttention(dim) + + +class MockFusedTransformer(torch.nn.Module): + @classmethod + def from_config(cls, config: dict, **kwargs) -> 'MockFusedTransformer': + return cls(dim=config['dim']) + + def __init__(self, dim: int): + super().__init__() + self.layers = torch.nn.ModuleList([MockFusedBlock(dim)]) + self.input_proj = torch.nn.Linear(dim, dim) + + +def fused_spec(): + from pipelines.ideogram import convert_ideogram4_transformer_checkpoint + return nt.TransformerSpec( + cls=MockFusedTransformer, + converter=convert_ideogram4_transformer_checkpoint, + converter_handles_quant=True, + ) + + +def fused_fixture(dim: int, quantized: bool = True) -> dict: + if quantized: + qkv = torch.randint(-128, 127, (3 * dim, dim), dtype=torch.int8) + else: + qkv = torch.randn(3 * dim, dim, dtype=torch.float16) + raw = { + 'model.diffusion_model.layers.0.attention.qkv.weight': qkv, + 'model.diffusion_model.layers.0.attention.o.weight': torch.randn(dim, dim, dtype=torch.float16), + 'model.diffusion_model.input_proj.weight': torch.randn(dim, dim, dtype=torch.float16), + 'model.diffusion_model.input_proj.bias': torch.zeros(dim, dtype=torch.float16), + } + if quantized: + raw['model.diffusion_model.layers.0.attention.qkv.weight_scale'] = torch.rand(3 * dim, 1, dtype=torch.float32) + raw['model.diffusion_model.layers.0.attention.qkv.comfy_quant'] = comfy_marker('int8_tensorwise') + return raw + + +def test_load_fused_comfy_end_to_end(): + """Quant-aware converter + comfy adoption: fused int8 qkv with row-wise + scales lands as three SDNQ linears holding the exact row slices.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = fused_fixture(dim, quantized=True) + write_fixture(raw, fd, path) + with ComfyTestEnv(dim): + transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}, dtype=torch.float32) + + attn = transformer.layers[0].attention + fused_w = raw['model.diffusion_model.layers.0.attention.qkv.weight'] + fused_s = raw['model.diffusion_model.layers.0.attention.qkv.weight_scale'] + for i, name in enumerate(('to_q', 'to_k', 'to_v')): + layer = getattr(attn, name) + assert layer.__class__.__name__ == 'SDNQLinear', f'{name} is {layer.__class__.__name__}' + assert layer.weight.dtype == torch.int8 + assert torch.equal(layer.weight.detach().cpu(), fused_w[i * dim:(i + 1) * dim]) + assert tuple(layer.scale.shape) == (dim, 1), f'{name} scale shape {layer.scale.shape}' + assert torch.equal(layer.scale.detach().cpu(), fused_s[i * dim:(i + 1) * dim]) + expected = fused_w[i * dim:(i + 1) * dim].float() * fused_s[i * dim:(i + 1) * dim] + dequantized = layer.sdnq_dequantizer(layer.weight, layer.scale, zero_point=None, svd_up=None, svd_down=None) + assert torch.equal(dequantized.detach().cpu(), expected) + # unmarked o.weight passes through the rename as a plain Linear + assert attn.to_out[0].__class__ is torch.nn.Linear + assert torch.allclose(attn.to_out[0].weight.detach().cpu(), raw['model.diffusion_model.layers.0.attention.o.weight'].float()) + assert transformer.input_proj.__class__ is torch.nn.Linear + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_fused_bf16_end_to_end(): + """Converter-first on a plain float fused file: no markers, standard load + path, split weights land on the class-native linears.""" + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = fused_fixture(dim, quantized=False) + write_fixture(raw, fd, path) + with ComfyTestEnv(dim): + transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}, dtype=torch.float32) + + attn = transformer.layers[0].attention + fused_w = raw['model.diffusion_model.layers.0.attention.qkv.weight'] + for i, name in enumerate(('to_q', 'to_k', 'to_v')): + layer = getattr(attn, name) + assert layer.__class__ is torch.nn.Linear + assert torch.allclose(layer.weight.detach().cpu(), fused_w[i * dim:(i + 1) * dim].float()) + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_load_fused_convrot_falls_back(): + """The one real Ideogram 4 civitai file is convrot-flagged: conversion + succeeds, then detection rejects it into the base-repo fallback.""" + import json as json_mod + fd, path = tempfile.mkstemp(suffix='.safetensors') + try: + dim = 8 + raw = fused_fixture(dim, quantized=True) + payload = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256, 'per_row': True}) + raw['model.diffusion_model.layers.0.attention.qkv.comfy_quant'] = torch.tensor(list(payload.encode()), dtype=torch.uint8) + write_fixture(raw, fd, path) + with ComfyTestEnv(dim): + raised = False + try: + nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}) + except nt.OverrideArchMismatch as e: + raised = True + assert 'convrot' in str(e) + assert raised, 'expected OverrideArchMismatch' + finally: + if os.path.exists(path): + os.unlink(path) + + +def test_build_component_comfy_skips_converter_when_not_quant_aware(): + """With converter_handles_quant=False (the default for every other arch), + a comfy state dict must reach the prequantized path without the converter + ever running; a crashing converter proves it was not invoked.""" + dim = 8 + raw = comfy_fixture(dim) + sd = {k[len('model.diffusion_model.'):]: v for k, v in raw.items()} + with ComfyTestEnv(dim): + component = nt.build_component( + component_name='transformer', + state_dict=sd, + config={'dim': dim}, + cls=MockMiniTransformer, + converter=crashing_converter, + acceptable_missing=('rope.',), + quant_args={}, + quant_type='SDNQConfig', + dtype=torch.float32, + ) + assert component.in_proj.__class__.__name__ == 'SDNQLinear' + + +def test_load_transformer_secondary_slot_syncs_tracker(): + """The secondary slot consumes sd_unet_secondary and syncs its own + tracker without touching the primary slot.""" + from modules import sd_unet, shared + from pipelines import generic_transformer as gt + + fd, path = tempfile.mkstemp(suffix='.safetensors') + dim = 8 + raw = { + 'model.diffusion_model.in_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.in_proj.bias': torch.zeros(dim), + 'model.diffusion_model.out_proj.weight': torch.randn(dim, dim), + 'model.diffusion_model.out_proj.bias': torch.zeros(dim), + } + write_fixture(raw, fd, path) + + orig_primary_opt = shared.opts.sd_unet + orig_secondary_opt = shared.opts.sd_unet_secondary + orig_primary = sd_unet.loaded_unet + orig_secondary = sd_unet.loaded_unet_secondary + sd_unet.unet_dict['mock-unet-2'] = path + shared.opts.data['sd_unet'] = 'Default' + shared.opts.data['sd_unet_secondary'] = 'mock-unet-2' + sd_unet.loaded_unet = None + sd_unet.loaded_unet_secondary = None + try: + with ComfyTestEnv(dim): + spec = nt.TransformerSpec(cls=MockMiniTransformer) + transformer = gt.load_transformer('fake/repo', cls_name=MockMiniTransformer, native_spec=spec, override_slot='secondary') + assert transformer is not None + assert sd_unet.loaded_unet_secondary == 'mock-unet-2' + assert sd_unet.loaded_unet is None, 'primary tracker must stay untouched' + finally: + sd_unet.unet_dict.pop('mock-unet-2', None) + shared.opts.data['sd_unet'] = orig_primary_opt + shared.opts.data['sd_unet_secondary'] = orig_secondary_opt + sd_unet.loaded_unet = orig_primary + sd_unet.loaded_unet_secondary = orig_secondary + if os.path.exists(path): + os.unlink(path) + + def test_build_component_comfy_preempts_sdnq_fresh_quant(): """When SDNQ on-load quant settings are active (quant_type=SDNQConfig), a comfy_quant file must still take the pre-quantized path: fresh quant of @@ -1395,7 +1691,26 @@ def run_all(): test_load_comfy_unsupported_format_raises_mismatch, test_load_comfy_marker_for_unknown_module_raises_mismatch, test_load_transformer_syncs_loaded_unet, + test_load_transformer_secondary_slot_syncs_tracker, test_build_component_comfy_preempts_sdnq_fresh_quant, + test_build_component_comfy_skips_converter_when_not_quant_aware, + ]: + run_test(cat, fn) + + log.warning('=== ideogram4 converter / fused load ===') + cat = category('ideogram4') + for fn in [ + test_ideogram4_converter_splits_fused_qkv_weight, + test_ideogram4_converter_splits_rowwise_scale, + test_ideogram4_converter_scalar_scale_copied_not_sliced, + test_ideogram4_converter_duplicates_marker, + test_ideogram4_converter_renames_o_with_sidecars, + test_ideogram4_converter_passthrough, + test_ideogram4_converter_defensive_bias_split, + test_ideogram4_converter_does_not_mutate_input, + test_load_fused_comfy_end_to_end, + test_load_fused_bf16_end_to_end, + test_load_fused_convrot_falls_back, ]: run_test(cat, fn) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 8ec96e8d4..21a90ce7b 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -1561,6 +1561,7 @@ {"id":"","label":"Use spaces","localized":"","hint":"Replace underscores with spaces in tag output.
Some prompt systems prefer spaces between words (e.g., 'long hair') while others use underscores (e.g., 'long_hair').","ui":"caption"}, {"id":"","label":"Username","localized":"","hint":"","ui":"component-8823"}, {"id":"","label":"UNET model","localized":"","hint":"","ui":"settings_sd"}, + {"id":"","label":"UNET model secondary","localized":"","hint":"Override for the second transformer of dual-transformer architectures:
- Ideogram 4: the unconditional transformer
- Wan combined stage: the second expert (transformer_2)

Shown only when the loaded model has a second transformer.
Default loads the component from the base model.","ui":"settings_sd"}, {"id":"","label":"Use torch streams","localized":"","hint":"In group offload, uses CUDA streams to prefetch the next group while the current one runs, hiding transfer latency.
Faster, but uses more VRAM.

Applies only to group offload.

Disabled by default.","reload":"model","ui":"settings_offload"}, {"id":"","label":"Use SVD quantization","localized":"","hint":"Adds a low-rank (SVDQuant) correction on top of SDNQ to recover accuracy lost at low bit widths, at the cost of extra size and compute.
Tuned by SVD rank size and SVD steps.

Disabled by default.","reload":"model","ui":"settings_quantization"}, {"id":"","label":"Use Dynamic quantization","localized":"","hint":"Picks a per-layer weight type automatically instead of one type everywhere. Each layer starts at the Quantization type (the minimum) and steps up to higher precision until its error meets the Dynamic loss threshold.
Protects error-sensitive layers at the cost of a larger model and slower load.

Disabled by default.","reload":"model","ui":"settings_quantization"}, diff --git a/webui.py b/webui.py index 92ebe6eff..480bb03ff 100644 --- a/webui.py +++ b/webui.py @@ -184,6 +184,7 @@ def load_model(): shared.opts.onchange("sd_model_refiner", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(op='refiner')), call=False) shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False) shared.opts.onchange("sd_unet", wrap_queued_call(lambda: modules.sd_unet.load_unet(shared.sd_model)), call=False) + shared.opts.onchange("sd_unet_secondary", wrap_queued_call(lambda: modules.sd_unet.load_unet_secondary(shared.sd_model)), call=False) shared.opts.onchange("sd_text_encoder", wrap_queued_call(lambda: modules.sd_models.reload_text_encoder()), call=False) shared.opts.onchange("temp_dir", modules.gr_tempdir.on_tmpdir_changed) timer.startup.record("onchange")