├── .github └── workflows │ └── tests.yml ├── .gitignore ├── LICENSE ├── README.md ├── comfyui_custom_nodes ├── webui_io.py ├── webui_proxy_nodes.py └── webui_save_image.py ├── comfyui_custom_scripts └── extensions │ ├── webuiEvents.js │ ├── webuiNodes.js │ ├── webuiPatches.js │ ├── webuiRequests.js │ └── webuiTypes.js ├── install.py ├── install_comfyui.py ├── javascript └── compute-height.js ├── lib_comfyui ├── argv_conversion.py ├── comfyui │ ├── iframe_requests.py │ ├── pre_main.py │ ├── queue_tracker.py │ ├── routes_extension.py │ └── type_conversion.py ├── comfyui_process.py ├── custom_extension_injector.py ├── default_workflow_types.py ├── external_code │ ├── __init__.py │ └── api.py ├── find_extensions.py ├── global_state.py ├── ipc │ ├── __init__.py │ ├── callback.py │ ├── payload.py │ └── strategies.py ├── platform_utils.py ├── torch_utils.py └── webui │ ├── accordion.py │ ├── callbacks.py │ ├── gradio_utils.py │ ├── patches.py │ ├── paths.py │ ├── proxies.py │ ├── reverse_proxy.py │ ├── settings.py │ └── tab.py ├── preload.py ├── requirements-test.txt ├── requirements.txt ├── resources └── front-page.gif ├── scripts └── comfyui.py ├── style.css └── tests ├── __init__.py ├── lib_comfyui_tests ├── __init__.py ├── argv_conversion_test.py ├── external_code_tests │ ├── __init__.py │ ├── import_test.py │ └── run_workflow_test.py └── ipc_tests │ ├── __init__.py │ └── payload_test.py └── utils.py /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: Run basic features tests on CPU with empty SD model 2 | 3 | on: 4 | - push 5 | - pull_request 6 | 7 | jobs: 8 | test: 9 | runs-on: ubuntu-latest 10 | steps: 11 | - name: Checkout sd-webui-comfyui 12 | uses: actions/checkout@v3 13 | 14 | - name: Checkout stable-diffusion-webui 15 | uses: actions/checkout@v3 16 | with: 17 | repository: 'AUTOMATIC1111/stable-diffusion-webui' 18 | path: 'stable-diffusion-webui' 19 | 20 | - name: Set up Python 3.10 21 | uses: actions/setup-python@v4 22 | with: 23 | python-version: 3.10.6 24 | cache: pip 25 | cache-dependency-path: | 26 | **/requirements*txt 27 | stable-diffusion-webui/launch.py 28 | 29 | - name: Install test dependencies 30 | run: | 31 | pip install -r requirements.txt -r requirements-test.txt 32 | env: 33 | PIP_DISABLE_PIP_VERSION_CHECK: "1" 34 | PIP_PROGRESS_BAR: "off" 35 | 36 | - name: Setup environment 37 | run: | 38 | cd stable-diffusion-webui 39 | python launch.py --skip-torch-cuda-test --exit 40 | env: 41 | PIP_DISABLE_PIP_VERSION_CHECK: "1" 42 | PIP_PROGRESS_BAR: "off" 43 | TORCH_INDEX_URL: https://download.pytorch.org/whl/cpu 44 | WEBUI_LAUNCH_LIVE_OUTPUT: "1" 45 | PYTHONUNBUFFERED: "1" 46 | 47 | - name: Run tests 48 | run: | 49 | python -m pytest -vv --junitxml=tests/results.xml tests 50 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | .idea/ 132 | 133 | /ComfyUI/ 134 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 ModelSurge 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # sd-webui-comfyui 2 | ## Overview 3 | sd-webui-comfyui is an extension for [A1111 webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) that embeds [ComfyUI](https://github.com/comfyanonymous/ComfyUI) workflows in different sections of the normal pipeline of the webui. This allows to create ComfyUI nodes that interact directly with some parts of the webui's normal pipeline. 4 | 5 | ![front-page-gif](/resources/front-page.gif) 6 | 7 | ## Features 8 | - Use ComfyUI directly into the Webui 9 | - Support for [loading custom nodes from other Webui extensions](https://github.com/ModelSurge/sd-webui-comfyui/wiki/Developing-custom-nodes-from-webui-extensions) 10 | - Integration of [ComfyUI workflows](https://github.com/ModelSurge/sd-webui-comfyui/wiki/Developing-custom-workflow-types) directly into the Webui's pipeline, such as `preprocess`, `preprocess (latent)`, `unet`, `postprocess (latent)`, `postprocess`, `transformer text encode`, etc. 11 | - Webui nodes for sharing resources and data, such as the model, the prompt, etc. 12 | 13 | For a full overview of all the advantageous features this extension adds to ComfyUI and to the Webui, check out the [wiki page](https://github.com/ModelSurge/sd-webui-comfyui/wiki). 14 | 15 | ## Officially supported versions 16 | - A1111 Webui >= `1.5.1` 17 | - ComfyUI == `latest` 18 | 19 | ## Installation 20 | 1) Go to Extensions > Available 21 | 2) Click the `Load from:` button 22 | 3) Enter "ComfyUI" in the search bar 23 | 4) Click the `Install` button of the ComfyUI Tab cell 24 | 5) Restart the webui 25 | 6) Go to the `ComfyUI` tab, and follow the instructions 26 | 27 | ## Remote users, reverse proxies, etc. 28 | The extension is now able to load comfyui for remote users using a local reverse proxy. 29 | This is necessary when the webui is started remotely, for example when: 30 | - using the command line arguments `--share`, or `--ngrok` 31 | - using reverse proxy options of the [sd-webui-tunnels](https://github.com/Bing-su/sd-webui-tunnels) extension 32 | 33 | If you want the extension to keep the reverse proxy disabled or always enable it for some reason, you can update your preferences in the settings tab. 34 | 35 | To start the reverse proxy, the extension needs the command line argument `--api` for the webui, which starts a fastapi server. 36 | Without fastapi, the extension will not be able to create a reverse proxy for comfyui, and then remote browsers will not be able to load comfyui iframes. 37 | 38 | In practice, if the webui url is `http://localhost:7860`, then the extension effectively creates two reverse proxies: 39 | - An HTTP reverse proxy at POST, GET, PUT and DELETE http://localhost:7860/sd-webui-comfyui/comfyui 40 | - A websockets reverse proxy at ws://localhost:7860/sd-webui-comfyui/comfyui/ws 41 | 42 | ## Contributing 43 | We welcome contributions from anyone who is interested in improving sd-webui-comfyui. If you would like to contribute, please follow these steps: 44 | 45 | 1) Fork the repository and create a new branch for your feature or bug fix. 46 | 2) Implement your changes, adding any necessary documentation and tests. 47 | 3) Submit a pull request. 48 | 4) We will review your contribution as soon as possible and provide feedback. 49 | 50 | ## License 51 | MIT 52 | 53 | ## Contact 54 | If you have any questions or concerns, please leave an issue, or start a thread in the discussions. 55 | 56 | Thank you for your interest! 57 | -------------------------------------------------------------------------------- /comfyui_custom_nodes/webui_io.py: -------------------------------------------------------------------------------- 1 | from lib_comfyui import global_state 2 | 3 | 4 | class StaticProperty(object): 5 | def __init__(self, f): 6 | self.f = f 7 | 8 | def __get__(self, *args): 9 | return self.f() 10 | 11 | 12 | class FromWebui: 13 | @classmethod 14 | def INPUT_TYPES(cls): 15 | return { 16 | "required": { 17 | "void": ("VOID", ), 18 | }, 19 | } 20 | 21 | @StaticProperty 22 | def RETURN_TYPES(): 23 | return getattr(global_state, "current_workflow_input_types", ()) 24 | 25 | RETURN_NAMES = () 26 | FUNCTION = "get_node_inputs" 27 | 28 | CATEGORY = "webui" 29 | 30 | @staticmethod 31 | def get_node_inputs(void): 32 | return global_state.node_inputs 33 | 34 | 35 | class ToWebui: 36 | @classmethod 37 | def INPUT_TYPES(cls): 38 | return { 39 | "required": {}, 40 | } 41 | RETURN_TYPES = () 42 | FUNCTION = "extend_node_outputs" 43 | 44 | CATEGORY = "webui" 45 | 46 | OUTPUT_NODE = True 47 | 48 | @staticmethod 49 | def extend_node_outputs(**outputs): 50 | global_state.node_outputs += [outputs] 51 | return () 52 | 53 | 54 | NODE_CLASS_MAPPINGS = { 55 | "FromWebui": FromWebui, 56 | "ToWebui": ToWebui, 57 | } 58 | 59 | NODE_DISPLAY_NAME_MAPPINGS = { 60 | "FromWebui": "From Webui", 61 | "ToWebui": "To Webui", 62 | } 63 | -------------------------------------------------------------------------------- /comfyui_custom_nodes/webui_proxy_nodes.py: -------------------------------------------------------------------------------- 1 | from lib_comfyui.webui import proxies 2 | from lib_comfyui import global_state 3 | 4 | 5 | class WebuiCheckpointLoader: 6 | @classmethod 7 | def INPUT_TYPES(s): 8 | return { 9 | "required": { 10 | "void": ("VOID", ), 11 | }, 12 | } 13 | RETURN_TYPES = ("MODEL", "CLIP", "VAE") 14 | FUNCTION = "load_checkpoint" 15 | 16 | CATEGORY = "loaders" 17 | 18 | def load_checkpoint(self, void): 19 | config = proxies.get_comfy_model_config() 20 | proxies.raise_on_unsupported_model_type(config) 21 | return ( 22 | proxies.ModelPatcher(proxies.Model()), 23 | proxies.ClipWrapper(proxies.Clip()), 24 | proxies.VaeWrapper(proxies.Vae()), 25 | ) 26 | 27 | 28 | class WebuiPrompts: 29 | @classmethod 30 | def INPUT_TYPES(s): 31 | return { 32 | "required": { 33 | "void": ("VOID", ), 34 | }, 35 | } 36 | RETURN_TYPES = ("STRING", "STRING") 37 | RETURN_NAMES = ("positive", "negative") 38 | FUNCTION = "get_prompts" 39 | 40 | CATEGORY = "text" 41 | 42 | def get_prompts(self, void): 43 | positive_prompts, _extra_networks = proxies.extra_networks_parse_prompts([getattr(global_state, 'last_positive_prompt', '')]) 44 | 45 | return ( 46 | positive_prompts[0], 47 | getattr(global_state, 'last_negative_prompt', ''), 48 | ) 49 | 50 | 51 | NODE_CLASS_MAPPINGS = { 52 | "WebuiCheckpointLoader": WebuiCheckpointLoader, 53 | "WebuiPrompts": WebuiPrompts, 54 | } 55 | 56 | NODE_DISPLAY_NAME_MAPPINGS = { 57 | "WebuiCheckpointLoader": 'Webui Checkpoint', 58 | "WebuiPrompts": "Webui Prompts", 59 | } 60 | -------------------------------------------------------------------------------- /comfyui_custom_nodes/webui_save_image.py: -------------------------------------------------------------------------------- 1 | from lib_comfyui.webui.settings import opts 2 | from torchvision.transforms.functional import to_pil_image 3 | from lib_comfyui.webui.paths import webui_save_image 4 | 5 | 6 | class WebuiSaveImage: 7 | @classmethod 8 | def INPUT_TYPES(cls): 9 | return { 10 | "required": { 11 | "location": (["txt2img-images", "img2img-images", "extras-images", "txt2img-grids", "img2img-grids", ], ), 12 | "images": ("IMAGE", ), 13 | }, 14 | } 15 | RETURN_TYPES = () 16 | FUNCTION = "save_image" 17 | 18 | CATEGORY = "image" 19 | 20 | OUTPUT_NODE = True 21 | 22 | def save_image(self, location, images): 23 | if 'txt2img' in location: 24 | output_dir = opts.outdir_samples or opts.outdir_txt2img_samples if 'images' in location else opts.outdir_grids or opts.outdir_txt2img_grids 25 | elif 'img2img' in location: 26 | output_dir = opts.outdir_samples or opts.outdir_img2img_samples if 'images' in location else opts.outdir_grids or opts.outdir_img2img_grids 27 | else: 28 | output_dir = opts.outdir_samples or opts.outdir_extras_samples 29 | 30 | for image in images: 31 | pil_image = to_pil_image(image.permute(2, 0, 1)) 32 | filename, _ = webui_save_image(image=pil_image, relative_path=output_dir, basename='') 33 | 34 | return [] 35 | 36 | 37 | NODE_CLASS_MAPPINGS = { 38 | "WebuiSaveImage": WebuiSaveImage, 39 | } 40 | 41 | NODE_DISPLAY_NAME_MAPPINGS = { 42 | "WebuiSaveImage": 'Webui Save Image', 43 | } 44 | -------------------------------------------------------------------------------- /comfyui_custom_scripts/extensions/webuiEvents.js: -------------------------------------------------------------------------------- 1 | import { app } from "/scripts/app.js"; 2 | import { api } from "/scripts/api.js"; 3 | 4 | 5 | const POLLING_TIMEOUT = 500; 6 | 7 | 8 | const iframeRegisteredEvent = new Promise(async (resolve, reject) => { 9 | const searchParams = new URLSearchParams(window.location.search); 10 | const workflowTypeId = searchParams.get("workflowTypeId"); 11 | const webuiClientId = searchParams.get("webuiClientId"); 12 | 13 | if (!workflowTypeId || !webuiClientId) { 14 | reject("Cannot identify comfyui client: search params missing."); 15 | } 16 | else { 17 | const workflowTypeInfo = await fetchWorkflowTypeInfo(workflowTypeId); 18 | resolve({ 19 | workflowTypeId, 20 | webuiClientId, 21 | ...workflowTypeInfo, 22 | }); 23 | } 24 | }); 25 | 26 | async function fetchWorkflowTypeInfo(workflowTypeId) { 27 | const response = await api.fetchApi("/sd-webui-comfyui/workflow_type?" + new URLSearchParams({ 28 | workflowTypeId, 29 | }), { 30 | method: "GET", 31 | headers: {"Content-Type": "application/json"}, 32 | cache: "no-store", 33 | }); 34 | return await response.json(); 35 | } 36 | 37 | const appReadyEvent = new Promise(resolve => { 38 | const appReadyOrRecursiveSetTimeout = () => { 39 | if (app.graph && window.name) { 40 | resolve(); 41 | } 42 | else { 43 | setTimeout(appReadyOrRecursiveSetTimeout, POLLING_TIMEOUT); 44 | } 45 | }; 46 | appReadyOrRecursiveSetTimeout(); 47 | }); 48 | 49 | 50 | export { 51 | iframeRegisteredEvent, 52 | appReadyEvent, 53 | } 54 | -------------------------------------------------------------------------------- /comfyui_custom_scripts/extensions/webuiNodes.js: -------------------------------------------------------------------------------- 1 | import { app } from "/scripts/app.js"; 2 | import { iframeRegisteredEvent } from "/extensions/webui_scripts/sd-webui-comfyui/extensions/webuiEvents.js"; 3 | import { isString, getTypesLength } from "/extensions/webui_scripts/sd-webui-comfyui/extensions/webuiTypes.js"; 4 | 5 | 6 | function createVoidWidget(node, name) { 7 | const widget = { 8 | type: "customtext", 9 | name, 10 | get value() { 11 | return `${Math.random()}${Date.now()}`; 12 | }, 13 | set value(x) {}, 14 | }; 15 | widget.parent = node; 16 | node.addCustomWidget(widget); 17 | 18 | return widget; 19 | } 20 | 21 | app.registerExtension({ 22 | name: "sd-webui-comfyui", 23 | async getCustomWidgets(app) { 24 | return { 25 | VOID(node, inputName) { 26 | createVoidWidget(node, inputName); 27 | }, 28 | }; 29 | }, 30 | async addCustomNodeDefs(defs) { 31 | let iframeInfo = null; 32 | 33 | try { 34 | iframeInfo = await iframeRegisteredEvent; 35 | } 36 | catch { 37 | return; 38 | } 39 | 40 | const nodes = webuiIoNodeNames.map(name => defs[name]); 41 | for (const node of nodes) { 42 | node.display_name = `${node.display_name} - ${iframeInfo.displayName}`; 43 | 44 | if (node.name === 'FromWebui') { 45 | let outputs = iframeInfo.webuiIoTypes.outputs; 46 | if (isString(outputs)) { 47 | outputs = [outputs]; 48 | } 49 | const are_types_array = Array.isArray(outputs); 50 | for (const k in outputs) { 51 | const v = outputs[k]; 52 | node.output_name.push(are_types_array ? v : k); 53 | node.output_is_list.push(false); 54 | node.output.push(v); 55 | } 56 | } 57 | else if (node.name === 'ToWebui') { 58 | let inputs = iframeInfo.webuiIoTypes.inputs; 59 | if (isString(inputs)) { 60 | node.input.required[inputs] = [inputs]; 61 | } 62 | else { 63 | for (const k in inputs) { 64 | const v = inputs[k]; 65 | node.input.required[k] = [v]; 66 | } 67 | } 68 | } 69 | } 70 | }, 71 | async nodeCreated(node) { 72 | let iframeInfo = null; 73 | 74 | try { 75 | iframeInfo = await iframeRegisteredEvent; 76 | } catch { 77 | return; 78 | } 79 | 80 | if (!webuiIoNodeNames.includes(node.type)) { 81 | return; 82 | } 83 | 84 | const maxIoLength = Math.max( 85 | getTypesLength(iframeInfo.webuiIoTypes.outputs), 86 | getTypesLength(iframeInfo.webuiIoTypes.inputs), 87 | ); 88 | // 260 and 40 are empirical values that seem to work 89 | node.size = [260, 40 + distanceBetweenIoSlots * maxIoLength]; 90 | }, 91 | async setup() { 92 | app.loadGraphData(); 93 | }, 94 | }); 95 | 96 | const webuiIoNodeNames = [ 97 | 'FromWebui', 98 | 'ToWebui', 99 | ]; 100 | 101 | const distanceBetweenIoSlots = 20; 102 | -------------------------------------------------------------------------------- /comfyui_custom_scripts/extensions/webuiPatches.js: -------------------------------------------------------------------------------- 1 | import { app } from "/scripts/app.js"; 2 | import { api } from "/scripts/api.js"; 3 | import { iframeRegisteredEvent } from "/extensions/webui_scripts/sd-webui-comfyui/extensions/webuiEvents.js"; 4 | import { getTypesLength } from "/extensions/webui_scripts/sd-webui-comfyui/extensions/webuiTypes.js"; 5 | 6 | 7 | async function patchUiEnv(iframeInfo) { 8 | if (iframeInfo.workflowTypeId.endsWith('_txt2img') || iframeInfo.workflowTypeId.endsWith('_img2img')) { 9 | const menuToHide = document.querySelector('.comfy-menu'); 10 | menuToHide.style.display = 'none'; 11 | patchSavingMechanism(); 12 | } 13 | 14 | await patchDefaultGraph(iframeInfo); 15 | } 16 | 17 | function patchSavingMechanism() { 18 | app.graph.original_serialize = app.graph.serialize; 19 | app.graph.patched_serialize = () => JSON.parse(localStorage.getItem('workflow')); 20 | app.graph.serialize = app.graph.patched_serialize; 21 | 22 | app.original_graphToPrompt = app.graphToPrompt; 23 | app.patched_graphToPrompt = () => { 24 | app.graph.serialize = app.graph.original_serialize; 25 | const result = app.original_graphToPrompt(); 26 | app.graph.serialize = app.graph.patched_serialize; 27 | return result; 28 | }; 29 | app.graphToPrompt = app.patched_graphToPrompt; 30 | 31 | const saveButton = document.querySelector('#comfy-save-button'); 32 | saveButton.removeAttribute('id'); 33 | const comfyParent = saveButton.parentElement; 34 | const muahahaButton = document.createElement('button'); 35 | muahahaButton.setAttribute('id', 'comfy-save-button'); 36 | comfyParent.appendChild(muahahaButton); 37 | muahahaButton.click = () => { 38 | app.graph.serialize = app.graph.original_serialize; 39 | saveButton.click(); 40 | app.graph.serialize = app.graph.patched_serialize; 41 | }; 42 | } 43 | 44 | async function patchDefaultGraph(iframeInfo) { 45 | // preserve the normal default graph 46 | if (!iframeInfo.defaultWorkflow) { 47 | return; 48 | } 49 | 50 | app.original_loadGraphData = app.loadGraphData; 51 | const doLoadGraphData = graphData => { 52 | if (graphData !== "auto") { 53 | return app.original_loadGraphData(graphData); 54 | } 55 | 56 | app.graph.clear(); 57 | 58 | const from_webui = LiteGraph.createNode("FromWebui"); 59 | const to_webui = LiteGraph.createNode("ToWebui"); 60 | 61 | if (!from_webui || !to_webui) { 62 | return; 63 | } 64 | 65 | app.graph.add(from_webui); 66 | app.graph.add(to_webui); 67 | 68 | const typesLength = getTypesLength(iframeInfo.webuiIoTypes.outputs); 69 | for (let i = 0; i < typesLength; ++i) { 70 | from_webui.connect(i, to_webui, i); 71 | } 72 | 73 | app.graph.arrange(); 74 | }; 75 | 76 | app.loadGraphData = (graphData) => { 77 | if (graphData) { 78 | return doLoadGraphData(graphData); 79 | } 80 | else { 81 | return doLoadGraphData(iframeInfo.defaultWorkflow); 82 | } 83 | }; 84 | 85 | app.loadGraphData(); 86 | } 87 | 88 | 89 | export { 90 | patchUiEnv, 91 | } 92 | -------------------------------------------------------------------------------- /comfyui_custom_scripts/extensions/webuiRequests.js: -------------------------------------------------------------------------------- 1 | import { app } from "/scripts/app.js"; 2 | import { api } from "/scripts/api.js"; 3 | import { patchUiEnv } from "/extensions/webui_scripts/sd-webui-comfyui/extensions/webuiPatches.js"; 4 | import { iframeRegisteredEvent, appReadyEvent } from "/extensions/webui_scripts/sd-webui-comfyui/extensions/webuiEvents.js"; 5 | 6 | 7 | async function setupWebuiRequestsEnvironment() { 8 | const iframeInfo = await iframeRegisteredEvent; 9 | await appReadyEvent; 10 | await patchUiEnv(iframeInfo); 11 | 12 | function addWebuiRequestListener(type, callback, options) { 13 | api.addEventListener(`webui_${type}`, async (data) => { 14 | api.fetchApi("/sd-webui-comfyui/webui_ws_response", { 15 | method: "POST", 16 | headers: {"Content-Type": "application/json"}, 17 | cache: "no-store", 18 | body: JSON.stringify({response: await callback(data)}), 19 | }); 20 | console.log(`[sd-webui-comfyui] WEBUI REQUEST - ${iframeInfo.workflowTypeId} - ${type}`); 21 | }, options); 22 | }; 23 | 24 | webuiRequests.forEach((request, type) => addWebuiRequestListener(type, request)); 25 | await registerClientToWebui(iframeInfo.workflowTypeId, iframeInfo.webuiClientId, window.name); 26 | console.log(`[sd-webui-comfyui][comfyui] INITIALIZED WS - ${iframeInfo.displayName}`); 27 | } 28 | 29 | async function registerClientToWebui(workflowTypeId, webuiClientId, sid) { 30 | await api.fetchApi("/sd-webui-comfyui/webui_register_client", { 31 | method: "POST", 32 | headers: {"Content-Type": "application/json"}, 33 | cache: "no-store", 34 | body: JSON.stringify({ 35 | workflowTypeId, 36 | webuiClientId, 37 | sid, 38 | }), 39 | }); 40 | } 41 | 42 | const webuiRequests = new Map([ 43 | ["queue_prompt", async (json) => { 44 | await app.queuePrompt(json.detail.queueFront ? -1 : 0, 1); 45 | }], 46 | ["serialize_graph", (json) => { 47 | return app.graph.original_serialize(); 48 | }], 49 | ["set_workflow", (json) => { 50 | app.loadGraphData(json.detail.workflow); 51 | }], 52 | ]); 53 | 54 | setupWebuiRequestsEnvironment(); 55 | -------------------------------------------------------------------------------- /comfyui_custom_scripts/extensions/webuiTypes.js: -------------------------------------------------------------------------------- 1 | function isString(value) { 2 | return typeof value === "string" || value instanceof String; 3 | } 4 | 5 | function getTypesLength(types) { 6 | if (isString(types)) { 7 | return 1; 8 | } 9 | else if (Array.isArray(types)) { 10 | return types.length; 11 | } 12 | else { 13 | return Object.keys(types).length; 14 | } 15 | } 16 | 17 | 18 | export { 19 | isString, 20 | getTypesLength, 21 | } 22 | -------------------------------------------------------------------------------- /install.py: -------------------------------------------------------------------------------- 1 | import launch 2 | import pathlib 3 | import pkg_resources 4 | import re 5 | import sys 6 | import traceback 7 | 8 | 9 | req_file = pathlib.Path(__file__).resolve().parent / "requirements.txt" 10 | req_re = re.compile('^([^=<>~]*)\s*(?:([=<>~])=\s*([^=<>~]*))?$') 11 | 12 | 13 | with open(req_file) as file: 14 | for package in file: 15 | try: 16 | package = package.strip() 17 | match = req_re.search(package) 18 | package_name = match.group(1) 19 | 20 | try: 21 | installed_version = pkg_resources.get_distribution(package_name).version 22 | except Exception: 23 | installed_version = None 24 | pass # package not installed, we still want to install it 25 | 26 | package_already_installed = installed_version is not None 27 | install_info = f"sd-webui-comfyui requirement: {package}" 28 | comparison, required_version = match.group(2, 3) 29 | 30 | if package_already_installed: 31 | install_info = f"sd-webui-comfyui requirement: changing {package_name} version from {installed_version} to {required_version}" 32 | if ( 33 | comparison == '~' or 34 | required_version is None or 35 | eval(f'"{installed_version}" {comparison}= "{required_version}"') 36 | ): 37 | continue 38 | 39 | launch.run_pip(f"install {package}", install_info) 40 | except Exception as e: 41 | print(traceback.format_exception_only(e)) 42 | print(f'Failed to install sd-webui-comfyui requirement: {package}', file=sys.stderr) 43 | -------------------------------------------------------------------------------- /install_comfyui.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | 5 | default_install_location = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'ComfyUI') 6 | 7 | 8 | def main(install_location, should_install_manager=False): 9 | repo_url = 'https://github.com/comfyanonymous/ComfyUI.git' 10 | install_repo(repo_url, install_location) 11 | 12 | if should_install_manager: 13 | manager_repo_url = 'https://github.com/ltdrdata/ComfyUI-Manager.git' 14 | manager_location = manager_location_from_comfyui_location(install_location) 15 | install_repo(manager_repo_url, manager_location) 16 | 17 | 18 | def manager_location_from_comfyui_location(comfyui_location): 19 | return os.path.join(comfyui_location, 'custom_nodes', 'ComfyUI-Manager') 20 | 21 | 22 | def install_repo(git_repo_url, install_location): 23 | import git 24 | os.mkdir(install_location) 25 | git.Repo.clone_from(git_repo_url, install_location) 26 | 27 | 28 | def update(install_location): 29 | print("[sd-webui-comfyui]", f"Updating comfyui at {install_location}...") 30 | if not install_location.is_dir() or not any(install_location.iterdir()): 31 | print("[sd-webui-comfyui]", f"Cannot update comfyui since it is not installed.", file=sys.stderr) 32 | return 33 | 34 | import git 35 | repo = git.Repo(install_location) 36 | current = repo.head.commit 37 | repo.remotes.origin.pull() 38 | if current == repo.head.commit: 39 | print("[sd-webui-comfyui]", "Already up to date.") 40 | else: 41 | print("[sd-webui-comfyui]", "Done updating comfyui.") 42 | 43 | 44 | if __name__ == '__main__': 45 | install_location = default_install_location 46 | if len(sys.argv) > 1: 47 | inistall_location = sys.argv[1] 48 | 49 | main(install_location) 50 | -------------------------------------------------------------------------------- /javascript/compute-height.js: -------------------------------------------------------------------------------- 1 | const SD_WEBUI_COMFYUI_POLLING_TIMEOUT = 500; 2 | 3 | 4 | function uuidv4() { 5 | return ([1e7]+-1e3+-4e3+-8e3+-1e11).replace(/[018]/g, c => (c ^ crypto.getRandomValues(new Uint8Array(1))[0] & 15 >> c / 4).toString(16)); 6 | } 7 | 8 | const WEBUI_CLIENT_ID = uuidv4(); 9 | 10 | function changeDisplayedWorkflowType(targetWorkflowType) { 11 | const targetIFrameElement = getWorkflowTypeIFrame(targetWorkflowType); 12 | const currentIFrameElement = targetIFrameElement.parentElement.querySelector(".comfyui-workflow-type-visible"); 13 | currentIFrameElement.classList.remove("comfyui-workflow-type-visible"); 14 | targetIFrameElement.classList.add("comfyui-workflow-type-visible"); 15 | } 16 | 17 | document.addEventListener("DOMContentLoaded", () => { 18 | onComfyuiTabLoaded(clearEnabledDisplayNames); 19 | onComfyuiTabLoaded(setupComfyuiTabEvents); 20 | }); 21 | 22 | function onComfyuiTabLoaded(callback) { 23 | if (getClearEnabledDisplayNamesButtons().some(e => e === null) || 24 | getWorkflowTypeIds() === null || 25 | getComfyuiContainer() === null || 26 | getTabNav() === null || 27 | getWebuiClientIdTextArea() === null 28 | ) { 29 | // webui not yet ready, try again in a bit 30 | setTimeout(() => { onComfyuiTabLoaded(callback); }, SD_WEBUI_COMFYUI_POLLING_TIMEOUT); 31 | return; 32 | } 33 | 34 | callback(); 35 | } 36 | 37 | function clearEnabledDisplayNames() { 38 | for (const clearButton of getClearEnabledDisplayNamesButtons()) { 39 | clearButton.click(); 40 | } 41 | } 42 | 43 | function setupComfyuiTabEvents() { 44 | setupWebuiClientId(); 45 | setupResizeTabEvent(); 46 | setupToggleFooterEvent(); 47 | 48 | updateComfyuiTabHeight(); 49 | 50 | getWorkflowTypeIds().forEach(id => setupIFrame(id)); 51 | } 52 | 53 | function reloadComfyuiIFrames() { 54 | getWorkflowTypeIds().forEach(id => { 55 | setupIFrame(id); 56 | }); 57 | } 58 | 59 | function setupWebuiClientId() { 60 | const textArea = getWebuiClientIdTextArea(); 61 | textArea.value = WEBUI_CLIENT_ID; 62 | textArea.dispatchEvent(new Event('input')); 63 | } 64 | 65 | function setupResizeTabEvent() { 66 | window.addEventListener("resize", updateComfyuiTabHeight); 67 | } 68 | 69 | function setupToggleFooterEvent() { 70 | new MutationObserver((mutationsList) => { 71 | for (const mutation of mutationsList) { 72 | if (mutation.type === 'attributes' && mutation.attributeName === 'style') { 73 | updateFooterStyle(); 74 | } 75 | } 76 | }) 77 | .observe(getComfyuiTab(), { attributes: true }); 78 | } 79 | 80 | function updateComfyuiTabHeight() { 81 | const container = getComfyuiContainer(); 82 | const tabNavBottom = getTabNav().getBoundingClientRect().bottom; 83 | container.style.height = `calc(100% - ${tabNavBottom}px)`; 84 | } 85 | 86 | function updateFooterStyle() { 87 | const tabDisplay = getComfyuiTab().style.display; 88 | const footer = getFooter(); 89 | 90 | if(footer === null) return; 91 | if(tabDisplay === 'block') { 92 | footer.classList.add('comfyui-remove-display'); 93 | } 94 | else { 95 | footer.classList.remove('comfyui-remove-display'); 96 | } 97 | } 98 | 99 | function getClearEnabledDisplayNamesButtons() { 100 | return [ 101 | document.getElementById("script_txt2img_comfyui_clear_enabled_display_names") ?? document.getElementById("script_txt2txt_comfyui_clear_enabled_display_names") ?? null, 102 | document.getElementById("script_img2img_comfyui_clear_enabled_display_names") ?? null, 103 | ]; 104 | } 105 | 106 | function getTabNav() { 107 | const tabs = document.getElementById("tabs") ?? null; 108 | return tabs ? tabs.querySelector(".tab-nav") : null; 109 | } 110 | 111 | function getComfyuiTab() { 112 | return document.getElementById("tab_comfyui_webui_root") ?? null; 113 | } 114 | 115 | function getComfyuiContainer() { 116 | return document.getElementById("comfyui_webui_container") ?? null; 117 | } 118 | 119 | function getWebuiClientIdTextArea() { 120 | return document.querySelector("#comfyui_webui_client_id textarea") ?? null; 121 | } 122 | 123 | function getFooter() { 124 | return document.querySelector('#footer') ?? null; 125 | } 126 | 127 | function getWorkflowTypeIFrame(workflowTypeId) { 128 | return document.querySelector(`[workflow_type_id="${workflowTypeId}"]`); 129 | } 130 | 131 | function getWorkflowTypeIds() { 132 | return getExtensionDynamicProperty('workflow_type_ids'); 133 | } 134 | 135 | function getExtensionDynamicProperty(key) { 136 | return JSON.parse(document.querySelector(`[sd_webui_comfyui_key="${key}"]`)?.innerText ?? "null"); 137 | } 138 | 139 | function reloadFrameElement(iframeElement) { 140 | iframeElement.src += ""; 141 | } 142 | 143 | function setupIFrame(workflowTypeId) { 144 | let messageToReceive = workflowTypeId; 145 | 146 | const iframeSearchParams = new URLSearchParams(); 147 | iframeSearchParams.set("workflowTypeId", workflowTypeId); 148 | iframeSearchParams.set("webuiClientId", WEBUI_CLIENT_ID); 149 | const iframe = getWorkflowTypeIFrame(workflowTypeId); 150 | const base_src = iframe.getAttribute("base_src"); 151 | const iframe_src = base_src + "?" + iframeSearchParams.toString(); 152 | if (iframe.src !== iframe_src) { 153 | iframe.src = iframe_src; 154 | } 155 | else { 156 | reloadFrameElement(iframe); 157 | } 158 | } 159 | -------------------------------------------------------------------------------- /lib_comfyui/argv_conversion.py: -------------------------------------------------------------------------------- 1 | from lib_comfyui import ipc 2 | from lib_comfyui.webui import settings 3 | 4 | 5 | COMFYUI_ARGV_PREFIX = 'comfyui_' 6 | 7 | 8 | def get_comfyui_args(): 9 | args = settings.get_additional_argv() + extract_comfyui_argv() 10 | deduplicate_comfyui_args(args) 11 | return args 12 | 13 | 14 | @ipc.restrict_to_process('webui') 15 | def extract_comfyui_argv(): 16 | from modules import shared 17 | result = [] 18 | for k, v in _items(shared.cmd_opts): 19 | if k.startswith(COMFYUI_ARGV_PREFIX): 20 | k = k.replace(COMFYUI_ARGV_PREFIX, '') 21 | result.extend(as_argv_list(k, v)) 22 | return result 23 | 24 | 25 | def as_argv_list(k, v): 26 | result = [] 27 | if is_used_argv(k, v): 28 | result.append(f'--{k.replace("_", "-")}') 29 | if is_paired_argv(k, v): 30 | result.append(str(v)) 31 | return result 32 | 33 | 34 | def deduplicate_comfyui_args(argv): 35 | seen_args = set() 36 | i = 0 37 | while i < len(argv): 38 | arg = argv[i] 39 | if arg in seen_args: 40 | if arg in ('--port',): 41 | len_to_remove = 2 42 | elif arg in ('--listen',): 43 | has_value = i + 1 < len(argv) and not argv[i + 1].startswith('--') 44 | len_to_remove = 1 + int(has_value) 45 | else: 46 | len_to_remove = 1 47 | argv[i:i + len_to_remove] = () 48 | else: 49 | if arg.startswith('--'): 50 | seen_args.add(arg) 51 | i += 1 52 | 53 | 54 | def _items(cmd_opts): 55 | return vars(cmd_opts).items() 56 | 57 | 58 | def is_used_argv(k, v): 59 | return v not in [False, None] 60 | 61 | 62 | def is_paired_argv(k, v): 63 | return v is not True 64 | -------------------------------------------------------------------------------- /lib_comfyui/comfyui/iframe_requests.py: -------------------------------------------------------------------------------- 1 | import json 2 | import multiprocessing 3 | from queue import Empty 4 | from typing import List, Any, Dict, Tuple, Optional 5 | from lib_comfyui import ipc, global_state, torch_utils, external_code 6 | from lib_comfyui.comfyui import queue_tracker 7 | 8 | 9 | class ComfyuiIFrameRequests: 10 | finished_comfyui_queue = multiprocessing.Queue() 11 | server_instance = None 12 | sid_map = {} 13 | 14 | @staticmethod 15 | @ipc.run_in_process('comfyui') 16 | def send(request, workflow_type, data=None): 17 | if data is None: 18 | data = {} 19 | 20 | cls = ComfyuiIFrameRequests 21 | if global_state.focused_webui_client_id is None: 22 | raise RuntimeError('No active webui connection') 23 | 24 | ws_client_ids = cls.sid_map[global_state.focused_webui_client_id] 25 | if workflow_type not in ws_client_ids: 26 | raise RuntimeError(f"The workflow type {workflow_type} has not been registered by the active webui client {global_state.focused_webui_client_id}") 27 | 28 | clear_queue(cls.finished_comfyui_queue) 29 | cls.server_instance.send_sync(request, data, ws_client_ids[workflow_type]) 30 | 31 | return cls.finished_comfyui_queue.get() 32 | 33 | @staticmethod 34 | @ipc.restrict_to_process('webui') 35 | def start_workflow_sync( 36 | batch_input_args: Tuple[Any, ...], 37 | workflow_type_id: str, 38 | workflow_input_types: List[str], 39 | queue_front: bool, 40 | ) -> List[Dict[str, Any]]: 41 | from modules import shared 42 | if shared.state.interrupted: 43 | raise RuntimeError('The workflow was not started because the webui has been interrupted') 44 | 45 | global_state.node_inputs = batch_input_args 46 | global_state.node_outputs = [] 47 | global_state.current_workflow_input_types = workflow_input_types 48 | 49 | try: 50 | queue_tracker.setup_tracker_id() 51 | 52 | # unsafe queue tracking 53 | ComfyuiIFrameRequests.send( 54 | request='webui_queue_prompt', 55 | workflow_type=workflow_type_id, 56 | data={ 57 | 'requiredNodeTypes': [], 58 | 'queueFront': queue_front, 59 | } 60 | ) 61 | 62 | if not queue_tracker.wait_until_done(): 63 | raise RuntimeError('The workflow has not returned normally') 64 | 65 | return global_state.node_outputs 66 | finally: 67 | global_state.current_workflow_input_types = () 68 | global_state.node_outputs = [] 69 | global_state.node_inputs = None 70 | 71 | @staticmethod 72 | @ipc.restrict_to_process('webui') 73 | def validate_amount_of_nodes_or_throw( 74 | workflow_type_id: str, 75 | max_amount_of_FromWebui_nodes: Optional[int], 76 | max_amount_of_ToWebui_nodes: Optional[int], 77 | ) -> None: 78 | workflow_graph = get_workflow_graph(workflow_type_id) 79 | all_nodes = workflow_graph['nodes'] 80 | enabled_nodes = [node for node in all_nodes if node['mode'] != 2] 81 | node_types = [node['type'] for node in enabled_nodes] 82 | amount_of_FromWebui_nodes = len([t for t in node_types if t == 'FromWebui']) 83 | amount_of_ToWebui_nodes = len([t for t in node_types if t == 'ToWebui']) 84 | max_FromWebui_nodes = max_amount_of_FromWebui_nodes if max_amount_of_FromWebui_nodes is not None else amount_of_FromWebui_nodes 85 | max_ToWebui_nodes = max_amount_of_ToWebui_nodes if max_amount_of_ToWebui_nodes is not None else amount_of_ToWebui_nodes 86 | 87 | if amount_of_FromWebui_nodes > max_FromWebui_nodes: 88 | raise TooManyFromWebuiNodesError(f'Unable to run the workflow {workflow_type_id}. ' 89 | f'Expected at most {max_FromWebui_nodes} FromWebui node(s), ' 90 | f'{amount_of_FromWebui_nodes} were found.') 91 | 92 | if amount_of_ToWebui_nodes > max_ToWebui_nodes: 93 | raise TooManyToWebuiNodesError(f'Unable to run the workflow {workflow_type_id}. ' 94 | f'Expected at most {max_ToWebui_nodes} ToWebui node(s), ' 95 | f'{amount_of_ToWebui_nodes} were found.') 96 | 97 | @staticmethod 98 | @ipc.restrict_to_process('comfyui') 99 | def register_client(request) -> None: 100 | workflow_type_id = request['workflowTypeId'] 101 | webui_client_id = request['webuiClientId'] 102 | sid = request['sid'] 103 | 104 | if webui_client_id not in ComfyuiIFrameRequests.sid_map: 105 | ComfyuiIFrameRequests.sid_map[webui_client_id] = {} 106 | 107 | ComfyuiIFrameRequests.sid_map[webui_client_id][workflow_type_id] = sid 108 | print(f'registered ws - {workflow_type_id} - {sid}') 109 | 110 | @staticmethod 111 | @ipc.restrict_to_process('comfyui') 112 | def handle_response(response): 113 | ComfyuiIFrameRequests.finished_comfyui_queue.put(response) 114 | 115 | 116 | def extend_infotext_with_comfyui_workflows(p, tab): 117 | workflows = {} 118 | for workflow_type in external_code.get_workflow_types(tab): 119 | workflow_type_id = workflow_type.get_ids(tab)[0] 120 | try: 121 | ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw( 122 | workflow_type_id, 123 | workflow_type.max_amount_of_FromWebui_nodes, 124 | workflow_type.max_amount_of_ToWebui_nodes, 125 | ) 126 | except RuntimeError: 127 | continue 128 | if not external_code.is_workflow_type_enabled(workflow_type_id): 129 | continue 130 | 131 | workflows[workflow_type.base_id] = get_workflow_graph(workflow_type_id) 132 | 133 | p.extra_generation_params['ComfyUI Workflows'] = json.dumps(workflows) 134 | 135 | 136 | def set_workflow_graph(workflow_json, workflow_type_id): 137 | return ComfyuiIFrameRequests.send( 138 | request='webui_set_workflow', 139 | workflow_type=workflow_type_id, 140 | data={'workflow': workflow_json} 141 | ) 142 | 143 | 144 | def get_workflow_graph(workflow_type_id): 145 | return ComfyuiIFrameRequests.send(request='webui_serialize_graph', workflow_type=workflow_type_id) 146 | 147 | 148 | def clear_queue(queue: multiprocessing.Queue): 149 | while not queue.empty(): 150 | try: 151 | queue.get(timeout=1) 152 | except Empty: 153 | pass 154 | 155 | 156 | class TooManyFromWebuiNodesError(RuntimeError): 157 | pass 158 | 159 | 160 | class TooManyToWebuiNodesError(RuntimeError): 161 | pass 162 | -------------------------------------------------------------------------------- /lib_comfyui/comfyui/pre_main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | 5 | def patch_sys_path(): 6 | comfyui_install_dir = os.getcwd() 7 | extension_dir = os.getenv("SD_WEBUI_COMFYUI_EXTENSION_DIR") 8 | if not comfyui_install_dir or not extension_dir: 9 | print("[sd-webui-comfyui]", f"Could not add new entries to sys.path. install_dir={comfyui_install_dir}, extension_dir={extension_dir}, sys.path={sys.path}", file=sys.stderr) 10 | print("[sd-webui-comfyui]", f"Exiting...", file=sys.stderr) 11 | exit(1) 12 | 13 | sys.path[:0] = (comfyui_install_dir, extension_dir) 14 | 15 | 16 | if __name__ == "__main__": 17 | patch_sys_path() 18 | 19 | 20 | import atexit 21 | import builtins 22 | import psutil 23 | import threading 24 | import time 25 | import runpy 26 | from lib_comfyui import ( 27 | custom_extension_injector, 28 | ipc, 29 | ) 30 | from lib_comfyui.comfyui import routes_extension, queue_tracker 31 | from lib_comfyui.webui import paths, settings 32 | 33 | 34 | original_print = builtins.print 35 | def comfyui_print(*args, **kwargs): 36 | return original_print('[ComfyUI]', *args, **kwargs) 37 | 38 | 39 | @ipc.restrict_to_process('comfyui') 40 | def main(): 41 | builtins.print = comfyui_print 42 | setup_ipc() 43 | patch_comfyui() 44 | start_comfyui() 45 | 46 | 47 | @ipc.restrict_to_process('comfyui') 48 | def setup_ipc(): 49 | print('[sd-webui-comfyui]', 'Setting up IPC...') 50 | ipc_strategy_class_name = os.getenv('SD_WEBUI_COMFYUI_IPC_STRATEGY_CLASS_NAME') 51 | print('[sd-webui-comfyui]', f'Using inter-process communication strategy: {settings.ipc_display_names[ipc_strategy_class_name]}') 52 | ipc_strategy_factory = getattr(ipc.strategies, os.getenv('SD_WEBUI_COMFYUI_IPC_STRATEGY_CLASS_NAME')) 53 | ipc.current_callback_listeners = {'comfyui': ipc.callback.CallbackWatcher(ipc.call_fully_qualified, 'comfyui', ipc_strategy_factory)} 54 | ipc.current_callback_proxies = {'webui': ipc.callback.CallbackProxy('webui', ipc_strategy_factory)} 55 | ipc.start_callback_listeners() 56 | atexit.register(ipc.stop_callback_listeners) 57 | 58 | parent_id = os.getppid() 59 | monitor_thread = threading.Thread(target=watch_webui_exit, args=(parent_id,)) 60 | monitor_thread.start() 61 | 62 | 63 | @ipc.restrict_to_process('comfyui') 64 | def watch_webui_exit(parent_id): 65 | while True: 66 | if not psutil.pid_exists(parent_id): 67 | print("[sd-webui-comfyui]", "The webui has exited, exiting comfyui.") 68 | exit() 69 | 70 | time.sleep(1) 71 | 72 | 73 | @ipc.restrict_to_process('comfyui') 74 | def patch_comfyui(): 75 | print('[sd-webui-comfyui]', 'Patching ComfyUI...') 76 | try: 77 | # workaround for newer versions of comfyui https://github.com/comfyanonymous/ComfyUI/commit/3039b08eb16777431946ed9ae4a63c5466336bff 78 | # remove the try-except to stop supporting older versions 79 | import comfy.options 80 | comfy.options.enable_args_parsing() 81 | except ImportError: 82 | pass 83 | 84 | paths.share_webui_folder_paths() 85 | custom_extension_injector.register_webui_extensions() 86 | routes_extension.patch_server_routes() 87 | queue_tracker.patch_prompt_queue() 88 | 89 | 90 | @ipc.restrict_to_process('comfyui') 91 | def start_comfyui(): 92 | print('[sd-webui-comfyui]', f'Launching ComfyUI with arguments: {" ".join(sys.argv[1:])}') 93 | runpy.run_path(os.path.join(os.getcwd(), 'main.py'), {'comfyui_print': comfyui_print}, '__main__') 94 | 95 | 96 | if __name__ == '__main__': 97 | ipc.current_process_id = 'comfyui' 98 | main() 99 | -------------------------------------------------------------------------------- /lib_comfyui/comfyui/queue_tracker.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import multiprocessing 3 | from lib_comfyui import ipc 4 | 5 | 6 | class PromptQueueTracker: 7 | done_event = multiprocessing.Event() 8 | put_event = multiprocessing.Event() 9 | tracked_id = None 10 | original_id = None 11 | queue_instance = None 12 | server_instance = None 13 | 14 | @staticmethod 15 | def patched__init__(self, server_instance): 16 | prompt_queue = self 17 | PromptQueueTracker.server_instance = server_instance 18 | PromptQueueTracker.queue_instance = self 19 | 20 | def patched_put(item, *args, original_put, **kwargs): 21 | with prompt_queue.mutex: 22 | if abs(item[0]) == PromptQueueTracker.tracked_id: 23 | PromptQueueTracker.put_event.set() 24 | 25 | return original_put(item, *args, **kwargs) 26 | 27 | prompt_queue.put = functools.partial(patched_put, original_put=prompt_queue.put) 28 | 29 | # task_done 30 | def patched_task_done(item_id, output, *args, original_task_done, **kwargs): 31 | with prompt_queue.mutex: 32 | v = prompt_queue.currently_running[item_id] 33 | if abs(v[0]) == PromptQueueTracker.tracked_id: 34 | PromptQueueTracker.done_event.set() 35 | 36 | return original_task_done(item_id, output, *args, **kwargs) 37 | 38 | prompt_queue.task_done = functools.partial(patched_task_done, original_task_done=prompt_queue.task_done) 39 | 40 | # wipe_queue 41 | def patched_wipe_queue(*args, original_wipe_queue, **kwargs): 42 | with prompt_queue.mutex: 43 | should_release_webui = True 44 | for v in prompt_queue.currently_running.values(): 45 | if abs(v[0]) == PromptQueueTracker.tracked_id: 46 | should_release_webui = False 47 | 48 | if should_release_webui: 49 | PromptQueueTracker.done_event.set() 50 | 51 | return original_wipe_queue(*args, **kwargs) 52 | 53 | prompt_queue.wipe_queue = functools.partial(patched_wipe_queue, original_wipe_queue=prompt_queue.wipe_queue) 54 | 55 | # delete_queue_item 56 | def patched_delete_queue_item(function, *args, original_delete_queue_item, **kwargs): 57 | def patched_function(x): 58 | res = function(x) 59 | if res and abs(x[0]) == PromptQueueTracker.tracked_id: 60 | PromptQueueTracker.done_event.set() 61 | return res 62 | 63 | return original_delete_queue_item(patched_function, *args, **kwargs) 64 | 65 | prompt_queue.delete_queue_item = functools.partial(patched_delete_queue_item, original_delete_queue_item=prompt_queue.delete_queue_item) 66 | 67 | 68 | @ipc.run_in_process('comfyui') 69 | def setup_tracker_id(): 70 | PromptQueueTracker.original_id = PromptQueueTracker.tracked_id 71 | PromptQueueTracker.tracked_id = PromptQueueTracker.server_instance.number 72 | PromptQueueTracker.put_event.clear() 73 | PromptQueueTracker.done_event.clear() 74 | 75 | 76 | @ipc.restrict_to_process('webui') 77 | def wait_until_done(): 78 | if not wait_until_put(): 79 | return False 80 | 81 | from modules import shared 82 | while True: 83 | has_been_set = check_done_event(timeout=1) 84 | if has_been_set: 85 | return True 86 | elif not tracked_id_present(): 87 | return False 88 | elif shared.state.interrupted: 89 | cancel_queued_workflow() 90 | return False 91 | 92 | 93 | @ipc.run_in_process('comfyui') 94 | def check_done_event(*args, **kwargs): 95 | return PromptQueueTracker.done_event.wait(*args, **kwargs) 96 | 97 | 98 | @ipc.run_in_process('comfyui') 99 | def wait_until_put(): 100 | was_put = PromptQueueTracker.put_event.wait(timeout=3) 101 | if not was_put: 102 | PromptQueueTracker.tracked_id = PromptQueueTracker.original_id 103 | return False 104 | 105 | return True 106 | 107 | 108 | @ipc.run_in_process('comfyui') 109 | def cancel_queued_workflow(): 110 | with PromptQueueTracker.queue_instance.mutex: 111 | is_running = False 112 | for v in PromptQueueTracker.queue_instance.currently_running.values(): 113 | if abs(v[0]) == PromptQueueTracker.tracked_id: 114 | is_running = True 115 | 116 | if is_running: 117 | import nodes 118 | nodes.interrupt_processing() 119 | return 120 | 121 | PromptQueueTracker.queue_instance.delete_queue_item(lambda a: a[1] == PromptQueueTracker.tracked_id) 122 | 123 | 124 | @ipc.run_in_process('comfyui') 125 | def tracked_id_present(): 126 | with PromptQueueTracker.queue_instance.mutex: 127 | for v in PromptQueueTracker.queue_instance.currently_running.values(): 128 | if abs(v[0]) == PromptQueueTracker.tracked_id: 129 | return True 130 | for x in PromptQueueTracker.queue_instance.queue: 131 | if abs(x[0]) == PromptQueueTracker.tracked_id: 132 | return True 133 | return False 134 | 135 | 136 | @ipc.restrict_to_process('comfyui') 137 | def add_queue__init__patch(callback): 138 | import execution 139 | original_init = execution.PromptQueue.__init__ 140 | 141 | def patched_PromptQueue__init__(self, server, *args, **kwargs): 142 | original_init(self, server, *args, **kwargs) 143 | callback(self, server, *args, **kwargs) 144 | 145 | execution.PromptQueue.__init__ = patched_PromptQueue__init__ 146 | 147 | 148 | @ipc.restrict_to_process('comfyui') 149 | def patch_prompt_queue(): 150 | add_queue__init__patch(PromptQueueTracker.patched__init__) 151 | -------------------------------------------------------------------------------- /lib_comfyui/comfyui/routes_extension.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import json 3 | from lib_comfyui import external_code 4 | from lib_comfyui.comfyui.iframe_requests import ComfyuiIFrameRequests 5 | 6 | 7 | def patch_server_routes(): 8 | add_server__init__patch(websocket_handler_patch) 9 | add_server__init__patch(workflow_type_ops_server_patch) 10 | 11 | 12 | def add_server__init__patch(callback): 13 | import server 14 | original_init = server.PromptServer.__init__ 15 | 16 | def patched_PromptServer__init__(self, loop: asyncio.AbstractEventLoop, *args, **kwargs): 17 | original_init(self, loop, *args, **kwargs) 18 | callback(self, loop, *args, **kwargs) 19 | 20 | server.PromptServer.__init__ = patched_PromptServer__init__ 21 | 22 | 23 | def websocket_handler_patch(instance, _loop): 24 | from aiohttp import web 25 | 26 | ComfyuiIFrameRequests.server_instance = instance 27 | 28 | @instance.routes.post("/sd-webui-comfyui/webui_register_client") 29 | async def webui_register_client(request): 30 | request = await request.json() 31 | 32 | ComfyuiIFrameRequests.register_client(request) 33 | 34 | return web.json_response() 35 | 36 | @instance.routes.post("/sd-webui-comfyui/webui_ws_response") 37 | async def webui_ws_response(response): 38 | response = await response.json() 39 | 40 | ComfyuiIFrameRequests.handle_response(response['response'] if 'response' in response else response) 41 | 42 | return web.json_response(status=200) 43 | 44 | 45 | def workflow_type_ops_server_patch(instance, _loop): 46 | from aiohttp import web 47 | 48 | @instance.routes.get("/sd-webui-comfyui/workflow_type") 49 | async def get_workflow_type(request): 50 | workflow_type_id = request.rel_url.query.get("workflowTypeId", None) 51 | workflow_type = next(iter( 52 | workflow_type 53 | for workflow_type in external_code.get_workflow_types() 54 | if workflow_type_id in workflow_type.get_ids() 55 | )) 56 | return web.json_response({ 57 | "displayName": workflow_type.display_name, 58 | "webuiIoTypes": { 59 | "inputs": list(workflow_type.input_types) if isinstance(workflow_type.input_types, tuple) else workflow_type.input_types, 60 | "outputs": list(workflow_type.types) if isinstance(workflow_type.types, tuple) else workflow_type.types, 61 | }, 62 | "defaultWorkflow": json.loads(external_code.get_default_workflow_json(workflow_type_id)) 63 | }) 64 | -------------------------------------------------------------------------------- /lib_comfyui/comfyui/type_conversion.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from PIL import Image 3 | import torchvision.transforms.functional as F 4 | from lib_comfyui import ipc 5 | from lib_comfyui.webui.proxies import get_comfy_model_config 6 | 7 | 8 | def webui_image_to_comfyui(batch): 9 | if isinstance(batch[0], Image.Image): 10 | batch = torch.stack([F.pil_to_tensor(image) / 255 for image in batch]) 11 | return batch.permute(0, 2, 3, 1) 12 | 13 | 14 | def comfyui_image_to_webui(batch, return_tensors=False): 15 | batch = batch.permute(0, 3, 1, 2) 16 | if return_tensors: 17 | return batch 18 | 19 | return [F.to_pil_image(image) for image in batch] 20 | 21 | 22 | @ipc.run_in_process('comfyui') 23 | def webui_latent_to_comfyui(batch): 24 | latent_format = get_comfy_model_config().latent_format 25 | return {'samples': latent_format.process_out(batch)} 26 | 27 | 28 | @ipc.run_in_process('comfyui') 29 | def comfyui_latent_to_webui(batch): 30 | latent_format = get_comfy_model_config().latent_format 31 | return latent_format.process_in(batch['samples']) 32 | -------------------------------------------------------------------------------- /lib_comfyui/comfyui_process.py: -------------------------------------------------------------------------------- 1 | import atexit 2 | import inspect 3 | import os 4 | import subprocess 5 | import sys 6 | from pathlib import Path 7 | 8 | from lib_comfyui import ipc, torch_utils, argv_conversion, global_state 9 | from lib_comfyui.webui import settings 10 | from lib_comfyui.comfyui import pre_main 11 | 12 | 13 | comfyui_process = None 14 | 15 | 16 | @ipc.restrict_to_process('webui') 17 | def start(): 18 | if not global_state.enabled: 19 | return 20 | 21 | install_location = settings.get_install_location() 22 | if not install_location.exists(): 23 | print('[sd-webui-comfyui]', f'Could not find ComfyUI under directory "{install_location}". The server will NOT be started.', file=sys.stderr) 24 | return 25 | 26 | ipc.current_callback_listeners = {'webui': ipc.callback.CallbackWatcher(ipc.call_fully_qualified, 'webui', global_state.ipc_strategy_class, clear_on_init=True)} 27 | ipc.current_callback_proxies = {'comfyui': ipc.callback.CallbackProxy('comfyui', global_state.ipc_strategy_class, clear_on_init=True)} 28 | ipc.start_callback_listeners() 29 | atexit.register(stop) 30 | start_comfyui_process(install_location) 31 | 32 | 33 | @ipc.restrict_to_process('webui') 34 | def start_comfyui_process(comfyui_install_location): 35 | global comfyui_process 36 | 37 | executable = get_comfyui_executable(comfyui_install_location) 38 | comfyui_env = get_comfyui_env(comfyui_install_location) 39 | install_comfyui_requirements(executable, comfyui_install_location, comfyui_env) 40 | args = [executable, inspect.getfile(pre_main)] + argv_conversion.get_comfyui_args() 41 | 42 | print('[sd-webui-comfyui]', 'Starting subprocess for comfyui...') 43 | comfyui_process = subprocess.Popen( 44 | args=args, 45 | executable=executable, 46 | cwd=str(comfyui_install_location), 47 | env=comfyui_env, 48 | ) 49 | 50 | 51 | def get_comfyui_executable(comfyui_install_location: Path) -> str: 52 | executable = sys.executable 53 | 54 | if os.name == 'nt': 55 | executable_paths = [ 56 | comfyui_install_location / 'venv' / 'scripts' / 'python.exe', 57 | comfyui_install_location.parent / 'python_embeded' / 'python.exe', 58 | ] 59 | else: 60 | executable_paths = [ 61 | comfyui_install_location / 'venv' / 'bin' / 'python', 62 | comfyui_install_location.parent / 'python_embeded' / 'python', 63 | ] 64 | 65 | for potential_executable in executable_paths: 66 | if potential_executable.exists(): 67 | executable = potential_executable 68 | print('[sd-webui-comfyui]', 'Detected custom ComfyUI venv:', executable) 69 | break 70 | 71 | return str(executable) 72 | 73 | 74 | def get_comfyui_env(comfyui_install_location): 75 | comfyui_env = os.environ.copy() 76 | if 'PYTHONPATH' in comfyui_env: 77 | del comfyui_env['PYTHONPATH'] 78 | 79 | problematic_windows_tmp_dirs = { 80 | k: v 81 | for k, v in comfyui_env.items() 82 | if k in ("TMP", "TEMP") and v == "tmp" 83 | } 84 | if os.name == "nt" and problematic_windows_tmp_dirs: 85 | problematic_windows_tmp_dirs = " and ".join(f'{k}="{v}"' for k, v in problematic_windows_tmp_dirs.items()) 86 | message_prefix = "[sd-webui-comfyui] " 87 | print( 88 | "\n".join([ 89 | f'{message_prefix}Found {problematic_windows_tmp_dirs} in the environment. On Windows, this is known to cause issues.', 90 | f'{message_prefix}If ComfyUI refuses to start, consider using TEMP="temp" or TMP="temp" as a workaround.', 91 | f'{message_prefix}For more information, see https://github.com/ModelSurge/sd-webui-comfyui/issues/170#issuecomment-1792009789', 92 | ]), 93 | file=sys.stderr, 94 | ) 95 | 96 | comfyui_env['SD_WEBUI_COMFYUI_EXTENSION_DIR'] = settings.get_extension_base_dir() 97 | comfyui_env['SD_WEBUI_COMFYUI_IPC_STRATEGY_CLASS_NAME'] = global_state.ipc_strategy_class.__name__ 98 | return comfyui_env 99 | 100 | 101 | def install_comfyui_requirements(executable, comfyui_install_location, comfyui_env): 102 | if executable == sys.executable: 103 | # requirements already installed in the webui by install.py 104 | return 105 | 106 | print('[sd-webui-comfyui]', 'Installing mandatory pip requirements in ComfyUI venv...') 107 | subprocess.check_call( 108 | args=[ 109 | executable, 110 | *(['-s'] if "python_embeded" in executable or "python_embedded" in executable else []), 111 | '-m', 112 | 'pip', 113 | 'install', 114 | '-r', 115 | str(Path(settings.get_extension_base_dir(), 'requirements.txt')), 116 | ], 117 | executable=executable, 118 | cwd=str(comfyui_install_location), 119 | env=comfyui_env, 120 | ) 121 | 122 | 123 | @ipc.restrict_to_process('webui') 124 | def stop(): 125 | atexit.unregister(stop) 126 | stop_comfyui_process() 127 | ipc.stop_callback_listeners() 128 | 129 | 130 | @ipc.restrict_to_process('webui') 131 | def stop_comfyui_process(): 132 | global comfyui_process 133 | if comfyui_process is None: 134 | return 135 | 136 | print('[sd-webui-comfyui]', 'Attempting to gracefully terminate the ComfyUI server...') 137 | comfyui_process.terminate() 138 | try: 139 | comfyui_process.wait(global_state.comfyui_graceful_termination_timeout) 140 | print('[sd-webui-comfyui]', 'The ComfyUI server was gracefully terminated') 141 | except subprocess.TimeoutExpired: 142 | print('[sd-webui-comfyui]', 'Graceful termination timed out. Killing the ComfyUI server...') 143 | comfyui_process.kill() 144 | print('[sd-webui-comfyui]', 'The ComfyUI server was killed') 145 | comfyui_process = None 146 | -------------------------------------------------------------------------------- /lib_comfyui/custom_extension_injector.py: -------------------------------------------------------------------------------- 1 | import os 2 | from lib_comfyui import find_extensions, ipc 3 | 4 | 5 | def register_webui_extensions(): 6 | node_paths, script_paths = find_extensions.get_extension_paths_to_load() 7 | register_custom_nodes(node_paths) 8 | register_custom_scripts(script_paths) 9 | 10 | 11 | @ipc.restrict_to_process('comfyui') 12 | def register_custom_nodes(custom_nodes_path_list): 13 | from folder_paths import add_model_folder_path 14 | 15 | for custom_nodes_path in custom_nodes_path_list: 16 | add_model_folder_path('custom_nodes', custom_nodes_path) 17 | 18 | 19 | @ipc.restrict_to_process('comfyui') 20 | def register_custom_scripts(custom_scripts_path_list): 21 | from nodes import EXTENSION_WEB_DIRS 22 | 23 | for custom_scripts_path in custom_scripts_path_list: 24 | name = f"webui_scripts/{os.path.basename(os.path.dirname(custom_scripts_path))}" 25 | dir = custom_scripts_path 26 | EXTENSION_WEB_DIRS[name] = dir 27 | -------------------------------------------------------------------------------- /lib_comfyui/default_workflow_types.py: -------------------------------------------------------------------------------- 1 | from lib_comfyui import external_code 2 | 3 | 4 | sandbox_tab_workflow_type = external_code.WorkflowType( 5 | base_id='sandbox', 6 | display_name='ComfyUI', 7 | tabs='tab', 8 | ) 9 | preprocess_workflow_type = external_code.WorkflowType( 10 | base_id='preprocess', 11 | display_name='Preprocess', 12 | tabs='img2img', 13 | default_workflow=external_code.AUTO_WORKFLOW, 14 | types='IMAGE', 15 | ) 16 | preprocess_latent_workflow_type = external_code.WorkflowType( 17 | base_id='preprocess_latent', 18 | display_name='Preprocess (latent)', 19 | tabs='img2img', 20 | default_workflow=external_code.AUTO_WORKFLOW, 21 | types='LATENT', 22 | ) 23 | postprocess_latent_workflow_type = external_code.WorkflowType( 24 | base_id='postprocess_latent', 25 | display_name='Postprocess (latent)', 26 | default_workflow=external_code.AUTO_WORKFLOW, 27 | types='LATENT', 28 | ) 29 | postprocess_workflow_type = external_code.WorkflowType( 30 | base_id='postprocess', 31 | display_name='Postprocess', 32 | default_workflow=external_code.AUTO_WORKFLOW, 33 | types='IMAGE', 34 | ) 35 | postprocess_image_workflow_type = external_code.WorkflowType( 36 | base_id='postprocess_image', 37 | display_name='Postprocess image', 38 | default_workflow=external_code.AUTO_WORKFLOW, 39 | types='IMAGE', 40 | max_amount_of_ToWebui_nodes=1, 41 | ) 42 | before_save_image_workflow_type = external_code.WorkflowType( 43 | base_id='before_save_image', 44 | display_name='Before save image', 45 | default_workflow=external_code.AUTO_WORKFLOW, 46 | types='IMAGE', 47 | max_amount_of_ToWebui_nodes=1, 48 | ) 49 | 50 | 51 | def add_default_workflow_types(): 52 | workflow_types = [ 53 | sandbox_tab_workflow_type, 54 | preprocess_workflow_type, 55 | preprocess_latent_workflow_type, 56 | postprocess_latent_workflow_type, 57 | postprocess_workflow_type, 58 | postprocess_image_workflow_type, 59 | before_save_image_workflow_type, 60 | ] 61 | 62 | for workflow_type in workflow_types: 63 | external_code.add_workflow_type(workflow_type) 64 | -------------------------------------------------------------------------------- /lib_comfyui/external_code/__init__.py: -------------------------------------------------------------------------------- 1 | def fix_path(): 2 | import sys 3 | from pathlib import Path 4 | 5 | extension_path = str(Path(__file__).parent.parent.parent) 6 | if extension_path not in sys.path: 7 | sys.path.append(extension_path) 8 | 9 | 10 | fix_path() 11 | from .api import * 12 | -------------------------------------------------------------------------------- /lib_comfyui/external_code/api.py: -------------------------------------------------------------------------------- 1 | import dataclasses 2 | import sys 3 | import traceback 4 | from pathlib import Path 5 | from typing import List, Tuple, Union, Any, Optional, Dict 6 | from lib_comfyui import global_state, ipc 7 | 8 | 9 | ALL_TABS = ... 10 | Tabs = Union[str, Tuple[str, ...]] 11 | AUTO_WORKFLOW = '"auto"' # json encoded string 12 | 13 | 14 | @dataclasses.dataclass 15 | class WorkflowType: 16 | """ 17 | Describes a unique type of ComfyUI workflow 18 | """ 19 | 20 | base_id: str 21 | display_name: str 22 | tabs: Tabs = ('txt2img', 'img2img') 23 | default_workflow: Union[str, Path] = "null" 24 | types: Union[str, Tuple[str, ...], Dict[str, str]] = dataclasses.field(default_factory=tuple) 25 | input_types: Union[str, Tuple[str, ...], Dict[str, str], None] = None 26 | max_amount_of_ToWebui_nodes: Optional[int] = None 27 | max_amount_of_FromWebui_nodes: Optional[int] = None 28 | 29 | def __post_init__(self): 30 | if isinstance(self.tabs, str): 31 | self.tabs = (self.tabs,) 32 | 33 | if self.input_types is None: 34 | self.input_types = self.types 35 | 36 | if not isinstance(self.types, (str, tuple, dict)): 37 | raise TypeError(f'types should be str, tuple or dict but it is {type(self.types)}') 38 | 39 | if not isinstance(self.input_types, (str, tuple, dict)): 40 | raise TypeError(f'input_types should be str, tuple or dict but it is {type(self.input_types)}') 41 | 42 | assert self.tabs, "tabs must not be empty" 43 | 44 | if isinstance(self.default_workflow, Path): 45 | with open(str(self.default_workflow), 'r') as f: 46 | self.default_workflow = f.read() 47 | elif self.default_workflow == AUTO_WORKFLOW: 48 | if not self.is_same_io(): 49 | raise ValueError('AUTO_WORKFLOW default workflow is only supported for same input and output types') 50 | 51 | def get_ids(self, tabs: Tabs = ALL_TABS) -> List[str]: 52 | if isinstance(tabs, str): 53 | tabs = (tabs,) 54 | 55 | return [ 56 | f'{self.base_id}_{tab}' 57 | for tab in self.tabs 58 | if tabs == ALL_TABS or tab in tabs 59 | ] 60 | 61 | def pretty_str(self): 62 | return f'"{self.display_name}" ({self.base_id})' 63 | 64 | def is_same_io(self): 65 | def normalize_to_tuple(types): 66 | if isinstance(types, dict): 67 | return tuple(types.values()) 68 | elif isinstance(types, str): 69 | return types, 70 | return types 71 | 72 | return normalize_to_tuple(self.input_types) == normalize_to_tuple(self.types) 73 | 74 | 75 | def get_workflow_types(tabs: Tabs = ALL_TABS) -> List[WorkflowType]: 76 | """ 77 | Get the list of currently registered workflow types 78 | To update the workflow types list, see `add_workflow_type` or `set_workflow_types` 79 | 80 | Args: 81 | tabs (Tabs): Whitelist of tabs 82 | Returns: 83 | List of workflow types that are defined on at least one of the given tabs 84 | """ 85 | if isinstance(tabs, str): 86 | tabs = (tabs,) 87 | 88 | return [ 89 | workflow_type 90 | for workflow_type in getattr(global_state, 'workflow_types', []) 91 | if tabs == ALL_TABS or any(tab in tabs for tab in workflow_type.tabs) 92 | ] 93 | 94 | 95 | def add_workflow_type(new_workflow_type: WorkflowType) -> None: 96 | """ 97 | Register a new workflow type 98 | You cannot call this function after the extension ui has been created 99 | 100 | Args: 101 | new_workflow_type (WorkflowType): new workflow type to add 102 | Raises: 103 | ValueError: If the id or display_name of new_workflow_type has already been registered 104 | NotImplementedError: if the workflow types list is modified after the ui has been instantiated 105 | """ 106 | workflow_types = get_workflow_types() 107 | 108 | for existing_workflow_type in workflow_types: 109 | if existing_workflow_type.base_id == new_workflow_type.base_id: 110 | raise ValueError(f'The id {new_workflow_type.base_id} already exists') 111 | if existing_workflow_type.display_name == new_workflow_type.display_name: 112 | raise ValueError(f'The display name {new_workflow_type.display_name} is already in use by workflow type {existing_workflow_type.base_id}') 113 | 114 | if getattr(global_state, 'is_ui_instantiated', False): 115 | raise NotImplementedError('Cannot modify workflow types after the ui has been instantiated') 116 | 117 | workflow_types.append(new_workflow_type) 118 | set_workflow_types(workflow_types) 119 | 120 | 121 | def set_workflow_types(workflow_types: List[WorkflowType]) -> None: 122 | """ 123 | Set the list of currently registered workflow types 124 | You cannot call this function after the extension ui has been created 125 | Args: 126 | workflow_types (List[WorkflowType]): the new workflow types 127 | Raises: 128 | NotImplementedError: if the workflow types list is modified after the ui has been instantiated 129 | Notes: 130 | A deep copy of workflow_types is used when calling this function from the comfyui process 131 | No copy is made when calling this function from the webui process 132 | """ 133 | if getattr(global_state, 'is_ui_instantiated', False): 134 | raise NotImplementedError('Cannot modify workflow types after the ui has been instantiated') 135 | 136 | global_state.workflow_types = workflow_types 137 | 138 | 139 | def clear_workflow_types() -> None: 140 | """ 141 | Clear the list of currently registered workflow types 142 | You cannot call this function after the extension ui has been created 143 | Raises: 144 | NotImplementedError: if the workflow types list is modified after the ui has been instantiated 145 | """ 146 | set_workflow_types([]) 147 | 148 | 149 | def get_workflow_type_ids(tabs: Tabs = ALL_TABS) -> List[str]: 150 | """ 151 | Get all workflow type ids of all currently registered workflow types 152 | Multiple ids can be assigned to each workflow type depending on how many tabs it is to be displayed on 153 | 154 | Args: 155 | tabs (Tabs): Whitelist of tabs for which to return the ids 156 | Returns: 157 | List of ids for the given tabs 158 | """ 159 | res = [] 160 | 161 | for workflow_type in get_workflow_types(tabs): 162 | res.extend(workflow_type.get_ids(tabs)) 163 | 164 | return res 165 | 166 | 167 | def get_workflow_type_display_names(tabs: Tabs = ALL_TABS) -> List[str]: 168 | """ 169 | Get the list of display names needed for the given tabs 170 | 171 | Args: 172 | tabs (Tabs): Whitelist of tabs for which to return the display names 173 | Returns: 174 | List of display names for the given tabs 175 | """ 176 | return [workflow_type.display_name for workflow_type in get_workflow_types(tabs)] 177 | 178 | 179 | def get_default_workflow_json(workflow_type_id: str) -> str: 180 | """ 181 | Get the default workflow for the given workflow type id 182 | 183 | Args: 184 | workflow_type_id (str): The workflow type id for which to get the default workflow 185 | Returns: 186 | The default workflow, or None if there is no default workflow for the given workflow type 187 | Raises: 188 | ValueError: If workflow_type_id does not exist 189 | """ 190 | for workflow_type in get_workflow_types(): 191 | if workflow_type_id in workflow_type.get_ids(): 192 | return workflow_type.default_workflow 193 | 194 | raise ValueError(workflow_type_id) 195 | 196 | 197 | def is_workflow_type_enabled(workflow_type_id: str) -> bool: 198 | return ( 199 | getattr(global_state, 'enable', True) and 200 | getattr(global_state, 'enabled_workflow_type_ids', {}).get(workflow_type_id, False) 201 | ) 202 | 203 | 204 | @ipc.restrict_to_process('webui') 205 | def run_workflow( 206 | workflow_type: WorkflowType, 207 | tab: str, 208 | batch_input: Any, 209 | queue_front: Optional[bool] = None, 210 | identity_on_error: Optional[bool] = False 211 | ) -> List[Any]: 212 | """ 213 | Run a comfyui workflow synchronously 214 | 215 | Args: 216 | workflow_type (WorkflowType): Target workflow type to run 217 | tab (str): The tab on which to run the workflow type. The workflow type must be present on the tab 218 | batch_input (Any): Batch object to pass as input to the workflow. The number of elements in this batch object will be the size of the comfyui batch. 219 | The particular type of batch_input depends on **workflow_type.input_types**: 220 | - if **input_types** is a dict, **batch_input** should be a dict. For each **k, v** pair of **batch_input**, **v** should match the type expected by **input_types[k]**. 221 | - if **input_types** is a tuple, **batch_input** should be a tuple. Each element **v** at index **i** of **batch_input** should match the type expected by **input_types[i]**. 222 | - if **input_types** is a str, **batch_input** should be a single value that should match the type expected by **input_types**. 223 | queue_front (Optional[bool]): Whether to queue the workflow before or after the currently queued workflows 224 | identity_on_error (Optional[bool]): Whether to return batch_input (converted to the type expected by workflow_type.types) instead of raising a RuntimeError when the workflow fails to run 225 | Returns: 226 | The outputs of the workflow, as a list. 227 | Each element of the list corresponds to one output node in the workflow. 228 | You can expect multiple values when a user has multiple output nodes in their workflow. 229 | The particular type of each returned element depends on workflow_type.types: 230 | - **types** is a dict: each element is a dict. For each **k, v** pair, **v** matches the type expected by **types[k]** 231 | - **types** is a tuple: each element is a tuple. For each element **v** at index **i**, **v** matches the type expected by **types[i]** 232 | - **types** is a str: each element is a single value that matches the type expected by **types** 233 | Each element of the list will have the same batch size as **batch_input** 234 | Raises: 235 | ValueError: If workflow_type is not present on the given tab 236 | TypeError: If the type of batch_input does not match the type expected by workflow_type.input_types 237 | RuntimeError: If identity_on_error is False and workflow execution fails for any reason 238 | AssertionError: If multiple candidate ids exist for workflow_type 239 | """ 240 | from lib_comfyui.comfyui.iframe_requests import ComfyuiIFrameRequests 241 | 242 | candidate_ids = workflow_type.get_ids(tab) 243 | assert len(candidate_ids) <= 1, ( 244 | f'Found multiple candidate workflow type ids for tab {tab} and workflow type {workflow_type.pretty_str()}: {candidate_ids}\n' 245 | 'The workflow type is likely to be incorrectly configured' 246 | ) 247 | 248 | if queue_front is None: 249 | queue_front = getattr(global_state, 'queue_front', True) 250 | 251 | batch_input_args, input_types = _normalize_to_tuple(batch_input, workflow_type.input_types) 252 | 253 | if not candidate_ids: 254 | raise ValueError(f'The workflow type {workflow_type.pretty_str()} does not exist on tab {tab}. ' 255 | f'Valid tabs for the given workflow type are: {workflow_type.tabs}') 256 | 257 | workflow_type_id = candidate_ids[0] 258 | 259 | try: 260 | ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw( 261 | workflow_type_id, 262 | workflow_type.max_amount_of_FromWebui_nodes, 263 | workflow_type.max_amount_of_ToWebui_nodes, 264 | ) 265 | except RuntimeError as e: 266 | if not identity_on_error: 267 | raise e 268 | 269 | traceback.print_exception(e) 270 | return batch_input_args 271 | 272 | try: 273 | if not is_workflow_type_enabled(workflow_type_id): 274 | raise WorkflowTypeDisabled(f'Workflow type {workflow_type.pretty_str()} is not enabled on tab {tab}') 275 | 276 | batch_output_params = ComfyuiIFrameRequests.start_workflow_sync( 277 | batch_input_args=batch_input_args, 278 | workflow_type_id=workflow_type_id, 279 | workflow_input_types=input_types, 280 | queue_front=queue_front, 281 | ) 282 | except RuntimeError as e: 283 | if not identity_on_error: 284 | raise e 285 | 286 | # don't print just because the workflow type is disabled 287 | if not isinstance(e, WorkflowTypeDisabled): 288 | traceback.print_exception(e) 289 | 290 | if not workflow_type.is_same_io(): 291 | print('[sd-webui-comfyui]', f'Returning input of type {workflow_type.input_types}, which likely does not match the expected output type {workflow_type.types}', file=sys.stderr) 292 | 293 | # denormalize tuple -> tuple|str|dict 294 | if isinstance(workflow_type.types, tuple): 295 | return [batch_input_args] 296 | elif isinstance(workflow_type.types, str): 297 | return [batch_input_args[0]] 298 | else: 299 | return [dict(zip(workflow_type.types.keys(), batch_input_args))] 300 | 301 | # denormalize dict -> tuple|str|dict 302 | if isinstance(workflow_type.types, tuple): 303 | return [tuple(params.values()) for params in batch_output_params] 304 | elif isinstance(workflow_type.types, str): 305 | return [next(iter(params.values())) for params in batch_output_params] 306 | else: 307 | return batch_output_params 308 | 309 | 310 | class WorkflowTypeDisabled(RuntimeError): 311 | pass 312 | 313 | 314 | def _normalize_to_tuple(batch_input, input_types): 315 | if isinstance(input_types, str): 316 | return (batch_input,), (input_types,) 317 | elif isinstance(input_types, tuple): 318 | if not isinstance(batch_input, tuple): 319 | raise TypeError(f'batch_input should be tuple but is instead {type(batch_input)}') 320 | 321 | if len(batch_input) != len(input_types): 322 | raise TypeError( 323 | f'batch_input received {len(batch_input)} values instead of {len(input_types)} (signature is {input_types})') 324 | 325 | return batch_input, input_types 326 | elif isinstance(input_types, dict): 327 | if not isinstance(batch_input, dict): 328 | raise TypeError(f'batch_input should be dict but is instead {type(batch_input)}') 329 | 330 | expected_keys = set(input_types.keys()) 331 | actual_keys = set(batch_input.keys()) 332 | if expected_keys - actual_keys: 333 | raise TypeError(f'batch_input is missing keys: {expected_keys - actual_keys}') 334 | 335 | # convert to tuple in the same order as the items in input_types 336 | return ( 337 | tuple(batch_input[k] for k in input_types.keys()), 338 | tuple(input_types.values()) 339 | ) 340 | else: 341 | raise TypeError(f'batch_input should be str, tuple or dict but is instead {type(batch_input)}') 342 | -------------------------------------------------------------------------------- /lib_comfyui/find_extensions.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | from lib_comfyui import ipc 3 | 4 | 5 | @ipc.run_in_process('webui') 6 | def get_extension_paths_to_load(): 7 | from modules.extensions import list_extensions, active 8 | list_extensions() 9 | active_paths = [e.path for e in active()] 10 | root_node_paths = [] 11 | root_script_paths = [] 12 | 13 | for path in active_paths: 14 | root_nodes = os.path.join(path, 'comfyui_custom_nodes') 15 | root_scripts = os.path.join(path, 'comfyui_custom_scripts') 16 | if os.path.exists(root_nodes): 17 | root_node_paths.append(root_nodes) 18 | if os.path.exists(root_scripts): 19 | root_script_paths.append(root_scripts) 20 | 21 | return root_node_paths, root_script_paths 22 | -------------------------------------------------------------------------------- /lib_comfyui/global_state.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from types import ModuleType 3 | from typing import List, Tuple, Optional, Dict, Any 4 | from lib_comfyui import ipc 5 | 6 | 7 | enabled: bool 8 | is_ui_instantiated: bool 9 | queue_front: bool 10 | focused_webui_client_id: Optional[str] = None 11 | 12 | workflow_types: List 13 | enabled_workflow_type_ids: Dict[str, bool] 14 | node_inputs: Tuple[Any, ...] 15 | node_outputs: List[Dict[str, Any]] 16 | current_workflow_input_types: Tuple[str, ...] 17 | 18 | ipc_strategy_class: type 19 | ipc_strategy_class_name: str 20 | comfyui_graceful_termination_timeout: float 21 | 22 | reverse_proxy_enabled: bool 23 | 24 | last_positive_prompt: str 25 | last_negative_prompt: str 26 | 27 | 28 | class GlobalState(ModuleType): 29 | __state = {} 30 | 31 | def __init__(self, glob): 32 | super().__init__(__name__) 33 | for k, v in glob.items(): 34 | setattr(self, k, v) 35 | 36 | def __getattr__(self, item): 37 | if item in ['__file__']: 38 | return globals()[item] 39 | 40 | return GlobalState.getattr(item) 41 | 42 | @staticmethod 43 | @ipc.run_in_process('webui') 44 | def getattr(item): 45 | try: 46 | return GlobalState.__state[item] 47 | except KeyError: 48 | raise AttributeError 49 | 50 | def __setattr__(self, item, value): 51 | GlobalState.setattr(item, value) 52 | 53 | @staticmethod 54 | @ipc.run_in_process('webui') 55 | def setattr(item, value): 56 | GlobalState.__state[item] = value 57 | 58 | def __delattr__(self, item): 59 | GlobalState.delattr(item) 60 | 61 | @staticmethod 62 | @ipc.run_in_process('webui') 63 | def delattr(item): 64 | del GlobalState.__state[item] 65 | 66 | def __contains__(self, item): 67 | return GlobalState.contains(item) 68 | 69 | @staticmethod 70 | @ipc.run_in_process('webui') 71 | def contains(item): 72 | return item in GlobalState.__state 73 | 74 | 75 | sys.modules[__name__] = GlobalState(globals()) 76 | -------------------------------------------------------------------------------- /lib_comfyui/ipc/__init__.py: -------------------------------------------------------------------------------- 1 | from . import callback 2 | from . import payload 3 | from . import strategies 4 | 5 | import gc 6 | import importlib 7 | import sys 8 | import time 9 | import logging 10 | 11 | 12 | def run_in_process(process_id): 13 | def annotation(function): 14 | def wrapper(*args, **kwargs): 15 | global current_process_id 16 | if process_id == current_process_id: 17 | return function(*args, **kwargs) 18 | else: 19 | start = time.time() 20 | res = current_callback_proxies[process_id].get(args=(function.__module__, function.__qualname__, args, kwargs)) 21 | logging.debug( 22 | '[sd-webui-comfyui] IPC call %s -> %s %s:\t%s', 23 | current_process_id, process_id, 24 | time.time() - start, 25 | f'{function.__module__}.{function.__qualname__}(*{args}, **{kwargs})' 26 | ) 27 | return res 28 | 29 | return wrapper 30 | 31 | return annotation 32 | 33 | 34 | def restrict_to_process(process_id): 35 | def annotation(function): 36 | def wrapper(*args, **kwargs): 37 | global current_process_id 38 | if process_id != current_process_id: 39 | raise RuntimeError(f'Can only call function {function.__module__}.{function.__qualname__} from process {process_id}. Current process is {current_process_id}') 40 | 41 | return function(*args, **kwargs) 42 | 43 | return wrapper 44 | 45 | return annotation 46 | 47 | 48 | def call_fully_qualified(module_name, qualified_name, args, kwargs): 49 | module_parts = module_name.split('.') 50 | try: 51 | module = sys.modules[module_parts[0]] 52 | for part in module_parts[1:]: 53 | module = getattr(module, part) 54 | except (AttributeError, KeyError): 55 | source_module = module_parts[-1] 56 | module = importlib.import_module(module_name, source_module) 57 | 58 | function = module 59 | for name in qualified_name.split('.'): 60 | function = getattr(function, name) 61 | return function(*args, **kwargs) 62 | 63 | 64 | current_process_id = 'webui' 65 | current_callback_listeners = {} 66 | current_callback_proxies = {} 67 | 68 | 69 | def start_callback_listeners(): 70 | assert not callback_listeners_started() 71 | for callback_listener in current_callback_listeners.values(): 72 | callback_listener.start() 73 | 74 | print('[sd-webui-comfyui]', 'Started callback listeners for process', current_process_id) 75 | 76 | 77 | def stop_callback_listeners(): 78 | assert callback_listeners_started() 79 | for callback_listener in current_callback_listeners.values(): 80 | callback_listener.stop() 81 | 82 | current_callback_proxies.clear() 83 | current_callback_listeners.clear() 84 | gc.collect() 85 | 86 | print('[sd-webui-comfyui]', 'Stopped callback listeners for process', current_process_id) 87 | 88 | 89 | def callback_listeners_started(): 90 | return any(callback_listener.is_running() for callback_listener in current_callback_listeners.values()) 91 | -------------------------------------------------------------------------------- /lib_comfyui/ipc/callback.py: -------------------------------------------------------------------------------- 1 | import threading 2 | from typing import Callable 3 | from lib_comfyui.ipc.payload import IpcSender, IpcReceiver 4 | from lib_comfyui.ipc.strategies import IpcStrategy 5 | 6 | 7 | class CallbackWatcher: 8 | def __init__(self, callback: Callable, name: str, strategy_factory: Callable[[str], IpcStrategy], clear_on_init: bool = False, clear_on_del: bool = True): 9 | self._callback = callback 10 | self._res_sender = IpcSender( 11 | f'res_{name}', 12 | strategy_factory, 13 | clear_on_init=clear_on_init, 14 | clear_on_del=clear_on_del, 15 | ) 16 | self._args_receiver = IpcReceiver( 17 | f'args_{name}', 18 | strategy_factory, 19 | clear_on_init=clear_on_init, 20 | clear_on_del=clear_on_del, 21 | ) 22 | self._producer_thread = None 23 | 24 | def start(self): 25 | def thread_loop(): 26 | while self._producer_thread.is_running(): 27 | self.attend_consumer(timeout=0.5) 28 | 29 | self._producer_thread = StoppableThread(target=thread_loop, daemon=True) 30 | self._producer_thread.start() 31 | 32 | def stop(self): 33 | if self._producer_thread is None: 34 | return 35 | 36 | self._producer_thread.join() 37 | self._producer_thread = None 38 | 39 | def is_running(self): 40 | return self._producer_thread and self._producer_thread.is_running() 41 | 42 | def attend_consumer(self, timeout: float = None): 43 | try: 44 | args, kwargs = self._args_receiver.recv(timeout=timeout) 45 | except TimeoutError: 46 | return 47 | 48 | try: 49 | self._res_sender.send(self._callback(*args, **kwargs)) 50 | except Exception as e: 51 | self._res_sender.send(RemoteError(e)) 52 | 53 | 54 | class CallbackProxy: 55 | def __init__(self, name: str, strategy_factory, clear_on_init: bool = False, clear_on_del: bool = True): 56 | self._res_receiver = IpcReceiver( 57 | f'res_{name}', 58 | strategy_factory, 59 | clear_on_init=clear_on_init, 60 | clear_on_del=clear_on_del, 61 | ) 62 | self._args_sender = IpcSender( 63 | f'args_{name}', 64 | strategy_factory, 65 | clear_on_init=clear_on_init, 66 | clear_on_del=clear_on_del, 67 | ) 68 | 69 | def get(self, args=None, kwargs=None): 70 | self._args_sender.send((args if args is not None else (), kwargs if kwargs is not None else {})) 71 | res = self._res_receiver.recv() 72 | if isinstance(res, RemoteError): 73 | raise res.error from res 74 | else: 75 | return res 76 | 77 | 78 | class RemoteError(Exception): 79 | def __init__(self, error): 80 | self.error = error 81 | 82 | 83 | class StoppableThread(threading.Thread): 84 | def __init__(self, *args, **kwargs): 85 | super().__init__(*args, **kwargs) 86 | self._stop_event = threading.Event() 87 | 88 | def stop(self): 89 | self._stop_event.set() 90 | 91 | def join(self, *args, **kwargs) -> None: 92 | self.stop() 93 | super().join(*args, **kwargs) 94 | 95 | def is_running(self): 96 | return not self._stop_event.is_set() 97 | -------------------------------------------------------------------------------- /lib_comfyui/ipc/payload.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import math 3 | import pickle 4 | import portalocker 5 | import tempfile 6 | import time 7 | import logging 8 | from pathlib import Path 9 | from typing import Optional, Any, Callable 10 | from lib_comfyui.ipc.strategies import IpcStrategy 11 | 12 | 13 | class IpcPayload: 14 | def __init__(self, name, strategy_factory: Callable[[str], IpcStrategy], lock_directory: str = None, clear_on_init: bool = False, clear_on_del: bool = True): 15 | self._name = name 16 | self._strategy = strategy_factory(f'ipc_payload_{name}') 17 | lock_directory = tempfile.gettempdir() if lock_directory is None else lock_directory 18 | self._lock_path = Path(lock_directory, f'ipc_payload_{name}') 19 | self._clear_on_del = clear_on_del 20 | 21 | if clear_on_init: 22 | with self.get_lock() as lock_file: 23 | self._strategy.clear(lock_file) 24 | 25 | def __del__(self): 26 | if self._clear_on_del: 27 | lock = self.get_lock() 28 | with lock as lock_file: 29 | self._strategy.clear(lock_file) 30 | 31 | def get_lock(self, timeout: Optional[float] = None, mode: str = 'wb+'): 32 | return portalocker.Lock( 33 | self._lock_path, 34 | mode=mode, 35 | timeout=timeout, 36 | check_interval=0.02, 37 | flags=portalocker.LOCK_EX | (portalocker.LOCK_NB * int(timeout is not None)), 38 | ) 39 | 40 | 41 | class IpcSender(IpcPayload): 42 | def send(self, value: Any): 43 | with self.get_lock() as lock_file: 44 | logging.debug(f'IPC payload {self._name}\tsend value: {value}') 45 | self._strategy.set_data(lock_file, pickle.dumps(value)) 46 | 47 | 48 | class IpcReceiver(IpcPayload): 49 | def recv(self, timeout: Optional[float] = None) -> Any: 50 | current_time = time.time() 51 | end_time = (current_time + timeout) if timeout is not None else math.inf 52 | 53 | while current_time < end_time: 54 | lock = self.get_lock(timeout=end_time - current_time, mode='rb+') 55 | 56 | try: 57 | with lock as lock_file: 58 | if self._strategy.is_empty(lock_file): 59 | raise FileEmptyException 60 | 61 | with self._strategy.get_data(lock_file) as data, restore_torch_load(): 62 | value = pickle.loads(data) 63 | del data 64 | 65 | logging.debug(f'IPC payload {self._name}\treceive value: {value}') 66 | return value 67 | except FileEmptyException: 68 | time.sleep(lock.check_interval) # yuck 69 | current_time = time.time() 70 | continue 71 | except portalocker.LockException: 72 | break 73 | 74 | raise TimeoutError 75 | 76 | 77 | class FileEmptyException(Exception): 78 | pass 79 | 80 | 81 | @contextlib.contextmanager 82 | def restore_torch_load(): 83 | from lib_comfyui import ipc 84 | import torch 85 | original_torch_load = torch.load 86 | 87 | if ipc.current_process_id == 'webui': 88 | try: 89 | from modules import safe 90 | torch.load = safe.unsafe_torch_load 91 | del safe 92 | except ImportError: 93 | pass 94 | 95 | yield 96 | 97 | torch.load = original_torch_load 98 | -------------------------------------------------------------------------------- /lib_comfyui/ipc/strategies.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import dataclasses 3 | import os 4 | import pickle 5 | import platform 6 | from abc import ABC, abstractmethod 7 | from multiprocessing.shared_memory import SharedMemory 8 | from typing import IO, Union 9 | 10 | 11 | BytesLike = Union[bytes, bytearray, memoryview] 12 | 13 | 14 | __all__ = [ 15 | 'FileSystemIpcStrategy', 16 | 'SharedMemoryIpcStrategy', 17 | 'OsFriendlyIpcStrategy', 18 | ] 19 | 20 | 21 | class IpcStrategy(ABC): 22 | @abstractmethod 23 | def is_empty(self, lock_file: IO) -> bool: 24 | pass 25 | 26 | @abstractmethod 27 | def set_data(self, lock_file: IO, data: BytesLike) -> None: 28 | pass 29 | 30 | @abstractmethod 31 | @contextlib.contextmanager 32 | def get_data(self, lock_file: IO) -> BytesLike: 33 | pass 34 | 35 | @abstractmethod 36 | def clear(self, lock_file: IO) -> None: 37 | pass 38 | 39 | 40 | class FileSystemIpcStrategy(IpcStrategy): 41 | def __init__(self, _name: str): 42 | pass 43 | 44 | def is_empty(self, lock_file: IO) -> bool: 45 | lock_file.seek(0, os.SEEK_END) 46 | return lock_file.tell() == 0 47 | 48 | def set_data(self, lock_file: IO, data: BytesLike): 49 | lock_file.write(data) 50 | 51 | @contextlib.contextmanager 52 | def get_data(self, lock_file: IO) -> BytesLike: 53 | lock_file.seek(0) 54 | yield lock_file.read() 55 | self.clear(lock_file) 56 | 57 | def clear(self, lock_file: IO): 58 | lock_file.seek(0) 59 | lock_file.truncate() 60 | 61 | 62 | class SharedMemoryIpcStrategy(IpcStrategy): 63 | def __init__(self, shm_name: str): 64 | self._shm_name = shm_name 65 | self._shm = None 66 | 67 | @dataclasses.dataclass 68 | class Metadata: 69 | is_empty: bool 70 | size: int 71 | 72 | def _get_metadata(self, lock_file: IO) -> Metadata: 73 | file_size = lock_file.seek(0, os.SEEK_END) 74 | if file_size == 0: 75 | return self.Metadata(is_empty=True, size=0) 76 | 77 | lock_file.seek(0) 78 | try: 79 | return pickle.loads(lock_file.read()) 80 | except (ModuleNotFoundError, AttributeError): 81 | return self.Metadata(is_empty=True, size=0) 82 | 83 | def _set_metadata(self, lock_file: IO, metadata: Metadata): 84 | lock_file.seek(0) 85 | serialized = pickle.dumps(metadata) 86 | lock_file.write(serialized) 87 | lock_file.truncate(len(serialized)) 88 | 89 | def is_empty(self, lock_file: IO) -> bool: 90 | return self._get_metadata(lock_file).is_empty 91 | 92 | def set_data(self, lock_file: IO, data: BytesLike): 93 | metadata = self._get_metadata(lock_file) 94 | assert metadata.is_empty, f'data of shared memory IPC payload {self._shm_name} has not yet been read' 95 | 96 | data_len = len(data) 97 | self._clear_shm() 98 | self._shm = SharedMemory(name=self._shm_name, create=True, size=data_len) 99 | self._shm.buf[:data_len] = data 100 | self._set_metadata(lock_file, self.Metadata(is_empty=False, size=data_len)) 101 | 102 | @contextlib.contextmanager 103 | def get_data(self, lock_file: IO) -> BytesLike: 104 | metadata = self._get_metadata(lock_file) 105 | assert not metadata.is_empty, f'metadata not found for shared memory IPC payload {self._shm_name}' 106 | 107 | shm = SharedMemory(name=self._shm_name) 108 | yield shm.buf[:metadata.size] 109 | shm.close() 110 | shm.unlink() 111 | self.clear(lock_file) 112 | 113 | def clear(self, lock_file: IO): 114 | self._set_metadata(lock_file, self.Metadata(is_empty=True, size=0)) 115 | 116 | def _clear_shm(self): 117 | try: 118 | if self._shm is None: 119 | self._shm = SharedMemory(name=self._shm_name) 120 | self._shm.close() 121 | self._shm.unlink() 122 | except FileNotFoundError: 123 | pass 124 | 125 | self._shm = None 126 | 127 | 128 | if platform.system().lower() == 'linux': 129 | OsFriendlyIpcStrategy = FileSystemIpcStrategy 130 | else: 131 | OsFriendlyIpcStrategy = SharedMemoryIpcStrategy 132 | -------------------------------------------------------------------------------- /lib_comfyui/platform_utils.py: -------------------------------------------------------------------------------- 1 | import platform 2 | 3 | 4 | def is_windows(): 5 | return 'windows' in platform.system().lower() 6 | -------------------------------------------------------------------------------- /lib_comfyui/torch_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Union 3 | 4 | 5 | def deep_to( 6 | tensor: Union[torch.Tensor, dict, list], 7 | *args, 8 | **kwargs, 9 | ) -> Union[torch.Tensor, dict, list]: 10 | if isinstance(tensor, torch.Tensor): 11 | tensor = tensor.to(*args, **kwargs) 12 | elif isinstance(tensor, dict): 13 | for k, v in tensor.items(): 14 | tensor[k] = deep_to(v, *args, **kwargs) 15 | elif isinstance(tensor, list): 16 | for i, v in enumerate(tensor): 17 | tensor[i] = deep_to(v, *args, **kwargs) 18 | elif isinstance(tensor, tuple): 19 | res = () 20 | for v in tensor: 21 | res += (deep_to(v, *args, **kwargs),) 22 | tensor = res 23 | 24 | return tensor 25 | -------------------------------------------------------------------------------- /lib_comfyui/webui/accordion.py: -------------------------------------------------------------------------------- 1 | import json 2 | import operator 3 | from typing import Tuple 4 | 5 | import gradio as gr 6 | from lib_comfyui import external_code, global_state 7 | from lib_comfyui.webui import gradio_utils, settings 8 | from lib_comfyui.comfyui import iframe_requests 9 | 10 | 11 | class AccordionInterface: 12 | def __init__(self, get_elem_id, tab): 13 | from modules import ui 14 | 15 | self.tab = tab 16 | 17 | self.workflow_types = external_code.get_workflow_types(self.tab) 18 | self.first_workflow_type = self.workflow_types[0] 19 | self.workflow_type_ids = { 20 | workflow_type.display_name: workflow_type.get_ids(self.tab)[0] 21 | for workflow_type in self.workflow_types 22 | } 23 | 24 | self.accordion = gr.Accordion( 25 | label='ComfyUI', 26 | open=False, 27 | elem_id=get_elem_id('accordion'), 28 | ) 29 | 30 | self.iframes = gr.HTML(value=self.get_iframes_html()) 31 | self.enabled_checkbox = gr.Checkbox( 32 | label='Enable', 33 | elem_id=get_elem_id('enable'), 34 | value=False, 35 | ) 36 | self.current_display_name = gr.Dropdown( 37 | label='Edit workflow type', 38 | choices=[workflow_type.display_name for workflow_type in self.workflow_types], 39 | value=self.first_workflow_type.display_name, 40 | elem_id=get_elem_id('current_display_name'), 41 | ) 42 | self.queue_front = gr.Checkbox( 43 | label='Queue front', 44 | elem_id=get_elem_id('queue_front'), 45 | value=True, 46 | ) 47 | self.refresh_button = gr.Button( 48 | value=f'{ui.refresh_symbol} Reload ComfyUI interfaces (client side)', 49 | elem_id=get_elem_id('refresh_button'), 50 | ) 51 | 52 | self.enabled_display_names = gradio_utils.ExtensionDynamicProperty( 53 | value=[], 54 | ) 55 | self.enabled_ids = gradio_utils.ExtensionDynamicProperty( 56 | value={ 57 | workflow_type_id: False 58 | for workflow_type_id in self.workflow_type_ids.values() 59 | }, 60 | ) 61 | self.clear_enabled_display_names_button = gr.Button( 62 | elem_id=get_elem_id('clear_enabled_display_names'), 63 | visible=False, 64 | ) 65 | 66 | self._rendered = False 67 | 68 | def get_iframes_html(self) -> str: 69 | comfyui_client_url = settings.get_comfyui_iframe_url() 70 | first_workflow_type_id = self.workflow_type_ids[self.first_workflow_type.display_name] 71 | 72 | iframes = [] 73 | for workflow_type_id in external_code.get_workflow_type_ids(self.tab): 74 | html_classes = [] 75 | if workflow_type_id == first_workflow_type_id: 76 | html_classes.append('comfyui-workflow-type-visible') 77 | 78 | iframes.append(f""" 79 | 85 | """) 86 | 87 | return f""" 88 |
89 | {''.join(iframes)} 90 |
91 | """ 92 | 93 | def arrange_components(self): 94 | if self._rendered: 95 | return 96 | 97 | with self.accordion.render(): 98 | with gr.Row(): 99 | self.iframes.render() 100 | 101 | with gr.Row(): 102 | with gr.Column(): 103 | self.enabled_checkbox.render() 104 | self.current_display_name.render() 105 | 106 | with gr.Column(): 107 | self.queue_front.render() 108 | self.refresh_button.render() 109 | 110 | self.enabled_display_names.render() 111 | self.enabled_ids.render() 112 | self.clear_enabled_display_names_button.render() 113 | 114 | def connect_events(self): 115 | if self._rendered: 116 | return 117 | 118 | self.refresh_button.click( 119 | fn=None, 120 | _js='reloadComfyuiIFrames' 121 | ) 122 | self.clear_enabled_display_names_button.click( 123 | fn=list, 124 | outputs=[self.enabled_display_names], 125 | ) 126 | 127 | self.activate_current_workflow_type() 128 | self.activate_enabled_workflow_types() 129 | self._rendered = True 130 | 131 | def get_script_ui_components(self) -> Tuple[gr.components.Component, ...]: 132 | return self.queue_front, self.enabled_ids 133 | 134 | def setup_infotext_fields(self, script): 135 | workflows_infotext_field = gr.HTML(visible=False) 136 | workflows_infotext_field.change( 137 | fn=self.on_infotext_change, 138 | inputs=[workflows_infotext_field, self.current_display_name], 139 | outputs=[workflows_infotext_field, self.enabled_display_names, self.enabled_checkbox], 140 | ) 141 | script.infotext_fields = [(workflows_infotext_field, 'ComfyUI Workflows')] 142 | 143 | def activate_current_workflow_type(self): 144 | current_workflow_type_id = gr.HTML( 145 | value=self.workflow_type_ids[self.first_workflow_type.display_name], 146 | visible=False, 147 | ) 148 | self.current_display_name.change( 149 | fn=self.workflow_type_ids.get, 150 | inputs=[self.current_display_name], 151 | outputs=[current_workflow_type_id], 152 | ) 153 | current_workflow_type_id.change( 154 | fn=None, 155 | _js='changeDisplayedWorkflowType', 156 | inputs=[current_workflow_type_id], 157 | ) 158 | 159 | def activate_enabled_workflow_types(self): 160 | self.enabled_display_names.change( 161 | fn=self.display_names_to_enabled_ids, 162 | inputs=[self.enabled_display_names], 163 | outputs=[self.enabled_ids], 164 | ) 165 | 166 | self.activate_enabled_checkbox() 167 | self.activate_enabled_display_names_colors() 168 | 169 | def activate_enabled_display_names_colors(self): 170 | style_body = '''{ 171 | color: greenyellow !important; 172 | font-weight: bold; 173 | }''' 174 | 175 | dropdown_input_style = gr.HTML() 176 | for comp in (self.enabled_display_names, self.current_display_name): 177 | comp.change( 178 | fn=lambda enabled_display_names, current_workflow_display_name: f'''''' if current_workflow_display_name in enabled_display_names else '', 181 | inputs=[self.enabled_display_names, self.current_display_name], 182 | outputs=[dropdown_input_style], 183 | ) 184 | 185 | dropdown_list_style = gr.HTML() 186 | self.enabled_display_names.change( 187 | fn=lambda enabled_display_names, current_workflow_display_name: f'''''', 193 | inputs=[self.enabled_display_names, self.current_display_name], 194 | outputs=[dropdown_list_style], 195 | ) 196 | 197 | def activate_enabled_checkbox(self): 198 | self.current_display_name.change( 199 | fn=operator.contains, 200 | inputs=[self.enabled_display_names, self.current_display_name], 201 | outputs=[self.enabled_checkbox], 202 | ) 203 | 204 | self.enabled_checkbox.select( 205 | fn=lambda enabled_display_names, current_workflow_display_name, enable: list( 206 | (operator.or_ if enable else operator.sub)( 207 | set(enabled_display_names), 208 | {current_workflow_display_name}, 209 | ) 210 | ), 211 | inputs=[self.enabled_display_names, self.current_display_name, self.enabled_checkbox], 212 | outputs=[self.enabled_display_names] 213 | ) 214 | 215 | def display_names_to_enabled_ids(self, enabled_display_names): 216 | return { 217 | self.workflow_type_ids[workflow_type.display_name]: workflow_type.display_name in enabled_display_names 218 | for workflow_type in self.workflow_types 219 | } 220 | 221 | def on_infotext_change(self, serialized_graphs, current_workflow_display_name) -> Tuple[dict, dict, dict]: 222 | if not serialized_graphs: 223 | return (gr.skip(),) * 3 224 | 225 | if not hasattr(global_state, 'enabled_workflow_type_ids'): 226 | global_state.enabled_workflow_type_ids = {} 227 | 228 | serialized_graphs = json.loads(serialized_graphs) 229 | workflow_graphs = { 230 | workflow_type.get_ids(self.tab)[0]: ( 231 | serialized_graphs.get(workflow_type.base_id, json.loads(workflow_type.default_workflow)), 232 | workflow_type, 233 | ) 234 | for workflow_type in self.workflow_types 235 | } 236 | 237 | new_enabled_display_names = [] 238 | for workflow_type_id, (graph, workflow_type) in workflow_graphs.items(): 239 | is_custom_workflow = workflow_type.base_id in serialized_graphs 240 | global_state.enabled_workflow_type_ids[workflow_type_id] = is_custom_workflow 241 | if is_custom_workflow: 242 | new_enabled_display_names.append(workflow_type.display_name) 243 | iframe_requests.set_workflow_graph(graph, workflow_type_id) 244 | 245 | return ( 246 | gr.update(value=''), 247 | gr.update(value=new_enabled_display_names), 248 | gr.update(value=current_workflow_display_name in new_enabled_display_names), 249 | ) 250 | -------------------------------------------------------------------------------- /lib_comfyui/webui/callbacks.py: -------------------------------------------------------------------------------- 1 | from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img 2 | from lib_comfyui import comfyui_process, ipc, global_state, external_code, default_workflow_types 3 | from lib_comfyui.webui import tab, settings, patches, reverse_proxy 4 | from lib_comfyui.comfyui import type_conversion 5 | 6 | 7 | @ipc.restrict_to_process('webui') 8 | def register_callbacks(): 9 | from modules import script_callbacks 10 | script_callbacks.on_ui_tabs(on_ui_tabs) 11 | script_callbacks.on_ui_settings(on_ui_settings) 12 | script_callbacks.on_after_component(on_after_component) 13 | script_callbacks.on_app_started(on_app_started) 14 | script_callbacks.on_script_unloaded(on_script_unloaded) 15 | script_callbacks.on_before_image_saved(on_before_image_saved) 16 | 17 | 18 | @ipc.restrict_to_process('webui') 19 | def on_ui_tabs(): 20 | return tab.create_tab() 21 | 22 | 23 | @ipc.restrict_to_process('webui') 24 | def on_ui_settings(): 25 | return settings.create_section() 26 | 27 | 28 | @ipc.restrict_to_process('webui') 29 | def on_after_component(*args, **kwargs): 30 | patches.watch_prompts(*args, **kwargs) 31 | settings.subscribe_update_button(*args, **kwargs) 32 | 33 | 34 | @ipc.restrict_to_process('webui') 35 | def on_app_started(_gr_root, fast_api): 36 | comfyui_process.start() 37 | reverse_proxy.create_comfyui_proxy(fast_api) 38 | 39 | 40 | @ipc.restrict_to_process('webui') 41 | def on_script_unloaded(): 42 | comfyui_process.stop() 43 | patches.clear_patches() 44 | global_state.is_ui_instantiated = False 45 | external_code.clear_workflow_types() 46 | 47 | 48 | @ipc.restrict_to_process('webui') 49 | def on_before_image_saved(params): 50 | if isinstance(params.p, StableDiffusionProcessingTxt2Img): 51 | tab = 'txt2img' 52 | elif isinstance(params.p, StableDiffusionProcessingImg2Img): 53 | tab = 'img2img' 54 | else: 55 | return 56 | 57 | if not external_code.is_workflow_type_enabled(default_workflow_types.before_save_image_workflow_type.get_ids(tab)[0]): 58 | return 59 | 60 | results = external_code.run_workflow( 61 | workflow_type=default_workflow_types.before_save_image_workflow_type, 62 | tab=tab, 63 | batch_input=type_conversion.webui_image_to_comfyui([params.image]), 64 | identity_on_error=True, 65 | ) 66 | 67 | params.image = type_conversion.comfyui_image_to_webui(results[0], return_tensors=False)[0] 68 | -------------------------------------------------------------------------------- /lib_comfyui/webui/gradio_utils.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | import gradio as gr 3 | import json 4 | 5 | 6 | def ExtensionDynamicProperty(value: Any, key: str = None, *, visible=False, **kwargs): 7 | extension_property = "sd_webui_comfyui" 8 | if key is not None: 9 | div_begin = f'
' 10 | div_end = '
' 11 | else: 12 | div_begin = '' 13 | div_end = '' 14 | 15 | def preprocess(x: Any) -> Any: 16 | return json.loads(x[len(div_begin):len(x) - len(div_end)]) 17 | 18 | def postprocess(y: str) -> Any: 19 | return f'{div_begin}{json.dumps(y)}{div_end}' 20 | 21 | component = gr.HTML( 22 | value=postprocess(value), 23 | visible=visible, 24 | **kwargs, 25 | ) 26 | component.preprocess = preprocess 27 | component.postprocess = postprocess 28 | 29 | return component 30 | -------------------------------------------------------------------------------- /lib_comfyui/webui/patches.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import re 3 | import sys 4 | 5 | from lib_comfyui import ipc, global_state, default_workflow_types, external_code 6 | from lib_comfyui.comfyui import type_conversion 7 | 8 | 9 | __original_create_sampler = None 10 | 11 | 12 | @ipc.restrict_to_process('webui') 13 | def apply_patches(): 14 | from modules import sd_samplers, generation_parameters_copypaste 15 | global __original_create_sampler 16 | 17 | __original_create_sampler = sd_samplers.create_sampler 18 | sd_samplers.create_sampler = functools.partial(create_sampler_hijack, original_function=sd_samplers.create_sampler) 19 | 20 | 21 | @ipc.restrict_to_process('webui') 22 | def watch_prompts(component, **kwargs): 23 | possible_elem_ids = { 24 | f'{tab}{negative}_prompt': bool(negative) 25 | for tab in ('txt2img', 'img2img') 26 | for negative in ('', '_neg') 27 | } 28 | event_listeners = ('change', 'blur') 29 | 30 | elem_id = getattr(component, 'elem_id', None) 31 | if elem_id in possible_elem_ids: 32 | attribute = f'last_{"negative" if possible_elem_ids[elem_id] else "positive"}_prompt' 33 | for event_listener in event_listeners: 34 | getattr(component, event_listener)( 35 | fn=lambda p: setattr(global_state, attribute, p), 36 | inputs=[component] 37 | ) 38 | 39 | 40 | @ipc.restrict_to_process('webui') 41 | def clear_patches(): 42 | from modules import sd_samplers, generation_parameters_copypaste 43 | global __original_create_sampler 44 | 45 | if __original_create_sampler is not None: 46 | sd_samplers.create_sampler = __original_create_sampler 47 | 48 | 49 | @ipc.restrict_to_process('webui') 50 | def create_sampler_hijack(name: str, model, original_function): 51 | sampler = original_function(name, model) 52 | sampler.sample_img2img = functools.partial(sample_img2img_hijack, original_function=sampler.sample_img2img) 53 | return sampler 54 | 55 | 56 | @ipc.restrict_to_process('webui') 57 | def sample_img2img_hijack(p, x, *args, original_function, **kwargs): 58 | from modules import processing 59 | workflow_type = default_workflow_types.preprocess_latent_workflow_type 60 | 61 | if not ( 62 | isinstance(p, processing.StableDiffusionProcessingImg2Img) and 63 | external_code.is_workflow_type_enabled(workflow_type.get_ids("img2img")[0]) 64 | ): 65 | return original_function(p, x, *args, **kwargs) 66 | 67 | processed_x = external_code.run_workflow( 68 | workflow_type=default_workflow_types.preprocess_latent_workflow_type, 69 | tab='img2img', 70 | batch_input=type_conversion.webui_latent_to_comfyui(x.to(device='cpu')), 71 | identity_on_error=True, 72 | ) 73 | verify_singleton(processed_x) 74 | x = type_conversion.comfyui_latent_to_webui(processed_x[0]).to(device=x.device) 75 | return original_function(p, x, *args, **kwargs) 76 | 77 | 78 | @ipc.restrict_to_process('webui') 79 | def patch_processing(p): 80 | from modules import processing 81 | 82 | p.sd_webui_comfyui_patches = getattr(p, 'sd_webui_comfyui_patches', set()) 83 | is_img2img = isinstance(p, processing.StableDiffusionProcessingImg2Img) 84 | 85 | if 'sample' not in p.sd_webui_comfyui_patches: 86 | p.sample = functools.partial(p_sample_patch, original_function=p.sample, is_img2img=is_img2img) 87 | p.sd_webui_comfyui_patches.add('sample') 88 | 89 | if is_img2img and 'init' not in p.sd_webui_comfyui_patches: 90 | p.init = functools.partial(p_img2img_init, original_function=p.init, p_ref=p) 91 | p.sd_webui_comfyui_patches.add('init') 92 | 93 | 94 | def p_sample_patch(*args, original_function, is_img2img, **kwargs): 95 | x = original_function(*args, **kwargs) 96 | tab = 'img2img' if is_img2img else 'txt2img' 97 | 98 | if not external_code.is_workflow_type_enabled(default_workflow_types.postprocess_latent_workflow_type.get_ids(tab)[0]): 99 | return x 100 | 101 | processed_x = external_code.run_workflow( 102 | workflow_type=default_workflow_types.postprocess_latent_workflow_type, 103 | tab=tab, 104 | batch_input=type_conversion.webui_latent_to_comfyui(x.to(device='cpu')), 105 | identity_on_error=True, 106 | ) 107 | verify_singleton(processed_x) 108 | return type_conversion.comfyui_latent_to_webui(processed_x[0]).to(device=x.device) 109 | 110 | 111 | def p_img2img_init(*args, original_function, p_ref, **kwargs): 112 | if not external_code.is_workflow_type_enabled(default_workflow_types.preprocess_workflow_type.get_ids("img2img")[0]): 113 | return original_function(*args, **kwargs) 114 | 115 | processed_images = external_code.run_workflow( 116 | workflow_type=default_workflow_types.preprocess_workflow_type, 117 | tab='img2img', 118 | batch_input=type_conversion.webui_image_to_comfyui(p_ref.init_images), 119 | identity_on_error=True, 120 | ) 121 | verify_singleton(processed_images) 122 | p_ref.init_images = type_conversion.comfyui_image_to_webui(processed_images[0]) 123 | return original_function(*args, **kwargs) 124 | 125 | 126 | def verify_singleton(l: list): 127 | if len(l) != 1: 128 | prefix = '\n[sd-webui-comfyui] ' 129 | print(f'{prefix}The last ComfyUI workflow returned {len(l)} batches instead of 1.' 130 | f'{prefix}This is likely due to the workflow not having exactly 1 "To Webui" node.' 131 | f'{prefix}Please verify that the workflow is valid.', file=sys.stderr) 132 | -------------------------------------------------------------------------------- /lib_comfyui/webui/paths.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import importlib 4 | from lib_comfyui import ipc 5 | 6 | 7 | @ipc.restrict_to_process('comfyui') 8 | def share_webui_folder_paths(): 9 | from folder_paths import add_model_folder_path 10 | webui_folder_paths = get_webui_folder_paths() 11 | for folder_id, folder_paths in webui_folder_paths.items(): 12 | for folder_path in folder_paths: 13 | add_model_folder_path(folder_id, folder_path) 14 | 15 | 16 | @ipc.run_in_process('webui') 17 | def get_webui_folder_paths() -> dict: 18 | from modules import paths, shared, sd_models 19 | return { 20 | 'checkpoints': [sd_models.model_path] + ([shared.cmd_opts.ckpt_dir] if shared.cmd_opts.ckpt_dir else []), 21 | 'vae': [os.path.join(paths.models_path, 'VAE')] + ([shared.cmd_opts.vae_dir] if shared.cmd_opts.vae_dir else []), 22 | 'vae_approx': [os.path.join(paths.models_path, "VAE-approx")], 23 | 'embeddings': [shared.cmd_opts.embeddings_dir], 24 | 'loras': [shared.cmd_opts.lora_dir], 25 | 'hypernetworks': [shared.cmd_opts.hypernetwork_dir], 26 | 'upscale_models': get_upscaler_paths(), 27 | 'controlnet': get_controlnet_paths() 28 | } 29 | 30 | 31 | # see https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/f865d3e11647dfd6c7b2cdf90dde24680e58acd8/modules/modelloader.py#L137 32 | @ipc.restrict_to_process('webui') 33 | def get_upscaler_paths(): 34 | from modules import shared, modelloader 35 | # We can only do this 'magic' method to dynamically load upscalers if they are referenced, 36 | # so we'll try to import any _model.py files before looking in __subclasses__ 37 | modules_dir = os.path.join(shared.script_path, "modules") 38 | for file in os.listdir(modules_dir): 39 | if "_model.py" in file: 40 | model_name = file.replace("_model.py", "") 41 | full_model = f"modules.{model_name}_model" 42 | try: 43 | importlib.import_module(full_model) 44 | except Exception: 45 | pass 46 | 47 | # some of upscaler classes will not go away after reloading their modules, and we'll end 48 | # up with two copies of those classes. The newest copy will always be the last in the list, 49 | # so we go from end to beginning and ignore duplicates 50 | used_classes = {} 51 | for cls in reversed(modelloader.Upscaler.__subclasses__()): 52 | class_name = str(cls) 53 | if class_name not in used_classes: 54 | used_classes[class_name] = cls 55 | 56 | upscaler_paths = set() 57 | for cls in reversed(used_classes.values()): 58 | name = cls.__name__ 59 | cmd_name = f"{name.lower().replace('upscaler', '')}_models_path" 60 | commandline_model_path = getattr(shared.cmd_opts, cmd_name, None) 61 | scaler = cls(commandline_model_path) 62 | scaler_path = commandline_model_path or scaler.model_path 63 | if scaler_path is not None and not scaler_path.endswith("None"): 64 | upscaler_paths.add(scaler_path) 65 | 66 | return upscaler_paths 67 | 68 | 69 | # see https://github.com/Mikubill/sd-webui-controlnet/blob/07bed6ccf8a468a45b2833cfdadc749927cbd575/scripts/global_state.py#L205 70 | @ipc.restrict_to_process('webui') 71 | def get_controlnet_paths(): 72 | from modules import shared 73 | controlnet_path = os.path.join(shared.extensions_dir, 'sd-webui-controlnet') 74 | try: 75 | sys.path.insert(0, controlnet_path) 76 | controlnet = importlib.import_module('extensions.sd-webui-controlnet.scripts.external_code', 'external_code') 77 | except: 78 | return [] 79 | finally: 80 | sys.path.pop(sys.path.index(controlnet_path)) 81 | 82 | ext_dirs = (shared.opts.data.get("control_net_models_path", None), getattr(shared.cmd_opts, 'controlnet_dir', None)) 83 | extra_lora_paths = (extra_lora_path for extra_lora_path in ext_dirs 84 | if extra_lora_path is not None and os.path.exists(extra_lora_path)) 85 | return [ 86 | controlnet.global_state.cn_models_dir, 87 | controlnet.global_state.cn_models_dir_old, 88 | *extra_lora_paths 89 | ] 90 | 91 | 92 | @ipc.run_in_process('webui') 93 | def webui_save_image(*args, relative_path=None, **kwargs): 94 | from modules import images, paths 95 | if relative_path is not None: 96 | kwargs['path'] = os.path.join(paths.data_path, relative_path) 97 | 98 | return images.save_image(*args, **kwargs) 99 | -------------------------------------------------------------------------------- /lib_comfyui/webui/proxies.py: -------------------------------------------------------------------------------- 1 | import functools 2 | import gc 3 | import sys 4 | import yaml 5 | import textwrap 6 | import torch 7 | from lib_comfyui import ipc, torch_utils 8 | from lib_comfyui.webui import settings 9 | 10 | 11 | class ModelPatcher: 12 | def __init__(self, model): 13 | self.model = model 14 | self.load_device = model.device 15 | self.offload_device = model.device 16 | self.model_options = {'transformer_options': {}} 17 | 18 | def model_size(self): 19 | # returning 0 means to manage the model with VRAMState.NORMAL_VRAM 20 | # https://github.com/comfyanonymous/ComfyUI/blob/ee8f8ee07fb141e5a5ce3abf602ed0fa2e50cf7b/comfy/model_management.py#L272-L276 21 | return 0 22 | 23 | def clone(self): 24 | return self 25 | 26 | def set_model_patch(self, *args, **kwargs): 27 | soft_raise('patching a webui resource is not yet supported') 28 | 29 | def set_model_patch_replace(self, *args, **kwargs): 30 | soft_raise('patching a webui resource is not yet supported') 31 | 32 | def model_patches_to(self, device): 33 | return 34 | 35 | def model_dtype(self): 36 | return self.model.dtype 37 | 38 | @property 39 | def current_device(self): 40 | return self.model.device 41 | 42 | def add_patches(self, *args, **kwargs): 43 | soft_raise('patching a webui resource is not yet supported') 44 | return [] 45 | 46 | def get_key_patches(self, *args, **kwargs): 47 | return {} 48 | 49 | def model_state_dict(self, *args, **kwargs): 50 | soft_raise('accessing the webui checkpoint state dict from comfyui is not yet suppported') 51 | return {} 52 | 53 | def patch_model(self, *args, **kwargs): 54 | return self.model 55 | 56 | def unpatch_model(self, *args, **kwargs): 57 | return 58 | 59 | def __getattr__(self, item): 60 | if item in self.__dict__: 61 | return self.__dict__[item] 62 | 63 | import comfy 64 | return functools.partial(getattr(comfy.sd.ModelPatcher, item), self) 65 | 66 | 67 | class Model: 68 | @property 69 | def model_config(self): 70 | return get_comfy_model_config() 71 | 72 | @property 73 | def model_type(self): 74 | import comfy 75 | return comfy.model_base.ModelType.EPS 76 | 77 | @property 78 | def latent_format(self): 79 | return get_comfy_model_config().latent_format 80 | 81 | def process_latent_in(self, latent, *args, **kwargs): 82 | return self.latent_format.process_in(latent) 83 | 84 | def process_latent_out(self, latent, *args, **kwargs): 85 | return self.latent_format.process_out(latent) 86 | 87 | def to(self, device, *args, **kwargs): 88 | assert str(device) == str(self.device), textwrap.dedent(f''' 89 | cannot move the webui unet to a different device 90 | comfyui attempted to move it from {self.device} to {device} 91 | ''') 92 | return self 93 | 94 | def is_adm(self, *args, **kwargs): 95 | adm_in_channels = get_comfy_model_config().unet_config.get('adm_in_channels', None) or 0 96 | return adm_in_channels > 0 97 | 98 | def encode_adm(self, *args, **kwargs): 99 | raise NotImplementedError('webui v-prediction checkpoints are not yet supported') 100 | 101 | def apply_model(self, *args, **kwargs): 102 | args = torch_utils.deep_to(args, device='cpu') 103 | del kwargs['transformer_options'] 104 | kwargs = torch_utils.deep_to(kwargs, device='cpu') 105 | return torch_utils.deep_to(Model.sd_model_apply(*args, **kwargs), device=self.device) 106 | 107 | @staticmethod 108 | @ipc.run_in_process('webui') 109 | def sd_model_apply(*args, **kwargs): 110 | from modules import shared, devices 111 | args = torch_utils.deep_to(args, shared.sd_model.device) 112 | kwargs = torch_utils.deep_to(kwargs, shared.sd_model.device) 113 | with devices.autocast(), torch.no_grad(): 114 | res = shared.sd_model.model(*args, **kwargs).cpu().share_memory_() 115 | free_webui_memory() 116 | return res 117 | 118 | def state_dict(self): 119 | soft_raise('accessing the webui checkpoint state dict from comfyui is not yet suppported') 120 | return {} 121 | 122 | def __getattr__(self, item): 123 | if item in self.__dict__: 124 | return self.__dict__[item] 125 | 126 | res = Model.sd_model_getattr(item) 127 | if item != "device": 128 | res = torch_utils.deep_to(res, device=self.device) 129 | 130 | return res 131 | 132 | @staticmethod 133 | @ipc.run_in_process('webui') 134 | def sd_model_getattr(item): 135 | from modules import shared 136 | res = getattr(shared.sd_model, item) 137 | res = torch_utils.deep_to(res, 'cpu') 138 | return res 139 | 140 | 141 | class ClipWrapper: 142 | def __init__(self, proxy): 143 | self.cond_stage_model = proxy 144 | self.patcher = ModelPatcher(self.cond_stage_model) 145 | 146 | @property 147 | def layer_idx(self): 148 | clip_skip = settings.opts.CLIP_stop_at_last_layers 149 | return -clip_skip if clip_skip > 1 else None 150 | 151 | def clone(self, *args, **kwargs): 152 | return self 153 | 154 | def load_from_state_dict(self, *args, **kwargs): 155 | return 156 | 157 | def clip_layer(self, layer_idx, *args, **kwargs): 158 | soft_raise(f'cannot control webui clip skip from comfyui. Tried to stop at layer {layer_idx}') 159 | return 160 | 161 | def tokenize(self, *args, **kwargs): 162 | args = torch_utils.deep_to(args, device='cpu') 163 | kwargs = torch_utils.deep_to(kwargs, device='cpu') 164 | return torch_utils.deep_to(ClipWrapper.sd_clip_tokenize_with_weights(*args, **kwargs), device=self.cond_stage_model.device) 165 | 166 | @staticmethod 167 | @ipc.run_in_process('webui') 168 | def sd_clip_tokenize_with_weights(text, return_word_ids=False): 169 | from modules import shared 170 | chunks, tokens_count, *_ = shared.sd_model.cond_stage_model.tokenize_line(text) 171 | weighted_tokens = [list(zip(chunk.tokens, chunk.multipliers)) for chunk in chunks] 172 | clip_max_len = shared.sd_model.cond_stage_model.wrapped.max_length 173 | if return_word_ids: 174 | padding_tokens_count = ((tokens_count // clip_max_len) + 1) * clip_max_len 175 | for token_i in range(padding_tokens_count): 176 | actual_id = token_i if token_i < tokens_count else 0 177 | weighted_tokens[token_i // clip_max_len][token_i % clip_max_len] += (actual_id,) 178 | 179 | return weighted_tokens 180 | 181 | def __getattr__(self, item): 182 | if item in self.__dict__: 183 | return self.__dict__[item] 184 | 185 | import comfy 186 | return functools.partial(getattr(comfy.sd.CLIP, item), self) 187 | 188 | 189 | class Clip: 190 | def clip_layer(self, layer_idx, *args, **kwargs): 191 | soft_raise(f'cannot control webui clip skip from comfyui. Tried to stop at layer {layer_idx}') 192 | return 193 | 194 | def reset_clip_layer(self, *args, **kwargs): 195 | return 196 | 197 | def encode_token_weights(self, *args, **kwargs): 198 | args = torch_utils.deep_to(args, device='cpu') 199 | kwargs = torch_utils.deep_to(kwargs, device='cpu') 200 | return torch_utils.deep_to(Clip.sd_clip_encode_token_weights(*args, **kwargs), device=self.device) 201 | 202 | @staticmethod 203 | @ipc.run_in_process('webui') 204 | def sd_clip_encode_token_weights(token_weight_pairs_list): 205 | from modules import shared 206 | tokens = [ 207 | [pair[0] for pair in token_weight_pairs] 208 | for token_weight_pairs in token_weight_pairs_list 209 | ] 210 | weights = [ 211 | [pair[1] for pair in token_weight_pairs] 212 | for token_weight_pairs in token_weight_pairs_list 213 | ] 214 | conds = [ 215 | shared.sd_model.cond_stage_model.process_tokens([tokens], [weights]) 216 | for tokens, weights in zip(tokens, weights) 217 | ] 218 | return torch.hstack(conds).cpu().share_memory_(), None 219 | 220 | def to(self, device): 221 | assert str(device) == str(self.device), textwrap.dedent(f''' 222 | cannot move the webui unet to a different device 223 | comfyui attempted to move it from {self.device} to {device} 224 | ''') 225 | return self 226 | 227 | def state_dict(self): 228 | soft_raise('accessing the webui checkpoint state dict from comfyui is not yet suppported') 229 | return {} 230 | 231 | def __getattr__(self, item): 232 | if item in self.__dict__: 233 | return self.__dict__[item] 234 | 235 | res = Clip.sd_clip_getattr(item) 236 | if item != "device": 237 | res = torch_utils.deep_to(res, device=self.device) 238 | 239 | return res 240 | 241 | @staticmethod 242 | @ipc.run_in_process('webui') 243 | def sd_clip_getattr(item): 244 | from modules import shared 245 | res = getattr(shared.sd_model.cond_stage_model.wrapped.transformer, item) 246 | res = torch_utils.deep_to(res, 'cpu') 247 | return res 248 | 249 | 250 | class VaeWrapper: 251 | def __init__(self, proxy): 252 | self.first_stage_model = proxy 253 | 254 | @property 255 | def vae_dtype(self): 256 | return self.first_stage_model.dtype 257 | 258 | @property 259 | def device(self): 260 | return self.first_stage_model.device 261 | 262 | @property 263 | def offload_device(self): 264 | return self.first_stage_model.device 265 | 266 | def __getattr__(self, item): 267 | if item in self.__dict__: 268 | return self.__dict__[item] 269 | 270 | import comfy 271 | return functools.partial(getattr(comfy.sd.VAE, item), self) 272 | 273 | 274 | class Vae: 275 | def state_dict(self): 276 | soft_raise('accessing the webui checkpoint state dict from comfyui is not yet suppported') 277 | return {} 278 | 279 | def encode(self, *args, **kwargs): 280 | args = torch_utils.deep_to(args, device='cpu') 281 | kwargs = torch_utils.deep_to(kwargs, device='cpu') 282 | res = torch_utils.deep_to(Vae.sd_vae_encode(*args, **kwargs), device=self.device) 283 | return DistributionProxy(res) 284 | 285 | @staticmethod 286 | @ipc.run_in_process('webui') 287 | def sd_vae_encode(*args, **kwargs): 288 | from modules import shared, devices 289 | args = torch_utils.deep_to(args, shared.sd_model.device) 290 | kwargs = torch_utils.deep_to(kwargs, shared.sd_model.device) 291 | with devices.autocast(), torch.no_grad(): 292 | res = shared.sd_model.first_stage_model.encode(*args, **kwargs).sample().cpu().share_memory_() 293 | free_webui_memory() 294 | return res 295 | 296 | def decode(self, *args, **kwargs): 297 | args = torch_utils.deep_to(args, device='cpu') 298 | kwargs = torch_utils.deep_to(kwargs, device='cpu') 299 | return torch_utils.deep_to(Vae.sd_vae_decode(*args, **kwargs), device=self.device) 300 | 301 | @staticmethod 302 | @ipc.run_in_process('webui') 303 | def sd_vae_decode(*args, **kwargs): 304 | from modules import shared, devices 305 | args = torch_utils.deep_to(args, shared.sd_model.device) 306 | kwargs = torch_utils.deep_to(kwargs, shared.sd_model.device) 307 | with devices.autocast(), torch.no_grad(): 308 | res = shared.sd_model.first_stage_model.decode(*args, **kwargs).cpu().share_memory_() 309 | free_webui_memory() 310 | return res 311 | 312 | def to(self, device): 313 | assert str(device) == str(self.device), textwrap.dedent(f''' 314 | cannot move the webui unet to a different device 315 | comfyui attempted to move it from {self.device} to {device} 316 | ''') 317 | return self 318 | 319 | def __getattr__(self, item): 320 | if item in self.__dict__: 321 | return self.__dict__[item] 322 | 323 | res = Vae.sd_vae_getattr(item) 324 | if item != "device": 325 | res = torch_utils.deep_to(res, device=self.device) 326 | 327 | return res 328 | 329 | @staticmethod 330 | @ipc.run_in_process('webui') 331 | def sd_vae_getattr(item): 332 | from modules import shared 333 | res = getattr(shared.sd_model.first_stage_model, item) 334 | res = torch_utils.deep_to(res, 'cpu') 335 | return res 336 | 337 | 338 | class DistributionProxy: 339 | def __init__(self, sample): 340 | self.sample_proxy = sample 341 | 342 | def sample(self, *args, **kwargs): 343 | return self.sample_proxy 344 | 345 | 346 | @ipc.run_in_process('webui') 347 | def free_webui_memory(): 348 | gc.collect(1) 349 | torch.cuda.empty_cache() 350 | 351 | 352 | @ipc.restrict_to_process('comfyui') 353 | def raise_on_unsupported_model_type(config): 354 | import comfy 355 | if type(config) not in ( 356 | comfy.supported_models.SD15, 357 | comfy.supported_models.SD20, 358 | ): 359 | raise NotImplementedError(f'Webui model type {type(config).__name__} is not yet supported') 360 | 361 | 362 | @ipc.restrict_to_process('comfyui') 363 | def get_comfy_model_config(): 364 | import comfy 365 | with open(sd_model_get_config()) as f: 366 | config_dict = yaml.safe_load(f) 367 | 368 | unet_config = config_dict['model']['params']['unet_config']['params'] 369 | unet_config['use_linear_in_transformer'] = unet_config.get('use_linear_in_transformer', False) 370 | unet_config['adm_in_channels'] = unet_config.get('adm_in_channels', None) 371 | unet_config['use_temporal_attention'] = unet_config.get('use_temporal_attention', False) 372 | return comfy.model_detection.model_config_from_unet_config(unet_config) 373 | 374 | 375 | @ipc.run_in_process('webui') 376 | def sd_model_get_config(): 377 | from modules import shared, sd_models, sd_models_config 378 | return sd_models_config.find_checkpoint_config(shared.sd_model.state_dict(), sd_models.select_checkpoint()) 379 | 380 | 381 | @ipc.run_in_process('webui') 382 | def extra_networks_parse_prompts(prompts): 383 | from modules import extra_networks 384 | return extra_networks.parse_prompts(prompts) 385 | 386 | 387 | def soft_raise(message): 388 | print(f'[sd-webui-comfyui] {message}', file=sys.stderr) 389 | -------------------------------------------------------------------------------- /lib_comfyui/webui/reverse_proxy.py: -------------------------------------------------------------------------------- 1 | from lib_comfyui.webui import settings 2 | from lib_comfyui import ipc, global_state 3 | 4 | 5 | @ipc.restrict_to_process("webui") 6 | def create_comfyui_proxy(fast_api): 7 | if not (global_state.enabled and global_state.reverse_proxy_enabled): 8 | return 9 | 10 | comfyui_url = settings.get_comfyui_server_url() 11 | proxy_route = settings.get_comfyui_reverse_proxy_route() 12 | 13 | create_http_reverse_proxy(fast_api, comfyui_url, proxy_route) 14 | create_ws_reverse_proxy(fast_api, comfyui_url, proxy_route) 15 | print("[sd-webui-comfyui]", f"Created a reverse proxy route to ComfyUI: {proxy_route}") 16 | 17 | 18 | def create_http_reverse_proxy(fast_api, comfyui_url, proxy_route): 19 | from starlette.requests import Request 20 | from starlette.responses import StreamingResponse, Response 21 | from starlette.background import BackgroundTask 22 | import httpx 23 | 24 | web_client = httpx.AsyncClient(base_url=comfyui_url) 25 | 26 | # src: https://github.com/tiangolo/fastapi/issues/1788#issuecomment-1071222163 27 | async def reverse_proxy(request: Request): 28 | base_path = request.url.path.replace(proxy_route, "", 1) 29 | url = httpx.URL(path=base_path, query=request.url.query.encode("utf-8")) 30 | rp_req = web_client.build_request(request.method, url, headers=request.headers.raw, content=await request.body()) 31 | try: 32 | rp_resp = await web_client.send(rp_req, stream=True) 33 | except httpx.ConnectError: 34 | return Response(status_code=404) 35 | else: 36 | return StreamingResponse( 37 | async_iter_raw_patched(rp_resp, proxy_route), 38 | status_code=rp_resp.status_code, 39 | headers=rp_resp.headers, 40 | background=BackgroundTask(rp_resp.aclose), 41 | ) 42 | 43 | fast_api.add_route(f"{proxy_route}/{{path:path}}", reverse_proxy, ["GET", "POST", "PUT", "DELETE"]) 44 | 45 | 46 | async def async_iter_raw_patched(response, proxy_route): 47 | proxy_route_bytes = proxy_route.encode("utf-8") 48 | import_paths_to_patch = [ 49 | "/scripts/", 50 | "/extensions/", 51 | "/extensions/webui_scripts/" 52 | ] 53 | patches = [ 54 | (b'/favicon', proxy_route_bytes + b'/favicon'), 55 | *( 56 | ( 57 | b'from "' + import_path.encode("utf-8"), 58 | b'from "' + proxy_route_bytes + import_path.encode("utf-8"), 59 | ) 60 | for import_path in import_paths_to_patch 61 | ), 62 | ] 63 | 64 | async for chunk in response.aiter_raw(): 65 | for substring, replacement in patches: 66 | chunk = chunk.replace(substring, replacement) 67 | yield chunk 68 | 69 | 70 | def create_ws_reverse_proxy(fast_api, comfyui_url, proxy_route): 71 | from fastapi import WebSocket 72 | import websockets 73 | import asyncio 74 | from starlette.websockets import WebSocketDisconnect 75 | from websockets.exceptions import ConnectionClosedOK 76 | 77 | ws_comfyui_url = http_to_ws(comfyui_url) 78 | 79 | @fast_api.websocket(f"{proxy_route}/ws") 80 | async def websocket_endpoint(ws_client: WebSocket): 81 | await ws_client.accept() 82 | async with websockets.connect(ws_comfyui_url) as ws_server: 83 | 84 | async def listen_to_client(): 85 | try: 86 | while True: 87 | data = await ws_client.receive_text() 88 | await ws_server.send(data) 89 | except WebSocketDisconnect: 90 | await ws_server.close() 91 | 92 | async def listen_to_server(): 93 | try: 94 | while True: 95 | data = await ws_server.recv() 96 | await ws_client.send_text(data) 97 | except ConnectionClosedOK: 98 | pass 99 | 100 | await asyncio.gather(listen_to_client(), listen_to_server()) 101 | 102 | 103 | def http_to_ws(url: str) -> str: 104 | from urllib.parse import urlparse, urlunparse 105 | 106 | parsed_url = urlparse(url) 107 | ws_scheme = 'wss' if parsed_url.scheme == 'https' else 'ws' 108 | ws_url = parsed_url._replace(scheme=ws_scheme) 109 | return f"{urlunparse(ws_url)}/ws" 110 | -------------------------------------------------------------------------------- /lib_comfyui/webui/settings.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import textwrap 3 | from pathlib import Path 4 | from lib_comfyui import ipc, global_state 5 | import install_comfyui 6 | 7 | 8 | @ipc.restrict_to_process('webui') 9 | def create_section(): 10 | from modules import shared 11 | import gradio as gr 12 | 13 | section = ('comfyui', "ComfyUI") 14 | shared.opts.add_option('comfyui_enabled', shared.OptionInfo(True, 'Enable sd-webui-comfyui extension', section=section)) 15 | 16 | shared.opts.add_option("comfyui_update_button", shared.OptionInfo( 17 | "Update comfyui (requires reload ui)", "Update comfyui", gr.Button, section=section)) 18 | 19 | shared.opts.add_option("comfyui_install_location", shared.OptionInfo( 20 | install_comfyui.default_install_location, "ComfyUI install location", section=section)) 21 | shared.opts.add_option("comfyui_additional_args", shared.OptionInfo( 22 | '', "Additional cli arguments to pass to ComfyUI (requires reload UI. Do NOT prepend --comfyui-, these are directly forwarded to comfyui)", section=section)) 23 | shared.opts.add_option("comfyui_client_address", shared.OptionInfo( 24 | '', 'Address of the ComfyUI server as seen from the webui. Only used by the extension to load the ComfyUI iframe (requires reload UI)', 25 | component_args={'placeholder': 'Leave empty to use the --listen address of the ComfyUI server'}, section=section)) 26 | 27 | shared.opts.onchange('comfyui_enabled', update_enabled) 28 | 29 | shared.opts.add_option("comfyui_ipc_strategy", shared.OptionInfo( 30 | next(iter(ipc_strategy_choices.keys())), "Interprocess communication strategy", gr.Dropdown, lambda: {"choices": list(ipc_strategy_choices.keys())}, section=section)) 31 | shared.opts.onchange('comfyui_ipc_strategy', update_ipc_strategy) 32 | update_ipc_strategy() 33 | 34 | shared.opts.add_option("comfyui_graceful_termination_timeout", shared.OptionInfo( 35 | 5, 'ComfyUI server graceful termination timeout (in seconds) when reloading the gradio UI (-1 to block until the ComfyUI server exits normally)', gr.Number, section=section)) 36 | shared.opts.onchange('comfyui_graceful_termination_timeout', update_comfyui_graceful_termination_timeout) 37 | update_comfyui_graceful_termination_timeout() 38 | 39 | shared.opts.add_option("comfyui_reverse_proxy_enabled", shared.OptionInfo( 40 | next(iter(reverse_proxy_choices.keys())), "Load ComfyUI iframes through a reverse proxy (requires reload UI. Needs --api. Default is on if webui is remote)", gr.Dropdown, lambda: {"choices": list(reverse_proxy_choices.keys())}, section=section)) 41 | shared.opts.onchange("comfyui_reverse_proxy_enabled", update_reverse_proxy_enabled) 42 | update_reverse_proxy_enabled() 43 | 44 | shared.opts.add_option("comfyui_reverse_proxy_disable_port", shared.OptionInfo(False,'Disable using the port when using reverse_proxy. Useful if using a nginx server',section=section)) 45 | 46 | 47 | @ipc.restrict_to_process('webui') 48 | def update_enabled(): 49 | from modules import shared 50 | global_state.enabled = shared.opts.data.get('comfyui_enabled', True) 51 | 52 | 53 | @ipc.restrict_to_process('webui') 54 | def update_ipc_strategy(): 55 | from modules import shared 56 | ipc_strategy_choice = shared.opts.data.get('comfyui_ipc_strategy', next(iter(ipc_strategy_choices.keys()))) 57 | global_state.ipc_strategy_class = ipc_strategy_choices[ipc_strategy_choice] 58 | global_state.ipc_strategy_class_name = global_state.ipc_strategy_class.__name__ 59 | 60 | 61 | @ipc.restrict_to_process('webui') 62 | def update_comfyui_graceful_termination_timeout(): 63 | from modules import shared 64 | timeout = shared.opts.data.get('comfyui_graceful_termination_timeout', 5) 65 | global_state.comfyui_graceful_termination_timeout = timeout if timeout >= 0 else None 66 | 67 | 68 | @ipc.restrict_to_process("webui") 69 | def update_reverse_proxy_enabled(): 70 | from modules import shared 71 | reverse_proxy_enabled = shared.opts.data.get('comfyui_reverse_proxy_enabled', next(iter(reverse_proxy_choices.keys()))) 72 | global_state.reverse_proxy_enabled = reverse_proxy_choices[reverse_proxy_enabled]() and getattr(shared.cmd_opts, "api", False) 73 | 74 | 75 | @ipc.restrict_to_process("webui") 76 | def subscribe_update_button(component, **kwargs): 77 | if getattr(component, "elem_id", None) == "setting_comfyui_update_button": 78 | component.click(fn=update_comfyui) 79 | 80 | 81 | @ipc.restrict_to_process("webui") 82 | def update_comfyui(): 83 | install_comfyui.update(get_install_location()) 84 | 85 | 86 | ipc_strategy_choices = { 87 | 'Default': ipc.strategies.OsFriendlyIpcStrategy, 88 | 'Shared memory': ipc.strategies.SharedMemoryIpcStrategy, 89 | 'File system': ipc.strategies.FileSystemIpcStrategy, 90 | } 91 | 92 | 93 | ipc_display_names = { 94 | v.__name__: k 95 | for k, v in ipc_strategy_choices.items() 96 | if k != 'Default' 97 | } 98 | 99 | 100 | @ipc.restrict_to_process('webui') 101 | def get_install_location() -> Path: 102 | from modules import shared 103 | install_location = install_comfyui.default_install_location 104 | install_location = shared.opts.data.get('comfyui_install_location', install_location).strip() 105 | return Path(install_location) 106 | 107 | 108 | @ipc.restrict_to_process('webui') 109 | def get_additional_argv(): 110 | from modules import shared 111 | return [arg.strip() for arg in shared.opts.data.get('comfyui_additional_args', '').split()] 112 | 113 | 114 | @ipc.restrict_to_process('webui') 115 | def get_setting_value(setting_key): 116 | webui_argv = get_additional_argv() 117 | index = webui_argv.index(setting_key) if setting_key in webui_argv else -1 118 | setting_value = webui_argv[index + 1] if 0 <= index < len(webui_argv) - 1 else None 119 | return setting_value 120 | 121 | 122 | @ipc.restrict_to_process('webui') 123 | def get_comfyui_iframe_url(): 124 | update_reverse_proxy_enabled() 125 | if global_state.reverse_proxy_enabled: 126 | return get_comfyui_reverse_proxy_url() 127 | else: 128 | return get_comfyui_client_url() 129 | 130 | 131 | @ipc.restrict_to_process('webui') 132 | def get_comfyui_reverse_proxy_url(): 133 | """ 134 | comfyui reverse proxy url, as seen from the browser 135 | """ 136 | return get_comfyui_reverse_proxy_route() 137 | 138 | 139 | def get_comfyui_reverse_proxy_route(): 140 | return "/sd-webui-comfyui/comfyui" 141 | 142 | 143 | @ipc.restrict_to_process('webui') 144 | def get_comfyui_client_url(): 145 | """ 146 | comfyui server direct url, as seen from the browser 147 | """ 148 | from modules import shared 149 | loopback_address = '127.0.0.1' 150 | server_url = get_setting_value('--listen') or getattr(shared.cmd_opts, 'comfyui_listen', loopback_address) 151 | client_url = shared.opts.data.get('comfyui_client_address', None) or getattr(shared.cmd_opts, 'webui_comfyui_client_address', None) or server_url 152 | client_url = canonicalize_url(client_url, get_port()) 153 | if client_url.startswith(('http://0.0.0.0', 'https://0.0.0.0')): 154 | print(textwrap.dedent(f""" 155 | [sd-webui-comfyui] changing the ComfyUI client address from {client_url} to {loopback_address} 156 | This does not change the --listen address passed to ComfyUI, but instead the address used by the extension to load the iframe 157 | To override this behavior, navigate to the extension settings or use the --webui-comfyui-client-address
cli argument 158 | """), sys.stderr) 159 | client_url = client_url.replace("0.0.0.0", "127.0.0.1", 1) 160 | 161 | return client_url 162 | 163 | 164 | @ipc.restrict_to_process('webui') 165 | def canonicalize_url(input_url: str, default_port: int = 8189) -> str: 166 | from urllib.parse import urlparse, urlunparse 167 | from modules import shared 168 | # Step 1: Prepend 'http://' if scheme is missing 169 | if not input_url.startswith(('http://', 'https://')): 170 | input_url = 'http://' + input_url 171 | 172 | # Step 2: Parse the modified URL 173 | parsed = urlparse(input_url) 174 | 175 | # Step 3: Add the missing scheme 176 | scheme = parsed.scheme if parsed.scheme else 'http' 177 | 178 | # Step 4: Add the missing port 179 | disable_port = shared.opts.data.get('comfyui_reverse_proxy_disable_port',False) 180 | netloc = parsed.netloc 181 | if ':' not in netloc and not disable_port: # Check if port is missing 182 | netloc += f":{default_port}" 183 | elif netloc.count(':') == 1 and parsed.scheme: # If port exists but scheme was present in input 184 | host, port = netloc.split(':') 185 | netloc = f"{host}:{port}" 186 | 187 | # Reconstruct the URL 188 | canonicalized_url = urlunparse((scheme, netloc, parsed.path, parsed.params, parsed.query, parsed.fragment)) 189 | return canonicalized_url 190 | 191 | 192 | @ipc.restrict_to_process('webui') 193 | def get_comfyui_server_url(): 194 | """ 195 | comfyui server url, as seen from the webui server 196 | """ 197 | return f"http://localhost:{get_port()}" 198 | 199 | 200 | @ipc.restrict_to_process('webui') 201 | def get_port(): 202 | from modules import shared 203 | return get_setting_value('--port') or getattr(shared.cmd_opts, 'comfyui_port', 8188) 204 | 205 | 206 | @ipc.restrict_to_process('webui') 207 | def is_webui_server_remote(): 208 | from modules import shared 209 | return any( 210 | bool(getattr(shared.cmd_opts, opt, False)) 211 | for opt in ( 212 | "listen", 213 | "share", 214 | "ngrok", 215 | 216 | # additional reverse proxy options from https://github.com/Bing-su/sd-webui-tunnels 217 | "cloudflared", 218 | "localhostrun", 219 | "remotemoe", 220 | "jprq", 221 | "bore", 222 | "googleusercontent", 223 | "tunnel_webhook", 224 | ) 225 | ) 226 | 227 | 228 | reverse_proxy_choices = { 229 | "Default": is_webui_server_remote, 230 | "Always": lambda: True, 231 | "Never": lambda: False, 232 | } 233 | 234 | 235 | class WebuiOptions: 236 | def __getattr__(self, item): 237 | return WebuiOptions.opts_getattr(item) 238 | 239 | @staticmethod 240 | @ipc.run_in_process('webui') 241 | def opts_getattr(item): 242 | from modules import shared 243 | return getattr(shared.opts, item) 244 | 245 | 246 | class WebuiSharedState: 247 | def __getattr__(self, item): 248 | return WebuiSharedState.shared_state_getattr(item) 249 | 250 | @staticmethod 251 | @ipc.run_in_process('webui') 252 | def shared_state_getattr(item): 253 | from modules import shared 254 | return getattr(shared.state, item) 255 | 256 | 257 | opts = WebuiOptions() 258 | shared_state = WebuiSharedState() 259 | 260 | 261 | __base_dir = None 262 | 263 | 264 | @ipc.run_in_process('webui') 265 | def get_extension_base_dir(): 266 | init_extension_base_dir() 267 | return __base_dir 268 | 269 | 270 | @ipc.restrict_to_process('webui') 271 | def init_extension_base_dir(): 272 | global __base_dir 273 | from modules import scripts 274 | if __base_dir is None: 275 | __base_dir = scripts.basedir() 276 | -------------------------------------------------------------------------------- /lib_comfyui/webui/tab.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import textwrap 4 | import gradio as gr 5 | import install_comfyui 6 | from lib_comfyui import external_code, ipc 7 | from lib_comfyui.webui import settings, gradio_utils 8 | from lib_comfyui.default_workflow_types import sandbox_tab_workflow_type 9 | 10 | 11 | webui_client_id = gr.Text( 12 | elem_id='comfyui_webui_client_id', 13 | visible=False, 14 | ) 15 | 16 | 17 | def create_tab(): 18 | install_location = settings.get_install_location() 19 | with gr.Blocks() as tab: 20 | if os.path.exists(install_location): 21 | gr.HTML(get_comfyui_app_html()) 22 | else: 23 | with gr.Row(): 24 | gr.Markdown(comfyui_install_instructions_markdown) 25 | 26 | with gr.Column(): 27 | with gr.Row(): 28 | install_manager = gr.Checkbox(label='Install with ComfyUI-Manager', value=True) 29 | 30 | with gr.Row(): 31 | install_path = gr.Textbox(placeholder=f'Leave empty to install at {install_comfyui.default_install_location}', label='Installation path') 32 | 33 | with gr.Row(): 34 | install_button = gr.Button('Install ComfyUI', variant='primary') 35 | 36 | with gr.Row(): 37 | installed_feedback = gr.Markdown() 38 | 39 | install_button.click(automatic_install_comfyui, inputs=[install_manager, install_path], outputs=[installed_feedback], show_progress=True) 40 | 41 | gradio_utils.ExtensionDynamicProperty( 42 | key='workflow_type_ids', 43 | value=external_code.get_workflow_type_ids(), 44 | ) 45 | webui_client_id.render() 46 | return [(tab, sandbox_tab_workflow_type.display_name, 'comfyui_webui_root')] 47 | 48 | 49 | @ipc.restrict_to_process('webui') 50 | def automatic_install_comfyui(should_install_manager, install_location): 51 | from modules import shared 52 | install_location = install_location.strip() 53 | if not install_location: 54 | install_location = install_comfyui.default_install_location 55 | 56 | if not can_install_at(install_location): 57 | message = 'Error! The provided path already exists. Please provide a path to an empty or non-existing directory.' 58 | print(message, file=sys.stderr) 59 | return gr.Markdown.update(message) 60 | 61 | install_comfyui.main(install_location, should_install_manager) 62 | shared.opts.comfyui_install_location = install_location 63 | 64 | return gr.Markdown.update('Installed! Now please reload the UI.') 65 | 66 | 67 | def can_install_at(path): 68 | is_empty_dir = os.path.isdir(path) and not os.listdir(path) 69 | return not os.path.exists(path) or is_empty_dir 70 | 71 | 72 | comfyui_install_instructions_markdown = ''' 73 | ## ComfyUI extension 74 | It looks like your ComfyUI installation isn't set up yet. 75 | If you already have ComfyUI installed on your computer, go to `Settings > ComfyUI`, and set the proper install location. 76 | 77 | Alternatively, if you don't have ComfyUI installed, you can install it here: 78 | ''' 79 | 80 | 81 | def get_comfyui_app_html(): 82 | return textwrap.dedent(f''' 83 |
84 | 88 |
89 | ''') 90 | -------------------------------------------------------------------------------- /preload.py: -------------------------------------------------------------------------------- 1 | def preload(parser): 2 | parser.add_argument("--comfyui-listen", type=str, default='127.0.0.1', nargs='?', const='0.0.0.0') 3 | parser.add_argument("--comfyui-port", type=int, default=8189) 4 | parser.add_argument("--comfyui-dont-upcast-attention", action='store_true', default=None) 5 | parser.add_argument("--comfyui-use-split-cross-attention", action='store_true', default=None) 6 | parser.add_argument("--comfyui-use-pytorch-cross-attention", action='store_true', default=None) 7 | parser.add_argument("--comfyui-disable-xformers", action='store_true', default=None) 8 | parser.add_argument("--comfyui-highvram", action='store_true', default=None) 9 | parser.add_argument("--comfyui-normalvram", action='store_true', default=None) 10 | parser.add_argument("--comfyui-lowvram", action='store_true', default=None) 11 | parser.add_argument("--comfyui-novram", action='store_true', default=None) 12 | parser.add_argument("--comfyui-cpu", action='store_true', default=None) 13 | parser.add_argument("--webui-comfyui-client-address", type=str, default=None) 14 | -------------------------------------------------------------------------------- /requirements-test.txt: -------------------------------------------------------------------------------- 1 | pytest~=7.3 2 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | portalocker>=2.7.0 2 | psutil 3 | -------------------------------------------------------------------------------- /resources/front-page.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ModelSurge/sd-webui-comfyui/10e95a42e27f7e7c2b15284491d1e09f2ea9a971/resources/front-page.gif -------------------------------------------------------------------------------- /scripts/comfyui.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from modules import scripts 4 | from lib_comfyui import global_state, platform_utils, external_code, default_workflow_types, comfyui_process 5 | from lib_comfyui.webui import callbacks, settings, patches, gradio_utils, accordion, tab 6 | from lib_comfyui.comfyui import iframe_requests, type_conversion 7 | 8 | 9 | class ComfyUIScript(scripts.Script): 10 | def __init__(self): 11 | # is_img2img is not available here. `accordion` is initialized below, in the is_img2img setter 12 | self.accordion = None 13 | self._is_img2img = None 14 | 15 | def title(self): 16 | return "ComfyUI" 17 | 18 | def show(self, is_img2img): 19 | return scripts.AlwaysVisible 20 | 21 | @property 22 | def is_img2img(self): 23 | return self._is_img2img 24 | 25 | @is_img2img.setter 26 | def is_img2img(self, is_img2img): 27 | self._is_img2img = is_img2img 28 | if self.accordion is None: 29 | # now, we can instantiate the accordion 30 | self.accordion = accordion.AccordionInterface(self.elem_id, self.get_tab()) 31 | 32 | def get_tab(self, is_img2img: bool = None): 33 | if is_img2img is None: 34 | is_img2img = self.is_img2img 35 | return "img2img" if is_img2img else "txt2img" 36 | 37 | def ui(self, is_img2img): 38 | global_state.is_ui_instantiated = True 39 | self.accordion.arrange_components() 40 | self.accordion.connect_events() 41 | self.accordion.setup_infotext_fields(self) 42 | return (tab.webui_client_id,) + self.accordion.get_script_ui_components() 43 | 44 | def process(self, p, webui_client_id, queue_front, enabled_workflow_type_ids, *args, **kwargs): 45 | if not getattr(global_state, 'enabled', True): 46 | return 47 | 48 | if not hasattr(global_state, 'enabled_workflow_type_ids'): 49 | global_state.enabled_workflow_type_ids = {} 50 | 51 | global_state.focused_webui_client_id = webui_client_id 52 | global_state.enabled_workflow_type_ids.update(enabled_workflow_type_ids) 53 | global_state.queue_front = queue_front 54 | patches.patch_processing(p) 55 | 56 | def postprocess_batch_list(self, p, pp, *args, **kwargs): 57 | iframe_requests.extend_infotext_with_comfyui_workflows(p, self.get_tab()) 58 | 59 | if not external_code.is_workflow_type_enabled(default_workflow_types.postprocess_workflow_type.get_ids(self.get_tab())[0]): 60 | return 61 | 62 | all_results = [] 63 | p_rescale_factor = 0 64 | for batch_input in extract_contiguous_buckets(pp.images, p.batch_size): 65 | batch_results = external_code.run_workflow( 66 | workflow_type=default_workflow_types.postprocess_workflow_type, 67 | tab=self.get_tab(), 68 | batch_input=type_conversion.webui_image_to_comfyui(torch.stack(batch_input).to('cpu')), 69 | identity_on_error=True, 70 | ) 71 | 72 | p_rescale_factor += len(batch_results) 73 | all_results.extend( 74 | image 75 | for batch in batch_results 76 | for image in type_conversion.comfyui_image_to_webui(batch, return_tensors=True)) 77 | 78 | p_rescale_factor = max(1, p_rescale_factor) 79 | for list_to_scale in [p.prompts, p.negative_prompts, p.seeds, p.subseeds]: 80 | list_to_scale[:] = list_to_scale * p_rescale_factor 81 | 82 | pp.images.clear() 83 | pp.images.extend(all_results) 84 | 85 | def postprocess_image(self, p, pp, *args): 86 | if not external_code.is_workflow_type_enabled( 87 | default_workflow_types.postprocess_image_workflow_type.get_ids(self.get_tab())[0]): 88 | return 89 | 90 | results = external_code.run_workflow( 91 | workflow_type=default_workflow_types.postprocess_image_workflow_type, 92 | tab=self.get_tab(), 93 | batch_input=type_conversion.webui_image_to_comfyui([pp.image]), 94 | identity_on_error=True, 95 | ) 96 | 97 | pp.image = type_conversion.comfyui_image_to_webui(results[0], return_tensors=False)[0] 98 | 99 | 100 | def extract_contiguous_buckets(images, batch_size): 101 | current_shape = None 102 | begin_index = 0 103 | 104 | for i, image in enumerate(images): 105 | if current_shape is None: 106 | current_shape = image.size() 107 | 108 | image_has_different_shape = image.size() != current_shape 109 | batch_is_full = i - begin_index >= batch_size 110 | is_last_image = i == len(images) - 1 111 | 112 | if image_has_different_shape or batch_is_full or is_last_image: 113 | end_index = i + 1 if is_last_image else i 114 | yield images[begin_index:end_index] 115 | begin_index = i 116 | current_shape = None 117 | 118 | 119 | callbacks.register_callbacks() 120 | default_workflow_types.add_default_workflow_types() 121 | settings.init_extension_base_dir() 122 | patches.apply_patches() 123 | -------------------------------------------------------------------------------- /style.css: -------------------------------------------------------------------------------- 1 | div#tab_comfyui_webui_root { 2 | border: none; 3 | } 4 | 5 | div#comfyui_webui_container { 6 | position: fixed; 7 | left: 0px; 8 | width: 100%; 9 | bottom: 0px; 10 | } 11 | 12 | div#comfyui_webui_container iframe { 13 | width: 100%; 14 | height: 100%; 15 | } 16 | 17 | div.comfyui_iframes iframe { 18 | display: none; 19 | width: 100%; 20 | height: 100%; 21 | } 22 | 23 | div.comfyui_iframes iframe.comfyui-workflow-type-visible { 24 | display: block; 25 | } 26 | 27 | .comfyui-remove-display { 28 | display: none; 29 | } 30 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ModelSurge/sd-webui-comfyui/10e95a42e27f7e7c2b15284491d1e09f2ea9a971/tests/__init__.py -------------------------------------------------------------------------------- /tests/lib_comfyui_tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ModelSurge/sd-webui-comfyui/10e95a42e27f7e7c2b15284491d1e09f2ea9a971/tests/lib_comfyui_tests/__init__.py -------------------------------------------------------------------------------- /tests/lib_comfyui_tests/argv_conversion_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from tests.utils import setup_test_env 3 | setup_test_env() 4 | from lib_comfyui import argv_conversion 5 | 6 | 7 | class DeduplicateArgvTest(unittest.TestCase): 8 | def setUp(self): 9 | self.argv = [] 10 | self.expected_argv = [] 11 | 12 | def tearDown(self): 13 | self.argv = [] 14 | self.expected_argv = [] 15 | 16 | def assert_deduplicated_equal_expected(self): 17 | argv_conversion.deduplicate_comfyui_args(self.argv) 18 | self.assertEqual(self.argv, self.expected_argv) 19 | 20 | def test_port_deduplicated(self): 21 | self.argv.extend(['--port', '1234', '--port', '5678']) 22 | self.expected_argv.extend(['--port', '1234']) 23 | self.assert_deduplicated_equal_expected() 24 | 25 | def test_port_mixed_deduplicated(self): 26 | self.argv.extend(['--port', '1234', '--port', '5678', '--comfyui-use-split-cross-attention', '--lowvram', '--port', '8765']) 27 | self.expected_argv.extend(['--port', '1234', '--comfyui-use-split-cross-attention', '--lowvram']) 28 | self.assert_deduplicated_equal_expected() 29 | 30 | def test_lowvram_deduplicated(self): 31 | self.argv.extend(['--lowvram', '--lowvram']) 32 | self.expected_argv.extend(['--lowvram']) 33 | self.assert_deduplicated_equal_expected() 34 | 35 | def test_lowvram_mixed_deduplicated(self): 36 | self.argv.extend(['--lowvram', '--port', '1234', '--lowvram', '--comfyui-use-split-cross-attention', '--lowvram']) 37 | self.expected_argv.extend(['--lowvram', '--port', '1234', '--comfyui-use-split-cross-attention']) 38 | self.assert_deduplicated_equal_expected() 39 | -------------------------------------------------------------------------------- /tests/lib_comfyui_tests/external_code_tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ModelSurge/sd-webui-comfyui/10e95a42e27f7e7c2b15284491d1e09f2ea9a971/tests/lib_comfyui_tests/external_code_tests/__init__.py -------------------------------------------------------------------------------- /tests/lib_comfyui_tests/external_code_tests/import_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from tests.utils import setup_test_env 3 | setup_test_env() 4 | 5 | import importlib 6 | import sys 7 | from lib_comfyui import external_code 8 | 9 | 10 | class ImportlibTest(unittest.TestCase): 11 | def test_importlib_uses_modules_cache(self): 12 | WorkflowType = external_code.WorkflowType 13 | del sys.modules['lib_comfyui.external_code'] 14 | comfyui_api = importlib.import_module('lib_comfyui.external_code', 'external_code') 15 | 16 | self.assertIs(WorkflowType, comfyui_api.WorkflowType) 17 | 18 | 19 | if __name__ == '__main__': 20 | unittest.main() 21 | -------------------------------------------------------------------------------- /tests/lib_comfyui_tests/external_code_tests/run_workflow_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from unittest.mock import patch 3 | from tests.utils import setup_test_env 4 | setup_test_env() 5 | 6 | from lib_comfyui import external_code, global_state 7 | 8 | 9 | class TestRunWorkflowBasicFunctionality(unittest.TestCase): 10 | def setUp(self) -> None: 11 | setattr(global_state, 'enabled_workflow_type_ids', { 12 | 'test_tab': True, 13 | }) 14 | 15 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 16 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 17 | def test_valid_workflow_with_dict_types(self, mock_start_workflow, _): 18 | mock_start_workflow.return_value = [{"key1": 0, "key2": 1}] 19 | workflow_type = external_code.WorkflowType( 20 | base_id="test", 21 | display_name="Test Tab", 22 | tabs="tab", 23 | types={"key1": 'IMAGE', "key2": 'LATENT'}, 24 | ) 25 | batch_input = {"key1": "value", "key2": 123} 26 | 27 | result = external_code.run_workflow(workflow_type, "tab", batch_input) 28 | 29 | mock_start_workflow.assert_called_once_with( 30 | batch_input_args=("value", 123), 31 | workflow_type_id="test_tab", 32 | workflow_input_types=('IMAGE', 'LATENT'), 33 | queue_front=True, 34 | ) 35 | self.assertEqual(result, mock_start_workflow.return_value) 36 | 37 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 38 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 39 | def test_valid_workflow_with_tuple_types(self, mock_start_workflow, _): 40 | mock_start_workflow.return_value = [{"key1": 0, "key2": 1}] 41 | workflow_type = external_code.WorkflowType( 42 | base_id="test", 43 | display_name="Test Tab", 44 | tabs="tab", 45 | types=("IMAGE", "LATENT"), 46 | ) 47 | batch_input = ("value", 123) 48 | 49 | result = external_code.run_workflow(workflow_type, "tab", batch_input) 50 | 51 | mock_start_workflow.assert_called_once_with( 52 | batch_input_args=("value", 123), 53 | workflow_type_id="test_tab", 54 | workflow_input_types=('IMAGE', 'LATENT'), 55 | queue_front=True, 56 | ) 57 | self.assertEqual(result, [tuple(batch.values()) for batch in mock_start_workflow.return_value]) 58 | 59 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 60 | @patch('lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync') 61 | def test_valid_workflow_with_str_types(self, mock_start_workflow, _): 62 | mock_start_workflow.return_value = [{"key1": 0}] 63 | workflow_type = external_code.WorkflowType( 64 | base_id="test", 65 | display_name="Test Tab", 66 | tabs="tab", 67 | types="IMAGE", 68 | ) 69 | batch_input = "value" 70 | 71 | result = external_code.run_workflow(workflow_type, "tab", batch_input) 72 | 73 | mock_start_workflow.assert_called_once_with( 74 | batch_input_args=("value",), 75 | workflow_type_id="test_tab", 76 | workflow_input_types=('IMAGE',), 77 | queue_front=True, 78 | ) 79 | self.assertEqual(result, [next(iter(batch.values())) for batch in mock_start_workflow.return_value]) 80 | 81 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 82 | @patch('lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync') 83 | def test_workflow_type_not_on_tab(self, mock_start_workflow, _): 84 | workflow_type = external_code.WorkflowType( 85 | base_id="test", 86 | display_name="Test Tab", 87 | tabs="tab", 88 | types="IMAGE", 89 | ) 90 | batch_input = "value" 91 | 92 | with self.assertRaises(ValueError): 93 | external_code.run_workflow(workflow_type, "not_tab", batch_input) 94 | 95 | mock_start_workflow.assert_not_called() 96 | 97 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 98 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 99 | def test_workflow_type_not_enabled(self, mock_start_workflow, _): 100 | setattr(global_state, 'enabled_workflow_type_ids', {}) 101 | 102 | workflow_type = external_code.WorkflowType( 103 | base_id="test", 104 | display_name="Test Tab", 105 | tabs="tab", 106 | types="IMAGE", 107 | ) 108 | batch_input = "value" 109 | 110 | with self.assertRaises(RuntimeError): 111 | external_code.run_workflow(workflow_type, "tab", batch_input) 112 | 113 | mock_start_workflow.assert_not_called() 114 | 115 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 116 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 117 | def test_multiple_outputs(self, mock_start_workflow, _): 118 | mock_start_workflow.return_value = [{"key1": 0, "key2": 1}, {"key1": 2, "key2": 3}] 119 | 120 | workflow_type = external_code.WorkflowType( 121 | base_id="test", 122 | display_name="Test Tab", 123 | tabs="tab", 124 | types={"key1": 'IMAGE', "key2": 'LATENT'}, 125 | ) 126 | batch_input = {"key1": "value", "key2": 123} 127 | 128 | result = external_code.run_workflow(workflow_type, "tab", batch_input) 129 | 130 | self.assertEqual(result, mock_start_workflow.return_value) 131 | 132 | 133 | class TestRunWorkflowInputValidation(unittest.TestCase): 134 | def setUp(self) -> None: 135 | setattr(global_state, 'enabled_workflow_type_ids', { 136 | 'test_tab': True, 137 | }) 138 | 139 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 140 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 141 | def test_batch_input_mismatch_dict(self, mock_start_workflow, _): 142 | workflow_type = external_code.WorkflowType( 143 | base_id="test", 144 | display_name="Test Tab", 145 | tabs="tab", 146 | types={"key1": 'IMAGE', "key2": 'LATENT'}, 147 | ) 148 | batch_input = {"key1": "value"} 149 | 150 | with self.assertRaises(TypeError): 151 | external_code.run_workflow(workflow_type, "tab", batch_input) 152 | 153 | mock_start_workflow.assert_not_called() 154 | 155 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 156 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 157 | def test_batch_input_mismatch_tuple(self, mock_start_workflow, _): 158 | workflow_type = external_code.WorkflowType( 159 | base_id="test", 160 | display_name="Test Tab", 161 | tabs="tab", 162 | types=(str, int) 163 | ) 164 | batch_input = ("value",) 165 | 166 | with self.assertRaises(TypeError): 167 | external_code.run_workflow(workflow_type, "tab", batch_input) 168 | 169 | mock_start_workflow.assert_not_called() 170 | 171 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 172 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 173 | def test_identity_on_error(self, mock_start_workflow, _): 174 | mock_start_workflow.side_effect = RuntimeError("Failed to execute workflow") 175 | workflow_type = external_code.WorkflowType( 176 | base_id="test", 177 | display_name="Test Tab", 178 | tabs="tab", 179 | types="IMAGE" 180 | ) 181 | batch_input = "value" 182 | 183 | # Testing identity_on_error=True 184 | result = external_code.run_workflow(workflow_type, "tab", batch_input, identity_on_error=True) 185 | self.assertEqual(result, [batch_input]) 186 | 187 | # Testing identity_on_error=False (the default value) 188 | with self.assertRaises(RuntimeError): 189 | external_code.run_workflow(workflow_type, "tab", batch_input) 190 | 191 | mock_start_workflow.assert_called_with( 192 | batch_input_args=("value",), 193 | workflow_type_id="test_tab", 194 | workflow_input_types=('IMAGE',), 195 | queue_front=True, 196 | ) 197 | 198 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 199 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 200 | def test_multiple_candidate_ids(self, mock_start_workflow, _): 201 | workflow_type = external_code.WorkflowType( 202 | base_id="test", 203 | display_name="Test Tab", 204 | tabs="tab", 205 | types="IMAGE" 206 | ) 207 | batch_input = "value" 208 | 209 | with patch.object(workflow_type, "get_ids", return_value=["id1", "id2"]): 210 | with self.assertRaises(AssertionError): 211 | external_code.run_workflow(workflow_type, "tab", batch_input) 212 | 213 | mock_start_workflow.assert_not_called() 214 | 215 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 216 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 217 | def test_invalid_batch_input_type(self, mock_start_workflow, _): 218 | workflow_type = external_code.WorkflowType( 219 | base_id="test", 220 | display_name="Test Tab", 221 | tabs="tab", 222 | types={"key1": 'IMAGE', "key2": 'LATENT'}, 223 | ) 224 | batch_input = ["Looking for a missing semi-colon? Here, you take this one with you --> ;", 123] 225 | 226 | with self.assertRaises(TypeError): 227 | external_code.run_workflow(workflow_type, "tab", batch_input) 228 | 229 | mock_start_workflow.assert_not_called() 230 | 231 | 232 | class TestRunWorkflowExecutionBehavior(unittest.TestCase): 233 | def setUp(self) -> None: 234 | setattr(global_state, 'enabled_workflow_type_ids', { 235 | 'test_tab': True, 236 | }) 237 | 238 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.validate_amount_of_nodes_or_throw") 239 | @patch("lib_comfyui.comfyui.iframe_requests.ComfyuiIFrameRequests.start_workflow_sync") 240 | def test_large_batch_input(self, mock_start_workflow, _): 241 | workflow_type = external_code.WorkflowType( 242 | base_id="test", 243 | display_name="Test Tab", 244 | tabs="tab", 245 | types=("IMAGE",) * 1000 246 | ) 247 | batch_input = ("value",) * 1000 248 | external_code.run_workflow(workflow_type, "tab", batch_input) 249 | mock_start_workflow.assert_called_once() 250 | 251 | 252 | if __name__ == '__main__': 253 | unittest.main() 254 | -------------------------------------------------------------------------------- /tests/lib_comfyui_tests/ipc_tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ModelSurge/sd-webui-comfyui/10e95a42e27f7e7c2b15284491d1e09f2ea9a971/tests/lib_comfyui_tests/ipc_tests/__init__.py -------------------------------------------------------------------------------- /tests/lib_comfyui_tests/ipc_tests/payload_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from tests.utils import setup_test_env, run_subprocess 3 | setup_test_env() 4 | 5 | import contextlib 6 | import pickle 7 | import time 8 | from typing import Any, Union 9 | from lib_comfyui.ipc.payload import IpcReceiver, IpcSender 10 | from lib_comfyui.ipc.strategies import FileSystemIpcStrategy 11 | 12 | 13 | class MockStrategy: 14 | storage = {} 15 | 16 | def __init__(self, name: str): 17 | self.name = name 18 | MockStrategy.storage[name] = None 19 | 20 | def is_empty(self, _lock_file: Any) -> bool: 21 | return self.storage[self.name] is None 22 | 23 | def set_data(self, _lock_file: Any, data: Union[bytes, bytearray, memoryview]): 24 | self.storage[self.name] = data 25 | 26 | @contextlib.contextmanager 27 | def get_data(self, _lock_file: Any) -> Union[bytes, bytearray, memoryview]: 28 | yield MockStrategy.storage[self.name] 29 | MockStrategy.storage[self.name] = None 30 | 31 | def clear(self, _lock_file: Any): 32 | MockStrategy.storage[self.name] = None 33 | 34 | 35 | class TestIpcSender(unittest.TestCase): 36 | def test_send_method(self): 37 | strategy: MockStrategy 38 | def strategy_factory(*args): 39 | nonlocal strategy 40 | strategy = MockStrategy(*args) 41 | return strategy 42 | 43 | sender = IpcSender("test_sender", strategy_factory) 44 | mock_value = "test_value" 45 | sender.send(mock_value) 46 | 47 | self.assertEqual(MockStrategy.storage[strategy.name], pickle.dumps(mock_value)) 48 | 49 | 50 | class TestIpcReceiver(unittest.TestCase): 51 | def test_receive_method(self): 52 | strategy: MockStrategy 53 | def strategy_factory(*args): 54 | nonlocal strategy 55 | strategy = MockStrategy(*args) 56 | return strategy 57 | 58 | receiver = IpcReceiver("test_receiver", strategy_factory, clear_on_init=True) 59 | 60 | mock_value = "test_value" 61 | strategy.set_data(None, pickle.dumps(mock_value)) # mock some data in the "shared memory" 62 | 63 | received_value = receiver.recv() 64 | self.assertEqual(received_value, mock_value) 65 | 66 | 67 | class TestFunctionalIpc(unittest.TestCase): 68 | def setUp(self) -> None: 69 | self.name = "test" 70 | 71 | def test_basic_send_and_receive(self): 72 | run_subprocess(__file__, sender_worker, self.name, FileSystemIpcStrategy, "test_value") 73 | received_value = run_subprocess(__file__, receiver_worker, self.name, FileSystemIpcStrategy) 74 | 75 | self.assertEqual(received_value, "test_value") 76 | 77 | def test_send_and_receive_complex_objects(self): 78 | complex_data = {"key": ["value1", "value2"], "num": 12345} 79 | 80 | run_subprocess(__file__, sender_worker, self.name, FileSystemIpcStrategy, complex_data) 81 | received_data = run_subprocess(__file__, receiver_worker, self.name, FileSystemIpcStrategy) 82 | 83 | self.assertEqual(received_data, complex_data) 84 | 85 | def test_concurrent_sends(self): 86 | run_subprocess(__file__, sender_worker, self.name, FileSystemIpcStrategy, "data1") 87 | run_subprocess(__file__, sender_worker, self.name, FileSystemIpcStrategy, "data2") 88 | 89 | # Only the last data sent should be received due to overwriting 90 | received_data = run_subprocess(__file__, receiver_worker, self.name, FileSystemIpcStrategy) 91 | self.assertEqual(received_data, "data2") 92 | 93 | def test_mismatched_names(self): 94 | run_subprocess(__file__, sender_worker, self.name, FileSystemIpcStrategy, "data") 95 | 96 | with self.assertRaises(Exception): # Expect an exception because names don't match 97 | run_subprocess(__file__, receiver_worker, "mismatched_name", FileSystemIpcStrategy) 98 | 99 | def test_timeout_behavior(self): 100 | start_time = time.time() 101 | with self.assertRaises(Exception): 102 | run_subprocess(__file__, receiver_worker, self.name, FileSystemIpcStrategy, 2) 103 | end_time = time.time() 104 | 105 | self.assertTrue(1.5 <= end_time - start_time <= 2.5) # Ensure timeout was approximately 2 seconds 106 | 107 | def test_queueing_behavior(self): 108 | run_subprocess(__file__, sender_worker, self.name, FileSystemIpcStrategy, "data1") 109 | run_subprocess(__file__, sender_worker, self.name, FileSystemIpcStrategy, "data2") 110 | 111 | # Since the send overwrites previous data, only "data2" should be received 112 | received_data = run_subprocess(__file__, receiver_worker, self.name, FileSystemIpcStrategy) 113 | self.assertEqual(received_data, "data2") 114 | 115 | def test_error_scenarios(self): 116 | # Test receiving with no data sent 117 | with self.assertRaises(Exception): 118 | run_subprocess(__file__, receiver_worker, self.name, FileSystemIpcStrategy) 119 | 120 | 121 | def sender_worker(name, strategy_cls, data): 122 | sender = IpcSender(name, strategy_cls, clear_on_init=True, clear_on_del=False) 123 | sender.send(data) 124 | 125 | 126 | def receiver_worker(name, strategy_cls, timeout=1): 127 | receiver = IpcReceiver(name, strategy_cls) 128 | return receiver.recv(timeout=timeout) 129 | 130 | 131 | if __name__ == "__main__": 132 | unittest.main() 133 | -------------------------------------------------------------------------------- /tests/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import subprocess 4 | import sys 5 | 6 | 7 | def setup_test_env(): 8 | extension_root = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) 9 | if extension_root not in sys.path: 10 | sys.path.append(extension_root) 11 | 12 | def worker_args(*args, **kwargs): 13 | return args, kwargs 14 | 15 | 16 | def subprocess_worker(args): 17 | args, kwargs = args 18 | return run_subprocess(*args, **kwargs) 19 | 20 | 21 | def run_subprocess(file, func, *args, **kwargs): 22 | # Read the source of the current file 23 | with open(file, 'r') as f: 24 | current_file_source = f.read() 25 | 26 | with open(__file__, 'r') as f: 27 | utils_source_file = f.read() 28 | 29 | args = f"pickle.loads({pickle.dumps(args)})" 30 | kwargs = f"pickle.loads({pickle.dumps(kwargs)})" 31 | dump_func_call_source = f"pickle.dumps({func.__name__}(*{args}, **{kwargs}))" 32 | 33 | # Craft the command to include the entire current file source, function definition, and function execution 34 | command = [ 35 | sys.executable, 36 | '-c', 37 | '\n'.join([ 38 | '__name__ += ".subprocess"', 39 | current_file_source, 40 | utils_source_file, 41 | f'sys.stdout.buffer.write({dump_func_call_source})', 42 | ]), 43 | ] 44 | 45 | # Start the subprocess with stdout piped 46 | proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) 47 | 48 | # Get the output and wait for the subprocess to finish 49 | out, err = proc.communicate() 50 | if proc.returncode != 0: 51 | raise RuntimeError(f"Subprocess failed with error code {proc.returncode} and message: {err.decode('utf-8')}") 52 | 53 | # Convert the pickled stdout data back to a Python object 54 | return pickle.loads(out) 55 | --------------------------------------------------------------------------------