From 0691c26d775b4ef3432e52d5b0a489015ebefcc4 Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Fri, 20 Mar 2026 20:50:46 -0700 Subject: [PATCH] =?UTF-8?q?feat:=20add=20Flux2=20support=20=E2=80=94=20CLI?= =?UTF-8?q?P=20type,=20model=20sampling,=20and=20override?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add "flux2" to clip_types and dual_clip_types lists so CLIPType.FLUX2 is correctly passed to ComfyUI's text encoder loader - Add "flux2" model_sampling_override option in Builder UI dropdown - Implement ModelSamplingFlux2 handler using fixed shift (default 2.02, matching ComfyUI's Flux2 supported_models.py) - Handle "flux2" in config_builder_node.py config generation with default shift value Users with Flux2 models (Klein, etc.) should set: - CLIP Type: flux2 - Model Sampling Override: Flux2 - For best results: enable Advanced Sampling with Flux2 Scheduler Co-Authored-By: Claude Opus 4.6 (1M context) --- config_builder_node.py | 6 ++++-- image_generation.py | 12 ++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/config_builder_node.py b/config_builder_node.py index ef1ef16..17ee2bd 100644 --- a/config_builder_node.py +++ b/config_builder_node.py @@ -644,6 +644,8 @@ class UltimateConfigBuilder: if model_sampling_override == "flux": config["model_sampling_flux_max_shift"] = model_sampling_flux_max_shift config["model_sampling_flux_base_shift"] = model_sampling_flux_base_shift + elif model_sampling_override == "flux2": + config["model_sampling_shift"] = model_sampling_shift if model_sampling_shift else "2.02" else: config["model_sampling_shift"] = model_sampling_shift config["use_advanced_sampling"] = use_advanced_sampling or False @@ -1057,10 +1059,10 @@ async def get_model_lists_endpoint(request): clip_types = [ "stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", - "hidream", "chroma", "ace", "flux" + "hidream", "chroma", "ace", "flux", "flux2" ] dual_clip_types = [ - "sdxl", "sd3", "flux", "hunyuan_video", "hidream", + "sdxl", "sd3", "flux", "flux2", "hunyuan_video", "hidream", "hunyuan_image", "hunyuan_video_15" ] diff --git a/image_generation.py b/image_generation.py index 8addbd3..cdff0ac 100644 --- a/image_generation.py +++ b/image_generation.py @@ -129,6 +129,18 @@ def generate_image( patched_model.add_object_patch("model_sampling", model_sampling) print(f"[GridTester] 🔧 Applied ModelSamplingFlux (max_shift={max_s}, base_shift={base_s}, computed_shift={shift:.4f})") + elif model_sampling_override == "flux2": + # Flux2: uses ModelSamplingFlux with fixed shift (default 2.02, matching ComfyUI's Flux2 supported_models) + sampling_base = comfy.model_sampling.ModelSamplingFlux + sampling_type = comfy.model_sampling.CONST + class ModelSamplingAdvanced(sampling_base, sampling_type): + pass + shift = float(model_sampling_shift) if model_sampling_shift else 2.02 + model_sampling = ModelSamplingAdvanced(patched_model.model.model_config) + model_sampling.set_parameters(shift=shift) + patched_model.add_object_patch("model_sampling", model_sampling) + print(f"[GridTester] 🔧 Applied ModelSamplingFlux2 (shift={shift})") + # === Flux Guidance === # Modify positive conditioning with guidance value (used by Flux 1 models) if flux_guidance_value and float(flux_guidance_value) > 0: