Make minimum steps configurable in sampler presets #662 #697 #483

Acly · May 13, 2024 · c8a709a · c8a709a
1 parent 4ab7b26
commit c8a709a
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/ai_diffusion/presets/samplers.json b/ai_diffusion/presets/samplers.json
@@ -3,50 +3,58 @@
         "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": {
         "sampler": "lcm",
         "scheduler": "sgm_uniform",
         "lora": "lcm",
         "steps": 6,
+        "minimum_steps": 4,
         "cfg": 1.8
     },
     "Realtime - Lightning": {
         "sampler": "euler",
         "scheduler": "sgm_uniform",
         "lora": "lightning",
         "steps": 8,
+        "minimum_steps": 4,
         "cfg": 1.5
     }
 }
diff --git 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
@@ -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