Skip to content

Commit 50c605e

Browse files
Add support for sqlite database (comfyanonymous#8444)
* Add support for sqlite database * fix
1 parent 9685d4f commit 50c605e

File tree

11 files changed

+345
-13
lines changed

11 files changed

+345
-13
lines changed

alembic.ini

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
# A generic, single database configuration.
2+
3+
[alembic]
4+
# path to migration scripts
5+
# Use forward slashes (/) also on windows to provide an os agnostic path
6+
script_location = alembic_db
7+
8+
# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
9+
# Uncomment the line below if you want the files to be prepended with date and time
10+
# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
11+
# for all available tokens
12+
# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
13+
14+
# sys.path path, will be prepended to sys.path if present.
15+
# defaults to the current working directory.
16+
prepend_sys_path = .
17+
18+
# timezone to use when rendering the date within the migration file
19+
# as well as the filename.
20+
# If specified, requires the python>=3.9 or backports.zoneinfo library and tzdata library.
21+
# Any required deps can installed by adding `alembic[tz]` to the pip requirements
22+
# string value is passed to ZoneInfo()
23+
# leave blank for localtime
24+
# timezone =
25+
26+
# max length of characters to apply to the "slug" field
27+
# truncate_slug_length = 40
28+
29+
# set to 'true' to run the environment during
30+
# the 'revision' command, regardless of autogenerate
31+
# revision_environment = false
32+
33+
# set to 'true' to allow .pyc and .pyo files without
34+
# a source .py file to be detected as revisions in the
35+
# versions/ directory
36+
# sourceless = false
37+
38+
# version location specification; This defaults
39+
# to alembic_db/versions. When using multiple version
40+
# directories, initial revisions must be specified with --version-path.
41+
# The path separator used here should be the separator specified by "version_path_separator" below.
42+
# version_locations = %(here)s/bar:%(here)s/bat:alembic_db/versions
43+
44+
# version path separator; As mentioned above, this is the character used to split
45+
# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep.
46+
# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas.
47+
# Valid values for version_path_separator are:
48+
#
49+
# version_path_separator = :
50+
# version_path_separator = ;
51+
# version_path_separator = space
52+
# version_path_separator = newline
53+
#
54+
# Use os.pathsep. Default configuration used for new projects.
55+
version_path_separator = os
56+
57+
# set to 'true' to search source files recursively
58+
# in each "version_locations" directory
59+
# new in Alembic version 1.10
60+
# recursive_version_locations = false
61+
62+
# the output encoding used when revision files
63+
# are written from script.py.mako
64+
# output_encoding = utf-8
65+
66+
sqlalchemy.url = sqlite:///user/comfyui.db
67+
68+
69+
[post_write_hooks]
70+
# post_write_hooks defines scripts or Python functions that are run
71+
# on newly generated revision scripts. See the documentation for further
72+
# detail and examples
73+
74+
# format using "black" - use the console_scripts runner, against the "black" entrypoint
75+
# hooks = black
76+
# black.type = console_scripts
77+
# black.entrypoint = black
78+
# black.options = -l 79 REVISION_SCRIPT_FILENAME
79+
80+
# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
81+
# hooks = ruff
82+
# ruff.type = exec
83+
# ruff.executable = %(here)s/.venv/bin/ruff
84+
# ruff.options = check --fix REVISION_SCRIPT_FILENAME

alembic_db/README.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
## Generate new revision
2+
3+
1. Update models in `/app/database/models.py`
4+
2. Run `alembic revision --autogenerate -m "{your message}"`

alembic_db/env.py

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
from sqlalchemy import engine_from_config
2+
from sqlalchemy import pool
3+
4+
from alembic import context
5+
6+
# this is the Alembic Config object, which provides
7+
# access to the values within the .ini file in use.
8+
config = context.config
9+
10+
11+
from app.database.models import Base
12+
target_metadata = Base.metadata
13+
14+
# other values from the config, defined by the needs of env.py,
15+
# can be acquired:
16+
# my_important_option = config.get_main_option("my_important_option")
17+
# ... etc.
18+
19+
20+
def run_migrations_offline() -> None:
21+
"""Run migrations in 'offline' mode.
22+
This configures the context with just a URL
23+
and not an Engine, though an Engine is acceptable
24+
here as well. By skipping the Engine creation
25+
we don't even need a DBAPI to be available.
26+
Calls to context.execute() here emit the given string to the
27+
script output.
28+
"""
29+
url = config.get_main_option("sqlalchemy.url")
30+
context.configure(
31+
url=url,
32+
target_metadata=target_metadata,
33+
literal_binds=True,
34+
dialect_opts={"paramstyle": "named"},
35+
)
36+
37+
with context.begin_transaction():
38+
context.run_migrations()
39+
40+
41+
def run_migrations_online() -> None:
42+
"""Run migrations in 'online' mode.
43+
In this scenario we need to create an Engine
44+
and associate a connection with the context.
45+
"""
46+
connectable = engine_from_config(
47+
config.get_section(config.config_ini_section, {}),
48+
prefix="sqlalchemy.",
49+
poolclass=pool.NullPool,
50+
)
51+
52+
with connectable.connect() as connection:
53+
context.configure(
54+
connection=connection, target_metadata=target_metadata
55+
)
56+
57+
with context.begin_transaction():
58+
context.run_migrations()
59+
60+
61+
if context.is_offline_mode():
62+
run_migrations_offline()
63+
else:
64+
run_migrations_online()

alembic_db/script.py.mako

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
"""${message}
2+
3+
Revision ID: ${up_revision}
4+
Revises: ${down_revision | comma,n}
5+
Create Date: ${create_date}
6+
7+
"""
8+
from typing import Sequence, Union
9+
10+
from alembic import op
11+
import sqlalchemy as sa
12+
${imports if imports else ""}
13+
14+
# revision identifiers, used by Alembic.
15+
revision: str = ${repr(up_revision)}
16+
down_revision: Union[str, None] = ${repr(down_revision)}
17+
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
18+
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
19+
20+
21+
def upgrade() -> None:
22+
"""Upgrade schema."""
23+
${upgrades if upgrades else "pass"}
24+
25+
26+
def downgrade() -> None:
27+
"""Downgrade schema."""
28+
${downgrades if downgrades else "pass"}

app/database/db.py

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
import logging
2+
import os
3+
import shutil
4+
from app.logger import log_startup_warning
5+
from utils.install_util import get_missing_requirements_message
6+
from comfy.cli_args import args
7+
8+
_DB_AVAILABLE = False
9+
Session = None
10+
11+
12+
try:
13+
from alembic import command
14+
from alembic.config import Config
15+
from alembic.runtime.migration import MigrationContext
16+
from alembic.script import ScriptDirectory
17+
from sqlalchemy import create_engine
18+
from sqlalchemy.orm import sessionmaker
19+
20+
_DB_AVAILABLE = True
21+
except ImportError as e:
22+
log_startup_warning(
23+
f"""
24+
------------------------------------------------------------------------
25+
Error importing dependencies: {e}
26+
{get_missing_requirements_message()}
27+
This error is happening because ComfyUI now uses a local sqlite database.
28+
------------------------------------------------------------------------
29+
""".strip()
30+
)
31+
32+
33+
def dependencies_available():
34+
"""
35+
Temporary function to check if the dependencies are available
36+
"""
37+
return _DB_AVAILABLE
38+
39+
40+
def can_create_session():
41+
"""
42+
Temporary function to check if the database is available to create a session
43+
During initial release there may be environmental issues (or missing dependencies) that prevent the database from being created
44+
"""
45+
return dependencies_available() and Session is not None
46+
47+
48+
def get_alembic_config():
49+
root_path = os.path.join(os.path.dirname(__file__), "../..")
50+
config_path = os.path.abspath(os.path.join(root_path, "alembic.ini"))
51+
scripts_path = os.path.abspath(os.path.join(root_path, "alembic_db"))
52+
53+
config = Config(config_path)
54+
config.set_main_option("script_location", scripts_path)
55+
config.set_main_option("sqlalchemy.url", args.database_url)
56+
57+
return config
58+
59+
60+
def get_db_path():
61+
url = args.database_url
62+
if url.startswith("sqlite:///"):
63+
return url.split("///")[1]
64+
else:
65+
raise ValueError(f"Unsupported database URL '{url}'.")
66+
67+
68+
def init_db():
69+
db_url = args.database_url
70+
logging.debug(f"Database URL: {db_url}")
71+
db_path = get_db_path()
72+
db_exists = os.path.exists(db_path)
73+
74+
config = get_alembic_config()
75+
76+
# Check if we need to upgrade
77+
engine = create_engine(db_url)
78+
conn = engine.connect()
79+
80+
context = MigrationContext.configure(conn)
81+
current_rev = context.get_current_revision()
82+
83+
script = ScriptDirectory.from_config(config)
84+
target_rev = script.get_current_head()
85+
86+
if target_rev is None:
87+
logging.warning("No target revision found.")
88+
elif current_rev != target_rev:
89+
# Backup the database pre upgrade
90+
backup_path = db_path + ".bkp"
91+
if db_exists:
92+
shutil.copy(db_path, backup_path)
93+
else:
94+
backup_path = None
95+
96+
try:
97+
command.upgrade(config, target_rev)
98+
logging.info(f"Database upgraded from {current_rev} to {target_rev}")
99+
except Exception as e:
100+
if backup_path:
101+
# Restore the database from backup if upgrade fails
102+
shutil.copy(backup_path, db_path)
103+
os.remove(backup_path)
104+
logging.exception("Error upgrading database: ")
105+
raise e
106+
107+
global Session
108+
Session = sessionmaker(bind=engine)
109+
110+
111+
def create_session():
112+
return Session()

app/database/models.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
from sqlalchemy.orm import declarative_base
2+
3+
Base = declarative_base()
4+
5+
6+
def to_dict(obj):
7+
fields = obj.__table__.columns.keys()
8+
return {
9+
field: (val.to_dict() if hasattr(val, "to_dict") else val)
10+
for field in fields
11+
if (val := getattr(obj, field))
12+
}
13+
14+
# TODO: Define models here

app/frontend_management.py

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,26 +16,17 @@
1616
import requests
1717
from typing_extensions import NotRequired
1818

19+
from utils.install_util import get_missing_requirements_message, requirements_path
20+
1921
from comfy.cli_args import DEFAULT_VERSION_STRING
2022
import app.logger
2123

22-
# The path to the requirements.txt file
23-
req_path = Path(__file__).parents[1] / "requirements.txt"
24-
2524

2625
def frontend_install_warning_message():
27-
"""The warning message to display when the frontend version is not up to date."""
28-
29-
extra = ""
30-
if sys.flags.no_user_site:
31-
extra = "-s "
3226
return f"""
33-
Please install the updated requirements.txt file by running:
34-
{sys.executable} {extra}-m pip install -r {req_path}
27+
{get_missing_requirements_message()}
3528
3629
This error is happening because the ComfyUI frontend is no longer shipped as part of the main repo but as a pip package instead.
37-
38-
If you are on the portable package you can run: update\\update_comfyui.bat to solve this problem
3930
""".strip()
4031

4132

@@ -48,7 +39,7 @@ def parse_version(version: str) -> tuple[int, int, int]:
4839
try:
4940
frontend_version_str = version("comfyui-frontend-package")
5041
frontend_version = parse_version(frontend_version_str)
51-
with open(req_path, "r", encoding="utf-8") as f:
42+
with open(requirements_path, "r", encoding="utf-8") as f:
5243
required_frontend = parse_version(f.readline().split("=")[-1])
5344
if frontend_version < required_frontend:
5445
app.logger.log_startup_warning(

comfy/cli_args.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,11 @@ def is_valid_directory(path: str) -> str:
203203
help="Set the base URL for the ComfyUI API. (default: https://api.comfy.org)",
204204
)
205205

206+
database_default_path = os.path.abspath(
207+
os.path.join(os.path.dirname(__file__), "..", "user", "comfyui.db")
208+
)
209+
parser.add_argument("--database-url", type=str, default=f"sqlite:///{database_default_path}", help="Specify the database URL, e.g. for an in-memory database you can use 'sqlite:///:memory:'.")
210+
206211
if comfy.options.args_parsing:
207212
args = parser.parse_args()
208213
else:

main.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,15 @@ def cleanup_temp():
238238
shutil.rmtree(temp_dir, ignore_errors=True)
239239

240240

241+
def setup_database():
242+
try:
243+
from app.database.db import init_db, dependencies_available
244+
if dependencies_available():
245+
init_db()
246+
except Exception as e:
247+
logging.error(f"Failed to initialize database. Please ensure you have installed the latest requirements. If the error persists, please report this as in future the database will be required: {e}")
248+
249+
241250
def start_comfyui(asyncio_loop=None):
242251
"""
243252
Starts the ComfyUI server using the provided asyncio event loop or creates a new one.
@@ -266,6 +275,7 @@ def start_comfyui(asyncio_loop=None):
266275
hook_breaker_ac10a0.restore_functions()
267276

268277
cuda_malloc_warning()
278+
setup_database()
269279

270280
prompt_server.add_routes()
271281
hijack_progress(prompt_server)

requirements.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ Pillow
1818
scipy
1919
tqdm
2020
psutil
21+
alembic
22+
SQLAlchemy
2123

2224
#non essential dependencies:
2325
kornia>=0.7.1

0 commit comments

Comments
 (0)