tool: add support for custom steps

1 files changed, 14 insertions, 2 deletions

diff --git a/harness.py b/harness.py
index 450df24..061ec02 100644
--- a/harness.py
+++ b/harness.py
@@ -58,7 +58,18 @@ async def main() -> None:
     parser.add_argument("--neg", default=None, help="Negative prompt")
     parser.add_argument("--size", default="256x256", help="Image size (WIDTHxHEIGHT)")
     parser.add_argument("--seed", default="123")
-    parser.add_argument("--steps", type=int, default=30, help="Global IMAGE_STEPS")
+    parser.add_argument(
+        "--steps",
+        type=int,
+        default=None,
+        help="Tool argument: sampling steps (default: omitted, use --default-steps)",
+    )
+    parser.add_argument(
+        "--default-steps",
+        type=int,
+        default=30,
+        help="Open WebUI's global IMAGE_STEPS, used when --steps is omitted",
+    )
     parser.add_argument(
         "--neg-nodes", default="459:452", help="Valve: negative_prompt_node_ids"
     )
@@ -78,7 +89,7 @@ async def main() -> None:
             workflow=args.workflow.read_text(encoding="utf-8"),
             nodes=json.loads(args.nodes.read_text(encoding="utf-8")),
             image_size="1024x1024",
-            image_steps=args.steps,
+            image_steps=args.default_steps,
             output_dir=args.out,
         )
     )
@@ -95,6 +106,7 @@ async def main() -> None:
             negative_prompt=args.neg,
             size=args.size,
             seed=args.seed,
+            steps=args.steps,
             __request__=cast(Any, SimpleNamespace()),
             __user__={"id": "local-dev"},
             __event_emitter__=console_event_emitter,