From aa52feb757e8a88165ba19eda49f25bf31d3b241 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 17 Nov 2024 12:43:46 -0500 Subject: [PATCH] dpm flowmatch Signed-off-by: Vladimir Mandic --- .pylintrc | 1 + .ruff.toml | 1 + CHANGELOG.md | 11 +- installer.py | 2 +- modules/api/helpers.py | 3 +- modules/flowmatch/__init__.py | 1 + modules/flowmatch/flowmatch_dpm.py | 897 +++++++++++++++++++++++++++++ modules/processing_args.py | 1 - modules/processing_class.py | 2 +- modules/sd_samplers.py | 33 +- modules/sd_samplers_diffusers.py | 32 +- wiki | 2 +- 12 files changed, 957 insertions(+), 29 deletions(-) create mode 100644 modules/flowmatch/__init__.py create mode 100644 modules/flowmatch/flowmatch_dpm.py diff --git a/.pylintrc b/.pylintrc index 2d8e4869b..534e7fdce 100644 --- a/.pylintrc +++ b/.pylintrc @@ -33,6 +33,7 @@ ignore-paths=/usr/lib/.*$, modules/omnigen, modules/instantir, modules/consistory, + modules/flowmatch, modules/pulid/eva_clip, repositories, extensions-builtin/sd-webui-agent-scheduler, diff --git a/.ruff.toml b/.ruff.toml index 2a29fb089..8852729f1 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -28,6 +28,7 @@ exclude = [ "modules/omnigen", "modules/instantir", "modules/consistory", + "modules/flowmatch", "modules/pulid/eva_clip", "repositories", "extensions-builtin/sd-extension-chainner/nodes", diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c25b49e6..dadb63d81 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,8 @@ # Change Log for SD.Next -## Update for 2024-11-15 +## Update for 2024-11-17 -### Highlights for 2024-11-15 +### Highlights for 2024-11-17 *What's New?* @@ -16,7 +16,7 @@ First, a massive update to docs including new UI top-level **info** tab with acc **Workflow Improvements**: - Native Docker support -- SD3x: ControlNets and all-in-one-safetensors +- SD3x & Flux.1: more ControlNets, all-in-one-safetensors, DPM samplers, etc. - XYZ grid: benchmarking, video creation, etc. - Enhanced prompt parsing - UI improvements @@ -26,7 +26,7 @@ And quite a few more improvements and fixes since the last update - for full det [README](https://github.com/vladmandic/automatic/blob/master/README.md) | [CHANGELOG](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [WiKi](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867) -### Details for 2024-11-15 +### Details for 2024-11-17 - Docs: - new top-level **info** tab with access to [changelog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) and [wiki](https://github.com/vladmandic/automatic/wiki) @@ -74,6 +74,9 @@ And quite a few more improvements and fixes since the last update - for full det - SD3: all-in-one safetensors - *examples*: [large](https://civitai.com/models/882666/sd35-large-google-flan?modelVersionId=1003031), [medium](https://civitai.com/models/900327) - *note*: enable *bnb* on-the-fly quantization for even bigger gains + - FlowMatch samplers: + - Applicable to SD 3.x and Flux.1 models + - Complete family: - Workflow improvements: - Native Docker support with pre-defined [Dockerfile](https://github.com/vladmandic/automatic/blob/dev/Dockerfile) diff --git a/installer.py b/installer.py index ba744b65f..e338a8aa7 100644 --- a/installer.py +++ b/installer.py @@ -504,7 +504,7 @@ def install_cuda(): def install_rocm_zluda(): if args.skip_all or args.skip_requirements: - return + return None from modules import rocm if not rocm.is_installed: log.warning('ROCm: could not find ROCm toolkit installed') diff --git a/modules/api/helpers.py b/modules/api/helpers.py index f13047b9b..d9a87537e 100644 --- a/modules/api/helpers.py +++ b/modules/api/helpers.py @@ -24,8 +24,7 @@ def decode_base64_to_image(encoding, quiet=False): shared.log.warning(f'API cannot decode image: {e}') if not quiet: raise HTTPException(status_code=500, detail="Invalid encoded image") from e - else: - return None + return None def encode_pil_to_base64(image): diff --git a/modules/flowmatch/__init__.py b/modules/flowmatch/__init__.py new file mode 100644 index 000000000..5a577cce8 --- /dev/null +++ b/modules/flowmatch/__init__.py @@ -0,0 +1 @@ +from .flowmatch_dpm import FlowMatchDPMSolverMultistepScheduler, FlowMatchDPMSolverMultistepSchedulerOutput diff --git a/modules/flowmatch/flowmatch_dpm.py b/modules/flowmatch/flowmatch_dpm.py new file mode 100644 index 000000000..83573105e --- /dev/null +++ b/modules/flowmatch/flowmatch_dpm.py @@ -0,0 +1,897 @@ +# Credits: @ukaprch + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +import torchsde + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.utils import BaseOutput, logging +from diffusers.utils.torch_utils import randn_tensor +from diffusers.schedulers.scheduling_utils import SchedulerMixin + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +class BatchedBrownianTree: + """A wrapper around torchsde.BrownianTree that enables batches of entropy.""" + + def __init__(self, x, t0, t1, seed=None, **kwargs): + t0, t1, self.sign = self.sort(t0, t1) + w0 = kwargs.get("w0", torch.zeros_like(x)) + if seed is None: + seed = torch.randint(0, 2**63 - 1, []).item() + self.batched = True + try: + assert len(seed) == x.shape[0] + w0 = w0[0] + except TypeError: + seed = [seed] + self.batched = False + self.trees = [ + torchsde.BrownianInterval( + t0=t0, + t1=t1, + size=w0.shape, + dtype=w0.dtype, + device=w0.device, + entropy=s, + tol=1e-6, + pool_size=24, + halfway_tree=True, + ) + for s in seed + ] + + @staticmethod + def sort(a, b): + return (a, b, 1) if a < b else (b, a, -1) + + def __call__(self, t0, t1): + t0, t1, sign = self.sort(t0, t1) + w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign) + return w if self.batched else w[0] + +class BrownianTreeNoiseSampler: + """A noise sampler backed by a torchsde.BrownianTree. + + Args: + x (Tensor): The tensor whose shape, device and dtype to use to generate + random samples. + sigma_min (float): The low end of the valid interval. + sigma_max (float): The high end of the valid interval. + seed (int or List[int]): The random seed. If a list of seeds is + supplied instead of a single integer, then the noise sampler will use one BrownianTree per batch item, each + with its own seed. + transform (callable): A function that maps sigma to the sampler's + internal timestep. + """ + + def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x): + self.transform = transform + t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max)) + self.tree = BatchedBrownianTree(x, t0, t1, seed) + + def __call__(self, sigma, sigma_next): + t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next)) + return self.tree(t0, t1) / (t1 - t0).abs().sqrt() + +@dataclass +class FlowMatchDPMSolverMultistepSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + +class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): + """ + `DPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + solver_order (`int`, defaults to 2): + The DPMSolver order which can be `2` or `3`. It is recommended to use `solver_order=2` for guided + sampling, and `solver_order=3` for unconditional sampling. + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + as Stable Diffusion. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True`. + algorithm_type (`str`, defaults to `dpmsolver++2M`): + Algorithm type for the solver; can be `dpmsolver2`, `dpmsolver2A`, `dpmsolver++2M`, `dpmsolver++2S`, `dpmsolver++sde`, `dpmsolver++2Msde`, + or `dpmsolver++3Msde`. + solver_type (`str`, defaults to `midpoint`): + Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the + sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers. + sigma_schedule (`str`, *optional*, defaults to None): Sigma schedule to compute the `sigmas`. Optionally, we use + the schedule "karras" introduced in the EDM paper (https://arxiv.org/abs/2206.00364). Other acceptable values are + "exponential". The exponential schedule was incorporated in this model: https://huggingface.co/stabilityai/cosxl. + Other acceptable values are "lambdas". The uniform-logSNR for step sizes proposed by Lu's DPM-Solver in the + noise schedule during the sampling process. The sigmas and time steps are determined according to a sequence of `lambda(t)`. + use_noise_sampler for BrownianTreeNoiseSampler (only valid for `dpmsolver++2S`, `dpmsolver++sde`, `dpmsolver++2Msde`, + or `dpmsolver++3Msde`): A noise sampler backed by a torchsde increasing the stability of convergence. Default strategy + (random noise) has it jumping all over the place, but Brownian sampling is more stable. Utilizes the model generation seed provided. + midpoint_ratio (`float`, *optional*, range: 0.4 to 0.6, default=0.5): Only valid for (`dpmsolver++sde`, `dpmsolver++2S`). + Higher values may result in smoothing, more vivid colors and less noise at the expense of more detail and effect. + s_noise (`float`, *optional*, defaults to 1.0): Sigma noise strength: range 0 - 1.1 (only valid for `dpmsolver++2S`, `dpmsolver++sde`, + `dpmsolver++2Msde`, or `dpmsolver++3Msde`). The amount of additional noise to counteract loss of detail during sampling. A + reasonable range is [1.000, 1.011]. Defaults to 1.0 from the original implementation. + use_SD35_sigmas: (`bool` defaults to False for FLUX and True for SD3). Based on original interpretation of using beta values for determining sigmas. + use_dynamic_shifting (`bool` defaults to False for SD3 and True for FLUX). When `True`, shift is ignored. + shift (`float`, defaults to 3.0): The shift value for the timestep schedule for SD3 when not using dynamic shifting + The remaining args are specific to Flux's dynamic shifting based on resolution + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + thresholding: Optional[bool] = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: Optional[float] = 1.0, + algorithm_type: str = "dpmsolver++2M", + solver_type: str = "midpoint", + sigma_schedule: Optional[str] = None, + shift: float = 3.0, + midpoint_ratio: Optional[float] = 0.5, + s_noise: Optional[float] = 1.0, + use_noise_sampler: Optional[bool] = True, + use_SD35_sigmas: Optional[bool] = False, + use_dynamic_shifting=False, + base_shift: Optional[float] = 0.5, + max_shift: Optional[float] = 1.15, + base_image_seq_len: Optional[int] = 256, + max_image_seq_len: Optional[int] = 4096, + ): + # settings for DPM-Solver + if algorithm_type not in ["dpmsolver2", "dpmsolver2A", "dpmsolver++2M", "dpmsolver++2S", "dpmsolver++sde", "dpmsolver++2Msde", "dpmsolver++3Msde"]: + raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}") + + if solver_type not in ["midpoint", "heun"]: + raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") + + # setable values + timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) + + sigmas = timesteps / num_train_timesteps + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.timesteps = sigmas * num_train_timesteps + self.h_last = None + self.h_1 = None + self.h_2 = None + self.noise_sampler = None + self._step_index = None + self._begin_index = None + self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.model_outputs = [None] * solver_order + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def set_timesteps(self, + num_inference_steps: int = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[float] = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + if self.config.use_dynamic_shifting and mu is None: + raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`") + + if sigmas is None: + self.use_SD35_sigmas = True + self.num_inference_steps = num_inference_steps + sigmas1 = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps, dtype=np.float64) + beta_start = 0.00085 + beta_end = 0.012 + betas = torch.linspace(beta_start**0.5, beta_end**0.5, self.config.num_train_timesteps, dtype=torch.float64) ** 2 + alphas = 1.0 - betas + alphas_cumprod = torch.cumprod(alphas, dim=0) + sigmas = np.array(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5) + del alphas_cumprod + del alphas + del betas + elif self.use_SD35_sigmas: + num_inference_steps = len(sigmas) + self.num_inference_steps = num_inference_steps + sigmas1 = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps, dtype=np.float64) + beta_start = 0.00085 + beta_end = 0.012 + betas = torch.linspace(beta_start**0.5, beta_end**0.5, self.config.num_train_timesteps, dtype=torch.float64) ** 2 + alphas = 1.0 - betas + alphas_cumprod = torch.cumprod(alphas, dim=0) + sigmas = np.array(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5) + del alphas_cumprod + del alphas + del betas + else: + num_inference_steps = len(sigmas) + self.num_inference_steps = num_inference_steps + + if self.config.sigma_schedule == "exponential": + if self.use_SD35_sigmas: + sigmas = np.flip(sigmas).copy() + sigma_min = sigmas[-1] + sigma_max = sigmas[0] + sigmas = self._convert_to_exponential(sigma_min, sigma_max, num_inference_steps=num_inference_steps) + OldRange = sigma_max - sigma_min + NewRange = 1.0 - sigma_min + sigmas = (((sigmas - sigma_min) * NewRange) / OldRange) + sigma_min + del sigmas1 + else: + sigma_min = sigmas[-1] + sigma_max = sigmas[0] + sigmas = self._convert_to_exponential(sigma_min, sigma_max, num_inference_steps=num_inference_steps) + elif self.config.sigma_schedule == "karras": + if self.use_SD35_sigmas: + sigmas = np.flip(sigmas).copy() + sigma_min = sigmas[-1] + sigma_max = sigmas[0] + sigmas = self._convert_to_karras(sigma_min, sigma_max, num_inference_steps=num_inference_steps) + OldRange = sigma_max - sigma_min + NewRange = 1.0 - sigma_min + sigmas = (((sigmas - sigma_min) * NewRange) / OldRange) + sigma_min + del sigmas1 + else: + sigma_min = sigmas[-1] + sigma_max = sigmas[0] + sigmas = self._convert_to_karras(sigma_min, sigma_max, num_inference_steps=num_inference_steps) + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device) + elif self.config.sigma_schedule == "lambdas": + if self.use_SD35_sigmas: + log_sigmas = np.log(sigmas) + lambdas = np.flip(log_sigmas.copy()) + lambdas = self._convert_to_lu(in_lambdas=lambdas, num_inference_steps=num_inference_steps) + sigmas = np.exp(lambdas) + sigma_min = sigmas[-1] + sigma_max = sigmas[0] + OldRange = sigma_max - sigma_min + NewRange = 1.0 - sigma_min + sigmas = (((sigmas - sigma_min) * NewRange) / OldRange) + sigma_min + del sigmas1 + del lambdas + del log_sigmas + else: + log_sigmas = np.log(sigmas) + lambdas = log_sigmas.copy() + lambdas = self._convert_to_lu(in_lambdas=lambdas, num_inference_steps=num_inference_steps) + sigmas = np.exp(lambdas) + del lambdas + del log_sigmas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device) + else: + if self.use_SD35_sigmas: + sigmas = np.flip(sigmas).copy() + sigma_min = sigmas[-1] + sigmas = np.linspace(1.0, sigma_min, num_inference_steps) + del sigmas1 + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device) + + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) + + timesteps = sigmas * self.config.num_train_timesteps + self.timesteps = timesteps.to(device=device) + self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + self.h_last = None + self.h_1 = None + self.h_2 = None + self.noise_sampler = None + self.model_outputs = [None] * self.config.solver_order + self._step_index = None + self._begin_index = None + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + def _convert_to_lu(self, in_lambdas: torch.Tensor, num_inference_steps) -> torch.Tensor: + """Constructs the noise schedule of Lu et al. (2022).""" + + lambda_min: float = in_lambdas[-1].item() + lambda_max: float = in_lambdas[0].item() + + rho = 1.0 # 1.0 is the value used in the paper + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = lambda_min ** (1 / rho) + max_inv_rho = lambda_max ** (1 / rho) + lambdas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + return lambdas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras + def _convert_to_karras(self, sigma_min, sigma_max, num_inference_steps) -> torch.Tensor: + rho = 7.0 # 7.0 is the value used in the paper + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = sigma_min ** (1 / rho) + max_inv_rho = sigma_max ** (1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + return sigmas + + def _convert_to_exponential(self, sigma_min, sigma_max, num_inference_steps) -> torch.Tensor: + sigmas = torch.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps).exp() + return sigmas + + def convert_model_output( + self, + model_output: torch.Tensor, + sample: torch.Tensor = None, + *args, + **kwargs, + ) -> torch.Tensor: + """ + Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is + designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an + integral of the data prediction model. + + + + The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise + prediction and data prediction models. + + + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + + Returns: + `torch.Tensor`: + The converted model output. + """ + timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) + if sample is None: + if len(args) > 1: + sample = args[1] + else: + raise ValueError("missing `sample` as a required keyward argument") + + # Flow Match needs to solve an integral of the data prediction model. + sigma = self.sigmas[self.step_index] + x0_pred = sample - sigma * model_output + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + generator: Optional[torch.Generator] = None, + variance_noise: Optional[torch.FloatTensor] = None, + return_dict: bool = True, + ) -> Union[FlowMatchDPMSolverMultistepSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with + the multistep DPMSolver. + + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + variance_noise (`torch.Tensor`): + Alternative to generating noise with `generator` by directly providing the noise for the variance + itself. Useful for methods such as [`LEdits++`]. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if self.step_index is None: + self._init_step_index(timestep) + + if self.config.algorithm_type in ["dpmsolver2", "dpmsolver2A"]: + pass + else: + model_output = self.convert_model_output(model_output, sample=sample) + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.model_outputs[-1] = model_output + + # Upcast to avoid precision issues when computing prev_sample + if sample.dtype != model_output.dtype: + sample = sample.to(model_output.dtype) + + if self.config.algorithm_type in ["dpmsolver2A", "dpmsolver++2S", "dpmsolver++sde", "dpmsolver++2Msde", "dpmsolver++3Msde"] and variance_noise is None: + # Create a noise sampler if it hasn't been created yet + if self.config.use_noise_sampler: + if self.noise_sampler is None: + min_sigma, max_sigma = self.sigmas[self.sigmas > 0].min(), self.sigmas.max() + self.noise_sampler = BrownianTreeNoiseSampler(sample, min_sigma, max_sigma, generator) + else: + noise = randn_tensor(model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype) + elif self.config.algorithm_type in ["dpmsolver2A", "dpmsolver++2S", "dpmsolver++sde", "dpmsolver++2Msde", "dpmsolver++3Msde"]: + noise = variance_noise.to(device=model_output.device, dtype=model_output.dtype) + else: + noise = None + + def sigma_fn(_t: torch.Tensor) -> torch.Tensor: + return _t.neg().exp() + def t_fn(_sigma: torch.Tensor) -> torch.Tensor: + return _sigma.log().neg() + sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] + sigma_prev = self.sigmas[self.step_index - 1] + if self.config.algorithm_type == "dpmsolver2": + if self.config.solver_order == 2: + if sigma_next == 0: + # Euler method + model_output = sample - sigma * model_output + d = (sample - model_output) / sigma + dt = sigma_next - sigma + sample = sample + d * dt + else: + # DPM-Solver2 + sigma_mid = sigma.log().lerp(sigma_next.log(), 0.5).exp() + + #using epsilon for new model output: + pred_original_sample = sample - sigma * model_output + # 2. Convert to an ODE derivative for 1st order + d = (sample - pred_original_sample) / sigma + # 3. delta timestep + dt = sigma_mid - sigma + x_2 = sample + d * dt + + #using epsilon for new model output: + denoised_2 = x_2 - sigma_mid * model_output + # 2. Convert to an ODE derivative for 2nd order + d = (x_2 - denoised_2) / sigma_mid + + # 3. delta timestep + dt = sigma_next - sigma + sample = sample + d * dt + + del pred_original_sample + del denoised_2 + del x_2 + del d + elif self.config.algorithm_type == "dpmsolver2A": + if self.config.solver_order == 2: + # get ancestral step + sigma_from = sigma + sigma_to = sigma_next + su = min(sigma_to, (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5) + sd = (sigma_to**2 - su**2) ** 0.5 + if sd == 0: + # Euler method + model_output = sample - sigma * model_output + d = (sample - model_output) / sigma + dt = sd - sigma + sample = sample + d * dt + else: + # DPM-Solver2A + sigma_mid = sigma.log().lerp(sd.log(), 0.5).exp() + + #using epsilon for new model output: + model_output = sample - sigma * model_output + # 2. Convert to an ODE derivative for 1st order + d = (sample - model_output) / sigma + dt = sd - sigma + sample = sample + d * dt + + #using epsilon for new model output: + pred_original_sample = sample - sigma * model_output + # 2. Convert to an ODE derivative for 1st order + d = (sample - pred_original_sample) / sigma + # 3. delta timestep + dt_1 = sigma_mid - sigma + x_2 = sample + d * dt_1 + + #using epsilon for new model output: + denoised_2 = x_2 - sigma_mid * model_output + # 2. Convert to an ODE derivative for 2nd order + d_2 = (x_2 - denoised_2) / sigma_mid + + # 3. delta timestep + dt_2 = sd - sigma_mid + sample = sample + d_2 * dt_2 + + if self.config.use_noise_sampler: + sample = sample + self.noise_sampler(sigma, sigma_next) * self.config.s_noise * su + else: + sample = sample + noise * self.config.s_noise * su + + del pred_original_sample + del denoised_2 + del x_2 + del d + elif self.config.algorithm_type == "dpmsolver++2M": + if self.config.solver_order == 2: + t, t_next = t_fn(sigma), t_fn(sigma_next) + h = t_next - t + if self.model_outputs[-2] is None or sigma_next == 0: + sample = (sigma_fn(t_next) / sigma_fn(t)) * sample - (-h).expm1() * model_output + else: + # DPM-Solver++(2M) + h_last = t - t_fn(sigma_prev) + r = h_last / h + denoised_d = (1 + 1 / (2 * r)) * model_output - (1 / (2 * r)) * self.model_outputs[-2] + sample = (sigma_fn(t_next) / sigma_fn(t)) * sample - (-h).expm1() * denoised_d + del denoised_d + elif self.config.algorithm_type == "dpmsolver++2S": + if self.config.solver_order == 2: + # get ancestral step + sigma_from = sigma + sigma_to = sigma_next + su = min(sigma_to, (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5) + sd = (sigma_to**2 - su**2) ** 0.5 + if sd == 0: + # Euler method + d = (sample - model_output) / sigma + dt = sd - sigma + sample = sample + d * dt + else: + # DPM-Solver++(2S) + t, t_next = t_fn(sigma), t_fn(sd) + r = self.config.midpoint_ratio + h = t_next - t + s = t + r * h + + # Euler method + d = (sample - model_output) / sigma + dt = sd - sigma + sample = sample + d * dt + + x_2 = (sigma_fn(s) / sigma_fn(t)) * sample - (-h * r).expm1() * model_output + + #using epsilon for new model output: + denoised_2 = x_2 - sigma_fn(s) * model_output + # 2. Convert to an ODE derivative for 2nd order + d = (x_2 - denoised_2) / sigma_fn(s) + dt = sd - sigma_next + sample = sample + d * dt + + del x_2 + del denoised_2 + del d + # Noise addition + if sigma_next > 0: + if self.config.use_noise_sampler: + sample = sample + self.noise_sampler(sigma, sigma_next) * self.config.s_noise * su + else: + sample = sample + noise * self.config.s_noise * su + elif self.config.algorithm_type == "dpmsolver++sde": + if self.config.solver_order == 2: + if sigma_next == 0: + # Euler method + d = (sample - model_output) / sigma + dt = sigma_next - sigma + sample = sample + d * dt + else: + # DPM-Solver++(SDE) + t, t_next = t_fn(sigma), t_fn(sigma_next) + r = self.config.midpoint_ratio + h = t_next - t + s = t + r * h + + # Euler method + d = (sample - model_output) / sigma + dt = sigma_next - sigma + sample = sample + d * dt + + # Step 1 + # get ancestral step + sigma_from = sigma_fn(t) + sigma_to = sigma_fn(s) + su = min(sigma_to, (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5) + sd = (sigma_to**2 - su**2) ** 0.5 + + # Euler method + d = (sample - model_output) / sigma + dt = sd - sigma + sample = sample + d * dt + + s_ = t_fn(sd) + x_2 = (sigma_fn(s_) / sigma_fn(t)) * sample - (t - s_).expm1() * model_output + if self.config.use_noise_sampler: + x_2 = x_2 + self.noise_sampler(sigma_fn(t), sigma_fn(s)) * self.config.s_noise * su + else: + x_2 = x_2 + noise * self.config.s_noise * su + + # Step 2 + # get ancestral step + sigma_from = sigma_fn(t) + sigma_to = sigma_fn(t_next) + su = min(sigma_to, (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5) + sd = (sigma_to**2 - su**2) ** 0.5 + + #using epsilon for new model output: + denoised_2 = x_2 - sigma_fn(s) * model_output + # 2. Convert to an ODE derivative for 2nd order + d = (x_2 - denoised_2) / sigma_fn(s) + dt = sd - sigma_next + sample = sample + d * dt + + if self.config.use_noise_sampler: + sample = sample + self.noise_sampler(sigma_fn(t), sigma_fn(t_next)) * self.config.s_noise * su + else: + sample = sample + noise * self.config.s_noise * su + del x_2 + del denoised_2 + del d + elif self.config.algorithm_type == "dpmsolver++2Msde": + if self.config.solver_order == 2: + if sigma_next == 0: + sample = model_output + else: + # DPM-Solver++(2M) SDE + t, s = -sigma.log(), -sigma_next.log() + h = s - t + eta_h = h * 1 + + # 3. Delta timestep + dt = sigma_next - sigma + sample = sample + model_output * dt + + sample = sigma_next / sigma * (-eta_h).exp() * sample + (-h - eta_h).expm1().neg() * model_output + + if self.model_outputs[-2] is not None: + r = self.h_last / h + if self.solver_type == 'heun': + sample = sample + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (model_output - self.model_outputs[-2]) + elif self.solver_type == 'midpoint': + sample = sample + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (model_output - self.model_outputs[-2]) + + if self.config.use_noise_sampler: + sample = sample + self.noise_sampler(sigma, sigma_next) * sigma_next * (-2 * eta_h).expm1().neg().sqrt() * self.config.s_noise + else: + sample = sample + noise * sigma_next * (-2 * eta_h).expm1().neg().sqrt() * self.config.s_noise + + self.h_last = h + elif self.config.algorithm_type == "dpmsolver++3Msde": + if self.config.solver_order == 3: + if sigma_next == 0: + sample = model_output + else: + # DPM-Solver++(3M) SDE + t, s = -sigma.log(), -sigma_next.log() + h = s - t + h_eta = h * 2 + + # 3. Delta timestep + dt = sigma_next - sigma + sample = sample + model_output * dt + + sample = torch.exp(-h_eta) * sample + (-h_eta).expm1().neg() * model_output + + if self.h_2 is not None: + r0 = self.h_1 / h + r1 = self.h_2 / h + d1_0 = (model_output - self.model_outputs[-2]) / r0 + d1_1 = (self.model_outputs[-2] - self.model_outputs[-3]) / r1 + d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1) + d2 = (d1_0 - d1_1) / (r0 + r1) + phi_2 = h_eta.neg().expm1() / h_eta + 1 + phi_3 = phi_2 / h_eta - 0.5 + sample = sample + phi_2 * d1 - phi_3 * d2 + del d1_0 + del d1_1 + del d1 + del d2 + del phi_2 + del phi_3 + elif self.h_1 is not None: + r = self.h_1 / h + d = (model_output - self.model_outputs[-2]) / r + phi_2 = h_eta.neg().expm1() / h_eta + 1 + sample = sample + phi_2 * d + del d + del phi_2 + + if self.config.use_noise_sampler: + sample = sample + self.noise_sampler(sigma, sigma_next) * sigma_next * (-2 * h).expm1().neg().sqrt() * self.config.s_noise + else: + sample = sample + noise * sigma_next * (-2 * h).expm1().neg().sqrt() * self.config.s_noise + + self.h_2 = self.h_1 + self.h_1 = h + if not self.config.use_noise_sampler and noise is not None: + del noise + prev_sample = sample + + # Cast sample back to expected dtype + prev_sample = prev_sample.to(model_output.dtype) + + # upon completion increase step index by one + self._step_index += 1 + + torch.cuda.empty_cache() + + if not return_dict: + return (prev_sample,) + + return FlowMatchDPMSolverMultistepSchedulerOutput(prev_sample=prev_sample) + + def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.Tensor`): + The input sample. + + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + noise: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + """ + Forward process in flow-matching + + Args: + sample (`torch.FloatTensor`): + The input sample. + timestep (`int`, *optional*): + The current timestep in the diffusion chain. + + Returns: + `torch.FloatTensor`: + A scaled input sample. + """ + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) + + if sample.device.type == "mps" and torch.is_floating_point(timestep): + # mps does not support float64 + schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) + timestep = timestep.to(sample.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(sample.device) + timestep = timestep.to(sample.device) + + # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timestep.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timestep.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(sample.shape): + sigma = sigma.unsqueeze(-1) + + sample = sigma * noise + (1.0 - sigma) * sample + + return sample + + def __len__(self): + return self.config.num_train_timesteps diff --git a/modules/processing_args.py b/modules/processing_args.py index 6c9a98369..ff766ec04 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -6,7 +6,6 @@ import time import inspect import torch import numpy as np -from PIL import Image from modules import shared, errors, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, prompt_parser_diffusers, timer from modules.processing_callbacks import diffusers_callback_legacy, diffusers_callback, set_callbacks_p from modules.processing_helpers import resize_hires, fix_prompts, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, get_generator, set_latents, apply_circular # pylint: disable=unused-import diff --git a/modules/processing_class.py b/modules/processing_class.py index b0b9e1f1b..79f51576f 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -232,7 +232,7 @@ class StableDiffusionProcessing: # a1111 compatibility items shared.opts.data['clip_skip'] = int(self.clip_skip) # for compatibility with a1111 sd_hijack_clip - self.seed_enable_extras: bool = True, + self.seed_enable_extras: bool = True self.is_using_inpainting_conditioning = False # a111 compatibility self.batch_index = 0 self.refiner_switch_at = 0 diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index bbc2f360b..e560744dd 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -47,6 +47,10 @@ def visible_sampler_names(): def create_sampler(name, model): + try: + current = model.scheduler.__class__.__name__ + except Exception: + current = None if name == 'Default' and hasattr(model, 'scheduler'): if getattr(model, "default_scheduler", None) is not None: model.scheduler = copy.deepcopy(model.default_scheduler) @@ -54,12 +58,13 @@ def create_sampler(name, model): model.prior_pipe.scheduler = copy.deepcopy(model.default_scheduler) model.prior_pipe.scheduler.config.clip_sample = False config = {k: v for k, v in model.scheduler.config.items() if not k.startswith('_')} - shared.log.debug(f'Sampler: sampler=default class={model.scheduler.__class__.__name__}: {config}') + shared.log.debug(f'Sampler: sampler=default class={current}: {config}') return model.scheduler config = find_sampler_config(name) if config is None or config.constructor is None: # shared.log.warning(f'Sampler: sampler="{name}" not found') return None + sampler = None if not shared.native: sampler = config.constructor(model) sampler.config = config @@ -68,24 +73,18 @@ def create_sampler(name, model): shared.log.debug(f'Sampler: sampler="{name}" config={config.options}') return sampler elif shared.native: - sampler = config.constructor(model) - if 'Flux' in model.__class__.__name__: - if 'base_image_seq_len' not in sampler.sampler.config or 'max_image_seq_len' not in sampler.sampler.config or 'base_shift' not in sampler.sampler.config or 'max_shift' not in sampler.sampler.config: - shared.log.warning(f'FLUX: sampler="{name}" unsupported') - # sampler.sampler.register_to_config(base_image_seq_len=256, max_image_seq_len=4096, base_shift=0.5, max_shift=1.15) - return None - if 'Lumina' in model.__class__.__name__: - shared.log.warning(f'AlphaVLLM-Lumina: sampler="{name}" unsupported') - return None - if 'StableDiffusion3' in model.__class__.__name__: - if sampler.name != 'Heun FlowMatch': - return None - return None - if 'AuraFlow' in model.__class__.__name__: - shared.log.warning(f'AuraFlow: sampler="{name}" unsupported') - return None + FlowModels = ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow'] if 'KDiffusion' in model.__class__.__name__: return None + if any(x in model.__class__.__name__ for x in FlowModels) and 'FlowMatch' not in name: + shared.log.warning(f'Sampler: default={current} target="{name}" class={model.__class__.__name__} linear scheduler unsupported') + return None + if not any(x in model.__class__.__name__ for x in FlowModels) and 'FlowMatch' in name: + shared.log.warning(f'Sampler: default={current} target="{name}" class={model.__class__.__name__} flow-match scheduler unsupported') + return None + sampler = config.constructor(model) + if sampler is None: + sampler = config.constructor(model) if not hasattr(model, 'scheduler_config'): model.scheduler_config = sampler.sampler.config.copy() if hasattr(sampler.sampler, 'config') else {} model.scheduler = sampler.sampler diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index c7eedf59b..22ec2be46 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -43,6 +43,9 @@ try: KDPM2DiscreteScheduler, KDPM2AncestralDiscreteScheduler, ) + + from modules.flowmatch import FlowMatchDPMSolverMultistepScheduler # pylint: disable=ungrouped-imports + except Exception as e: import diffusers shared.log.error(f'Diffusers import error: version={diffusers.__version__} error: {e}') @@ -70,7 +73,15 @@ config = { 'DPM++ 2M SDE': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "sde-dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 2 }, 'DPM++ 2M EDM': { 'solver_order': 2, 'solver_type': 'midpoint', 'final_sigmas_type': 'zero', 'algorithm_type': 'dpmsolver++' }, 'DPM++ Cosine': { 'solver_order': 2, 'sigma_schedule': "exponential", 'prediction_type': "v-prediction" }, - 'DPM SDE': { 'use_karras_sigmas': False, 'noise_sampler_seed': None, 'timestep_spacing': 'linspace', 'steps_offset': 0 }, + 'DPM SDE': { 'use_karras_sigmas': False, 'noise_sampler_seed': None, 'timestep_spacing': 'linspace', 'steps_offset': 0, }, + + 'DPM2 FlowMatch': { 'shift': 1, 'use_dynamic_shifting': False, 'solver_order': 2, 'sigma_schedule': None, 'use_SD35_sigmas': False, 'algorithm_type': 'dpmsolver2', 'use_noise_sampler': True }, + 'DPM2a FlowMatch': { 'shift': 1, 'use_dynamic_shifting': False, 'solver_order': 2, 'sigma_schedule': None, 'use_SD35_sigmas': False, 'algorithm_type': 'dpmsolver2A', 'use_noise_sampler': True }, + 'DPM2++ 2M FlowMatch': { 'shift': 1, 'use_dynamic_shifting': False, 'solver_order': 2, 'sigma_schedule': None, 'use_SD35_sigmas': False, 'algorithm_type': 'dpmsolver++2M', 'use_noise_sampler': True }, + 'DPM2++ 2S FlowMatch': { 'shift': 1, 'use_dynamic_shifting': False, 'solver_order': 2, 'sigma_schedule': None, 'use_SD35_sigmas': False, 'algorithm_type': 'dpmsolver++2S', 'use_noise_sampler': True }, + 'DPM2++ SDE FlowMatch': { 'shift': 1, 'use_dynamic_shifting': False, 'solver_order': 2, 'sigma_schedule': None, 'use_SD35_sigmas': False, 'algorithm_type': 'dpmsolver++sde', 'use_noise_sampler': True }, + 'DPM2++ 2M SDE FlowMatch': { 'shift': 1, 'use_dynamic_shifting': False, 'solver_order': 2, 'sigma_schedule': None, 'use_SD35_sigmas': False, 'algorithm_type': 'dpmsolver++2Msde', 'use_noise_sampler': True }, + 'DPM2++ 3M SDE FlowMatch': { 'shift': 1, 'use_dynamic_shifting': False, 'solver_order': 3, 'sigma_schedule': None, 'use_SD35_sigmas': False, 'algorithm_type': 'dpmsolver++3Msde', 'use_noise_sampler': True }, 'Heun': { 'use_beta_sigmas': False, 'use_karras_sigmas': False, 'timestep_spacing': 'linspace' }, 'Heun FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1 }, @@ -112,6 +123,14 @@ samplers_data_diffusers = [ sd_samplers_common.SamplerData('DPM++ Cosine', lambda model: DiffusionSampler('DPM++ 2M EDM', CosineDPMSolverMultistepScheduler, model), [], {}), sd_samplers_common.SamplerData('DPM SDE', lambda model: DiffusionSampler('DPM SDE', DPMSolverSDEScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM2 FlowMatch', lambda model: DiffusionSampler('DPM2 FlowMatch', FlowMatchDPMSolverMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM2a FlowMatch', lambda model: DiffusionSampler('DPM2a FlowMatch', FlowMatchDPMSolverMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM2++ 2M FlowMatch', lambda model: DiffusionSampler('DPM2++ 2M FlowMatch', FlowMatchDPMSolverMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM2++ 2S FlowMatch', lambda model: DiffusionSampler('DPM2++ 2S FlowMatch', FlowMatchDPMSolverMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM2++ SDE FlowMatch', lambda model: DiffusionSampler('DPM2++ SDE FlowMatch', FlowMatchDPMSolverMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM2++ 2M SDE FlowMatch', lambda model: DiffusionSampler('DPM2++ 2M SDE FlowMatch', FlowMatchDPMSolverMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('DPM2++ 3M SDE FlowMatch', lambda model: DiffusionSampler('DPM2++ 3M SDE FlowMatch', FlowMatchDPMSolverMultistepScheduler, model), [], {}), + sd_samplers_common.SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}), sd_samplers_common.SamplerData('Heun FlowMatch', lambda model: DiffusionSampler('Heun FlowMatch', FlowMatchHeunDiscreteScheduler, model), [], {}), @@ -183,6 +202,12 @@ class DiffusionSampler: self.config['use_karras_sigmas'] = shared.opts.schedulers_sigma == 'karras' if 'use_exponential_sigmas' in self.config: self.config['use_exponential_sigmas'] = shared.opts.schedulers_sigma == 'exponential' + if 'sigma_schedule' in self.config and shared.opts.schedulers_sigma == 'beta': + self.config['sigma_schedule'] = 'beta' + if 'sigma_schedule' in self.config and shared.opts.schedulers_sigma == 'karras': + self.config['sigma_schedule'] = 'karras' + if 'sigma_schedule' in self.config and shared.opts.schedulers_sigma == 'exponential': + self.config['sigma_schedule'] = 'exponential' else: pass # timesteps are set using set_timesteps in set_pipeline_args @@ -201,7 +226,10 @@ class DiffusionSampler: if 'shift' in self.config: self.config['shift'] = shared.opts.schedulers_shift if 'use_dynamic_shifting' in self.config: - self.config['use_dynamic_shifting'] = shared.opts.schedulers_dynamic_shift + if 'Flux' in model.__class__.__name__: + self.config['use_dynamic_shifting'] = shared.opts.schedulers_dynamic_shift + if 'use_SD35_sigmas' in self.config: + self.config['use_SD35_sigmas'] = 'StableDiffusion3' in model.__class__.__name__ if 'rescale_betas_zero_snr' in self.config: self.config['rescale_betas_zero_snr'] = shared.opts.schedulers_rescale_betas if 'timestep_spacing' in self.config and shared.opts.schedulers_timestep_spacing != 'default' and shared.opts.schedulers_timestep_spacing is not None: diff --git a/wiki b/wiki index 82ca84bc9..713906e92 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 82ca84bc91dd1c915393a96207cd3e5d848a388f +Subproject commit 713906e920e02607ea04951858aabeff7ce641f2