diff options
| -rw-r--r-- | .cspell/custom-dictionary-workspace.txt | 4 | ||||
| -rw-r--r-- | .vscode/settings.json | 10 | ||||
| -rw-r--r-- | comfier.py | 35 | ||||
| -rw-r--r-- | harness.py | 16 | ||||
| -rw-r--r-- | tests/test_generate.py | 23 |
5 files changed, 76 insertions, 12 deletions
diff --git a/.cspell/custom-dictionary-workspace.txt b/.cspell/custom-dictionary-workspace.txt new file mode 100644 index 0000000..636e006 --- /dev/null +++ b/.cspell/custom-dictionary-workspace.txt @@ -0,0 +1,4 @@ +# Custom Dictionary Words +COMFYUI +Morgner +webui diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..097d610 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,10 @@ +{ + "cSpell.customDictionaries": { + "custom-dictionary-workspace": { + "name": "custom-dictionary-workspace", + "path": "${workspaceFolder:comfier}/.cspell/custom-dictionary-workspace.txt", + "addWords": true, + "scope": "workspace" + } + } +}
\ No newline at end of file @@ -35,7 +35,7 @@ class Tools: Uses the workflow, connection settings and node mappings (prompt, seed, width, height, steps) configured in Open WebUI's image generation settings, and adds - negative prompts, per-call size and reproducible seeds on top. + negative prompts, per-call steps and reproducible seeds on top. """ class Valves(BaseModel): @@ -47,6 +47,11 @@ class Tools: default="negative_prompt", description="Name of the workflow node input that receives the negative prompt (e.g. 'negative_prompt' or 'text').", ) + max_steps: int = Field( + default=100, + ge=1, + description="Maximum number of sampling steps the model may request per image.", + ) negative_prompt_node_ids: str = Field( default="", description="Comma-separated IDs of the workflow nodes that receive the negative prompt (e.g. '459:452'). Leave empty to disable negative prompting.", @@ -61,6 +66,7 @@ class Tools: negative_prompt: Optional[str] = None, size: Optional[Annotated[str, Field(pattern=r"^\d+[xX]\d+$")]] = None, seed: Optional[int | str] = None, + steps: Optional[Annotated[int, Field(ge=1)]] = None, __request__: Optional[Request] = None, __user__: Optional[dict[str, Any]] = None, __event_emitter__: Optional[EventEmitter] = None, @@ -68,17 +74,18 @@ class Tools: ) -> dict[str, Any]: """ Generate an image from a text prompt with ComfyUI. - Optionally accepts a negative prompt (what to avoid), an image size and a seed for reproducible results. + Optionally accepts a negative prompt (what to avoid), an image size, a seed for reproducible results and a number of sampling steps. The generated image is automatically displayed to the user in the chat, so do not embed or link it again; just tell the user it has been created. - Generation may take a while. The result contains the seed that was used (random if none was given). - To reproduce an image, pass that seed again with the same prompt and size. + Generation may take a while. The result contains the seed and the number of steps that were used. + To reproduce an image, pass that seed again with the same prompt, size and steps. :param positive_prompt: Descriptive text prompt for the image to generate. :param negative_prompt: Optional text describing what to avoid in the image. If omitted, the administrator's default negative prompt is used; pass an empty string for none. :param size: Optional image size in pixels as 'WIDTHxHEIGHT' (e.g. '1024x768'). If omitted, the default size configured in Open WebUI is used. :param seed: Optional non-negative integer seed (or its decimal string) for reproducible generation. If omitted, a random seed is used. - :return: A dict with 'status', the 'seed' used, a 'message' for the model and the list of stored image URLs under 'images'. + :param steps: Optional number of sampling steps. More steps usually add detail but take proportionally longer. If omitted, the default configured in Open WebUI is used. Values above the administrator's limit are rejected. + :return: A dict with 'status', the 'seed' and 'steps' used, a 'message' for the model and the list of stored image URLs under 'images'. """ if __request__ is None: @@ -99,7 +106,7 @@ class Tools: else self.valves.negative_prompt ) width, height = await self._resolve_dimensions(size) - steps = await self._resolve_steps() + resolved_steps = await self._resolve_steps(steps) parameter_values: dict[str, Any] = { "positive_prompt": positive_prompt, @@ -107,7 +114,7 @@ class Tools: "seed": normalized_seed, "width": width, "height": height, - "steps": steps, + "steps": resolved_steps, } self._apply_workflow_parameters( @@ -132,7 +139,7 @@ class Tools: "seed": str(normalized_seed), "width": width, "height": height, - "steps": steps, + "steps": resolved_steps, }, event_emitter=__event_emitter__, ) @@ -178,6 +185,7 @@ class Tools: return { "status": "success", "seed": str(normalized_seed), + "steps": resolved_steps, "message": "The image has been successfully generated and is already visible to the user in the chat. You do not need to display or embed the image again - just acknowledge that it has been created.", "images": image_urls, } @@ -654,8 +662,15 @@ class Tools: return resolved - async def _resolve_steps(self) -> int: - """Load the global IMAGE_STEPS setting from Open WebUI config.""" + async def _resolve_steps(self, steps: Optional[int]) -> int: + """Validate the requested step count, or fall back to the global IMAGE_STEPS setting.""" + + if steps is not None: + if not 1 <= steps <= self.valves.max_steps: + raise ValueError( + f"Invalid steps {steps}. Must be between 1 and {self.valves.max_steps}." + ) + return steps image_cfg = await get_image_config() raw_steps = getattr(image_cfg, "IMAGE_STEPS", None) @@ -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, diff --git a/tests/test_generate.py b/tests/test_generate.py index 8eaf44b..0bb0b5e 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -129,3 +129,26 @@ async def test_stalled_backend_still_queued_keeps_waiting( await comfy.release() result, _ = await asyncio.wait_for(task, timeout=5) assert result["status"] == "success" + + +async def test_steps_requested_and_default( + make_tools: Callable[[], Any], comfy: FakeComfy +) -> None: + tools, _ = make_tools() + result, _ = await _run(tools, steps=12) + assert comfy.submitted[0]["459:458"]["inputs"]["steps"] == 12 + assert result["steps"] == 12 + + result, _ = await _run(tools) # falls back to the global setting (FakeEnv: 25) + assert comfy.submitted[1]["459:458"]["inputs"]["steps"] == 25 + assert result["steps"] == 25 + + +@pytest.mark.parametrize("steps", [0, -3, 101]) +async def test_steps_out_of_range_rejected( + make_tools: Callable[[], Any], comfy: FakeComfy, steps: int +) -> None: + tools, _ = make_tools() + with pytest.raises(ValueError, match="between 1 and 100"): + await _run(tools, steps=steps) + assert not comfy.submitted |