Skip to content

Add Dataset Library backend: curate, upload, train custom AI art models #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
guillaumelauzier merged 1 commit into from
Aug 21, 2026
+1,978 −1
Merged

Add Dataset Library backend: curate, upload, train custom AI art models #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions workers/api/.dev.vars.example
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,12 @@ FAL_KEY=""
# fal.ai key. Production must NEVER set RENDER_MOCK.
RENDER_MOCK="1"

# --- Dataset Library / training jobs ---------------------------------------
# OPT-IN ONLY, same convention as RENDER_MOCK. A training job completes
# synchronously in-request with canned weights instead of submitting to
# fal.ai's queue API. Production must NEVER set TRAINING_MOCK.
TRAINING_MOCK="1"

# --- Task #20: observability (all optional; unset = the surface no-ops) ----
# ADMIN_HANDLES gates /v1/internal/*; empty means those routes 403 with
# `admin_unconfigured`. Set it to your own handle to use /v1/internal/stats
Expand Down
114 changes: 114 additions & 0 deletions workers/api/migrations/0020_dataset_library.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
-- Dataset Library: a creator's own curated images/video, and the
-- training jobs that turn a dataset + base model into a published
-- render-model version (see migrations/0019_custom_model_provider.sql).
--
-- Lifecycle: a dataset is uploaded to and curated privately, then a
-- training job consumes it to produce an immutable render_model_versions
-- row (weights_ref pointing at the trained checkpoint/LoRA). The dataset
-- itself is never deleted or mutated by training — the same item set can
-- back many training jobs across many base models (this is the "one
-- dataset, many models" comparison feature).

CREATE TABLE IF NOT EXISTS datasets (
id INTEGER PRIMARY KEY AUTOINCREMENT,
owner_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
slug TEXT NOT NULL,
title TEXT NOT NULL,
-- The creator's curatorial statement — what this visual world is and
-- why it was collected. Becomes the public dataset_note on any model
-- version trained from this set (see publishVersionHandler).
description TEXT,
-- own | licensed | public_domain — a required, plain-language
-- self-declaration, not a verification step. See handler comment for
-- why there's no CHECK constraint enforcing the exact string set here.
rights_declaration TEXT NOT NULL,
-- private | public. Private by default — a dataset is source material,
-- not a publication. Publishing a dataset (making its item grid
-- visible) is a distinct, deliberate act from publishing a model
-- trained on it.
visibility TEXT NOT NULL DEFAULT 'private',
-- Denormalised so the library grid renders without a per-card COUNT
-- query. Kept in sync by insertDatasetItem/deleteDatasetItem in the
-- same statement, not via a trigger, so it stays visible in git blame.
item_count INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
UNIQUE (owner_id, slug)
);

CREATE INDEX IF NOT EXISTS idx_datasets_owner
ON datasets(owner_id, updated_at DESC);

CREATE TABLE IF NOT EXISTS dataset_items (
id INTEGER PRIMARY KEY AUTOINCREMENT,
dataset_id INTEGER NOT NULL REFERENCES datasets(id) ON DELETE CASCADE,
r2_key TEXT NOT NULL UNIQUE,
-- image | video
kind TEXT NOT NULL,
caption TEXT,
-- sha256 of the stored bytes. UNIQUE per dataset is the platform's
-- dedup mechanism — an exact re-upload (the common case: the same
-- file dragged in twice, or re-imported from a URL list that
-- overlaps a previous import) is rejected at the DB layer rather
-- than silently duplicated. Near-duplicate (perceptually similar,
-- not byte-identical) detection is a deliberate v2 scope cut — it
-- needs a perceptual hash, not a content hash — flagged in the
-- handler, not silently absent.
content_hash TEXT NOT NULL,
byte_size INTEGER NOT NULL,
width INTEGER,
height INTEGER,
duration_seconds REAL,
-- ready | flagged. There's no async "processing" stage yet — an item
-- is either accepted (ready) or rejected with a reason (flagged) at
-- upload/import time, since this Worker can't run a background
-- thumbnailing/transcoding pipeline. `flag_reason` is the specific
-- cause (unsupported codec, corrupt file, too small) referenced by
-- the design's "flagged items" state.
status TEXT NOT NULL DEFAULT 'ready',
flag_reason TEXT,
created_at INTEGER NOT NULL,
UNIQUE (dataset_id, content_hash)
);

CREATE INDEX IF NOT EXISTS idx_dataset_items_dataset
ON dataset_items(dataset_id, id DESC);

CREATE TABLE IF NOT EXISTS training_jobs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
dataset_id INTEGER NOT NULL REFERENCES datasets(id) ON DELETE CASCADE,
base_model TEXT NOT NULL,
-- lora | dreambooth | full_finetune (see TRAINING_METHODS in
-- src/db/render.ts — reused here, not redefined).
training_method TEXT NOT NULL,
price_tokens INTEGER NOT NULL,
-- What the resulting model version should charge per render, chosen
-- by the creator at train time so it's ready to publish the moment
-- training succeeds rather than defaulting to free.
render_price_tokens INTEGER NOT NULL DEFAULT 25,
-- queued | training | succeeded | failed
status TEXT NOT NULL DEFAULT 'queued',
-- fal.ai's queue request_id once submitted. Null while queued.
provider_job_id TEXT,
-- Populated on submit: the render_models shell this job will publish
-- a version onto. Created eagerly (status='private') so a job in
-- progress has somewhere to point before it succeeds.
model_id INTEGER REFERENCES render_models(id),
-- Populated on success: the immutable version the job produced.
model_version_id INTEGER REFERENCES render_model_versions(id),
error_code TEXT,
idempotency_key TEXT NOT NULL UNIQUE,
created_at INTEGER NOT NULL,
started_at INTEGER,
finished_at INTEGER
);

CREATE INDEX IF NOT EXISTS idx_training_jobs_user
ON training_jobs(user_id, created_at DESC);
CREATE INDEX IF NOT EXISTS idx_training_jobs_dataset
ON training_jobs(dataset_id, created_at DESC);
-- Cron dispatch/poll scans by status; this keeps both bounded lookups
-- (getQueuedTrainingJobs, getInFlightTrainingJobs) index-only.
CREATE INDEX IF NOT EXISTS idx_training_jobs_status
ON training_jobs(status, id ASC);
209 changes: 209 additions & 0 deletions workers/api/scripts/smoke_api.mjs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -869,6 +869,215 @@ cookie = myCookie;
}
}

// ---- Dataset Library + training (Task #21) --------------------------
// Fresh account so the 200-token signup grant is intact — training a
// LoRA costs 150, deliberately priced to fit inside that grant.
{
const trainer = privateKeyToAccount(generatePrivateKey());
await login(trainer);

let datasetSlug = null;
{
const r = await api("POST", "/v1/datasets", {
title: `Smoke Dataset ${RUN}`,
description: "A hand-collected archive of smoke-test textures.",
rights_declaration: "own",
});
datasetSlug = r.data?.dataset?.slug ?? null;
record("POST /v1/datasets", r.status === 201 && !!datasetSlug, JSON.stringify(r.data).slice(0, 200));
}
{
const r = await api("POST", "/v1/datasets", { title: "Bad rights" });
record(
"POST /v1/datasets rejects missing rights_declaration",
r.status === 400 && r.data?.error === "invalid_rights_declaration",
JSON.stringify(r.data),
);
}
{
const r = await api("GET", "/v1/datasets/mine");
record(
"GET /v1/datasets/mine",
r.status === 200 && r.data?.datasets?.some((d) => d.slug === datasetSlug),
JSON.stringify(r.data).slice(0, 200),
);
}

// A 1x1 PNG, re-encoded with a distinct trailing comment byte per
// upload so each item hashes uniquely (dedup is exact-content-hash).
const onePxPngB64 =
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==";
function pngDataUrl(tag) {
// Append a harmless PNG tEXt-like tail so distinct calls hash
// differently; real decoders ignore trailing garbage bytes but we
// never decode these ourselves — only R2-store and hash them.
const bytes = Buffer.from(onePxPngB64, "base64");
const tagged = Buffer.concat([bytes, Buffer.from(`smoke-${tag}`)]);
return `data:image/png;base64,${tagged.toString("base64")}`;
}

let firstItemId = null;
const itemIds = [];
for (let i = 0; i < 5; i++) {
const r = await api("POST", `/v1/datasets/${datasetSlug}/items`, {
data_url: pngDataUrl(i),
caption: `texture ${i}`,
});
if (i === 0) firstItemId = r.data?.item?.id ?? null;
if (r.status === 201) itemIds.push(r.data.item.id);
record(`POST /v1/datasets/:slug/items (${i + 1}/5)`, r.status === 201, JSON.stringify(r.data).slice(0, 150));
}
{
// Re-upload the first item's exact bytes — must dedup, not double-count.
const r = await api("POST", `/v1/datasets/${datasetSlug}/items`, {
data_url: pngDataUrl(0),
});
record(
"duplicate item upload dedups instead of double-counting",
r.status === 200 && r.data?.duplicate === true && r.data?.item?.id === firstItemId,
JSON.stringify(r.data),
);
}
{
const r = await api("POST", `/v1/datasets/${datasetSlug}/items`, {
data_url: "data:text/plain;base64,aGVsbG8=",
});
record(
"unsupported mime rejected",
r.status === 400 && r.data?.error === "unsupported_mime",
JSON.stringify(r.data),
);
}
{
const r = await api("GET", `/v1/datasets/${datasetSlug}/items`);
record(
"GET /v1/datasets/:slug/items reflects item_count exactly (dedup excluded)",
r.status === 200 && r.data?.items?.length === 5,
`count=${r.data?.items?.length}`,
);
}
{
const r = await api("POST", `/v1/datasets/${datasetSlug}/items/import-urls`, { urls: [] });
record(
"import-urls rejects an empty list",
r.status === 400 && r.data?.error === "invalid_urls",
JSON.stringify(r.data),
);
}
{
const tooMany = Array.from({ length: 26 }, (_, i) => `https://example.com/${i}.png`);
const r = await api("POST", `/v1/datasets/${datasetSlug}/items/import-urls`, { urls: tooMany });
record(
"import-urls rejects more than 25 urls",
r.status === 400 && r.data?.error === "too_many_urls",
JSON.stringify(r.data),
);
}
{
const r = await api("GET", `/v1/datasets/${datasetSlug}/items/${firstItemId}/file`);
record(
"GET dataset item file (owner) returns the stored bytes",
r.status === 200,
`status=${r.status}`,
);
}

let trainingJobId = null;
let publishedModelId = null;
{
// Below MIN_TRAIN_ITEMS on a brand-new, item-less dataset.
const r0 = await api("POST", "/v1/datasets", { title: "Too small", rights_declaration: "own" });
const emptySlug = r0.data?.dataset?.slug;
const r = await api("POST", `/v1/datasets/${emptySlug}/train`, {
base_model: "flux-dev",
training_method: "lora",
});
record(
"train rejects a dataset below the minimum item count",
r.status === 422 && r.data?.error === "not_enough_items",
JSON.stringify(r.data),
);
}
{
const r = await api("POST", `/v1/datasets/${datasetSlug}/train`, {
base_model: "made-up-model",
training_method: "lora",
});
record(
"train rejects an unlisted base model",
r.status === 400 && r.data?.error === "invalid_base_model",
JSON.stringify(r.data),
);
}
{
const r = await api("POST", `/v1/datasets/${datasetSlug}/train`, {
base_model: "flux-dev",
training_method: "lora",
});
trainingJobId = r.data?.job?.id ?? null;
publishedModelId = r.data?.job?.model_id ?? null;
record(
"POST train (TRAINING_MOCK) completes synchronously and succeeds",
r.status === 201 && r.data?.job?.status === "succeeded" && !!r.data?.job?.model_version_id,
JSON.stringify(r.data).slice(0, 250),
);
}
{
const r = await api("GET", "/v1/tokens/account");
record(
"training debit landed (150 tokens spent on top of prior renders)",
r.status === 200 && r.data?.lifetime_spent >= 150,
JSON.stringify(r.data),
);
}
{
const r = await api("GET", `/v1/training/${trainingJobId}`);
record("GET /v1/training/:id (owner)", r.status === 200 && r.data?.job?.id === trainingJobId, JSON.stringify(r.data).slice(0, 200));
}
{
const r = await api("GET", `/v1/datasets/${datasetSlug}/training`);
record(
"GET /v1/datasets/:slug/training lists the job",
r.status === 200 && r.data?.jobs?.some((j) => j.id === trainingJobId),
JSON.stringify(r.data).slice(0, 200),
);
}
{
// The trained model version carries the dataset's provenance and
// is renderable exactly like any other fal_custom model.
const r = await api("GET", "/v1/models/mine");
const model = r.data?.models?.find((mo) => mo.id === publishedModelId);
record(
"trained model appears in /v1/models/mine with provenance",
r.status === 200 && !!model && model.latest_version === 1,
JSON.stringify(model ?? {}).slice(0, 250),
);
}
{
const r = await api("DELETE", `/v1/datasets/${datasetSlug}/items/${itemIds[itemIds.length - 1]}`);
record("DELETE dataset item (owner)", r.status === 200, JSON.stringify(r.data));
}
{
const r = await api("GET", `/v1/datasets/${datasetSlug}`);
record(
"item_count decremented after delete",
r.status === 200 && r.data?.dataset?.item_count === 4,
`item_count=${r.data?.dataset?.item_count}`,
);
}
{
// A second account can't read a private dataset's items.
const stranger = privateKeyToAccount(generatePrivateKey());
await login(stranger);
const r = await api("GET", `/v1/datasets/${datasetSlug}`);
record(
"private dataset invisible to a non-owner",
r.status === 404,
`status=${r.status}`,
);
}
}

const failed = results.filter((r) => !r.ok);
console.log(`\n=== ${results.length - failed.length}/${results.length} passed ===`);
if (failed.length) {
Expand Down
Loading
Loading