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: