diff options
Diffstat (limited to 'comfier.py')
| -rw-r--r-- | comfier.py | 35 |
1 files changed, 25 insertions, 10 deletions
@@ -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) |