diff --git a/CHANGELOG.md b/CHANGELOG.md index 47e1e6f7e..9c6352a1e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -202,11 +202,17 @@ For full details, see [ChangeLog](https://github.com/vladmandic/automatic/blob/m - **Ernie-Image** add native *LoRA* support, *img2img* and *inpaint* workflows - **Chroma** add native *LoRA* support - **Flux.2** add native *LoRA* support - - **Prompt enhance** add info to image metadata - custom **VAE** loader for all pipelines *note*: vae still needs to be compatible with the model + - **Schedulers** new option in ui: *fallback on invalid* + if you choose scheduler that is not compatible with the model and fallback is not enabled, it will raise an error, + if fallback is enabled, it will try to find closest scheduler that is compatible with the model instead of just default scheduler + any change of requested-vs-active is logged as warning + plus add *beta start, beta end, steps offset* params to most schedulers + - **Prompt enhance** add info to image metadata - **CivitAI** downloaded thumbnails now include metadata - **Installer** support for `git+http` style references + - **XYZ Grid** add option *continue on error* to allow processing to continue even if one of the grid cells fails - **UI** - **Networks** using networks to load model or auto-download a reference model will now be reflected in the UI - ability to manually reorient *input/output* panels diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 914200c23..083925b85 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -218,8 +218,7 @@ def process_base(p: processing.StableDiffusionProcessing): if isinstance(v, torch.Tensor): err_args[k] = f'{v.device}:{v.dtype}:{v.shape}' log.error(f'Processing: args={err_args} {e}') - if shared.cmd_opts.debug: - errors.display(e, 'Processing') + errors.display(e, 'Processing') except RuntimeError as e: shared.state.interrupted = True err_args = base_args.copy() diff --git a/modules/res4lyf/abnorsett_scheduler.py b/modules/res4lyf/abnorsett_scheduler.py index 810b0deab..7a225d8d8 100644 --- a/modules/res4lyf/abnorsett_scheduler.py +++ b/modules/res4lyf/abnorsett_scheduler.py @@ -97,14 +97,40 @@ class ABNorsettScheduler(SchedulerMixin, ConfigMixin): def set_begin_index(self, begin_index: int = 0) -> None: self._begin_index = begin_index - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, get_sigmas_beta, get_sigmas_exponential, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps @@ -295,22 +321,22 @@ class ABNorsettScheduler(SchedulerMixin, ConfigMixin): x_next = sample + dt * (c0 * v_n + c1 * v_nm1) elif curr_order >= 3: - # For now, fallback to AB2 (variable) for higher orders to ensure stability - # given the complexity of variable-step AB3/4 formulas inline. - # The user specifically requested abnorsett_2m. - v_nm1 = self.model_outputs[-2] - sigma_prev = self.prev_sigmas[-2] - dt_prev = sigma_curr - sigma_prev + # For now, fallback to AB2 (variable) for higher orders to ensure stability + # given the complexity of variable-step AB3/4 formulas inline. + # The user specifically requested abnorsett_2m. + v_nm1 = self.model_outputs[-2] + sigma_prev = self.prev_sigmas[-2] + dt_prev = sigma_curr - sigma_prev - if abs(dt_prev) < 1e-8: - x_next = sample + dt * v_n - else: + if abs(dt_prev) < 1e-8: + x_next = sample + dt * v_n + else: r = dt / dt_prev c0 = 1 + 0.5 * r c1 = -0.5 * r x_next = sample + dt * (c0 * v_n + c1 * v_nm1) else: - x_next = sample + dt * v_n + x_next = sample + dt * v_n else: x_next = torch.exp(-h) * sample + h * res diff --git a/modules/res4lyf/bong_tangent_scheduler.py b/modules/res4lyf/bong_tangent_scheduler.py index d3b7eaa84..280186ffa 100644 --- a/modules/res4lyf/bong_tangent_scheduler.py +++ b/modules/res4lyf/bong_tangent_scheduler.py @@ -105,7 +105,7 @@ class BongTangentScheduler(SchedulerMixin, ConfigMixin): sample = sample / ((sigma**2 + 1) ** 0.5) return sample - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -113,7 +113,33 @@ class BongTangentScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps timestep_spacing = getattr(self.config, "timestep_spacing", "linspace") diff --git a/modules/res4lyf/common_sigma_scheduler.py b/modules/res4lyf/common_sigma_scheduler.py index 98e2d4162..561a38444 100644 --- a/modules/res4lyf/common_sigma_scheduler.py +++ b/modules/res4lyf/common_sigma_scheduler.py @@ -99,7 +99,7 @@ class CommonSigmaScheduler(SchedulerMixin, ConfigMixin): def set_begin_index(self, begin_index: int = 0) -> None: self._begin_index = begin_index - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -107,7 +107,33 @@ class CommonSigmaScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/deis_scheduler_alt.py b/modules/res4lyf/deis_scheduler_alt.py index bcb3a266e..f2b00a925 100644 --- a/modules/res4lyf/deis_scheduler_alt.py +++ b/modules/res4lyf/deis_scheduler_alt.py @@ -85,10 +85,36 @@ class RESDEISMultistepScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/etdrk_scheduler.py b/modules/res4lyf/etdrk_scheduler.py index cc7f693fe..bce0eae2c 100644 --- a/modules/res4lyf/etdrk_scheduler.py +++ b/modules/res4lyf/etdrk_scheduler.py @@ -97,7 +97,7 @@ class ETDRKScheduler(SchedulerMixin, ConfigMixin): def set_begin_index(self, begin_index: int = 0) -> None: self._begin_index = begin_index - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -105,7 +105,33 @@ class ETDRKScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/gauss_legendre_scheduler.py b/modules/res4lyf/gauss_legendre_scheduler.py index 0cbb5ea03..7486515be 100644 --- a/modules/res4lyf/gauss_legendre_scheduler.py +++ b/modules/res4lyf/gauss_legendre_scheduler.py @@ -145,9 +145,35 @@ class GaussLegendreScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps # 1. Spacing diff --git a/modules/res4lyf/langevin_dynamics_scheduler.py b/modules/res4lyf/langevin_dynamics_scheduler.py index be4a9c1f9..9106eef91 100644 --- a/modules/res4lyf/langevin_dynamics_scheduler.py +++ b/modules/res4lyf/langevin_dynamics_scheduler.py @@ -98,8 +98,8 @@ class LangevinDynamicsScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, generator: torch.Generator | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( @@ -109,7 +109,33 @@ class LangevinDynamicsScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/lawson_scheduler.py b/modules/res4lyf/lawson_scheduler.py index 3631024bf..7f4140393 100644 --- a/modules/res4lyf/lawson_scheduler.py +++ b/modules/res4lyf/lawson_scheduler.py @@ -95,7 +95,7 @@ class LawsonScheduler(SchedulerMixin, ConfigMixin): def set_begin_index(self, begin_index: int = 0) -> None: self._begin_index = begin_index - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -103,7 +103,33 @@ class LawsonScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/linear_rk_scheduler.py b/modules/res4lyf/linear_rk_scheduler.py index 955e4af8e..48fd17fb8 100644 --- a/modules/res4lyf/linear_rk_scheduler.py +++ b/modules/res4lyf/linear_rk_scheduler.py @@ -101,9 +101,35 @@ class LinearRKScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps # 1. Spacing diff --git a/modules/res4lyf/lobatto_scheduler.py b/modules/res4lyf/lobatto_scheduler.py index e1698b935..b66e36a40 100644 --- a/modules/res4lyf/lobatto_scheduler.py +++ b/modules/res4lyf/lobatto_scheduler.py @@ -101,9 +101,35 @@ class LobattoScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps # 1. Spacing diff --git a/modules/res4lyf/pec_scheduler.py b/modules/res4lyf/pec_scheduler.py index d5951b937..45710d8a2 100644 --- a/modules/res4lyf/pec_scheduler.py +++ b/modules/res4lyf/pec_scheduler.py @@ -99,8 +99,8 @@ class PECScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32, ): @@ -111,7 +111,33 @@ class PECScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/radau_iia_scheduler.py b/modules/res4lyf/radau_iia_scheduler.py index 4d072205b..f7efe07d2 100644 --- a/modules/res4lyf/radau_iia_scheduler.py +++ b/modules/res4lyf/radau_iia_scheduler.py @@ -135,9 +135,35 @@ class RadauIIAScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps # 1. Spacing diff --git a/modules/res4lyf/res_multistep_scheduler.py b/modules/res4lyf/res_multistep_scheduler.py index 081e83307..cec85e0ef 100644 --- a/modules/res4lyf/res_multistep_scheduler.py +++ b/modules/res4lyf/res_multistep_scheduler.py @@ -120,14 +120,40 @@ class RESMultistepScheduler(SchedulerMixin, ConfigMixin): sigma = self.sigmas[self._step_index] return sample / ((sigma**2 + 1) ** 0.5) - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, get_sigmas_beta, get_sigmas_exponential, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/res_multistep_sde_scheduler.py b/modules/res4lyf/res_multistep_sde_scheduler.py index adc40f832..5b0155c12 100644 --- a/modules/res4lyf/res_multistep_sde_scheduler.py +++ b/modules/res4lyf/res_multistep_sde_scheduler.py @@ -111,7 +111,7 @@ class RESMultistepSDEScheduler(SchedulerMixin, ConfigMixin): sample = sample / ((sigma**2 + 1) ** 0.5) return sample - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -119,7 +119,33 @@ class RESMultistepSDEScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/res_singlestep_scheduler.py b/modules/res4lyf/res_singlestep_scheduler.py index 86d10fd24..3eb0e82a7 100644 --- a/modules/res4lyf/res_singlestep_scheduler.py +++ b/modules/res4lyf/res_singlestep_scheduler.py @@ -96,7 +96,7 @@ class RESSinglestepScheduler(SchedulerMixin, ConfigMixin): sigma = self.sigmas[self._step_index] return sample / ((sigma**2 + 1) ** 0.5) - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -104,7 +104,33 @@ class RESSinglestepScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/res_singlestep_sde_scheduler.py b/modules/res4lyf/res_singlestep_sde_scheduler.py index a83b5b403..036c275be 100644 --- a/modules/res4lyf/res_singlestep_sde_scheduler.py +++ b/modules/res4lyf/res_singlestep_sde_scheduler.py @@ -99,7 +99,7 @@ class RESSinglestepSDEScheduler(SchedulerMixin, ConfigMixin): sample = sample / ((sigma**2 + 1) ** 0.5) return sample - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -107,7 +107,33 @@ class RESSinglestepSDEScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/res_unified_scheduler.py b/modules/res4lyf/res_unified_scheduler.py index 2b1d1bf0a..84cec78d1 100644 --- a/modules/res4lyf/res_unified_scheduler.py +++ b/modules/res4lyf/res_unified_scheduler.py @@ -92,14 +92,40 @@ class RESUnifiedScheduler(SchedulerMixin, ConfigMixin): sigma = self.sigmas[self._step_index] return sample / ((sigma**2 + 1) ** 0.5) - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, get_sigmas_beta, get_sigmas_exponential, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps timestep_spacing = getattr(self.config, "timestep_spacing", "linspace") diff --git a/modules/res4lyf/riemannian_flow_scheduler.py b/modules/res4lyf/riemannian_flow_scheduler.py index 882ad54a3..4c3f5f6df 100644 --- a/modules/res4lyf/riemannian_flow_scheduler.py +++ b/modules/res4lyf/riemannian_flow_scheduler.py @@ -95,7 +95,7 @@ class RiemannianFlowScheduler(SchedulerMixin, ConfigMixin): def set_begin_index(self, begin_index: int = 0) -> None: self._begin_index = begin_index - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -103,7 +103,33 @@ class RiemannianFlowScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps timestep_spacing = getattr(self.config, "timestep_spacing", "linspace") diff --git a/modules/res4lyf/rungekutta_44s_scheduler.py b/modules/res4lyf/rungekutta_44s_scheduler.py index d18941e01..20be7b564 100644 --- a/modules/res4lyf/rungekutta_44s_scheduler.py +++ b/modules/res4lyf/rungekutta_44s_scheduler.py @@ -68,7 +68,33 @@ class RungeKutta44Scheduler(SchedulerMixin, ConfigMixin): self._sigmas_cpu = None self._step_index = None - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps # 1. Base sigmas diff --git a/modules/res4lyf/rungekutta_57s_scheduler.py b/modules/res4lyf/rungekutta_57s_scheduler.py index 5d118bff7..97443bfd1 100644 --- a/modules/res4lyf/rungekutta_57s_scheduler.py +++ b/modules/res4lyf/rungekutta_57s_scheduler.py @@ -70,9 +70,35 @@ class RungeKutta57Scheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps # 1. Spacing diff --git a/modules/res4lyf/rungekutta_67s_scheduler.py b/modules/res4lyf/rungekutta_67s_scheduler.py index 55af1b16c..861e13487 100644 --- a/modules/res4lyf/rungekutta_67s_scheduler.py +++ b/modules/res4lyf/rungekutta_67s_scheduler.py @@ -70,9 +70,35 @@ class RungeKutta67Scheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, - num_inference_steps: int, - device: str | torch.device = None, + num_inference_steps: int | None = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps # 1. Spacing diff --git a/modules/res4lyf/scheduler_utils.py b/modules/res4lyf/scheduler_utils.py index 608ac11ad..3d401baf6 100644 --- a/modules/res4lyf/scheduler_utils.py +++ b/modules/res4lyf/scheduler_utils.py @@ -84,6 +84,27 @@ def get_sigmas_flow(n, sigma_min, sigma_max, device="cpu", dtype: torch.dtype = def apply_shift(sigmas, shift): return shift * sigmas / (1 + (shift - 1) * sigmas) + +def get_base_sigmas(alphas_cumprod: torch.Tensor) -> np.ndarray: + """Return the standard diffusion sigmas derived from alphas_cumprod.""" + return np.array(((1 - alphas_cumprod.cpu().numpy()) / alphas_cumprod.cpu().numpy()) ** 0.5, dtype=np.float32) + + +def sigma_to_t(sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: + """Convert sigma values to corresponding discrete timesteps via interpolation.""" + sigma = np.array(sigma, dtype=np.float32) + log_sigma = np.log(np.maximum(sigma, 1e-10)) + dists = log_sigma - log_sigmas[:, np.newaxis] + low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) + high_idx = low_idx + 1 + low = log_sigmas[low_idx] + high = log_sigmas[high_idx] + w = (low - log_sigma) / (low - high) + w = np.clip(w, 0, 1) + t = (1 - w) * low_idx + w * high_idx + return t.reshape(sigma.shape).astype(np.float32) + + def get_dynamic_shift(mu, base_shift, max_shift, base_seq_len, max_seq_len): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len @@ -117,3 +138,43 @@ def add_noise_to_sample( noisy_samples = original_samples + sigma * noise return noisy_samples + + +def validate_custom_schedule_args(timesteps, sigmas): + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` should be set.") + if timesteps is None and sigmas is None: + raise ValueError("Must pass exactly one of `timesteps` or `sigmas`.") + + +def prepare_res4lyf_timesteps_and_sigmas( + config, + alphas_cumprod: torch.Tensor, + num_inference_steps: int | None = None, + timesteps=None, + sigmas=None, + device: str | torch.device = None, + dtype: torch.dtype = torch.float32, +): + validate_custom_schedule_args(timesteps, sigmas) + + base_sigmas = get_base_sigmas(alphas_cumprod) + + if timesteps is not None: + if getattr(config, "use_karras_sigmas", False) or getattr(config, "use_exponential_sigmas", False) or getattr(config, "use_beta_sigmas", False): + raise ValueError("Cannot set `timesteps` when karras/exponential/beta sigmas are enabled.") + if getattr(config, "use_flow_sigmas", False): + raise ValueError("Cannot set `timesteps` when `use_flow_sigmas` is enabled.") + + timesteps_array = np.array(timesteps, dtype=np.float32) + sigmas_array = np.interp(timesteps_array, np.arange(len(base_sigmas), dtype=np.float32), base_sigmas) + num_inference_steps = len(timesteps_array) + sigmas_array = np.concatenate([sigmas_array, [0.0]]) + return num_inference_steps, timesteps_array, sigmas_array + + sigmas_array = np.array(sigmas, dtype=np.float32) + if num_inference_steps is None: + num_inference_steps = len(sigmas_array) - 1 + log_sigmas = np.log(base_sigmas) + timesteps_array = np.array([sigma_to_t(sigma, log_sigmas) for sigma in sigmas_array[:-1]], dtype=np.float32) + return num_inference_steps, timesteps_array, sigmas_array diff --git a/modules/res4lyf/simple_exponential_scheduler.py b/modules/res4lyf/simple_exponential_scheduler.py index 01a901e12..eddd60acf 100644 --- a/modules/res4lyf/simple_exponential_scheduler.py +++ b/modules/res4lyf/simple_exponential_scheduler.py @@ -95,7 +95,7 @@ class SimpleExponentialScheduler(SchedulerMixin, ConfigMixin): def set_begin_index(self, begin_index: int = 0) -> None: self._begin_index = begin_index - def set_timesteps(self, num_inference_steps: int, device: str | torch.device = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + def set_timesteps(self, num_inference_steps: int | None = None, device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): from .scheduler_utils import ( apply_shift, get_dynamic_shift, @@ -103,7 +103,33 @@ class SimpleExponentialScheduler(SchedulerMixin, ConfigMixin): get_sigmas_exponential, get_sigmas_flow, get_sigmas_karras, + prepare_res4lyf_timesteps_and_sigmas, ) + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps diff --git a/modules/res4lyf/specialized_rk_scheduler.py b/modules/res4lyf/specialized_rk_scheduler.py index 33b6df815..75820e3cb 100644 --- a/modules/res4lyf/specialized_rk_scheduler.py +++ b/modules/res4lyf/specialized_rk_scheduler.py @@ -106,8 +106,34 @@ class SpecializedRKScheduler(SchedulerMixin, ConfigMixin): def set_timesteps( self, num_inference_steps: int, - device: str | torch.device = None, + device: str | torch.device = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, mu: float | None = None, dtype: torch.dtype = torch.float32): + from .scheduler_utils import prepare_res4lyf_timesteps_and_sigmas + if timesteps is not None or sigmas is not None: + num_inference_steps, timesteps, sigmas = prepare_res4lyf_timesteps_and_sigmas( + self.config, + self.alphas_cumprod, + num_inference_steps, + timesteps=timesteps, + sigmas=sigmas, + device=device, + dtype=dtype, + ) + self.num_inference_steps = num_inference_steps + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=dtype) + self.sigmas = torch.from_numpy(sigmas).to(device=device, dtype=dtype) + self.init_noise_sigma = self.sigmas.max().item() if self.sigmas.numel() > 0 else 1.0 + for attr in ("_step_index", "_begin_index", "model_outputs", "x0_outputs", "prev_sigmas", "lower_order_nums", "sample_at_start_of_step"): + if hasattr(self, attr): + if attr in ("_step_index", "_begin_index"): + setattr(self, attr, None) + elif attr == "lower_order_nums": + setattr(self, attr, 0) + elif attr == "sample_at_start_of_step": + setattr(self, attr, None) + else: + setattr(self, attr, []) + return self.num_inference_steps = num_inference_steps # 1. Spacing diff --git a/modules/schedulers/scheduler_dpm_flowmatch.py b/modules/schedulers/scheduler_dpm_flowmatch.py index 3f6c6827c..8289aa2f5 100644 --- a/modules/schedulers/scheduler_dpm_flowmatch.py +++ b/modules/schedulers/scheduler_dpm_flowmatch.py @@ -181,7 +181,7 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): if solver_type not in ["midpoint", "heun"]: raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") - if sigma_schedule not in [None, "karras", "exponential", "lambdas", "betas"]: + if sigma_schedule not in [None, "karras", "exponential", "lambdas", "betas", "flowmatch"]: raise NotImplementedError(f"{sigma_schedule} is not implemented for {self.__class__}") if beta_schedule not in ["linear", "scaled linear"]: diff --git a/modules/sd_hijack_schedulers.py b/modules/sd_hijack_schedulers.py new file mode 100644 index 000000000..fd79b3430 --- /dev/null +++ b/modules/sd_hijack_schedulers.py @@ -0,0 +1,324 @@ +from __future__ import annotations +import inspect +import torch +import diffusers +import numpy as np + + +_scheduled_classes = [ + "DEISMultistepScheduler", + "DPMSolverMultistepScheduler", + "DPMSolverMultistepInverseScheduler", + "DPMSolverSinglestepScheduler", + "FlowMatchHeunDiscreteScheduler", + "SASolverScheduler", + "UniPCMultistepScheduler", +] +_patched_schedulers = set() +_orig_unipc_set_timesteps = None + + +def init_hijack(): + for class_name in _scheduled_classes: + scheduler_cls = getattr(diffusers, class_name, None) + if scheduler_cls is None: + continue + _patch_scheduler_set_timesteps(scheduler_cls) + + +def _patch_scheduler_set_timesteps(scheduler_cls): + if scheduler_cls in _patched_schedulers: + return + + scheduler_cls.original_set_timesteps = scheduler_cls.set_timesteps + + # @wraps(scheduler_cls.original_set_timesteps) + def set_timesteps(self, num_inference_steps=None, device=None, timesteps=None, sigmas=None, mu=None, **kwargs): + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed.") + + if timesteps is None and sigmas is None: + return _call_original_set_timesteps( + scheduler_cls.original_set_timesteps, + self, + num_inference_steps=num_inference_steps, + device=device, + mu=mu, + **kwargs, + ) + + if timesteps is not None: + if "timesteps" in set(inspect.signature(scheduler_cls.original_set_timesteps).parameters.keys()): + return _call_original_set_timesteps( + scheduler_cls.original_set_timesteps, + self, + num_inference_steps=None, + device=device, + mu=mu, + timesteps=timesteps, + **kwargs, + ) + if "sigmas" in set(inspect.signature(scheduler_cls.original_set_timesteps).parameters.keys()) and getattr(self.config, "use_flow_sigmas", False): + sigmas_values = _invert_unipc_timesteps(self, timesteps) + return _call_original_set_timesteps( + scheduler_cls.original_set_timesteps, + self, + num_inference_steps=None, + device=device, + mu=mu, + sigmas=sigmas_values, + **kwargs, + ) + num_inference_steps, timesteps_array, sigmas_array = _prepare_custom_schedule_from_timesteps( + self, + timesteps, + ) + _assign_custom_schedule(self, num_inference_steps, timesteps_array, sigmas_array, device) + return + + if sigmas is not None: + if "sigmas" in set(inspect.signature(scheduler_cls.original_set_timesteps).parameters.keys()): + return _call_original_set_timesteps( + scheduler_cls.original_set_timesteps, + self, + num_inference_steps=None, + device=device, + mu=mu, + sigmas=sigmas, + **kwargs, + ) + num_inference_steps, timesteps_array, sigmas_array = _prepare_custom_schedule_from_sigmas(self, sigmas, mu=mu) + _assign_custom_schedule(self, num_inference_steps, timesteps_array, sigmas_array, device) + return + + scheduler_cls.set_timesteps = set_timesteps + _patched_schedulers.add(scheduler_cls) + + +def _call_original_set_timesteps(original, self, num_inference_steps=None, device=None, mu=None, **kwargs): + signature = inspect.signature(original) + call_args = {} + + if "num_inference_steps" in signature.parameters and num_inference_steps is not None: + call_args["num_inference_steps"] = num_inference_steps + if "device" in signature.parameters: + call_args["device"] = device + if "mu" in signature.parameters and mu is not None: + call_args["mu"] = mu + + if "sigmas" in signature.parameters and "sigmas" in kwargs: + sigmas_value = kwargs["sigmas"] + if not isinstance(sigmas_value, (np.ndarray, torch.Tensor)): + kwargs["sigmas"] = np.array(sigmas_value, dtype=np.float32) + if "timesteps" in signature.parameters and "timesteps" in kwargs: + timesteps_value = kwargs["timesteps"] + if not isinstance(timesteps_value, (np.ndarray, torch.Tensor)): + kwargs["timesteps"] = np.array(timesteps_value, dtype=np.int64) + call_args.update(kwargs) + return original(self, **call_args) + + +def _get_base_sigmas(scheduler) -> np.ndarray: + if hasattr(scheduler, "alphas_cumprod"): + alphas_cumprod = scheduler.alphas_cumprod.cpu().numpy() + return np.array(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, dtype=np.float32) + if hasattr(scheduler, "sigmas"): + return np.array(scheduler.sigmas.cpu().numpy(), dtype=np.float32) + raise ValueError("Scheduler does not expose alphas_cumprod or sigmas for custom schedule conversion.") + + +def _get_final_sigma(scheduler) -> float: + if getattr(scheduler.config, "final_sigmas_type", None) == "zero": + return 0.0 + base_sigmas = _get_base_sigmas(scheduler) + return float(base_sigmas[0]) + + +def _sigma_to_t(sigma: np.ndarray, log_sigmas: np.ndarray) -> np.ndarray: + sigma = np.array(sigma, dtype=np.float32) + log_sigma = np.log(np.maximum(sigma, 1e-10)) + dists = log_sigma - log_sigmas[:, np.newaxis] + low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2) + high_idx = low_idx + 1 + low = log_sigmas[low_idx] + high = log_sigmas[high_idx] + w = (low - log_sigma) / (low - high) + w = np.clip(w, 0, 1) + t = (1 - w) * low_idx + w * high_idx + return t.reshape(sigma.shape).astype(np.float32) + + +def _compute_timesteps_from_sigmas(scheduler, sigmas_array: np.ndarray) -> np.ndarray: + base_sigmas = _get_base_sigmas(scheduler) + log_sigmas = np.log(base_sigmas) + if hasattr(scheduler, "_sigma_to_t"): + sigma_to_t = scheduler._sigma_to_t + parameters = list(inspect.signature(sigma_to_t).parameters) + if len(parameters) == 2: + results = [] + for sigma in sigmas_array: + value = sigma_to_t(np.array([sigma], dtype=np.float32), log_sigmas) + value = np.array(value, dtype=np.float32) + results.append(float(value.reshape(-1)[0])) + timesteps = np.array(results, dtype=np.float32) + else: + timesteps = np.array( + [sigma_to_t(float(sigma)) for sigma in sigmas_array], + dtype=np.float32, + ) + else: + timesteps = _sigma_to_t(sigmas_array, log_sigmas) + return np.round(timesteps).astype(np.int64) + + +def _prepare_custom_schedule_from_sigmas(scheduler, sigmas, mu=None): + sigmas_array = np.array(sigmas, dtype=np.float32) + if sigmas_array.ndim != 1: + raise ValueError("`sigmas` must be a 1D sequence.") + if sigmas_array.size == 0: + raise ValueError("`sigmas` cannot be empty.") + + final_sigma = _get_final_sigma(scheduler) + if sigmas_array.size > 1 and np.isclose(sigmas_array[-1], final_sigma, atol=1e-6): + sigmas_values = sigmas_array[:-1] + appended_sigma = float(sigmas_array[-1]) + else: + sigmas_values = sigmas_array + appended_sigma = final_sigma + + if sigmas_values.size == 0: + raise ValueError("`sigmas` must contain at least one non-final sigma value.") + + if getattr(scheduler.config, "use_flow_sigmas", False): + flow_sigmas = sigmas_values.astype(np.float32) + if getattr(scheduler.config, "use_dynamic_shifting", False): + if mu is None: + raise ValueError("`mu` is required for flow sigmas when use_dynamic_shifting is enabled.") + if not hasattr(scheduler, "time_shift"): + raise ValueError("Scheduler does not support dynamic shifting for custom sigmas.") + flow_sigmas = scheduler.time_shift(mu, 1.0, flow_sigmas) + else: + flow_shift = float(getattr(scheduler.config, "flow_shift", 1.0)) + flow_sigmas = flow_shift * flow_sigmas / (1 + (flow_shift - 1.0) * flow_sigmas) + + if getattr(scheduler.config, "shift_terminal", False) and hasattr(scheduler, "stretch_shift_to_terminal"): + flow_sigmas = scheduler.stretch_shift_to_terminal(flow_sigmas) + + eps = 1e-6 + if np.fabs(flow_sigmas[0] - 1) < eps: + flow_sigmas[0] -= eps + + timesteps = np.round(flow_sigmas * float(scheduler.config.num_train_timesteps)).astype(np.int64) + elif getattr(scheduler.config, "use_karras_sigmas", False) or getattr(scheduler.config, "use_exponential_sigmas", False) or getattr(scheduler.config, "use_beta_sigmas", False) or getattr(scheduler.config, "use_lu_lambdas", False): + raise ValueError("Custom sigmas are not supported when the scheduler uses a specialized sigma schedule configuration.") + else: + timesteps = _compute_timesteps_from_sigmas(scheduler, sigmas_values) + + return int(timesteps.shape[0]), timesteps, np.concatenate([sigmas_values, [appended_sigma]]).astype(np.float32) + + +def _prepare_custom_schedule_from_timesteps(scheduler, timesteps): + timesteps_array = np.array(timesteps, dtype=np.float32) + if timesteps_array.ndim != 1: + raise ValueError("`timesteps` must be a 1D sequence.") + if timesteps_array.size == 0: + raise ValueError("`timesteps` cannot be empty.") + + if getattr(scheduler.config, "use_karras_sigmas", False) or getattr(scheduler.config, "use_exponential_sigmas", False) or getattr(scheduler.config, "use_beta_sigmas", False) or getattr(scheduler.config, "use_flow_sigmas", False) or getattr(scheduler.config, "use_lu_lambdas", False): + raise ValueError("Cannot set custom timesteps when the scheduler uses a specialized sigma schedule configuration.") + + base_sigmas = _get_base_sigmas(scheduler) + sigma_values = np.interp(timesteps_array, np.arange(base_sigmas.shape[0], dtype=np.float32), base_sigmas).astype(np.float32) + final_sigma = _get_final_sigma(scheduler) + return int(timesteps_array.shape[0]), timesteps_array.astype(np.int64), np.concatenate([sigma_values, [final_sigma]]).astype(np.float32) + + +def _invert_unipc_timesteps(scheduler, timesteps): + timesteps_array = np.array(timesteps, dtype=np.float32) + num_train_timesteps = float(scheduler.config.num_train_timesteps) + transformed_sigmas = timesteps_array / num_train_timesteps + if np.any(transformed_sigmas <= 0) or np.any(transformed_sigmas >= 1): + raise ValueError("Custom timesteps for UniPCMultistepScheduler must be within the valid flow sigma range (0, num_train_timesteps).") + if getattr(scheduler.config, "use_dynamic_shifting", False): + raise ValueError("Cannot convert custom timesteps to sigmas for UniPCMultistepScheduler when use_dynamic_shifting is enabled.") + if getattr(scheduler.config, "shift_terminal", False): + raise ValueError("Cannot convert custom timesteps to sigmas for UniPCMultistepScheduler when shift_terminal is enabled.") + flow_shift = float(getattr(scheduler.config, "flow_shift", 1.0)) + if flow_shift != 1.0: + sigmas = transformed_sigmas / (flow_shift - (flow_shift - 1.0) * transformed_sigmas) + else: + sigmas = transformed_sigmas + return sigmas.astype(np.float32) + + +def _assign_custom_schedule(scheduler, num_inference_steps, timesteps, sigmas, device): + scheduler.timesteps = torch.from_numpy(np.array(timesteps, dtype=np.int64)).to(device=device, dtype=torch.int64) + scheduler.sigmas = torch.from_numpy(np.array(sigmas, dtype=np.float32)) + scheduler.num_inference_steps = int(num_inference_steps) + + if hasattr(scheduler, "config") and hasattr(scheduler.config, "solver_order"): + scheduler.model_outputs = [None] * scheduler.config.solver_order + scheduler.lower_order_nums = 0 + + scheduler._step_index = None + scheduler._begin_index = None + scheduler.sigmas = scheduler.sigmas.to("cpu") + + +def hijack_unipc(): + global _orig_unipc_set_timesteps # pylint: disable=global-statement + + from diffusers import UniPCMultistepScheduler + _orig_unipc_set_timesteps = UniPCMultistepScheduler.set_timesteps + + def _unipc_set_timesteps_device_fix(self, num_inference_steps=None, device=None, **kwargs): + _orig_unipc_set_timesteps(self, num_inference_steps=num_inference_steps, device=device, **kwargs) + if device is not None: + self.sigmas = self.sigmas.to(device) + + UniPCMultistepScheduler.set_timesteps = _unipc_set_timesteps_device_fix + + +def attach_scale_noise_if_missing(sampler): + def scale_noise(sample, timestep, noise=None): + if noise is None: + raise ValueError("`scale_noise` requires a `noise` tensor") + + sigmas = sampler.sigmas.to(device=sample.device, dtype=sample.dtype) + schedule_timesteps = sampler.timesteps.to(device=sample.device, dtype=sample.dtype) + + if isinstance(timestep, torch.Tensor): + timestep_tensor = timestep.to(device=sample.device, dtype=sample.dtype) + else: + timestep_tensor = torch.tensor([timestep], device=sample.device, dtype=sample.dtype) + + if timestep_tensor.ndim == 0: + timestep_tensor = timestep_tensor.unsqueeze(0) + + if schedule_timesteps.ndim == 0: + sigma = sigmas + else: + step_indices = [] + if getattr(sampler, "begin_index", None) is None: + for t in timestep_tensor: + indices = (schedule_timesteps == t).nonzero(as_tuple=False) + if len(indices) > 1: + step_indices.append(indices[1].item()) + elif len(indices) == 1: + step_indices.append(indices[0].item()) + else: + step_indices.append(torch.argmin(torch.abs(schedule_timesteps - t)).item()) + elif getattr(sampler, "step_index", None) is not None: + step_indices = [sampler.step_index] * timestep_tensor.shape[0] + else: + step_indices = [sampler.begin_index] * timestep_tensor.shape[0] + + step_indices = torch.tensor(step_indices, device=schedule_timesteps.device, dtype=torch.int64) + sigma = sigmas[step_indices] + + while sigma.ndim < noise.ndim: + sigma = sigma.unsqueeze(-1) + + return sigma * noise + (1.0 - sigma) * sample + + sampler.scale_noise = scale_noise diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index cda0b01bd..d38ab0acf 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -4,8 +4,7 @@ from modules import shared from modules.logger import log -debug = log.trace if os.environ.get('SD_SAMPLER_DEBUG', None) is not None else lambda *args, **kwargs: None -debug('Trace: SAMPLER') +debug = os.environ.get('SD_SAMPLER_DEBUG', None) all_samplers = [] all_samplers_map = {} samplers = all_samplers @@ -19,9 +18,7 @@ def find_sampler(name:str): return all_samplers_map.get("UniPC", None) for sampler in all_samplers: if sampler.name.lower() == name.lower() or name in sampler.aliases: - debug(f'Find sampler: name="{name}" found={sampler.name}') return sampler - debug(f'Find sampler: name="{name}" found=None') return None @@ -38,7 +35,6 @@ def list_samplers(): samplers_for_img2img = all_samplers samplers_map = {} return all_samplers - # log.debug(f'Available samplers: {[x.name for x in all_samplers]}') def find_sampler_config(name): @@ -49,7 +45,7 @@ def find_sampler_config(name): return config -def restore_default(model): +def restore_default(model, requested="Default"): if model is None: return None if getattr(model, "default_scheduler", None) is not None and getattr(model, "scheduler", None) is not None: @@ -62,7 +58,10 @@ def restore_default(model): shared.state.prediction_type = "flow_prediction" elif hasattr(model.scheduler, "config") and hasattr(model.scheduler.config, "prediction_type"): shared.state.prediction_type = model.scheduler.config.prediction_type - log.debug(f'Sampler: "Default" cls={model.scheduler.__class__.__name__} config={config}') + if requested != "Default": + log.warning(f'Sampler: requested="{requested}" set="Default" cls={model.scheduler.__class__.__name__} config={config}') + else: + log.debug(f'Sampler: Default cls={model.scheduler.__class__.__name__} config={config}') return model.scheduler @@ -85,10 +84,31 @@ def create_sampler(name, model, scheduler_overrides=None): if name == 'Default' and hasattr(model, 'scheduler'): return restore_default(model) + config = None + + # switch to flow variant when applicable + if not is_flexible and config is None and requires_flow and 'Flow' not in name and shared.opts.schedulers_fallback: + redirect = f'{name} FlowMatch' + config = find_sampler_config(redirect) + if config is not None: + log.warning(f'Sampler: requested="{name}" redirected="{redirect}"') + name = redirect + + # switch to discrete variant when applicable + if not is_flexible and config is None and not requires_flow and 'Flow' in name and shared.opts.schedulers_fallback: + redirect = name.replace(' FlowMatch', '').strip() + config = find_sampler_config(redirect) + if config is not None: + log.warning(f'Sampler: requested="{name}" redirected="{redirect}"') + name = redirect + # create sampler - config = find_sampler_config(name) + if config is None: + config = find_sampler_config(name) + if config is None or config.constructor is None: - return restore_default(model) + return restore_default(model, name) + from modules import sd_samplers_diffusers sd_samplers_diffusers.scheduler_overrides = scheduler_overrides or {} try: @@ -96,18 +116,26 @@ def create_sampler(name, model, scheduler_overrides=None): finally: sd_samplers_diffusers.scheduler_overrides = {} if sampler.sampler is None: - return restore_default(model) - is_flow = ('FlowMatch' in sampler.sampler.__class__.__name__) or (getattr(sampler.sampler.config, 'prediction_type', None) == 'flow_prediction') + return restore_default(model, name) + + pred_type = getattr(sampler.sampler.config, 'prediction_type', None) + is_flow = ('FlowMatch' in sampler.sampler.__class__.__name__) or (pred_type == 'flow_prediction') # validate sampler prediction type - if (model is not None) and is_flexible: + if (model is None) or is_flexible: pass elif (model is not None) and (is_flow and not requires_flow): - log.error(f'Sampler: "{sampler.name}" cls={sampler.sampler.__class__.__name__} pipe={model.__class__.__name__} model requires sampler with discrete prediction') - return restore_default(model) + log.error(f'Sampler: "{sampler.name}" cls={sampler.sampler.__class__.__name__} pipe={model.__class__.__name__} type={pred_type} model requires sampler with discrete prediction') + if not debug: + return restore_default(model, name) + else: + raise ValueError(f'Sampler: name="{sampler.name}" cls={sampler.sampler.__class__.__name__} type={pred_type} model requires sampler with discrete prediction') elif (model is not None) and (not is_flow and requires_flow): - log.error(f'Sampler: "{sampler.name}" cls={sampler.sampler.__class__.__name__} pipe={model.__class__.__name__} model requires sampler with flow prediction') - return restore_default(model) + log.error(f'Sampler: "{sampler.name}" cls={sampler.sampler.__class__.__name__} pipe={model.__class__.__name__} type={pred_type} model requires sampler with flow prediction') + if not debug: + return restore_default(model, name) + else: + raise ValueError(f'Sampler: name="{sampler.name}" cls={sampler.sampler.__class__.__name__} type={pred_type} model requires sampler with flow prediction') # assign sampler if model is not None: diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index 1e83a1429..d3a250756 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -5,7 +5,7 @@ import inspect import diffusers from modules import shared, errors from modules.logger import log -from modules.sd_samplers_hijack import init_samplers_hijack +from modules.sd_hijack_schedulers import init_hijack, hijack_unipc, attach_scale_noise_if_missing # pylint: disable=unused-import from modules.sd_samplers_common import SamplerData, flow_models @@ -13,7 +13,8 @@ debug = os.environ.get('SD_SAMPLER_DEBUG', None) is not None debug_log = log.trace if debug else lambda *args, **kwargs: None scheduler_overrides = {} # set by sd_samplers.create_sampler() before constructor call flow_exclude = ['PeRFlow'] -init_samplers_hijack() +hijack_unipc() +# init_hijack() # Diffusers schedulers try: @@ -371,16 +372,17 @@ def get_sampler_capability(): # TODO enso-required return {} -def get_override(key, default=None): +def get_override(key): if key in scheduler_overrides: return scheduler_overrides[key] - return getattr(shared.opts, key, default) + return getattr(shared.opts, key, None) class DiffusionSampler: def __init__(self, name, constructor, model, **kwargs): if name == 'Default': return + self.name = name self.config = {} self.sampler = None @@ -508,31 +510,46 @@ class DiffusionSampler: self.sampler = None return - if self.config.get('prediction_type') == 'flow_prediction' and 'FlowMatch' not in constructor.__name__: + if (self.config.get('prediction_type') == 'flow_prediction') and ('FlowMatch' not in constructor.__name__): try: cls_source = inspect.getsource(constructor) if '"flow_prediction"' not in cls_source and "'flow_prediction'" not in cls_source: - log.warning(f'Sampler: "{name}" does not support flow_prediction') - self.sampler = None - return + if debug or not shared.opts.schedulers_fallback: + raise ValueError(f'Sampler: name="{name}" does not appear to support flow_prediction') + else: + log.warning(f'Sampler: name="{name}" does not support flow_prediction') + self.sampler = None + return except (TypeError, OSError): pass if hasattr(sampler, 'set_timesteps'): + # if not hasattr(sampler, "scale_noise") and hasattr(sampler, "timesteps") and hasattr(sampler, "sigmas"): + # attach_scale_noise_if_missing(sampler) + accept_sigmas = "sigmas" in set(inspect.signature(sampler.set_timesteps).parameters.keys()) accepts_timesteps = "timesteps" in set(inspect.signature(sampler.set_timesteps).parameters.keys()) accept_scale_noise = hasattr(sampler, "scale_noise") - debug_log(f'Sampler: "{name}" sigmas={accept_sigmas} timesteps={accepts_timesteps}') - default_accept_sigmas = model is not None and hasattr(model.default_scheduler, 'set_timesteps') and "sigmas" in set(inspect.signature(model.default_scheduler.set_timesteps).parameters.keys()) - default_accept_scale_noise = model is not None and hasattr(model.default_scheduler, "scale_noise") + debug_log(f'Sampler: name="{name}" sigmas={accept_sigmas} timesteps={accepts_timesteps} scale_noise={accept_scale_noise}') + + default_accept_sigmas = (model is not None) and hasattr(model.default_scheduler, 'set_timesteps') and "sigmas" in set(inspect.signature(model.default_scheduler.set_timesteps).parameters.keys()) if default_accept_sigmas and not accept_sigmas: - log.warning(f'Sampler: "{name}" does not accept sigmas') - self.sampler = None - return + if debug or not shared.opts.schedulers_fallback: + raise ValueError(f'Sampler: name="{name}" does not accept sigmas') + else: + log.warning(f'Sampler: name="{name}" does not accept sigmas') + self.sampler = None + return + + default_accept_scale_noise = (model is not None) and hasattr(model.default_scheduler, "scale_noise") if default_accept_scale_noise and not accept_scale_noise: - log.warning(f'Sampler: "{name}" does not implement scale noise') - self.sampler = None - return + log.warning(f'Sampler: name="{name}" does not implement scale noise') + if debug or not shared.opts.schedulers_fallback: + raise ValueError(f'Sampler: name="{name}" does not implement scale noise') + else: + log.warning(f'Sampler: name="{name}" does not implement scale noise') + self.sampler = None + return # monkey-patch to allow sdxl pipeline to execute flowmatch samplers if not hasattr(sampler, 'scale_model_input'): @@ -541,6 +558,4 @@ class DiffusionSampler: sampler.init_noise_sigma = 1.0 self.sampler = sampler - - # log.debug_log(f'Sampler: class="{self.sampler.__class__.__name__}" config={self.sampler.config}') self.sampler.name = name diff --git a/modules/sd_samplers_hijack.py b/modules/sd_samplers_hijack.py deleted file mode 100644 index 750065ba1..000000000 --- a/modules/sd_samplers_hijack.py +++ /dev/null @@ -1,15 +0,0 @@ -_orig_unipc_set_timesteps = None - - -def init_samplers_hijack(): - global _orig_unipc_set_timesteps # pylint: disable=global-statement - - from diffusers import UniPCMultistepScheduler - _orig_unipc_set_timesteps = UniPCMultistepScheduler.set_timesteps - - def _unipc_set_timesteps_device_fix(self, num_inference_steps=None, device=None, **kwargs): - _orig_unipc_set_timesteps(self, num_inference_steps=num_inference_steps, device=device, **kwargs) - if device is not None: - self.sigmas = self.sigmas.to(device) - - UniPCMultistepScheduler.set_timesteps = _unipc_set_timesteps_device_fix diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 950308273..695a1f5c6 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -752,6 +752,7 @@ def create_settings(cmd_opts): "schedulers_beta_end": OptionInfo(0, "Beta end", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001}), "schedulers_timesteps_range": OptionInfo(1000, "Timesteps range", gr.Slider, {"minimum": 250, "maximum": 4000, "step": 1}), "schedulers_shift": OptionInfo(3, "Sampler shift", gr.Slider, {"minimum": 0.1, "maximum": 10, "step": 0.1, "visible": False}), + "schedulers_fallback": OptionInfo(True, "Sampler fallback on invalid", gr.Checkbox, {"visible": False}), "schedulers_dynamic_shift": OptionInfo(False, "Sampler dynamic shift", gr.Checkbox, {"visible": False}), "schedulers_sigma_adjust": OptionInfo(1.0, "Sigma adjust", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01, "visible": False}), "schedulers_sigma_adjust_min": OptionInfo(0.2, "Sigma adjust start", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), diff --git a/modules/ui_sections.py b/modules/ui_sections.py index c215bdcde..5dd003400 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -244,6 +244,11 @@ def create_sampler_options(tabname): log.debug(f'Sampler set options: {sampler_options}') shared.opts.save(silent=True) + def set_sampler_fallback(fallback): + log.debug(f'Sampler set options: fallback={fallback}') + shared.opts.schedulers_fallback = fallback + shared.opts.save(silent=True) + def set_sampler_timesteps(timesteps): log.debug(f'Sampler set options: timesteps={timesteps}') shared.opts.schedulers_timesteps = timesteps @@ -323,6 +328,8 @@ def create_sampler_options(tabname): values += ['dynamic'] if shared.opts.data.get('schedulers_dynamic_shift', False) else [] values += ['rescale'] if shared.opts.data.get('schedulers_rescale_betas', False) else [] sampler_options = gr.CheckboxGroup(label='Options', elem_id=f"{tabname}_sampler_options", choices=options, value=values, type='value') + with gr.Row(elem_classes=['flex-break']): + sampler_fallback = gr.Checkbox(label='Fallback on invalid', value=shared.opts.schedulers_fallback, elem_id=f"{tabname}_sampler_fallback") sampler_sigma.change(fn=set_sampler_sigma, inputs=[sampler_sigma], outputs=[]) sampler_spacing.change(fn=set_sampler_spacing, inputs=[sampler_spacing], outputs=[]) @@ -333,6 +340,7 @@ def create_sampler_options(tabname): sampler_order.change(fn=set_sampler_order, inputs=[sampler_order], outputs=[]) sampler_shift.change(fn=set_sampler_shift, inputs=[sampler_shift, sampler_base_shift, sampler_max_shift], outputs=[]) sampler_options.change(fn=set_sampler_options, inputs=[sampler_options], outputs=[]) + sampler_fallback.change(fn=set_sampler_fallback, inputs=[sampler_fallback], outputs=[]) sampler_sigma_adjust_val.change(fn=set_sigma_adjust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[]) sampler_sigma_adjust_min.change(fn=set_sigma_adjust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[]) sampler_sigma_adjust_max.change(fn=set_sigma_adjust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[]) diff --git a/scripts/xyz_grid_on.py b/scripts/xyz_grid_on.py index 992208ebe..7bbb67368 100644 --- a/scripts/xyz_grid_on.py +++ b/scripts/xyz_grid_on.py @@ -73,6 +73,7 @@ class XYZGridScript(scripts_manager.Script): include_subgrids = gr.Checkbox(label='Include sub grids', value=False, elem_id=self.elem_id("include_sub_grids"), container=False) include_images = gr.Checkbox(label='Include images', value=False, elem_id=self.elem_id("include_lone_images"), container=False) create_video = gr.Checkbox(label='Create video', value=False, elem_id=self.elem_id("xyz_create_video"), container=False) + continue_on_error = gr.Checkbox(label='Continue on error', value=False, elem_id=self.elem_id("xyz_continue_on_error"), container=False) with gr.Row(visible=False) as ui_video: video_type, video_duration, video_loop, video_pad, video_interpolate = create_video_inputs(tab='img2img' if is_img2img else 'txt2img') @@ -167,6 +168,7 @@ class XYZGridScript(scripts_manager.Script): include_grid, include_subgrids, include_images, include_time, include_text, margin_size, create_video, video_type, video_duration, video_loop, video_pad, video_interpolate, + continue_on_error, ] def process(self, p, @@ -178,6 +180,7 @@ class XYZGridScript(scripts_manager.Script): include_grid, include_subgrids, include_images, include_time, include_text, margin_size, create_video, video_type, video_duration, video_loop, video_pad, video_interpolate, + continue_on_error = False, ): # pylint: disable=W0221 global active, xyz_results_cache # pylint: disable=W0603 xyz_results_cache = None @@ -337,8 +340,9 @@ class XYZGridScript(scripts_manager.Script): def cell(x, y, z, ix, iy, iz): if shared.state.interrupted: - log.warning('XYZ grid: Interrupted') - return processing.Processed(p, [], p.seed, ""), 0 + if not continue_on_error: + log.warning('XYZ grid: Interrupted') + return processing.Processed(p, [], p.seed, ""), 0 p.xyz = True pc = copy(p) pc.override_settings_restore_afterwards = False diff --git a/wiki b/wiki index 7526ee878..b8499e3e6 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 7526ee87821e8cf69e3a1e4e26f5ecc7d3508cee +Subproject commit b8499e3e6e74b7d03b80b8a02b5dca24432257cf