Add model management and database
- use sqlalchemy + alembic + sqlite for db - extract model data and previews - endpoints for db interactions - add tests
This commit is contained in:
pythongosssss
@@ -1,19 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
import logging
|
||||
from app.database.db import create_session
|
||||
import folder_paths
|
||||
import glob
|
||||
import comfy.utils
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
from folder_paths import map_legacy, filter_files_extensions, filter_files_content_types
|
||||
from folder_paths import map_legacy, filter_files_extensions, get_full_path
|
||||
from app.database.models import Tag, Model
|
||||
from app.model_processor import get_model_previews, model_processor
|
||||
from utils.web import dumps
|
||||
from sqlalchemy.orm import joinedload
|
||||
import sqlalchemy.exc
|
||||
|
||||
|
||||
def bad_request(message: str):
|
||||
return web.json_response({"error": message}, status=400)
|
||||
|
||||
def missing_field(field: str):
|
||||
return bad_request(f"{field} is required")
|
||||
|
||||
def not_found(message: str):
|
||||
return web.json_response({"error": message + " not found"}, status=404)
|
||||
|
||||
class ModelFileManager:
|
||||
def __init__(self) -> None:
|
||||
self.cache: dict[str, tuple[list[dict], dict[str, float], float]] = {}
|
||||
@@ -62,7 +73,7 @@ class ModelFileManager:
|
||||
folder = folders[0][path_index]
|
||||
full_filename = os.path.join(folder, filename)
|
||||
|
||||
previews = self.get_model_previews(full_filename)
|
||||
previews = get_model_previews(full_filename)
|
||||
default_preview = previews[0] if len(previews) > 0 else None
|
||||
if default_preview is None or (isinstance(default_preview, str) and not os.path.isfile(default_preview)):
|
||||
return web.Response(status=404)
|
||||
@@ -76,6 +87,183 @@ class ModelFileManager:
|
||||
except:
|
||||
return web.Response(status=404)
|
||||
|
||||
@routes.get("/v2/models")
|
||||
async def get_models(request):
|
||||
with create_session() as session:
|
||||
model_path = request.query.get("path", None)
|
||||
model_type = request.query.get("type", None)
|
||||
query = session.query(Model).options(joinedload(Model.tags))
|
||||
if model_path:
|
||||
query = query.filter(Model.path == model_path)
|
||||
if model_type:
|
||||
query = query.filter(Model.type == model_type)
|
||||
models = query.all()
|
||||
if model_path and model_type:
|
||||
if len(models) == 0:
|
||||
return not_found("Model")
|
||||
return web.json_response(models[0].to_dict(), dumps=dumps)
|
||||
|
||||
return web.json_response([model.to_dict() for model in models], dumps=dumps)
|
||||
|
||||
@routes.post("/v2/models")
|
||||
async def add_model(request):
|
||||
with create_session() as session:
|
||||
data = await request.json()
|
||||
model_type = data.get("type", None)
|
||||
model_path = data.get("path", None)
|
||||
|
||||
if not model_type:
|
||||
return missing_field("type")
|
||||
if not model_path:
|
||||
return missing_field("path")
|
||||
|
||||
tags = data.pop("tags", [])
|
||||
fields = Model.metadata.tables["model"].columns.keys()
|
||||
|
||||
# Validate keys are valid model fields
|
||||
for key in data.keys():
|
||||
if key not in fields:
|
||||
return bad_request(f"Invalid field: {key}")
|
||||
|
||||
# Validate file exists
|
||||
if not get_full_path(model_type, model_path):
|
||||
return not_found(f"File '{model_type}/{model_path}'")
|
||||
|
||||
model = Model()
|
||||
for field in fields:
|
||||
if field in data:
|
||||
setattr(model, field, data[field])
|
||||
|
||||
model.tags = session.query(Tag).filter(Tag.id.in_(tags)).all()
|
||||
for tag in tags:
|
||||
if tag not in [t.id for t in model.tags]:
|
||||
return not_found(f"Tag '{tag}'")
|
||||
|
||||
try:
|
||||
session.add(model)
|
||||
session.commit()
|
||||
except sqlalchemy.exc.IntegrityError as e:
|
||||
session.rollback()
|
||||
return bad_request(e.orig.args[0])
|
||||
|
||||
model_processor.run()
|
||||
|
||||
return web.json_response(model.to_dict(), dumps=dumps)
|
||||
|
||||
@routes.delete("/v2/models")
|
||||
async def delete_model(request):
|
||||
with create_session() as session:
|
||||
model_path = request.query.get("path", None)
|
||||
model_type = request.query.get("type", None)
|
||||
if not model_path:
|
||||
return missing_field("path")
|
||||
if not model_type:
|
||||
return missing_field("type")
|
||||
|
||||
full_path = get_full_path(model_type, model_path)
|
||||
if full_path:
|
||||
return bad_request("Model file exists, please delete the file before deleting the model record.")
|
||||
|
||||
model = session.query(Model).filter(Model.path == model_path, Model.type == model_type).first()
|
||||
if not model:
|
||||
return not_found("Model")
|
||||
session.delete(model)
|
||||
session.commit()
|
||||
return web.Response(status=204)
|
||||
|
||||
@routes.get("/v2/tags")
|
||||
async def get_tags(request):
|
||||
with create_session() as session:
|
||||
tags = session.query(Tag).all()
|
||||
return web.json_response(
|
||||
[{"id": tag.id, "name": tag.name} for tag in tags]
|
||||
)
|
||||
|
||||
@routes.post("/v2/tags")
|
||||
async def create_tag(request):
|
||||
with create_session() as session:
|
||||
data = await request.json()
|
||||
name = data.get("name", None)
|
||||
if not name:
|
||||
return missing_field("name")
|
||||
tag = Tag(name=name)
|
||||
session.add(tag)
|
||||
session.commit()
|
||||
return web.json_response({"id": tag.id, "name": tag.name})
|
||||
|
||||
@routes.delete("/v2/tags")
|
||||
async def delete_tag(request):
|
||||
with create_session() as session:
|
||||
tag_id = request.query.get("id", None)
|
||||
if not tag_id:
|
||||
return missing_field("id")
|
||||
tag = session.query(Tag).filter(Tag.id == tag_id).first()
|
||||
if not tag:
|
||||
return not_found("Tag")
|
||||
session.delete(tag)
|
||||
session.commit()
|
||||
return web.Response(status=204)
|
||||
|
||||
@routes.post("/v2/models/tags")
|
||||
async def add_model_tag(request):
|
||||
with create_session() as session:
|
||||
data = await request.json()
|
||||
tag_id = data.get("tag", None)
|
||||
model_path = data.get("path", None)
|
||||
model_type = data.get("type", None)
|
||||
|
||||
if tag_id is None:
|
||||
return missing_field("tag")
|
||||
if model_path is None:
|
||||
return missing_field("path")
|
||||
if model_type is None:
|
||||
return missing_field("type")
|
||||
|
||||
try:
|
||||
tag_id = int(tag_id)
|
||||
except ValueError:
|
||||
return bad_request("Invalid tag id")
|
||||
|
||||
tag = session.query(Tag).filter(Tag.id == tag_id).first()
|
||||
model = session.query(Model).filter(Model.path == model_path, Model.type == model_type).first()
|
||||
if not model:
|
||||
return not_found("Model")
|
||||
model.tags.append(tag)
|
||||
session.commit()
|
||||
return web.json_response(model.to_dict(), dumps=dumps)
|
||||
|
||||
@routes.delete("/v2/models/tags")
|
||||
async def delete_model_tag(request):
|
||||
with create_session() as session:
|
||||
tag_id = request.query.get("tag", None)
|
||||
model_path = request.query.get("path", None)
|
||||
model_type = request.query.get("type", None)
|
||||
|
||||
if tag_id is None:
|
||||
return missing_field("tag")
|
||||
if model_path is None:
|
||||
return missing_field("path")
|
||||
if model_type is None:
|
||||
return missing_field("type")
|
||||
|
||||
try:
|
||||
tag_id = int(tag_id)
|
||||
except ValueError:
|
||||
return bad_request("Invalid tag id")
|
||||
|
||||
model = session.query(Model).filter(Model.path == model_path, Model.type == model_type).first()
|
||||
if not model:
|
||||
return not_found("Model")
|
||||
model.tags = [tag for tag in model.tags if tag.id != tag_id]
|
||||
session.commit()
|
||||
return web.Response(status=204)
|
||||
|
||||
|
||||
|
||||
@routes.get("/v2/models/missing")
|
||||
async def get_missing_models(request):
|
||||
return web.json_response(model_processor.missing_models)
|
||||
|
||||
def get_model_file_list(self, folder_name: str):
|
||||
folder_name = map_legacy(folder_name)
|
||||
folders = folder_paths.folder_names_and_paths[folder_name]
|
||||
@@ -146,39 +334,5 @@ class ModelFileManager:
|
||||
|
||||
return [{"name": f, "pathIndex": pathIndex} for f in result], dirs, time.perf_counter()
|
||||
|
||||
def get_model_previews(self, filepath: str) -> list[str | BytesIO]:
|
||||
dirname = os.path.dirname(filepath)
|
||||
|
||||
if not os.path.exists(dirname):
|
||||
return []
|
||||
|
||||
basename = os.path.splitext(filepath)[0]
|
||||
match_files = glob.glob(f"{basename}.*", recursive=False)
|
||||
image_files = filter_files_content_types(match_files, "image")
|
||||
safetensors_file = next(filter(lambda x: x.endswith(".safetensors"), match_files), None)
|
||||
safetensors_metadata = {}
|
||||
|
||||
result: list[str | BytesIO] = []
|
||||
|
||||
for filename in image_files:
|
||||
_basename = os.path.splitext(filename)[0]
|
||||
if _basename == basename:
|
||||
result.append(filename)
|
||||
if _basename == f"{basename}.preview":
|
||||
result.append(filename)
|
||||
|
||||
if safetensors_file:
|
||||
safetensors_filepath = os.path.join(dirname, safetensors_file)
|
||||
header = comfy.utils.safetensors_header(safetensors_filepath, max_size=8*1024*1024)
|
||||
if header:
|
||||
safetensors_metadata = json.loads(header)
|
||||
safetensors_images = safetensors_metadata.get("__metadata__", {}).get("ssmd_cover_images", None)
|
||||
if safetensors_images:
|
||||
safetensors_images = json.loads(safetensors_images)
|
||||
for image in safetensors_images:
|
||||
result.append(BytesIO(base64.b64decode(image)))
|
||||
|
||||
return result
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.clear_cache()
|
||||
|
||||
Reference in New Issue
Block a user