feat: Align Your Step Scheduler for WebUI

2024-04-25 02:36:19 +08:00
parent f0b8729958
commit 88f2311337
3 changed files with 126 additions and 3 deletions
@@ -55,8 +55,13 @@ def sample_dpmpp_2m_alt(model, x, sigmas, extra_args=None, callback=None, disabl


 def add_sample_dpmpp_2m_alt_webui() -> None:
+    """Adds DPM-Solver++(2M) alternative sampler to the list of available samplers."""
    try:
-        from modules import sd_samplers, sd_samplers_common, sd_samplers_kdiffusion  # type: ignore
+        from modules import (  # type: ignore
+            sd_samplers,
+            sd_samplers_common,
+            sd_samplers_kdiffusion,
+        )
    except ImportError:
        return

@@ -0,0 +1,110 @@
+# Copyright 2024 SLAPaper
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+
+import numpy as np
+import torch
+
+
+def loglinear_interp(t_steps: list[float], num_steps: int) -> np.ndarray:
+    """
+    Performs log-linear interpolation of a given array of decreasing numbers.
+    """
+    xs = np.linspace(0, 1, len(t_steps))
+    ys = np.log(t_steps[::-1])
+
+    new_xs = np.linspace(0, 1, num_steps)
+    new_ys = np.interp(new_xs, xs, ys)
+
+    interped_ys = np.exp(new_ys)[::-1].copy()
+    return interped_ys
+
+
+def align_your_step_scheduler_v15(
+    n: int, sigma_min: float, sigma_max: float, device
+) -> torch.Tensor:
+    """SD15 AYS scheduler from https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html"""
+    NOISE_LEVELS = [
+        14.615,
+        6.475,
+        3.861,
+        2.697,
+        1.886,
+        1.396,
+        0.963,
+        0.652,
+        0.399,
+        0.152,
+        0.029,
+    ]
+    TIMESTEP_INDICES = [999, 850, 736, 645, 545, 455, 343, 233, 124, 24, 0]
+    sigs = [sigma for sigma in loglinear_interp(NOISE_LEVELS, n)]
+    sigs.append(0.0)
+
+    logging.info(f"AYS scheduler: {sigs=}")
+    return torch.FloatTensor(sigs).to(device)
+
+
+def align_your_step_scheduler_xl(
+    n: int, sigma_min: float, sigma_max: float, device
+) -> torch.Tensor:
+    """SDXL AYS scheduler from https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html"""
+    NOISE_LEVELS = [
+        14.615,
+        6.315,
+        3.771,
+        2.181,
+        1.342,
+        0.862,
+        0.555,
+        0.380,
+        0.234,
+        0.113,
+        0.029,
+    ]
+    TIMESTEP_INDICES = [999, 845, 730, 587, 443, 310, 193, 116, 53, 13, 0]
+    sigs = [sigma for sigma in loglinear_interp(NOISE_LEVELS, n)]
+    sigs.append(0.0)
+
+    logging.info(f"AYS scheduler: {sigs=}")
+    return torch.FloatTensor(sigs).to(device)
+
+
+def add_align_your_step_scheduler() -> None:
+    """Add AYS scheduler to the list of schedulers"""
+    try:
+        from modules import sd_schedulers  # type: ignore
+    except ImportError:
+        return
+
+    scheduler_v15 = sd_schedulers.Scheduler(
+        "ays_v15", "Align Your Step SD15", align_your_step_scheduler_v15
+    )
+    scheduler_xl = sd_schedulers.Scheduler(
+        "ays_xl", "Align Your Step SDXL", align_your_step_scheduler_xl
+    )
+
+    sd_schedulers.schedulers.append(scheduler_v15)
+    sd_schedulers.schedulers.append(scheduler_xl)
+
+    sd_schedulers.schedulers_map[scheduler_v15.name] = scheduler_v15
+    sd_schedulers.schedulers_map[scheduler_v15.label] = scheduler_v15
+
+    sd_schedulers.schedulers_map[scheduler_xl.name] = scheduler_xl
+    sd_schedulers.schedulers_map[scheduler_xl.label] = scheduler_xl
+
+
+add_align_your_step_scheduler()