Spaces:
Running
Running
File size: 2,217 Bytes
e302830 3bbf2c7 e302830 3bbf2c7 e302830 3bbf2c7 e302830 3bbf2c7 e302830 3bbf2c7 |
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 |
import sqlite3
from pathlib import Path
from typing import Any, Dict
from pydantic import BaseModel
class PersistentSettings(BaseModel):
"""
This pydantic model will try to initialize itself from
the database upon every instantiation
It further supplies an update function, that allows to write
back any changes into the database, under its key.
"""
def __init__(self, **data: Any):
# Connect to the SQLite database
self.conn = sqlite3.connect("config.db")
# Create a table for settings if it doesn't exist
self.conn.execute("""
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value TEXT
)
""")
# Fetch settings from the database and initialize
super().__init__(**self.fetch_settings(), **data)
def fetch_settings(self) -> Dict[str, Any]:
"""
Retrieve settings from the database
"""
cursor = self.conn.cursor()
cursor.execute("SELECT key, value FROM settings")
settings = dict(cursor.fetchall())
cursor.close()
return settings
def update(self, **data: Any) -> None:
"""
Persist the pydantic-dict that represents the model
"""
cursor = self.conn.cursor()
# Update or insert each key-value pair into the database
for key, value in {**self.dict(), **data}.items():
cursor.execute(
"INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)",
(key, value)
)
# Commit the changes to the database
self.conn.commit()
cursor.close()
def close(self) -> None:
"""
Close the database connection
"""
self.conn.close()
class TortoiseConfig(PersistentSettings):
EXTRA_VOICES_DIR: str = ""
AR_CHECKPOINT: str = "."
DIFF_CHECKPOINT: str = "."
LOW_VRAM: bool = True
def __init__(self, **data: Any):
super().__init__(**data)
if not Path(self.AR_CHECKPOINT).is_file():
self.AR_CHECKPOINT = "."
if not Path(self.DIFF_CHECKPOINT).is_file():
self.DIFF_CHECKPOINT = "."
|