File size: 4,190 Bytes
9748021 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import sqlite3
import json
import os
from datetime import datetime
class DatabaseManager:
def __init__(self, db_path="lumaforge.db"):
self.db_path = db_path
self.init_db()
def _get_connection(self):
"""Creates a new SQLite database connection."""
conn = sqlite3.connect(self.db_path, check_same_thread=False)
conn.row_factory = sqlite3.Row
return conn
def init_db(self):
"""Initializes the database and creates the generations table if it doesn't exist."""
query = """
CREATE TABLE IF NOT EXISTS generations (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT UNIQUE,
prompt TEXT NOT NULL,
expanded_prompt TEXT, -- JSON string
negative_prompt TEXT,
steps INTEGER,
guidance_scale REAL,
seed INTEGER,
aspect_ratio TEXT,
device TEXT,
latency_sec REAL,
memory_used_mb REAL,
status TEXT, -- 'completed', 'refused', 'error'
image_path TEXT, -- Path to the stored file on disk
error_message TEXT, -- Holds error logs if failed
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
"""
conn = self._get_connection()
try:
with conn:
conn.execute(query)
print(f"[DatabaseManager] Database initialized at {os.path.abspath(self.db_path)}")
finally:
conn.close()
def log_generation(self, session_id: str, prompt: str, status: str,
expanded_prompt: dict = None, negative_prompt: str = None,
steps: int = None, guidance_scale: float = None,
seed: int = -1, aspect_ratio: str = "1:1", device: str = "mps",
latency_sec: float = 0.0, memory_used_mb: float = 0.0,
image_path: str = None, error_message: str = None):
"""Logs a generation run (success or failure) to the SQLite database."""
query = """
INSERT OR REPLACE INTO generations (
session_id, prompt, expanded_prompt, negative_prompt, steps,
guidance_scale, seed, aspect_ratio, device, latency_sec,
memory_used_mb, status, image_path, error_message
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
"""
# Serialize dict to JSON string
expanded_prompt_json = json.dumps(expanded_prompt) if expanded_prompt else None
conn = self._get_connection()
try:
with conn:
conn.execute(query, (
session_id, prompt, expanded_prompt_json, negative_prompt, steps,
guidance_scale, seed, aspect_ratio, device, latency_sec,
memory_used_mb, status, image_path, error_message
))
print(f"[DatabaseManager] Logged session {session_id} to database (status: {status})")
except Exception as e:
print(f"[DatabaseManager Error] Failed to log session {session_id}: {e}")
finally:
conn.close()
def get_history(self, limit: int = 50) -> list:
"""Retrieves historical generation logs, returning a list of dicts."""
query = "SELECT * FROM generations ORDER BY created_at DESC LIMIT ?;"
conn = self._get_connection()
try:
cursor = conn.cursor()
cursor.execute(query, (limit,))
rows = cursor.fetchall()
history = []
for row in rows:
item = dict(row)
# Deserialize expanded prompt back to dict
if item["expanded_prompt"]:
try:
item["expanded_prompt"] = json.loads(item["expanded_prompt"])
except Exception:
pass
history.append(item)
return history
except Exception as e:
print(f"[DatabaseManager Error] Failed to retrieve history: {e}")
return []
finally:
conn.close()
|