Skip to content

Developing custom workflow types

ljleb edited this page Sep 21, 2023 · 12 revisions

It is possible to register new workflow types from another extension. To do this, you can import the external_code module from lib_comfyui:

# file extensions/sd-webui-example/scripts/my_extension_script.py:

import importlib

lib_comfyui = importlib.import_module("extensions.sd-webui-comfyui.lib_comfyui.external_code", "external_code")

# new workflows to display in the txt2img and/or img2img tabs
additional_workflow_types = [
    lib_comfyui.WorkflowType(
        base_id="example_workflow",
        display_name="Example Workflow",
        types="LATENT",

        # Remove either "txt2img" or "img2img" if your workflow should only be added to a particular tab
        # The default is both tabs:
        #tabs=("txt2img", "img2img"),

        # The default json-serialized workflow to load on startup. If user presses CTRL+D in comfyui, it will default to this workflow as well
        # You can also pass a `pathlib.Path` object, the constructor will read the contents of the file right away
        # The default is `"null"`, i.e. do not replace the existing default workflow in comfyui
        #default_workflow = "null",
    ),
    ...
]

# register the workflows with the extension
for workflow_type in additional_workflow_types:
    lib_comfyui.add_workflow_type(workflow_type)

Note that add_workflow_type() must be called before the gradio interface is created. Otherwise you will get an error.

This will register new workflows in the gradio interface. Users will have access to a new interface, which they will be able to modify to their taste:

image

For more examples of workflow type registration code, have a look at the registration code for the sd-webui-comfyui extension: https://github.com/ModelSurge/sd-webui-comfyui/blob/main/lib_comfyui/default_workflow_types.py

To launch the workflow at any time after the gradio user interface has been created (for example, during the generation process), you can use the external_code.run_workflow() function:

# file extensions/sd-webui-example/scripts/my_extension_script.py:

import importlib

lib_comfyui = importlib.import_module("extensions.sd-webui-comfyui.lib_comfyui.external_code", "external_code")

...  # other imports, our example workflow type registration code above, and possibly other code of your liking

example_workflow = additional_workflow_types[0]  # our example custom WorkflowType object created above

class ExampleScript(scripts.Script):
    ...  # other script methods

    def postprocess_batch(p: processing.StableDiffusionProcessing, *args, images, **kwargs):
        # run the workflow and update the batch images with the result
        # since workflows can have multiple output nodes, `run_workflow()` returns a list of batches: one per output node.
        images[:] = lib_comfyui.run_workflow(
            workflow_type=example_workflow,
            tab="txt2img" if isinstance(p, processing.StableDiffusionProcessingTxt2Img) else "img2img",
            batch_input=images,
        )[0]

For more examples, see the postprocess_batch_list implementation here:
https://github.com/ModelSurge/sd-webui-comfyui/blob/main/scripts/comfyui.py

And the other predefined workflow hijacks here:
https://github.com/ModelSurge/sd-webui-comfyui/blob/main/lib_comfyui/webui/patches.py