From fe99d3fe5dd9fb91c9774e7fdfe95147c10aa9bb Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 16 Jan 2026 01:16:06 +0000 Subject: [PATCH 1/5] feat: add FLUX.2 Klein model support Add support for FLUX.2 Klein distilled models (4B and 9B variants): - Add pipeline loader for Flux2KleinPipeline - Add model detection for 'flux.2' + 'klein' patterns - Add pipeline mapping in shared_items - Add shared Qwen3ForCausalLM text encoder handling: - 4B variants use Z-Image-Turbo's Qwen3-8B - 9B variants use FLUX.2-klein-9B's Qwen3-14B - Add reference entries for distilled (4B, 9B) and base models - Update diffusers commit for Flux2KleinPipeline support --- html/reference-distilled.json | 20 ++++++++++++++++ html/reference.json | 18 +++++++++++++++ installer.py | 2 +- modules/sd_detect.py | 2 ++ modules/sd_models.py | 4 ++++ modules/shared_items.py | 2 ++ pipelines/generic.py | 17 ++++++++++++++ pipelines/model_flux2_klein.py | 42 ++++++++++++++++++++++++++++++++++ 8 files changed, 106 insertions(+), 1 deletion(-) create mode 100644 pipelines/model_flux2_klein.py diff --git a/html/reference-distilled.json b/html/reference-distilled.json index d31e3fdc5..36e583ec2 100644 --- a/html/reference-distilled.json +++ b/html/reference-distilled.json @@ -161,5 +161,25 @@ "preview": "Tencent-Hunyuan--HunyuanDiT-v1.1-Diffusers-Distilled.jpg", "tags": "distilled", "extras": "sampler: Default, cfg_scale: 2.0" + }, + "Black Forest Labs FLUX.2 Klein 4B": { + "path": "black-forest-labs/FLUX.2-klein-4B", + "preview": "black-forest-labs--FLUX.2-klein-4B.jpg", + "desc": "FLUX.2-klein-4B is a 4 billion parameter size-distilled version of FLUX.2-dev optimized for consumer GPUs. Achieves sub-second inference with 4 steps while fitting in ~13GB VRAM. Supports both text-to-image generation and multi-reference image editing. Apache 2.0 licensed.", + "skip": true, + "tags": "distilled", + "extras": "sampler: Default, cfg_scale: 4.0, steps: 4", + "size": 8.5, + "date": "2025 January" + }, + "Black Forest Labs FLUX.2 Klein 9B": { + "path": "black-forest-labs/FLUX.2-klein-9B", + "preview": "black-forest-labs--FLUX.2-klein-9B.jpg", + "desc": "FLUX.2-klein-9B is a 9 billion parameter size-distilled version of FLUX.2-dev. Higher quality than 4B variant with sub-second inference using 4 steps. Requires ~29GB VRAM. Supports text-to-image and multi-reference editing. Non-commercial license.", + "skip": true, + "tags": "distilled", + "extras": "sampler: Default, cfg_scale: 4.0, steps: 4", + "size": 18.5, + "date": "2025 January" } } \ No newline at end of file diff --git a/html/reference.json b/html/reference.json index d86214467..63cd80124 100644 --- a/html/reference.json +++ b/html/reference.json @@ -124,6 +124,24 @@ "size": 104.74, "date": "2025 November" }, + "Black Forest Labs FLUX.2 Klein Base 4B": { + "path": "black-forest-labs/FLUX.2-klein-base-4B", + "preview": "black-forest-labs--FLUX.2-klein-base-4B.jpg", + "desc": "FLUX.2-klein-base-4B is the undistilled 4 billion parameter base model of FLUX.2-klein. Requires 50 inference steps for full quality but offers flexibility for fine-tuning. Fits in ~13GB VRAM. Supports text-to-image and multi-reference editing. Apache 2.0 licensed.", + "skip": true, + "extras": "sampler: Default, cfg_scale: 4.0, steps: 50", + "size": 8.5, + "date": "2025 January" + }, + "Black Forest Labs FLUX.2 Klein Base 9B": { + "path": "black-forest-labs/FLUX.2-klein-base-9B", + "preview": "black-forest-labs--FLUX.2-klein-base-9B.jpg", + "desc": "FLUX.2-klein-base-9B is the undistilled 9 billion parameter base model of FLUX.2-klein. Requires 50 inference steps for full quality but offers flexibility for fine-tuning. Requires ~29GB VRAM. Supports text-to-image and multi-reference editing. Non-commercial license.", + "skip": true, + "extras": "sampler: Default, cfg_scale: 4.0, steps: 50", + "size": 18.5, + "date": "2025 January" + }, "Z-Image-Turbo": { "path": "Tongyi-MAI/Z-Image-Turbo", diff --git a/installer.py b/installer.py index 913c7d357..08a4bc3e0 100644 --- a/installer.py +++ b/installer.py @@ -648,7 +648,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all: return - sha = '5efb81fa711863fdece9136ad10788440e658b40' # diffusers commit hash + sha = '61f175660a8ac54f1470a74a810e6c38fb4795d5' # diffusers commit hash # if args.use_rocm or args.use_zluda or args.use_directml: # sha = '043ab2520f6a19fce78e6e060a68dbc947edb9f9' # lock diffusers versions for now pkg = pkg_resources.working_set.by_key.get('diffusers', None) diff --git a/modules/sd_detect.py b/modules/sd_detect.py index 0e93b6d20..a1cb6e913 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -92,6 +92,8 @@ def guess_by_name(fn, current_guess): new_guess = 'HiDream' elif 'chroma' in fn.lower() and 'xl' not in fn.lower(): new_guess = 'Chroma' + elif 'flux.2' in fn.lower() and 'klein' in fn.lower(): + new_guess = 'FLUX2 Klein' elif 'flux.2' in fn.lower(): new_guess = 'FLUX2' elif 'flux' in fn.lower() or 'flex.1' in fn.lower(): diff --git a/modules/sd_models.py b/modules/sd_models.py index f0a832dd6..8db759e87 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -359,6 +359,10 @@ def load_diffuser_force(detected_model_type, checkpoint_info, diffusers_load_con from pipelines.model_flux2 import load_flux2 sd_model = load_flux2(checkpoint_info, diffusers_load_config) allow_post_quant = False + elif model_type in ['FLUX2 Klein']: + from pipelines.model_flux2_klein import load_flux2_klein + sd_model = load_flux2_klein(checkpoint_info, diffusers_load_config) + allow_post_quant = False elif model_type in ['FLEX']: from pipelines.model_flex import load_flex sd_model = load_flex(checkpoint_info, diffusers_load_config) diff --git a/modules/shared_items.py b/modules/shared_items.py index b5df05390..3186b3d12 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -48,6 +48,8 @@ pipelines = { 'Qwen': getattr(diffusers, 'QwenImagePipeline', None), 'HunyuanImage': getattr(diffusers, 'HunyuanImagePipeline', None), 'Z-Image': getattr(diffusers, 'ZImagePipeline', None), + 'FLUX2': getattr(diffusers, 'Flux2Pipeline', None), + 'FLUX2 Klein': getattr(diffusers, 'Flux2KleinPipeline', None), 'LongCat': getattr(diffusers, 'LongCatImagePipeline', None), 'GLM-Image': getattr(diffusers, 'GlmImagePipeline', None), # dynamically imported and redefined later diff --git a/pipelines/generic.py b/pipelines/generic.py index d6dad46c8..0a51c1cec 100644 --- a/pipelines/generic.py +++ b/pipelines/generic.py @@ -200,6 +200,23 @@ def load_text_encoder(repo_id, cls_name, load_config=None, subfolder="text_encod **load_args, **quant_args, ) + # 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: + 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' + shared.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, + ) # load from repo if text_encoder is None: diff --git a/pipelines/model_flux2_klein.py b/pipelines/model_flux2_klein.py new file mode 100644 index 000000000..9b8b05ac8 --- /dev/null +++ b/pipelines/model_flux2_klein.py @@ -0,0 +1,42 @@ +import transformers +import diffusers +from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae +from pipelines import generic + + +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) + + 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}" 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 Qwen3ForCausalLM (8B), shared across all Klein variants + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3ForCausalLM, load_config=diffusers_load_config) + + pipe = diffusers.Flux2KleinPipeline.from_pretrained( + repo_id, + transformer=transformer, + text_encoder=text_encoder, + cache_dir=shared.opts.diffusers_dir, + **load_args, + ) + pipe.task_args = { + 'output_type': 'np', + } + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flux2klein"] = diffusers.Flux2KleinPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux2klein"] = diffusers.Flux2KleinPipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux2klein"] = diffusers.Flux2KleinPipeline + + del text_encoder + del transformer + sd_hijack_te.init_hijack(pipe) + sd_hijack_vae.init_hijack(pipe) + + devices.torch_gc(force=True, reason='load') + return pipe From c605a1bb626ca1b6aa7ac996fa7d7f67cbf45fec Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 16 Jan 2026 03:01:52 +0000 Subject: [PATCH 2/5] feat(taesd): add FLUX.2 preview support Enable live preview during FLUX.2 and FLUX.2 Klein image generation using the TAE FLUX.2 decoder from madebyollin/taesd. - Add dedicated TAE entries (FLUX.1, FLUX.2, SD3) that auto-select based on model type, making the dropdown only affect SD/SDXL models - Add FLUX.2 latent unpacking in callback to convert packed [B, seq_len, 128] format to spatial [B, 32, H, W] for preview - Support FLUX.2's 32 latent channels (vs 16 for FLUX.1/SD3) --- modules/processing_callbacks.py | 33 ++++++++++++++++++++++++++++++++- modules/sd_vae_taesd.py | 20 +++++++++++++++++--- modules/taesd/taesd.py | 6 +++++- 3 files changed, 54 insertions(+), 5 deletions(-) diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index d62d464dc..327fb8100 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -116,7 +116,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {} if current_noise_pred is None: current_noise_pred = kwargs.get("predicted_image_embedding", None) - if hasattr(pipe, "_unpack_latents") and hasattr(pipe, "vae_scale_factor"): # FLUX + if hasattr(pipe, "_unpack_latents") and hasattr(pipe, "vae_scale_factor"): # FLUX.1 if p.hr_resize_mode > 0 and (p.hr_upscaler != 'None' or p.hr_resize_mode == 5) and p.is_hr_pass: width = max(getattr(p, 'width', 0), getattr(p, 'hr_upscale_to_x', 0)) height = max(getattr(p, 'height', 0), getattr(p, 'hr_upscale_to_y', 0)) @@ -128,6 +128,37 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {} shared.state.current_noise_pred = pipe._unpack_latents(current_noise_pred, height, width, pipe.vae_scale_factor) # pylint: disable=protected-access else: shared.state.current_noise_pred = current_noise_pred + elif hasattr(pipe, "_unpatchify_latents"): # FLUX.2 - unpack [B, seq, patch_ch] to [B, ch, H, W] + # Get dimensions for unpacking, same logic as FLUX.1 + vae_scale = getattr(pipe, 'vae_scale_factor', 8) + if p.hr_resize_mode > 0 and (p.hr_upscaler != 'None' or p.hr_resize_mode == 5) and p.is_hr_pass: + width = max(getattr(p, 'width', 0), getattr(p, 'hr_upscale_to_x', 0)) + height = max(getattr(p, 'height', 0), getattr(p, 'hr_upscale_to_y', 0)) + else: + width = getattr(p, 'width', 1024) + height = getattr(p, 'height', 1024) + latents = kwargs['latents'] + if len(latents.shape) == 3: # packed format [B, seq_len, patch_channels] + b, seq_len, patch_ch = latents.shape + channels = patch_ch // 4 # 4 = 2x2 patch + h_patches = height // vae_scale // 2 + w_patches = width // vae_scale // 2 + if h_patches * w_patches != seq_len: # fallback to square assumption + h_patches = w_patches = int(seq_len ** 0.5) + # [B, h*w, C*4] -> [B, h, w, C, 2, 2] -> [B, C, h, 2, w, 2] -> [B, C, H, W] + latents = latents.view(b, h_patches, w_patches, channels, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5).reshape(b, channels, h_patches * 2, w_patches * 2) + shared.state.current_latent = latents + if current_noise_pred is not None and len(current_noise_pred.shape) == 3: + b, seq_len, patch_ch = current_noise_pred.shape + channels = patch_ch // 4 + h_patches = height // vae_scale // 2 + w_patches = width // vae_scale // 2 + if h_patches * w_patches != seq_len: + h_patches = w_patches = int(seq_len ** 0.5) + current_noise_pred = current_noise_pred.view(b, h_patches, w_patches, channels, 2, 2) + current_noise_pred = current_noise_pred.permute(0, 3, 1, 4, 2, 5).reshape(b, channels, h_patches * 2, w_patches * 2) + shared.state.current_noise_pred = current_noise_pred else: shared.state.current_latent = kwargs['latents'] shared.state.current_noise_pred = current_noise_pred diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index 276dae80f..d076dacbb 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -17,6 +17,9 @@ TAESD_MODELS = { 'TAESD 1.2 Chocolate-Dipped Shortbread': { 'fn': 'taesd_12_', 'uri': 'https://github.com/madebyollin/taesd/raw/8909b44e3befaa0efa79c5791e4fe1c4d4f7884e', 'model': None }, 'TAESD 1.1 Fruit Loops': { 'fn': 'taesd_11_', 'uri': 'https://github.com/madebyollin/taesd/raw/3e8a8a2ab4ad4079db60c1c7dc1379b4cc0c6b31', 'model': None }, 'TAESD 1.0': { 'fn': 'taesd_10_', 'uri': 'https://github.com/madebyollin/taesd/raw/88012e67cf0454e6d90f98911fe9d4aef62add86', 'model': None }, + 'TAE FLUX.1': { 'fn': 'taef1.pth', 'uri': 'https://github.com/madebyollin/taesd/raw/main/taef1_decoder.pth', 'model': None }, + 'TAE FLUX.2': { 'fn': 'taef2.pth', 'uri': 'https://github.com/madebyollin/taesd/raw/main/taef2_decoder.pth', 'model': None }, + 'TAE SD3': { 'fn': 'taesd3.pth', 'uri': 'https://github.com/madebyollin/taesd/raw/main/taesd3_decoder.pth', 'model': None }, 'TAE HunyuanVideo': { 'fn': 'taehv.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taehv.pth', 'model': None }, 'TAE WanVideo': { 'fn': 'taew1.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taew2_1.pth', 'model': None }, 'TAE MochiVideo': { 'fn': 'taem1.pth', 'uri': 'https://github.com/madebyollin/taem1/raw/refs/heads/main/taem1.pth', 'model': None }, @@ -38,7 +41,7 @@ prev_cls = '' prev_type = '' prev_model = '' lock = threading.Lock() -supported = ['sd', 'sdxl', 'sd3', 'f1', 'h1', 'zimage', 'lumina2', 'hunyuanvideo', 'wanai', 'chrono', 'cosmos', 'mochivideo', 'pixartsigma', 'pixartalpha', 'hunyuandit', 'omnigen', 'qwen', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4'] +supported = ['sd', 'sdxl', 'sd3', 'f1', 'f2', 'h1', 'zimage', 'lumina2', 'hunyuanvideo', 'wanai', 'chrono', 'cosmos', 'mochivideo', 'pixartsigma', 'pixartalpha', 'hunyuandit', 'omnigen', 'qwen', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4'] def warn_once(msg, variant=None): @@ -59,8 +62,14 @@ def get_model(model_type = 'decoder', variant = None): model_cls = 'sd' elif model_cls in {'pixartsigma', 'hunyuandit', 'omnigen', 'auraflow'}: model_cls = 'sdxl' - elif model_cls in {'h1', 'zimage', 'lumina2', 'chroma', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4'}: + elif model_cls in {'f1', 'h1', 'zimage', 'lumina2', 'chroma', 'longcat', 'omnigen2', 'flite', 'ovis', 'kandinsky5', 'glmimage', 'cogview3', 'cogview4'}: model_cls = 'f1' + variant = 'TAE FLUX.1' + elif model_cls == 'f2': + model_cls = 'f2' + variant = 'TAE FLUX.2' + elif model_cls == 'sd3': + variant = 'TAE SD3' elif model_cls in {'wanai', 'qwen', 'chrono', 'cosmos'}: variant = variant or 'TAE WanVideo' elif model_cls not in supported: @@ -149,7 +158,12 @@ def decode(latents): dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16 tensor = latents.unsqueeze(0) if len(latents.shape) == 3 else latents tensor = tensor.detach().clone().to(devices.device, dtype=dtype) - if variant.startswith('TAESD'): + shared.log.debug(f'Decode: type="taesd" variant="{variant}" input={latents.shape} tensor={tensor.shape}') + # FLUX.2 has 128 latent channels that need reshaping to 32 channels for TAESD + if variant == 'TAE FLUX.2' and len(tensor.shape) == 4 and tensor.shape[1] == 128: + b, c, h, w = tensor.shape + tensor = tensor.reshape(b, 32, h * 2, w * 2) + if variant.startswith('TAESD') or variant in {'TAE FLUX.1', 'TAE FLUX.2', 'TAE SD3'}: image = vae.decoder(tensor).clamp(0, 1).detach() image = image[0] else: diff --git a/modules/taesd/taesd.py b/modules/taesd/taesd.py index 8e391a8fb..f066f4cfd 100644 --- a/modules/taesd/taesd.py +++ b/modules/taesd/taesd.py @@ -77,7 +77,11 @@ class TAESD(nn.Module): # pylint: disable=abstract-method self.decoder = self.decoder.to(devices.device, dtype=self.dtype) def guess_latent_channels(self, decoder_path, encoder_path): - return 16 if ("f1" in encoder_path or "f1" in decoder_path) or ("sd3" in encoder_path or "sd3" in decoder_path) else 4 + if "f2" in encoder_path or "f2" in decoder_path: + return 32 # FLUX.2 uses 32 latent channels + if ("f1" in encoder_path or "f1" in decoder_path) or ("sd3" in encoder_path or "sd3" in decoder_path): + return 16 + return 4 @staticmethod def scale_latents(x): From 749371364b01e1cb17eac4595406d63c092ce529 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 16 Jan 2026 03:03:47 +0000 Subject: [PATCH 3/5] lint --- html/reference-distilled.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/html/reference-distilled.json b/html/reference-distilled.json index 36e583ec2..123d8e6e4 100644 --- a/html/reference-distilled.json +++ b/html/reference-distilled.json @@ -182,4 +182,4 @@ "size": 18.5, "date": "2025 January" } -} \ No newline at end of file +} From 5e2bc013679f7bf3b027113216986a8e19a3ff57 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 16 Jan 2026 03:06:48 +0000 Subject: [PATCH 4/5] fix(taesd): remove unused variable warning --- modules/sd_vae_taesd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index d076dacbb..80cf5de69 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -161,7 +161,7 @@ def decode(latents): shared.log.debug(f'Decode: type="taesd" variant="{variant}" input={latents.shape} tensor={tensor.shape}') # FLUX.2 has 128 latent channels that need reshaping to 32 channels for TAESD if variant == 'TAE FLUX.2' and len(tensor.shape) == 4 and tensor.shape[1] == 128: - b, c, h, w = tensor.shape + b, _c, h, w = tensor.shape tensor = tensor.reshape(b, 32, h * 2, w * 2) if variant.startswith('TAESD') or variant in {'TAE FLUX.1', 'TAE FLUX.2', 'TAE SD3'}: image = vae.decoder(tensor).clamp(0, 1).detach() From eaa8dbcd429c09373ccec92069be8a754434874a Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 16 Jan 2026 03:24:39 +0000 Subject: [PATCH 5/5] fix: correct comments and cleanup model descriptions - Fix Klein text encoder comment to specify correct sizes per variant - Lock TAESD decode logging behind SD_PREVIEW_DEBUG env var - Fix misleading comment about FLUX.2 128-channel reshape (is fallback) - Remove VRAM requirements from model descriptions in reference files --- html/reference-distilled.json | 4 ++-- html/reference.json | 6 +++--- modules/sd_vae_taesd.py | 8 ++++++-- pipelines/model_flux2_klein.py | 2 +- 4 files changed, 12 insertions(+), 8 deletions(-) diff --git a/html/reference-distilled.json b/html/reference-distilled.json index 123d8e6e4..e77c3b4b6 100644 --- a/html/reference-distilled.json +++ b/html/reference-distilled.json @@ -165,7 +165,7 @@ "Black Forest Labs FLUX.2 Klein 4B": { "path": "black-forest-labs/FLUX.2-klein-4B", "preview": "black-forest-labs--FLUX.2-klein-4B.jpg", - "desc": "FLUX.2-klein-4B is a 4 billion parameter size-distilled version of FLUX.2-dev optimized for consumer GPUs. Achieves sub-second inference with 4 steps while fitting in ~13GB VRAM. Supports both text-to-image generation and multi-reference image editing. Apache 2.0 licensed.", + "desc": "FLUX.2-klein-4B is a 4 billion parameter size-distilled version of FLUX.2-dev optimized for consumer GPUs. Achieves sub-second inference with 4 steps. Supports both text-to-image generation and multi-reference image editing. Apache 2.0 licensed.", "skip": true, "tags": "distilled", "extras": "sampler: Default, cfg_scale: 4.0, steps: 4", @@ -175,7 +175,7 @@ "Black Forest Labs FLUX.2 Klein 9B": { "path": "black-forest-labs/FLUX.2-klein-9B", "preview": "black-forest-labs--FLUX.2-klein-9B.jpg", - "desc": "FLUX.2-klein-9B is a 9 billion parameter size-distilled version of FLUX.2-dev. Higher quality than 4B variant with sub-second inference using 4 steps. Requires ~29GB VRAM. Supports text-to-image and multi-reference editing. Non-commercial license.", + "desc": "FLUX.2-klein-9B is a 9 billion parameter size-distilled version of FLUX.2-dev. Higher quality than 4B variant with sub-second inference using 4 steps. Supports text-to-image and multi-reference editing. Non-commercial license.", "skip": true, "tags": "distilled", "extras": "sampler: Default, cfg_scale: 4.0, steps: 4", diff --git a/html/reference.json b/html/reference.json index 63cd80124..2f1f6562b 100644 --- a/html/reference.json +++ b/html/reference.json @@ -127,7 +127,7 @@ "Black Forest Labs FLUX.2 Klein Base 4B": { "path": "black-forest-labs/FLUX.2-klein-base-4B", "preview": "black-forest-labs--FLUX.2-klein-base-4B.jpg", - "desc": "FLUX.2-klein-base-4B is the undistilled 4 billion parameter base model of FLUX.2-klein. Requires 50 inference steps for full quality but offers flexibility for fine-tuning. Fits in ~13GB VRAM. Supports text-to-image and multi-reference editing. Apache 2.0 licensed.", + "desc": "FLUX.2-klein-base-4B is the undistilled 4 billion parameter base model of FLUX.2-klein. Requires 50 inference steps for full quality but offers flexibility for fine-tuning. Supports text-to-image and multi-reference editing. Apache 2.0 licensed.", "skip": true, "extras": "sampler: Default, cfg_scale: 4.0, steps: 50", "size": 8.5, @@ -136,7 +136,7 @@ "Black Forest Labs FLUX.2 Klein Base 9B": { "path": "black-forest-labs/FLUX.2-klein-base-9B", "preview": "black-forest-labs--FLUX.2-klein-base-9B.jpg", - "desc": "FLUX.2-klein-base-9B is the undistilled 9 billion parameter base model of FLUX.2-klein. Requires 50 inference steps for full quality but offers flexibility for fine-tuning. Requires ~29GB VRAM. Supports text-to-image and multi-reference editing. Non-commercial license.", + "desc": "FLUX.2-klein-base-9B is the undistilled 9 billion parameter base model of FLUX.2-klein. Requires 50 inference steps for full quality but offers flexibility for fine-tuning. Supports text-to-image and multi-reference editing. Non-commercial license.", "skip": true, "extras": "sampler: Default, cfg_scale: 4.0, steps: 50", "size": 18.5, @@ -146,7 +146,7 @@ "Z-Image-Turbo": { "path": "Tongyi-MAI/Z-Image-Turbo", "preview": "Tongyi-MAI--Z-Image-Turbo.jpg", - "desc": "Z-Image-Turbo, a distilled version of Z-Image that matches or exceeds leading competitors with only 8 NFEs (Number of Function Evaluations). It offers sub-second inference latency on enterprise-grade H800 GPUs and fits comfortably within 16G VRAM consumer devices. It excels in photorealistic image generation, bilingual text rendering (English & Chinese), and robust instruction adherence.", + "desc": "Z-Image-Turbo, a distilled version of Z-Image that matches or exceeds leading competitors with only 8 NFEs (Number of Function Evaluations). It excels in photorealistic image generation, bilingual text rendering (English & Chinese), and robust instruction adherence.", "skip": true, "extras": "sampler: Default, cfg_scale: 1.0, steps: 9", "size": 20.3, diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index 80cf5de69..1c5dd0dd0 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -12,6 +12,9 @@ import torch from modules import devices, paths, shared +debug = os.environ.get('SD_PREVIEW_DEBUG', None) is not None + + TAESD_MODELS = { 'TAESD 1.3 Mocha Croissant': { 'fn': 'taesd_13_', 'uri': 'https://github.com/madebyollin/taesd/raw/7f572ca629c9b0d3c9f71140e5f501e09f9ea280', 'model': None }, 'TAESD 1.2 Chocolate-Dipped Shortbread': { 'fn': 'taesd_12_', 'uri': 'https://github.com/madebyollin/taesd/raw/8909b44e3befaa0efa79c5791e4fe1c4d4f7884e', 'model': None }, @@ -158,8 +161,9 @@ def decode(latents): dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16 tensor = latents.unsqueeze(0) if len(latents.shape) == 3 else latents tensor = tensor.detach().clone().to(devices.device, dtype=dtype) - shared.log.debug(f'Decode: type="taesd" variant="{variant}" input={latents.shape} tensor={tensor.shape}') - # FLUX.2 has 128 latent channels that need reshaping to 32 channels for TAESD + if debug: + shared.log.debug(f'Decode: type="taesd" variant="{variant}" input={latents.shape} tensor={tensor.shape}') + # Fallback: reshape packed 128-channel latents to 32 channels if not already unpacked if variant == 'TAE FLUX.2' and len(tensor.shape) == 4 and tensor.shape[1] == 128: b, _c, h, w = tensor.shape tensor = tensor.reshape(b, 32, h * 2, w * 2) diff --git a/pipelines/model_flux2_klein.py b/pipelines/model_flux2_klein.py index 9b8b05ac8..d810821d9 100644 --- a/pipelines/model_flux2_klein.py +++ b/pipelines/model_flux2_klein.py @@ -16,7 +16,7 @@ def load_flux2_klein(checkpoint_info, diffusers_load_config=None): # 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 Qwen3ForCausalLM (8B), shared across all Klein variants + # Load text encoder - Klein uses Qwen3 (4B for Klein-4B, 8B for Klein-9B) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3ForCausalLM, load_config=diffusers_load_config) pipe = diffusers.Flux2KleinPipeline.from_pretrained(