diff --git a/CHANGELOG.md b/CHANGELOG.md index 257976dd5..f809ec118 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,8 @@ # Change Log for SD.Next -## Update for 2026-05-29 +## Update for 2026-06-01 -### Highlights for 2026-05-29 +### Highlights for 2026-06-01 *What's New?* - **Anima** made it to release version, Microsoft joins the game with **Lens** @@ -14,7 +14,7 @@ And we have new [Home](https://vladmandic.github.io/sdnext/) page and new [Contr Plus continued work on modernization of codebase: UI is now fully TypeScript based and we have a new modular LoRA loader -### Details for 2026-05-29 +### Details for 2026-06-01 - **Models** - [CircleStone Anima 1.0](https://huggingface.co/circlestone-labs/Anima) in *Base* and *Turbo* (distilled) variants @@ -63,6 +63,9 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas see [backends](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends) for list of available attention backends *note* compatiblity matrix between torch backend, torch version and model specifics is relatively small at the moment *note* does not replace existing *attention* settings + - **Shared components** additional support for shared model components + avoids unnecessary downloads and allows to share components between different models + enabled by default, see *settings -> text encoder -> use shared instance* - **Changes** - all **Guidance** params are now set to *-1* by default to allow using model defaults and avoid confusion with different model behaviour log will print default values used by model if not set by user @@ -76,6 +79,8 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas - Automated fixes using `/check-` skills - Automated syntax, spelling and readability improvements to `/wiki` pages - **Internal** + - massive new codebase/refactor to use native transformers loader! + - refactor shared components loader - update `torch==2.12` for *CUDA, ROCm, IPEX* - complete refactor of `core` JavaScript codebase to TypeScript! - complete refactor of `modernui` JavaScript codebase to TypeScript! @@ -105,6 +110,9 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas - `kanvas` image change notification - `reinstall` force reinstal of transformers and diffusers - `ipex` torch install error, thanks @liutyi + - `taesd` preview constant size with reduced layers + - `output path` use correct base folder for initial folders + - `ltx` prompt embeds move to device, thanks @ryanmeador ## Update for 2026-05-13 diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 63cc1907e..bb968ab75 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -205,7 +205,7 @@ def create_settings(cmd_opts): "sd_textencder_linebreak": OptionInfo(True, "Use line break as prompt segment marker", gr.Checkbox), "diffusers_zeros_prompt_pad": OptionInfo(False, "Use zeros for prompt padding", gr.Checkbox), "te_optional_sep": OptionInfo("

Optional

", "", gr.HTML), - "te_shared_t5": OptionInfo(True, "T5: Use shared instance of text encoder"), + "te_shared_te": OptionInfo(True, "Use shared instance of text encoder"), "te_pooled_embeds": OptionInfo(False, "SDXL: Use weighted pooled embeds"), "te_complex_human_instruction": OptionInfo(True, "Sana: Use complex human instructions"), "te_use_mask": OptionInfo(True, "Lumina: Use mask in transformers"), diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py index 819706767..c6d60ccb9 100644 --- a/modules/video_models/video_load.py +++ b/modules/video_models/video_load.py @@ -69,26 +69,26 @@ def load_model(selected: models_def.Model): load_args, quant_args = model_quant.get_dit_args({}, module='TE', device_map=True) # loader deduplication of text-encoder models - if selected.te_cls.__name__ == 'T5EncoderModel' and shared.opts.te_shared_t5: + if selected.te_cls.__name__ == 'T5EncoderModel' and shared.opts.te_shared_te: selected.te = 'Disty0/t5-xxl' selected.te_folder = '' selected.te_revision = None - if selected.te_cls.__name__ == 'UMT5EncoderModel' and shared.opts.te_shared_t5: + if selected.te_cls.__name__ == 'UMT5EncoderModel' and shared.opts.te_shared_te: if 'SDNQ' in selected.name: selected.te = 'Disty0/Wan2.2-T2V-A14B-SDNQ-uint4-svd-r32' else: selected.te = 'Wan-AI/Wan2.2-TI2V-5B-Diffusers' selected.te_folder = 'text_encoder' selected.te_revision = None - if selected.te_cls.__name__ == 'LlamaModel' and shared.opts.te_shared_t5: + if selected.te_cls.__name__ == 'LlamaModel' and shared.opts.te_shared_te: selected.te = 'hunyuanvideo-community/HunyuanVideo' selected.te_folder = 'text_encoder' selected.te_revision = None - if selected.te_cls.__name__ == 'Qwen2_5_VLForConditionalGeneration' and shared.opts.te_shared_t5: + if selected.te_cls.__name__ == 'Qwen2_5_VLForConditionalGeneration' and shared.opts.te_shared_te: selected.te = 'ai-forever/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers' selected.te_folder = 'text_encoder' selected.te_revision = None - if selected.te_cls.__name__ == 'Gemma3ForConditionalGeneration' and shared.opts.te_shared_t5: + if selected.te_cls.__name__ == 'Gemma3ForConditionalGeneration' and shared.opts.te_shared_te: if 'SDNQ' in selected.name: selected.te = 'OzzyGT/LTX-2.3-sdnq-dynamic-int4' else: diff --git a/modules/video_models/video_overrides.py b/modules/video_models/video_overrides.py index 582acc92e..028e544ba 100644 --- a/modules/video_models/video_overrides.py +++ b/modules/video_models/video_overrides.py @@ -37,9 +37,9 @@ def load_override(selected: Model, **load_args): ltx2_connectors_cls = LTX2TextConnectors except ImportError as e: log.warning(f'Video load: LTX2TextConnectors unavailable ({e}); dedup of LTX-2.3 connectors disabled') - if ('LTXVideo 2.3' in selected.name and shared.opts.te_shared_t5 and ltx2_connectors_cls is not None): + if ('LTXVideo 2.3' in selected.name and shared.opts.te_shared_te and ltx2_connectors_cls is not None): conn_repo = 'OzzyGT/LTX-2.3-sdnq-dynamic-int4' if 'SDNQ' in selected.name else 'OzzyGT/LTX-2.3' - log.debug(f'Video load: module=connectors repo="{conn_repo}" cls={ltx2_connectors_cls.__name__} shared={shared.opts.te_shared_t5}') + log.debug(f'Video load: module=connectors repo="{conn_repo}" cls={ltx2_connectors_cls.__name__} shared={shared.opts.te_shared_te}') kwargs['connectors'] = ltx2_connectors_cls.from_pretrained( conn_repo, subfolder='connectors', diff --git a/pipelines/generic.py b/pipelines/generic.py index 8ba75214b..7265ddb86 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -1,320 +1,6 @@ -import os -import sys -import json -import diffusers -import transformers -from modules import shared, devices, errors, sd_models, model_quant -from modules.logger import log +from pipelines.generic_transformer import load_transformer +from pipelines.generic_text_encoder import load_text_encoder +from pipelines.generic_vae import load_vae_override -debug = os.environ.get('SD_LOAD_DEBUG', None) is not None - - -def _loader(component): - """Return loader type for log messages.""" - if sys.platform != 'linux': - return 'default' - if component == 'diffusers': - return 'runai' if shared.opts.runai_streamer_diffusers else 'default' - return 'runai' if shared.opts.runai_streamer_transformers else 'default' - - -def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, native_spec=None, **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. - - With ``native_spec`` set and a .safetensors override selected, dispatches - to :func:`pipelines.native_transformer.load`. Without a spec, a single-file - override falls back to ``from_single_file``. - """ - if shared.state.interrupted: - return None - transformer = None - if load_config is None: - load_config = {} - if modules_to_not_convert is None: - modules_to_not_convert = [] - if modules_dtype_dict is None: - modules_dtype_dict = {} - jobid = shared.state.begin('Load DiT') - try: - load_args, quant_args = model_quant.get_dit_args(load_config, module='Model', device_map=True, allow_quant=allow_quant, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) - quant_type = model_quant.get_quant_type(quant_args) - dtype = dtype or devices.dtype - - local_file = None - if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - from modules import sd_unet - 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] - - if local_file is not None and local_file.lower().endswith('.gguf'): - log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}') - from modules import ggml - ggml.install_gguf() - loader = cls_name.from_single_file if hasattr(cls_name, 'from_single_file') else cls_name.from_pretrained - transformer = loader( - local_file, - quantization_config=diffusers.GGUFQuantizationConfig(compute_dtype=dtype), - cache_dir=shared.opts.hfcache_dir, - **load_args, - ) - transformer = model_quant.do_post_load_quant(transformer, allow=quant_type is not None) - elif local_file is not None and local_file.lower().endswith('.safetensors') and native_spec is not None: - from pipelines import native_transformer - log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader=native args={load_args}') - transformer, _ = native_transformer.load( - local_file, repo_id, native_spec, load_config, - allow_quant=allow_quant, - dtype=dtype, - modules_to_not_convert=modules_to_not_convert, - modules_dtype_dict=modules_dtype_dict, - quant_args=quant_args, - quant_type=quant_type, - **kwargs, - ) - elif local_file is not None and local_file.lower().endswith('.safetensors'): - log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}') - if dtype is not None: - load_args['torch_dtype'] = dtype - load_args.pop('device_map', None) # single-file uses different syntax - loader = cls_name.from_single_file if hasattr(cls_name, 'from_single_file') else cls_name.from_pretrained - transformer = loader( - local_file, - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - **kwargs, - ) - else: - log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} quant="{quant_type}" loader={_loader("diffusers")} args={load_args}') - if 'sdnq-' in repo_id.lower(): - quant_args = {} - if dtype is not None: - load_args['torch_dtype'] = dtype - if subfolder is not None: - load_args['subfolder'] = subfolder - if variant is not None: - load_args['variant'] = variant - transformer = cls_name.from_pretrained( - repo_id, - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - **kwargs, - ) - - sd_models.allow_post_quant = False # we already handled it - if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: - sd_models.move_model(transformer, devices.cpu) - - if transformer is not None and not hasattr(transformer, 'quantization_config'): # attach quantization_config - if hasattr(transformer, 'config') and hasattr(transformer.config, 'quantization_config'): - transformer.quantization_config = transformer.config.quantization_config - elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): - transformer.quantization_config = quant_args.get('quantization_config', None) - - except Exception as e: - log.error(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} {e}') - errors.display(e, 'Load') - raise - - devices.torch_gc() - shared.state.end(jobid) - return transformer - - -def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encoder", allow_quant=True, allow_shared=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, **kwargs): - if shared.state.interrupted: - return None - text_encoder = None - if load_config is None: - load_config = {} - if modules_to_not_convert is None: - modules_to_not_convert = [] - if modules_dtype_dict is None: - modules_dtype_dict = {} - jobid = shared.state.begin('Load TE') - try: - load_args, quant_args = model_quant.get_dit_args(load_config, module='TE', device_map=True, allow_quant=allow_quant, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) - quant_type = model_quant.get_quant_type(quant_args) - load_args.pop('torch_dtype', None) - dtype = dtype or devices.dtype - load_args['dtype'] = dtype - - # load from local file if specified - local_file = None - if shared.opts.sd_text_encoder is not None and shared.opts.sd_text_encoder != 'Default': - from modules import model_te - if shared.opts.sd_text_encoder not in list(model_te.te_dict): - log.error(f'Load module: type=te file="{shared.opts.sd_text_encoder}" not found') - elif os.path.exists(model_te.te_dict[shared.opts.sd_text_encoder]): - local_file = model_te.te_dict[shared.opts.sd_text_encoder] - - # load from local file gguf - if local_file is not None and local_file.lower().endswith('.gguf'): - log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")}') - """ - from modules import ggml - ggml.install_gguf() - text_encoder = cls_name.from_pretrained( - gguf_file=local_file, - quantization_config=diffusers.GGUFQuantizationConfig(compute_dtype=dtype), - cache_dir=shared.opts.hfcache_dir, - **load_args, - ) - text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) - """ - text_encoder = model_te.load_t5(local_file) - text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) - - # load from local file safetensors - elif local_file is not None and local_file.lower().endswith('.safetensors'): - log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")}') - from modules import model_te - text_encoder = model_te.load_t5(local_file) - text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) - - # use shared t5 if possible - elif cls_name == transformers.T5EncoderModel and allow_shared and shared.opts.te_shared_t5: - if model_quant.check_nunchaku('TE'): - import nunchaku - repo_id = 'nunchaku-ai/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors' - cls_name = nunchaku.NunchakuT5EncoderModel - log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="SVDQuant" loader={_loader("transformers")}') - text_encoder = nunchaku.NunchakuT5EncoderModel.from_pretrained( - repo_id, - torch_dtype=dtype, - **kwargs, - ) - text_encoder.quantization_method = 'SVDQuant' - else: - if 'sdnq-uint4-svd' in repo_id.lower(): - repo_id = 'Disty0/FLUX.1-dev-SDNQ-uint4-svd-r32' - load_args['subfolder'] = 'text_encoder_2' - else: - repo_id = 'Disty0/t5-xxl' - with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f: - load_args['config'] = transformers.T5Config(**json.load(f)) - log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}') - text_encoder = cls_name.from_pretrained( - repo_id, - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - **kwargs, - ) - elif cls_name == transformers.UMT5EncoderModel and allow_shared and shared.opts.te_shared_t5: - if 'sdnq-uint4-svd' in repo_id.lower(): - repo_id = 'Disty0/Wan2.2-T2V-A14B-SDNQ-uint4-svd-r32' - else: - repo_id = 'Wan-AI/Wan2.1-T2V-1.3B-Diffusers' - subfolder = 'text_encoder' - log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}') - text_encoder = cls_name.from_pretrained( - repo_id, - cache_dir=shared.opts.hfcache_dir, - subfolder=subfolder, - **load_args, - **quant_args, - **kwargs, - ) - elif cls_name == transformers.Qwen2_5_VLForConditionalGeneration and allow_shared and shared.opts.te_shared_t5: - repo_id = 'hunyuanvideo-community/HunyuanImage-2.1-Diffusers' - subfolder = 'text_encoder' - log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}') - text_encoder = cls_name.from_pretrained( - repo_id, - cache_dir=shared.opts.hfcache_dir, - subfolder=subfolder, - **load_args, - **quant_args, - **kwargs, - ) - # Qwen3ForCausalLM - shared text encoders by hidden_size: - # - Z-Image, Klein-4B: Qwen3-4B (hidden_size=2560) - # - Klein-9B: Qwen3-8B (hidden_size=4096) - # SDNQ repos for Klein and Z-Image contain text encoders pre-quantized with different quantization methods, skip shared loading - elif cls_name == transformers.Qwen3ForCausalLM and allow_shared and shared.opts.te_shared_t5 and 'sdnq' not in repo_id.lower(): - if '-9b' in repo_id.lower(): - shared_repo = 'black-forest-labs/FLUX.2-klein-9B' # 9B variants use Qwen3-8B - else: - shared_repo = 'Tongyi-MAI/Z-Image-Turbo' # 4B variants and Z-Image use Qwen3-4B - subfolder = 'text_encoder' - log.debug(f'Load model: text_encoder="{shared_repo}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}') - text_encoder = cls_name.from_pretrained( - shared_repo, - cache_dir=shared.opts.hfcache_dir, - subfolder=subfolder, - **load_args, - **quant_args, - **kwargs, - ) - - # load from repo - if text_encoder is None: - log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={_loader("transformers")} shared={shared.opts.te_shared_t5}') - if subfolder is not None: - load_args['subfolder'] = subfolder - if variant is not None: - load_args['variant'] = variant - text_encoder = cls_name.from_pretrained( - repo_id, - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - **kwargs, - ) - - sd_models.allow_post_quant = False # we already handled it - if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None: - sd_models.move_model(text_encoder, devices.cpu) - - if text_encoder is not None and not hasattr(text_encoder, 'quantization_config'): # attach quantization_config - if hasattr(text_encoder, 'config') and hasattr(text_encoder.config, 'quantization_config'): - text_encoder.quantization_config = text_encoder.config.quantization_config - elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): - text_encoder.quantization_config = quant_args.get('quantization_config', None) - - except Exception as e: - log.error(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} {e}') - errors.display(e, 'Load') - raise - - devices.torch_gc() - shared.state.end(jobid) - return text_encoder - - -def load_vae_override(pipe, load_config=None, override_cls=None, override_args={}): - if shared.state.interrupted: - return - if (shared.opts.sd_vae in [None, 'None', 'Default', 'Automatic']): - return - if (pipe is None) or (getattr(pipe, 'vae', None) is None): - return - if load_config is None: - load_config = {} - - cls = override_cls or pipe.vae.__class__ - if not hasattr(cls, 'from_single_file'): - log.error(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} safetensors=unsupported') - return - load_args, quant_args = model_quant.get_dit_args(load_config, module='VAE') - log.info(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} args={load_args} quant={quant_args}') - try: - fn = os.path.join(shared.opts.vae_dir, shared.opts.sd_vae) - vae = cls.from_single_file( - fn, - cache_dir=shared.opts.hfcache_dir, - **override_args, - **load_args, - **quant_args, - ) - if vae is not None: - pipe.vae = vae - except Exception as e: - log.error(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} {e}') - # errors.display(e, 'Load') +__all__ = ["load_transformer", "load_text_encoder", "load_vae_override"] diff --git a/pipelines/generic_shared.py b/pipelines/generic_shared.py new file mode 100644 index 000000000..588c51602 --- /dev/null +++ b/pipelines/generic_shared.py @@ -0,0 +1,98 @@ +import os +import transformers + + +shared_te_map = { + 'T5-XXL SDNQ-UInt4': { + 'cls': transformers.T5EncoderModel, + 'identifier': 'sdnq-uint4', + 'target_repo': 'Disty0/FLUX.1-dev-SDNQ-uint4-svd-r32', + }, + 'T5-XXL Base': { # template + 'cls': transformers.T5EncoderModel, # desired model class, used as primary matching criteria + 'identifier': None, # additional identifier to match in repo_id or None to ignore + 'target_repo': 'Disty0/t5-xxl', # repo to load from instead of original repo_id + 'target_subfolder': None, # subfolder in repo to load from, None to ignore + 'config_class': transformers.T5Config, # config class to use for loading or None to ignore + 'config_path': os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), # path to config file to use for loading or None to ignore + }, + + 'UMT5 SDNQ-UInt4': { + 'cls': transformers.UMT5EncoderModel, + 'identifier': 'sdnq-uint4', + 'target_repo': 'Disty0/Wan2.2-T2V-A14B-SDNQ-uint4-svd-r32', + 'target_subfolder': 'text_encoder', + }, + 'UMT5 Base': { + 'cls': transformers.UMT5EncoderModel, + 'target_repo': 'Wan-AI/Wan2.1-T2V-1.3B-Diffusers', + 'target_subfolder': 'text_encoder', + }, + + 'Qwen-2.5 SDNQ-4Bit': { + 'cls': transformers.Qwen2_5_VLForConditionalGeneration, + 'identifier': 'sdnq-4bit', + 'target_repo': 'Disty0/Qwen-Image-2512-SDNQ-uint4-svd-r32', + 'target_subfolder': 'text_encoder', + }, + 'Qwen-2.5 SDNQ-UInt4': { + 'cls': transformers.Qwen2_5_VLForConditionalGeneration, + 'identifier': 'sdnq-uint4', + 'target_repo': 'Disty0/Qwen-Image-2512-SDNQ-uint4-svd-r32', + 'target_subfolder': 'text_encoder', + }, + 'Qwen-2.5 Base': { + 'cls': transformers.Qwen2_5_VLForConditionalGeneration, + 'target_repo': 'hunyuanvideo-community/HunyuanImage-2.1-Diffusers', + 'target_subfolder': 'text_encoder', + }, + + 'Qwen-3 9B SDNQ-4bit': { + 'cls': transformers.Qwen3ForCausalLM, + 'identifier': '9b-sdnq-4bit', + 'target_repo': 'Disty0/FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32', + 'target_subfolder': 'text_encoder', + }, + 'Qwen-3 9B SDNQ-UInt4': { + 'cls': transformers.Qwen3ForCausalLM, + 'identifier': '9b-sdnq-uint4', + 'target_repo': 'Disty0/FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32', + 'target_subfolder': 'text_encoder', + }, + 'Qwen-3 9B Base': { + 'cls': transformers.Qwen3ForCausalLM, + 'identifier': '9b', + 'target_repo': 'black-forest-labs/FLUX.2-klein-9B', + 'target_subfolder': 'text_encoder', + }, + + 'Qwen-3 4B SDNQ-4Bit': { # match after 9b + 'cls': transformers.Qwen3ForCausalLM, + 'identifier': 'sdnq-4bit', + 'target_repo': 'Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32', + 'target_subfolder': 'text_encoder', + }, + 'Qwen-3 4B SDNQ-UInt4': { + 'cls': transformers.Qwen3ForCausalLM, + 'identifier': 'sdnq-uint4', + 'target_repo': 'Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32', + 'target_subfolder': 'text_encoder', + }, + 'Qwen-3 4B Base': { + 'cls': transformers.Qwen3ForCausalLM, + 'target_repo': 'Tongyi-MAI/Z-Image-Turbo', + 'target_subfolder': 'text_encoder', + }, + + 'Qwen-3 0.5B SDNQ-UInt4': { + 'cls': transformers.Qwen3Model, + 'identifier': 'uint4', + 'target_repo': 'vladmandic/Anima-1.0-Base-sdnq-svd-dynamic-uint4', + 'target_subfolder': 'text_encoder', + }, + 'Qwen-3 0.5B Base': { + 'cls': transformers.Qwen3Model, + 'target_repo': 'vladmandic/Anima-1.0-Base', + 'target_subfolder': 'text_encoder', + }, +} diff --git a/pipelines/generic_text_encoder.py b/pipelines/generic_text_encoder.py new file mode 100644 index 000000000..a9c393c2f --- /dev/null +++ b/pipelines/generic_text_encoder.py @@ -0,0 +1,141 @@ +import os +import json +import transformers +from modules import shared, devices, errors, sd_models, model_quant +from modules.logger import log +from pipelines.generic_util import get_loader +from pipelines.generic_shared import shared_te_map + + +debug = os.environ.get('SD_LOAD_DEBUG', None) is not None + + +def get_shared(cls, repo_id, subfolder=None, variant=None): + args = {} + if variant is not None: + args['variant'] = variant + for name, item in shared_te_map.items(): + if item['cls'] == cls and (item['identifier'] is None or item['identifier'].lower() in repo_id.lower()): + if item.get('config_class', None) is not None and item.get('config_path', None) is not None: + with open(item['config_path'], encoding='utf8') as f: + args['config'] = item['config_class'](**json.load(f)) + if item.get('target_subfolder', None) is not None: + args['subfolder'] = item['target_subfolder'] + log.debug(f'Load model: text_encoder="{repo_id}" cls={cls.__name__} target="{item["target_repo"]}" args={args} shared="{name}"') + return item['target_repo'], args + if subfolder is not None: # use default provided subfolder + args['subfolder'] = subfolder + return repo_id, args + + +def load_local_file(local_file, cls_name, quant_type): # t5-only + from modules import model_te + text_encoder = None + + # 1. load from local file gguf + if local_file.lower().endswith('.gguf'): + log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={get_loader("transformers")} file=gguf') + text_encoder = model_te.load_t5(local_file) + text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) + + # 2. t5 - load from local file safetensors + elif local_file.lower().endswith('.safetensors'): + log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={get_loader("transformers")} file=safetensors') + text_encoder = model_te.load_t5(local_file) + text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) + + return text_encoder + + +def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encoder", allow_quant=True, allow_shared=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, **kwargs): + if shared.state.interrupted: + return None + text_encoder = None + allow_shared = allow_shared and shared.opts.te_shared_te + if load_config is None: + load_config = {} + if modules_to_not_convert is None: + modules_to_not_convert = [] + if modules_dtype_dict is None: + modules_dtype_dict = {} + jobid = shared.state.begin('Load TE') + try: + load_args, quant_args = model_quant.get_dit_args(load_config, module='TE', device_map=True, allow_quant=allow_quant, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) + quant_type = model_quant.get_quant_type(quant_args) + load_args.pop('torch_dtype', None) + dtype = dtype or devices.dtype + load_args['dtype'] = dtype + + # 1. load override from local file + if (shared.opts.sd_text_encoder is not None) and (shared.opts.sd_text_encoder != 'Default') and (text_encoder is None): + local_file = None + from modules import model_te + if shared.opts.sd_text_encoder not in list(model_te.te_dict): + log.error(f'Load module: type=te file="{shared.opts.sd_text_encoder}" not found') + elif os.path.exists(model_te.te_dict[shared.opts.sd_text_encoder]): + local_file = model_te.te_dict[shared.opts.sd_text_encoder] + if local_file is not None: + text_encoder = load_local_file(local_file, cls_name, quant_type) + + # 2. load override from repo + if (shared.opts.sd_text_encoder is not None) and (shared.opts.sd_text_encoder != 'Default') and (text_encoder is None): + repo_id = shared.opts.sd_text_encoder + if '/' in repo_id: # shared.opts.sd_text_encoder can be in format org/repo or org/repo/subfolder + parts = repo_id.split('/') + if len(parts) >= 3: + repo_id = '/'.join(parts[:2]) + load_args['subfolder'] = '/'.join(parts[2:]) + log.debug(f'Load model: text_encoder="{repo_id}" quant="{quant_type}" loader={get_loader("transformers")} type=override') + text_encoder = transformers.AutoModel.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + **kwargs, + ) + + # 3. load shared from repo + if allow_shared and (text_encoder is None): + log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={get_loader("transformers")}') + target_repo, extra_args = get_shared(cls_name, repo_id, subfolder=subfolder, variant=variant) + text_encoder = cls_name.from_pretrained( + target_repo, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + **extra_args, + ) + + # 4. load default from repo + if text_encoder is None: + log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" loader={get_loader("transformers")}') + if subfolder is not None: + load_args['subfolder'] = subfolder + if variant is not None: + load_args['variant'] = variant + text_encoder = cls_name.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + **kwargs, + ) + + sd_models.allow_post_quant = False # we already handled it + if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None: + sd_models.move_model(text_encoder, devices.cpu) + + if text_encoder is not None and not hasattr(text_encoder, 'quantization_config'): # attach quantization_config + if hasattr(text_encoder, 'config') and hasattr(text_encoder.config, 'quantization_config'): + text_encoder.quantization_config = text_encoder.config.quantization_config + elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): + text_encoder.quantization_config = quant_args.get('quantization_config', None) + + except Exception as e: + log.error(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} {e}') + errors.display(e, 'Load') + raise + + devices.torch_gc() + shared.state.end(jobid) + return text_encoder diff --git a/pipelines/generic_transformer.py b/pipelines/generic_transformer.py new file mode 100644 index 000000000..5651a5eac --- /dev/null +++ b/pipelines/generic_transformer.py @@ -0,0 +1,119 @@ +import os +import diffusers +from modules import shared, devices, errors, sd_models, model_quant +from modules.logger import log +from pipelines.generic_util import get_loader + + +debug = os.environ.get('SD_LOAD_DEBUG', None) is not None + + +def load_transformer(repo_id, cls_name, load_config=None, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=None, modules_dtype_dict=None, native_spec=None, **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. + + With ``native_spec`` set and a .safetensors override selected, dispatches + to :func:`pipelines.native_transformer.load`. Without a spec, a single-file + override falls back to ``from_single_file``. + """ + if shared.state.interrupted: + return None + transformer = None + if load_config is None: + load_config = {} + if modules_to_not_convert is None: + modules_to_not_convert = [] + if modules_dtype_dict is None: + modules_dtype_dict = {} + jobid = shared.state.begin('Load DiT') + try: + load_args, quant_args = model_quant.get_dit_args(load_config, module='Model', device_map=True, allow_quant=allow_quant, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) + quant_type = model_quant.get_quant_type(quant_args) + dtype = dtype or devices.dtype + + local_file = None + if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': + from modules import sd_unet + 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] + + # 1. load gguf + if local_file is not None and local_file.lower().endswith('.gguf'): + log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={get_loader("diffusers")} args={load_args}') + from modules import ggml + ggml.install_gguf() + loader = cls_name.from_single_file if hasattr(cls_name, 'from_single_file') else cls_name.from_pretrained + transformer = loader( + local_file, + quantization_config=diffusers.GGUFQuantizationConfig(compute_dtype=dtype), + cache_dir=shared.opts.hfcache_dir, + **load_args, + ) + transformer = model_quant.do_post_load_quant(transformer, allow=quant_type is not None) + # 2. load safetensors with native loader if spec is available + elif local_file is not None and local_file.lower().endswith('.safetensors') and native_spec is not None: + from pipelines import native_transformer + log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader=native args={load_args}') + transformer, _ = native_transformer.load( + local_file, repo_id, native_spec, load_config, + allow_quant=allow_quant, + dtype=dtype, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + quant_args=quant_args, + quant_type=quant_type, + **kwargs, + ) + # 3. load safetensors with diffusers loader + elif local_file is not None and local_file.lower().endswith('.safetensors'): + log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" loader={get_loader("diffusers")} args={load_args}') + if dtype is not None: + load_args['torch_dtype'] = dtype + load_args.pop('device_map', None) # single-file uses different syntax + loader = cls_name.from_single_file if hasattr(cls_name, 'from_single_file') else cls_name.from_pretrained + transformer = loader( + local_file, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + **kwargs, + ) + # 4. default loading from diffusers repo + else: + log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} subfolder={subfolder} quant="{quant_type}" loader={get_loader("diffusers")} args={load_args}') + if 'sdnq-' in repo_id.lower(): + quant_args = {} + if dtype is not None: + load_args['torch_dtype'] = dtype + if subfolder is not None: + load_args['subfolder'] = subfolder + if variant is not None: + load_args['variant'] = variant + transformer = cls_name.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + **kwargs, + ) + + sd_models.allow_post_quant = False # we already handled it + if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: + sd_models.move_model(transformer, devices.cpu) + + if transformer is not None and not hasattr(transformer, 'quantization_config'): # attach quantization_config + if hasattr(transformer, 'config') and hasattr(transformer.config, 'quantization_config'): + transformer.quantization_config = transformer.config.quantization_config + elif (quant_type is not None) and (quant_args.get('quantization_config', None) is not None): + transformer.quantization_config = quant_args.get('quantization_config', None) + + except Exception as e: + log.error(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} {e}') + errors.display(e, 'Load') + raise + + devices.torch_gc() + shared.state.end(jobid) + return transformer diff --git a/pipelines/generic_util.py b/pipelines/generic_util.py new file mode 100644 index 000000000..92d026d1e --- /dev/null +++ b/pipelines/generic_util.py @@ -0,0 +1,11 @@ +import sys +from modules import shared + + +def get_loader(component): + """Return loader type for log messages.""" + if sys.platform != 'linux': + return 'default' + if component == 'diffusers': + return 'runai' if shared.opts.runai_streamer_diffusers else 'default' + return 'runai' if shared.opts.runai_streamer_transformers else 'default' diff --git a/pipelines/generic_vae.py b/pipelines/generic_vae.py new file mode 100644 index 000000000..27f48fac4 --- /dev/null +++ b/pipelines/generic_vae.py @@ -0,0 +1,38 @@ +import os +from modules import shared, model_quant +from modules.logger import log + + +debug = os.environ.get('SD_LOAD_DEBUG', None) is not None + + +def load_vae_override(pipe, load_config=None, override_cls=None, override_args={}): + if shared.state.interrupted: + return + if (shared.opts.sd_vae in [None, 'None', 'Default', 'Automatic']): + return + if (pipe is None) or (getattr(pipe, 'vae', None) is None): + return + if load_config is None: + load_config = {} + + cls = override_cls or pipe.vae.__class__ + if not hasattr(cls, 'from_single_file'): + log.error(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} safetensors=unsupported') + return + load_args, quant_args = model_quant.get_dit_args(load_config, module='VAE') + log.info(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} args={load_args} quant={quant_args}') + try: + fn = os.path.join(shared.opts.vae_dir, shared.opts.sd_vae) + vae = cls.from_single_file( + fn, + cache_dir=shared.opts.hfcache_dir, + **override_args, + **load_args, + **quant_args, + ) + if vae is not None: + pipe.vae = vae + except Exception as e: + log.error(f'Load model: vae="{shared.opts.sd_vae}" cls={cls.__name__} {e}') + # errors.display(e, 'Load') diff --git a/pipelines/model_anima.py b/pipelines/model_anima.py index 4ed2228fb..2274d8afb 100644 --- a/pipelines/model_anima.py +++ b/pipelines/model_anima.py @@ -39,7 +39,12 @@ def load_transformer_components(repo_id, diffusers_load_config, adapter_cls): log.error(f'Load model: type=Anima custom transformer="{local_file}": {e}') errors.display(e, 'Load') return None, None - transformer = generic.load_transformer(repo_id, cls_name=diffusers.CosmosTransformer3DModel, load_config=diffusers_load_config, subfolder="transformer") + transformer = generic.load_transformer( + repo_id, + cls_name=diffusers.CosmosTransformer3DModel, + load_config=diffusers_load_config, + subfolder="transformer" + ) return transformer, None @@ -89,7 +94,13 @@ def load_anima(checkpoint_info, diffusers_load_config=None): transformer, llm_adapter = load_transformer_components(repo_id, diffusers_load_config, AnimaLLMAdapter) if transformer is None: return None - text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3Model, load_config=diffusers_load_config, subfolder="text_encoder", allow_shared=False) + text_encoder = generic.load_text_encoder( + repo_id, + cls_name=transformers.Qwen3Model, + load_config=diffusers_load_config, + subfolder="text_encoder", + allow_shared=False + ) if llm_adapter is None: shared.state.begin('Load adapter') diff --git a/pipelines/model_qwen.py b/pipelines/model_qwen.py index 094d23ed3..04708470f 100644 --- a/pipelines/model_qwen.py +++ b/pipelines/model_qwen.py @@ -67,8 +67,11 @@ def load_qwen(checkpoint_info, diffusers_load_config=None): native_spec=QWEN_SPEC, ) - repo_te = 'Qwen/Qwen-Image' - text_encoder = generic.load_text_encoder(repo_te, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) + text_encoder = generic.load_text_encoder( + repo_id, + cls_name=transformers.Qwen2_5_VLForConditionalGeneration, + load_config=diffusers_load_config + ) repo_id, repo_subfolder = qwen.check_qwen_pruning(repo_id, repo_subfolder) if repo_subfolder is not None and repo_subfolder.startswith('nunchaku'):