Merge pull request #52 from jkoelker/feat/issue-51-expose-generation-params

feat: expose steps and guidance_scale in /dream and /model commands

Author: jkoelker

src/oneiro/bot.py

@@ -34,6 +34,14 @@
 MIN_LORA_WEIGHT = -2.0
 MAX_LORA_WEIGHT = 2.0
 
+# Steps validation limits (for /dream and /model commands)
+MIN_STEPS = 1
+MAX_STEPS = 100
+
+# Guidance scale validation limits (for /dream and /model commands)
+MIN_GUIDANCE_SCALE = 0.0
+MAX_GUIDANCE_SCALE = 15.0
+
 
 def validate_lora_weight(weight: float, lora_name: str) -> None:
     """Validate that a LoRA weight is within acceptable bounds.
@@ -346,6 +354,22 @@ async def on_raw_reaction_add(payload: discord.RawReactionActionEvent):
         description="Random seed (-1 for random)",
         required=False,
     )
+    @option(
+        "steps",
+        int,
+        description="Number of inference steps (default: model-specific)",
+        required=False,
+        min_value=MIN_STEPS,
+        max_value=MAX_STEPS,
+    )
+    @option(
+        "guidance_scale",
+        float,
+        description="CFG scale - prompt adherence (default: model-specific)",
+        required=False,
+        min_value=MIN_GUIDANCE_SCALE,
+        max_value=MAX_GUIDANCE_SCALE,
+    )
     @option(
         "lora",
         str,
@@ -372,6 +396,8 @@ async def dream(
         width: int = 1024,
         height: int = 1024,
         seed: int = -1,
+        steps: int | None = None,
+        guidance_scale: float | None = None,
         lora: str | None = None,
         scheduler: str | None = None,
     ):
@@ -408,12 +434,26 @@ async def dream(
         current_model = pipeline_manager.current_model or "zimage-turbo"
         model_config = config.get("models", current_model, default={}) if config else {}
         pipeline_type = model_config.get("type") if model_config else None
-        steps = model_config.get("steps", 9)
-        guidance_scale = model_config.get("guidance_scale", 0.0)
+
+        # Get model config defaults
+        model_steps = model_config.get("steps", 9)
+        model_guidance = model_config.get("guidance_scale", 0.0)
 
         # Handle Qwen's true_cfg_scale
         if model_config.get("true_cfg_scale"):
-            guidance_scale = model_config.get("true_cfg_scale", 4.0)
+            model_guidance = model_config["true_cfg_scale"]
+
+        # Check for model-specific overrides set via /model command
+        model_overrides = config.get("model_overrides", current_model, default={}) if config else {}
+        if model_overrides:
+            if "steps" in model_overrides:
+                model_steps = model_overrides["steps"]
+            if "guidance_scale" in model_overrides:
+                model_guidance = model_overrides["guidance_scale"]
+
+        # User-provided values take priority over model defaults
+        actual_steps = steps if steps is not None else model_steps
+        actual_guidance = guidance_scale if guidance_scale is not None else model_guidance
 
         # Resolve LoRAs if specified
         lora_configs: list[LoraConfig] = []
@@ -507,8 +547,8 @@ async def dream(
             "width": width,
             "height": height,
             "seed": seed,
-            "steps": steps,
-            "guidance_scale": guidance_scale,
+            "steps": actual_steps,
+            "guidance_scale": actual_guidance,
         }
 
         if scheduler:
@@ -575,6 +615,8 @@ async def on_complete(result: GenerationResult | Exception) -> None:
             embed.add_field(name="Seed", value=str(result.seed), inline=True)
             embed.add_field(name="Time", value=f"{elapsed:.1f}s", inline=True)
             embed.add_field(name="Model", value=f"`{current_model}`", inline=True)
+            embed.add_field(name="Steps", value=str(result.steps), inline=True)
+            embed.add_field(name="CFG", value=f"{result.guidance_scale:.1f}", inline=True)
             if is_img2img:
                 embed.add_field(name="Strength", value=f"{strength:.2f}", inline=True)
             if lora_configs:
@@ -667,10 +709,28 @@ async def queue_status(ctx: discord.ApplicationContext):
         required=False,
         choices=SCHEDULER_CHOICES,
     )
+    @option(
+        "steps",
+        int,
+        description="Override default steps for this model",
+        required=False,
+        min_value=MIN_STEPS,
+        max_value=MAX_STEPS,
+    )
+    @option(
+        "guidance_scale",
+        float,
+        description="Override default CFG scale for this model",
+        required=False,
+        min_value=MIN_GUIDANCE_SCALE,
+        max_value=MAX_GUIDANCE_SCALE,
+    )
     async def model_command(
         ctx: discord.ApplicationContext,
         model: str,
         scheduler: str | None = None,
+        steps: int | None = None,
+        guidance_scale: float | None = None,
     ):
         """Switch the active diffusion model and optionally override its scheduler.
 
@@ -698,14 +758,13 @@ async def model_command(
 
         # Check if already loaded
         if pipeline_manager.current_model == model:
+            # Model is already active - handle overrides only
+            overrides_applied = []
+
             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
+                    overrides_applied.append(f"scheduler=`{scheduler}`")
                 else:
                     await ctx.respond(
                         f"✅ Model `{model}` is already active.\n"
@@ -714,10 +773,25 @@ async def model_command(
                     )
                     return
 
-            await ctx.respond(
-                f"✅ Model `{model}` is already active.",
-                ephemeral=True,
-            )
+            # Save steps/guidance_scale overrides to state
+            if config.state_path:
+                if steps is not None:
+                    config.set("model_overrides", model, "steps", value=steps)
+                    overrides_applied.append(f"steps={steps}")
+                if guidance_scale is not None:
+                    config.set("model_overrides", model, "guidance_scale", value=guidance_scale)
+                    overrides_applied.append(f"guidance_scale={guidance_scale}")
+
+            if overrides_applied:
+                await ctx.respond(
+                    f"✅ Model `{model}` already active. Set: {', '.join(overrides_applied)}",
+                    ephemeral=True,
+                )
+            else:
+                await ctx.respond(
+                    f"✅ Model `{model}` is already active.",
+                    ephemeral=True,
+                )
             return
 
         # Defer for model loading (can be slow)
@@ -733,10 +807,22 @@ async def model_command(
 
             if config.state_path:
                 config.set("defaults", "model", value=model)
+                # Save steps/guidance_scale overrides to state
+                if steps is not None:
+                    config.set("model_overrides", model, "steps", value=steps)
+                if guidance_scale is not None:
+                    config.set("model_overrides", model, "guidance_scale", value=guidance_scale)
 
             msg = f"✅ Switched to model `{model}`"
+            overrides = []
             if scheduler:
-                msg += f" with scheduler `{scheduler}`"
+                overrides.append(f"scheduler=`{scheduler}`")
+            if steps is not None:
+                overrides.append(f"steps={steps}")
+            if guidance_scale is not None:
+                overrides.append(f"guidance_scale={guidance_scale}")
+            if overrides:
+                msg += f" with {', '.join(overrides)}"
             await loading_msg.edit(content=msg)
         except Exception as e:
             await ctx.followup.send(f"❌ Failed to load model: {e}", ephemeral=True)