radames's picture
db for cache
ee20381
raw
history blame
1.78 kB
import sqlite3
from pathlib import Path
class Database:
def __init__(self, db_path=None):
if db_path is None:
raise ValueError("db_path must be provided")
self.db_path = db_path
self.db_file = self.db_path / "cache.db"
if not self.db_file.exists():
print("Creating database")
print("DB_FILE", self.db_file)
db = sqlite3.connect(self.db_file)
with open(Path("schema.sql"), "r") as f:
db.executescript(f.read())
db.commit()
db.close()
def get_db(self):
db = sqlite3.connect(self.db_file, check_same_thread=False)
db.row_factory = sqlite3.Row
return db
def __enter__(self):
self.db = self.get_db()
return self.db
def __exit__(self, exc_type, exc_value, traceback):
self.db.close()
def __call__(self):
return self
def insert(self, prompt: str, negative_prompt: str, image_path: str, seed: int):
with self() as db:
cursor = db.cursor()
cursor.execute(
"INSERT INTO cache (prompt, negative_prompt, image_path, seed) VALUES (?, ?, ?, ?)",
(prompt, negative_prompt, image_path, seed),
)
db.commit()
def check(self, prompt: str, negative_prompt: str, seed: int):
with self() as db:
cursor = db.cursor()
cursor.execute(
"SELECT image_path FROM cache WHERE prompt = ? AND negative_prompt = ? AND seed = ? ORDER BY RANDOM() LIMIT 1",
(prompt, negative_prompt, seed),
)
image_path = cursor.fetchone()
if image_path:
return image_path
return False