Merge pull request #50 from jkoelker/jk/scheduler

feat: add scheduler configuration for civitai checkpoints

Author: jkoelker

config.toml

@@ -143,6 +143,21 @@ verify_hashes = true
 # - SD 3, SD 3.5 Medium/Large (1024x1024)
 # - PixArt Alpha/Sigma, Kolors, Hunyuan DiT, AuraFlow
 #
+# SCHEDULER OPTIONS:
+# The scheduler controls the denoising process. Available options:
+#   - dpm++_karras: DPM++ with Karras sigmas (recommended for SD/SDXL/Pony)
+#   - dpm++: DPM++ without Karras sigmas
+#   - euler_a: Euler Ancestral (good variety, slightly random)
+#   - euler: Euler (deterministic)
+#   - heun: Heun (higher quality, 2x slower)
+#   - ddim: DDIM (classic, deterministic)
+#   - default: Keep the model's built-in scheduler
+#
+# Defaults by model type:
+#   - SD/SDXL/Pony/Illustrious: dpm++_karras
+#   - Turbo/Lightning/LCM/Hyper: default (keep optimized scheduler)
+#   - Flux/SD3/PixArt/etc: default (flow-based models)
+#
 # Example 1: CivitAI checkpoint by model ID (downloads automatically)
 # [models.civitai-example]
 # type = "civitai"
@@ -181,3 +196,9 @@ verify_hashes = true
 # steps = 30
 # guidance_scale = 5.0
 # supports_negative_prompt = true
+#
+# Example 6: CivitAI checkpoint with custom scheduler
+# [models.realistic-xl]
+# type = "civitai"
+# civitai_model_id = 12345
+# scheduler = "euler_a"           # Override default scheduler

src/oneiro/bot.py

@@ -12,7 +12,14 @@
 from oneiro.civitai import CivitaiClient, CivitaiError
 from oneiro.config import Config
 from oneiro.filters import ContentFilter
-from oneiro.pipelines import GenerationResult, LoraConfig, LoraSource, PipelineManager
+from oneiro.pipelines import (
+    SCHEDULER_CHOICES,
+    GenerationResult,
+    LoraConfig,
+    LoraSource,
+    PipelineManager,
+)
+from oneiro.pipelines.civitai_checkpoint import CivitaiCheckpointPipeline
 from oneiro.pipelines.lora import is_lora_compatible, parse_civitai_url
 from oneiro.queue import GenerationQueue, QueueStatus
 
@@ -349,6 +356,13 @@ async def on_raw_reaction_add(payload: discord.RawReactionActionEvent):
         required=False,
         autocomplete=get_lora_choices,
     )
+    @option(
+        "scheduler",
+        str,
+        description="Scheduler for denoising (dpm++_karras, euler_a, etc.)",
+        required=False,
+        choices=SCHEDULER_CHOICES,
+    )
     async def dream(
         ctx: discord.ApplicationContext,
         prompt: str,
@@ -359,6 +373,7 @@ async def dream(
         height: int = 1024,
         seed: int = -1,
         lora: str | None = None,
+        scheduler: str | None = None,
     ):
         """Generate an image from a text prompt."""
         global config, pipeline_manager, generation_queue, content_filter, civitai_client
@@ -486,7 +501,6 @@ async def dream(
                     await ctx.followup.send(f"❌ Invalid LoRA specification: {e}", ephemeral=True)
                     return
 
-        # Build generation request
         request: dict[str, Any] = {
             "prompt": prompt,
             "negative_prompt": negative_prompt,
@@ -497,7 +511,9 @@ async def dream(
             "guidance_scale": guidance_scale,
         }
 
-        # Add LoRA configs to request if any
+        if scheduler:
+            request["scheduler"] = scheduler
+
         if lora_configs:
             request["loras"] = lora_configs
 
@@ -566,6 +582,8 @@ async def on_complete(result: GenerationResult | Exception) -> None:
                 if len(lora_display) > 1024:
                     lora_display = lora_display[:1021] + "..."
                 embed.add_field(name="LoRA", value=lora_display, inline=True)
+            if scheduler:
+                embed.add_field(name="Scheduler", value=f"`{scheduler}`", inline=True)
             embed.set_image(url="attachment://dream.png")
             embed.set_footer(
                 text=f"Requested by {ctx.author.name} • React ❌ to delete",
@@ -642,11 +660,25 @@ async def queue_status(ctx: discord.ApplicationContext):
         required=True,
         autocomplete=get_model_choices,
     )
+    @option(
+        "scheduler",
+        str,
+        description="Override default scheduler for this model",
+        required=False,
+        choices=SCHEDULER_CHOICES,
+    )
     async def model_command(
         ctx: discord.ApplicationContext,
         model: str,
+        scheduler: str | None = None,
     ):
-        """Switch the active generation model."""
+        """Switch the active diffusion model and optionally override its scheduler.
+
+        This command changes which configured model is used for image generation.
+        If a scheduler is provided and supported by the loaded pipeline, it will be
+        applied after the model is loaded. The selected model may also be stored as
+        the default in the persistent configuration if a state path is configured.
+        """
         global config, pipeline_manager
 
         if pipeline_manager is None or config is None:
@@ -666,6 +698,22 @@ async def model_command(
 
         # Check if already loaded
         if pipeline_manager.current_model == model:
+            if scheduler and pipeline_manager.pipeline is not None:
+                if isinstance(pipeline_manager.pipeline, CivitaiCheckpointPipeline):
+                    pipeline_manager.pipeline.configure_scheduler(scheduler)
+                    await ctx.respond(
+                        f"✅ Model `{model}` already active, scheduler set to `{scheduler}`.",
+                        ephemeral=True,
+                    )
+                    return
+                else:
+                    await ctx.respond(
+                        f"✅ Model `{model}` is already active.\n"
+                        f"⚠️ Scheduler override is not supported for this pipeline type.",
+                        ephemeral=True,
+                    )
+                    return
+
             await ctx.respond(
                 f"✅ Model `{model}` is already active.",
                 ephemeral=True,
@@ -679,11 +727,17 @@ async def model_command(
             loading_msg = await ctx.followup.send(f"⏳ Loading model `{model}`...")
             await pipeline_manager.load_model(model)
 
-            # Persist model choice for next restart
+            if scheduler and pipeline_manager.pipeline is not None:
+                if isinstance(pipeline_manager.pipeline, CivitaiCheckpointPipeline):
+                    pipeline_manager.pipeline.configure_scheduler(scheduler)
+
             if config.state_path:
                 config.set("defaults", "model", value=model)
 
-            await loading_msg.edit(content=f"✅ Switched to model `{model}`")
+            msg = f"✅ Switched to model `{model}`"
+            if scheduler:
+                msg += f" with scheduler `{scheduler}`"
+            await loading_msg.edit(content=msg)
         except Exception as e:
             await ctx.followup.send(f"❌ Failed to load model: {e}", ephemeral=True)
 

src/oneiro/pipelines/__init__.py

@@ -9,6 +9,8 @@
 from oneiro.pipelines.base import BasePipeline, GenerationResult
 from oneiro.pipelines.civitai_checkpoint import (
     CIVITAI_BASE_MODEL_PIPELINE_MAP,
+    SCHEDULER_CHOICES,
+    SCHEDULER_MAP,
     CivitaiCheckpointPipeline,
     PipelineConfig,
     get_pipeline_config_for_base_model,
@@ -42,6 +44,8 @@
     "CivitaiCheckpointPipeline",
     "PipelineConfig",
     "CIVITAI_BASE_MODEL_PIPELINE_MAP",
+    "SCHEDULER_CHOICES",
+    "SCHEDULER_MAP",
     "get_pipeline_config_for_base_model",
     "LoraConfig",
     "LoraSource",