From c8a709a6d420137ce627d8ea1337ba441790912f Mon Sep 17 00:00:00 2001
From: Acly <aclysia@gmail.com>
Date: Mon, 13 May 2024 13:03:23 +0200
Subject: [PATCH] Make minimum steps configurable in sampler presets #662 #697
 #483

---
 ai_diffusion/presets/samplers.json | 8 ++++++++
 ai_diffusion/style.py              | 2 ++
 ai_diffusion/workflow.py           | 6 ++----
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/ai_diffusion/presets/samplers.json b/ai_diffusion/presets/samplers.json
index b1af399a4..008dc9b27 100644
--- a/ai_diffusion/presets/samplers.json
+++ b/ai_diffusion/presets/samplers.json
@@ -3,36 +3,42 @@
         "sampler": "dpmpp_2m",
         "scheduler": "karras",
         "steps": 20,
+        "minimum_steps": 4,
         "cfg": 7.0
     },
     "Alternative - Euler A": {
         "sampler": "euler_ancestral",
         "scheduler": "normal",
         "steps": 20,
+        "minimum_steps": 4,
         "cfg": 6.0
     },
     "Creative - DPM++ 2M SDE": {
         "sampler": "dpmpp_2m_sde_gpu",
         "scheduler": "karras",
         "steps": 20,
+        "minimum_steps": 4,
         "cfg": 6.0
     },
     "Fast - UniPC BH2": {
         "sampler": "uni_pc_bh2",
         "scheduler": "ddim_uniform",
         "steps": 12,
+        "minimum_steps": 4,
         "cfg": 4.0
     },
     "Turbo/Lightning Merge - DPM++ SDE": {
         "sampler": "dpmpp_sde_gpu",
         "scheduler": "karras",
         "steps": 6,
+        "minimum_steps": 4,
         "cfg": 2.0
     },
     "Lightning Merge - Euler A Uniform": {
         "sampler": "euler_ancestral",
         "scheduler": "sgm_uniform",
         "steps": 8,
+        "minimum_steps": 4,
         "cfg": 1.5
     },
     "Realtime - LCM": {
@@ -40,6 +46,7 @@
         "scheduler": "sgm_uniform",
         "lora": "lcm",
         "steps": 6,
+        "minimum_steps": 4,
         "cfg": 1.8
     },
     "Realtime - Lightning": {
@@ -47,6 +54,7 @@
         "scheduler": "sgm_uniform",
         "lora": "lightning",
         "steps": 8,
+        "minimum_steps": 4,
         "cfg": 1.5
     }
 }
\ No newline at end of file
diff --git a/ai_diffusion/style.py b/ai_diffusion/style.py
index be0674485..23a2f4218 100644
--- a/ai_diffusion/style.py
+++ b/ai_diffusion/style.py
@@ -292,6 +292,7 @@ class SamplerPreset(NamedTuple):
     steps: int
     cfg: float
     lora: str | None = None
+    minimum_steps: int = 4
 
 
 class SamplerPresets:
@@ -421,6 +422,7 @@ def names(self):
         "sampler": "dpmpp_3m_sde",
         "scheduler": "exponential",
         "steps": 20,
+        "minimum_steps": 4,
         "cfg": 7.0
     }
 }
diff --git a/ai_diffusion/workflow.py b/ai_diffusion/workflow.py
index 8f30bdd23..8ec66e80a 100644
--- a/ai_diffusion/workflow.py
+++ b/ai_diffusion/workflow.py
@@ -42,9 +42,7 @@ def _sampling_from_style(style: Style, strength: float, is_live: bool):
         total_steps=total_steps or preset.steps,
     )
     if strength < 1.0:
-        # Unless we have something like a 1-step turbo model, ensure there are at least 4 steps
-        # even at very low strength with low total steps of 4-8 (like Lightning/LCM).
-        min_steps = min(4, total_steps)
+        min_steps = min(preset.minimum_steps, total_steps)
         result.total_steps, result.start_step = _apply_strength(strength, total_steps, min_steps)
     return result
 
@@ -53,7 +51,7 @@ def _apply_strength(strength: float, steps: int, min_steps: int = 0) -> tuple[in
     start_at_step = round(steps * (1 - strength))
 
     if min_steps and steps - start_at_step < min_steps:
-        steps = math.floor(min_steps * 1 / strength)
+        steps = math.floor(min_steps / strength)
         start_at_step = steps - min_steps
 
     return steps, start_at_step