Spaces:
Running
Running
File size: 8,466 Bytes
ed37502 27fea48 ed37502 27fea48 739bb39 b341f22 c231a93 b341f22 c231a93 ed37502 b341f22 ed37502 | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | """SQLAlchemy database models for the content catalog and job queue."""
from __future__ import annotations
from datetime import datetime
from sqlalchemy import (
Boolean,
DateTime,
Float,
Index,
Integer,
String,
Text,
func,
)
from sqlalchemy.ext.asyncio import AsyncAttrs, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from content_engine.config import settings
class Base(AsyncAttrs, DeclarativeBase):
pass
class Character(Base):
__tablename__ = "characters"
id: Mapped[str] = mapped_column(String(64), primary_key=True)
name: Mapped[str] = mapped_column(String(128), nullable=False)
trigger_word: Mapped[str] = mapped_column(String(128), nullable=False)
lora_filename: Mapped[str] = mapped_column(String(256), nullable=False)
lora_strength: Mapped[float] = mapped_column(Float, default=0.85)
default_checkpoint: Mapped[str | None] = mapped_column(String(256))
description: Mapped[str | None] = mapped_column(Text)
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.now()
)
class Image(Base):
__tablename__ = "images"
id: Mapped[str] = mapped_column(String(36), primary_key=True)
batch_id: Mapped[str | None] = mapped_column(String(36), index=True)
character_id: Mapped[str | None] = mapped_column(String(64), index=True)
template_id: Mapped[str | None] = mapped_column(String(128))
content_rating: Mapped[str] = mapped_column(String(8), index=True) # sfw | nsfw
# Generation parameters
positive_prompt: Mapped[str | None] = mapped_column(Text)
negative_prompt: Mapped[str | None] = mapped_column(Text)
checkpoint: Mapped[str | None] = mapped_column(String(256))
loras_json: Mapped[str | None] = mapped_column(Text) # JSON array
seed: Mapped[int | None] = mapped_column(Integer)
steps: Mapped[int | None] = mapped_column(Integer)
cfg: Mapped[float | None] = mapped_column(Float)
sampler: Mapped[str | None] = mapped_column(String(64))
scheduler: Mapped[str | None] = mapped_column(String(64))
width: Mapped[int | None] = mapped_column(Integer)
height: Mapped[int | None] = mapped_column(Integer)
# Searchable variation attributes
pose: Mapped[str | None] = mapped_column(String(128))
outfit: Mapped[str | None] = mapped_column(String(128))
emotion: Mapped[str | None] = mapped_column(String(128))
camera_angle: Mapped[str | None] = mapped_column(String(128))
lighting: Mapped[str | None] = mapped_column(String(128))
scene: Mapped[str | None] = mapped_column(String(128))
# File info
file_path: Mapped[str] = mapped_column(String(512), nullable=False)
file_hash: Mapped[str | None] = mapped_column(String(64))
file_size: Mapped[int | None] = mapped_column(Integer)
generation_backend: Mapped[str | None] = mapped_column(String(32)) # local | cloud
comfyui_prompt_id: Mapped[str | None] = mapped_column(String(36))
generation_time_seconds: Mapped[float | None] = mapped_column(Float)
# Quality and publishing
quality_score: Mapped[float | None] = mapped_column(Float)
is_approved: Mapped[bool] = mapped_column(Boolean, default=False)
is_published: Mapped[bool] = mapped_column(Boolean, default=False)
published_platform: Mapped[str | None] = mapped_column(String(64))
published_at: Mapped[datetime | None] = mapped_column(DateTime)
scheduled_at: Mapped[datetime | None] = mapped_column(DateTime)
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.now()
)
__table_args__ = (
Index("idx_images_approved", "is_approved", postgresql_where=(is_approved == True)), # noqa: E712
Index(
"idx_images_unpublished",
"is_published",
"is_approved",
),
)
class GenerationJob(Base):
__tablename__ = "generation_jobs"
id: Mapped[str] = mapped_column(String(36), primary_key=True)
batch_id: Mapped[str | None] = mapped_column(String(36), index=True)
character_id: Mapped[str | None] = mapped_column(String(64))
template_id: Mapped[str | None] = mapped_column(String(128))
content_rating: Mapped[str | None] = mapped_column(String(8))
variables_json: Mapped[str | None] = mapped_column(Text)
workflow_json: Mapped[str | None] = mapped_column(Text)
backend: Mapped[str | None] = mapped_column(String(32)) # local | replicate | runpod
status: Mapped[str] = mapped_column(
String(16), default="pending", index=True
) # pending | queued | running | completed | failed
comfyui_prompt_id: Mapped[str | None] = mapped_column(String(36))
cloud_job_id: Mapped[str | None] = mapped_column(String(128))
result_image_id: Mapped[str | None] = mapped_column(String(36))
error_message: Mapped[str | None] = mapped_column(Text)
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.now()
)
started_at: Mapped[datetime | None] = mapped_column(DateTime)
completed_at: Mapped[datetime | None] = mapped_column(DateTime)
class ScheduledPost(Base):
__tablename__ = "scheduled_posts"
id: Mapped[str] = mapped_column(String(36), primary_key=True)
image_id: Mapped[str] = mapped_column(String(36), nullable=False)
platform: Mapped[str] = mapped_column(String(64), nullable=False)
scheduled_at: Mapped[datetime] = mapped_column(DateTime, nullable=False)
caption: Mapped[str | None] = mapped_column(Text)
is_teaser: Mapped[bool] = mapped_column(Boolean, default=False)
status: Mapped[str] = mapped_column(
String(16), default="pending"
) # pending | published | failed | cancelled
published_at: Mapped[datetime | None] = mapped_column(DateTime)
error_message: Mapped[str | None] = mapped_column(Text)
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.now()
)
__table_args__ = (
Index("idx_scheduled_pending", "status", "scheduled_at"),
)
class TrainingJob(Base):
__tablename__ = "training_jobs"
id: Mapped[str] = mapped_column(String(36), primary_key=True)
name: Mapped[str] = mapped_column(String(128), nullable=False)
status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
progress: Mapped[float] = mapped_column(Float, default=0.0)
current_epoch: Mapped[int] = mapped_column(Integer, default=0)
total_epochs: Mapped[int] = mapped_column(Integer, default=0)
current_step: Mapped[int] = mapped_column(Integer, default=0)
total_steps: Mapped[int] = mapped_column(Integer, default=0)
loss: Mapped[float | None] = mapped_column(Float)
started_at: Mapped[float | None] = mapped_column(Float)
completed_at: Mapped[float | None] = mapped_column(Float)
output_path: Mapped[str | None] = mapped_column(String(512))
error: Mapped[str | None] = mapped_column(Text)
log_text: Mapped[str | None] = mapped_column(Text) # newline-separated log lines
pod_id: Mapped[str | None] = mapped_column(String(64))
gpu_type: Mapped[str | None] = mapped_column(String(64))
backend: Mapped[str] = mapped_column(String(16), default="runpod")
base_model: Mapped[str | None] = mapped_column(String(64))
model_type: Mapped[str | None] = mapped_column(String(16))
trigger_word: Mapped[str | None] = mapped_column(String(128))
image_upload_dir: Mapped[str | None] = mapped_column(String(512))
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.now()
)
# --- Engine / Session factories ---
# Ensure database directories exist (critical for HF Spaces first run)
import logging as _logging
from pathlib import Path as _Path
_db_url = settings.database.url
_db_path = _db_url.replace("sqlite+aiosqlite:///", "")
_db_dir = _Path(_db_path).parent
_db_dir.mkdir(parents=True, exist_ok=True)
_logging.getLogger(__name__).info("DB path: %s (dir exists: %s)", _db_path, _db_dir.exists())
_catalog_engine = create_async_engine(
_db_url,
echo=False,
connect_args={"check_same_thread": False}, # SQLite specific
)
catalog_session_factory = async_sessionmaker(
_catalog_engine, expire_on_commit=False
)
async def init_db() -> None:
"""Create all tables. Call once at startup."""
async with _catalog_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
|