diff --git a/pipelines/generic.py b/pipelines/generic.py index c3fd8fd42..6ad9a3fd5 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -203,7 +203,8 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod # Qwen3ForCausalLM - shared text encoders by hidden_size: # - Z-Image, Klein-4B: Qwen3-4B (hidden_size=2560) # - Klein-9B: Qwen3-8B (hidden_size=4096) - elif cls_name == transformers.Qwen3ForCausalLM and allow_shared and shared.opts.te_shared_t5: + # 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: diff --git a/pipelines/model_flux2_klein.py b/pipelines/model_flux2_klein.py index 0b6ce47e0..d810821d9 100644 --- a/pipelines/model_flux2_klein.py +++ b/pipelines/model_flux2_klein.py @@ -1,58 +1,23 @@ -import os -import shutil import transformers import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae from pipelines import generic -def ensure_tokenizer_files(checkpoint_info): - """Ensure tokenizer files are compatible with Transformers v4.""" - local_path = checkpoint_info.path - if not os.path.isdir(local_path): - return - tokenizer_path = os.path.join(local_path, 'tokenizer') - if not os.path.isdir(tokenizer_path): - return - - # Check all required files exist (v5 tokenizers may be missing these or have incompatible configs) - required = ['vocab.json', 'merges.txt', 'tokenizer_config.json'] - missing = [f for f in required if not os.path.exists(os.path.join(tokenizer_path, f))] - if not missing: - return - - # Download v4-compatible tokenizer files from Z-Image-Turbo - shared.log.debug(f'Load model: fetching v4-compatible tokenizer from Z-Image-Turbo (missing: {missing})') - try: - from huggingface_hub import hf_hub_download - for f in required: - src = hf_hub_download('Tongyi-MAI/Z-Image-Turbo', f'tokenizer/{f}', cache_dir=shared.opts.hfcache_dir) - shutil.copy(src, os.path.join(tokenizer_path, f)) - except Exception as e: - shared.log.warning(f'Load model: failed to fetch tokenizer files: {e}') - - def load_flux2_klein(checkpoint_info, diffusers_load_config=None): if diffusers_load_config is None: diffusers_load_config = {} repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - # Ensure tokenizer files exist for models created with Transformers v5 - ensure_tokenizer_files(checkpoint_info) - - # Detect SDNQ pre-quantized repo - disable shared text encoder to use pre-quantized weights from SDNQ repo - is_sdnq = 'sdnq' in repo_id.lower() - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) - shared.log.debug(f'Load model: type=Flux2Klein repo="{repo_id}" prequant={"sdnq" if is_sdnq else "none"} config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + shared.log.debug(f'Load model: type=Flux2Klein repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') # Load transformer - Klein uses Flux2Transformer2DModel (same class as Flux2, different size) transformer = generic.load_transformer(repo_id, cls_name=diffusers.Flux2Transformer2DModel, load_config=diffusers_load_config) # Load text encoder - Klein uses Qwen3 (4B for Klein-4B, 8B for Klein-9B) - # For SDNQ repos, disable shared text encoder to use pre-quantized weights from the repo - text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3ForCausalLM, load_config=diffusers_load_config, allow_shared=not is_sdnq) + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3ForCausalLM, load_config=diffusers_load_config) pipe = diffusers.Flux2KleinPipeline.from_pretrained( repo_id,