Как мы создаём Visionatrix: упрощая ComfyUI

Добрый день всем.

Хотелось бы поделиться нашим опытом разработки Visionatrix — надстройки над ComfyUI, которую мы создаём вдвоём, и что из этого получается.

Мы часто читаем Хабр на протяжении более десяти лет, но руки до написания статьи никак не доходили. Наконец, решили рассказать о нашем проекте, который, возможно, будет интересен сообществу разработчиков.

Сразу оговоримся, что статья не будет объяснять, что такое ComfyUI и его особенности работы. Предполагается, что вы уже знакомы с этим инструментом.

Если нет, вы можете перейти к разделам 2-3 статьи, где мы рассказываем о процессе разработки в сжатые сроки, имея очень ограниченное время (мы работаем полный рабочий день, и Visionatrix — наш хобби-проект).

Что включает в себя эта статья:

1. Описание проблем ComfyUI и почему мы решили создать Visionatrix

2. Как мы подходили к разработке и что делали в первую очередь

3. Что получилось в итоге

Надеюсь, вам будет интересно.

Вот ссылка на проект на GitHub-e

Проблемы ComfyUI

Интерфейс ComfyUI предназначен для создания workflow-ов, а не для использования

Среднестатистическому человеку, не являющемуся программистом, довольно сложно разобраться в том, что такое KSampler, как на результат влияет параметры CFG и значение latent noise, и в остальных нюансах диффузии.

Нам хотелось создать что-то, что может использовать обычный человек, знакомый с компьютером, но не знакомый с тонкостями диффузии.

Кто-то может сказать, что есть A1111, но мы хотели ещё более простую альтернативу, и по некоторым причинам решили основываться именно на ComfyUI, так как у него большое сообщество разработчиков workflow-ов.

Также мы не хотели, чтобы при открытии программы пользователь видел множество блоков с вводом, требующих детального изучения. Понимая, что для обычного пользователя сложно разобраться в сложных настройках и блоках, мы решили дать возможность гибко управлять тем, какие параметры отображать в интерфейсе, а какие скрывать.

Поскольку бэкенд ComfyUI умеет исполнять только workflow-ы в формате API, в котором не сохраняются многие метаданные от исходного workflow-а, нам пришлось написать свои небольшие опциональные ноды. Наш бэкенд на Python парсит их и на их основе генерирует динамический UI.

Мы также оставили возможность не использовать специальные ноды, а указывать всю информацию в "названиях" нод, чтобы бэкенд мог понять, что ему отдавать в UI. Формат примерно такой:

input;display name;flag1;flag2;...

Например:

input;Photo of a person;optional;advanced;order=2

И на основе этого бэкенд формирует список входных параметров workflow-а, который UI должен отобразить. Конечно, можно подключать любой другой UI; мы понимаем, что наш может показаться не идеальным для кого-то, и стараемся всё делать модульно.

Наш интерфейс, основываясь на флагах optional (означает, что параметр для workflow можно выставлять или не выставлять) и advanced, показывает, будет ли параметр скрыт по умолчанию или нет, чтобы workflows не перегружали интерфейс.

Обычный пользователь хочет просто сгенерировать картинку на основе промпта или ввода и не хочет регулировать количество шагов в диффузии, верно?

Парсер был написан довольно быстро и достаточно прост; мы старались не усложнять проект и не добавлять лишнего, получилось следующие опции у входных параметров:

• input - ключевое слово, по которому бекенд понимает что это input параметр

• display name - имя параметра в UI

• optional - является ли параметр опциональным, для валидации на стороне UI что все требуемые параметры заполнены

• advanced - флаг определяющий скрывать ли по умолчанию параметр

• order=1 - порядок отображения в UI

• custom_id=custom_name - the custom id of the input param, по которому оно будет передавать на бекенд

Результаты работы workflow-ов в ComfyUI

Здесь мы столкнулись с другой проблемой: большинство нод имеют выходы, но их не нужно выводить, так как они чаще используются для отладки процесса workflow-а.

Пример: workflow внутри обращается к VLM для получения описания input картинки, далее с этим текстом вы в workflowе что то делаете и используеться ноду "ShowText" - что бы отобразить результат.

Полностью эту проблему мы не решили, но она есть в списке TODO на ближайшее будущее.

Чтобы решить проблему с выводом ненужных выходных данных и сделать интерфейс более удобным, мы приняли решение, что результаты класса SaveImage, а также VHS_VideoCombine для поддержки видео, всегда будут отображаться в UI. Это позволяет пользователю фокусироваться только на ключевых результатах и снижает информационную перегрузку.

Когда будет время, мы планируем добавить флаг debug для workflow-ов. Если он передан в момент создания таска, мы будем собирать все выходы и выводить их в UI в виде отдельно выпадающего списка — это будет полезно для отладки.

На данный момент мы пока просто добавили возможность скачивать сгенерированный ComfyUI workflow из UI, его всегда можно загрузить в ComfyUI и увидеть что работает не так.

История генераций и поддержка нескольких пользователей

ComfyUI не поддерживает нормально историю генераций и работу с несколькими пользователями, а для нас это было очень важно.

Удобно же, когда в семье есть один мощный компьютер, вы на нём запустили софт, и каждый член семьи использует свой аккаунт, не мешая друг другу.

Установка workflow-ов

В ComfyUI с этим большие проблемы. Чаще всего вы видите отличный workflow на openart.ai, но даже потратив 30 минут на установку всего необходимого, запустить его не получается, т.к. вам нужно установить 20 нод и откуда то взять 5 новых моделей.

Кроме того, сегодня установка workflow-а может работать, а завтра перестать, потому что, например, используемая в нём модель была удалена.

Мы частично решили эту проблему, создав временный каталог моделей в формате JSON и написали GitHub CI для проверки доступности моделей (такое часто бывает).

Похоже, что скоро в ComfyUI frontend появится поддержка указания хэшей в нодах для загрузки моделей, и тогда мы избавимся от нашего каталога, так как сможем находить модель на CivitAI по хэшу и написать парсер для HuggingFace по имени и сравнению хэшей.

Также скоро выйдет новая версия ComfyUI с поддержкой ComfyUI registry для нод. Правда, установка нод от этого не станет намного проще и не решит проблему Python-пакетов с несовместимыми зависимостями, но в целом ситуация немного улучшится.

В Visionatrix мы создали "небольшой магазин workflow-ов", используя GitHub Pages и версионирование по веткам. Также у workflow-ов есть версия и расширенные метаданные.

Получилось просто и эффективно: не нужно ничего арендовать, всё хранится в репозитории и собирается на CI гитхаба.

Удобство установки и масштабирование

Нашим главным требованием было, чтобы всё было легко устанавливать и использовать.

Также должно было быть легко масштабировать решение.

Поэтому было решено сделать всё комплексно: установка в один клик для домашнего использования использует SQLite, но при необходимости поддерживается PostgreSQL для масштабирования.

Одна инсталяция может работать в стандартном режиме, включая в себя и серверную часть, и часть, исполняющую workflow-ы (в ComfyUI эти части неразделимы). Это же можно легко изменить, подключив один инстанс к другому, где подключаемый становится внешним рабочим исполнителем тасков работающим в режиме worker-а.

О том, как мы всё это реализовали, возможно, расскажем позже, если это будет кому-то интересно. Возможно, вам будет достаточно того, что в свободное время мы пишем и "генерируем" подобие документации: https://visionatrix.github.io/VixFlowsDocs/

Почему именно подобие? Потому что проект быстро развивается, параллельно развивается ComfyUI, причём разработчики ComfyUI не особо информируют о будущих изменениях, поэтому нужно быть готовыми к резким неожиданным изменениям.

Мы также видим, как быстро и динамично развивается область генеративного ИИ в целом, и если вы ограничены по времени, вам лучше сосредоточиться на тех частях, которые в ближайшем будущем меняться не будут.

Как мы подходили к разработке и что делали в первую очередь

Использование LLM для помощи в разработке

Сразу оговорюсь: в программировании конкретно автор этой статьи более 10 лет, начал программировать ещё в школе на ассемблере и Delphi, а затем долгое время писал на C/C++.

Идея Visionatrix родилась в начале 2024 года, когда уже была доступна ChatGPT-4.

Мой коллега и я оба работаем полный рабочий день в сфере open-source и интеграций с ИИ, не касаясь диффузии и ComfyUI. После 8-часового рабочего дня остаётся мало сил, и на хобби-проект удаётся выделять 2-4 часа в сутки. Иногда удаётся найти немного времени на выходных, но его всё равно недостаточно для крупных проектов.

С учётом ограниченного времени и желания ускорить разработку, мы решили использовать генеративный ИИ для написания кода бэкенда. Основная цель состояла в том, чтобы ChatGPT/Claude мог эффективно работать с самыми популярными Python-библиотеками, такими как httpx, FastAPI, Pydantic и SQLAlchemy.

Также было важно, чтобы проект поддерживал и синхронный, и асинхронный режим. Те, кто знаком с Python, понимают, что это создаёт серьёзные трудности, так как необходимость поддержки обоих режимов значительно увеличивает сложность и объём кода.

Почему нужна синхронность помимо асинхронности? Потому что в будущем, возможно, с популяризацией Python версии 3.13+ синхронные проекты получат вторую жизнь; всё-таки настоящая многопоточность есть многопоточность.

У нас один человек писал бэкенд часть, а второй — в основном UI и занимался тестированием. ИИ пока что использовали в основном для бэкенда.

Весь проект был разделён на файлы, которые удобно «скармливать» ИИ, а именно:

База данных:

• database.py

• db_queries.py

• db_queries_async.py

Pydantic-модели вынесены в отдельный файл pydantic_models.py — ChatGPT очень любит получать их для контекста.

Вся логика работы с базой данных, отношениями между тасками, дочерними тасками (да, мы сделали и это; удобно, когда результат таска созданного из текущего можно увидеть, не покидая страницу текущего workflow-а) и всем прочим тоже вынесена в отдельные файлы:

• tasks_engine.py

• tasks_engine_async.py

Это позволило не писать функции для работы с БД самим; нужно было только написать первые 4-5 штук, а далее, при запросах на расширение функционала, писать что-то вроде:

"Пожалуйста, на основе файлов database.py, pydantic_models.py, db_queries.py, tasks_engine.py добавь возможность иметь таскам приоритет, добавь или расширь существующие синхронные функции для этого. Сохраняй стиль написания как в оригинальных файлах.

pydantic_models.py:

содержимое файла

db_queries.py:

содержимое файла"

А затем отдельно следующий запрос, но уже без pydantic_models.py и database.py, а просто сбрасывать в тот же чат файлы tasks_engine_async.py и db_queries_async.py и писать: "а теперь давай тоже самое, но для асинхронных реализаций".

Хотя за ChatGPT-4 приходилось многое править вручную, начиная с начала лета она достаточно поумнела и начала выдавать код гораздо более высокого качества, требующего меньше правок.

Этой осенью с таким подходом это позволило, тратя гораздо меньше усилий, расширять новый функционал, добавляя разные штуки, даже после работы, уставшим, лёжа на диване, так как вы просто просите, что вам нужно, и потом, если что, просите переделать.

Сейчас, когда доступна o1-preview, ситуация одновременно стала и лучше, и хуже. Как так?

Запросы для o1-preview составить гораздо сложнее, но при хорошем промпте она выдаёт код гораздо лучше, чем ChatGPT-4o — но результаты гораздо тяжелее оценивать, так как код, сгенерированный o1-preview, сложнее, напоминает написанный быстро человеком код, а потом сразу переписанный с оптимизациями.

Но в целом подход остался таким же: чёткое разделение по файлам функционала, чтобы было максимально удобно предоставлять код ИИ и потом его оценивать.

Также LLM очень любит, когда код максимально типизирован, и это позволяет быстрее отлавливать ошибки с помощью утилит ruff или pylint.

Всё-таки pre-commit в Python — это то, за что мы любим этот язык, когда множество ошибок можно отловить, даже не запуская код.

Использование LLM для написания тестов

Изначально мы вообще не собирались писать тесты, так как это хобби-проект и не было на них времени.

Но потом ситуация немного изменилась, так как, во-первых, захотелось измерить производительность между версиями PyTorch, самого ComfyUI, параметров, с которыми это всё можно запускать, да и протестировать это всё на разном железе.

Изначально пару раз делали это вручную, но уже после нескольких попыток стало понятно, что так продолжаться далее не может.

Это было уже этой осенью, во времена ChatGPT-4o и o1-preview.

У нас были файлы с исходниками бэкенда, openapi.json, доступ к LLM и неимоверное желание иметь тесты...

Мы дали всё это LLM, и, объяснив ситуацию, получили то, что ожидали, — почти рабочий код бенчмаркинга с первого раза.

Кому интересно, вот практически весь этот файл был написан ChatGPT:

https://github.com/Visionatrix/VixFlowsDocs/blob/main/scripts/benchmarks/benchmark.py

Что же было выяснено интересного в процессе написания?

Пришло осознание, что нам понадобилось более 5 запросов к LLM для генерации рабочего кода.

Может быть, это из-за недостатков LLM? Возможно, но, возможно, причина в том, что документация на бэкенд (openapi.json) и сами алгоритмы бэкенда не достаточно хороши.

LLM можно представить как разработчика, который зашёл в репозиторий и попробовал использовать API проекта. У него не выходит с первой, второй, третьей попытки, и он идёт дальше, искать софт на который написать интеграцию проще.

И это показатель того, насколько другим разработчикам легко будет использовать ваш продукт.

Хотя тесты и были в итоге написаны нормально с помощью ИИ, мы арендовали пару инстансов с разными видеокартами типа 4090, A100, H100 и прогнали их в черновую. Для их написания пришлось объяснять LLM, что для создания таска из ComfyUI workflow-а нужно передавать все аргументы и все файлы по очереди как параметры в форме, какой эндпоинт вызывать, чтобы получить статус таска, как получать результаты.

Это всё оказалось довольно неочевидным с первого раза для LLM (и, судя по всему, будет неочевидно и для человека) и требует дальнейших улучшений и доработки.

Главное, что было выяснено: насколько чётко составлен проект и его документация OpenAPI, можно проверить, используя LLM, попросив написать на него скрипт для тестов или просто UI на Gradio.

Чем меньше вопросов вам придётся отвечать ИИ, тем лучше вы справились, и это же уменьшит время в будущем на поддержку проекта, так как вы сможете отдать больше задач на аутсорс джуниорам ИИ и ускорить процесс разработки, избавив себя от рутины.

Что получилось в итоге

Мы создали решение, который легко устанавливается как на Linux, так и на macOS, а также работает на Windows.

Использует базу данных для сохранения истории тасков, которые пользователь создаёт из workflows.

Поддерживает работу с базами данных SQLite и PostgreSQL.

Поддерживается подключение к одному инстансу другого инстанса, установленного как локально, так и удалённо.

Поскольку всё это сделано на ComfyUI, поддерживается большинство возможностей ComfyUI.

Реализована поддержка перевода промптов с помощью Gemini и Ollama.

Написан небольшой скрипт для бенчмаркинга, который используется также для тестов.

Написан небольшой скрипт для GitHub-a, который проверяет, всё ли в порядке с моделям которые используются в workflow-ах.

Чего не хватает проекту?

• Расширение промпта: например, модель для рисования не знает, кто такой Ведьмак Геральт из Ривии; LLM должна преобразовать и упростить промпт для получения более лучшего результата.

• Более удобный API для создания тасков, а также спецификации OpenAPI для каждого flow — это позволит гораздо проще писать интеграции для проекта, а также создавать демки на основе Gradio с помощью одного запроса к LLM.

• Автоматическая генерация документации и примеров для workflows, с помощью всё того же LLM.

• Разработка универсального метода ввода в формате чата, когда на основе какого-либо результата вы просите переделать то, что вам не нравится, возможно словами, возможно выбором объектов, после чего LLM сама должна выбрать, какие workflows использовать для достижения результата (максимальная лёгкость использования, насколько это возможно, когда вы сами изначально не знаете, что хотите).

Заключение

Разработка Visionatrix стала для нас интересным опытом создания надстройки над ComfyUI, направленной на упрощение использования и расширение возможностей для обычных пользователей. Мы продолжаем активно развивать проект и будем рады, если он заинтересует сообщество.

Если у вас есть идеи, предложения или вы хотите внести свой вклад, мы будем рады вашему участию. Проект открыт для контрибьюторов, и любая помощь приветствуется. Также, если вам понравился проект, не забудьте поставить ⭐️ на GitHub — нам будет очень приятно.

Если вас заинтересовали какие-либо из затронутых в статье тем, дайте знать в комментариях.