Skip to content

Commit 1dc3da6

Browse files
authored
Add most basic Asset support for models (Comfy-Org#11315)
* Brought over minimal elements from PR 10045 to reproduce seed_assets and register_assets_system without adding anything to the DB or server routes yet, for now making everything sync (can introduce async once everything is cleaned up and brought over) * Added db script to insert assets stuff, cleaned up some code; assets (models) now get added/rescanned * Added support for 5 http endpoints for assets * Replaced Optional with | None in schemas_in.py and schemas_out.py * Remove two routes that will not be relevant yet in this PR: HEAD /api/assets/hash/<hash> and PUT /api/assets/<id>/preview * Remove some functions the two deleted endpoints were using * Don't show assets scan message upon calling /object_info endpoint * removed unsued import to satisfy ruff * Simplified hashing function tpye hint and _hash_file_obj * Satisfied ruff
1 parent 114fc73 commit 1dc3da6

16 files changed

Lines changed: 1847 additions & 9 deletions

File tree

alembic_db/versions/0001_assets.py

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
"""
2+
Initial assets schema
3+
Revision ID: 0001_assets
4+
Revises: None
5+
Create Date: 2025-12-10 00:00:00
6+
"""
7+
8+
from alembic import op
9+
import sqlalchemy as sa
10+
11+
revision = "0001_assets"
12+
down_revision = None
13+
branch_labels = None
14+
depends_on = None
15+
16+
17+
def upgrade() -> None:
18+
# ASSETS: content identity
19+
op.create_table(
20+
"assets",
21+
sa.Column("id", sa.String(length=36), primary_key=True),
22+
sa.Column("hash", sa.String(length=256), nullable=True),
23+
sa.Column("size_bytes", sa.BigInteger(), nullable=False, server_default="0"),
24+
sa.Column("mime_type", sa.String(length=255), nullable=True),
25+
sa.Column("created_at", sa.DateTime(timezone=False), nullable=False),
26+
sa.CheckConstraint("size_bytes >= 0", name="ck_assets_size_nonneg"),
27+
)
28+
op.create_index("uq_assets_hash", "assets", ["hash"], unique=True)
29+
op.create_index("ix_assets_mime_type", "assets", ["mime_type"])
30+
31+
# ASSETS_INFO: user-visible references
32+
op.create_table(
33+
"assets_info",
34+
sa.Column("id", sa.String(length=36), primary_key=True),
35+
sa.Column("owner_id", sa.String(length=128), nullable=False, server_default=""),
36+
sa.Column("name", sa.String(length=512), nullable=False),
37+
sa.Column("asset_id", sa.String(length=36), sa.ForeignKey("assets.id", ondelete="RESTRICT"), nullable=False),
38+
sa.Column("preview_id", sa.String(length=36), sa.ForeignKey("assets.id", ondelete="SET NULL"), nullable=True),
39+
sa.Column("user_metadata", sa.JSON(), nullable=True),
40+
sa.Column("created_at", sa.DateTime(timezone=False), nullable=False),
41+
sa.Column("updated_at", sa.DateTime(timezone=False), nullable=False),
42+
sa.Column("last_access_time", sa.DateTime(timezone=False), nullable=False),
43+
sa.UniqueConstraint("asset_id", "owner_id", "name", name="uq_assets_info_asset_owner_name"),
44+
)
45+
op.create_index("ix_assets_info_owner_id", "assets_info", ["owner_id"])
46+
op.create_index("ix_assets_info_asset_id", "assets_info", ["asset_id"])
47+
op.create_index("ix_assets_info_name", "assets_info", ["name"])
48+
op.create_index("ix_assets_info_created_at", "assets_info", ["created_at"])
49+
op.create_index("ix_assets_info_last_access_time", "assets_info", ["last_access_time"])
50+
op.create_index("ix_assets_info_owner_name", "assets_info", ["owner_id", "name"])
51+
52+
# TAGS: normalized tag vocabulary
53+
op.create_table(
54+
"tags",
55+
sa.Column("name", sa.String(length=512), primary_key=True),
56+
sa.Column("tag_type", sa.String(length=32), nullable=False, server_default="user"),
57+
sa.CheckConstraint("name = lower(name)", name="ck_tags_lowercase"),
58+
)
59+
op.create_index("ix_tags_tag_type", "tags", ["tag_type"])
60+
61+
# ASSET_INFO_TAGS: many-to-many for tags on AssetInfo
62+
op.create_table(
63+
"asset_info_tags",
64+
sa.Column("asset_info_id", sa.String(length=36), sa.ForeignKey("assets_info.id", ondelete="CASCADE"), nullable=False),
65+
sa.Column("tag_name", sa.String(length=512), sa.ForeignKey("tags.name", ondelete="RESTRICT"), nullable=False),
66+
sa.Column("origin", sa.String(length=32), nullable=False, server_default="manual"),
67+
sa.Column("added_at", sa.DateTime(timezone=False), nullable=False),
68+
sa.PrimaryKeyConstraint("asset_info_id", "tag_name", name="pk_asset_info_tags"),
69+
)
70+
op.create_index("ix_asset_info_tags_tag_name", "asset_info_tags", ["tag_name"])
71+
op.create_index("ix_asset_info_tags_asset_info_id", "asset_info_tags", ["asset_info_id"])
72+
73+
# ASSET_CACHE_STATE: N:1 local cache rows per Asset
74+
op.create_table(
75+
"asset_cache_state",
76+
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
77+
sa.Column("asset_id", sa.String(length=36), sa.ForeignKey("assets.id", ondelete="CASCADE"), nullable=False),
78+
sa.Column("file_path", sa.Text(), nullable=False), # absolute local path to cached file
79+
sa.Column("mtime_ns", sa.BigInteger(), nullable=True),
80+
sa.Column("needs_verify", sa.Boolean(), nullable=False, server_default=sa.text("false")),
81+
sa.CheckConstraint("(mtime_ns IS NULL) OR (mtime_ns >= 0)", name="ck_acs_mtime_nonneg"),
82+
sa.UniqueConstraint("file_path", name="uq_asset_cache_state_file_path"),
83+
)
84+
op.create_index("ix_asset_cache_state_file_path", "asset_cache_state", ["file_path"])
85+
op.create_index("ix_asset_cache_state_asset_id", "asset_cache_state", ["asset_id"])
86+
87+
# ASSET_INFO_META: typed KV projection of user_metadata for filtering/sorting
88+
op.create_table(
89+
"asset_info_meta",
90+
sa.Column("asset_info_id", sa.String(length=36), sa.ForeignKey("assets_info.id", ondelete="CASCADE"), nullable=False),
91+
sa.Column("key", sa.String(length=256), nullable=False),
92+
sa.Column("ordinal", sa.Integer(), nullable=False, server_default="0"),
93+
sa.Column("val_str", sa.String(length=2048), nullable=True),
94+
sa.Column("val_num", sa.Numeric(38, 10), nullable=True),
95+
sa.Column("val_bool", sa.Boolean(), nullable=True),
96+
sa.Column("val_json", sa.JSON(), nullable=True),
97+
sa.PrimaryKeyConstraint("asset_info_id", "key", "ordinal", name="pk_asset_info_meta"),
98+
)
99+
op.create_index("ix_asset_info_meta_key", "asset_info_meta", ["key"])
100+
op.create_index("ix_asset_info_meta_key_val_str", "asset_info_meta", ["key", "val_str"])
101+
op.create_index("ix_asset_info_meta_key_val_num", "asset_info_meta", ["key", "val_num"])
102+
op.create_index("ix_asset_info_meta_key_val_bool", "asset_info_meta", ["key", "val_bool"])
103+
104+
# Tags vocabulary
105+
tags_table = sa.table(
106+
"tags",
107+
sa.column("name", sa.String(length=512)),
108+
sa.column("tag_type", sa.String()),
109+
)
110+
op.bulk_insert(
111+
tags_table,
112+
[
113+
{"name": "models", "tag_type": "system"},
114+
{"name": "input", "tag_type": "system"},
115+
{"name": "output", "tag_type": "system"},
116+
117+
{"name": "configs", "tag_type": "system"},
118+
{"name": "checkpoints", "tag_type": "system"},
119+
{"name": "loras", "tag_type": "system"},
120+
{"name": "vae", "tag_type": "system"},
121+
{"name": "text_encoders", "tag_type": "system"},
122+
{"name": "diffusion_models", "tag_type": "system"},
123+
{"name": "clip_vision", "tag_type": "system"},
124+
{"name": "style_models", "tag_type": "system"},
125+
{"name": "embeddings", "tag_type": "system"},
126+
{"name": "diffusers", "tag_type": "system"},
127+
{"name": "vae_approx", "tag_type": "system"},
128+
{"name": "controlnet", "tag_type": "system"},
129+
{"name": "gligen", "tag_type": "system"},
130+
{"name": "upscale_models", "tag_type": "system"},
131+
{"name": "hypernetworks", "tag_type": "system"},
132+
{"name": "photomaker", "tag_type": "system"},
133+
{"name": "classifiers", "tag_type": "system"},
134+
135+
{"name": "encoder", "tag_type": "system"},
136+
{"name": "decoder", "tag_type": "system"},
137+
138+
{"name": "missing", "tag_type": "system"},
139+
{"name": "rescan", "tag_type": "system"},
140+
],
141+
)
142+
143+
144+
def downgrade() -> None:
145+
op.drop_index("ix_asset_info_meta_key_val_bool", table_name="asset_info_meta")
146+
op.drop_index("ix_asset_info_meta_key_val_num", table_name="asset_info_meta")
147+
op.drop_index("ix_asset_info_meta_key_val_str", table_name="asset_info_meta")
148+
op.drop_index("ix_asset_info_meta_key", table_name="asset_info_meta")
149+
op.drop_table("asset_info_meta")
150+
151+
op.drop_index("ix_asset_cache_state_asset_id", table_name="asset_cache_state")
152+
op.drop_index("ix_asset_cache_state_file_path", table_name="asset_cache_state")
153+
op.drop_constraint("uq_asset_cache_state_file_path", table_name="asset_cache_state")
154+
op.drop_table("asset_cache_state")
155+
156+
op.drop_index("ix_asset_info_tags_asset_info_id", table_name="asset_info_tags")
157+
op.drop_index("ix_asset_info_tags_tag_name", table_name="asset_info_tags")
158+
op.drop_table("asset_info_tags")
159+
160+
op.drop_index("ix_tags_tag_type", table_name="tags")
161+
op.drop_table("tags")
162+
163+
op.drop_constraint("uq_assets_info_asset_owner_name", table_name="assets_info")
164+
op.drop_index("ix_assets_info_owner_name", table_name="assets_info")
165+
op.drop_index("ix_assets_info_last_access_time", table_name="assets_info")
166+
op.drop_index("ix_assets_info_created_at", table_name="assets_info")
167+
op.drop_index("ix_assets_info_name", table_name="assets_info")
168+
op.drop_index("ix_assets_info_asset_id", table_name="assets_info")
169+
op.drop_index("ix_assets_info_owner_id", table_name="assets_info")
170+
op.drop_table("assets_info")
171+
172+
op.drop_index("uq_assets_hash", table_name="assets")
173+
op.drop_index("ix_assets_mime_type", table_name="assets")
174+
op.drop_table("assets")

app/assets/api/routes.py

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import logging
2+
import uuid
3+
from aiohttp import web
4+
5+
from pydantic import ValidationError
6+
7+
import app.assets.manager as manager
8+
from app import user_manager
9+
from app.assets.api import schemas_in
10+
from app.assets.helpers import get_query_dict
11+
12+
ROUTES = web.RouteTableDef()
13+
USER_MANAGER: user_manager.UserManager | None = None
14+
15+
# UUID regex (canonical hyphenated form, case-insensitive)
16+
UUID_RE = r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"
17+
18+
def register_assets_system(app: web.Application, user_manager_instance: user_manager.UserManager) -> None:
19+
global USER_MANAGER
20+
USER_MANAGER = user_manager_instance
21+
app.add_routes(ROUTES)
22+
23+
def _error_response(status: int, code: str, message: str, details: dict | None = None) -> web.Response:
24+
return web.json_response({"error": {"code": code, "message": message, "details": details or {}}}, status=status)
25+
26+
27+
def _validation_error_response(code: str, ve: ValidationError) -> web.Response:
28+
return _error_response(400, code, "Validation failed.", {"errors": ve.json()})
29+
30+
31+
@ROUTES.get("/api/assets")
32+
async def list_assets(request: web.Request) -> web.Response:
33+
"""
34+
GET request to list assets.
35+
"""
36+
query_dict = get_query_dict(request)
37+
try:
38+
q = schemas_in.ListAssetsQuery.model_validate(query_dict)
39+
except ValidationError as ve:
40+
return _validation_error_response("INVALID_QUERY", ve)
41+
42+
payload = manager.list_assets(
43+
include_tags=q.include_tags,
44+
exclude_tags=q.exclude_tags,
45+
name_contains=q.name_contains,
46+
metadata_filter=q.metadata_filter,
47+
limit=q.limit,
48+
offset=q.offset,
49+
sort=q.sort,
50+
order=q.order,
51+
owner_id=USER_MANAGER.get_request_user_id(request),
52+
)
53+
return web.json_response(payload.model_dump(mode="json"))
54+
55+
56+
@ROUTES.get(f"/api/assets/{{id:{UUID_RE}}}")
57+
async def get_asset(request: web.Request) -> web.Response:
58+
"""
59+
GET request to get an asset's info as JSON.
60+
"""
61+
asset_info_id = str(uuid.UUID(request.match_info["id"]))
62+
try:
63+
result = manager.get_asset(
64+
asset_info_id=asset_info_id,
65+
owner_id=USER_MANAGER.get_request_user_id(request),
66+
)
67+
except ValueError as e:
68+
return _error_response(404, "ASSET_NOT_FOUND", str(e), {"id": asset_info_id})
69+
except Exception:
70+
logging.exception(
71+
"get_asset failed for asset_info_id=%s, owner_id=%s",
72+
asset_info_id,
73+
USER_MANAGER.get_request_user_id(request),
74+
)
75+
return _error_response(500, "INTERNAL", "Unexpected server error.")
76+
return web.json_response(result.model_dump(mode="json"), status=200)
77+
78+
79+
@ROUTES.get("/api/tags")
80+
async def get_tags(request: web.Request) -> web.Response:
81+
"""
82+
GET request to list all tags based on query parameters.
83+
"""
84+
query_map = dict(request.rel_url.query)
85+
86+
try:
87+
query = schemas_in.TagsListQuery.model_validate(query_map)
88+
except ValidationError as e:
89+
return web.json_response(
90+
{"error": {"code": "INVALID_QUERY", "message": "Invalid query parameters", "details": e.errors()}},
91+
status=400,
92+
)
93+
94+
result = manager.list_tags(
95+
prefix=query.prefix,
96+
limit=query.limit,
97+
offset=query.offset,
98+
order=query.order,
99+
include_zero=query.include_zero,
100+
owner_id=USER_MANAGER.get_request_user_id(request),
101+
)
102+
return web.json_response(result.model_dump(mode="json"))

app/assets/api/schemas_in.py

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
import json
2+
import uuid
3+
from typing import Any, Literal
4+
5+
from pydantic import (
6+
BaseModel,
7+
ConfigDict,
8+
Field,
9+
conint,
10+
field_validator,
11+
)
12+
13+
14+
class ListAssetsQuery(BaseModel):
15+
include_tags: list[str] = Field(default_factory=list)
16+
exclude_tags: list[str] = Field(default_factory=list)
17+
name_contains: str | None = None
18+
19+
# Accept either a JSON string (query param) or a dict
20+
metadata_filter: dict[str, Any] | None = None
21+
22+
limit: conint(ge=1, le=500) = 20
23+
offset: conint(ge=0) = 0
24+
25+
sort: Literal["name", "created_at", "updated_at", "size", "last_access_time"] = "created_at"
26+
order: Literal["asc", "desc"] = "desc"
27+
28+
@field_validator("include_tags", "exclude_tags", mode="before")
29+
@classmethod
30+
def _split_csv_tags(cls, v):
31+
# Accept "a,b,c" or ["a","b"] (we are liberal in what we accept)
32+
if v is None:
33+
return []
34+
if isinstance(v, str):
35+
return [t.strip() for t in v.split(",") if t.strip()]
36+
if isinstance(v, list):
37+
out: list[str] = []
38+
for item in v:
39+
if isinstance(item, str):
40+
out.extend([t.strip() for t in item.split(",") if t.strip()])
41+
return out
42+
return v
43+
44+
@field_validator("metadata_filter", mode="before")
45+
@classmethod
46+
def _parse_metadata_json(cls, v):
47+
if v is None or isinstance(v, dict):
48+
return v
49+
if isinstance(v, str) and v.strip():
50+
try:
51+
parsed = json.loads(v)
52+
except Exception as e:
53+
raise ValueError(f"metadata_filter must be JSON: {e}") from e
54+
if not isinstance(parsed, dict):
55+
raise ValueError("metadata_filter must be a JSON object")
56+
return parsed
57+
return None
58+
59+
60+
class TagsListQuery(BaseModel):
61+
model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)
62+
63+
prefix: str | None = Field(None, min_length=1, max_length=256)
64+
limit: int = Field(100, ge=1, le=1000)
65+
offset: int = Field(0, ge=0, le=10_000_000)
66+
order: Literal["count_desc", "name_asc"] = "count_desc"
67+
include_zero: bool = True
68+
69+
@field_validator("prefix")
70+
@classmethod
71+
def normalize_prefix(cls, v: str | None) -> str | None:
72+
if v is None:
73+
return v
74+
v = v.strip()
75+
return v.lower() or None
76+
77+
78+
class SetPreviewBody(BaseModel):
79+
"""Set or clear the preview for an AssetInfo. Provide an Asset.id or null."""
80+
preview_id: str | None = None
81+
82+
@field_validator("preview_id", mode="before")
83+
@classmethod
84+
def _norm_uuid(cls, v):
85+
if v is None:
86+
return None
87+
s = str(v).strip()
88+
if not s:
89+
return None
90+
try:
91+
uuid.UUID(s)
92+
except Exception:
93+
raise ValueError("preview_id must be a UUID")
94+
return s

0 commit comments

Comments
 (0)