| import os |
| import sqlite3 |
| import threading |
| from typing import List |
| from modules import shared |
|
|
| from scripts.mo.data.storage import Storage |
| from scripts.mo.environment import env, logger |
| from scripts.mo.models import Record, ModelType |
|
|
| _DB_FILE = 'database.sqlite' |
| _DB_VERSION = 6 |
| _DB_TIMEOUT = 30 |
|
|
|
|
| def map_row_to_record(row) -> Record: |
| return Record( |
| id_=row[0], |
| name=row[1], |
| model_type=ModelType.by_value(row[2]), |
| download_url=row[3], |
| url=row[4], |
| download_path=row[5], |
| download_filename=row[6], |
| preview_url=row[7], |
| description=row[8], |
| positive_prompts=row[9], |
| negative_prompts=row[10], |
| sha256_hash=row[11], |
| md5_hash=row[12], |
| created_at=row[13], |
| groups=row[14].split(',') if row[14] else [], |
| subdir=row[15], |
| location=row[16], |
| weight=row[17] |
| ) |
|
|
|
|
| class SQLiteStorage(Storage): |
|
|
| def __init__(self): |
| self.local = threading.local() |
| self._initialize() |
|
|
| def _connection(self): |
| if not hasattr(self.local, "connection"): |
| mo_database_dir = getattr(shared.cmd_opts, "mo_database_dir") |
| database_dir = mo_database_dir if mo_database_dir is not None else env.script_dir |
| db_file_path = os.path.join(database_dir, _DB_FILE) |
| self.local.connection = sqlite3.connect(db_file_path, _DB_TIMEOUT) |
| return self.local.connection |
|
|
| def _initialize(self): |
| cursor = self._connection().cursor() |
|
|
| cursor.execute('''CREATE TABLE IF NOT EXISTS Record |
| (id INTEGER PRIMARY KEY, |
| _name TEXT, |
| model_type TEXT, |
| download_url TEXT, |
| url TEXT DEFAULT '', |
| download_path TEXT DEFAULT '', |
| download_filename TEXT DEFAULT '', |
| preview_url TEXT DEFAULT '', |
| description TEXT DEFAULT '', |
| positive_prompts TEXT DEFAULT '', |
| negative_prompts TEXT DEFAULT '', |
| sha256_hash TEXT DEFAULT '', |
| md5_hash TEXT DEFAULT '', |
| created_at INTEGER DEFAULT 0, |
| groups TEXT DEFAULT '', |
| subdir TEXT DEFAULT '', |
| location TEXT DEFAULT '', |
| weight REAL DEFAULT 1) |
| ''') |
|
|
| cursor.execute(f'''CREATE TABLE IF NOT EXISTS Version |
| (version INTEGER DEFAULT {_DB_VERSION})''') |
| self._connection().commit() |
| self._check_database_version() |
|
|
| def _check_database_version(self): |
| cursor = self._connection().cursor() |
| cursor.execute('SELECT * FROM Version ', ) |
| row = cursor.fetchone() |
|
|
| if row is None: |
| cursor.execute(f'INSERT INTO Version VALUES ({_DB_VERSION})') |
| self._connection().commit() |
|
|
| version = _DB_VERSION if row is None else row[0] |
| if version != _DB_VERSION: |
| self._run_migration(version) |
|
|
| def _run_migration(self, current_version): |
| for ver in range(current_version, _DB_VERSION): |
| if ver == 1: |
| self._migrate_1_to_2() |
| elif ver == 2: |
| self._migrate_2_to_3() |
| elif ver == 3: |
| self._migrate_3_to_4() |
| elif ver == 4: |
| self._migrate_4_to_5() |
| elif ver == 5: |
| self._migrage_5_to_6() |
| else: |
| raise Exception(f'Missing SQLite migration from {ver} to {_DB_VERSION}') |
|
|
| def _migrate_1_to_2(self): |
| cursor = self._connection().cursor() |
| cursor.execute('ALTER TABLE Record ADD COLUMN created_at INTEGER DEFAULT 0;') |
| cursor.execute("DELETE FROM Version") |
| cursor.execute('INSERT INTO Version VALUES (2)') |
| self._connection().commit() |
|
|
| def _migrate_2_to_3(self): |
| cursor = self._connection().cursor() |
| cursor.execute("ALTER TABLE Record ADD COLUMN groups TEXT DEFAULT '';") |
| cursor.execute("DELETE FROM Version") |
| cursor.execute('INSERT INTO Version VALUES (3)') |
| self._connection().commit() |
|
|
| def _migrate_3_to_4(self): |
| cursor = self._connection().cursor() |
| cursor.execute("ALTER TABLE Record RENAME COLUMN model_hash TO sha256_hash;") |
| cursor.execute("ALTER TABLE Record ADD COLUMN subdir TEXT DEFAULT '';") |
| cursor.execute("DELETE FROM Version") |
| cursor.execute('INSERT INTO Version VALUES (4)') |
| self._connection().commit() |
|
|
| def _migrate_4_to_5(self): |
| cursor = self._connection().cursor() |
| cursor.execute("ALTER TABLE Record ADD COLUMN location TEXT DEFAULT '';") |
| cursor.execute("DELETE FROM Version") |
| cursor.execute('INSERT INTO Version VALUES (5)') |
| self._connection().commit() |
| |
| def _migrage_5_to_6(self): |
| cursor = self._connection().cursor() |
| cursor.execute("ALTER TABLE Record ADD COLUMN weight REAL DEFAULT 1;") |
| cursor.execute("DELETE FROM Version") |
| cursor.execute('INSERT INTO Version VALUES (6)') |
| self._connection().commit() |
|
|
| def get_all_records(self) -> List: |
| cursor = self._connection().cursor() |
| cursor.execute('SELECT * FROM Record') |
| rows = cursor.fetchall() |
| result = [] |
| for row in rows: |
| result.append(map_row_to_record(row)) |
| return result |
|
|
| def query_records(self, name_query: str = None, groups=None, model_types=None, show_downloaded=True, |
| show_not_downloaded=True) -> List: |
|
|
| query = 'SELECT * FROM Record' |
|
|
| is_where_appended = False |
| append_and = False |
|
|
| if name_query is not None and name_query: |
| if not is_where_appended: |
| query += ' WHERE' |
| is_where_appended = True |
|
|
| query += f" LOWER(_name) LIKE '%{name_query}%'" |
| append_and = True |
|
|
| if model_types is not None and len(model_types) > 0: |
| if not is_where_appended: |
| query += ' WHERE' |
| is_where_appended = True |
|
|
| if append_and: |
| query += ' AND' |
|
|
| query += ' (' |
| append_or = False |
| for model_type in model_types: |
| if append_or: |
| query += ' OR' |
| query += f" model_type='{model_type}'" |
| append_or = True |
|
|
| query += ')' |
|
|
| append_and = True |
| pass |
|
|
| if groups is not None and len(groups) > 0: |
| if not is_where_appended: |
| query += ' WHERE' |
|
|
| for group in groups: |
| if append_and: |
| query += ' AND' |
| query += f" LOWER(groups) LIKE '%{group}%'" |
| append_and = True |
|
|
| logger.debug(f'query: {query}') |
| cursor = self._connection().cursor() |
| cursor.execute(query) |
| rows = cursor.fetchall() |
| result = [] |
| for row in rows: |
| record = map_row_to_record(row) |
| is_downloaded = bool(record.location) and os.path.exists(record.location) |
|
|
| if show_downloaded and is_downloaded: |
| result.append(record) |
| elif show_not_downloaded and not is_downloaded: |
| result.append(record) |
|
|
| return result |
|
|
| def get_record_by_id(self, id_) -> Record: |
| cursor = self._connection().cursor() |
| cursor.execute('SELECT * FROM Record WHERE id=?', (id_,)) |
| row = cursor.fetchone() |
| return None if row is None else map_row_to_record(row) |
|
|
| def get_records_by_group(self, group: str) -> List: |
| cursor = self._connection().cursor() |
| cursor.execute(f"SELECT * FROM Record WHERE LOWER(groups) LIKE '%{group}%'") |
| rows = cursor.fetchall() |
| result = [] |
| for row in rows: |
| result.append(map_row_to_record(row)) |
| return result |
| |
| def get_records_by_query(self, query: str) -> List: |
| cursor = self._connection().cursor() |
| cursor.execute(query) |
| rows = cursor.fetchall() |
| result = [] |
| for row in rows: |
| result.append(map_row_to_record(row)) |
| return result |
|
|
| def add_record(self, record: Record): |
| cursor = self._connection().cursor() |
| data = ( |
| record.name, |
| record.model_type.value, |
| record.download_url, |
| record.url, |
| record.download_path, |
| record.download_filename, |
| record.preview_url, |
| record.description, |
| record.positive_prompts, |
| record.negative_prompts, |
| record.sha256_hash, |
| record.md5_hash, |
| record.created_at, |
| ",".join(record.groups), |
| record.subdir, |
| record.location, |
| record.weight |
| ) |
| cursor.execute( |
| """INSERT INTO Record( |
| _name, |
| model_type, |
| download_url, |
| url, |
| download_path, |
| download_filename, |
| preview_url, |
| description, |
| positive_prompts, |
| negative_prompts, |
| sha256_hash, |
| md5_hash, |
| created_at, |
| groups, |
| subdir, |
| location, |
| weight) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", |
| data) |
| self._connection().commit() |
|
|
| def update_record(self, record: Record): |
| cursor = self._connection().cursor() |
| data = ( |
| record.name, |
| record.model_type.value, |
| record.download_url, |
| record.url, |
| record.download_path, |
| record.download_filename, |
| record.preview_url, |
| record.description, |
| record.positive_prompts, |
| record.negative_prompts, |
| record.sha256_hash, |
| record.md5_hash, |
| ",".join(record.groups), |
| record.subdir, |
| record.location, |
| record.weight, |
| record.id_ |
| ) |
| cursor.execute( |
| """UPDATE Record SET |
| _name=?, |
| model_type=?, |
| download_url=?, |
| url=?, |
| download_path=?, |
| download_filename=?, |
| preview_url=?, |
| description=?, |
| positive_prompts=?, |
| negative_prompts=?, |
| sha256_hash=?, |
| md5_hash=?, |
| groups=?, |
| subdir=?, |
| location=?, |
| weight=? |
| WHERE id=? |
| """, data |
| ) |
|
|
| self._connection().commit() |
|
|
| def remove_record(self, _id): |
| cursor = self._connection().cursor() |
| cursor.execute("DELETE FROM Record WHERE id=?", (_id,)) |
| self._connection().commit() |
|
|
| def get_available_groups(self) -> List: |
| cursor = self._connection().cursor() |
| cursor.execute('SELECT groups FROM Record') |
| rows = cursor.fetchall() |
| result = [] |
| for row in rows: |
| if row[0]: |
| result.extend(row[0].split(",")) |
|
|
| result = list(set(result)) |
| return list(filter(None, result)) |
|
|
| def get_all_records_locations(self) -> List: |
| cursor = self._connection().cursor() |
| cursor.execute('SELECT location FROM Record') |
| rows = cursor.fetchall() |
| result = [] |
| for row in rows: |
| if row[0]: |
| result.append(row[0]) |
|
|
| return result |
|
|
