1580158
/
automatic
mirror of https://github.com/vladmandic/automatic
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix(model_flux2_klein): update generic.py and pipeline

pull/4556/head
CalamitousFelicitousness 2026-01-17 21:29:55 +00:00
parent 33ee04a9f3
commit dd9001b20d
2 changed files with 4 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pipelines/generic.py
View File

@ -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:

39
pipelines/model_flux2_klein.py
View File

@ -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,