Add additional db model metadata fields and model downloading function

app/database/db.py

@@ -1,29 +1,50 @@
 import logging
 import os
 import shutil
+from app.logger import log_startup_warning
 from utils.install_util import get_missing_requirements_message
 from comfy.cli_args import args
 
+_DB_AVAILABLE = False
 Session = None
 
 
-def can_create_session():
-    return Session is not None
-
-
 try:
-    import alembic
-    import sqlalchemy
-except ImportError as e:
-    logging.error(get_missing_requirements_message())
-    raise e
+    from alembic import command
+    from alembic.config import Config
+    from alembic.runtime.migration import MigrationContext
+    from alembic.script import ScriptDirectory
+    from sqlalchemy import create_engine
+    from sqlalchemy.orm import sessionmaker
 
-from alembic import command
-from alembic.config import Config
-from alembic.runtime.migration import MigrationContext
-from alembic.script import ScriptDirectory
-from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker
+    _DB_AVAILABLE = True
+except ImportError as e:
+    log_startup_warning(
+        f"""
+------------------------------------------------------------------------
+Error importing dependencies: {e}
+
+{get_missing_requirements_message()}
+
+This error is happening because ComfyUI now uses a local sqlite database.
+------------------------------------------------------------------------
+""".strip()
+    )
+
+
+def dependencies_available():
+    """
+    Temporary function to check if the dependencies are available
+    """
+    return _DB_AVAILABLE
+
+
+def can_create_session():
+    """
+    Temporary function to check if the database is available to create a session
+    During initial release there may be environmental issues (or missing dependencies) that prevent the database from being created
+    """
+    return dependencies_available() and Session is not None
 
 
 def get_alembic_config():
@@ -49,6 +70,8 @@ def get_db_path():
 def init_db():
     db_url = args.database_url
     logging.debug(f"Database URL: {db_url}")
+    db_path = get_db_path()
+    db_exists = os.path.exists(db_path)
 
     config = get_alembic_config()
 
@@ -64,9 +87,8 @@ def init_db():
 
     if current_rev != target_rev:
         # Backup the database pre upgrade
-        db_path = get_db_path()
         backup_path = db_path + ".bkp"
-        if os.path.exists(db_path):
+        if db_exists:
             shutil.copy(db_path, backup_path)
         else:
             backup_path = None

app/database/models.py

@@ -1,5 +1,6 @@
 from sqlalchemy import (
     Column,
+    Integer,
     Text,
     DateTime,
 )
@@ -20,7 +21,7 @@ def to_dict(obj):
 
 class Model(Base):
     """
-    SQLAlchemy model representing a model file in the system.
+    sqlalchemy model representing a model file in the system.
 
     This class defines the database schema for storing information about model files,
     including their type, path, hash, and when they were added to the system.
@@ -28,7 +29,11 @@ class Model(Base):
     Attributes:
         type (Text): The type of the model, this is the name of the folder in the models folder (primary key)
         path (Text): The file path of the model relative to the type folder (primary key)
-        hash (Text): A sha256 hash of the model file
+        file_name (Text): The name of the model file
+        file_size (Integer): The size of the model file in bytes
+        hash (Text): A hash of the model file
+        hash_algorithm (Text): The algorithm used to generate the hash
+        source_url (Text): The URL of the model file
         date_added (DateTime): Timestamp of when the model was added to the system
     """
 
@@ -36,7 +41,11 @@ class Model(Base):
 
     type = Column(Text, primary_key=True)
     path = Column(Text, primary_key=True)
+    file_name = Column(Text)
+    file_size = Column(Integer)
     hash = Column(Text)
+    hash_algorithm = Column(Text)
+    source_url = Column(Text)
     date_added = Column(DateTime, server_default=func.now())
 
     def to_dict(self):