diff --git a/docs/src/content/docs/Users Guide/04.Image Generation/upscale.mdx b/docs/src/content/docs/Users Guide/04.Image Generation/upscale.mdx index cf21a90f328..ddeccdf7898 100644 --- a/docs/src/content/docs/Users Guide/04.Image Generation/upscale.mdx +++ b/docs/src/content/docs/Users Guide/04.Image Generation/upscale.mdx @@ -19,7 +19,7 @@ The panel is split into sections, top to bottom: - **Scale**: the output size is the source size multiplied by this value. The panel shows the resulting input and output size and warns when the output is large enough to need much more time and VRAM. - **Treatment preset**, or **Creativity** and **Structure** directly. Higher Creativity invents more detail and gives the prompt more influence. Higher Structure keeps the source composition more strictly; it is only shown when a Tile ControlNet is in use. - **Detail guidance**: the prompt. It is the project prompt shared with Generate. For upscaling, describe the medium, texture and style rather than new content. -- **Generation**: the **Main model** for the detail pass (**SD1.5**, **SDXL**, **FLUX.1**, **FLUX.2**, **Z-Image**, **Anima**, **Krea-2** or **Qwen Image**), steps, CFG scale, scheduler, seed, and LoRAs that match the main model. FLUX.1, FLUX.2, Z-Image, Krea-2 and Qwen Image always run **Euler**, so they show no scheduler. +- **Generation**: the **Main model** for the detail pass (**SD1.5**, **SDXL**, **FLUX.1**, **FLUX.2**, **Z-Image**, **Anima**, **Krea-2** or **Qwen Image**), steps, CFG scale, scheduler, seed, and LoRAs that match the main model. The scheduler list only holds schedulers that work for upscaling (see below). - **Advanced**: the components the main model needs (see below), tile size and tile overlap, VAE and VAE precision, and CLIP skip for SD1.5. ## What each model family needs @@ -44,7 +44,7 @@ Z-Image, Anima, Krea-2 and Qwen Image count the start of the detail pass in step FLUX.2 starts where FLUX.1 does and then spreads all of your Steps over the rest of the pass, so 30 steps run 30 times. Picking a FLUX.2 model therefore sets its own step count: 4 for Klein 4B and 9B, 28 for Klein Base and [dev]. Switching to another family afterwards sets Steps back to 30. -The panel offers only schedulers that start where Creativity puts the pass: **Euler** on FLUX.1, FLUX.2, Z-Image, Krea-2 and Qwen Image, and **Euler**, **DPM++ 2M**, **DPM++ 2M SDE** and **ER-SDE** on Anima. Heun and LCM build their own noise schedule, so they would treat the enlarged image as pure noise or leave noise in it. +The panel offers only schedulers that start where Creativity puts the pass: **Euler** and **Heun** on FLUX.1 and Z-Image, and **Euler**, **DPM++ 2M**, **DPM++ 2M SDE** and **ER-SDE** on Anima. LCM builds its own noise schedule, so it would treat the enlarged image as pure noise. Anima's Heun starts below that point and leaves noise in the image. FLUX.2, Krea-2 and Qwen Image always run **Euler**, so they show no scheduler. The Z-Image Tile ControlNet is optional because the detail pass runs over the whole image at once: there are no tile seams for it to hold together. Adding it lets **Structure** pull the result back toward the enlarged image, at the cost of a second network's memory and time. diff --git a/invokeai/app/invocations/flux2/flux2_denoise.py b/invokeai/app/invocations/flux2/flux2_denoise.py index 928d2e1d117..58228653d53 100644 --- a/invokeai/app/invocations/flux2/flux2_denoise.py +++ b/invokeai/app/invocations/flux2/flux2_denoise.py @@ -31,7 +31,6 @@ from invokeai.backend.flux2.extensions.regional_prompting_extension import Flux2RegionalPromptingExtension from invokeai.backend.flux2.ref_image_extension import Flux2RefImageExtension from invokeai.backend.flux2.sampling_utils import ( - FLUX2_TXT2IMG_SCHEDULER_KWARGS, compute_empirical_mu, generate_img_ids_flux2, get_noise_flux2, @@ -379,36 +378,24 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: transformer_config = context.models.get_config(self.transformer.transformer) assert transformer_config.base == BaseModelType.Flux2 and transformer_config.type == ModelType.Main - # Calculate the timestep schedule using FLUX.2 specific schedule - # This matches diffusers' Flux2Pipeline implementation - # Note: Schedule shifting is handled by the scheduler via mu parameter + # The FLUX.2 schedule: diffusers' Flux2KleinPipeline linear sigmas with the exponential shift + # from mu, which its scheduler applies internally. image_seq_len = packed_h * packed_w - timesteps = get_schedule_flux2( - num_steps=self.num_steps, - image_seq_len=image_seq_len, - ) - # Compute mu for dynamic schedule shifting (used by FlowMatchEulerDiscreteScheduler) mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=self.num_steps) + timesteps = time_shift_flux2(get_schedule_flux2(num_steps=self.num_steps, image_seq_len=image_seq_len), mu) + + # Every path runs this shifted schedule, and it is shifted before clipping so the clip selects + # from the schedule the model is actually run on. That keeps denoising_start and denoising_end + # on one scale: a pass that ends at denoising_end = x stops at the sigma a later pass with + # denoising_start = x begins at. The clip works by sigma value, not by step count, so shifting + # first only changes how many steps fall below a given sigma. + timesteps = clip_timestep_schedule_fractional(timesteps, self.denoising_start, self.denoising_end) - # img2img and inpainting step this schedule manually (see the scheduler setup below), while - # txt2img hands it to the scheduler, which applies the exponential shift from mu itself. Apply - # the same shift here so both paths follow one schedule, and apply it before clipping so the - # clip selects from the schedule the model is actually run on. The clip works by sigma value, - # not by step count, so shifting first does not change what a given denoising_start means -- - # it only changes how many steps fall below that sigma. - # - # Keyed on the same condition that decides whether a scheduler gets built below. Manual Euler - # also runs when the requested scheduler is unavailable (e.g. "lcm" when FlowMatchLCMScheduler - # cannot be imported); keying on denoise_mask/denoising_start alone would leave that run on the - # unshifted linear schedule -- the one this fix exists to avoid. + # Manual Euler steps img2img and inpainting, and also runs when the requested scheduler is + # unavailable (e.g. "lcm" when FlowMatchLCMScheduler cannot be imported). uses_manual_euler = ( self.denoise_mask is not None or self.denoising_start > 1e-5 or self.scheduler not in FLUX_SCHEDULER_MAP ) - if uses_manual_euler: - timesteps = time_shift_flux2(timesteps, mu) - - # Clip the timesteps schedule based on denoising_start and denoising_end - timesteps = clip_timestep_schedule_fractional(timesteps, self.denoising_start, self.denoising_end) # Clipping keeps only the sigmas that fall inside the denoising window, so a narrow window # runs fewer steps than were requested. Under the shift that is severe: at 4 steps and @@ -424,13 +411,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: if self.add_noise: assert noise is not None # Noise the init latents using the first timestep from the clipped - # InvokeAI schedule. - # - # Known limitation: if a scheduler later uses a different first - # effective timestep/sigma than this precomputed schedule, the - # img2img preblend below may not match that scheduler exactly. - # This is an existing pipeline limitation and applies to both - # seed-generated noise and externally supplied noise. + # InvokeAI schedule, which manual Euler steps from. t_0 = timesteps[0] x = t_0 * noise + (1.0 - t_0) * init_latents else: @@ -495,27 +476,12 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: num_steps = len(timesteps) - 1 cfg_scale_list = [self.cfg_scale] * num_steps - # Create scheduler with FLUX.2 Klein configuration - # For inpainting/img2img, use manual Euler stepping to preserve the exact - # clipped timestep schedule used for the initial latent/noise preblend. - # For txt2img, use the scheduler with dynamic shifting for optimal results. - # - # This split is intentional. Reusing a scheduler for img2img here can - # change the first effective timestep/sigma and break parity with the - # preblend computed above. + # txt2img uses the selected Diffusers scheduler; img2img and inpainting keep manual Euler on the + # respaced schedule above. The schedule is already shifted, so the scheduler must not shift it + # again: no dynamic shifting, and the default static shift of 1.0. scheduler = None if not uses_manual_euler: - # Only use scheduler for txt2img - use manual Euler for inpainting to preserve exact timesteps - scheduler_class = FLUX_SCHEDULER_MAP[self.scheduler] - # FlowMatchHeunDiscreteScheduler only supports num_train_timesteps and shift parameters - # FlowMatchEulerDiscreteScheduler and FlowMatchLCMScheduler support dynamic shifting - if self.scheduler == "heun": - scheduler = scheduler_class( - num_train_timesteps=1000, - shift=3.0, - ) - else: - scheduler = scheduler_class(**FLUX2_TXT2IMG_SCHEDULER_KWARGS) + scheduler = FLUX_SCHEDULER_MAP[self.scheduler](num_train_timesteps=1000) # Prepare reference image extension for FLUX.2 Klein built-in editing ref_image_extension = None @@ -638,7 +604,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: neg_txt=neg_txt, neg_txt_ids=neg_txt_ids, scheduler=scheduler, - mu=mu, inpaint_extension=inpaint_extension, img_cond_seq=img_cond_seq, img_cond_seq_ids=img_cond_seq_ids, diff --git a/invokeai/app/invocations/z_image/z_image_denoise.py b/invokeai/app/invocations/z_image/z_image_denoise.py index a96d0962b40..cab698269dc 100644 --- a/invokeai/app/invocations/z_image/z_image_denoise.py +++ b/invokeai/app/invocations/z_image/z_image_denoise.py @@ -1,4 +1,3 @@ -import inspect import math from contextlib import ExitStack from typing import Callable, Iterator, Optional @@ -27,7 +26,12 @@ from invokeai.app.invocations.vae.z_image_image_to_latents import ZImageImageToLatentsInvocation from invokeai.app.invocations.z_image.z_image_control import ZImageControlField from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.backend.flux.schedulers import ZIMAGE_SCHEDULER_LABELS, ZIMAGE_SCHEDULER_MAP, ZIMAGE_SCHEDULER_NAME_VALUES +from invokeai.backend.flux.schedulers import ( + ZIMAGE_SCHEDULER_LABELS, + ZIMAGE_SCHEDULER_MAP, + ZIMAGE_SCHEDULER_NAME_VALUES, + set_heun_timesteps_from_sigmas, +) from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_TRANSFORMER_PREFIX @@ -407,6 +411,16 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: end_idx = int(self.denoising_end * (total_sigmas - 1)) + 1 sigmas = sigmas[start_idx:end_idx] + is_heun = self.scheduler == "heun" + if is_heun: + # Heun divides by every sigma it evaluates the model at, so an interval of zero length at sigma 0 is + # 0/0 and turns the latents into NaN; an explicit shift of 0 schedules [1, 0, ..., 0]. Such an interval + # has nothing to integrate (Euler steps through it unchanged), so Heun runs only the others and counts + # its progress over them. With none left, the pass returns the blended latents, as Euler would. + sigmas = [ + sigma for sigma, sigma_next in zip(sigmas[:-1], sigmas[1:], strict=True) if sigma != sigma_next + ] + sigmas[-1:] + total_steps = len(sigmas) - 1 # Load input latents if provided (image-to-image) @@ -428,13 +442,8 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: if self.add_noise: assert noise is not None # Noise the init latents using the first sigma from the clipped - # InvokeAI schedule. - # - # Known limitation: if the selected scheduler later starts from a - # different first effective sigma/timestep than sigmas[0], the - # img2img preblend below may not match that scheduler exactly. - # This is an existing pipeline limitation and affects both - # internally generated noise and externally supplied noise. + # InvokeAI schedule. Euler and Heun start from this sigma; LCM + # creates its own schedule, whose first sigma can differ. s_0 = sigmas[0] latents = s_0 * noise + (1.0 - s_0) * init_latents else: @@ -474,16 +483,13 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: num_train_timesteps=1000, shift=1.0, ) - # Set timesteps - LCM uses its own sigma schedule (num_inference_steps), - # while other schedulers can use custom sigmas if supported - is_lcm = self.scheduler == "lcm" - set_timesteps_sig = inspect.signature(scheduler.set_timesteps) - if not is_lcm and "sigmas" in set_timesteps_sig.parameters: - scheduler.set_timesteps(sigmas=sigmas, device=device) + if is_heun: + # Heun runs the shifted, clipped schedule above, so it starts at the sigma the img2img + # preblend used and stops at the requested end. + set_heun_timesteps_from_sigmas(scheduler, sigmas[:-1], sigmas[-1], device) else: - # LCM or a scheduler without custom-sigma support computes its own - # schedule from num_inference_steps. That can diverge from sigmas[0] - # used in the img2img preblend above. + # LCM owns its schedule: clipping only sets its step count. Its first sigma can + # differ from sigmas[0] used in the img2img preblend above. scheduler.set_timesteps(num_inference_steps=total_steps, device=device) # For Heun scheduler, the number of actual steps may differ @@ -685,7 +691,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: sigma_curr = sched_timestep.item() / scheduler.config.num_train_timesteps # For Heun scheduler, track if we're in first or second order step - is_heun = hasattr(scheduler, "state_in_first_order") in_first_order = scheduler.state_in_first_order if is_heun else True # Timestep tensor for Z-Image model @@ -766,9 +771,10 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: if inpaint_extension is not None: latents = inpaint_extension.merge_intermediate_latents_with_init_latents(latents, sigma_prev) - # For Heun, only increment user step after second-order step completes + # For Heun, a user step completes with its second-order evaluation. The final + # interval has none: Diffusers ends it with a first-order step to the terminal sigma. if is_heun: - if not in_first_order: + if not in_first_order or step_index == num_scheduler_steps - 1: user_step += 1 # Only call step_callback if we haven't exceeded total_steps if user_step <= total_steps: @@ -776,7 +782,8 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: step_callback( PipelineIntermediateState( step=user_step, - order=2, + # Progress over completed user steps, not the solver's order. + order=1, total_steps=total_steps, timestep=int(sigma_curr * 1000), latents=latents, diff --git a/invokeai/backend/flux/denoise.py b/invokeai/backend/flux/denoise.py index 2245d752d67..30766a559a7 100644 --- a/invokeai/backend/flux/denoise.py +++ b/invokeai/backend/flux/denoise.py @@ -13,7 +13,11 @@ from invokeai.backend.flux.extensions.xlabs_controlnet_extension import XLabsControlNetExtension from invokeai.backend.flux.extensions.xlabs_ip_adapter_extension import XLabsIPAdapterExtension from invokeai.backend.flux.model import Flux -from invokeai.backend.flux.schedulers import is_flow_match_lcm_scheduler, set_heun_timesteps_from_sigmas +from invokeai.backend.flux.schedulers import ( + is_flow_match_lcm_scheduler, + set_heun_timesteps_from_sigmas, + set_timesteps_from_sigmas, +) from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.util.devices import TorchDevice @@ -59,10 +63,7 @@ def denoise( set_heun_timesteps_from_sigmas(scheduler, timesteps[:-1], timesteps[-1], img.device) elif isinstance(scheduler, FlowMatchEulerDiscreteScheduler): # InvokeAI supplies model-evaluation sigmas plus its requested terminal endpoint. - scheduler.set_timesteps(sigmas=timesteps[:-1], device=img.device) - scheduler.sigmas[-1] = torch.as_tensor( - timesteps[-1], dtype=scheduler.sigmas.dtype, device=scheduler.sigmas.device - ) + set_timesteps_from_sigmas(scheduler, timesteps[:-1], timesteps[-1], img.device) elif is_flow_match_lcm_scheduler(scheduler): # LCM owns its schedule; clipping controls the requested evaluation count, # while its exact start/end sigmas remain scheduler-defined. diff --git a/invokeai/backend/flux/schedulers.py b/invokeai/backend/flux/schedulers.py index 64ab9a69932..08d808c99b0 100644 --- a/invokeai/backend/flux/schedulers.py +++ b/invokeai/backend/flux/schedulers.py @@ -68,6 +68,19 @@ def set_heun_timesteps_from_sigmas( scheduler.sample = None +def set_timesteps_from_sigmas( + scheduler: SchedulerMixin, + sigmas: list[float], + terminal_sigma: float, + device: str | torch.device | None, +) -> None: + """Set InvokeAI's shifted model sigmas on a scheduler whose set_timesteps() takes them (Euler, LCM).""" + # set_timesteps() appends a terminal sigma of 0, and step() reads sigmas[step_index + 1], so the + # requested terminal sigma has to replace it for a pass that ends before the schedule does. + scheduler.set_timesteps(sigmas=sigmas, device=device) + scheduler.sigmas[-1] = torch.as_tensor(terminal_sigma, dtype=scheduler.sigmas.dtype, device=scheduler.sigmas.device) + + def is_flow_match_lcm_scheduler(scheduler: SchedulerMixin) -> bool: """Check the optional LCM scheduler without requiring it in older Diffusers builds.""" return _HAS_LCM and isinstance(scheduler, FlowMatchLCMScheduler) diff --git a/invokeai/backend/flux2/denoise.py b/invokeai/backend/flux2/denoise.py index 3699db70ef3..d39688136d6 100644 --- a/invokeai/backend/flux2/denoise.py +++ b/invokeai/backend/flux2/denoise.py @@ -4,14 +4,19 @@ which use Qwen3 as the text encoder instead of CLIP+T5. """ -import inspect import math from typing import Any, Callable -import numpy as np import torch +from diffusers import FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler +from diffusers.schedulers.scheduling_utils import SchedulerMixin from tqdm import tqdm +from invokeai.backend.flux.schedulers import ( + is_flow_match_lcm_scheduler, + set_heun_timesteps_from_sigmas, + set_timesteps_from_sigmas, +) from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.util.devices import TorchDevice @@ -32,10 +37,8 @@ def denoise( # Negative conditioning for CFG neg_txt: torch.Tensor | None = None, neg_txt_ids: torch.Tensor | None = None, - # Scheduler for stepping (e.g., FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler) - scheduler: Any = None, - # Dynamic shifting parameter for FLUX.2 Klein (computed from image resolution) - mu: float | None = None, + # Diffusers Euler, Heun or LCM scheduler that runs `timesteps`. Manual Euler when omitted. + scheduler: SchedulerMixin | None = None, # Inpainting extension for merging latents during denoising inpaint_extension: RectifiedFlowInpaintExtension | None = None, # Reference image conditioning (multi-reference image editing) @@ -64,16 +67,16 @@ def denoise( img_ids: Image position IDs tensor. txt: Text encoder hidden states (Qwen3 embeddings). txt_ids: Text position IDs tensor. - timesteps: List of timesteps for denoising schedule (linear sigmas from 1.0 to 1/n). + timesteps: The shifted, clipped sigma schedule: the sigmas the model is evaluated at, followed + by the terminal sigma the last step ends at. step_callback: Callback function for progress updates. guidance: Guidance strength. Inert for all current FLUX.2 Klein variants (their guidance_embeds projection weights are absent/zero). cfg_scale: List of CFG scale values per step. neg_txt: Negative text embeddings for CFG (optional). neg_txt_ids: Negative text position IDs (optional). - scheduler: Optional diffusers scheduler (Euler, Heun, LCM). If None, uses manual Euler. - mu: Dynamic shifting parameter computed from image resolution. Required when scheduler - has use_dynamic_shifting=True. + scheduler: Optional diffusers scheduler (Euler, Heun, LCM) without a shift of its own; it runs + `timesteps` as given. If None, uses manual Euler. Returns: Denoised latent tensor. @@ -105,32 +108,18 @@ def denoise( # Use scheduler if provided use_scheduler = scheduler is not None if use_scheduler: - # Set up scheduler with sigmas and mu for dynamic shifting - # Convert timesteps (0-1 range) to sigmas for the scheduler - # The scheduler will apply dynamic shifting internally using mu (if enabled in scheduler config) - sigmas = np.array(timesteps[:-1], dtype=np.float32) # Exclude final 0.0 - - # Check if scheduler supports sigmas parameter using inspect.signature - # FlowMatchHeunDiscreteScheduler and FlowMatchLCMScheduler don't support sigmas - set_timesteps_sig = inspect.signature(scheduler.set_timesteps) - supports_sigmas = "sigmas" in set_timesteps_sig.parameters - if supports_sigmas and mu is not None: - # Pass mu if provided - it will only be used if scheduler has use_dynamic_shifting=True - scheduler.set_timesteps(sigmas=sigmas.tolist(), mu=mu, device=img.device) - elif supports_sigmas: - scheduler.set_timesteps(sigmas=sigmas.tolist(), device=img.device) + # Every scheduler runs InvokeAI's schedule: the model-evaluation sigmas plus the requested + # terminal sigma, which is not 0 when denoising_end < 1. Unlike FLUX.1, LCM gets it too: the + # FLUX.2 schedule is shifted from mu and these schedulers do not shift, so a schedule LCM built + # itself would be the unshifted linear one. + is_heun = isinstance(scheduler, FlowMatchHeunDiscreteScheduler) + if is_heun: + set_heun_timesteps_from_sigmas(scheduler, timesteps[:-1], timesteps[-1], img.device) + elif isinstance(scheduler, FlowMatchEulerDiscreteScheduler) or is_flow_match_lcm_scheduler(scheduler): + set_timesteps_from_sigmas(scheduler, timesteps[:-1], timesteps[-1], img.device) else: - # Scheduler doesn't support sigmas (e.g., Heun, LCM) - use num_inference_steps - # - # Important for img2img callers: if the initial latent/noise blend was - # computed from a separate pre-scheduler schedule, that preblend may not - # match this scheduler's true first step exactly. - scheduler_kwargs: dict[str, Any] = {"num_inference_steps": len(sigmas), "device": img.device} - if mu is not None and "mu" in set_timesteps_sig.parameters: - scheduler_kwargs["mu"] = mu - scheduler.set_timesteps(**scheduler_kwargs) + raise TypeError(f"Unsupported FLUX.2 scheduler: {type(scheduler).__name__}") num_scheduler_steps = len(scheduler.timesteps) - is_heun = hasattr(scheduler, "state_in_first_order") user_step = 0 pbar = tqdm(total=total_steps, desc=f"Denoising{TorchDevice.get_session_device_label()}") @@ -187,6 +176,14 @@ def denoise( neg_pred = neg_pred[:, :original_seq_len] pred = neg_pred + step_cfg_scale * (pred - neg_pred) + # Use the sample and sigma of this model evaluation; after the step, `img` is already at + # the next sigma. For Heun the interval's preview comes from its second-order evaluation, or + # from the final first-order step, which has none. + should_emit_preview = ( + not is_heun or not in_first_order or step_index == num_scheduler_steps - 1 + ) and user_step < total_steps + preview_img = img - t_curr * pred if should_emit_preview else None + # Use scheduler.step() for the update. `generator` matters for LCM, which re-noises every # step. Euler and Heun don't use step noise here: with `s_churn` at 0 it is never applied. step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img, generator=generator) @@ -202,13 +199,14 @@ def denoise( if inpaint_extension is not None: img = inpaint_extension.merge_intermediate_latents_with_init_latents(img, t_prev) - # For Heun, only increment user step after second-order step completes + # For Heun, a user step completes with its second-order evaluation. The final interval has + # none: Diffusers ends it with a first-order step to the terminal sigma. if is_heun: - if not in_first_order: + if not in_first_order or step_index == num_scheduler_steps - 1: user_step += 1 if user_step <= total_steps: pbar.update(1) - preview_img = img - t_curr * pred + assert preview_img is not None if inpaint_extension is not None: preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents( preview_img, 0.0 @@ -216,7 +214,8 @@ def denoise( step_callback( PipelineIntermediateState( step=user_step, - order=2, + # Progress over completed user steps, not the solver's order. + order=1, total_steps=total_steps, timestep=int(t_curr * 1000), latents=preview_img, @@ -226,7 +225,7 @@ def denoise( user_step += 1 if user_step <= total_steps: pbar.update(1) - preview_img = img - t_curr * pred + assert preview_img is not None if inpaint_extension is not None: preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_img, 0.0) step_callback( diff --git a/invokeai/backend/flux2/sampling_utils.py b/invokeai/backend/flux2/sampling_utils.py index 88649603fb7..025daa3e06b 100644 --- a/invokeai/backend/flux2/sampling_utils.py +++ b/invokeai/backend/flux2/sampling_utils.py @@ -109,21 +109,6 @@ def unpack_flux2(x: torch.Tensor, height: int, width: int) -> torch.Tensor: ) -# Config for the scheduler the FLUX.2 txt2img path builds. Shared with the schedule-shift test so the -# manual shift used by img2img/inpainting is always checked against the scheduler the node actually -# uses, rather than against a copy that can drift out of sync. -FLUX2_TXT2IMG_SCHEDULER_KWARGS = { - "num_train_timesteps": 1000, - "shift": 3.0, - "use_dynamic_shifting": True, - "base_shift": 0.5, - "max_shift": 1.15, - "base_image_seq_len": 256, - "max_image_seq_len": 4096, - "time_shift_type": "exponential", -} - - def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: """Compute mu for FLUX.2 schedule shifting. @@ -154,14 +139,13 @@ def get_schedule_flux2( ) -> list[float]: """Get linear timestep schedule for FLUX.2. - Returns a linear sigma schedule from 1.0 to 1/num_steps. - The actual schedule shifting is handled by the FlowMatchEulerDiscreteScheduler - using the mu parameter and use_dynamic_shifting=True. + Returns a linear sigma schedule from 1.0 to 1/num_steps, unshifted. Callers shift it with + ``time_shift_flux2()``. Args: num_steps: Number of denoising steps. image_seq_len: Number of image tokens (packed_h * packed_w). Currently unused, - but kept for API compatibility. The scheduler computes shifting internally. + but kept for API compatibility. Returns: List of linear sigmas from 1.0 to 1/num_steps, plus final 0.0. @@ -169,7 +153,6 @@ def get_schedule_flux2( import numpy as np # Create linear sigmas from 1.0 to 1/num_steps - # The scheduler will apply dynamic shifting using mu parameter sigmas = np.linspace(1.0, 1 / num_steps, num_steps) sigmas_list = [float(s) for s in sigmas] @@ -180,18 +163,15 @@ def get_schedule_flux2( def time_shift_flux2(sigmas: list[float], mu: float) -> list[float]: - """Apply the exponential schedule shift that the txt2img scheduler applies internally. + """Apply the exponential schedule shift that diffusers' Flux2KleinPipeline scheduler applies. - ``get_schedule_flux2()`` returns an unshifted linear schedule because the - FlowMatchEulerDiscreteScheduler used for txt2img shifts it itself from ``mu``. Code paths that - step the schedule manually (img2img, inpainting) must apply the shift themselves, otherwise they - run the model on a completely different sigma trajectory than txt2img does: with 9 steps the - shifted schedule bottoms out at sigma 0.485, the linear one at 0.111. A distilled model such as - FLUX.2 Klein is never trained at the low end of the linear schedule and leaves a grainy residue - there. + ``get_schedule_flux2()`` returns an unshifted linear schedule. Every FLUX.2 path runs it shifted, + otherwise the model sees a completely different sigma trajectory: with 9 steps the shifted + schedule bottoms out at sigma 0.485, the linear one at 0.111. A distilled model such as FLUX.2 + Klein is never trained at the low end of the linear schedule and leaves a grainy residue there. Mirrors diffusers' ``FlowMatchEulerDiscreteScheduler._time_shift_exponential`` with an exponent - of 1.0, which is how the scheduler invokes it:: + of 1.0, which is how the scheduler invokes it under dynamic shifting:: sigma' = exp(mu) / (exp(mu) + (1 / sigma - 1)) diff --git a/invokeai/frontend/webv2/src/app/UpscaleCapabilityGate.browser.test.tsx b/invokeai/frontend/webv2/src/app/UpscaleCapabilityGate.browser.test.tsx index 40c29fbab67..9e7d27aa04b 100644 --- a/invokeai/frontend/webv2/src/app/UpscaleCapabilityGate.browser.test.tsx +++ b/invokeai/frontend/webv2/src/app/UpscaleCapabilityGate.browser.test.tsx @@ -152,8 +152,19 @@ describe('Upscale without the served architecture table', () => { ?.click() ); - // Euler is all Z-Image has that keeps to the blend, so there is no choice to show. + // Euler and Heun both start where the blend left the latents on Z-Image; LCM does not. expect(host.textContent).toContain('widgets.upscale.cfgScale'); - expect(host.querySelector('[aria-label="widgets.upscale.scheduler"]')).toBeNull(); + const scheduler = host.querySelector('input[aria-label="widgets.upscale.scheduler"]'); + expect(scheduler?.value).toBe('Euler'); + + await settle(() => + scheduler + ?.closest('[data-scope="combobox"][data-part="root"]') + ?.querySelector('button[aria-label="common.openSelector"]') + ?.click() + ); + + const offered = [...document.querySelectorAll('[role="option"]')].map((option) => option.textContent); + expect(offered).toEqual(['Euler', 'Heun (2nd order)']); }); }); diff --git a/invokeai/frontend/webv2/src/features/upscale/core/__snapshots__/upscaleIndexClippedStarts.json b/invokeai/frontend/webv2/src/features/upscale/core/__snapshots__/upscaleIndexClippedStarts.json index 5c238a8fa37..1de557686a8 100644 --- a/invokeai/frontend/webv2/src/features/upscale/core/__snapshots__/upscaleIndexClippedStarts.json +++ b/invokeai/frontend/webv2/src/features/upscale/core/__snapshots__/upscaleIndexClippedStarts.json @@ -1528,7 +1528,8 @@ "euler" ], "z-image": [ - "euler" + "euler", + "heun" ] } } diff --git a/invokeai/frontend/webv2/src/features/upscale/core/graph.test.ts b/invokeai/frontend/webv2/src/features/upscale/core/graph.test.ts index 819acfeac8f..7e65acbe778 100644 --- a/invokeai/frontend/webv2/src/features/upscale/core/graph.test.ts +++ b/invokeai/frontend/webv2/src/features/upscale/core/graph.test.ts @@ -716,12 +716,15 @@ describe('Creativity on index-clipping architectures', () => { }); describe('schedulers on the upscale pass', () => { - // A stored Heun or LCM on a flow model -- left from SD, recalled, or picked before the list was - // narrowed -- would start the pass from pure noise, or on Anima below the blend. + // A stored LCM on a flow model, or Heun on Anima -- left from SD, recalled, or picked before the + // list was narrowed -- would start the pass from pure noise, or on Anima below the blend. FLUX.2 + // steps with Euler whenever the pass starts above pure noise. it.each([ ['flux', () => createFluxValues(), 'lcm', 'euler'], + ['flux', () => createFluxValues(), 'heun', 'heun'], ['flux2', () => createFlux2Values('klein_4b'), 'heun', 'euler'], - ['z-image', () => createZImageValues(), 'heun', 'euler'], + ['z-image', () => createZImageValues(), 'lcm', 'euler'], + ['z-image', () => createZImageValues(), 'heun', 'heun'], ['anima', () => createAnimaValues(), 'heun', 'euler'], ['anima', () => createAnimaValues(), 'dpmpp_2m', 'dpmpp_2m'], ['sd-1', () => createValues('sd-1'), 'heun', 'heun'], diff --git a/invokeai/frontend/webv2/src/features/upscale/core/settings.test.ts b/invokeai/frontend/webv2/src/features/upscale/core/settings.test.ts index ee67a52626c..afa713e6ff9 100644 --- a/invokeai/frontend/webv2/src/features/upscale/core/settings.test.ts +++ b/invokeai/frontend/webv2/src/features/upscale/core/settings.test.ts @@ -654,9 +654,9 @@ describe('schedulers per architecture', () => { const offered = (base: string) => upscaleSchedulerOptions(model(base, 'main', base) as never).map((option) => option.value); - // A single scheduler hides the field. - expect(offered('flux')).toEqual(['euler']); - expect(offered('z-image')).toEqual(['euler']); + // LCM builds its own schedule from pure noise, so it would ignore Creativity. + expect(offered('flux')).toEqual(['euler', 'heun']); + expect(offered('z-image')).toEqual(['euler', 'heun']); expect(offered('anima')).toEqual(['euler', 'dpmpp_2m', 'dpmpp_2m_sde', 'er_sde']); // `krea2_denoise` and `qwen_image_denoise` have no scheduler input. expect(offered('krea-2')).toEqual(['euler']); diff --git a/invokeai/frontend/webv2/src/features/upscale/core/settings.ts b/invokeai/frontend/webv2/src/features/upscale/core/settings.ts index 0a71d06e59f..b7809643b5f 100644 --- a/invokeai/frontend/webv2/src/features/upscale/core/settings.ts +++ b/invokeai/frontend/webv2/src/features/upscale/core/settings.ts @@ -308,8 +308,8 @@ export interface UpscaleArchitecture { readonly indexClipping: UpscaleIndexClipping | null; /** * The schedulers that start where Creativity puts the pass; null where all of the base's do. The - * img2img blend noises the frame to the clipped schedule's first sigma, but Heun and LCM on the - * flow models build their own schedule from pure noise and denoise the frame as if it were noise. + * img2img blend noises the frame to the clipped schedule's first sigma, but LCM on FLUX.1, Z-Image + * and Anima builds its own schedule from pure noise and denoises the frame as if it were noise. * Anima's Heun clips its own schedule, which runs from 1 to nearly 0, so it starts below the blend * and leaves noise behind. `test_upscale_index_clipped_starts.py` checks the index-clipped lists. */ @@ -369,8 +369,8 @@ export const UPSCALE_ARCHITECTURES = { label: 'FLUX.1', negativePrompt: 'never', // `flux/denoise.py` hands Heun and LCM to the same diffusers schedulers as `z_image_denoise`, so - // when the backend contract finds Z-Image's Heun starting at the blend, this list is due as well. - schedulers: ['euler'], + // this list follows the one the backend contract checks for Z-Image. + schedulers: ['euler', 'heun'], takesModelGuidance: false, takesModelSteps: false, tileControlNet: null, @@ -408,7 +408,7 @@ export const UPSCALE_ARCHITECTURES = { }, label: 'Z-Image', negativePrompt: 'cfg-gated', - schedulers: ['euler'], + schedulers: ['euler', 'heun'], takesModelGuidance: true, takesModelSteps: false, tileControlNet: { matchesName: nameIncludes('tile'), required: false }, diff --git a/tests/app/invocations/test_denoise_noise_inputs.py b/tests/app/invocations/test_denoise_noise_inputs.py index e90e2f530ec..10b47dea7d1 100644 --- a/tests/app/invocations/test_denoise_noise_inputs.py +++ b/tests/app/invocations/test_denoise_noise_inputs.py @@ -753,60 +753,6 @@ def test_flux2_partial_denoise_short_circuit_uses_first_clipped_timestep(): assert torch.equal(result, expected) -def test_flux2_lcm_scheduler_setup_passes_mu(): - from invokeai.backend.flux2.denoise import denoise - - class DummyScheduler: - def __init__(self) -> None: - self.received_mu = None - self.timesteps = torch.tensor([750.0, 500.0], dtype=torch.float32) - self.sigmas = torch.tensor([0.75, 0.5, 0.0], dtype=torch.float32) - self.config = SimpleNamespace(num_train_timesteps=1000) - - def set_timesteps(self, num_inference_steps: int, device: str | torch.device, mu: float | None = None) -> None: - self.received_mu = mu - - def step( - self, - model_output: torch.Tensor, - timestep: torch.Tensor, - sample: torch.Tensor, - generator: torch.Generator | None = None, - ): - return SimpleNamespace(prev_sample=sample) - - class DummyModel(torch.nn.Module): - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - timestep: torch.Tensor, - img_ids: torch.Tensor, - txt_ids: torch.Tensor, - guidance: torch.Tensor, - joint_attention_kwargs: dict | None = None, - return_dict: bool = False, - ): - return (torch.zeros_like(hidden_states),) - - scheduler = DummyScheduler() - denoise( - model=DummyModel(), - img=torch.zeros(1, 4, 8), - img_ids=torch.zeros(1, 4, 4, dtype=torch.long), - txt=torch.zeros(1, 4, 8), - txt_ids=torch.zeros(1, 4, 4, dtype=torch.long), - timesteps=[0.75, 0.5, 0.0], - step_callback=lambda _: None, - guidance=1.0, - cfg_scale=[1.0, 1.0], - scheduler=scheduler, - mu=0.42, - ) - - assert scheduler.received_mu == pytest.approx(0.42) - - @pytest.mark.skipif(FLUX_SCHEDULER_MAP.get("lcm") is None, reason="FlowMatchLCMScheduler not available") def test_flux2_lcm_step_noise_follows_the_generator(): """LCM re-noises the sample on every step, so FLUX.2 must draw that noise from the seeded @@ -1035,17 +981,11 @@ def forward(self, x: list[torch.Tensor], t: torch.Tensor, cap_feats: list[torch. _assert_step_generator(step_generators[0], expected_seed) +# Heun is covered through the node in test_z_image_denoise_schedulers.py. @pytest.mark.parametrize( "scheduler_name", [ "euler", - pytest.param( - "heun", - marks=pytest.mark.xfail( - reason="Known img2img preblend mismatch for Z-Image with scheduler-defined first step.", - strict=True, - ), - ), pytest.param( "lcm", marks=pytest.mark.xfail( diff --git a/tests/app/invocations/test_flux2_denoise_wiring.py b/tests/app/invocations/test_flux2_denoise_wiring.py index 3d82eeb1dc3..86c30b852a3 100644 --- a/tests/app/invocations/test_flux2_denoise_wiring.py +++ b/tests/app/invocations/test_flux2_denoise_wiring.py @@ -6,25 +6,31 @@ pass the whole suite: * dropping the ``time_shift_flux2`` call, so img2img runs the unshifted linear schedule again; -* shifting unconditionally, which shifts txt2img twice (the scheduler already applies ``mu``); +* leaving txt2img to a scheduler that shifts the schedule itself, after the clip: a txt2img pass + with ``denoising_end`` then stops at a different sigma than the img2img pass that continues it + starts at, and Heun ran a static shift of its own; * re-normalizing ``noise_packed`` before the inpaint extension is built, bringing back the half of the noise-scaling bug that lives in the unmasked region; * handing ``_prepare_normalized_start_latents`` a ``t_0`` from the unclipped schedule; * skipping the respacing of the clipped schedule, so img2img at low strength runs one or two steps instead of the requested number and the final Euler jump flattens the source image's detail. -These tests drive the real ``_run_diffusion`` up to the transformer load and assert on what the -wiring actually produced. +These tests drive the real ``_run_diffusion`` up to the transformer load, or through the denoising loop +with a stand-in transformer, and assert on what the wiring actually produced. """ +from contextlib import nullcontext +from typing import Any from unittest.mock import MagicMock, patch import pytest import torch +from diffusers import FlowMatchHeunDiscreteScheduler import invokeai.app.invocations.flux2_denoise as flux2_denoise from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional +from invokeai.backend.flux.schedulers import FLUX_SCHEDULER_MAP, set_heun_timesteps_from_sigmas from invokeai.backend.flux2.sampling_utils import ( compute_empirical_mu, get_schedule_flux2, @@ -51,6 +57,19 @@ class _StopBeforeLoad(Exception): """Raised in place of entering the transformer's device context, to end _run_diffusion early.""" +class _RecordingTransformer(torch.nn.Module): + """Stands in for the FLUX.2 transformer: records the sigma of every evaluation, predicts zero.""" + + def __init__(self) -> None: + super().__init__() + self.sigmas: list[float] = [] + + def forward(self, hidden_states: torch.Tensor, timestep: torch.Tensor, **kwargs: Any) -> tuple[torch.Tensor]: + del kwargs + self.sigmas.append(float(timestep[0])) + return (torch.zeros_like(hidden_states),) + + def _bn_stats() -> tuple[torch.Tensor, torch.Tensor]: return ( torch.full((PACKED_CHANNELS,), BN_MEAN_VALUE), @@ -72,8 +91,13 @@ def _drive( with_mask: bool = False, scheduler: str = "euler", bn: bool = True, + transformer: torch.nn.Module | None = None, ): - """Run `_run_diffusion` far enough to build the schedule, start latents and inpaint extension.""" + """Run `_run_diffusion` far enough to build the schedule, start latents and inpaint extension. + + With a `transformer`, the run continues through the real denoising loop instead of stopping at the + transformer load. + """ from invokeai.backend.model_manager.taxonomy import BaseModelType, Flux2VariantType, ModelFormat, ModelType from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( ConditioningFieldData, @@ -83,10 +107,21 @@ def _drive( noise = _latents(7) init_latents = _latents(8) mask = torch.full_like(noise, 0.5) - captured: dict = {"shift_calls": [], "start_kwargs": [], "inpaint": [], "redensify_calls": []} + captured: dict = { + "shift_calls": [], + "clip_calls": [], + "start_kwargs": [], + "inpaint": [], + "redensify_calls": [], + "denoise_kwargs": [], + } transformer_info = MagicMock() - transformer_info.model_on_device = MagicMock(side_effect=_StopBeforeLoad) + if transformer is None: + transformer_info.model_on_device = MagicMock(side_effect=_StopBeforeLoad) + else: + transformer_info.model = transformer + transformer_info.model_on_device = MagicMock(return_value=nullcontext((None, transformer))) context = MagicMock() context.models.load.return_value = transformer_info @@ -128,6 +163,12 @@ def recording_shift(timesteps, mu): captured["shift_calls"].append((list(timesteps), mu)) return real_shift(timesteps, mu) + real_clip = flux2_denoise.clip_timestep_schedule_fractional + + def recording_clip(timesteps, denoising_start, denoising_end): + captured["clip_calls"].append(list(timesteps)) + return real_clip(timesteps, denoising_start, denoising_end) + real_redensify = flux2_denoise.redensify_schedule_flux2 def recording_redensify(timesteps, num_steps, mu): @@ -135,6 +176,12 @@ def recording_redensify(timesteps, num_steps, mu): captured["redensify_calls"].append({"clipped": list(timesteps), "num_steps": num_steps, "result": result}) return result + real_denoise = flux2_denoise.denoise + + def recording_denoise(**kwargs): + captured["denoise_kwargs"].append(kwargs) + return real_denoise(**kwargs) + real_prepare = Flux2DenoiseInvocation._prepare_normalized_start_latents def recording_prepare(self, **kwargs): @@ -153,7 +200,9 @@ def __init__(self, *, init_latents, inpaint_mask, noise): patch.object(Flux2DenoiseInvocation, "_prepare_normalized_start_latents", recording_prepare), patch("invokeai.backend.util.devices.TorchDevice.choose_torch_device", return_value=torch.device("cpu")), patch.object(flux2_denoise, "time_shift_flux2", recording_shift), + patch.object(flux2_denoise, "clip_timestep_schedule_fractional", recording_clip), patch.object(flux2_denoise, "redensify_schedule_flux2", recording_redensify), + patch.object(flux2_denoise, "denoise", recording_denoise), patch.object(flux2_denoise, "RectifiedFlowInpaintExtension", _RecordingInpaintExtension), ): try: @@ -168,42 +217,82 @@ def __init__(self, *, init_latents, inpaint_mask, noise): return captured -def _expected_schedule(denoising_start: float, denoising_end: float = 1.0, shifted: bool = True) -> list[float]: +def _expected_schedule(denoising_start: float, denoising_end: float = 1.0) -> list[float]: linear = get_schedule_flux2(num_steps=NUM_STEPS, image_seq_len=IMAGE_SEQ_LEN) mu = compute_empirical_mu(image_seq_len=IMAGE_SEQ_LEN, num_steps=NUM_STEPS) - timesteps = time_shift_flux2(linear, mu) if shifted else linear - clipped = clip_timestep_schedule_fractional(timesteps, denoising_start, denoising_end) - return redensify_schedule_flux2(clipped, NUM_STEPS, mu) if shifted else clipped - - -def test_txt2img_does_not_shift_the_schedule() -> None: - """The txt2img scheduler applies mu itself; shifting here too would shift twice.""" - captured = _drive(denoising_start=0.0, with_init=False) - - assert captured["shift_calls"] == [] - assert captured["reached_transformer"] - - -@pytest.mark.parametrize(("denoising_start", "with_mask"), [(0.5, False), (0.0, True)]) -def test_manual_euler_paths_shift_the_schedule_once(denoising_start: float, with_mask: bool) -> None: - """img2img and inpainting step the schedule themselves, so they must apply the shift.""" - captured = _drive(denoising_start=denoising_start, with_mask=with_mask) + clipped = clip_timestep_schedule_fractional(time_shift_flux2(linear, mu), denoising_start, denoising_end) + return redensify_schedule_flux2(clipped, NUM_STEPS, mu) + + +@pytest.mark.parametrize( + ("denoising_start", "with_init", "with_mask", "scheduler"), + [ + pytest.param(0.0, False, False, "euler", id="txt2img"), + pytest.param(0.0, False, False, "heun", id="txt2img-heun"), + pytest.param(0.5, True, False, "euler", id="img2img"), + pytest.param(0.0, True, True, "euler", id="inpaint"), + # Manual Euler also runs when the requested scheduler is unavailable (e.g. lcm without LCM + # support). + pytest.param(0.0, False, False, "not_a_real_scheduler", id="scheduler-fallback"), + ], +) +def test_every_path_shifts_the_full_schedule_once_before_clipping( + denoising_start: float, with_init: bool, with_mask: bool, scheduler: str +) -> None: + captured = _drive(denoising_start=denoising_start, with_init=with_init, with_mask=with_mask, scheduler=scheduler) assert len(captured["shift_calls"]) == 1 - shifted_timesteps, mu = captured["shift_calls"][0] + shifted_input, mu = captured["shift_calls"][0] assert mu == pytest.approx(compute_empirical_mu(image_seq_len=IMAGE_SEQ_LEN, num_steps=NUM_STEPS)) - # Shifted before clipping, so the call still sees the full linear schedule. - assert shifted_timesteps == pytest.approx(get_schedule_flux2(num_steps=NUM_STEPS, image_seq_len=IMAGE_SEQ_LEN)) - - -def test_a_scheduler_that_falls_back_to_manual_euler_still_shifts() -> None: - """Manual Euler also runs when the requested scheduler is unavailable (e.g. lcm without LCM support). - - Keying the shift on denoise_mask/denoising_start alone left that path on the unshifted schedule. + # Shifted before clipping, so the shift sees the full linear schedule and the clip its output. + assert shifted_input == pytest.approx(get_schedule_flux2(num_steps=NUM_STEPS, image_seq_len=IMAGE_SEQ_LEN)) + assert captured["clip_calls"] == [time_shift_flux2(shifted_input, mu)] + + +@pytest.mark.parametrize( + "scheduler", + [ + "euler", + "heun", + pytest.param( + "lcm", marks=pytest.mark.skipif("lcm" not in FLUX_SCHEDULER_MAP, reason="FlowMatchLCMScheduler unavailable") + ), + ], +) +def test_txt2img_schedulers_run_the_schedule_the_node_shifted(scheduler: str) -> None: + """The node shifts the schedule once. A scheduler that shifted it again, or built its own the way Heun + did with a static shift of 3, would evaluate the model elsewhere.""" + transformer = _RecordingTransformer() + captured = _drive(with_init=False, scheduler=scheduler, transformer=transformer, bn=False) + + expected = _expected_schedule(denoising_start=0.0) + if scheduler == "heun": + reference = FlowMatchHeunDiscreteScheduler(num_train_timesteps=1000) + set_heun_timesteps_from_sigmas(reference, expected[:-1], expected[-1], "cpu") + evaluated = (reference.timesteps / 1000).tolist() + else: + evaluated = expected[:-1] + + (denoise_kwargs,) = captured["denoise_kwargs"] + assert (denoise_kwargs["scheduler"].timesteps / 1000).tolist() == pytest.approx(evaluated, abs=1e-6) + # The transformer runs in bfloat16, so the sigmas it saw are rounded to bfloat16. + assert transformer.sigmas == pytest.approx(evaluated, abs=4e-3) + + +def test_a_pass_ending_at_denoising_end_stops_where_the_continuing_pass_starts() -> None: + """denoising_end and denoising_start meet at one sigma, so two chained passes split one schedule. + + The txt2img pass used to clip the unshifted schedule and then append a terminal sigma of 0: it + finished the image instead of stopping, at a sigma the continuing img2img pass never assumed. """ - captured = _drive(denoising_start=0.0, with_init=False, scheduler="not_a_real_scheduler") - - assert len(captured["shift_calls"]) == 1 + split = 0.6 + first = _drive(with_init=False, denoising_end=split, transformer=_RecordingTransformer(), bn=False) + second = _drive(denoising_start=split) + + (denoise_kwargs,) = first["denoise_kwargs"] + terminal_sigma = float(denoise_kwargs["scheduler"].sigmas[-1]) + assert terminal_sigma == pytest.approx(1.0 - split) + assert terminal_sigma == pytest.approx(second["start_kwargs"][0]["t_0"]) def test_start_latents_use_the_first_sigma_of_the_clipped_schedule() -> None: diff --git a/tests/app/invocations/test_z_image_denoise_schedulers.py b/tests/app/invocations/test_z_image_denoise_schedulers.py new file mode 100644 index 00000000000..3bf005be828 --- /dev/null +++ b/tests/app/invocations/test_z_image_denoise_schedulers.py @@ -0,0 +1,187 @@ +"""Z-Image's Heun scheduler runs the node's shifted, clipped schedule. + +Heun used to rebuild its own schedule from the step count with a shift of 1.0: it ran linear sigmas +instead of the shifted ones, started img2img at sigma 1.0 whatever the img2img blend used, ignored +`denoising_end`, and its progress stopped one step short of the end. + +These tests drive the real `_run_diffusion` through the denoising loop with a stand-in transformer. +The node's own manual Euler path is the reference for the schedule: it steps the clipped sigmas +directly, so Heun has to evaluate the same sigmas in Diffusers' Heun layout. An interval of zero length, +which an explicit shift of 0 leaves at sigma 0, is the exception: Heun would divide by zero there. +""" + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from invokeai.app.invocations.z_image.z_image_denoise import ZImageDenoiseInvocation +from invokeai.app.util.step_callback import calc_percentage +from invokeai.backend.flux.schedulers import ZIMAGE_SCHEDULER_MAP +from invokeai.backend.model_manager.taxonomy import ModelFormat + +SIZE = 128 # a 16x16 latent +STEPS = 8 +VELOCITY = 0.25 + + +class _ConstantVelocityTransformer(torch.nn.Module): + """Predicts a constant velocity and records each evaluation's sigma and sample level. + + With a constant velocity every flow-matching solver moves the sample by exactly + (sigma_end - sigma_start) * velocity, so the result shows where the schedule ended. + """ + + def __init__(self) -> None: + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(1)) + self.sigmas: list[float] = [] + self.sample_levels: list[float] = [] + + def forward(self, x: list[torch.Tensor], t: torch.Tensor, cap_feats: list[torch.Tensor]): + del cap_feats + # The model's t runs from 0 (noise) to 1 (clean); the node negates the prediction. + self.sigmas.append(1.0 - float(t[0])) + self.sample_levels.append(float(torch.stack(x).mean())) + return ([torch.full_like(sample, -VELOCITY) for sample in x],) + + +def _drive(scheduler: str, denoising_start: float = 0.0, denoising_end: float = 1.0, shift: float | None = None): + """Run the node with zero init latents and unit noise, so a blended sample's value is its sigma.""" + transformer = _ConstantVelocityTransformer() + transformer_info = MagicMock(model=transformer) + transformer_info.model_on_device.return_value = nullcontext((None, transformer)) + + context = MagicMock() + context.models.load.return_value = transformer_info + context.models.get_config.return_value = SimpleNamespace(format=ModelFormat.Diffusers) + context.tensors.load.return_value = torch.zeros(1, 16, SIZE // 8, SIZE // 8) + + node = ZImageDenoiseInvocation.model_construct( + id="test", + latents=MagicMock(latents_name="latents") if denoising_start > 0 else None, + add_noise=True, + width=SIZE, + height=SIZE, + steps=STEPS, + denoising_start=denoising_start, + denoising_end=denoising_end, + shift=shift, + guidance_scale=1.0, + positive_conditioning=SimpleNamespace(conditioning_name="positive", mask=None), + negative_conditioning=None, + transformer=MagicMock(transformer="transformer", loras=[]), + control=None, + denoise_mask=None, + seed=0, + scheduler=scheduler, + ) + regional_extension = SimpleNamespace( + regional_text_conditioning=SimpleNamespace(prompt_embeds=torch.zeros(4, 16)), + regional_attn_mask=None, + ) + callbacks = [] + + with ( + patch( + "invokeai.app.invocations.z_image.z_image_denoise.TorchDevice.choose_torch_device", + return_value=torch.device("cpu"), + ), + patch( + "invokeai.app.invocations.z_image.z_image_denoise.TorchDevice.choose_bfloat16_safe_dtype", + return_value=torch.float32, + ), + patch( + "invokeai.app.invocations.z_image.z_image_denoise.ZImageRegionalPromptingExtension.from_text_conditionings", + return_value=regional_extension, + ), + patch.object(node, "_load_text_conditioning", return_value=[]), + patch.object(node, "_prepare_noise_tensor", return_value=torch.ones(1, 16, SIZE // 8, SIZE // 8)), + patch.object(node, "_estimate_working_memory", return_value=0), + patch.object(node, "_build_step_callback", return_value=callbacks.append), + ): + result = node._run_diffusion(context) + + return transformer, result, callbacks + + +def _heun_layout(sigmas: list[float]) -> list[float]: + """Diffusers' Heun evaluates the first sigma once and every later one twice.""" + return [sigmas[0]] + [sigma for sigma in sigmas[1:] for _ in range(2)] + + +@pytest.mark.parametrize( + ("denoising_start", "denoising_end"), + [ + pytest.param(0.0, 1.0, id="txt2img"), + pytest.param(0.5, 1.0, id="img2img"), + pytest.param(0.0, 0.5, id="partial"), + ], +) +def test_heun_evaluates_the_sigmas_euler_steps(denoising_start: float, denoising_end: float): + euler, _, _ = _drive("euler", denoising_start, denoising_end) + heun, _, _ = _drive("heun", denoising_start, denoising_end) + + assert heun.sigmas == pytest.approx(_heun_layout(euler.sigmas), abs=1e-6) + + +def test_heun_img2img_starts_at_the_sigma_the_blend_used(): + heun, _, _ = _drive("heun", denoising_start=0.5) + + # Zero init latents and unit noise: the blended sample's value is the sigma it was blended at. + assert heun.sigmas[0] == pytest.approx(heun.sample_levels[0], abs=1e-6) + assert heun.sigmas[0] < 1.0 + + +@pytest.mark.parametrize("scheduler", ["euler", "heun"]) +def test_a_pass_ending_at_denoising_end_stops_where_the_continuing_pass_starts(scheduler: str): + _, result, _ = _drive(scheduler, denoising_end=0.5) + continuing, _, _ = _drive(scheduler, denoising_start=0.5) + + # From unit noise at sigma 1.0, a constant velocity leaves the sample at 1 + (sigma_end - 1) * v. + end_sigma = 1.0 + (float(result.mean()) - 1.0) / VELOCITY + assert end_sigma == pytest.approx(continuing.sample_levels[0], abs=1e-5) + assert end_sigma > 0.0 + + +@pytest.mark.skipif("lcm" not in ZIMAGE_SCHEDULER_MAP, reason="FlowMatchLCMScheduler unavailable") +def test_lcm_keeps_its_own_schedule(): + """LCM builds its schedule from the step count; only Heun is handed the node's sigmas.""" + lcm, _, callbacks = _drive("lcm") + + assert len(lcm.sigmas) == STEPS + assert lcm.sigmas[0] == pytest.approx(1.0) + assert [callback.step for callback in callbacks] == list(range(1, STEPS + 1)) + + +def test_heun_progress_reaches_every_step(): + _, _, callbacks = _drive("heun") + + assert [callback.step for callback in callbacks] == list(range(1, STEPS + 1)) + assert [calc_percentage(callback) for callback in callbacks] == pytest.approx( + [step / STEPS for step in range(1, STEPS + 1)] + ) + + +def test_heun_at_a_shift_of_zero_takes_the_single_step_to_zero_that_euler_takes(): + # A shift of 0 schedules [1, 0, ..., 0]: one interval, then intervals of zero length at sigma 0, where Heun + # would divide 0 by 0. + heun, result, callbacks = _drive("heun", shift=0.0) + + # One step from unit noise at sigma 1 to sigma 0 at a constant velocity. + assert torch.allclose(result, torch.full_like(result, 1.0 - VELOCITY)) + assert heun.sigmas == [1.0] + assert [callback.step for callback in callbacks] == [1] + assert calc_percentage(callbacks[-1]) == pytest.approx(1.0) + + +def test_heun_img2img_at_a_shift_of_zero_returns_the_blend_when_only_sigma_zero_is_left(): + # Past the first step a shift of 0 leaves only sigma 0: the blend uses it, and no interval remains to denoise. + heun, result, callbacks = _drive("heun", denoising_start=0.5, shift=0.0) + + # Zero init latents blended at sigma 0 stay zero. + assert torch.equal(result, torch.zeros_like(result)) + assert heun.sigmas == [] + assert callbacks == [] diff --git a/tests/backend/flux2/test_denoise_ref_image_conditioning.py b/tests/backend/flux2/test_denoise_ref_image_conditioning.py index b53c5624d72..c91aff620e5 100644 --- a/tests/backend/flux2/test_denoise_ref_image_conditioning.py +++ b/tests/backend/flux2/test_denoise_ref_image_conditioning.py @@ -61,7 +61,7 @@ def _inputs() -> dict[str, Any]: def _scheduler() -> FlowMatchEulerDiscreteScheduler: - return FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0) + return FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) @pytest.mark.parametrize("use_scheduler", [False, True], ids=["euler", "scheduler"]) diff --git a/tests/backend/flux2/test_denoise_schedulers.py b/tests/backend/flux2/test_denoise_schedulers.py new file mode 100644 index 00000000000..bd04314fda0 --- /dev/null +++ b/tests/backend/flux2/test_denoise_schedulers.py @@ -0,0 +1,153 @@ +"""The FLUX.2 Diffusers scheduler path runs the schedule it is handed. + +The node passes the shifted, clipped schedule: the sigmas the model is evaluated at plus the terminal +sigma the last step ends at. Euler used to drop that terminal sigma (Diffusers appends 0, so a pass +with denoising_end < 1 finished the image instead of stopping), and Heun replaced the whole schedule +with its own static shift. +""" + +from typing import Any + +import pytest +import torch +from diffusers import FlowMatchHeunDiscreteScheduler + +from invokeai.app.util.step_callback import calc_percentage +from invokeai.backend.flux.schedulers import _HAS_LCM, FLUX_SCHEDULER_MAP +from invokeai.backend.flux2.denoise import denoise + +VELOCITY = 0.25 + + +class _ConstantVelocityModel: + """Predicts a constant velocity and records the sigma of every evaluation. + + With a constant velocity every first- or second-order flow-matching solver moves the sample by + exactly (sigma_end - sigma_start) * velocity, so the result shows where the schedule ended. + """ + + def __init__(self) -> None: + self.sigmas: list[float] = [] + + def __call__(self, hidden_states: torch.Tensor, timestep: torch.Tensor, **kwargs: Any) -> tuple[torch.Tensor]: + del kwargs + self.sigmas.append(float(timestep[0])) + return (torch.full_like(hidden_states, VELOCITY),) + + +class _SampleDependentModel: + """Predicts from the sample, so each evaluation's clean estimate differs.""" + + def __init__(self) -> None: + self.evaluations: list[tuple[float, torch.Tensor]] = [] + + def __call__(self, hidden_states: torch.Tensor, timestep: torch.Tensor, **kwargs: Any) -> tuple[torch.Tensor]: + del kwargs + sigma = float(timestep[0]) + self.evaluations.append((sigma, hidden_states.clone())) + return (0.1 * hidden_states + 0.2 * sigma,) + + +def _run(scheduler, timesteps: list[float], model=None): + model = model or _ConstantVelocityModel() + img = torch.tensor([[[1.0], [-0.5]]]) + callbacks = [] + result = denoise( + model=model, # type: ignore[arg-type] + img=img, + img_ids=torch.zeros(1, 2, 4, dtype=torch.long), + txt=torch.zeros(1, 3, 1), + txt_ids=torch.zeros(1, 3, 4, dtype=torch.long), + timesteps=timesteps, + step_callback=callbacks.append, + guidance=1.0, + cfg_scale=[1.0] * (len(timesteps) - 1), + scheduler=scheduler, + ) + return result, img, model, callbacks + + +@pytest.mark.parametrize("terminal_sigma", [0.0, 0.4]) +def test_euler_runs_the_given_sigmas_and_stops_at_the_terminal_sigma(terminal_sigma: float): + schedule = [1.0, 0.95, 0.85, 0.6, terminal_sigma] + scheduler = FLUX_SCHEDULER_MAP["euler"](num_train_timesteps=1000) + + result, img, model, callbacks = _run(scheduler, schedule) + + assert model.sigmas == pytest.approx(schedule[:-1]) + assert torch.equal(scheduler.sigmas[-1], torch.tensor(terminal_sigma)) + assert torch.allclose(result, img + (terminal_sigma - schedule[0]) * VELOCITY) + assert [callback.step for callback in callbacks] == [1, 2, 3, 4] + assert all(callback.order == 1 and callback.total_steps == 4 for callback in callbacks) + + +def test_heun_runs_the_given_sigmas_including_the_terminal_sigma(): + schedule = [1.0, 0.95, 0.85, 0.6, 0.4] + scheduler = FLUX_SCHEDULER_MAP["heun"](num_train_timesteps=1000) + + result, img, model, callbacks = _run(scheduler, schedule) + + # Diffusers' Heun layout: the first sigma once, every later one twice, then the terminal sigma. + expected_sigmas = [1.0, 0.95, 0.95, 0.85, 0.85, 0.6, 0.6, 0.4] + assert scheduler.sigmas.tolist() == pytest.approx(expected_sigmas) + assert model.sigmas == pytest.approx(expected_sigmas[:-1]) + assert torch.allclose(result, img + (schedule[-1] - schedule[0]) * VELOCITY, atol=1e-6, rtol=0) + + +@pytest.mark.parametrize("step_count", [4, 5]) +def test_heun_progress_reaches_every_user_step(step_count: int): + schedule = torch.linspace(1.0, 0.0, step_count + 1).tolist() + + _, _, _, callbacks = _run(FlowMatchHeunDiscreteScheduler(num_train_timesteps=1000), schedule) + + assert [callback.step for callback in callbacks] == list(range(1, step_count + 1)) + assert [calc_percentage(callback) for callback in callbacks] == pytest.approx( + [step / step_count for step in range(1, step_count + 1)] + ) + + +@pytest.mark.parametrize("scheduler_name", ["euler", "heun"]) +def test_previews_use_the_sample_and_sigma_of_the_model_evaluation(scheduler_name: str): + scheduler = FLUX_SCHEDULER_MAP[scheduler_name](num_train_timesteps=1000) + + _, img, _, callbacks = _run(scheduler, [1.0, 0.7, 0.35, 0.0]) + + # With a constant velocity the clean estimate x - sigma * v is the same at every sigma. + assert len(callbacks) == 3 + assert all(torch.allclose(callback.latents, img - VELOCITY) for callback in callbacks) + + +def test_heun_previews_come_from_each_intervals_second_order_evaluation(): + schedule = [0.9, 0.6, 0.3, 0.1] + model = _SampleDependentModel() + + _, _, _, callbacks = _run(FLUX_SCHEDULER_MAP["heun"](num_train_timesteps=1000), schedule, model=model) + + # Evaluations: 0.9 (first order), 0.6 (second order), 0.6 (first order), 0.3 (second order), then the + # final interval's first-order step from 0.3, which has no second-order evaluation. + assert [sigma for sigma, _ in model.evaluations] == pytest.approx([0.9, 0.6, 0.6, 0.3, 0.3]) + for callback, evaluation_index in zip(callbacks, [1, 3, 4], strict=True): + sigma, sample = model.evaluations[evaluation_index] + assert torch.allclose(callback.latents, sample - sigma * (0.1 * sample + 0.2 * sigma)) + + +@pytest.mark.skipif(not _HAS_LCM, reason="FlowMatchLCMScheduler not available in this Diffusers version") +@pytest.mark.parametrize("terminal_sigma", [0.0, 0.4]) +def test_lcm_runs_the_given_sigmas_and_renoises_to_the_terminal_sigma(terminal_sigma: float): + schedule = [1.0, 0.95, 0.85, 0.6, terminal_sigma] + scheduler = FLUX_SCHEDULER_MAP["lcm"](num_train_timesteps=1000) + + _, _, model, callbacks = _run(scheduler, schedule) + + assert model.sigmas == pytest.approx(schedule[:-1]) + assert torch.equal(scheduler.sigmas[-1], torch.tensor(terminal_sigma)) + assert [callback.step for callback in callbacks] == [1, 2, 3, 4] + + +def test_an_unrelated_scheduler_is_rejected(): + class _UnrelatedScheduler: + def set_timesteps(self, sigmas: list[float], device: torch.device) -> None: + raise AssertionError("must not be configured") + + with pytest.raises(TypeError, match="Unsupported FLUX.2 scheduler"): + _run(_UnrelatedScheduler(), [1.0, 0.5, 0.0]) diff --git a/tests/backend/flux2/test_flux2_schedule_shift.py b/tests/backend/flux2/test_flux2_schedule_shift.py index 3a7f0c6fd76..e1b7c774d88 100644 --- a/tests/backend/flux2/test_flux2_schedule_shift.py +++ b/tests/backend/flux2/test_flux2_schedule_shift.py @@ -1,10 +1,9 @@ -"""Regression tests for the FLUX.2 schedule shift shared by the txt2img and img2img paths. +"""Regression tests for the FLUX.2 schedule shift every denoising path runs on. -``get_schedule_flux2()`` returns an unshifted linear schedule because the txt2img scheduler applies -the exponential shift from ``mu`` itself. img2img and inpainting step the schedule manually and must -apply the same shift, otherwise they run the model on a sigma trajectory txt2img never visits -- with -9 steps the shifted schedule bottoms out at 0.485 while the linear one runs down to 0.111, and a -distilled model like FLUX.2 Klein leaves a grainy residue down there. +``get_schedule_flux2()`` returns an unshifted linear schedule; the node shifts it with +``time_shift_flux2()`` before clipping. Without the shift the model runs on a sigma trajectory it was +never meant to see -- with 9 steps the shifted schedule bottoms out at 0.485 while the linear one runs +down to 0.111, and a distilled model like FLUX.2 Klein leaves a grainy residue down there. """ import numpy as np @@ -13,7 +12,6 @@ from invokeai.backend.flux.sampling_utils import clip_timestep_schedule_fractional from invokeai.backend.flux2.sampling_utils import ( - FLUX2_TXT2IMG_SCHEDULER_KWARGS, compute_empirical_mu, get_schedule_flux2, redensify_schedule_flux2, @@ -21,24 +19,27 @@ unshift_flux2, ) - -def _txt2img_scheduler() -> FlowMatchEulerDiscreteScheduler: - """The scheduler exactly as Flux2DenoiseInvocation builds it for txt2img. - - Built from the same config the node uses rather than a copy of it, so the two cannot drift apart - without this test noticing. - """ - return FlowMatchEulerDiscreteScheduler(**FLUX2_TXT2IMG_SCHEDULER_KWARGS) +# The scheduler config of diffusers' Flux2KleinPipeline, which shifts the linear schedule from mu itself. +REFERENCE_SCHEDULER_KWARGS = { + "num_train_timesteps": 1000, + "shift": 3.0, + "use_dynamic_shifting": True, + "base_shift": 0.5, + "max_shift": 1.15, + "base_image_seq_len": 256, + "max_image_seq_len": 4096, + "time_shift_type": "exponential", +} @pytest.mark.parametrize("num_steps", [4, 9, 20, 30]) -def test_shift_matches_the_txt2img_scheduler(num_steps: int) -> None: - """The manual shift must reproduce the sigmas the txt2img scheduler produces.""" +def test_shift_matches_the_reference_scheduler(num_steps: int) -> None: + """The node's shift must reproduce the sigmas the reference pipeline's scheduler produces.""" image_seq_len = 64 * 64 timesteps = get_schedule_flux2(num_steps=num_steps, image_seq_len=image_seq_len) mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_steps) - scheduler = _txt2img_scheduler() + scheduler = FlowMatchEulerDiscreteScheduler(**REFERENCE_SCHEDULER_KWARGS) scheduler.set_timesteps(sigmas=timesteps[:-1], mu=mu) shifted = time_shift_flux2(timesteps, mu)