87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
import os
|
|
import sqlite3
|
|
|
|
# --- DB Settings ---
|
|
DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "data")
|
|
os.makedirs(DATA_DIR, exist_ok=True)
|
|
DB_PATH = os.path.join(DATA_DIR, "default.db")
|
|
|
|
def init_db():
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS images (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
filename TEXT,
|
|
subfolder TEXT,
|
|
created_at REAL
|
|
)
|
|
''')
|
|
|
|
# Add is_deleted column if not exists
|
|
try:
|
|
cursor.execute("ALTER TABLE images ADD COLUMN is_deleted INTEGER DEFAULT 0")
|
|
except sqlite3.OperationalError:
|
|
pass
|
|
|
|
# Normalized Tables
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS positive_prompts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT UNIQUE
|
|
)
|
|
''')
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS negative_prompts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT UNIQUE
|
|
)
|
|
''')
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS tags (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT UNIQUE
|
|
)
|
|
''')
|
|
|
|
# Relations
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS positive_prompt_image_relations (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
image_id INTEGER,
|
|
positive_prompt_id INTEGER,
|
|
FOREIGN KEY(image_id) REFERENCES images(id),
|
|
FOREIGN KEY(positive_prompt_id) REFERENCES positive_prompts(id)
|
|
)
|
|
''')
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS negative_prompt_image_relations (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
image_id INTEGER,
|
|
negative_prompt_id INTEGER,
|
|
FOREIGN KEY(image_id) REFERENCES images(id),
|
|
FOREIGN KEY(negative_prompt_id) REFERENCES negative_prompts(id)
|
|
)
|
|
''')
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS tag_image_relations (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
image_id INTEGER,
|
|
tag_id INTEGER,
|
|
FOREIGN KEY(image_id) REFERENCES images(id),
|
|
FOREIGN KEY(tag_id) REFERENCES tags(id)
|
|
)
|
|
''')
|
|
|
|
# Indices
|
|
cursor.execute("CREATE INDEX IF NOT EXISTS idx_images_created_at ON images(created_at)")
|
|
cursor.execute("CREATE INDEX IF NOT EXISTS idx_pp_name ON positive_prompts(name)")
|
|
cursor.execute("CREATE INDEX IF NOT EXISTS idx_np_name ON negative_prompts(name)")
|
|
cursor.execute("CREATE INDEX IF NOT EXISTS idx_tag_name ON tags(name)")
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def get_db_connection():
|
|
return sqlite3.connect(DB_PATH)
|