InvokeAI/invokeai/app/services/db.ipynb
psychedelicious 9c89d3452c feat(nodes): add high-level images service
feat(nodes): add ResultsServiceABC & SqliteResultsService

**Doesn't actually work bc of circular imports. Can't even test it.**

- add a base class for ResultsService and SQLite implementation
- use `graph_execution_manager` `on_changed` callback to keep `results` table in sync

fix(nodes): fix results service bugs

chore(ui): regen api

fix(ui): fix type guards

feat(nodes): add `result_type` to results table, fix types

fix(nodes): do not shadow `list` builtin

feat(nodes): add results router

It doesn't work due to circular imports still

fix(nodes): Result class should use outputs classes, not fields

feat(ui): crude results router

fix(ui): send to canvas in currentimagebuttons not working

feat(nodes): add core metadata builder

feat(nodes): add design doc

feat(nodes): wip latents db stuff

feat(nodes): images_db_service and resources router

feat(nodes): wip images db & router

feat(nodes): update image related names

feat(nodes): update urlservice

feat(nodes): add high-level images service
2023-05-24 11:30:47 -04:00

579 lines
19 KiB
Plaintext

{
"cells": [
{
"cell_type": "code",
"execution_count": 40,
"metadata": {},
"outputs": [],
"source": [
"from abc import ABC, abstractmethod\n",
"from enum import Enum\n",
"import enum\n",
"import sqlite3\n",
"import threading\n",
"from typing import Optional, Type, TypeVar, Union\n",
"from PIL.Image import Image as PILImage\n",
"from pydantic import BaseModel, Field\n",
"from torch import Tensor"
]
},
{
"cell_type": "code",
"execution_count": 41,
"metadata": {},
"outputs": [],
"source": [
"\n",
"class ResourceOrigin(str, Enum):\n",
" \"\"\"The origin of a resource (eg image or tensor).\"\"\"\n",
"\n",
" RESULTS = \"results\"\n",
" UPLOADS = \"uploads\"\n",
" INTERMEDIATES = \"intermediates\"\n",
"\n",
"\n",
"class ImageKind(str, Enum):\n",
" \"\"\"The kind of an image. Use ImageKind.OTHER for non-default kinds.\"\"\"\n",
"\n",
" IMAGE = \"image\"\n",
" CONTROL_IMAGE = \"control_image\"\n",
" OTHER = \"other\"\n",
"\n",
"\n",
"class TensorKind(str, Enum):\n",
" \"\"\"The kind of a tensor. Use TensorKind.OTHER for non-default kinds.\"\"\"\n",
"\n",
" IMAGE_LATENTS = \"image_latents\"\n",
" CONDITIONING = \"conditioning\"\n",
" OTHER = \"other\"\n"
]
},
{
"cell_type": "code",
"execution_count": 42,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def create_sql_values_string_from_string_enum(enum: Type[Enum]):\n",
" \"\"\"\n",
" Creates a string of the form \"('value1'), ('value2'), ..., ('valueN')\" from a StrEnum.\n",
" \"\"\"\n",
"\n",
" delimiter = \", \"\n",
" values = [f\"('{e.value}')\" for e in enum]\n",
" return delimiter.join(values)\n",
"\n",
"\n",
"def create_sql_table_from_enum(\n",
" enum: Type[Enum],\n",
" table_name: str,\n",
" primary_key_name: str,\n",
" conn: sqlite3.Connection,\n",
" cursor: sqlite3.Cursor,\n",
" lock: threading.Lock,\n",
"):\n",
" \"\"\"\n",
" Creates and populates a table to be used as a functional enum.\n",
" \"\"\"\n",
"\n",
" try:\n",
" lock.acquire()\n",
"\n",
" values_string = create_sql_values_string_from_string_enum(enum)\n",
"\n",
" cursor.execute(\n",
" f\"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS {table_name} (\n",
" {primary_key_name} TEXT PRIMARY KEY\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" f\"\"\"--sql\n",
" INSERT OR IGNORE INTO {table_name} ({primary_key_name}) VALUES {values_string};\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n",
"\n",
"\n",
"\"\"\"\n",
"`resource_origins` functions as an enum for the ResourceOrigin model.\n",
"\"\"\"\n",
"\n",
"\n",
"# def create_resource_origins_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
"# create_sql_table_from_enum(\n",
"# enum=ResourceOrigin,\n",
"# table_name=\"resource_origins\",\n",
"# primary_key_name=\"origin_name\",\n",
"# conn=conn,\n",
"# cursor=cursor,\n",
"# lock=lock,\n",
"# )\n",
"\n",
"\n",
"\"\"\"\n",
"`image_kinds` functions as an enum for the ImageType model.\n",
"\"\"\"\n",
"\n",
"\n",
"# def create_image_kinds_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" # create_sql_table_from_enum(\n",
" # enum=ImageKind,\n",
" # table_name=\"image_kinds\",\n",
" # primary_key_name=\"kind_name\",\n",
" # conn=conn,\n",
" # cursor=cursor,\n",
" # lock=lock,\n",
" # )\n",
"\n",
"\n",
"\"\"\"\n",
"`tensor_kinds` functions as an enum for the TensorType model.\n",
"\"\"\"\n",
"\n",
"\n",
"# def create_tensor_kinds_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" # create_sql_table_from_enum(\n",
" # enum=TensorKind,\n",
" # table_name=\"tensor_kinds\",\n",
" # primary_key_name=\"kind_name\",\n",
" # conn=conn,\n",
" # cursor=cursor,\n",
" # lock=lock,\n",
" # )\n",
"\n",
"\n",
"\"\"\"\n",
"`images` stores all images, regardless of type\n",
"\"\"\"\n",
"\n",
"\n",
"def create_images_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" try:\n",
" lock.acquire()\n",
"\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS images (\n",
" id TEXT PRIMARY KEY,\n",
" origin TEXT,\n",
" image_kind TEXT,\n",
" created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\n",
" FOREIGN KEY(origin) REFERENCES resource_origins(origin_name),\n",
" FOREIGN KEY(image_kind) REFERENCES image_kinds(kind_name)\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE UNIQUE INDEX IF NOT EXISTS idx_images_id ON images(id);\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE INDEX IF NOT EXISTS idx_images_origin ON images(origin);\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE INDEX IF NOT EXISTS idx_images_image_kind ON images(image_kind);\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n",
"\n",
"\n",
"\"\"\"\n",
"`images_results` stores additional data specific to `results` images.\n",
"\"\"\"\n",
"\n",
"\n",
"def create_images_results_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" try:\n",
" lock.acquire()\n",
"\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS images_results (\n",
" images_id TEXT PRIMARY KEY,\n",
" session_id TEXT NOT NULL,\n",
" node_id TEXT NOT NULL,\n",
" FOREIGN KEY(images_id) REFERENCES images(id) ON DELETE CASCADE\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE UNIQUE INDEX IF NOT EXISTS idx_images_results_images_id ON images_results(images_id);\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n",
"\n",
"\n",
"\"\"\"\n",
"`images_intermediates` stores additional data specific to `intermediates` images\n",
"\"\"\"\n",
"\n",
"\n",
"def create_images_intermediates_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" try:\n",
" lock.acquire()\n",
"\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS images_intermediates (\n",
" images_id TEXT PRIMARY KEY,\n",
" session_id TEXT NOT NULL,\n",
" node_id TEXT NOT NULL,\n",
" FOREIGN KEY(images_id) REFERENCES images(id) ON DELETE CASCADE\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE UNIQUE INDEX IF NOT EXISTS idx_images_intermediates_images_id ON images_intermediates(images_id);\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n",
"\n",
"\n",
"\"\"\"\n",
"`images_metadata` stores basic metadata for any image type\n",
"\"\"\"\n",
"\n",
"\n",
"def create_images_metadata_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" try:\n",
" lock.acquire()\n",
"\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS images_metadata (\n",
" images_id TEXT PRIMARY KEY,\n",
" metadata TEXT,\n",
" FOREIGN KEY(images_id) REFERENCES images(id) ON DELETE CASCADE\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE UNIQUE INDEX IF NOT EXISTS idx_images_metadata_images_id ON images_metadata(images_id);\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n",
"\n",
"\n",
"# `tensors` table: stores references to tensor\n",
"\n",
"\n",
"def create_tensors_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" try:\n",
" lock.acquire()\n",
"\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS tensors (\n",
" id TEXT PRIMARY KEY,\n",
" origin TEXT,\n",
" tensor_kind TEXT,\n",
" created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\n",
" FOREIGN KEY(origin) REFERENCES resource_origins(origin_name),\n",
" FOREIGN KEY(tensor_kind) REFERENCES tensor_kinds(kind_name)\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE UNIQUE INDEX IF NOT EXISTS idx_tensors_id ON tensors(id);\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE INDEX IF NOT EXISTS idx_tensors_origin ON tensors(origin);\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE INDEX IF NOT EXISTS idx_tensors_tensor_kind ON tensors(tensor_kind);\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n",
"\n",
"\n",
"# `tensors_results` stores additional data specific to `result` tensor\n",
"\n",
"\n",
"def create_tensors_results_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" try:\n",
" lock.acquire()\n",
"\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS tensors_results (\n",
" tensors_id TEXT PRIMARY KEY,\n",
" session_id TEXT NOT NULL,\n",
" node_id TEXT NOT NULL,\n",
" FOREIGN KEY(tensors_id) REFERENCES tensors(id) ON DELETE CASCADE\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE UNIQUE INDEX IF NOT EXISTS idx_tensors_results_tensors_id ON tensors_results(tensors_id);\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n",
"\n",
"\n",
"# `tensors_intermediates` stores additional data specific to `intermediate` tensor\n",
"\n",
"\n",
"def create_tensors_intermediates_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" try:\n",
" lock.acquire()\n",
"\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS tensors_intermediates (\n",
" tensors_id TEXT PRIMARY KEY,\n",
" session_id TEXT NOT NULL,\n",
" node_id TEXT NOT NULL,\n",
" FOREIGN KEY(tensors_id) REFERENCES tensors(id) ON DELETE CASCADE\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE UNIQUE INDEX IF NOT EXISTS idx_tensors_intermediates_tensors_id ON tensors_intermediates(tensors_id);\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n",
"\n",
"\n",
"# `tensors_metadata` table: stores generated/transformed metadata for tensor\n",
"\n",
"\n",
"def create_tensors_metadata_table(conn: sqlite3.Connection, cursor: sqlite3.Cursor, lock: threading.Lock):\n",
" try:\n",
" lock.acquire()\n",
"\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE TABLE IF NOT EXISTS tensors_metadata (\n",
" tensors_id TEXT PRIMARY KEY,\n",
" metadata TEXT,\n",
" FOREIGN KEY(tensors_id) REFERENCES tensors(id) ON DELETE CASCADE\n",
" );\n",
" \"\"\"\n",
" )\n",
" cursor.execute(\n",
" \"\"\"--sql\n",
" CREATE UNIQUE INDEX IF NOT EXISTS idx_tensors_metadata_tensors_id ON tensors_metadata(tensors_id);\n",
" \"\"\"\n",
" )\n",
" conn.commit()\n",
" finally:\n",
" lock.release()\n"
]
},
{
"cell_type": "code",
"execution_count": 43,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"db_path = '/home/bat/Documents/Code/outputs/test.db'\n",
"if (os.path.exists(db_path)):\n",
" os.remove(db_path)\n",
"\n",
"conn = sqlite3.connect(\n",
" db_path, check_same_thread=False\n",
")\n",
"cursor = conn.cursor()\n",
"lock = threading.Lock()"
]
},
{
"cell_type": "code",
"execution_count": 44,
"metadata": {},
"outputs": [],
"source": [
"create_sql_table_from_enum(\n",
" enum=ResourceOrigin,\n",
" table_name=\"resource_origins\",\n",
" primary_key_name=\"origin_name\",\n",
" conn=conn,\n",
" cursor=cursor,\n",
" lock=lock,\n",
")\n",
"\n",
"create_sql_table_from_enum(\n",
" enum=ImageKind,\n",
" table_name=\"image_kinds\",\n",
" primary_key_name=\"kind_name\",\n",
" conn=conn,\n",
" cursor=cursor,\n",
" lock=lock,\n",
")\n",
"\n",
"create_sql_table_from_enum(\n",
" enum=TensorKind,\n",
" table_name=\"tensor_kinds\",\n",
" primary_key_name=\"kind_name\",\n",
" conn=conn,\n",
" cursor=cursor,\n",
" lock=lock,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 45,
"metadata": {},
"outputs": [],
"source": [
"create_images_table(conn, cursor, lock)\n",
"create_images_results_table(conn, cursor, lock)\n",
"create_images_intermediates_table(conn, cursor, lock)\n",
"create_images_metadata_table(conn, cursor, lock)"
]
},
{
"cell_type": "code",
"execution_count": 46,
"metadata": {},
"outputs": [],
"source": [
"create_tensors_table(conn, cursor, lock)\n",
"create_tensors_results_table(conn, cursor, lock)\n",
"create_tensors_intermediates_table(conn, cursor, lock)\n",
"create_tensors_metadata_table(conn, cursor, lock)"
]
},
{
"cell_type": "code",
"execution_count": 59,
"metadata": {},
"outputs": [],
"source": [
"\n",
"from pydantic import StrictStr\n",
"\n",
"\n",
"class GeneratedImageOrLatentsMetadata(BaseModel):\n",
" \"\"\"Core generation metadata for an image/tensor generated in InvokeAI.\n",
"\n",
" Generated by traversing the execution graph, collecting the parameters of the nearest ancestors of a given node.\n",
"\n",
" Full metadata may be accessed by querying for the session in the `graph_executions` table.\n",
" \"\"\"\n",
"\n",
" positive_conditioning: Optional[StrictStr] = Field(\n",
" default=None, description=\"The positive conditioning.\"\n",
" )\n",
" negative_conditioning: Optional[str] = Field(\n",
" default=None, description=\"The negative conditioning.\"\n",
" )\n",
" width: Optional[int] = Field(\n",
" default=None, description=\"Width of the image/tensor in pixels.\"\n",
" )\n",
" height: Optional[int] = Field(\n",
" default=None, description=\"Height of the image/tensor in pixels.\"\n",
" )\n",
" seed: Optional[int] = Field(\n",
" default=None, description=\"The seed used for noise generation.\"\n",
" )\n",
" cfg_scale: Optional[float] = Field(\n",
" default=None, description=\"The classifier-free guidance scale.\"\n",
" )\n",
" steps: Optional[int] = Field(\n",
" default=None, description=\"The number of steps used for inference.\"\n",
" )\n",
" scheduler: Optional[str] = Field(\n",
" default=None, description=\"The scheduler used for inference.\"\n",
" )\n",
" model: Optional[str] = Field(\n",
" default=None, description=\"The model used for inference.\"\n",
" )\n",
" strength: Optional[float] = Field(\n",
" default=None,\n",
" description=\"The strength used for image-to-image/tensor-to-tensor.\",\n",
" )\n",
" image: Optional[str] = Field(\n",
" default=None, description=\"The ID of the initial image.\"\n",
" )\n",
" tensor: Optional[str] = Field(\n",
" default=None, description=\"The ID of the initial tensor.\"\n",
" )\n",
" # Pending model refactor:\n",
" # vae: Optional[str] = Field(default=None,description=\"The VAE used for decoding.\")\n",
" # unet: Optional[str] = Field(default=None,description=\"The UNet used dor inference.\")\n",
" # clip: Optional[str] = Field(default=None,description=\"The CLIP Encoder used for conditioning.\")\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 61,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"GeneratedImageOrLatentsMetadata(positive_conditioning='123', negative_conditioning=None, width=None, height=None, seed=None, cfg_scale=None, steps=None, scheduler=None, model=None, strength=None, image=None, tensor=None)"
]
},
"execution_count": 61,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"GeneratedImageOrLatentsMetadata(positive_conditioning='123')"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.6"
},
"orig_nbformat": 4
},
"nbformat": 4,
"nbformat_minor": 2
}