mirror of
https://github.com/civitai/civitai.git
synced 2026-09-20 22:08:18 +08:00
Merge pull request #4878 from civitai/feat/huggingface-model-import
Feat/huggingface model import
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: official-model-admin
|
||||
description: Create a Draft CivitaiOfficial model and version through the API, work out whether the version is API-only or needs hosted model files (and walk the user through uploading them), update an existing model's description, or transfer a model to CivitaiOfficial. Every description write requires the user's approval of the exact text, enforced by an approval hash. Use when setting up an official model or version for review before publishing, or when an official model's description needs to change. Called by onboard-generator-model; usable on its own.
|
||||
description: Create a Draft CivitaiOfficial model and version through the API, work out whether the version is API-only or needs hosted model files (and either import those from Hugging Face or walk the user through uploading them), update an existing model's description, or transfer a model to CivitaiOfficial. Every description write requires the user's approval of the exact text, enforced by an approval hash. Use when setting up an official model or version for review before publishing, or when an official model's description needs to change. Called by onboard-generator-model; usable on its own.
|
||||
---
|
||||
|
||||
# Official Model Admin
|
||||
@@ -82,14 +82,46 @@ node .claude/skills/official-model-admin/model.mjs create-version --model-id <id
|
||||
|
||||
### 4. Hosted weights only: get the files uploaded
|
||||
|
||||
`create-version` prints the upload link, `/models/<modelId>/model-versions/<versionId>/wizard?step=2`. This is the one step that can't be scripted.
|
||||
Two routes. When the weights live on Hugging Face, our servers can fetch them and you attach the
|
||||
result yourself; otherwise the user uploads through the wizard.
|
||||
|
||||
#### From Hugging Face
|
||||
|
||||
Ask the user to queue the repo at **`/moderator/huggingface-import`** — paste the model URL, check the
|
||||
**Group name** (prefilled from the repo; it is what the batch is filed under, and nothing renames it
|
||||
after Import), tick the files, Import. The transfer runs server-side on a cron, so it takes as long as it takes; nothing
|
||||
downloads to anyone's machine. Then:
|
||||
|
||||
```bash
|
||||
node .claude/skills/official-model-admin/model.mjs hf-imports --repo <owner/name>
|
||||
node .claude/skills/official-model-admin/model.mjs attach-import --import <id> --version <id> --type Model --writable
|
||||
```
|
||||
|
||||
`hf-imports` lists each transferred file with its size, state, group and a suggested type;
|
||||
`attach-import` creates the model file on your version, and scanning and hashing follow on their own.
|
||||
|
||||
Both filters run on the server. `--repo` is an **exact** match against the repo id Hugging Face
|
||||
returned rather than what was pasted, so its casing must be HF's; `--group` is a case-insensitive
|
||||
substring of the group name and is the forgiving one to reach for.
|
||||
|
||||
🔴 **You choose `--type`, and it decides whether the version can load at all.** The suggested type is
|
||||
advisory and deliberately never names the primary weights — a mislabelled weight file passes every
|
||||
check and produces a version nothing can load. Read the filename:
|
||||
`ae.safetensors` is a `VAE`, anything under `text_encoder/` is a `Text Encoder`, and the large
|
||||
`.safetensors` at the repo root is the weight file (`Model`, or `Diffusion Model` / `UNet` when the
|
||||
repo splits them). If the repo's layout does not make a file's role obvious, ask the user rather than
|
||||
guessing — a mislabelled weight file passes every check here and produces a version nothing can load.
|
||||
|
||||
#### Uploaded by hand
|
||||
|
||||
`create-version` prints the upload link, `/models/<modelId>/model-versions/<versionId>/wizard?step=2`.
|
||||
|
||||
**Give the user that link and ask them to upload the model files there.** They can also use the model page: the version menu → **Manage files**. Tell them what the files need:
|
||||
|
||||
- **At least one weight file** with type `Model`, `Pruned Model`, `Diffusion Model`, `UNet`, `Negative` or `VAE`. Supporting files such as text encoders can also be uploaded.
|
||||
- **`SafeTensor` format** for a checkpoint.
|
||||
|
||||
Then wait. When they say the upload is done, run:
|
||||
Then wait. When the files are attached or the user says the upload is done, run:
|
||||
|
||||
```bash
|
||||
node .claude/skills/official-model-admin/model.mjs files --version <id>
|
||||
|
||||
@@ -17,7 +17,7 @@ import { queryDb } from '../generation-coverage/coverage.mjs';
|
||||
const OFFICIAL_USER_ID = 12042163;
|
||||
const KINDS = { 'api-only': 'ExternalGeneration', 'hosted-weights': 'Download' };
|
||||
const ECOSYSTEMS_DIR = resolve(projectRoot, 'src/server/services/orchestrator/ecosystems');
|
||||
// Mirrors checkLoadable in src/server/services/resource-load.service.ts.
|
||||
// Mirrors LOADABLE_FILE_TYPES in src/utils/file-display-helpers.ts, which checkLoadable applies.
|
||||
const LOADABLE_FILE_TYPES = ['Model', 'Pruned Model', 'Diffusion Model', 'UNet', 'Negative', 'VAE'];
|
||||
// Closed models reachable only through their provider's API, plus fal, which hosts third-party models.
|
||||
const EXTERNAL_ENGINES = [
|
||||
@@ -303,6 +303,69 @@ async function files() {
|
||||
console.log('\nREADY: the weight files are uploaded and scanned. Next: the coverage row.');
|
||||
}
|
||||
|
||||
|
||||
// The transfer itself is queued from /moderator/huggingface-import.
|
||||
|
||||
const formatBytes = (bytes) => {
|
||||
if (!bytes) return '—';
|
||||
const units = ['B', 'KB', 'MB', 'GB', 'TB'];
|
||||
const i = Math.min(Math.floor(Math.log(bytes) / Math.log(1024)), units.length - 1);
|
||||
return `${(bytes / 1024 ** i).toFixed(i ? 2 : 0)} ${units[i]}`;
|
||||
};
|
||||
|
||||
async function hfImports() {
|
||||
// Filtered server-side: a client-side filter over this page silently misses anything past the
|
||||
// limit, and reports it as "no imports" — indistinguishable from never having imported it.
|
||||
const input = { limit: 100 };
|
||||
if (flags.repo) input.repo = flags.repo;
|
||||
if (flags.group) input.groupName = flags.group;
|
||||
const shown = await trpcCall('huggingFaceImport.getAll', input, 'GET');
|
||||
if (!shown.length) {
|
||||
const what = flags.repo ?? flags.group;
|
||||
return console.log(what ? `No imports matching ${what}.` : 'No imports.');
|
||||
}
|
||||
|
||||
for (const row of shown) {
|
||||
const attached = row.modelFileId
|
||||
? `attached → file ${row.modelFileId} on version ${row.modelVersionId}`
|
||||
: row.status === 'Completed'
|
||||
? 'ready to attach'
|
||||
: '';
|
||||
console.log(
|
||||
[
|
||||
String(row.id).padStart(5),
|
||||
row.status.padEnd(12),
|
||||
formatBytes(row.sizeBytes).padStart(10),
|
||||
`${row.repo}/${row.filename}`,
|
||||
`[${row.groupName}]`,
|
||||
row.suggestedType ? `(suggests ${row.suggestedType})` : '',
|
||||
attached,
|
||||
]
|
||||
.filter(Boolean)
|
||||
.join(' ')
|
||||
);
|
||||
if (row.error) console.log(` ${row.error}`);
|
||||
}
|
||||
console.log(
|
||||
`
|
||||
Attach with: attach-import --import <id> --version <id> --type <${LOADABLE_FILE_TYPES.join('|')}|Text Encoder|Config|...>`
|
||||
);
|
||||
console.log('The type decides whether the version can load — read the filename, do not guess.');
|
||||
}
|
||||
|
||||
async function attachImport() {
|
||||
const id = requiredInt('import');
|
||||
const modelVersionId = requiredInt('version');
|
||||
const type = required('type');
|
||||
if (!writable) return dryRun('huggingFaceImport.attach', { id, modelVersionId, type });
|
||||
|
||||
const result = await trpcCall('huggingFaceImport.attach', { id, modelVersionId, type });
|
||||
console.log(
|
||||
`Attached import ${id} as model file ${result.modelFileId} on version ${result.modelVersionId}.`
|
||||
);
|
||||
console.log('Scanning and hashing start on their own. Check with: files --version ' + modelVersionId);
|
||||
}
|
||||
|
||||
const HELP = `Usage: node .claude/skills/official-model-admin/model.mjs <command> [flags]
|
||||
|
||||
whoami
|
||||
@@ -313,6 +376,8 @@ const HELP = `Usage: node .claude/skills/official-model-admin/model.mjs <command
|
||||
create-version --model-id <id> --name <n> --base-model <name> --kind <api-only|hosted-weights>
|
||||
[--no-download] [--writable]
|
||||
files --version <id> are the uploaded files ready?
|
||||
hf-imports [--repo <owner/name>] [--group <name>] Hugging Face transfers and their state
|
||||
attach-import --import <id> --version <id> --type <type> [--writable]
|
||||
|
||||
Description writes need --approved <hash>, printed by the dry run of the same command.
|
||||
--kind api-only → ExternalGeneration (no files); hosted-weights → Download, or Generation with --no-download.
|
||||
@@ -327,6 +392,8 @@ dispatch(
|
||||
evidence,
|
||||
'create-version': createVersion,
|
||||
files,
|
||||
'hf-imports': hfImports,
|
||||
'attach-import': attachImport,
|
||||
},
|
||||
HELP
|
||||
);
|
||||
|
||||
@@ -30,7 +30,7 @@ Supporting skills: `deploy-status` (the deploy stops), `postgres-query` (used by
|
||||
|
||||
Ask for the model name and a reference link, plus the Civitai model URL if the model already exists.
|
||||
|
||||
First work out the **kind**: `api-only` (the provider runs it, no files) or `hosted-weights` (we run it from files the user uploads). Settle it with the "Versions: API-only or hosted weights" steps in `official-model-admin`: gather the evidence, then have the user confirm. Never guess it. Settle the kind before the case, because the kind decides the case.
|
||||
First work out the **kind**: `api-only` (the provider runs it, no files) or `hosted-weights` (we run it from files on the version). Settle it with the "Versions: API-only or hosted weights" steps in `official-model-admin`: gather the evidence, then have the user confirm. Never guess it. Settle the kind before the case, because the kind decides the case.
|
||||
|
||||
Then work out the **case**:
|
||||
|
||||
@@ -52,7 +52,7 @@ Adding a base model or an ecosystem costs a second deploy and is hard to undo on
|
||||
|
||||
Show the user which rule matched and the evidence for it: the kind, whether the line already has an ecosystem, and for hosted weights, why existing resources do or don't work on it. The user confirms the case. Never guess it, just as you never guess the kind.
|
||||
|
||||
The kind also changes Phase 3. A hosted-weights version adds a stop while the user uploads its files.
|
||||
The kind also changes Phase 3. A hosted-weights version adds a stop for its files — a server-side Hugging Face import, or a browser upload by the user.
|
||||
|
||||
Show the user the case, the kind, the phases and where the run will stop, for deploys and for uploads. Get their confirmation before continuing.
|
||||
|
||||
@@ -67,7 +67,7 @@ Show the user the case, the kind, the phases and where the run will stop, for de
|
||||
|
||||
3. **Version and coverage.**
|
||||
1. `official-model-admin create-version --kind <kind>`, using the kind from Phase 0.
|
||||
2. **Hosted weights only:** give the user the upload link `create-version` prints, and **stop until they say the upload is done**. Then run `official-model-admin files --version <id>` until it reports READY. If it reports NOT READY, pass its reason on to the user.
|
||||
2. **Hosted weights only:** get the files onto the version. If the weights are on Hugging Face, ask the user to queue the repo at `/moderator/huggingface-import`, then attach each transferred file with `official-model-admin attach-import` — see "From Hugging Face" in that skill's step 4; you choose the file type, so read the filenames rather than guessing. Otherwise give the user the upload link `create-version` prints and **stop until they say the upload is done**. Either way, then run `official-model-admin files --version <id>` until it reports READY. If it reports NOT READY, pass its reason on to the user.
|
||||
3. `generation-coverage add`, labelled with the base model name.
|
||||
|
||||
4. **Gate, before the deploy.** `generation-gate-rules add`:
|
||||
|
||||
+6
-1
@@ -2,7 +2,8 @@
|
||||
# Keep this file up-to-date when you add new variables to `.env`.
|
||||
# This file will be committed to version control, so make sure not to have any secrets in it.
|
||||
# If you are cloning this repo, create a copy of this file named `.env` and populate it with your secrets.
|
||||
# When adding additional env variables, the schema in /env/schema.mjs should be updated accordingly
|
||||
# When adding additional env variables, update src/env/server-schema.ts (or client-schema.ts for
|
||||
# NEXT_PUBLIC_* vars) accordingly
|
||||
# The default values for Prisma, Redis, S3, and Email are set to work with the docker-compose setup
|
||||
|
||||
# Database
|
||||
@@ -51,6 +52,10 @@ DISCORD_BOT_TOKEN=
|
||||
DISCORD_GUILD_ID=
|
||||
DISCORD_WEBHOOK_MOD_ALERTS=
|
||||
|
||||
# Optional. Without it, only public ungated Hugging Face repos can be imported at
|
||||
# /moderator/huggingface-import; with it, repos this token's account has accepted the terms for.
|
||||
# HUGGING_FACE_TOKEN=
|
||||
|
||||
# File uploading
|
||||
S3_UPLOAD_KEY=REFER_TO_README
|
||||
S3_UPLOAD_SECRET=REFER_TO_README
|
||||
|
||||
@@ -65,7 +65,7 @@ MiniMax H3 appears in both rows, one version each, so the base model alone doesn
|
||||
|
||||
The deciding question is whether the provider publishes weights we run, or only an API. `official-model-admin evidence` collects these signals, and the user confirms the kind.
|
||||
|
||||
The upload wizard (`/models/<modelId>/model-versions/<versionId>/wizard`) skips its files step for `ExternalGeneration`. For hosted weights, files go in at `?step=2`, or through the **Manage files** item in the version menu. This is the only step that has to be done in the browser.
|
||||
The upload wizard (`/models/<modelId>/model-versions/<versionId>/wizard`) skips its files step for `ExternalGeneration`. For hosted weights there are two routes: weights already on Hugging Face are queued at `/moderator/huggingface-import` and fetched server-side, then attached with `official-model-admin attach-import`; anything else is uploaded by hand at `?step=2`, or through the **Manage files** item in the version menu. Either way a human drives it — this is the one step no skill completes on its own.
|
||||
|
||||
A hosted-weights version is ready to cover once it has a scanned file of a weight type, and a checkpoint also needs a SafeTensor file. That's the same rule `checkLoadable` in `src/server/services/resource-load.service.ts` applies. `official-model-admin files` checks it.
|
||||
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
# Importing models from Hugging Face
|
||||
|
||||
**Status:** built, not deployed. The migration has not been applied to any database.
|
||||
|
||||
## The goal
|
||||
|
||||
Give a moderator a Hugging Face repo URL and have *our* servers fetch the weights into our storage —
|
||||
instead of a person downloading 20 GB and re-uploading it through the browser wizard. That upload was the one
|
||||
part of `official-model-admin` the skill had to hand back to a human; step 4 of that skill now offers
|
||||
this as the first route.
|
||||
|
||||
## What was built
|
||||
|
||||
| Piece | Where |
|
||||
| --- | --- |
|
||||
| `HuggingFaceImport` table + status enum | `packages/civitai-db-schema/prisma/schema.full.prisma`, migration `20260914120000_huggingface_import` |
|
||||
| HF API client — parse a URL, resolve a branch to a commit sha, list files with sizes and LFS sha256, ranged reads | `src/server/services/huggingface.service.ts` |
|
||||
| Queue + the resumable transfer | `src/server/services/huggingface-import.service.ts` |
|
||||
| The runner | `src/server/jobs/process-huggingface-imports.ts`, registered in the `jobs` array in `run-jobs` |
|
||||
| tRPC surface (`getAll` — filterable by `groupName`/`repo`, and by `unattached` — `getCounts`, `lookup`, `enqueue`, `attach`, `detach`, `delete`, `renameGroup`, `retry`, `cancel`) | `src/server/routers/huggingface-import.router.ts` |
|
||||
| Moderator page | `src/pages/moderator/huggingface-import.tsx` + `src/components/Moderation/HuggingFaceImport/` |
|
||||
| The *Manage files* picker (moderator-only) | `src/components/Moderation/HuggingFaceImport/AddFromImportsModal.tsx`, opened via `src/components/Dialog/triggers/add-from-hugging-face-imports.ts` from `AddFromImportsButton.tsx` in `src/components/Resource/Files.tsx` |
|
||||
| Server-side multipart helpers (`createMultipartUpload`, `uploadPart`) | `src/utils/s3-utils.ts` |
|
||||
| The shared key builder (`buildUploadKey`) | `src/utils/upload-key.ts` — used by `/api/upload` **and** the import |
|
||||
|
||||
Config: `HUGGING_FACE_TOKEN` is **optional**. Without it, public ungated repos import fine; with it,
|
||||
repos whose terms that token's account has accepted.
|
||||
|
||||
## Why a resumable transfer, and why a cron job
|
||||
|
||||
A 20 GB transfer cannot live inside one request or one job run — jobs here hold a Redis lock sized in
|
||||
minutes (5 typical, 30 at the longest), and a pod rolls on every deploy. So the transfer is a sequence
|
||||
of **independent parts**:
|
||||
|
||||
- HF's CDN honours HTTP `Range`, so each part is its own request.
|
||||
- S3/B2 multipart lets a *different process* upload part N, as long as it holds the `uploadId`.
|
||||
|
||||
The row stores `uploadId`, `partSize` and the `parts` written so far. Each job run moves as many parts as fit in
|
||||
its budget (2 minutes under a 5-minute lock, leaving room for the slowest part still in flight) and
|
||||
stops; the next run continues the same file.
|
||||
A deploy costs one part, not the file. `advanceImport` re-reads the row's status between parts, so a
|
||||
cancel takes effect within one part rather than at the end.
|
||||
|
||||
Claiming is `FOR UPDATE SKIP LOCKED`. A run that yields cleanly clears its claim, so the row is
|
||||
eligible on the very next tick; a run that *dies* leaves its claim, and the 20-minute stale window is
|
||||
what recovers it. That asymmetry is deliberate — it is what tells "out of budget" from "the pod went
|
||||
away."
|
||||
|
||||
## Attaching to a version
|
||||
|
||||
`attach` turns a finished import into a `ModelFile` on a version, going through `createFileHandler`
|
||||
rather than writing the row directly — that is what gets the storage-resolver registration and the
|
||||
inline scan submission, so scanning and hashing follow on their own.
|
||||
|
||||
Four ways in, one path underneath: the Attach control on a completed row, the **Add from Hugging
|
||||
Face imports** picker inside a version's *Manage files*, the `huggingFaceImport.attach` procedure,
|
||||
and two commands on `official-model-admin` — `hf-imports` (what transferred) and `attach-import`
|
||||
(put one on a version).
|
||||
|
||||
What that changes for the skill: the 20GB re-upload a person used to perform is gone, and attaching is
|
||||
scriptable. A human still queues the repo on the page, and still confirms the file type when the
|
||||
filename does not settle it — `suggestFileType` is returned on both `lookup` and `getAll`, but it is
|
||||
advisory and refuses to name the primary weights.
|
||||
|
||||
🔴 **The file type is explicit, never inferred.** `suggestFileType` offers a guess from the filename
|
||||
for accessories (VAE, text encoder, config) and deliberately returns null for the primary weights,
|
||||
because a mislabelled weight file passes every check here and produces a version nothing can load.
|
||||
|
||||
## Naming, and what a group is
|
||||
|
||||
**An imported object is keyed exactly like a browser upload.** Both paths call the same
|
||||
`buildUploadKey` and resolve the same backend, so a file transferred from Hugging Face is stored at
|
||||
`model/<userId>/<name>.<token><ext>` in the bucket `/api/upload` would have used — B2 whenever
|
||||
`S3_UPLOAD_B2_ENDPOINT` is set — with the same `filenamize` and the same 4-character token. That shape is load-bearing: `/api/upload/sign-part` authorises a part by reading the
|
||||
userId out of segment 1, so nothing may be inserted ahead of it.
|
||||
|
||||
🔴 **Nothing about the import appears in the key** — not the repo, not the revision, not the group.
|
||||
A key is immutable once the object exists, so anything encoded there could never be corrected without
|
||||
copying the bytes. **The `HuggingFaceImport` row is the index**: `repo`, `revision`, `filename` and
|
||||
`groupName` are columns, and columns can be fixed.
|
||||
|
||||
**The group** is the batch a moderator filed the import under:
|
||||
|
||||
- `repo` is taken from **Hugging Face's own response**, never from the pasted URL — the same repo typed
|
||||
with different casing would otherwise be filed under two groups nothing could merge.
|
||||
- `groupName` defaults to the repo's own name (`black-forest-labs/FLUX.1-Krea-dev` → `FLUX.1-Krea-dev`)
|
||||
and is typed on the lookup screen before Import is pressed.
|
||||
- It can be renamed later, at any status, from the group header on the **Unattached** tab. The name
|
||||
never reaches a storage key, so a rename desynchronises nothing. `renameGroup` is scoped by the
|
||||
group's current name as well as repo and revision, because one repo at one revision can be two
|
||||
batches.
|
||||
|
||||
### Finding a group again
|
||||
|
||||
`getAll` filters **on the server** — `groupName` as a case-insensitive substring, `repo` as an exact
|
||||
match on the id Hugging Face returned. Both the page's filter box and the skill's `--group`/`--repo`
|
||||
pass through to it.
|
||||
|
||||
🔴 That is not a convenience. Both callers take at most 100 rows, and filtering those client-side
|
||||
meant a group older than that window returned nothing — which reads exactly like "never imported".
|
||||
`@@index([groupName])` serves the equality lookups; the substring search does not use it, and at this
|
||||
table's size that is a few milliseconds of seq scan. If it ever stops being one, the answer is a
|
||||
trigram index rather than a narrower filter.
|
||||
|
||||
## Throughput
|
||||
|
||||
Segmenting is nearly free; doing the segments **serially** was not. Each part is two hops — a ranged
|
||||
read from HF, a part written to the bucket — and one-at-a-time left both idle half the time, making the
|
||||
import roughly half the speed of a single stream.
|
||||
|
||||
Parts now move concurrently: `PARTS_IN_FLIGHT` (3) per file, two files per run. Parallel range
|
||||
requests also get more out of Hugging Face than one connection does, which is the same reason
|
||||
`hf_transfer` exists. **Memory is the ceiling, not the network** —
|
||||
`PARTS_IN_FLIGHT × file concurrency × partSize` = 3 × 2 × 16 MB ≈ 96 MB resident today. That is the dial.
|
||||
|
||||
Because parts finish out of order, the completed set has holes: the resume point is the set of missing
|
||||
part numbers, never `parts.length + 1`, and `completeMultipartUpload` gets them sorted.
|
||||
|
||||
No real-world measurement exists yet — nothing has run against a live bucket.
|
||||
|
||||
## What it does not do yet
|
||||
|
||||
**Nothing checks that the file you attached is the file you meant.** The scan catches malware, not
|
||||
mislabelling.
|
||||
|
||||
The direction is a **pull**: an import usually happens before anyone knows which version will want
|
||||
it, so the version draws from the pool rather than the import pushing at a version. Both surfaces
|
||||
that make that work now exist — the **Unattached** tab on the import page and the **Add from Hugging
|
||||
Face imports** picker inside a version's *Manage files*, both grouped and filtered by `groupName`.
|
||||
|
||||
The only two exits from the unattached list are *attached* and *deleted*: there is deliberately no
|
||||
dismissed-but-stored state, because a hidden row still costs storage and would let the count
|
||||
understate what we hold.
|
||||
|
||||
🔴 **Deleting refuses whenever a `ModelFile` still points at the object**, resolved through
|
||||
`urlsSafeToDelete` over `ModelFile.url` — not through the import row's `modelFileId`, which detach
|
||||
clears while leaving the `ModelFile` alive. Judging from the row alone destroys the bytes a
|
||||
published version is serving, two clicks after a detach.
|
||||
|
||||
**No quota.** Moderator-only at the router; everything underneath is already scoped per owner, so
|
||||
opening it up needs a per-user quota — size and count — and a `userId` in the
|
||||
`(repo, revision, filename)` unique index — the header comment in `huggingface-import.router.ts` is
|
||||
the prerequisite list.
|
||||
|
||||
**Licenses are shown, not enforced.** The page surfaces the declared license and flags gated repos.
|
||||
|
||||
## What HF gives us before any bytes move
|
||||
|
||||
`/api/models/{repo}/tree/{revision}?recursive=true` returns each file's `size` and, for LFS files
|
||||
(every weight file), `lfs.oid` — **the content sha256**. So the page can mark a file we already store
|
||||
(`ModelFileHash` lookup) and offer to skip it, and the transfer knows its part count up front. The
|
||||
commit `sha` is pinned rather than `main`, so a re-import is reproducible and the provenance record
|
||||
means something.
|
||||
|
||||
## The old importer is retired
|
||||
|
||||
The previous importer (`src/server/importers/*`, `GET /api/import` and the hourly `processImportsJob`)
|
||||
is deleted. It created a `Model` with no versions, hardcoded `baseModel` to `SD 1.5`, and pointed
|
||||
`ModelFile.url` at huggingface.co, so it never transferred anything — and it carried a second HF
|
||||
client. Removing `processImportsJob` from the `jobs` array is what stops the scheduler running it.
|
||||
|
||||
The `Import` table, its `ImportStatus` enum, and the `fromImportId` columns on `Model` and
|
||||
`ModelVersion` are left in place: dropping them is a migration over existing rows, and nothing reads
|
||||
them any more.
|
||||
|
||||
## Open questions
|
||||
|
||||
1. **Who may import**, and under what quota — see above.
|
||||
2. **Whether redistribution rights should be enforced** rather than displayed.
|
||||
3. ~~**Does deleting an unattached import remove its row, or leave a tombstone?**~~ **Settled: hard
|
||||
delete.** `(repo, revision, filename)` is unique, so a tombstone would block re-importing that
|
||||
exact file — and a deliberate deletion is precisely the case where you might want it back. The
|
||||
cost is losing the record that we once held those bytes.
|
||||
4. **Attribution.** The row records repo, revision and filename, and an attached file links back to it.
|
||||
Nothing surfaces that on the model page yet.
|
||||
@@ -37,7 +37,7 @@ Tiering reflects head-moderator guidance on what's actually used day-to-day.
|
||||
- **[Tier 2 — Low priority](#tier-2--low-priority):** pages the head moderator doesn't use or doesn't know about. Real features, but defer until Tier 1 is done.
|
||||
- **[Excluded — will not migrate](#excluded--will-not-migrate):** Paddle (per decision: no Paddle pages in the moderator app) + dev/test scaffolds.
|
||||
|
||||
> **Counts:** Tier 1 ≈ **31** pages · Tier 2 ≈ **30** pages · Excluded **6** (2 Paddle + 4 scaffolds).
|
||||
> **Counts:** Tier 1 ≈ **31** pages · Tier 2 ≈ **31** pages · Excluded **6** (2 Paddle + 4 scaffolds).
|
||||
|
||||
### Suggested order within Tier 1
|
||||
|
||||
@@ -489,6 +489,13 @@ Real, working features the head moderator doesn't use or doesn't know about. Mig
|
||||
- Infra: **Postgres (raw SQL on `Image`) + Redis (sysRedis `RATINGS_SANITY_IDS` set)**
|
||||
- Notes: not on the head-mod list; commented out in nav + research-only — parked here pending confirmation (could be excluded).
|
||||
|
||||
- [ ] **`/moderator/huggingface-import`** — `huggingface-import.tsx` + `components/Moderation/HuggingFaceImport/` — flag: none (`requireModerator`) — **added 2026-09-14, after this inventory was taken**
|
||||
- Procedures: `huggingFaceImport.getAll`, `getCounts`, `getConfig` (queries); `lookup`, `enqueue`, `attach`, `detach`, `delete`, `renameGroup`, `retry`, `cancel`, `setConfig` (mutations)
|
||||
- Services: `huggingface.service.ts` (HF API client), `huggingface-import.service.ts` (queue + resumable multipart transfer), `huggingface-import-config.service.ts` (transfer settings)
|
||||
- Schemas: `huggingface-import.schema.ts`
|
||||
- Infra: **Postgres (`HuggingFaceImport`) + S3/B2 multipart + Redis (sysRedis transfer config) + the `process-huggingface-imports` cron**
|
||||
- Notes: a port moves the page, not the transfer — the cron and `createFileHandler` (scan + hash submission) stay in the main app, so this is delegate-shaped if it moves at all.
|
||||
|
||||
---
|
||||
|
||||
# Excluded — will not migrate
|
||||
|
||||
+46
@@ -0,0 +1,46 @@
|
||||
-- CreateEnum
|
||||
CREATE TYPE "HuggingFaceImportStatus" AS ENUM ('Queued', 'Transferring', 'Completed', 'Failed', 'Canceled');
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "HuggingFaceImport" (
|
||||
"id" SERIAL NOT NULL,
|
||||
"repo" TEXT NOT NULL,
|
||||
"revision" TEXT NOT NULL,
|
||||
"filename" TEXT NOT NULL,
|
||||
"groupName" TEXT NOT NULL,
|
||||
"sourceUrl" TEXT NOT NULL,
|
||||
"sizeBytes" BIGINT,
|
||||
"sourceSha256" TEXT,
|
||||
"status" "HuggingFaceImportStatus" NOT NULL DEFAULT 'Queued',
|
||||
"bytesTransferred" BIGINT NOT NULL DEFAULT 0,
|
||||
"uploadId" TEXT,
|
||||
"partSize" INTEGER,
|
||||
"parts" JSONB,
|
||||
"bucket" TEXT,
|
||||
"key" TEXT,
|
||||
"url" TEXT,
|
||||
"error" TEXT,
|
||||
"attempts" INTEGER NOT NULL DEFAULT 0,
|
||||
"nextAttemptAt" TIMESTAMP(3),
|
||||
"userId" INTEGER,
|
||||
"modelVersionId" INTEGER,
|
||||
"modelFileId" INTEGER,
|
||||
"claimedBy" TEXT,
|
||||
"claimedAt" TIMESTAMP(3),
|
||||
"heartbeatAt" TIMESTAMP(3),
|
||||
"startedAt" TIMESTAMP(3),
|
||||
"completedAt" TIMESTAMP(3),
|
||||
"createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updatedAt" TIMESTAMP(3) NOT NULL,
|
||||
|
||||
CONSTRAINT "HuggingFaceImport_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "HuggingFaceImport_repo_revision_filename_key" ON "HuggingFaceImport"("repo", "revision", "filename");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "HuggingFaceImport_status_createdAt_idx" ON "HuggingFaceImport"("status", "createdAt");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "HuggingFaceImport_groupName_idx" ON "HuggingFaceImport"("groupName");
|
||||
@@ -848,6 +848,63 @@ model Import {
|
||||
importId Int?
|
||||
}
|
||||
|
||||
enum HuggingFaceImportStatus {
|
||||
Queued
|
||||
Transferring
|
||||
Completed
|
||||
Failed
|
||||
Canceled
|
||||
}
|
||||
|
||||
/// One file pulled from Hugging Face into our storage. The row IS the provenance record: the
|
||||
/// uploaded object (bucket/key/url) is traceable back to the exact repo revision it came from.
|
||||
model HuggingFaceImport {
|
||||
id Int @id @default(autoincrement())
|
||||
repo String
|
||||
/// Commit sha, never a branch name — an import must name the exact bytes it took.
|
||||
revision String
|
||||
filename String
|
||||
/// What a moderator calls this batch, defaulting to the repo's own name. Never appears in a
|
||||
/// storage key — the key is the ordinary upload shape — so this is free to be corrected.
|
||||
groupName String
|
||||
sourceUrl String
|
||||
/// Both come from the HF tree API before any bytes move: size, and lfs.oid which is the content
|
||||
/// sha256 for LFS files. The sha is what lets us skip a file we already store under the same hash;
|
||||
/// the transfer itself is not verified against it.
|
||||
sizeBytes BigInt?
|
||||
sourceSha256 String?
|
||||
status HuggingFaceImportStatus @default(Queued)
|
||||
bytesTransferred BigInt @default(0)
|
||||
/// The resume point. A transfer is a sequence of ranged reads from HF written as multipart parts,
|
||||
/// and `uploadId` + `parts` is what lets a LATER job run continue one an earlier run left unfinished
|
||||
/// instead of starting the file again.
|
||||
uploadId String?
|
||||
partSize Int?
|
||||
parts Json?
|
||||
bucket String?
|
||||
key String?
|
||||
url String?
|
||||
error String?
|
||||
attempts Int @default(0)
|
||||
nextAttemptAt DateTime?
|
||||
userId Int?
|
||||
modelVersionId Int?
|
||||
modelFileId Int?
|
||||
/// Worker lease. A transfer outlives any one job run, so a claim plus a heartbeat is what stops two
|
||||
/// runs moving the same file and what lets the next run tell "in flight" from "abandoned".
|
||||
claimedBy String?
|
||||
claimedAt DateTime?
|
||||
heartbeatAt DateTime?
|
||||
startedAt DateTime?
|
||||
completedAt DateTime?
|
||||
createdAt DateTime @default(now())
|
||||
updatedAt DateTime @updatedAt
|
||||
|
||||
@@unique([repo, revision, filename])
|
||||
@@index([status, createdAt])
|
||||
@@index([groupName])
|
||||
}
|
||||
|
||||
enum ModelStatus {
|
||||
/// saved but incomplete
|
||||
Draft
|
||||
|
||||
@@ -174,6 +174,17 @@ export const ImportStatus = {
|
||||
|
||||
export type ImportStatus = (typeof ImportStatus)[keyof typeof ImportStatus];
|
||||
|
||||
export const HuggingFaceImportStatus = {
|
||||
Queued: 'Queued',
|
||||
Transferring: 'Transferring',
|
||||
Completed: 'Completed',
|
||||
Failed: 'Failed',
|
||||
Canceled: 'Canceled',
|
||||
} as const;
|
||||
|
||||
export type HuggingFaceImportStatus =
|
||||
(typeof HuggingFaceImportStatus)[keyof typeof HuggingFaceImportStatus];
|
||||
|
||||
export const ModelStatus = {
|
||||
Draft: 'Draft',
|
||||
Training: 'Training',
|
||||
|
||||
@@ -143,6 +143,15 @@ export const ImportStatus = {
|
||||
Completed: 'Completed',
|
||||
} as const;
|
||||
export type ImportStatus = (typeof ImportStatus)[keyof typeof ImportStatus];
|
||||
export const HuggingFaceImportStatus = {
|
||||
Queued: 'Queued',
|
||||
Transferring: 'Transferring',
|
||||
Completed: 'Completed',
|
||||
Failed: 'Failed',
|
||||
Canceled: 'Canceled',
|
||||
} as const;
|
||||
export type HuggingFaceImportStatus =
|
||||
(typeof HuggingFaceImportStatus)[keyof typeof HuggingFaceImportStatus];
|
||||
export const ModelStatus = {
|
||||
Draft: 'Draft',
|
||||
Training: 'Training',
|
||||
|
||||
@@ -20,6 +20,7 @@ import type {
|
||||
LinkType,
|
||||
ModelType,
|
||||
ImportStatus,
|
||||
HuggingFaceImportStatus,
|
||||
ModelStatus,
|
||||
TrainingStatus,
|
||||
CommercialUse,
|
||||
@@ -2379,6 +2380,58 @@ export type HomeBlock = {
|
||||
permanent: Generated<boolean>;
|
||||
sourceId: number | null;
|
||||
};
|
||||
export type HuggingFaceImport = {
|
||||
id: Generated<number>;
|
||||
repo: string;
|
||||
/**
|
||||
* Commit sha, never a branch name — an import must name the exact bytes it took.
|
||||
*/
|
||||
revision: string;
|
||||
filename: string;
|
||||
/**
|
||||
* What a moderator calls this batch, defaulting to the repo's own name. Never appears in a
|
||||
* storage key — the key is the ordinary upload shape — so this is free to be corrected.
|
||||
*/
|
||||
groupName: string;
|
||||
sourceUrl: string;
|
||||
/**
|
||||
* Both come from the HF tree API before any bytes move: size, and lfs.oid which is the content
|
||||
* sha256 for LFS files. The sha is what lets us skip a file we already store under the same hash;
|
||||
* the transfer itself is not verified against it.
|
||||
*/
|
||||
sizeBytes: string | null;
|
||||
sourceSha256: string | null;
|
||||
status: Generated<HuggingFaceImportStatus>;
|
||||
bytesTransferred: Generated<string>;
|
||||
/**
|
||||
* The resume point. A transfer is a sequence of ranged reads from HF written as multipart parts,
|
||||
* and `uploadId` + `parts` is what lets a LATER job run continue one an earlier run left unfinished
|
||||
* instead of starting the file again.
|
||||
*/
|
||||
uploadId: string | null;
|
||||
partSize: number | null;
|
||||
parts: unknown | null;
|
||||
bucket: string | null;
|
||||
key: string | null;
|
||||
url: string | null;
|
||||
error: string | null;
|
||||
attempts: Generated<number>;
|
||||
nextAttemptAt: Timestamp | null;
|
||||
userId: number | null;
|
||||
modelVersionId: number | null;
|
||||
modelFileId: number | null;
|
||||
/**
|
||||
* Worker lease. A transfer outlives any one job run, so a claim plus a heartbeat is what stops two
|
||||
* runs moving the same file and what lets the next run tell "in flight" from "abandoned".
|
||||
*/
|
||||
claimedBy: string | null;
|
||||
claimedAt: Timestamp | null;
|
||||
heartbeatAt: Timestamp | null;
|
||||
startedAt: Timestamp | null;
|
||||
completedAt: Timestamp | null;
|
||||
createdAt: Generated<Timestamp>;
|
||||
updatedAt: Timestamp;
|
||||
};
|
||||
export type Image = {
|
||||
id: Generated<number>;
|
||||
pHash: string | null;
|
||||
@@ -4471,6 +4524,7 @@ export type DB = {
|
||||
GenerationPreset: GenerationPreset;
|
||||
GenerationServiceProvider: GenerationServiceProvider;
|
||||
HomeBlock: HomeBlock;
|
||||
HuggingFaceImport: HuggingFaceImport;
|
||||
Image: Image;
|
||||
ImageConnection: ImageConnection;
|
||||
ImageEngagement: ImageEngagement;
|
||||
|
||||
@@ -48,6 +48,7 @@ export const UPDATED_AT_TABLES = new Set<keyof DB>([
|
||||
'EntityModeration',
|
||||
'GenerationPreset',
|
||||
'HomeBlock',
|
||||
'HuggingFaceImport',
|
||||
'Image',
|
||||
'ImageReaction',
|
||||
'LicensingRoot',
|
||||
|
||||
@@ -30,6 +30,8 @@ export type ModelType = "Checkpoint" | "TextualInversion" | "Hypernetwork" | "Ae
|
||||
|
||||
export type ImportStatus = "Pending" | "Processing" | "Failed" | "Completed";
|
||||
|
||||
export type HuggingFaceImportStatus = "Queued" | "Transferring" | "Completed" | "Failed" | "Canceled";
|
||||
|
||||
export type ModelStatus = "Draft" | "Training" | "Published" | "Scheduled" | "Unpublished" | "UnpublishedViolation" | "GatherInterest" | "Deleted";
|
||||
|
||||
export type TrainingStatus = "Pending" | "Submitted" | "Paused" | "Denied" | "Processing" | "InReview" | "Failed" | "Approved" | "Expired";
|
||||
@@ -834,6 +836,38 @@ export interface Import {
|
||||
importId: number | null;
|
||||
}
|
||||
|
||||
export interface HuggingFaceImport {
|
||||
id: number;
|
||||
repo: string;
|
||||
revision: string;
|
||||
filename: string;
|
||||
groupName: string;
|
||||
sourceUrl: string;
|
||||
sizeBytes: bigint | null;
|
||||
sourceSha256: string | null;
|
||||
status: HuggingFaceImportStatus;
|
||||
bytesTransferred: bigint;
|
||||
uploadId: string | null;
|
||||
partSize: number | null;
|
||||
parts: JsonValue | null;
|
||||
bucket: string | null;
|
||||
key: string | null;
|
||||
url: string | null;
|
||||
error: string | null;
|
||||
attempts: number;
|
||||
nextAttemptAt: Date | null;
|
||||
userId: number | null;
|
||||
modelVersionId: number | null;
|
||||
modelFileId: number | null;
|
||||
claimedBy: string | null;
|
||||
claimedAt: Date | null;
|
||||
heartbeatAt: Date | null;
|
||||
startedAt: Date | null;
|
||||
completedAt: Date | null;
|
||||
createdAt: Date;
|
||||
updatedAt: Date;
|
||||
}
|
||||
|
||||
export interface Model {
|
||||
id: number;
|
||||
name: string;
|
||||
|
||||
@@ -2129,6 +2129,9 @@ export const REDIS_SYS_KEYS = {
|
||||
DAILY_CHALLENGE: {
|
||||
CONFIG: 'daily-challenge:config',
|
||||
},
|
||||
HUGGING_FACE_IMPORT: {
|
||||
CONFIG: 'hugging-face-import:config',
|
||||
},
|
||||
COLLECTION: {
|
||||
RANDOM_SEED: 'collection:random-seed',
|
||||
},
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
import dynamic from 'next/dynamic';
|
||||
import { dialogStore } from '~/components/Dialog/dialogStore';
|
||||
import type { AddFromImportsModalProps } from '~/components/Moderation/HuggingFaceImport/AddFromImportsModal';
|
||||
|
||||
const AddFromImportsModal = dynamic(
|
||||
() => import('~/components/Moderation/HuggingFaceImport/AddFromImportsModal'),
|
||||
{ ssr: false }
|
||||
);
|
||||
|
||||
export function openAddFromImportsModal(props: AddFromImportsModalProps) {
|
||||
dialogStore.trigger({ component: AddFromImportsModal, props });
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
import { describe, expect, test, vi } from 'vitest';
|
||||
import { page } from 'vitest/browser';
|
||||
import type * as FilesProviderModule from '~/components/Resource/FilesProvider';
|
||||
import { renderWithProviders } from '../../../../test/component-setup';
|
||||
|
||||
const { mockOpen, mockAdoptFiles, user } = vi.hoisted(() => ({
|
||||
mockOpen: vi.fn(),
|
||||
mockAdoptFiles: vi.fn(),
|
||||
user: { current: { isModerator: true } as { isModerator: boolean } | null },
|
||||
}));
|
||||
|
||||
// Pins the dialog-store route; the reason is on AddFromImportsButton.
|
||||
vi.mock('~/components/Dialog/triggers/add-from-hugging-face-imports', () => ({
|
||||
openAddFromImportsModal: mockOpen,
|
||||
}));
|
||||
vi.mock('~/components/Resource/FilesProvider', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof FilesProviderModule>()),
|
||||
useFilesContext: () => ({ adoptFiles: mockAdoptFiles, modelType: 'Checkpoint' }),
|
||||
}));
|
||||
vi.mock('~/hooks/useCurrentUser', () => ({ useCurrentUser: () => user.current }));
|
||||
|
||||
import { AddFromImportsButton } from '~/components/Moderation/HuggingFaceImport/AddFromImportsButton';
|
||||
|
||||
describe('AddFromImportsButton', () => {
|
||||
test('opens the picker through the dialog store, with the provider handles', async () => {
|
||||
user.current = { isModerator: true };
|
||||
renderWithProviders(<AddFromImportsButton modelVersionId={42} />);
|
||||
|
||||
await page.getByRole('button', { name: 'Add from Hugging Face imports' }).click();
|
||||
|
||||
expect(mockOpen).toHaveBeenCalledWith({
|
||||
modelVersionId: 42,
|
||||
modelType: 'Checkpoint',
|
||||
adoptFiles: mockAdoptFiles,
|
||||
});
|
||||
});
|
||||
|
||||
test('renders nothing for a non-moderator', async () => {
|
||||
user.current = { isModerator: false };
|
||||
renderWithProviders(
|
||||
<div data-testid="host">
|
||||
<AddFromImportsButton modelVersionId={42} />
|
||||
</div>
|
||||
);
|
||||
|
||||
await expect.element(page.getByTestId('host')).toBeInTheDocument();
|
||||
expect(page.getByTestId('host').element().childElementCount).toBe(0);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,26 @@
|
||||
import { Button } from '@mantine/core';
|
||||
import { IconCloudDownload } from '@tabler/icons-react';
|
||||
import { openAddFromImportsModal } from '~/components/Dialog/triggers/add-from-hugging-face-imports';
|
||||
import { useFilesContext } from '~/components/Resource/FilesProvider';
|
||||
import { useCurrentUser } from '~/hooks/useCurrentUser';
|
||||
|
||||
/**
|
||||
* 🔴 Through the dialog store, not an inline `<Modal>`: Manage files is a store dialog (z-index
|
||||
* 300+), and a plain Modal renders at Mantine's default 200 — behind it.
|
||||
*/
|
||||
export function AddFromImportsButton({ modelVersionId }: { modelVersionId: number }) {
|
||||
const currentUser = useCurrentUser();
|
||||
const { adoptFiles, modelType } = useFilesContext();
|
||||
if (!currentUser?.isModerator) return null;
|
||||
|
||||
return (
|
||||
<Button
|
||||
size="compact-sm"
|
||||
variant="light"
|
||||
leftSection={<IconCloudDownload size={14} />}
|
||||
onClick={() => openAddFromImportsModal({ modelVersionId, modelType, adoptFiles })}
|
||||
>
|
||||
Add from Hugging Face imports
|
||||
</Button>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
import { beforeEach, describe, expect, test, vi } from 'vitest';
|
||||
import { page } from 'vitest/browser';
|
||||
import { DialogProvider } from '~/components/Dialog/DialogProvider';
|
||||
import { dialogStore, useDialogStore } from '~/components/Dialog/dialogStore';
|
||||
import type * as NotificationsModule from '~/utils/notifications';
|
||||
import type * as TrpcModule from '~/utils/trpc';
|
||||
import { renderWithProviders } from '../../../../test/component-setup';
|
||||
|
||||
/**
|
||||
* Pins what only this file can see: `adoptFiles` receives the ids `attach` CREATED, after the
|
||||
* version query is refreshed. Import ids are also `number[]`, so the types cannot catch a swap.
|
||||
*/
|
||||
|
||||
const calls = vi.hoisted(() => [] as string[]);
|
||||
const listed = vi.hoisted(() => ({ rows: [] as unknown[] }));
|
||||
const { mockAdoptFiles, mockMutateAsync, mockSuccess, mockError } = vi.hoisted(() => ({
|
||||
mockAdoptFiles: vi.fn(),
|
||||
mockMutateAsync: vi.fn(),
|
||||
mockSuccess: vi.fn(),
|
||||
mockError: vi.fn(),
|
||||
}));
|
||||
|
||||
const allImports = [
|
||||
{
|
||||
id: 11,
|
||||
groupName: 'flux-krea',
|
||||
repo: 'black-forest-labs/FLUX.1-Krea-dev',
|
||||
revision: 'aaaaaaa1',
|
||||
filename: 'flux1-krea-dev.safetensors',
|
||||
sizeBytes: 1000,
|
||||
suggestedType: null,
|
||||
createdAt: new Date('2026-09-01'),
|
||||
},
|
||||
{
|
||||
id: 12,
|
||||
groupName: 'flux-krea',
|
||||
repo: 'black-forest-labs/FLUX.1-Krea-dev',
|
||||
revision: 'aaaaaaa1',
|
||||
filename: 'ae.safetensors',
|
||||
sizeBytes: 500,
|
||||
suggestedType: 'VAE',
|
||||
createdAt: new Date('2026-09-01'),
|
||||
},
|
||||
];
|
||||
|
||||
const invalidate = (name: string) => () => {
|
||||
calls.push(`invalidate:${name}`);
|
||||
return Promise.resolve();
|
||||
};
|
||||
|
||||
vi.mock('~/utils/trpc', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof TrpcModule>()),
|
||||
trpc: {
|
||||
useUtils: () => ({
|
||||
modelVersion: { getByIdForEdit: { invalidate: invalidate('getByIdForEdit') } },
|
||||
huggingFaceImport: {
|
||||
getAll: { invalidate: invalidate('getAll') },
|
||||
getCounts: { invalidate: invalidate('getCounts') },
|
||||
},
|
||||
}),
|
||||
huggingFaceImport: {
|
||||
getAll: { useQuery: () => ({ data: listed.rows, isLoading: false }) },
|
||||
attach: { useMutation: () => ({ mutateAsync: mockMutateAsync }) },
|
||||
},
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('~/utils/notifications', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof NotificationsModule>()),
|
||||
showSuccessNotification: mockSuccess,
|
||||
showErrorNotification: mockError,
|
||||
}));
|
||||
|
||||
import AddFromImportsModal from '~/components/Moderation/HuggingFaceImport/AddFromImportsModal';
|
||||
|
||||
/** Opened the way the button opens it — through the store, stacked above whatever is showing. */
|
||||
async function openPicker() {
|
||||
renderWithProviders(<DialogProvider />);
|
||||
dialogStore.trigger({
|
||||
component: AddFromImportsModal,
|
||||
props: { modelVersionId: 42, modelType: 'Checkpoint', adoptFiles: mockAdoptFiles },
|
||||
});
|
||||
await expect.element(page.getByText('Add from Hugging Face imports')).toBeVisible();
|
||||
}
|
||||
|
||||
async function pickType(filename: string, label: string) {
|
||||
const row = page.getByText(filename).element().parentElement as HTMLElement;
|
||||
(row.querySelector('input') as HTMLInputElement).click();
|
||||
await page.getByRole('option', { name: label, exact: true }).click();
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
useDialogStore.getState().closeAll();
|
||||
calls.length = 0;
|
||||
listed.rows = allImports;
|
||||
mockAdoptFiles.mockReset().mockImplementation(async (ids: number[]) => {
|
||||
calls.push(`adopt:${ids.join(',')}`);
|
||||
});
|
||||
mockMutateAsync.mockReset().mockImplementation(async ({ id }: { id: number }) => {
|
||||
calls.push(`attach:${id}`);
|
||||
return { modelFileId: 900 + id };
|
||||
});
|
||||
mockSuccess.mockReset();
|
||||
mockError.mockReset();
|
||||
});
|
||||
|
||||
describe('AddFromImportsModal', () => {
|
||||
test('adopts the CREATED file ids, after the version query is refreshed', async () => {
|
||||
await openPicker();
|
||||
|
||||
await pickType('flux1-krea-dev.safetensors', 'Checkpoint');
|
||||
await pickType('ae.safetensors', 'VAE');
|
||||
await page.getByRole('button', { name: 'Attach 2' }).click();
|
||||
|
||||
await expect.poll(() => mockSuccess.mock.calls.length).toBe(1);
|
||||
expect(mockMutateAsync.mock.calls.map(([input]) => input)).toEqual([
|
||||
{ id: 11, modelVersionId: 42, type: 'Model' },
|
||||
{ id: 12, modelVersionId: 42, type: 'VAE' },
|
||||
]);
|
||||
expect(mockAdoptFiles).toHaveBeenCalledTimes(1);
|
||||
expect(mockAdoptFiles).toHaveBeenCalledWith([911, 912]);
|
||||
// Adopting before the refresh would read the cached version, which lacks the new files.
|
||||
expect(calls.indexOf('adopt:911,912')).toBeGreaterThan(
|
||||
calls.indexOf('invalidate:getByIdForEdit')
|
||||
);
|
||||
});
|
||||
|
||||
test('attaches nothing until a type is picked — no row is pre-selected', async () => {
|
||||
await openPicker();
|
||||
|
||||
// The VAE row carries a suggestion; it must stay a hint, not a selection.
|
||||
await expect.element(page.getByPlaceholder('Suggested: VAE')).toBeVisible();
|
||||
await expect.element(page.getByRole('button', { name: /^Attach/ })).toBeDisabled();
|
||||
});
|
||||
|
||||
test('cannot be clicked again while a batch is running', async () => {
|
||||
let release!: () => void;
|
||||
mockMutateAsync.mockImplementation(
|
||||
() => new Promise((resolve) => (release = () => resolve({ modelFileId: 911 })))
|
||||
);
|
||||
await openPicker();
|
||||
|
||||
await pickType('flux1-krea-dev.safetensors', 'Checkpoint');
|
||||
await page.getByRole('button', { name: 'Attach 1' }).click();
|
||||
await expect.poll(() => mockMutateAsync.mock.calls.length).toBe(1);
|
||||
|
||||
// Mid-flight the button is loading and the rows are frozen, so a second click cannot re-submit.
|
||||
await expect.element(page.getByRole('button', { name: /^Attach/ })).toBeDisabled();
|
||||
expect((page.getByPlaceholder('Pick a file type').element() as HTMLInputElement).disabled).toBe(
|
||||
true
|
||||
);
|
||||
|
||||
release();
|
||||
await expect.poll(() => mockSuccess.mock.calls.length).toBe(1);
|
||||
expect(mockMutateAsync).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test('drops a row from the selection once it leaves the list', async () => {
|
||||
// 11 attaches and drops out of the refreshed list; 12 fails, so the modal stays open with both
|
||||
// still selected in state. Counting 11 would re-attach it on the next click.
|
||||
mockMutateAsync.mockImplementation(async ({ id }: { id: number }) => {
|
||||
if (id === 12) throw new Error('Scan submission failed.');
|
||||
listed.rows = allImports.filter((row) => row.id !== 11);
|
||||
return { modelFileId: 911 };
|
||||
});
|
||||
await openPicker();
|
||||
|
||||
await pickType('flux1-krea-dev.safetensors', 'Checkpoint');
|
||||
await pickType('ae.safetensors', 'VAE');
|
||||
await page.getByRole('button', { name: 'Attach 2' }).click();
|
||||
await expect.poll(() => mockError.mock.calls.length).toBe(1);
|
||||
|
||||
await expect.element(page.getByRole('button', { name: 'Attach 1' })).toBeVisible();
|
||||
});
|
||||
|
||||
test('keeps every attach failure when the refresh also fails', async () => {
|
||||
mockMutateAsync.mockImplementation(async ({ id }: { id: number }) => {
|
||||
if (id === 12) throw new Error('Created model file 555, but this import was attached first.');
|
||||
return { modelFileId: 900 + id };
|
||||
});
|
||||
mockAdoptFiles.mockRejectedValue(new Error('refresh failed'));
|
||||
await openPicker();
|
||||
|
||||
await pickType('flux1-krea-dev.safetensors', 'Checkpoint');
|
||||
await pickType('ae.safetensors', 'VAE');
|
||||
await page.getByRole('button', { name: 'Attach 2' }).click();
|
||||
|
||||
await expect.poll(() => mockError.mock.calls.length).toBe(1);
|
||||
const messages = (mockError.mock.calls[0][0].error as { message: string }[]).map(
|
||||
(error) => error.message
|
||||
);
|
||||
// The lost-claim message is the only place file 555 is ever named.
|
||||
expect(messages).toContain('Created model file 555, but this import was attached first.');
|
||||
expect(messages).toContain('refresh failed');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,185 @@
|
||||
import { Alert, Badge, Button, Group, Modal, Select, Stack, Text, TextInput } from '@mantine/core';
|
||||
import { useDebouncedValue } from '@mantine/hooks';
|
||||
import { useState } from 'react';
|
||||
import { useDialogContext } from '~/components/Dialog/DialogProvider';
|
||||
import { attachImports } from '~/components/Moderation/HuggingFaceImport/attach-imports';
|
||||
import { byGroup } from '~/components/Moderation/HuggingFaceImport/utils';
|
||||
import type { ModelFileType } from '~/server/common/constants';
|
||||
import type { ModelType } from '~/shared/utils/prisma/enums';
|
||||
import { getModelFileTypeOptions } from '~/utils/file-display-helpers';
|
||||
import { formatBytes } from '~/utils/number-helpers';
|
||||
import { showErrorNotification, showSuccessNotification } from '~/utils/notifications';
|
||||
import { trpc } from '~/utils/trpc';
|
||||
|
||||
export type AddFromImportsModalProps = {
|
||||
modelVersionId: number;
|
||||
modelType?: ModelType | null;
|
||||
/** From the `FilesProvider` that opened this — a dialog mounts outside that subtree. */
|
||||
adoptFiles: (modelFileIds: number[]) => Promise<void>;
|
||||
};
|
||||
|
||||
export default function AddFromImportsModal({
|
||||
modelVersionId,
|
||||
modelType,
|
||||
adoptFiles,
|
||||
}: AddFromImportsModalProps) {
|
||||
const dialog = useDialogContext();
|
||||
const [filter, setFilter] = useState('');
|
||||
const [debouncedFilter] = useDebouncedValue(filter, 300);
|
||||
const [selection, setSelection] = useState<Record<number, ModelFileType | null>>({});
|
||||
const [attaching, setAttaching] = useState(false);
|
||||
const queryUtils = trpc.useUtils();
|
||||
|
||||
const { data = [], isLoading } = trpc.huggingFaceImport.getAll.useQuery({
|
||||
limit: 200,
|
||||
unattached: true,
|
||||
groupName: debouncedFilter.trim() || undefined,
|
||||
});
|
||||
|
||||
const attach = trpc.huggingFaceImport.attach.useMutation();
|
||||
|
||||
const listed = new Set(data.map((row) => row.id));
|
||||
// Only rows still on screen: an import attached in the last batch drops out of the list but
|
||||
// would otherwise stay selected, and the next click would attach it a second time.
|
||||
const chosen = Object.entries(selection)
|
||||
.map(([id, type]) => ({ id: Number(id), type }))
|
||||
.filter(
|
||||
(item): item is { id: number; type: ModelFileType } => !!item.type && listed.has(item.id)
|
||||
);
|
||||
|
||||
const onAttach = async () => {
|
||||
setAttaching(true);
|
||||
const { modelFileIds, failures } = await attachImports({
|
||||
chosen,
|
||||
modelVersionId,
|
||||
attachOne: (input) => attach.mutateAsync(input),
|
||||
});
|
||||
try {
|
||||
await queryUtils.modelVersion.getByIdForEdit.invalidate({
|
||||
id: modelVersionId,
|
||||
withFiles: true,
|
||||
});
|
||||
await Promise.all([
|
||||
queryUtils.huggingFaceImport.getAll.invalidate(),
|
||||
queryUtils.huggingFaceImport.getCounts.invalidate(),
|
||||
]);
|
||||
await adoptFiles(modelFileIds);
|
||||
|
||||
if (failures.length) {
|
||||
// No auto-close: a lost import claim leaves a created file whose id appears only in this message.
|
||||
showErrorNotification({
|
||||
title: `Attached ${modelFileIds.length} of ${chosen.length}`,
|
||||
error: failures.map((message) => ({ message })),
|
||||
autoClose: false,
|
||||
});
|
||||
return;
|
||||
}
|
||||
showSuccessNotification({
|
||||
title: 'Attached',
|
||||
message: `${modelFileIds.length} file(s) added. Scanning starts on its own.`,
|
||||
});
|
||||
dialog.onClose();
|
||||
} catch (error) {
|
||||
// The attaches may have succeeded; silence here invites a second click.
|
||||
showErrorNotification({
|
||||
title: `Attached ${modelFileIds.length} of ${chosen.length}, but the file list could not be refreshed`,
|
||||
error: [
|
||||
error instanceof Error ? error : { message: String(error) },
|
||||
...failures.map((message) => ({ message })),
|
||||
],
|
||||
autoClose: false,
|
||||
});
|
||||
} finally {
|
||||
setAttaching(false);
|
||||
}
|
||||
};
|
||||
|
||||
const groups = byGroup(data);
|
||||
|
||||
return (
|
||||
<Modal {...dialog} title="Add from Hugging Face imports" size="lg" centered>
|
||||
<Stack gap="md">
|
||||
<TextInput
|
||||
placeholder="Filter by group name"
|
||||
value={filter}
|
||||
onChange={(event) => setFilter(event.currentTarget.value)}
|
||||
/>
|
||||
|
||||
{!groups.length ? (
|
||||
<Text c="dimmed" size="sm">
|
||||
{isLoading
|
||||
? 'Loading…'
|
||||
: debouncedFilter.trim()
|
||||
? `Nothing unattached matches "${debouncedFilter.trim()}".`
|
||||
: 'Nothing unattached — every transferred file is already on a version.'}
|
||||
</Text>
|
||||
) : (
|
||||
groups.map((group) => (
|
||||
<Stack key={`${group.groupName}:${group.repo}:${group.revision}`} gap={6}>
|
||||
<Group gap="xs" wrap="wrap">
|
||||
<Text fw={600} size="sm">
|
||||
{group.groupName}
|
||||
</Text>
|
||||
<Badge size="xs" variant="light" color="gray">
|
||||
{group.repo}
|
||||
</Badge>
|
||||
<Badge size="xs" variant="light">
|
||||
{group.revision.slice(0, 7)}
|
||||
</Badge>
|
||||
</Group>
|
||||
|
||||
{group.items.map((row) => {
|
||||
const options = getModelFileTypeOptions(row.filename, { modelType });
|
||||
const type = selection[row.id] ?? null;
|
||||
const suggested = options.find((option) => option.value === row.suggestedType);
|
||||
return (
|
||||
<Group key={row.id} gap="xs" pl="sm" wrap="nowrap" align="center">
|
||||
<Text size="xs" ff="monospace" lineClamp={1} style={{ flex: 1 }}>
|
||||
{row.filename}
|
||||
</Text>
|
||||
<Text size="xs" c="dimmed">
|
||||
{row.sizeBytes ? formatBytes(row.sizeBytes) : '—'}
|
||||
</Text>
|
||||
{/* The type IS the selection, so a suggestion is only ever a hint: pre-filling
|
||||
it would attach every suggested file on the next click. */}
|
||||
<Select
|
||||
size="xs"
|
||||
w={190}
|
||||
clearable
|
||||
data={options}
|
||||
disabled={attaching}
|
||||
placeholder={suggested ? `Suggested: ${suggested.label}` : 'Pick a file type'}
|
||||
value={type}
|
||||
onChange={(value) =>
|
||||
setSelection((prev) => ({
|
||||
...prev,
|
||||
[row.id]: (value as ModelFileType | null) ?? null,
|
||||
}))
|
||||
}
|
||||
/>
|
||||
</Group>
|
||||
);
|
||||
})}
|
||||
</Stack>
|
||||
))
|
||||
)}
|
||||
|
||||
<Alert color="gray">
|
||||
<Text size="xs">
|
||||
A file is attached only once you give it a type. Nothing is pre-selected, and weights
|
||||
get no suggestion — the type decides whether this version loads.
|
||||
</Text>
|
||||
</Alert>
|
||||
|
||||
<Group justify="flex-end">
|
||||
<Button variant="default" onClick={dialog.onClose}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button loading={attaching} disabled={!chosen.length} onClick={onAttach}>
|
||||
Attach {chosen.length || ''}
|
||||
</Button>
|
||||
</Group>
|
||||
</Stack>
|
||||
</Modal>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
import { Anchor, Button, Group, NumberInput, Popover, Select, Stack, Text } from '@mantine/core';
|
||||
import { useState } from 'react';
|
||||
import type { ModelFileType } from '~/server/common/constants';
|
||||
import { getModelFileTypeOptions } from '~/utils/file-display-helpers';
|
||||
import { showErrorNotification, showSuccessNotification } from '~/utils/notifications';
|
||||
import { trpc } from '~/utils/trpc';
|
||||
|
||||
export function AttachControl({
|
||||
importId,
|
||||
filename,
|
||||
suggestedType,
|
||||
modelFileId,
|
||||
modelVersionId,
|
||||
}: {
|
||||
importId: number;
|
||||
filename: string;
|
||||
suggestedType: string | null;
|
||||
modelFileId: number | null;
|
||||
modelVersionId: number | null;
|
||||
}) {
|
||||
const [opened, setOpened] = useState(false);
|
||||
const [versionId, setVersionId] = useState<number | ''>('');
|
||||
const options = getModelFileTypeOptions(filename);
|
||||
// No fallback to the list's first entry: `suggestFileType` leaves primary weights unsuggested on
|
||||
// purpose, and that label decides whether the version loads.
|
||||
const [type, setType] = useState<ModelFileType | null>(
|
||||
options.some((option) => option.value === suggestedType)
|
||||
? (suggestedType as ModelFileType)
|
||||
: null
|
||||
);
|
||||
const queryUtils = trpc.useUtils();
|
||||
|
||||
const attach = trpc.huggingFaceImport.attach.useMutation({
|
||||
onSuccess: async (result) => {
|
||||
showSuccessNotification({
|
||||
title: 'Attached',
|
||||
message: `File ${result.modelFileId} added to version ${result.modelVersionId}. The scan starts on its own.`,
|
||||
});
|
||||
setOpened(false);
|
||||
await queryUtils.huggingFaceImport.getAll.invalidate();
|
||||
},
|
||||
onError: (error) =>
|
||||
showErrorNotification({ title: 'Could not attach', error: new Error(error.message) }),
|
||||
});
|
||||
|
||||
if (modelFileId)
|
||||
return (
|
||||
<Anchor
|
||||
size="xs"
|
||||
target="_blank"
|
||||
href={`/models/v/${modelVersionId}`}
|
||||
title={`Model file ${modelFileId}`}
|
||||
>
|
||||
version {modelVersionId}
|
||||
</Anchor>
|
||||
);
|
||||
|
||||
return (
|
||||
<Popover
|
||||
opened={opened}
|
||||
onChange={setOpened}
|
||||
position="bottom-end"
|
||||
shadow="md"
|
||||
// The theme defaults Popover to withinPortal: false, and this renders inside a Card and a
|
||||
// scroll container — both of which clip it.
|
||||
withinPortal
|
||||
>
|
||||
<Popover.Target>
|
||||
<Button size="compact-xs" variant="light" onClick={() => setOpened((o) => !o)}>
|
||||
Attach
|
||||
</Button>
|
||||
</Popover.Target>
|
||||
<Popover.Dropdown>
|
||||
<Stack gap="xs" w={240}>
|
||||
<Text size="xs" c="dimmed">
|
||||
Creates a model file on that version. Scanning and hashing follow automatically.
|
||||
</Text>
|
||||
<NumberInput
|
||||
size="xs"
|
||||
label="Model version ID"
|
||||
value={versionId}
|
||||
onChange={(value) => setVersionId(typeof value === 'number' ? value : '')}
|
||||
min={1}
|
||||
allowDecimal={false}
|
||||
hideControls
|
||||
/>
|
||||
<Select
|
||||
size="xs"
|
||||
label="File type"
|
||||
placeholder="Pick a file type"
|
||||
data={options}
|
||||
value={type}
|
||||
onChange={(value) => setType(value as ModelFileType | null)}
|
||||
/>
|
||||
<Group justify="flex-end">
|
||||
<Button
|
||||
size="compact-xs"
|
||||
loading={attach.isPending}
|
||||
disabled={!versionId || !type}
|
||||
onClick={() =>
|
||||
versionId &&
|
||||
type &&
|
||||
attach.mutate({ id: importId, modelVersionId: versionId, type })
|
||||
}
|
||||
>
|
||||
Attach
|
||||
</Button>
|
||||
</Group>
|
||||
</Stack>
|
||||
</Popover.Dropdown>
|
||||
</Popover>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
import { Alert, Button, Card, Group, NumberInput, Stack, Switch, Text, Title } from '@mantine/core';
|
||||
import { useEffect, useState } from 'react';
|
||||
import { formatBytes } from '~/utils/number-helpers';
|
||||
import { showErrorNotification, showSuccessNotification } from '~/utils/notifications';
|
||||
import { trpc } from '~/utils/trpc';
|
||||
|
||||
/** Mirrors `PART_SIZE_BYTES`. Only used to show the operator what their numbers cost. */
|
||||
const PART_SIZE = 16 * 1024 * 1024;
|
||||
/**
|
||||
* Measured: RSS grows roughly twice the retained payload, because `arrayBuffer()` leaves undici's
|
||||
* concat buffer alive off-heap where it barely pressures GC. Showing the retained figure alone would
|
||||
* understate what a pod needs by half.
|
||||
*/
|
||||
const RESIDENT_MULTIPLIER = 2;
|
||||
|
||||
export function ImportConfigSection() {
|
||||
const queryUtils = trpc.useUtils();
|
||||
const { data: config } = trpc.huggingFaceImport.getConfig.useQuery();
|
||||
|
||||
const [draft, setDraft] = useState<{
|
||||
enabled: boolean;
|
||||
filesInParallel: number;
|
||||
partsInFlight: number;
|
||||
workBudgetSeconds: number;
|
||||
} | null>(null);
|
||||
|
||||
// Seeds ONCE, deliberately: re-seeding on every `config` change would discard edits in progress
|
||||
// the moment a background refetch landed. The cost is that a change made elsewhere is not picked
|
||||
// up until reload — acceptable for a panel one moderator opens at a time.
|
||||
useEffect(() => {
|
||||
if (config && !draft) setDraft({ ...config });
|
||||
}, [config, draft]);
|
||||
|
||||
const save = trpc.huggingFaceImport.setConfig.useMutation({
|
||||
onSuccess: async (saved) => {
|
||||
setDraft({ ...saved });
|
||||
showSuccessNotification({ title: 'Saved', message: 'Applies from the next job tick.' });
|
||||
await queryUtils.huggingFaceImport.getConfig.invalidate();
|
||||
},
|
||||
onError: (error) =>
|
||||
showErrorNotification({ title: 'Could not save', error: new Error(error.message) }),
|
||||
});
|
||||
|
||||
// Rendering nothing while the query is in flight left a silent gap where the panel belongs, which
|
||||
// reads as "this page has no settings" rather than "not loaded yet".
|
||||
if (!draft)
|
||||
return (
|
||||
<Card withBorder padding="lg">
|
||||
<Stack gap="xs">
|
||||
<Title order={4}>Transfer settings</Title>
|
||||
<Text c="dimmed" size="sm">
|
||||
Loading…
|
||||
</Text>
|
||||
</Stack>
|
||||
</Card>
|
||||
);
|
||||
|
||||
const residentBytes =
|
||||
draft.filesInParallel * draft.partsInFlight * PART_SIZE * RESIDENT_MULTIPLIER;
|
||||
const dirty =
|
||||
!!config &&
|
||||
(['enabled', 'filesInParallel', 'partsInFlight', 'workBudgetSeconds'] as const).some(
|
||||
(key) => draft[key] !== config[key]
|
||||
);
|
||||
|
||||
return (
|
||||
<Card withBorder padding="lg">
|
||||
<Stack gap="md">
|
||||
<Group justify="space-between" align="flex-start">
|
||||
<Stack gap={2}>
|
||||
<Title order={4}>Transfer settings</Title>
|
||||
<Text c="dimmed" size="sm">
|
||||
Applies from the next job tick. Turning transfers off leaves queued files untouched.
|
||||
</Text>
|
||||
</Stack>
|
||||
<Switch
|
||||
checked={draft.enabled}
|
||||
onChange={(event) => setDraft({ ...draft, enabled: event.currentTarget.checked })}
|
||||
label={draft.enabled ? 'Transfers on' : 'Transfers off'}
|
||||
/>
|
||||
</Group>
|
||||
|
||||
<Group grow align="flex-start">
|
||||
<NumberInput
|
||||
label="Files at once"
|
||||
description="Across the whole fleet"
|
||||
min={1}
|
||||
max={4}
|
||||
value={draft.filesInParallel}
|
||||
onChange={(value) =>
|
||||
setDraft({ ...draft, filesInParallel: typeof value === 'number' ? value : 1 })
|
||||
}
|
||||
/>
|
||||
<NumberInput
|
||||
label="Parts per file"
|
||||
description="Concurrent 16 MB reads"
|
||||
min={1}
|
||||
max={6}
|
||||
value={draft.partsInFlight}
|
||||
onChange={(value) =>
|
||||
setDraft({ ...draft, partsInFlight: typeof value === 'number' ? value : 1 })
|
||||
}
|
||||
/>
|
||||
<NumberInput
|
||||
label="Seconds per run"
|
||||
description="Must stay inside the 5-minute lock"
|
||||
min={15}
|
||||
max={240}
|
||||
value={draft.workBudgetSeconds}
|
||||
onChange={(value) =>
|
||||
setDraft({ ...draft, workBudgetSeconds: typeof value === 'number' ? value : 120 })
|
||||
}
|
||||
/>
|
||||
</Group>
|
||||
|
||||
{/* The number these three knobs actually buy. Without it they read as speed dials, and the
|
||||
one that matters is memory on a pod that is also serving traffic. */}
|
||||
<Alert color={residentBytes > 400 * 1024 * 1024 ? 'orange' : 'gray'}>
|
||||
<Text size="sm">
|
||||
Roughly <strong>{formatBytes(residentBytes)}</strong> resident on the pod running the
|
||||
transfer — {draft.filesInParallel} × {draft.partsInFlight} × 16 MB, doubled,
|
||||
because the fetch leaves a copy alive off-heap.
|
||||
</Text>
|
||||
</Alert>
|
||||
|
||||
<Group justify="flex-end">
|
||||
<Button
|
||||
variant="default"
|
||||
disabled={!dirty}
|
||||
onClick={() => config && setDraft({ ...config })}
|
||||
>
|
||||
Reset
|
||||
</Button>
|
||||
<Button loading={save.isPending} disabled={!dirty} onClick={() => save.mutate(draft)}>
|
||||
Save
|
||||
</Button>
|
||||
</Group>
|
||||
</Stack>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,320 @@
|
||||
import { useState } from 'react';
|
||||
import { useDebouncedValue } from '@mantine/hooks';
|
||||
import {
|
||||
ActionIcon,
|
||||
Anchor,
|
||||
Badge,
|
||||
Card,
|
||||
Group,
|
||||
Progress,
|
||||
Stack,
|
||||
SegmentedControl,
|
||||
Table,
|
||||
Text,
|
||||
TextInput,
|
||||
Title,
|
||||
Tooltip,
|
||||
} from '@mantine/core';
|
||||
import { IconPlayerStop, IconRefresh, IconTrash, IconUnlink } from '@tabler/icons-react';
|
||||
import { openConfirmModal } from '@mantine/modals';
|
||||
import { CopyButton } from '~/components/CopyButton/CopyButton';
|
||||
import { UnattachedSection } from '~/components/Moderation/HuggingFaceImport/UnattachedSection';
|
||||
import { AttachControl } from '~/components/Moderation/HuggingFaceImport/AttachControl';
|
||||
import { formatBytes } from '~/utils/number-helpers';
|
||||
import { showErrorNotification } from '~/utils/notifications';
|
||||
import { trpc } from '~/utils/trpc';
|
||||
|
||||
const STATUS_COLOR: Record<string, string> = {
|
||||
Queued: 'gray',
|
||||
Transferring: 'blue',
|
||||
Completed: 'teal',
|
||||
Failed: 'red',
|
||||
Canceled: 'orange',
|
||||
};
|
||||
|
||||
const ACTIVE = new Set(['Queued', 'Transferring']);
|
||||
|
||||
export function ImportQueueTable() {
|
||||
const queryUtils = trpc.useUtils();
|
||||
const [tab, setTab] = useState<'all' | 'unattached'>('all');
|
||||
const [groupFilter, setGroupFilter] = useState('');
|
||||
// Debounced so a keystroke is not a query; the filter is server-side because a client-side one
|
||||
// over a capped page silently stops finding older groups.
|
||||
const [debouncedFilter] = useDebouncedValue(groupFilter, 300);
|
||||
// Counts come from their own query: a page of rows cannot say how many exist outside it, which
|
||||
// is what the client-side filter got wrong. Filtered the same way the rows are, so a tab label
|
||||
// never counts a population the list beneath it is not showing.
|
||||
const { data: counts } = trpc.huggingFaceImport.getCounts.useQuery({
|
||||
groupName: debouncedFilter.trim() || undefined,
|
||||
});
|
||||
const { data = [], isLoading: isPending } = trpc.huggingFaceImport.getAll.useQuery(
|
||||
{ limit: 100, groupName: debouncedFilter.trim() || undefined },
|
||||
{
|
||||
// A transfer advances a part at a time on a cron; polling is how the bar moves without a
|
||||
// websocket, and it stops as soon as nothing is in flight.
|
||||
refetchInterval: (query) =>
|
||||
query.state.data?.some((row) => ACTIVE.has(row.status)) ? 5000 : false,
|
||||
}
|
||||
);
|
||||
|
||||
const onError = (error: { message: string }) =>
|
||||
showErrorNotification({ title: 'Action failed', error: new Error(error.message) });
|
||||
|
||||
// `ok: false` means the row moved on before the click landed — a cancel on a row that just
|
||||
// completed, say. Invalidating and saying nothing renders a refusal as a successful no-op.
|
||||
const onSettled = (result: { ok: boolean } | undefined, action: string) => {
|
||||
if (result && !result.ok)
|
||||
showErrorNotification({
|
||||
title: `Could not ${action}`,
|
||||
error: new Error('The import is no longer in a state where that applies. Refreshed.'),
|
||||
});
|
||||
return Promise.all([
|
||||
queryUtils.huggingFaceImport.getAll.invalidate(),
|
||||
queryUtils.huggingFaceImport.getCounts.invalidate(),
|
||||
]);
|
||||
};
|
||||
|
||||
// Wires the path `buildAttachInput`'s refusal recommends. Until this existed, "detach or delete
|
||||
// that file first" named something no moderator could do.
|
||||
const detach = trpc.huggingFaceImport.detach.useMutation({
|
||||
onError,
|
||||
onSuccess: (result) => onSettled(result, 'detach'),
|
||||
});
|
||||
// A Failed row can still hold an uploadId, and its parts are billed until something aborts them.
|
||||
// Retry was the only exit, so abandoning a transfer meant paying for it indefinitely.
|
||||
const remove = trpc.huggingFaceImport.delete.useMutation({
|
||||
onError,
|
||||
onSuccess: (result) => onSettled(result, 'delete'),
|
||||
});
|
||||
const retry = trpc.huggingFaceImport.retry.useMutation({
|
||||
onError,
|
||||
onSuccess: (result) => onSettled(result, 'retry'),
|
||||
});
|
||||
const cancel = trpc.huggingFaceImport.cancel.useMutation({
|
||||
onError,
|
||||
onSuccess: (result) => onSettled(result, 'cancel'),
|
||||
});
|
||||
|
||||
return (
|
||||
<Card withBorder padding="lg">
|
||||
<Stack gap="md">
|
||||
<Group justify="space-between" align="center" wrap="nowrap">
|
||||
<Group gap="lg" wrap="nowrap">
|
||||
<Title order={4}>Imports</Title>
|
||||
<SegmentedControl
|
||||
size="xs"
|
||||
value={tab}
|
||||
onChange={(value) => setTab(value as 'all' | 'unattached')}
|
||||
data={[
|
||||
{ value: 'all', label: `All${counts ? ` (${counts.total})` : ''}` },
|
||||
{
|
||||
value: 'unattached',
|
||||
label: `Unattached${counts ? ` (${counts.unattached})` : ''}`,
|
||||
},
|
||||
]}
|
||||
/>
|
||||
</Group>
|
||||
<TextInput
|
||||
size="xs"
|
||||
w={260}
|
||||
placeholder="Filter by group name…"
|
||||
value={groupFilter}
|
||||
onChange={(event) => setGroupFilter(event.currentTarget.value)}
|
||||
/>
|
||||
</Group>
|
||||
|
||||
{tab === 'unattached' && <UnattachedSection filter={debouncedFilter} />}
|
||||
|
||||
{tab === 'unattached' ? null : !data.length ? (
|
||||
<Text c="dimmed" size="sm">
|
||||
{isPending
|
||||
? 'Loading…'
|
||||
: debouncedFilter.trim()
|
||||
? `No imports match "${debouncedFilter.trim()}".`
|
||||
: 'Nothing imported yet.'}
|
||||
</Text>
|
||||
) : (
|
||||
<Table.ScrollContainer minWidth={900}>
|
||||
<Table verticalSpacing="xs" fz="sm">
|
||||
<Table.Thead>
|
||||
<Table.Tr>
|
||||
<Table.Th>Source</Table.Th>
|
||||
<Table.Th>Size</Table.Th>
|
||||
<Table.Th>Progress</Table.Th>
|
||||
<Table.Th>Uploaded file</Table.Th>
|
||||
<Table.Th>Attached to</Table.Th>
|
||||
<Table.Th />
|
||||
</Table.Tr>
|
||||
</Table.Thead>
|
||||
<Table.Tbody>
|
||||
{data.map((row) => {
|
||||
const pct = row.sizeBytes
|
||||
? Math.min(100, Math.round((row.bytesTransferred / row.sizeBytes) * 100))
|
||||
: 0;
|
||||
return (
|
||||
<Table.Tr key={row.id}>
|
||||
<Table.Td>
|
||||
<Stack gap={2}>
|
||||
<Anchor
|
||||
size="sm"
|
||||
target="_blank"
|
||||
href={`https://huggingface.co/${row.repo}/blob/${row.revision}/${row.filename}`}
|
||||
>
|
||||
{row.repo} / {row.filename}
|
||||
</Anchor>
|
||||
<Group gap={6}>
|
||||
<Badge size="xs" color={STATUS_COLOR[row.status] ?? 'gray'}>
|
||||
{row.status}
|
||||
</Badge>
|
||||
<Text size="xs" c="dimmed">
|
||||
{row.revision.slice(0, 7)}
|
||||
</Text>
|
||||
</Group>
|
||||
{row.error && (
|
||||
<Text size="xs" c="red.4" lineClamp={2}>
|
||||
{row.error}
|
||||
</Text>
|
||||
)}
|
||||
</Stack>
|
||||
</Table.Td>
|
||||
<Table.Td>
|
||||
<Text size="xs">{row.sizeBytes ? formatBytes(row.sizeBytes) : '—'}</Text>
|
||||
</Table.Td>
|
||||
<Table.Td miw={160}>
|
||||
<Stack gap={2}>
|
||||
<Progress
|
||||
value={row.status === 'Completed' ? 100 : pct}
|
||||
color={STATUS_COLOR[row.status] ?? 'gray'}
|
||||
/>
|
||||
<Text size="xs" c="dimmed">
|
||||
{formatBytes(row.bytesTransferred)}
|
||||
</Text>
|
||||
</Stack>
|
||||
</Table.Td>
|
||||
<Table.Td>
|
||||
{row.url ? (
|
||||
<Group gap={4} wrap="nowrap">
|
||||
<Text size="xs" lineClamp={1} className="max-w-[260px]">
|
||||
{row.url}
|
||||
</Text>
|
||||
<CopyButton value={row.url}>
|
||||
{({ copied, copy, Icon, color }) => (
|
||||
<Tooltip label={copied ? 'Copied' : 'Copy file URL'}>
|
||||
<ActionIcon
|
||||
variant="subtle"
|
||||
size="sm"
|
||||
color={color}
|
||||
onClick={copy}
|
||||
>
|
||||
<Icon size={14} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</CopyButton>
|
||||
</Group>
|
||||
) : (
|
||||
<Text size="xs" c="dimmed">
|
||||
—
|
||||
</Text>
|
||||
)}
|
||||
</Table.Td>
|
||||
<Table.Td>
|
||||
{row.status === 'Completed' ? (
|
||||
<Group gap={4} wrap="nowrap">
|
||||
<AttachControl
|
||||
importId={row.id}
|
||||
filename={row.filename}
|
||||
suggestedType={row.suggestedType}
|
||||
modelFileId={row.modelFileId}
|
||||
modelVersionId={row.modelVersionId}
|
||||
/>
|
||||
{row.modelFileId && (
|
||||
<Tooltip label="Detach — leaves the model file in place">
|
||||
<ActionIcon
|
||||
variant="subtle"
|
||||
size="sm"
|
||||
loading={detach.isPending && detach.variables?.id === row.id}
|
||||
onClick={() => detach.mutate({ id: row.id })}
|
||||
>
|
||||
<IconUnlink size={14} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
</Group>
|
||||
) : (
|
||||
<Text size="xs" c="dimmed">
|
||||
—
|
||||
</Text>
|
||||
)}
|
||||
</Table.Td>
|
||||
<Table.Td>
|
||||
<Group gap={4} wrap="nowrap" justify="flex-end">
|
||||
{ACTIVE.has(row.status) && (
|
||||
<Tooltip label="Cancel">
|
||||
<ActionIcon
|
||||
variant="subtle"
|
||||
color="orange"
|
||||
size="sm"
|
||||
loading={cancel.isPending && cancel.variables?.id === row.id}
|
||||
onClick={() => cancel.mutate({ id: row.id })}
|
||||
>
|
||||
<IconPlayerStop size={14} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
)}
|
||||
{(row.status === 'Failed' || row.status === 'Canceled') && (
|
||||
<>
|
||||
<Tooltip label="Retry from the start">
|
||||
<ActionIcon
|
||||
variant="subtle"
|
||||
size="sm"
|
||||
loading={retry.isPending && retry.variables?.id === row.id}
|
||||
onClick={() => retry.mutate({ id: row.id })}
|
||||
>
|
||||
<IconRefresh size={14} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
<Tooltip label="Delete — frees any parts already uploaded">
|
||||
<ActionIcon
|
||||
variant="subtle"
|
||||
color="red"
|
||||
size="sm"
|
||||
loading={remove.isPending && remove.variables?.id === row.id}
|
||||
onClick={() =>
|
||||
openConfirmModal({
|
||||
title: 'Delete this import?',
|
||||
centered: true,
|
||||
labels: { confirm: 'Delete', cancel: 'Cancel' },
|
||||
confirmProps: { color: 'red' },
|
||||
children: (
|
||||
<Text size="sm">
|
||||
Aborts the upload of{' '}
|
||||
<Text span ff="monospace" size="sm">
|
||||
{row.filename}
|
||||
</Text>{' '}
|
||||
and removes anything already stored for it. Re-importing
|
||||
means transferring it again from {row.repo}.
|
||||
</Text>
|
||||
),
|
||||
onConfirm: () => remove.mutate({ id: row.id }),
|
||||
})
|
||||
}
|
||||
>
|
||||
<IconTrash size={14} />
|
||||
</ActionIcon>
|
||||
</Tooltip>
|
||||
</>
|
||||
)}
|
||||
</Group>
|
||||
</Table.Td>
|
||||
</Table.Tr>
|
||||
);
|
||||
})}
|
||||
</Table.Tbody>
|
||||
</Table>
|
||||
</Table.ScrollContainer>
|
||||
)}
|
||||
</Stack>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
import { beforeEach, describe, expect, test, vi } from 'vitest';
|
||||
import { page, userEvent } from 'vitest/browser';
|
||||
import type * as NotificationsModule from '~/utils/notifications';
|
||||
import type * as TrpcModule from '~/utils/trpc';
|
||||
import { renderWithProviders } from '../../../../test/component-setup';
|
||||
|
||||
const { mockMutate } = vi.hoisted(() => ({ mockMutate: vi.fn() }));
|
||||
|
||||
vi.mock('~/utils/trpc', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof TrpcModule>()),
|
||||
trpc: {
|
||||
useUtils: () => ({
|
||||
huggingFaceImport: {
|
||||
getAll: { invalidate: vi.fn() },
|
||||
getCounts: { invalidate: vi.fn() },
|
||||
},
|
||||
}),
|
||||
huggingFaceImport: {
|
||||
renameGroup: { useMutation: () => ({ mutate: mockMutate, isPending: false }) },
|
||||
},
|
||||
},
|
||||
}));
|
||||
vi.mock('~/utils/notifications', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof NotificationsModule>()),
|
||||
showSuccessNotification: vi.fn(),
|
||||
showErrorNotification: vi.fn(),
|
||||
}));
|
||||
|
||||
import { RenameGroupControl } from '~/components/Moderation/HuggingFaceImport/RenameGroupControl';
|
||||
|
||||
async function openRename() {
|
||||
renderWithProviders(
|
||||
<RenameGroupControl repo="owner/name" revision="abc123" groupName="flux-krea" />
|
||||
);
|
||||
await page.getByRole('button', { name: 'Rename group' }).click();
|
||||
const input = page.getByLabelText('Group name');
|
||||
await expect.element(input).toHaveValue('flux-krea');
|
||||
return input;
|
||||
}
|
||||
|
||||
beforeEach(() => mockMutate.mockReset());
|
||||
|
||||
describe('RenameGroupControl', () => {
|
||||
test('renames the group it was opened on, from its current name', async () => {
|
||||
const input = await openRename();
|
||||
await input.fill(' FLUX Krea ');
|
||||
await page.getByRole('button', { name: 'Rename', exact: true }).click();
|
||||
|
||||
expect(mockMutate).toHaveBeenCalledTimes(1);
|
||||
expect(mockMutate).toHaveBeenCalledWith({
|
||||
repo: 'owner/name',
|
||||
revision: 'abc123',
|
||||
from: 'flux-krea',
|
||||
groupName: 'FLUX Krea',
|
||||
});
|
||||
});
|
||||
|
||||
test('submits on Enter', async () => {
|
||||
const input = await openRename();
|
||||
await input.fill('FLUX Krea');
|
||||
await userEvent.keyboard('{Enter}');
|
||||
|
||||
expect(mockMutate).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test('will not submit an empty or unchanged name', async () => {
|
||||
const input = await openRename();
|
||||
const save = page.getByRole('button', { name: 'Rename', exact: true });
|
||||
|
||||
await expect.element(save).toBeDisabled();
|
||||
await input.fill(' ');
|
||||
await expect.element(save).toBeDisabled();
|
||||
await userEvent.keyboard('{Enter}');
|
||||
|
||||
expect(mockMutate).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,90 @@
|
||||
import { ActionIcon, Button, Group, Popover, Stack, TextInput } from '@mantine/core';
|
||||
import { IconPencil } from '@tabler/icons-react';
|
||||
import { useState } from 'react';
|
||||
import { showErrorNotification, showSuccessNotification } from '~/utils/notifications';
|
||||
import { trpc } from '~/utils/trpc';
|
||||
|
||||
export function RenameGroupControl({
|
||||
repo,
|
||||
revision,
|
||||
groupName,
|
||||
}: {
|
||||
repo: string;
|
||||
revision: string;
|
||||
groupName: string;
|
||||
}) {
|
||||
const [opened, setOpened] = useState(false);
|
||||
const [draft, setDraft] = useState(groupName);
|
||||
const queryUtils = trpc.useUtils();
|
||||
|
||||
const rename = trpc.huggingFaceImport.renameGroup.useMutation({
|
||||
onSuccess: async (result) => {
|
||||
setOpened(false);
|
||||
showSuccessNotification({
|
||||
title: 'Renamed',
|
||||
message: `${result.renamed} file(s) are now in "${result.groupName}".`,
|
||||
});
|
||||
await queryUtils.huggingFaceImport.getAll.invalidate();
|
||||
await queryUtils.huggingFaceImport.getCounts.invalidate();
|
||||
},
|
||||
onError: (error) =>
|
||||
showErrorNotification({ title: 'Could not rename', error: new Error(error.message) }),
|
||||
});
|
||||
|
||||
const name = draft.trim();
|
||||
const submit = () => {
|
||||
if (!name || name === groupName) return;
|
||||
rename.mutate({ repo, revision, from: groupName, groupName: name });
|
||||
};
|
||||
|
||||
return (
|
||||
<Popover
|
||||
opened={opened}
|
||||
onChange={setOpened}
|
||||
position="bottom-start"
|
||||
shadow="md"
|
||||
// The theme defaults Popover to withinPortal: false, and this renders inside a Card.
|
||||
withinPortal
|
||||
>
|
||||
<Popover.Target>
|
||||
<ActionIcon
|
||||
variant="subtle"
|
||||
size="sm"
|
||||
aria-label="Rename group"
|
||||
title="Rename group"
|
||||
onClick={() => {
|
||||
setDraft(groupName);
|
||||
setOpened((open) => !open);
|
||||
}}
|
||||
>
|
||||
<IconPencil size={14} />
|
||||
</ActionIcon>
|
||||
</Popover.Target>
|
||||
<Popover.Dropdown>
|
||||
<Stack gap="xs" w={260}>
|
||||
<TextInput
|
||||
size="xs"
|
||||
label="Group name"
|
||||
value={draft}
|
||||
maxLength={120}
|
||||
data-autofocus
|
||||
onChange={(event) => setDraft(event.currentTarget.value)}
|
||||
onKeyDown={(event) => {
|
||||
if (event.key === 'Enter') submit();
|
||||
}}
|
||||
/>
|
||||
<Group justify="flex-end">
|
||||
<Button
|
||||
size="compact-xs"
|
||||
loading={rename.isPending}
|
||||
disabled={!name || name === groupName}
|
||||
onClick={submit}
|
||||
>
|
||||
Rename
|
||||
</Button>
|
||||
</Group>
|
||||
</Stack>
|
||||
</Popover.Dropdown>
|
||||
</Popover>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
import {
|
||||
Alert,
|
||||
Anchor,
|
||||
Badge,
|
||||
Button,
|
||||
Card,
|
||||
Checkbox,
|
||||
Group,
|
||||
Stack,
|
||||
Text,
|
||||
TextInput,
|
||||
Title,
|
||||
} from '@mantine/core';
|
||||
import { useState } from 'react';
|
||||
import { showErrorNotification, showSuccessNotification } from '~/utils/notifications';
|
||||
import { formatBytes } from '~/utils/number-helpers';
|
||||
import { trpc } from '~/utils/trpc';
|
||||
|
||||
const WEIGHT_EXTENSIONS = /\.(safetensors|sft|ckpt|pt|pth|bin|gguf|onnx)$/i;
|
||||
|
||||
export function RepoLookupSection() {
|
||||
const [source, setSource] = useState('');
|
||||
const [selected, setSelected] = useState<string[]>([]);
|
||||
const [groupName, setGroupName] = useState('');
|
||||
const queryUtils = trpc.useUtils();
|
||||
|
||||
const lookup = trpc.huggingFaceImport.lookup.useMutation({
|
||||
onSuccess: (data) => {
|
||||
// Weights we don't already hold are what "import this model" means; configs and duplicates opt in.
|
||||
// Prefilled from the repo and settable only here: no UI renames a group after Import.
|
||||
setGroupName(data.repo.split('/').pop() ?? data.repo);
|
||||
setSelected(
|
||||
data.files.filter((f) => WEIGHT_EXTENSIONS.test(f.path) && !f.existing).map((f) => f.path)
|
||||
);
|
||||
},
|
||||
onError: (error) =>
|
||||
showErrorNotification({ title: 'Lookup failed', error: new Error(error.message) }),
|
||||
});
|
||||
|
||||
const enqueue = trpc.huggingFaceImport.enqueue.useMutation({
|
||||
onSuccess: async (result) => {
|
||||
showSuccessNotification({
|
||||
title: 'Queued',
|
||||
message: `${result.queued} file(s) queued${
|
||||
result.skipped ? `, ${result.skipped} already queued` : ''
|
||||
}.`,
|
||||
});
|
||||
await queryUtils.huggingFaceImport.getAll.invalidate();
|
||||
},
|
||||
onError: (error) =>
|
||||
showErrorNotification({ title: 'Could not queue', error: new Error(error.message) }),
|
||||
});
|
||||
|
||||
const repo = lookup.data;
|
||||
|
||||
return (
|
||||
<Card withBorder padding="lg">
|
||||
<Stack gap="md">
|
||||
<Stack gap={4}>
|
||||
<Title order={4}>Import from Hugging Face</Title>
|
||||
<Text c="dimmed" size="sm">
|
||||
Paste a model URL. Files transfer server-side; nothing downloads to your machine.
|
||||
</Text>
|
||||
</Stack>
|
||||
|
||||
<Group align="flex-end" wrap="nowrap">
|
||||
<TextInput
|
||||
className="flex-1"
|
||||
label="Model URL"
|
||||
placeholder="https://huggingface.co/owner/name"
|
||||
value={source}
|
||||
onChange={(event) => setSource(event.currentTarget.value)}
|
||||
onKeyDown={(event) => {
|
||||
if (event.key === 'Enter' && source.trim()) lookup.mutate({ source });
|
||||
}}
|
||||
/>
|
||||
<Button
|
||||
onClick={() => lookup.mutate({ source })}
|
||||
loading={lookup.isPending}
|
||||
disabled={!source.trim()}
|
||||
>
|
||||
Look up
|
||||
</Button>
|
||||
</Group>
|
||||
|
||||
{repo && (
|
||||
<Stack gap="sm">
|
||||
<Group gap="xs">
|
||||
<Anchor href={`https://huggingface.co/${repo.repo}`} target="_blank" size="sm">
|
||||
{repo.repo}
|
||||
</Anchor>
|
||||
<Badge size="sm" variant="light">
|
||||
{repo.revision.slice(0, 7)}
|
||||
</Badge>
|
||||
{repo.license && (
|
||||
<Badge size="sm" variant="light" color="gray">
|
||||
{repo.license}
|
||||
</Badge>
|
||||
)}
|
||||
{repo.gated && (
|
||||
<Badge size="sm" color="orange">
|
||||
gated
|
||||
</Badge>
|
||||
)}
|
||||
</Group>
|
||||
|
||||
{repo.gated && (
|
||||
<Alert color="orange">
|
||||
This repo is gated. Importing it needs a server token whose Hugging Face account has
|
||||
accepted its terms — check the license before mirroring it here.
|
||||
</Alert>
|
||||
)}
|
||||
|
||||
<TextInput
|
||||
label="Group name"
|
||||
description="What this batch is called when attaching its files later. It can't be changed after Import."
|
||||
value={groupName}
|
||||
onChange={(event) => setGroupName(event.currentTarget.value)}
|
||||
/>
|
||||
|
||||
<Checkbox.Group value={selected} onChange={setSelected}>
|
||||
<Stack gap={4}>
|
||||
{repo.files.map((file) => (
|
||||
<Checkbox
|
||||
key={file.path}
|
||||
value={file.path}
|
||||
label={
|
||||
<Group gap="xs" wrap="nowrap">
|
||||
<Text size="sm">{file.path}</Text>
|
||||
<Text size="xs" c="dimmed">
|
||||
{formatBytes(file.size)}
|
||||
</Text>
|
||||
{file.existing && (
|
||||
<Badge size="xs" color="teal" variant="light">
|
||||
already stored as {file.existing.name}
|
||||
</Badge>
|
||||
)}
|
||||
</Group>
|
||||
}
|
||||
/>
|
||||
))}
|
||||
</Stack>
|
||||
</Checkbox.Group>
|
||||
|
||||
<Group justify="flex-end">
|
||||
<Button
|
||||
loading={enqueue.isPending}
|
||||
disabled={!selected.length || !groupName.trim()}
|
||||
onClick={() =>
|
||||
enqueue.mutate({
|
||||
repo: repo.repo,
|
||||
revision: repo.revision,
|
||||
paths: selected,
|
||||
groupName: groupName.trim() || undefined,
|
||||
})
|
||||
}
|
||||
>
|
||||
Import {selected.length} file{selected.length === 1 ? '' : 's'}
|
||||
</Button>
|
||||
</Group>
|
||||
</Stack>
|
||||
)}
|
||||
</Stack>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
import { Alert, Badge, Button, Group, Stack, Text } from '@mantine/core';
|
||||
import { openConfirmModal } from '@mantine/modals';
|
||||
import { DaysFromNow } from '~/components/Dates/DaysFromNow';
|
||||
import { RenameGroupControl } from '~/components/Moderation/HuggingFaceImport/RenameGroupControl';
|
||||
import { byGroup } from '~/components/Moderation/HuggingFaceImport/utils';
|
||||
import dayjs from '~/shared/utils/dayjs';
|
||||
import { formatBytes } from '~/utils/number-helpers';
|
||||
import { showErrorNotification, showSuccessNotification } from '~/utils/notifications';
|
||||
import { trpc } from '~/utils/trpc';
|
||||
import type { HuggingFaceImportView } from '~/server/services/huggingface-import.service';
|
||||
|
||||
const STALE_DAYS = 60;
|
||||
|
||||
export function UnattachedSection({ filter }: { filter: string }) {
|
||||
const queryUtils = trpc.useUtils();
|
||||
const groupName = filter.trim() || undefined;
|
||||
|
||||
const { data = [], isLoading } = trpc.huggingFaceImport.getAll.useQuery({
|
||||
limit: 200,
|
||||
unattached: true,
|
||||
groupName,
|
||||
});
|
||||
|
||||
const refresh = async () => {
|
||||
await Promise.all([
|
||||
queryUtils.huggingFaceImport.getAll.invalidate(),
|
||||
queryUtils.huggingFaceImport.getCounts.invalidate(),
|
||||
]);
|
||||
};
|
||||
|
||||
const remove = trpc.huggingFaceImport.delete.useMutation({
|
||||
onError: (error) =>
|
||||
showErrorNotification({ title: 'Could not delete', error: new Error(error.message) }),
|
||||
});
|
||||
|
||||
const confirmDelete = (rows: HuggingFaceImportView[]) =>
|
||||
openConfirmModal({
|
||||
title: `Delete ${rows.length} file${rows.length === 1 ? '' : 's'}?`,
|
||||
centered: true,
|
||||
labels: { confirm: 'Delete', cancel: 'Cancel' },
|
||||
confirmProps: { color: 'red' },
|
||||
children: (
|
||||
<Stack gap="sm">
|
||||
<Text size="sm">
|
||||
Frees{' '}
|
||||
<strong>{formatBytes(rows.reduce((sum, r) => sum + (r.sizeBytes ?? 0), 0))}</strong>.
|
||||
Re-importing means transferring them again from{' '}
|
||||
<Text span ff="monospace" size="sm">
|
||||
{rows[0]?.repo}
|
||||
</Text>
|
||||
.
|
||||
</Text>
|
||||
<Alert color="gray">
|
||||
<Text size="xs">
|
||||
Any file a model version still points at is refused, so a detached import that is
|
||||
still in use cannot be deleted from here.
|
||||
</Text>
|
||||
</Alert>
|
||||
</Stack>
|
||||
),
|
||||
onConfirm: async () => {
|
||||
// One call per file rather than a bulk endpoint: each deletes a distinct stored object, and
|
||||
// a partial failure should leave the rest deleted rather than rolling back freed bytes.
|
||||
let deleted = 0;
|
||||
for (const row of rows) {
|
||||
const ok = await remove
|
||||
.mutateAsync({ id: row.id })
|
||||
.then(() => true)
|
||||
.catch(() => false);
|
||||
if (ok) deleted++;
|
||||
}
|
||||
|
||||
// Counted, not assumed: the server refuses a delete it cannot make safe, and this toast is
|
||||
// the last thing a moderator reconciling storage reads.
|
||||
const failed = rows.length - deleted;
|
||||
if (failed)
|
||||
showErrorNotification({
|
||||
title: `Deleted ${deleted} of ${rows.length}`,
|
||||
error: new Error(`${failed} could not be deleted — see the errors above.`),
|
||||
});
|
||||
else showSuccessNotification({ title: 'Deleted', message: `${deleted} file(s) removed.` });
|
||||
await refresh();
|
||||
},
|
||||
});
|
||||
|
||||
const groups = byGroup(data);
|
||||
|
||||
return (
|
||||
<Stack gap="md">
|
||||
{!groups.length ? (
|
||||
<Text c="dimmed" size="sm">
|
||||
{isLoading
|
||||
? 'Loading…'
|
||||
: groupName
|
||||
? `Nothing unattached matches "${groupName}".`
|
||||
: 'Nothing unattached — every transferred file is on a model version.'}
|
||||
</Text>
|
||||
) : (
|
||||
groups.map((group) => {
|
||||
const stale = dayjs().diff(dayjs(group.oldest), 'day') >= STALE_DAYS;
|
||||
return (
|
||||
<Stack key={`${group.groupName}:${group.repo}:${group.revision}`} gap={6}>
|
||||
<Group gap="xs" wrap="wrap">
|
||||
<Text fw={600} size="sm">
|
||||
{group.groupName}
|
||||
</Text>
|
||||
<RenameGroupControl
|
||||
repo={group.repo}
|
||||
revision={group.revision}
|
||||
groupName={group.groupName}
|
||||
/>
|
||||
<Badge size="xs" variant="light" color="gray">
|
||||
{group.repo}
|
||||
</Badge>
|
||||
<Badge size="xs" variant="light">
|
||||
{group.revision.slice(0, 7)}
|
||||
</Badge>
|
||||
<Text size="xs" c="dimmed">
|
||||
{group.items.length} file{group.items.length === 1 ? '' : 's'} ·{' '}
|
||||
{formatBytes(group.bytes)} · imported <DaysFromNow date={group.oldest} />
|
||||
</Text>
|
||||
{stale && (
|
||||
<Badge size="xs" color="red">
|
||||
stale
|
||||
</Badge>
|
||||
)}
|
||||
<Button
|
||||
size="compact-xs"
|
||||
variant="subtle"
|
||||
color="red"
|
||||
ml="auto"
|
||||
onClick={() => confirmDelete(group.items)}
|
||||
>
|
||||
Delete {group.items.length}
|
||||
</Button>
|
||||
</Group>
|
||||
|
||||
{group.items.map((row) => (
|
||||
<Group key={row.id} gap="xs" pl="sm" wrap="nowrap">
|
||||
<Text size="xs" ff="monospace" lineClamp={1}>
|
||||
{row.filename}
|
||||
</Text>
|
||||
<Text size="xs" c="dimmed">
|
||||
{row.sizeBytes ? formatBytes(row.sizeBytes) : '—'}
|
||||
</Text>
|
||||
{row.suggestedType && (
|
||||
<Badge size="xs" variant="light">
|
||||
{row.suggestedType}
|
||||
</Badge>
|
||||
)}
|
||||
<Button
|
||||
size="compact-xs"
|
||||
variant="subtle"
|
||||
color="red"
|
||||
ml="auto"
|
||||
onClick={() => confirmDelete([row])}
|
||||
>
|
||||
Delete
|
||||
</Button>
|
||||
</Group>
|
||||
))}
|
||||
</Stack>
|
||||
);
|
||||
})
|
||||
)}
|
||||
</Stack>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
import { attachImports } from '~/components/Moderation/HuggingFaceImport/attach-imports';
|
||||
|
||||
const chosen = [
|
||||
{ id: 1, type: 'Model' as const },
|
||||
{ id: 2, type: 'VAE' as const },
|
||||
{ id: 3, type: 'Config' as const },
|
||||
];
|
||||
|
||||
describe('attachImports', () => {
|
||||
it('attaches every chosen import and reports the files it created', async () => {
|
||||
const attachOne = vi.fn(async ({ id }: { id: number }) => ({ modelFileId: id + 100 }));
|
||||
|
||||
const result = await attachImports({ chosen, modelVersionId: 42, attachOne });
|
||||
|
||||
expect(attachOne).toHaveBeenCalledTimes(3);
|
||||
expect(attachOne.mock.calls.map(([input]) => input)).toEqual([
|
||||
{ id: 1, modelVersionId: 42, type: 'Model' },
|
||||
{ id: 2, modelVersionId: 42, type: 'VAE' },
|
||||
{ id: 3, modelVersionId: 42, type: 'Config' },
|
||||
]);
|
||||
expect(result).toEqual({ modelFileIds: [101, 102, 103], failures: [] });
|
||||
});
|
||||
|
||||
it('keeps going after a failure, and reports the files that WERE created', async () => {
|
||||
// The router creates the ModelFile before it claims the import, so a lost claim leaves a real
|
||||
// file behind. Counting successes instead of naming ids loses the only handle to it.
|
||||
const attachOne = vi.fn(async ({ id }: { id: number }) => {
|
||||
if (id === 2) throw new Error('Created model file 555, but someone else attached it first.');
|
||||
return { modelFileId: id + 100 };
|
||||
});
|
||||
|
||||
const result = await attachImports({ chosen, modelVersionId: 42, attachOne });
|
||||
|
||||
// Three calls, not two: an early break or a `.every` would stop at the failure and silently
|
||||
// skip the rest of what the moderator asked for.
|
||||
expect(attachOne).toHaveBeenCalledTimes(3);
|
||||
expect(result.modelFileIds).toEqual([101, 103]);
|
||||
expect(result.failures).toEqual([
|
||||
'Created model file 555, but someone else attached it first.',
|
||||
]);
|
||||
});
|
||||
|
||||
it('collects EVERY failure message, not just the first', async () => {
|
||||
const attachOne = vi.fn(async ({ id }: { id: number }) => {
|
||||
throw new Error(`failed ${id}`);
|
||||
});
|
||||
|
||||
const result = await attachImports({ chosen, modelVersionId: 42, attachOne });
|
||||
|
||||
expect(result.modelFileIds).toEqual([]);
|
||||
expect(result.failures).toEqual(['failed 1', 'failed 2', 'failed 3']);
|
||||
});
|
||||
|
||||
it('starts each call only after the previous one settled', async () => {
|
||||
// Held open until released, so "waited for the previous call" cannot be confused with "waited a
|
||||
// moment" — a staggered parallel loop still starts all three while the first is unresolved.
|
||||
const releases: Array<() => void> = [];
|
||||
const attachOne = vi.fn(
|
||||
({ id }: { id: number }) =>
|
||||
new Promise<{ modelFileId: number }>((resolve) =>
|
||||
releases.push(() => resolve({ modelFileId: id }))
|
||||
)
|
||||
);
|
||||
const settle = () => new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
const done = attachImports({ chosen, modelVersionId: 42, attachOne });
|
||||
for (let started = 1; started <= chosen.length; started++) {
|
||||
await settle();
|
||||
expect(attachOne).toHaveBeenCalledTimes(started);
|
||||
releases[started - 1]();
|
||||
}
|
||||
await expect(done).resolves.toEqual({ modelFileIds: [1, 2, 3], failures: [] });
|
||||
});
|
||||
|
||||
it('reports a rejection that is not an Error by its value', async () => {
|
||||
const attachOne = vi.fn(async () => {
|
||||
throw 'storage resolver unavailable';
|
||||
});
|
||||
|
||||
const result = await attachImports({
|
||||
chosen: chosen.slice(0, 1),
|
||||
modelVersionId: 42,
|
||||
attachOne,
|
||||
});
|
||||
|
||||
expect(result.failures).toEqual(['storage resolver unavailable']);
|
||||
});
|
||||
|
||||
it('does nothing when nothing was chosen', async () => {
|
||||
const attachOne = vi.fn();
|
||||
const result = await attachImports({ chosen: [], modelVersionId: 42, attachOne });
|
||||
|
||||
expect(attachOne).not.toHaveBeenCalled();
|
||||
expect(result).toEqual({ modelFileIds: [], failures: [] });
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,48 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { byGroup } from '~/components/Moderation/HuggingFaceImport/utils';
|
||||
import type { HuggingFaceImportView } from '~/server/services/huggingface-import.service';
|
||||
|
||||
const row = (over: Partial<HuggingFaceImportView>) =>
|
||||
({
|
||||
id: 1,
|
||||
groupName: 'flux-krea',
|
||||
repo: 'black-forest-labs/FLUX.1-Krea-dev',
|
||||
revision: 'aaaaaaa1',
|
||||
filename: 'model.safetensors',
|
||||
sizeBytes: 100,
|
||||
createdAt: new Date('2026-09-01'),
|
||||
...over,
|
||||
} as HuggingFaceImportView);
|
||||
|
||||
describe('byGroup', () => {
|
||||
it('keeps a group name that contains spaces intact', () => {
|
||||
// Group names are free text a moderator types, so anything that re-parses a joined key renders
|
||||
// the wrong label and collides two groups onto one React key.
|
||||
const groups = byGroup([
|
||||
row({ id: 1, groupName: 'FLUX Krea batch' }),
|
||||
row({ id: 2, groupName: 'FLUX Dev batch' }),
|
||||
]);
|
||||
expect(groups.map((g) => g.groupName).sort()).toEqual(['FLUX Dev batch', 'FLUX Krea batch']);
|
||||
});
|
||||
|
||||
it('separates two revisions of the same repo', () => {
|
||||
// The default group name is derived from the repo, so both revisions carry the same name.
|
||||
const groups = byGroup([
|
||||
row({ id: 1, revision: 'aaaaaaa1' }),
|
||||
row({ id: 2, revision: 'bbbbbbb2' }),
|
||||
]);
|
||||
expect(groups).toHaveLength(2);
|
||||
expect(groups.map((g) => g.revision).sort()).toEqual(['aaaaaaa1', 'bbbbbbb2']);
|
||||
});
|
||||
|
||||
it('sums bytes and reports the oldest import in the group', () => {
|
||||
const groups = byGroup([
|
||||
row({ id: 1, sizeBytes: 100, createdAt: new Date('2026-09-05') }),
|
||||
row({ id: 2, sizeBytes: 250, createdAt: new Date('2026-09-02') }),
|
||||
row({ id: 3, sizeBytes: null, createdAt: new Date('2026-09-09') }),
|
||||
]);
|
||||
expect(groups).toHaveLength(1);
|
||||
expect(groups[0].bytes).toBe(350);
|
||||
expect(groups[0].oldest).toEqual(new Date('2026-09-02'));
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,48 @@
|
||||
import type { ModelFileType } from '~/server/common/constants';
|
||||
|
||||
export type AttachSelection = { id: number; type: ModelFileType };
|
||||
|
||||
export type AttachOutcome = {
|
||||
/** Model file ids created, in the order they were created. */
|
||||
modelFileIds: number[];
|
||||
failures: string[];
|
||||
};
|
||||
|
||||
/**
|
||||
* Attaches each chosen import in turn.
|
||||
*
|
||||
* Sequential on purpose: each attach awaits a storage-resolver registration and a scan
|
||||
* submission, so a parallel batch multiplies that load for no user-visible gain.
|
||||
*
|
||||
* 🔴 Every item is attempted even after one fails, and the caller is told which files were
|
||||
* CREATED, not merely how many succeeded. A failure here is not "nothing happened": the router
|
||||
* creates the `ModelFile` first and only then claims the import, so losing that claim leaves a
|
||||
* real file behind whose id is the only way to find it again.
|
||||
*/
|
||||
export async function attachImports({
|
||||
chosen,
|
||||
modelVersionId,
|
||||
attachOne,
|
||||
}: {
|
||||
chosen: AttachSelection[];
|
||||
modelVersionId: number;
|
||||
attachOne: (input: {
|
||||
id: number;
|
||||
modelVersionId: number;
|
||||
type: ModelFileType;
|
||||
}) => Promise<{ modelFileId: number }>;
|
||||
}): Promise<AttachOutcome> {
|
||||
const modelFileIds: number[] = [];
|
||||
const failures: string[] = [];
|
||||
|
||||
for (const item of chosen) {
|
||||
try {
|
||||
const result = await attachOne({ id: item.id, modelVersionId, type: item.type });
|
||||
modelFileIds.push(result.modelFileId);
|
||||
} catch (error) {
|
||||
failures.push(error instanceof Error ? error.message : String(error));
|
||||
}
|
||||
}
|
||||
|
||||
return { modelFileIds, failures };
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
import { groupBy } from 'lodash-es';
|
||||
import type { HuggingFaceImportView } from '~/server/services/huggingface-import.service';
|
||||
|
||||
/**
|
||||
* Groups by the batch a moderator would recognise. Revision is part of the identity: a group name
|
||||
* defaults to the repo's last path segment, so re-importing the same repo at a new revision reuses
|
||||
* the name — and collapsing those would put one revision's badge over both revisions' files, under
|
||||
* a `Delete N` that spans them.
|
||||
*/
|
||||
export function byGroup(rows: HuggingFaceImportView[]) {
|
||||
const groups = groupBy(rows, (row) => JSON.stringify([row.groupName, row.repo, row.revision]));
|
||||
return Object.values(groups).map((items) => ({
|
||||
groupName: items[0].groupName,
|
||||
repo: items[0].repo,
|
||||
revision: items[0].revision,
|
||||
items,
|
||||
bytes: items.reduce((sum, item) => sum + (item.sizeBytes ?? 0), 0),
|
||||
oldest: items.reduce(
|
||||
(oldest, item) => (item.createdAt < oldest ? item.createdAt : oldest),
|
||||
items[0].createdAt
|
||||
),
|
||||
}));
|
||||
}
|
||||
@@ -21,6 +21,7 @@ export function ModerationNav() {
|
||||
hidden: !features.comicCreator,
|
||||
},
|
||||
{ label: 'Models', href: '/moderator/models' },
|
||||
{ label: 'HuggingFace Import', href: '/moderator/huggingface-import' },
|
||||
{ label: 'Training Models', href: '/moderator/training-models' },
|
||||
{ label: 'Training Data Review', href: '/moderator/review/training-data' },
|
||||
// Migrated to the moderator app — the /moderator/* route redirects there (see the moderator
|
||||
|
||||
@@ -39,6 +39,7 @@ import { isEqual, startCase } from 'lodash-es';
|
||||
import { useEffect, useState } from 'react';
|
||||
|
||||
import { UploadNotice } from '~/components/UploadNotice/UploadNotice';
|
||||
import { AddFromImportsButton } from '~/components/Moderation/HuggingFaceImport/AddFromImportsButton';
|
||||
import type { FileFromContextProps } from '~/components/Resource/FilesProvider';
|
||||
import { useFilesContext } from '~/components/Resource/FilesProvider';
|
||||
import type { ModelFileType, ZipModelFileType } from '~/server/common/constants';
|
||||
@@ -50,13 +51,9 @@ import { useS3UploadStore } from '~/store/s3-upload.store';
|
||||
import { removeDuplicates } from '~/utils/array-helpers';
|
||||
import { showErrorNotification } from '~/utils/notifications';
|
||||
import { formatBytes, formatKBytes, formatSeconds } from '~/utils/number-helpers';
|
||||
import { getDisplayName, getFileExtension, sanitizeDownloadFilename } from '~/utils/string-helpers';
|
||||
import { getFileExtension, sanitizeDownloadFilename } from '~/utils/string-helpers';
|
||||
import { trpc } from '~/utils/trpc';
|
||||
import {
|
||||
comfyFileTypeLabels,
|
||||
filterFileTypeByExtension,
|
||||
UNQUANTIZED_QUANT_TYPE,
|
||||
} from '~/utils/file-display-helpers';
|
||||
import { getModelFileTypeOptions, UNQUANTIZED_QUANT_TYPE } from '~/utils/file-display-helpers';
|
||||
import classes from './Files.module.scss';
|
||||
import { LegacyActionIcon } from '~/components/LegacyActionIcon/LegacyActionIcon';
|
||||
import { isAndroidDevice } from '~/utils/device-helpers';
|
||||
@@ -144,6 +141,7 @@ export function Files({ showRenameOnPrimary }: { showRenameOnPrimary?: boolean }
|
||||
onDrop,
|
||||
dropzoneConfig,
|
||||
baseModel,
|
||||
modelVersionId,
|
||||
addLinkedComponent,
|
||||
removeLinkedComponent,
|
||||
} = useFilesContext();
|
||||
@@ -312,11 +310,14 @@ export function Files({ showRenameOnPrimary }: { showRenameOnPrimary?: boolean }
|
||||
py="md"
|
||||
style={{ borderColor: 'rgba(34, 139, 230, 0.2)' }}
|
||||
>
|
||||
<Group gap="xs">
|
||||
<IconFile3d size={20} style={{ color: 'var(--mantine-color-blue-4)' }} />
|
||||
<Text fw={600} c="white">
|
||||
Model Files
|
||||
</Text>
|
||||
<Group gap="xs" justify="space-between" wrap="nowrap">
|
||||
<Group gap="xs">
|
||||
<IconFile3d size={20} style={{ color: 'var(--mantine-color-blue-4)' }} />
|
||||
<Text fw={600} c="white">
|
||||
Model Files
|
||||
</Text>
|
||||
</Group>
|
||||
{modelVersionId && <AddFromImportsButton modelVersionId={modelVersionId} />}
|
||||
</Group>
|
||||
<Text size="sm" c="dimmed" mt={4}>
|
||||
The main model files users will download. We'll show the best match based on their
|
||||
@@ -1034,11 +1035,6 @@ function FileEditForm({
|
||||
}
|
||||
};
|
||||
|
||||
// Keep the file's own type selectable even when it's no longer offered, so a
|
||||
// legacy type doesn't render as a blank Select.
|
||||
const filterByFileExtension = (value: ModelFileType) =>
|
||||
value === versionFile.type || filterFileTypeByExtension(value, versionFile.name);
|
||||
|
||||
const handleReset = () => {
|
||||
updateFile(versionFile.uuid, {
|
||||
type: initialFile.type,
|
||||
@@ -1078,12 +1074,11 @@ function FileEditForm({
|
||||
w={160}
|
||||
placeholder="Type"
|
||||
error={error?.type?._errors[0]}
|
||||
data={fileTypes.filter(filterByFileExtension).map((x) => ({
|
||||
label:
|
||||
comfyFileTypeLabels[x] ??
|
||||
getDisplayName(x === 'Model' ? versionFile.modelType ?? x : x),
|
||||
value: x,
|
||||
}))}
|
||||
data={getModelFileTypeOptions(versionFile.name, {
|
||||
types: fileTypes,
|
||||
currentType: versionFile.type,
|
||||
modelType: versionFile.modelType,
|
||||
})}
|
||||
value={versionFile.type ?? null}
|
||||
onChange={(value) => {
|
||||
const newType = value as ModelFileType | null;
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
import { beforeEach, describe, expect, test, vi } from 'vitest';
|
||||
import { page } from 'vitest/browser';
|
||||
import { useS3UploadStore } from '~/store/s3-upload.store';
|
||||
import type * as TrpcModule from '~/utils/trpc';
|
||||
import { renderWithProviders } from '../../../test/component-setup';
|
||||
|
||||
/**
|
||||
* `adoptFiles` is the only way a file created outside this provider reaches `files`, and nothing
|
||||
* else in the suite exercises the real one.
|
||||
*
|
||||
* The fixture is deliberately un-permissive: populated metadata, a distinct `overrideName`, and a
|
||||
* real pending row seeded through the provider's own upload-store path. A fixture where every
|
||||
* mapped field is `undefined` cannot observe a mapping that drops one.
|
||||
*/
|
||||
|
||||
const noMutation = vi.hoisted(() => () => ({
|
||||
mutateAsync: vi.fn(),
|
||||
mutate: vi.fn(),
|
||||
isLoading: false,
|
||||
}));
|
||||
|
||||
const { mockForEditFetch } = vi.hoisted(() => ({ mockForEditFetch: vi.fn() }));
|
||||
|
||||
vi.mock('~/utils/trpc', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof TrpcModule>()),
|
||||
trpc: {
|
||||
useUtils: () => ({
|
||||
modelFile: { hasOfficialFileOfSize: { fetch: vi.fn() } },
|
||||
modelVersion: { getByIdForEdit: { fetch: mockForEditFetch } },
|
||||
}),
|
||||
modelFile: {
|
||||
getOptions: { useQuery: () => ({ data: undefined }) },
|
||||
create: { useMutation: noMutation },
|
||||
},
|
||||
modelVersion: {
|
||||
setLinkedComponents: { useMutation: noMutation },
|
||||
linkOfficialFileByHash: { useMutation: noMutation },
|
||||
addLinkedComponent: { useMutation: noMutation },
|
||||
publish: { useMutation: noMutation },
|
||||
},
|
||||
model: { publish: { useMutation: noMutation } },
|
||||
},
|
||||
}));
|
||||
vi.mock('~/hooks/useFileHash', () => ({ useFileHash: () => ({ hashFile: vi.fn() }) }));
|
||||
vi.mock('~/components/Resource/official-match', () => ({ resolveOfficialFileHash: vi.fn() }));
|
||||
|
||||
import type { FileFromContextProps } from '~/components/Resource/FilesProvider';
|
||||
import { FilesProvider, useFilesContext } from '~/components/Resource/FilesProvider';
|
||||
|
||||
const VERSION_ID = 42;
|
||||
|
||||
/** A server row with every mapped field populated, so a dropped or swapped field is observable. */
|
||||
const serverFile = (id: number, name: string) => ({
|
||||
id,
|
||||
name,
|
||||
overrideName: `${name}-override`,
|
||||
type: 'Model',
|
||||
sizeKB: 1024,
|
||||
metadata: {
|
||||
size: 'pruned',
|
||||
fp: 'fp16',
|
||||
format: 'SafeTensor',
|
||||
quantType: 'Q8_0',
|
||||
isRequired: true,
|
||||
},
|
||||
});
|
||||
|
||||
/** Every field `toFileFromContext` maps, serialised so one `toBe` pins the whole row. */
|
||||
const project = (file: FileFromContextProps) =>
|
||||
[
|
||||
file.name,
|
||||
file.overrideName,
|
||||
file.type,
|
||||
file.fp,
|
||||
file.format,
|
||||
file.size,
|
||||
file.quantType,
|
||||
file.isRequired,
|
||||
file.sizeKB,
|
||||
file.versionId,
|
||||
file.modelType,
|
||||
]
|
||||
.map((value) => value ?? '-')
|
||||
.join('|');
|
||||
|
||||
/** What `project` prints for a `serverFile` mapped under this harness's version and model. */
|
||||
const row = (name: string, fp = 'fp16') =>
|
||||
`${name}|${name}-override|Model|${fp}|SafeTensor|pruned|Q8_0|true|1024|${VERSION_ID}|Checkpoint`;
|
||||
|
||||
function Harness({ adopt }: { adopt: number[] }) {
|
||||
const { files, adoptFiles, updateFile } = useFilesContext();
|
||||
return (
|
||||
<div>
|
||||
<button onClick={() => void adoptFiles(adopt)}>adopt</button>
|
||||
{/* An edit held in provider state and not yet saved — what a creator mid-form has. */}
|
||||
<button onClick={() => files[0] && updateFile(files[0].uuid, { fp: 'bf16' })}>edit</button>
|
||||
<span data-testid="files">{files.map(project).join(' / ') || 'none'}</span>
|
||||
<span data-testid="uuids">{files.map((file) => file.uuid).join(',')}</span>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function renderHarness({
|
||||
versionFiles = [] as unknown[],
|
||||
adopt = [] as number[],
|
||||
pending = false,
|
||||
} = {}) {
|
||||
if (pending)
|
||||
// Through the provider's real pending-row path (it reads the upload store when it seeds).
|
||||
useS3UploadStore.setState({
|
||||
items: [
|
||||
{
|
||||
name: 'pending.safetensors',
|
||||
size: 2048,
|
||||
file: new File([], 'pending.safetensors'),
|
||||
meta: { versionId: VERSION_ID, uuid: 'pending-uuid', type: 'Model' },
|
||||
},
|
||||
] as never,
|
||||
});
|
||||
|
||||
renderWithProviders(
|
||||
<FilesProvider
|
||||
model={{ id: 1, type: 'Checkpoint' }}
|
||||
version={{ id: VERSION_ID, files: versionFiles as never, baseModel: 'Flux.1 D' }}
|
||||
>
|
||||
<Harness adopt={adopt} />
|
||||
</FilesProvider>
|
||||
);
|
||||
}
|
||||
|
||||
const filesText = () => page.getByTestId('files').element().textContent;
|
||||
const uuidsText = () => page.getByTestId('uuids').element().textContent ?? '';
|
||||
|
||||
/** Budget trimmed from 15 s: the state is ABSORBING, so a revert reports in 2 s. */
|
||||
const untilAdopted = (name: string) =>
|
||||
expect.element(page.getByTestId('files'), { timeout: 2000 }).toHaveTextContent(name);
|
||||
|
||||
async function editFirstRow() {
|
||||
await page.getByRole('button', { name: 'edit' }).click();
|
||||
await expect.element(page.getByTestId('files')).toHaveTextContent('|bf16|');
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
mockForEditFetch.mockReset();
|
||||
// Module-scope store, shared across every test in this file.
|
||||
useS3UploadStore.setState({ items: [] as never });
|
||||
});
|
||||
|
||||
describe('FilesProvider.adoptFiles', () => {
|
||||
test('adds a named file with every mapped field, and asks for the files', async () => {
|
||||
renderHarness({ versionFiles: [], adopt: [8] });
|
||||
mockForEditFetch.mockResolvedValue({ files: [serverFile(8, 'imported.safetensors')] });
|
||||
|
||||
await page.getByRole('button', { name: 'adopt' }).click();
|
||||
await untilAdopted('imported.safetensors');
|
||||
|
||||
expect(filesText()).toBe(row('imported.safetensors'));
|
||||
// The REQUEST, not the mock's answer: `withFiles: false` returns no files, so the list would
|
||||
// never update — invisible if only the response is read.
|
||||
expect(mockForEditFetch).toHaveBeenCalledTimes(1);
|
||||
expect(mockForEditFetch).toHaveBeenCalledWith({ id: VERSION_ID, withFiles: true });
|
||||
});
|
||||
|
||||
test('adopts ONLY the ids it was given', async () => {
|
||||
renderHarness({ versionFiles: [], adopt: [8] });
|
||||
mockForEditFetch.mockResolvedValue({
|
||||
files: [serverFile(8, 'imported.safetensors'), serverFile(9, 'someone-elses.safetensors')],
|
||||
});
|
||||
|
||||
await page.getByRole('button', { name: 'adopt' }).click();
|
||||
await untilAdopted('imported.safetensors');
|
||||
|
||||
expect(filesText()).toBe(row('imported.safetensors'));
|
||||
});
|
||||
|
||||
test('leaves an existing row untouched, including edits not yet saved', async () => {
|
||||
renderHarness({ versionFiles: [serverFile(7, 'existing.safetensors')], adopt: [8] });
|
||||
// The server still reports the ORIGINAL precision; state holds the unsaved edit.
|
||||
mockForEditFetch.mockResolvedValue({
|
||||
files: [serverFile(7, 'existing.safetensors'), serverFile(8, 'imported.safetensors')],
|
||||
});
|
||||
|
||||
await editFirstRow();
|
||||
const [existingUuid] = uuidsText().split(',');
|
||||
await page.getByRole('button', { name: 'adopt' }).click();
|
||||
await untilAdopted('imported.safetensors');
|
||||
|
||||
expect(filesText()).toBe(
|
||||
`${row('existing.safetensors', 'bf16')} / ${row('imported.safetensors')}`
|
||||
);
|
||||
// `Files.tsx` binds upload progress and the edit form's dirty baseline to the uuid.
|
||||
expect(uuidsText().split(',')[0]).toBe(existingUuid);
|
||||
});
|
||||
|
||||
test('keeps a file still uploading, which the server cannot report', async () => {
|
||||
renderHarness({
|
||||
versionFiles: [serverFile(7, 'existing.safetensors')],
|
||||
adopt: [8],
|
||||
pending: true,
|
||||
});
|
||||
mockForEditFetch.mockResolvedValue({
|
||||
files: [serverFile(7, 'existing.safetensors'), serverFile(8, 'imported.safetensors')],
|
||||
});
|
||||
|
||||
await page.getByRole('button', { name: 'adopt' }).click();
|
||||
await untilAdopted('imported.safetensors');
|
||||
|
||||
expect(filesText()).toContain('pending.safetensors');
|
||||
});
|
||||
|
||||
test('re-naming a file the list already holds neither duplicates nor rebuilds it', async () => {
|
||||
renderHarness({ versionFiles: [serverFile(7, 'existing.safetensors')], adopt: [7, 8] });
|
||||
mockForEditFetch.mockResolvedValue({
|
||||
files: [serverFile(7, 'existing.safetensors'), serverFile(8, 'imported.safetensors')],
|
||||
});
|
||||
|
||||
await editFirstRow();
|
||||
const [existingUuid] = uuidsText().split(',');
|
||||
await page.getByRole('button', { name: 'adopt' }).click();
|
||||
await untilAdopted('imported.safetensors');
|
||||
|
||||
// Rebuilding id 7 from the server would revert the edit and mint a new uuid.
|
||||
expect(filesText()).toBe(
|
||||
`${row('existing.safetensors', 'bf16')} / ${row('imported.safetensors')}`
|
||||
);
|
||||
expect(uuidsText().split(',')[0]).toBe(existingUuid);
|
||||
});
|
||||
|
||||
test('keeps an edit made while the fetch was in flight', async () => {
|
||||
renderHarness({ versionFiles: [serverFile(7, 'existing.safetensors')], adopt: [8] });
|
||||
let resolveFetch!: (value: unknown) => void;
|
||||
mockForEditFetch.mockReturnValue(new Promise((resolve) => (resolveFetch = resolve)));
|
||||
|
||||
await page.getByRole('button', { name: 'adopt' }).click();
|
||||
await expect.poll(() => mockForEditFetch.mock.calls.length).toBe(1);
|
||||
await editFirstRow();
|
||||
resolveFetch({
|
||||
files: [serverFile(7, 'existing.safetensors'), serverFile(8, 'imported.safetensors')],
|
||||
});
|
||||
await untilAdopted('imported.safetensors');
|
||||
|
||||
// Building from the list as it was when `adoptFiles` was called would drop the edit.
|
||||
expect(filesText()).toBe(
|
||||
`${row('existing.safetensors', 'bf16')} / ${row('imported.safetensors')}`
|
||||
);
|
||||
});
|
||||
|
||||
test('asks for nothing when no ids were created', async () => {
|
||||
renderHarness({ versionFiles: [serverFile(7, 'existing.safetensors')], adopt: [] });
|
||||
|
||||
await page.getByRole('button', { name: 'adopt' }).click();
|
||||
await expect.element(page.getByTestId('files')).toHaveTextContent('existing.safetensors');
|
||||
expect(mockForEditFetch).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -86,6 +86,8 @@ type FilesContextState = {
|
||||
files: FileFromContextProps[];
|
||||
linkedComponents: LinkedComponent[];
|
||||
modelId?: number;
|
||||
modelVersionId?: number;
|
||||
modelType?: ModelType | null;
|
||||
baseModel?: string;
|
||||
usageControl?: ModelUsageControl | null;
|
||||
dropzoneConfig: DropzoneOptions;
|
||||
@@ -98,6 +100,7 @@ type FilesContextState = {
|
||||
addLinkedComponent: (
|
||||
component: LinkedComponent | Omit<LinkedComponent, 'fileId' | 'fileName' | 'sizeKB'>
|
||||
) => Promise<void>;
|
||||
adoptFiles: (modelFileIds: number[]) => Promise<void>;
|
||||
removeLinkedComponent: (versionId: number) => void;
|
||||
};
|
||||
|
||||
@@ -110,6 +113,27 @@ type FilesProviderProps = {
|
||||
children: React.ReactNode;
|
||||
};
|
||||
|
||||
function toFileFromContext(
|
||||
file: NonNullable<ModelVersionById['files']>[number],
|
||||
{ versionId, modelType, uuid }: { versionId?: number; modelType?: ModelType | null; uuid: string }
|
||||
) {
|
||||
return {
|
||||
id: file.id,
|
||||
name: file.name,
|
||||
overrideName: file.overrideName ?? null,
|
||||
type: file.type as ModelFileType,
|
||||
sizeKB: file.sizeKB,
|
||||
size: file.metadata?.size,
|
||||
fp: file.metadata?.fp,
|
||||
format: file.metadata?.format,
|
||||
quantType: file.metadata?.quantType,
|
||||
isRequired: file.metadata?.isRequired ?? null,
|
||||
versionId,
|
||||
uuid,
|
||||
modelType,
|
||||
} as FileFromContextProps;
|
||||
}
|
||||
|
||||
const FilesContext = createContext<FilesContextState | null>(null);
|
||||
export const useFilesContext = () => {
|
||||
const context = useContext(FilesContext);
|
||||
@@ -126,21 +150,13 @@ export function FilesProvider({ model, version, children }: FilesProviderProps)
|
||||
|
||||
const [errors, setErrors] = useState<FileErrors | null>(null);
|
||||
const [files, setFiles] = useState<FileFromContextProps[]>(() => {
|
||||
const initialFiles = (version?.files?.map((file) => ({
|
||||
id: file.id,
|
||||
name: file.name,
|
||||
overrideName: file.overrideName ?? null,
|
||||
type: file.type as ModelFileType,
|
||||
sizeKB: file.sizeKB,
|
||||
size: file.metadata?.size,
|
||||
fp: file.metadata?.fp,
|
||||
format: file.metadata?.format,
|
||||
quantType: file.metadata?.quantType,
|
||||
isRequired: file.metadata?.isRequired ?? null,
|
||||
versionId: version.id,
|
||||
uuid: randomId(),
|
||||
modelType: model?.type ?? null,
|
||||
})) ?? []) as FileFromContextProps[];
|
||||
const initialFiles = (version?.files ?? []).map((file) =>
|
||||
toFileFromContext(file, {
|
||||
versionId: version?.id,
|
||||
modelType: model?.type ?? null,
|
||||
uuid: randomId(),
|
||||
})
|
||||
);
|
||||
const uploading = useS3UploadStore
|
||||
.getState()
|
||||
.items.filter((x) => x.meta?.versionId === version?.id)
|
||||
@@ -167,6 +183,38 @@ export function FilesProvider({ model, version, children }: FilesProviderProps)
|
||||
// effect doesn't start the same file twice across renders.
|
||||
const startedUploadsRef = useRef<Set<string>>(new Set());
|
||||
|
||||
/**
|
||||
* For files created outside this provider. `files` is seeded once in a `useState` initializer,
|
||||
* so no query invalidation reaches it.
|
||||
*
|
||||
* 🔴 Append-only, and only the named ids. Re-seeding would revert unsaved metadata edits and
|
||||
* duplicate an upload whose row the server committed before this client learned its id.
|
||||
*
|
||||
* `getByIdForEdit`, not `getById`: it reads the primary (`forceWriteDb`), so a file created a
|
||||
* moment ago is not lost to replica lag.
|
||||
*/
|
||||
const adoptFiles = async (modelFileIds: number[]) => {
|
||||
if (!version?.id || !modelFileIds.length) return;
|
||||
const fresh = await queryUtils.modelVersion.getByIdForEdit.fetch({
|
||||
id: version.id,
|
||||
withFiles: true,
|
||||
});
|
||||
const wanted = new Set(modelFileIds);
|
||||
setFiles((state) => {
|
||||
const present = new Set(state.map((file) => file.id).filter(isDefined));
|
||||
const added = (fresh?.files ?? [])
|
||||
.filter((file) => wanted.has(file.id) && !present.has(file.id))
|
||||
.map((file) =>
|
||||
toFileFromContext(file, {
|
||||
versionId: version.id,
|
||||
modelType: model?.type ?? null,
|
||||
uuid: randomId(),
|
||||
})
|
||||
);
|
||||
return added.length ? [...state, ...added] : state;
|
||||
});
|
||||
};
|
||||
|
||||
const handleUpdateFile = (uuid: string, file: Partial<FileFromContextProps>) => {
|
||||
setFiles((state) => state.map((x) => (x.uuid === uuid ? { ...x, ...file } : x)));
|
||||
};
|
||||
@@ -796,9 +844,12 @@ export function FilesProvider({ model, version, children }: FilesProviderProps)
|
||||
removeFile,
|
||||
dropzoneConfig,
|
||||
modelId: model?.id,
|
||||
modelVersionId: version?.id,
|
||||
modelType: model?.type ?? null,
|
||||
baseModel: version?.baseModel ?? undefined,
|
||||
usageControl: version?.usageControl,
|
||||
validationCheck: checkValidation,
|
||||
adoptFiles,
|
||||
addLinkedComponent,
|
||||
removeLinkedComponent,
|
||||
}}
|
||||
|
||||
Vendored
+3
-1
@@ -457,6 +457,9 @@ export const serverSchema = z
|
||||
// Local-dev opt-in for /api/training-studio/host to hand the shared ORCHESTRATOR_ACCESS_TOKEN
|
||||
// (the ORCHESTRATOR_MODE=dev arm of getOrchestratorToken) to the browser. Never set in prod.
|
||||
ALLOW_DEV_ORCHESTRATOR_TOKEN_PASSTHROUGH: zc.booleanString.optional().default(false),
|
||||
// Optional. Without it only public, ungated Hugging Face repos can be imported; with it, repos
|
||||
// this token's account has accepted the terms for.
|
||||
HUGGING_FACE_TOKEN: z.string().optional(),
|
||||
AXIOM_TOKEN: z.string().optional(),
|
||||
AXIOM_ORG_ID: z.string().optional(),
|
||||
AXIOM_DATASTREAM: z.string().optional(),
|
||||
@@ -755,7 +758,6 @@ export const serverSchema = z
|
||||
FRESHDESK_DOMAIN: z.string().optional(),
|
||||
FRESHDESK_TOKEN: z.string().optional(),
|
||||
FRESHDESK_AGENT_ID: z.coerce.number().optional(),
|
||||
UPLOAD_PROHIBITED_EXTENSIONS: commaDelimitedStringArray().optional(),
|
||||
// Enforce the post-completion object-existence check in /api/upload/complete.
|
||||
// 🔴 Defaults to FALSE = observe-only: the probe still runs and its verdict is
|
||||
// logged, but a missing object does NOT fail the request. That ordering is
|
||||
|
||||
@@ -1,39 +0,0 @@
|
||||
import type { NextApiRequest, NextApiResponse } from 'next';
|
||||
import { dbWrite } from '~/server/db/client';
|
||||
import * as z from 'zod';
|
||||
import { processImport } from '~/server/importers/importRouter';
|
||||
import { ModEndpoint } from '~/server/utils/endpoint-helpers';
|
||||
|
||||
const importSchema = z.object({
|
||||
source: z.string().trim().url(),
|
||||
wait: z
|
||||
.preprocess((x) => x == 'true', z.boolean())
|
||||
.optional()
|
||||
.default(false),
|
||||
data: z.any().optional(),
|
||||
});
|
||||
|
||||
export default ModEndpoint(
|
||||
async function importSource(req: NextApiRequest, res: NextApiResponse) {
|
||||
const { source, wait, data } = importSchema.parse(req.query);
|
||||
const userId = -1; //Default civitai user id
|
||||
|
||||
const { id } = await dbWrite.import.create({
|
||||
data: {
|
||||
source,
|
||||
userId,
|
||||
data: data,
|
||||
},
|
||||
select: { id: true },
|
||||
});
|
||||
|
||||
if (wait) {
|
||||
const result = await processImport({ id, source, userId, data });
|
||||
res.status(200).json(result);
|
||||
} else {
|
||||
res.status(200).json({ id });
|
||||
await processImport({ id, source, userId, data });
|
||||
}
|
||||
},
|
||||
['GET']
|
||||
);
|
||||
@@ -2,8 +2,6 @@ import type { NextApiRequest, NextApiResponse } from 'next';
|
||||
import { instrumentApiResponse } from '~/server/prom/http-errors';
|
||||
import { getServerAuthSession } from '~/server/auth/get-server-auth-session';
|
||||
import { UploadType } from '~/server/common/enums';
|
||||
import { extname } from 'node:path';
|
||||
import { filenamize, generateToken } from '~/utils/string-helpers';
|
||||
import {
|
||||
getMultipartPutUrl,
|
||||
getUploadS3Client,
|
||||
@@ -11,6 +9,7 @@ import {
|
||||
getUploadChunkSize,
|
||||
} from '~/utils/s3-utils';
|
||||
import type { UploadBackend } from '~/utils/s3-utils';
|
||||
import { buildUploadKey } from '~/utils/upload-key';
|
||||
import { env } from '~/env/server';
|
||||
import { logToAxiom } from '~/server/logging/client';
|
||||
|
||||
@@ -26,15 +25,9 @@ const upload = async (req: NextApiRequest, res: NextApiResponse) => {
|
||||
}
|
||||
|
||||
const { filename: fullFilename } = req.body;
|
||||
const ext = extname(fullFilename);
|
||||
const filename = filenamize(fullFilename.replace(ext, ''));
|
||||
let { type } = req.body;
|
||||
if (!type || !Object.values(UploadType).includes(type)) type = UploadType.Default;
|
||||
|
||||
if (env.UPLOAD_PROHIBITED_EXTENSIONS?.includes(ext)) {
|
||||
return res.status(400).json({ error: 'File type not allowed' });
|
||||
}
|
||||
|
||||
// Determine upload backend: B2 for model/training uploads when the B2
|
||||
// endpoint is configured (no Flipt flag — see below).
|
||||
let backend: UploadBackend = 'default';
|
||||
@@ -54,7 +47,7 @@ const upload = async (req: NextApiRequest, res: NextApiResponse) => {
|
||||
backend = 'b2';
|
||||
}
|
||||
|
||||
const key = `${type ?? UploadType.Default}/${userId}/${filename}.${generateToken(4)}${ext}`;
|
||||
const key = buildUploadKey(type ?? UploadType.Default, userId, fullFilename);
|
||||
const s3 = backend === 'b2' ? getUploadS3Client('b2') : null;
|
||||
const bucket = backend === 'b2' ? getUploadBucket('b2') : null;
|
||||
|
||||
|
||||
@@ -78,7 +78,7 @@ import { leaderboardJobs } from '~/server/jobs/prepare-leaderboard';
|
||||
// import { processCreatorProgramImageGenerationRewards } from '~/server/jobs/process-creator-program-image-generation-rewards';
|
||||
import { csamJobs } from '~/server/jobs/process-csam';
|
||||
import { processingEngingEarlyAccess } from '~/server/jobs/process-ending-early-access';
|
||||
import { processImportsJob } from '~/server/jobs/process-imports';
|
||||
import { processHuggingFaceImportsJob } from '~/server/jobs/process-huggingface-imports';
|
||||
import { processRewards, rewardsDailyReset } from '~/server/jobs/process-rewards';
|
||||
import { processScheduledPublishing } from '~/server/jobs/process-scheduled-publishing';
|
||||
import { processSubscriptionsRequiringRenewal } from '~/server/jobs/process-subscriptions-requiring-renewal';
|
||||
@@ -129,7 +129,7 @@ import { booleanString } from '~/utils/zod-helpers';
|
||||
|
||||
export const jobs: Job[] = [
|
||||
scanFilesFallbackJob,
|
||||
processImportsJob,
|
||||
processHuggingFaceImportsJob,
|
||||
sendNotificationsJob,
|
||||
notificationCursorMonitor,
|
||||
sendWebhooksJob,
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
import { Container, Stack, Text, Title } from '@mantine/core';
|
||||
import { Page } from '~/components/AppLayout/Page';
|
||||
import { Meta } from '~/components/Meta/Meta';
|
||||
import { ImportQueueTable } from '~/components/Moderation/HuggingFaceImport/ImportQueueTable';
|
||||
import { ImportConfigSection } from '~/components/Moderation/HuggingFaceImport/ImportConfigSection';
|
||||
import { RepoLookupSection } from '~/components/Moderation/HuggingFaceImport/RepoLookupSection';
|
||||
import { createServerSideProps } from '~/server/utils/server-side-helpers';
|
||||
|
||||
export const getServerSideProps = createServerSideProps({ requireModerator: true });
|
||||
|
||||
function HuggingFaceImportPage() {
|
||||
return (
|
||||
<>
|
||||
<Meta title="Hugging Face Import" deIndex />
|
||||
<Container size="lg" py="lg">
|
||||
<Stack gap="xl">
|
||||
<Stack gap={4}>
|
||||
<Title order={2}>Hugging Face Import</Title>
|
||||
<Text c="dimmed" size="sm">
|
||||
Queue model files for transfer from Hugging Face into our storage. A cron job moves
|
||||
them a part at a time, so a transfer survives a deploy and resumes where it stopped.
|
||||
</Text>
|
||||
</Stack>
|
||||
|
||||
<RepoLookupSection />
|
||||
<ImportQueueTable />
|
||||
<ImportConfigSection />
|
||||
</Stack>
|
||||
</Container>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
export default Page(HuggingFaceImportPage);
|
||||
@@ -45,6 +45,8 @@ vi.mock('~/server/auth/get-server-auth-session', () => ({
|
||||
}));
|
||||
|
||||
import handler from '~/pages/api/upload/sign-part';
|
||||
import { UploadType } from '~/server/common/enums';
|
||||
import { buildUploadKey } from '~/utils/upload-key';
|
||||
// Mocked in ~/__tests__/setup as vi.fn(); imported so log ORDER is observable and the logger
|
||||
// can be made to REJECT.
|
||||
import { logToAxiom } from '~/server/logging/client';
|
||||
@@ -86,9 +88,13 @@ function makeReq() {
|
||||
method: 'POST',
|
||||
body: {
|
||||
bucket: 'civitai-modelfiles',
|
||||
// The handler scopes re-signing to `key.split('/')[1] === String(userId)`, so the 42
|
||||
// here is load-bearing: any other value 403s.
|
||||
key: 'model/42/thing.safetensors',
|
||||
// The handler scopes re-signing to `key.split('/')[1] === String(userId)`, so the 42 here is
|
||||
// load-bearing: any other value 403s.
|
||||
//
|
||||
// Built rather than written out, so the minting side and this parsing side cannot drift apart
|
||||
// with both suites green — a literal here is a COPY of the shape, and the shape now lives in
|
||||
// another module.
|
||||
key: buildUploadKey(UploadType.Model, 42, 'thing.safetensors'),
|
||||
uploadId: 'test-upload-id',
|
||||
partNumber: 3,
|
||||
backend: 'b2',
|
||||
|
||||
@@ -1,45 +0,0 @@
|
||||
import { ImportStatus } from '~/shared/utils/prisma/enums';
|
||||
import { createImporter } from '~/server/importers/importer';
|
||||
|
||||
// Find match for URL like: https://huggingface.co/nitrosocke/Arcane-Diffusion
|
||||
const hfAuthorRegex = /^https:\/\/huggingface\.co\/([\w\-]+)$/;
|
||||
export const hfAuthorImporter = createImporter(
|
||||
(source) => {
|
||||
return hfAuthorRegex.test(source);
|
||||
},
|
||||
async ({ source }) => {
|
||||
// Get the author and model name from the URL
|
||||
const [, author] = hfAuthorRegex.exec(source) ?? [];
|
||||
|
||||
// Get the model from HuggingFace
|
||||
const hfModels = await getHuggingFaceModels(author);
|
||||
|
||||
return {
|
||||
status: ImportStatus.Completed,
|
||||
dependencies: hfModels.map((hfModel) => ({
|
||||
source: `https://huggingface.co/${hfModel.id}`,
|
||||
data: hfModel,
|
||||
})),
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
async function getHuggingFaceModels(author: string) {
|
||||
const result = (await fetch(`https://huggingface.co/api/models?author=${author}&full=true`).then(
|
||||
(r) => r.json()
|
||||
)) as HuggingFaceModel[];
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
type HuggingFaceModel = {
|
||||
id: string;
|
||||
author: string;
|
||||
lastModified: string;
|
||||
tags: string[];
|
||||
downloads: number;
|
||||
likes: number;
|
||||
siblings: {
|
||||
rfilename: string;
|
||||
}[];
|
||||
};
|
||||
@@ -1,223 +0,0 @@
|
||||
import type { Prisma } from '@prisma/client';
|
||||
|
||||
import { ImportStatus, ModelType } from '~/shared/utils/prisma/enums';
|
||||
import { createImporter } from '~/server/importers/importer';
|
||||
import { dbWrite } from '~/server/db/client';
|
||||
import { uploadViaUrl } from '~/utils/cf-images-utils';
|
||||
import { markdownToHtml } from '~/utils/markdown-helpers';
|
||||
import { bytesToKB } from '~/utils/number-helpers';
|
||||
import { imageToBlurhash } from '~/utils/image-utils';
|
||||
import { getModelFileFormat } from '~/utils/file-helpers';
|
||||
|
||||
// Find match for URL like: https://huggingface.co/nitrosocke/Arcane-Diffusion
|
||||
const hfModelRegex = /^https:\/\/huggingface\.co\/([\w\-\.]+)\/([\w\-\.]+)/;
|
||||
export const hfModelImporter = createImporter(
|
||||
(source) => {
|
||||
return hfModelRegex.test(source);
|
||||
},
|
||||
async ({ id, source, data, userId }) => {
|
||||
userId ??= -1;
|
||||
// Get the author and model name from the URL
|
||||
const [, author, modelName] = hfModelRegex.exec(source) ?? [];
|
||||
|
||||
// Get the model from HuggingFace
|
||||
let hfModel: HuggingFaceModel | undefined = data;
|
||||
if (!hfModel) {
|
||||
try {
|
||||
hfModel = await getHuggingFaceModel(author, modelName);
|
||||
} catch (error) {
|
||||
throw new Error(`Could not find model ${author}/${modelName}`);
|
||||
}
|
||||
}
|
||||
|
||||
await importModelFromHuggingFace(hfModel, { id, source, userId });
|
||||
|
||||
return {
|
||||
status: ImportStatus.Completed,
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
// TODO.remove
|
||||
async function importModelFromHuggingFace(
|
||||
{ id, siblings, author }: HuggingFaceModel,
|
||||
{ id: importId, source, userId }: { id?: number; source?: string; userId: number }
|
||||
) {
|
||||
const hfRootUrl = `https://huggingface.co/${id}/resolve/main/`;
|
||||
const files = siblings.map((x) => ({
|
||||
name: x.rfilename,
|
||||
url: hfRootUrl + x.rfilename,
|
||||
}));
|
||||
|
||||
// check for previous models imported from same hfModel.id
|
||||
let model = await dbWrite.model.findFirst({
|
||||
where: { fromImport: { source } },
|
||||
select: { id: true, modelVersions: { select: { files: true } } },
|
||||
});
|
||||
|
||||
// Prepare modelVersions files
|
||||
// for each file in the model, create a modelVersion on the model
|
||||
const modelVersions: Prisma.ModelVersionUncheckedCreateInput[] = [];
|
||||
let type: ModelType = ModelType.Checkpoint;
|
||||
for (const { name, url } of files) {
|
||||
// TODO Import: Improve this to handle models that aren't saved as `.ckpt` or `.safetensors`
|
||||
// Example: https://huggingface.co/sd-dreambooth-library/the-witcher-game-ciri/tree/main
|
||||
if (!isModelFile(name)) continue;
|
||||
|
||||
const existingVersion = model?.modelVersions.find((v) => v.files.some((f) => f.name === name));
|
||||
if (existingVersion) continue;
|
||||
|
||||
// HEAD the file to get the size
|
||||
const { headers } = await fetch(url, { method: 'HEAD' });
|
||||
const size = bytesToKB(parseInt(headers.get('Content-Length') ?? '0'));
|
||||
type = fileToModelType(name, size);
|
||||
|
||||
modelVersions.push({
|
||||
modelId: 0,
|
||||
name: filenameToVersionName(name, id),
|
||||
fromImportId: importId,
|
||||
baseModel: 'SD 1.5',
|
||||
files: {
|
||||
create: [
|
||||
{
|
||||
url,
|
||||
sizeKB: size,
|
||||
name,
|
||||
type: 'Model',
|
||||
format: getModelFileFormat(name),
|
||||
} as Prisma.ModelFileCreateWithoutModelVersionInput,
|
||||
],
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// If there aren't versions, there's nothing for us to do...
|
||||
if (modelVersions.length === 0) return;
|
||||
|
||||
// Prep image and description if needed
|
||||
const imagesToCreate: Prisma.ImageUncheckedCreateInput[] = [];
|
||||
let description = `<p>Originally posted to <a href="https://huggingface.co/${id}">HuggingFace by ${author}</a></p>`;
|
||||
if (!model) {
|
||||
// Get README
|
||||
try {
|
||||
const readme = await fetch(hfRootUrl + 'README.md').then((r) => r.text());
|
||||
description += await markdownToHtml(readme);
|
||||
} catch (error) {
|
||||
// This is fine... 🔥
|
||||
}
|
||||
|
||||
// Upload images
|
||||
const imageFiles = files.filter((f) => isImage(f.name));
|
||||
if (imageFiles.length === 0)
|
||||
// if no images, use the default
|
||||
imageFiles.push({
|
||||
name: 'default.png',
|
||||
url: `https://thumbnails.huggingface.co/social-thumbnails/models/${id}.png`,
|
||||
});
|
||||
|
||||
// Process images (limit to 20)
|
||||
for (const { name, url } of imageFiles.slice(0, 20)) {
|
||||
try {
|
||||
const { hash, height, width } = await imageToBlurhash(url);
|
||||
const { id } = await uploadViaUrl(url, {
|
||||
userId,
|
||||
source: 'huggingface',
|
||||
});
|
||||
imagesToCreate.push({ name, url: id, userId, hash, height, width });
|
||||
} catch (error) {
|
||||
console.error(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
await dbWrite.$transaction(
|
||||
async (tx) => {
|
||||
// if it doesn't exist, create it
|
||||
if (!model) {
|
||||
// Create model
|
||||
model = await tx.model.create({
|
||||
data: {
|
||||
name: id.split('/').pop() ?? id,
|
||||
description,
|
||||
fromImportId: importId,
|
||||
type,
|
||||
userId,
|
||||
lastVersionAt: new Date(),
|
||||
},
|
||||
select: { id: true, modelVersions: { select: { files: true } } },
|
||||
});
|
||||
}
|
||||
|
||||
// ! - commented out for type issues
|
||||
// for (const data of modelVersions) {
|
||||
// const versionImages = [];
|
||||
// for (const data of imagesToCreate) {
|
||||
// const image = await tx.image.create({
|
||||
// data,
|
||||
// select: { id: true },
|
||||
// });
|
||||
// versionImages.push(image);
|
||||
// }
|
||||
|
||||
// data.modelId = model.id;
|
||||
// data.images = {
|
||||
// create: versionImages.map((image, index) => ({ imageId: image.id, index })),
|
||||
// };
|
||||
// await tx.modelVersion.create({ data });
|
||||
// }
|
||||
},
|
||||
{
|
||||
maxWait: 10000,
|
||||
timeout: 30000,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
async function getHuggingFaceModel(author: string, modelName: string) {
|
||||
const result = (await fetch(`https://huggingface.co/api/models/${author}/${modelName}`).then(
|
||||
(r) => r.json()
|
||||
)) as HuggingFaceModel;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
function filenameToVersionName(filename: string, hfModelId: string) {
|
||||
const modelName = hfModelId.split('/')[1];
|
||||
const versionName = filename
|
||||
.replace(modelName, '')
|
||||
.replace(modelFileRegex, '')
|
||||
.replace(/[-_]/g, ' ')
|
||||
.trim();
|
||||
return versionName;
|
||||
}
|
||||
|
||||
function isImage(filename: string) {
|
||||
return /\.(png|gif|jpg|jpeg)$/.test(filename);
|
||||
}
|
||||
|
||||
const modelFileRegex = /\.(ckpt|pt|bin|safetensors)$/;
|
||||
function isModelFile(filename: string) {
|
||||
if (filename.endsWith('pytorch_model.bin')) return false;
|
||||
return modelFileRegex.test(filename);
|
||||
}
|
||||
|
||||
function fileToModelType(filename: string, sizeKB: number) {
|
||||
if (/\.(pt|bin)$/.test(filename)) {
|
||||
if (sizeKB > 10 * 1000) return ModelType.Hypernetwork;
|
||||
if (sizeKB < 1000) return ModelType.TextualInversion;
|
||||
// TODO ModelType Importing: determine some way of determining if something is a Aesthetic Gradient or TI
|
||||
}
|
||||
return ModelType.Checkpoint;
|
||||
}
|
||||
|
||||
type HuggingFaceModel = {
|
||||
id: string;
|
||||
author: string;
|
||||
lastModified: string;
|
||||
tags: string[];
|
||||
downloads: number;
|
||||
likes: number;
|
||||
siblings: {
|
||||
rfilename: string;
|
||||
}[];
|
||||
};
|
||||
@@ -1,67 +0,0 @@
|
||||
import { Prisma } from '@prisma/client';
|
||||
import { hfModelImporter } from '~/server/importers/huggingFaceModel';
|
||||
import { dbWrite } from '~/server/db/client';
|
||||
import { ImportStatus } from '~/shared/utils/prisma/enums';
|
||||
import { hfAuthorImporter } from '~/server/importers/huggingFaceAuthor';
|
||||
import type { ImportDependency, ImportRunInput } from '~/server/importers/importer';
|
||||
import { chunk } from 'lodash-es';
|
||||
|
||||
const importers = [hfModelImporter, hfAuthorImporter];
|
||||
|
||||
export async function processImport(input: ImportRunInput) {
|
||||
const { id, source } = input;
|
||||
const importer = importers.find((i) => i.canHandle(source));
|
||||
|
||||
const updateStatus = async (status: ImportStatus, data: any = null) => {
|
||||
// eslint-disable-line
|
||||
await dbWrite.import.update({
|
||||
where: { id },
|
||||
data: { status, data: data ?? Prisma.JsonNull },
|
||||
});
|
||||
return { id, status, data };
|
||||
};
|
||||
|
||||
if (!importer) {
|
||||
return await updateStatus(ImportStatus.Failed, { error: 'No importer found' });
|
||||
}
|
||||
|
||||
await updateStatus(ImportStatus.Processing);
|
||||
try {
|
||||
const { status, data, dependencies } = await importer.run(input);
|
||||
if (dependencies) await processDependencies(input, dependencies);
|
||||
return await updateStatus(status, data);
|
||||
} catch (error: any) {
|
||||
// eslint-disable-line
|
||||
console.error(error);
|
||||
return await updateStatus(ImportStatus.Failed, { error: error.message, stack: error.stack });
|
||||
}
|
||||
}
|
||||
|
||||
async function processDependencies(
|
||||
{ userId, id: parentId }: ImportRunInput,
|
||||
deps: ImportDependency[]
|
||||
) {
|
||||
// Add the import jobs
|
||||
for (const batch of chunk(deps, 900)) {
|
||||
await dbWrite.import.createMany({
|
||||
data: batch.map(({ source, data }) => ({
|
||||
source,
|
||||
userId,
|
||||
parentId,
|
||||
data,
|
||||
})),
|
||||
});
|
||||
}
|
||||
|
||||
const childJobs = await dbWrite.import.findMany({
|
||||
where: {
|
||||
parentId,
|
||||
},
|
||||
});
|
||||
|
||||
for (const batch of chunk(childJobs, 10)) {
|
||||
try {
|
||||
await Promise.all(batch.map((job) => processImport(job)));
|
||||
} catch (e) {} // We handle this inside the processImport...
|
||||
}
|
||||
}
|
||||
@@ -1,31 +0,0 @@
|
||||
import type { ImportStatus } from '~/shared/utils/prisma/enums';
|
||||
|
||||
type Importer = {
|
||||
canHandle: (source: string) => boolean;
|
||||
run: (input: ImportRunInput) => Promise<ImportResult>;
|
||||
};
|
||||
|
||||
export type ImportRunInput = {
|
||||
id: number;
|
||||
source: string;
|
||||
userId?: number | null;
|
||||
data?: any; // eslint-disable-line
|
||||
};
|
||||
|
||||
export type ImportDependency = {
|
||||
source: string;
|
||||
data?: any; // eslint-disable-line
|
||||
};
|
||||
|
||||
type ImportResult = {
|
||||
status: ImportStatus;
|
||||
data?: any; // eslint-disable-line
|
||||
dependencies?: ImportDependency[];
|
||||
};
|
||||
|
||||
export function createImporter(canHandle: Importer['canHandle'], run: Importer['run']): Importer {
|
||||
return {
|
||||
canHandle,
|
||||
run,
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
import { randomUUID } from 'crypto';
|
||||
import { createJob } from './job';
|
||||
import { getHuggingFaceImportConfig } from '~/server/services/huggingface-import-config.service';
|
||||
import { processImportQueue } from '~/server/services/huggingface-import.service';
|
||||
|
||||
/**
|
||||
* The per-run budget is operator-set (`workBudgetSeconds`, capped at 240) and must leave room for the
|
||||
* SLOWEST part still in flight when it expires — otherwise the lock lapses mid-run and the next tick
|
||||
* starts a second run. This lock is the ceiling that cap is chosen against.
|
||||
*/
|
||||
const LOCK_EXPIRATION_SECONDS = 5 * 60;
|
||||
|
||||
export const processHuggingFaceImportsJob = createJob(
|
||||
'process-huggingface-imports',
|
||||
'* * * * *',
|
||||
async () => {
|
||||
const config = await getHuggingFaceImportConfig();
|
||||
// The kill switch stops this run claiming anything; queued rows are left exactly as they are, so
|
||||
// turning it back on resumes rather than restarts.
|
||||
if (!config.enabled) return { skipped: 'disabled' };
|
||||
|
||||
const { moved, bytes } = await processImportQueue({
|
||||
deadline: Date.now() + config.workBudgetSeconds * 1000,
|
||||
worker: `${process.env.HOSTNAME ?? 'local'}:${randomUUID().slice(0, 8)}`,
|
||||
concurrency: config.filesInParallel,
|
||||
partsInFlight: config.partsInFlight,
|
||||
});
|
||||
return { moved, bytes };
|
||||
},
|
||||
// `keepLockOnDisconnect`: this run legitimately outlives the scheduler's client timeout. A second
|
||||
// run would not collide — the lock is fleet-wide, and `SKIP LOCKED` plus the heartbeat send it to
|
||||
// different files — it would simply transfer more at once, on whichever pod the scheduler hit.
|
||||
{ lockExpiration: LOCK_EXPIRATION_SECONDS, keepLockOnDisconnect: true }
|
||||
);
|
||||
@@ -1,23 +0,0 @@
|
||||
import { createJob } from './job';
|
||||
import { dbWrite } from '~/server/db/client';
|
||||
import { ImportStatus } from '~/shared/utils/prisma/enums';
|
||||
import dayjs from '~/shared/utils/dayjs';
|
||||
import { chunk } from 'lodash-es';
|
||||
import { processImport } from '~/server/importers/importRouter';
|
||||
|
||||
export const processImportsJob = createJob('process-imports', '1 */1 * * *', async () => {
|
||||
// Get pending import jobs that are older than 30 minutes
|
||||
const importJobs = await dbWrite.import.findMany({
|
||||
where: {
|
||||
status: ImportStatus.Pending,
|
||||
createdAt: { lt: dayjs().add(-30, 'minutes').toDate() },
|
||||
},
|
||||
});
|
||||
|
||||
// Process the pending jobs
|
||||
for (const batch of chunk(importJobs, 10)) {
|
||||
try {
|
||||
await Promise.all(batch.map((job) => processImport(job)));
|
||||
} catch (e) {} // We handle this inside the processImport...
|
||||
}
|
||||
});
|
||||
@@ -0,0 +1,141 @@
|
||||
import { createFileHandler } from '~/server/controllers/model-file.controller';
|
||||
import { getByIdSchema } from '~/server/schema/base.schema';
|
||||
import {
|
||||
attachHuggingFaceImportSchema,
|
||||
renameHuggingFaceGroupSchema,
|
||||
setHuggingFaceImportConfigSchema,
|
||||
enqueueHuggingFaceImportSchema,
|
||||
getHuggingFaceImportCountsSchema,
|
||||
getHuggingFaceImportsSchema,
|
||||
lookupHuggingFaceRepoSchema,
|
||||
} from '~/server/schema/huggingface-import.schema';
|
||||
import {
|
||||
buildAttachInput,
|
||||
cancelImport,
|
||||
deleteImport,
|
||||
detachImport,
|
||||
enqueueImports,
|
||||
getImportCounts,
|
||||
getImports,
|
||||
linkImportToFile,
|
||||
renameGroup,
|
||||
resolveRepoForImport,
|
||||
retryImport,
|
||||
} from '~/server/services/huggingface-import.service';
|
||||
import { HuggingFaceError, parseHuggingFaceRepo } from '~/server/services/huggingface.service';
|
||||
import {
|
||||
getHuggingFaceImportConfig,
|
||||
setHuggingFaceImportConfig,
|
||||
} from '~/server/services/huggingface-import-config.service';
|
||||
import { moderatorProcedure, router } from '~/server/trpc';
|
||||
import { throwBadRequestError } from '~/server/utils/errorHandling';
|
||||
|
||||
/**
|
||||
* 🔴 Opening this beyond moderators needs more than a procedure swap: a per-user quota, and a
|
||||
* `userId` in the `(repo, revision, filename)` unique index — without that last one, one user's
|
||||
* queued import silently swallows another's.
|
||||
*/
|
||||
export const huggingFaceImportRouter = router({
|
||||
getAll: moderatorProcedure.input(getHuggingFaceImportsSchema).query(({ input, ctx }) =>
|
||||
getImports({
|
||||
...input,
|
||||
userId: ctx.user.id,
|
||||
isModerator: !!ctx.user.isModerator,
|
||||
})
|
||||
),
|
||||
|
||||
lookup: moderatorProcedure.input(lookupHuggingFaceRepoSchema).mutation(async ({ input }) => {
|
||||
const target = parseHuggingFaceRepo(input.source);
|
||||
if (!target)
|
||||
throw throwBadRequestError(
|
||||
'Could not read an owner/name out of that. Paste the model page URL.'
|
||||
);
|
||||
|
||||
try {
|
||||
return await resolveRepoForImport(target);
|
||||
} catch (error) {
|
||||
if (error instanceof HuggingFaceError) throw throwBadRequestError(error.message);
|
||||
throw error;
|
||||
}
|
||||
}),
|
||||
|
||||
enqueue: moderatorProcedure
|
||||
.input(enqueueHuggingFaceImportSchema)
|
||||
.mutation(async ({ input, ctx }) => {
|
||||
// The tree is re-read rather than trusting the sizes and hashes the client posted back: they
|
||||
// decide what we store and what we skip as already held.
|
||||
const repo = await resolveRepoForImport({ repo: input.repo, revision: input.revision });
|
||||
const result = await enqueueImports({
|
||||
repo,
|
||||
paths: input.paths,
|
||||
userId: ctx.user.id,
|
||||
groupName: input.groupName,
|
||||
});
|
||||
if (!result.queued && !result.skipped)
|
||||
throw throwBadRequestError('None of those files exist at that revision.');
|
||||
return result;
|
||||
}),
|
||||
|
||||
attach: moderatorProcedure
|
||||
.input(attachHuggingFaceImportSchema)
|
||||
.mutation(async ({ input, ctx }) => {
|
||||
const { importId, ...fileInput } = await buildAttachInput({
|
||||
...input,
|
||||
userId: ctx.user.id,
|
||||
isModerator: !!ctx.user.isModerator,
|
||||
});
|
||||
const file = await createFileHandler({ input: fileInput, ctx });
|
||||
const linked = await linkImportToFile({
|
||||
id: importId,
|
||||
modelFileId: file.id,
|
||||
modelVersionId: input.modelVersionId,
|
||||
});
|
||||
if (!linked)
|
||||
throw throwBadRequestError(
|
||||
`Created model file ${file.id}, but this import was attached by someone else first. Delete file ${file.id}.`
|
||||
);
|
||||
return { modelFileId: file.id, modelVersionId: input.modelVersionId };
|
||||
}),
|
||||
|
||||
getCounts: moderatorProcedure
|
||||
.input(getHuggingFaceImportCountsSchema)
|
||||
.query(({ input, ctx }) =>
|
||||
getImportCounts({ ...input, userId: ctx.user.id, isModerator: !!ctx.user.isModerator })
|
||||
),
|
||||
|
||||
delete: moderatorProcedure
|
||||
.input(getByIdSchema)
|
||||
.mutation(({ input, ctx }) =>
|
||||
deleteImport({ id: input.id, userId: ctx.user.id, isModerator: !!ctx.user.isModerator })
|
||||
),
|
||||
|
||||
getConfig: moderatorProcedure.query(() => getHuggingFaceImportConfig()),
|
||||
|
||||
setConfig: moderatorProcedure
|
||||
.input(setHuggingFaceImportConfigSchema)
|
||||
.mutation(({ input }) => setHuggingFaceImportConfig(input)),
|
||||
|
||||
renameGroup: moderatorProcedure
|
||||
.input(renameHuggingFaceGroupSchema)
|
||||
.mutation(({ input, ctx }) =>
|
||||
renameGroup({ ...input, userId: ctx.user.id, isModerator: !!ctx.user.isModerator })
|
||||
),
|
||||
|
||||
detach: moderatorProcedure
|
||||
.input(getByIdSchema)
|
||||
.mutation(({ input, ctx }) =>
|
||||
detachImport({ id: input.id, userId: ctx.user.id, isModerator: !!ctx.user.isModerator })
|
||||
),
|
||||
|
||||
retry: moderatorProcedure
|
||||
.input(getByIdSchema)
|
||||
.mutation(({ input, ctx }) =>
|
||||
retryImport({ id: input.id, userId: ctx.user.id, isModerator: !!ctx.user.isModerator })
|
||||
),
|
||||
|
||||
cancel: moderatorProcedure
|
||||
.input(getByIdSchema)
|
||||
.mutation(({ input, ctx }) =>
|
||||
cancelImport({ id: input.id, userId: ctx.user.id, isModerator: !!ctx.user.isModerator })
|
||||
),
|
||||
});
|
||||
@@ -36,6 +36,9 @@ export const appRouter = router({
|
||||
download: lazy(() => import('./download.router').then((m) => m.downloadRouter)),
|
||||
feedback: lazy(() => import('./feedback.router').then((m) => m.feedbackRouter)),
|
||||
homeBlock: lazy(() => import('./home-block.router').then((m) => m.homeBlockRouter)),
|
||||
huggingFaceImport: lazy(() =>
|
||||
import('./huggingface-import.router').then((m) => m.huggingFaceImportRouter)
|
||||
),
|
||||
image: lazy(() => import('./image.router').then((m) => m.imageRouter)),
|
||||
merch: lazy(() => import('./merch.router').then((m) => m.merchRouter)),
|
||||
model: lazy(() => import('./model.router').then((m) => m.modelRouter)),
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
import * as z from 'zod';
|
||||
import { constants } from '~/server/common/constants';
|
||||
import { huggingFaceImportConfigSchema } from '~/server/services/huggingface-import-config.service';
|
||||
|
||||
export type LookupHuggingFaceRepoInput = z.infer<typeof lookupHuggingFaceRepoSchema>;
|
||||
export const lookupHuggingFaceRepoSchema = z.object({
|
||||
source: z.string().trim().min(1),
|
||||
});
|
||||
|
||||
export type EnqueueHuggingFaceImportInput = z.infer<typeof enqueueHuggingFaceImportSchema>;
|
||||
export const enqueueHuggingFaceImportSchema = z.object({
|
||||
repo: z.string().trim().min(1),
|
||||
revision: z.string().trim().min(1),
|
||||
paths: z.array(z.string().min(1)).min(1).max(100),
|
||||
groupName: z.string().trim().min(1).max(120).optional(),
|
||||
});
|
||||
|
||||
export type GetHuggingFaceImportsInput = z.infer<typeof getHuggingFaceImportsSchema>;
|
||||
export const getHuggingFaceImportsSchema = z.object({
|
||||
limit: z.number().int().min(1).max(200).default(100),
|
||||
/** Substring, case-insensitive — what a moderator types into the filter box. */
|
||||
groupName: z.string().trim().min(1).max(120).optional(),
|
||||
/** Exact `owner/name`, as Hugging Face reports it. What the skill filters by. */
|
||||
repo: z.string().trim().min(1).optional(),
|
||||
/** Completed transfers no model version has claimed. */
|
||||
unattached: z.boolean().optional(),
|
||||
});
|
||||
|
||||
export type GetHuggingFaceImportCountsInput = z.infer<typeof getHuggingFaceImportCountsSchema>;
|
||||
export const getHuggingFaceImportCountsSchema = getHuggingFaceImportsSchema.pick({
|
||||
groupName: true,
|
||||
repo: true,
|
||||
});
|
||||
|
||||
export type AttachHuggingFaceImportInput = z.infer<typeof attachHuggingFaceImportSchema>;
|
||||
export const attachHuggingFaceImportSchema = z.object({
|
||||
id: z.number().int().positive(),
|
||||
modelVersionId: z.number().int().positive(),
|
||||
// Explicit, never inferred: this is what decides whether the version is loadable.
|
||||
type: z.enum(constants.modelFileTypes),
|
||||
});
|
||||
|
||||
export type RenameHuggingFaceGroupInput = z.infer<typeof renameHuggingFaceGroupSchema>;
|
||||
export const renameHuggingFaceGroupSchema = z.object({
|
||||
repo: z.string().trim().min(1),
|
||||
revision: z.string().trim().min(1),
|
||||
/** The group's current name. Not trimmed: it must match the stored value exactly. */
|
||||
from: z.string().min(1).max(120),
|
||||
groupName: z.string().trim().min(1).max(120),
|
||||
});
|
||||
|
||||
export type SetHuggingFaceImportConfigInput = z.infer<typeof setHuggingFaceImportConfigSchema>;
|
||||
export const setHuggingFaceImportConfigSchema = huggingFaceImportConfigSchema.partial();
|
||||
@@ -0,0 +1,890 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import { dbMock } from '~/__tests__/mocks/db.mock';
|
||||
import { redisMock } from '~/__tests__/mocks/redis.mock';
|
||||
import type * as HuggingFaceService from '~/server/services/huggingface.service';
|
||||
import type * as S3Utils from '~/utils/s3-utils';
|
||||
|
||||
const {
|
||||
mockReadRange,
|
||||
mockHeadFile,
|
||||
mockUploadPart,
|
||||
mockCreateMultipart,
|
||||
mockComplete,
|
||||
mockAbort,
|
||||
mockDeleteObject,
|
||||
mockUrlsSafeToDelete,
|
||||
} = vi.hoisted(() => ({
|
||||
mockReadRange: vi.fn(),
|
||||
mockHeadFile: vi.fn(),
|
||||
mockUploadPart: vi.fn(),
|
||||
mockCreateMultipart: vi.fn(),
|
||||
mockComplete: vi.fn(),
|
||||
mockAbort: vi.fn(),
|
||||
mockDeleteObject: vi.fn(),
|
||||
mockUrlsSafeToDelete: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('~/server/services/huggingface.service', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof HuggingFaceService>()),
|
||||
readHuggingFaceRange: mockReadRange,
|
||||
headHuggingFaceFile: mockHeadFile,
|
||||
}));
|
||||
|
||||
vi.mock('~/utils/s3-utils', async (importOriginal) => ({
|
||||
...(await importOriginal<typeof S3Utils>()),
|
||||
getS3Client: () => ({}),
|
||||
getUploadS3Client: () => ({}),
|
||||
getUploadBucket: () => 'model-bucket',
|
||||
getBucket: async () => 'model-bucket',
|
||||
getGetUrlByKey: async (key: string, opts: { bucket?: string }) => ({
|
||||
key,
|
||||
bucket: opts.bucket,
|
||||
url: `https://s3.example/${opts.bucket}/${key}?X-Amz-Signature=abc`,
|
||||
}),
|
||||
createMultipartUpload: mockCreateMultipart,
|
||||
uploadPart: mockUploadPart,
|
||||
completeMultipartUpload: mockComplete,
|
||||
abortMultipartUpload: mockAbort,
|
||||
deleteObject: mockDeleteObject,
|
||||
urlsSafeToDelete: mockUrlsSafeToDelete,
|
||||
}));
|
||||
|
||||
import { parseHuggingFaceRepo, suggestFileType } from '~/server/services/huggingface.service';
|
||||
import {
|
||||
getHuggingFaceImportConfig,
|
||||
HUGGING_FACE_IMPORT_DEFAULTS,
|
||||
setHuggingFaceImportConfig,
|
||||
} from '~/server/services/huggingface-import-config.service';
|
||||
import {
|
||||
deleteImport,
|
||||
getImportCounts,
|
||||
getImports,
|
||||
PART_SIZE_BYTES,
|
||||
processImportQueue,
|
||||
renameGroup,
|
||||
} from '~/server/services/huggingface-import.service';
|
||||
|
||||
/** The width under test. Passed in rather than read from config, so these tests do not depend on
|
||||
* what an operator has set in Redis. */
|
||||
const TEST_PARTS_IN_FLIGHT = 3;
|
||||
|
||||
const sysRedisMock = redisMock.sysRedis;
|
||||
const dbWrite = dbMock.dbWrite;
|
||||
const dbRead = dbMock.dbRead;
|
||||
|
||||
// Tracks the real constant rather than restating it — a part-size change should not need a test edit.
|
||||
const PART_SIZE = PART_SIZE_BYTES;
|
||||
|
||||
/** A claim that yields the given row once, then nothing — so the drain loop always terminates. */
|
||||
function claimOnce(row: Record<string, unknown>) {
|
||||
let served = false;
|
||||
dbWrite.$queryRaw.mockImplementation(async () => {
|
||||
if (served) return [];
|
||||
served = true;
|
||||
return [row];
|
||||
});
|
||||
}
|
||||
|
||||
function baseRow(overrides: Record<string, unknown> = {}) {
|
||||
return {
|
||||
id: 1,
|
||||
repo: 'owner/name',
|
||||
filename: 'model.safetensors',
|
||||
sourceUrl: 'https://huggingface.co/owner/name/resolve/abc123/model.safetensors',
|
||||
sizeBytes: BigInt(PART_SIZE * 2 + 100),
|
||||
status: 'Transferring',
|
||||
uploadId: null,
|
||||
partSize: null,
|
||||
parts: null,
|
||||
bucket: null,
|
||||
key: null,
|
||||
attempts: 0,
|
||||
userId: 7,
|
||||
claimedBy: 'test-worker',
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe('parseHuggingFaceRepo', () => {
|
||||
it.each([
|
||||
['https://huggingface.co/owner/name', 'owner/name', undefined],
|
||||
['https://huggingface.co/owner/name/tree/abc123', 'owner/name', 'abc123'],
|
||||
['https://huggingface.co/owner/name/blob/abc123/model.safetensors', 'owner/name', 'abc123'],
|
||||
['https://huggingface.co/models/owner/name', 'owner/name', undefined],
|
||||
['owner/name', 'owner/name', undefined],
|
||||
])('reads %s', (input, repo, revision) => {
|
||||
expect(parseHuggingFaceRepo(input)).toEqual(revision ? { repo, revision } : { repo });
|
||||
});
|
||||
|
||||
it.each(['', 'https://huggingface.co/owner', 'not a url'])('rejects %s', (input) => {
|
||||
expect(parseHuggingFaceRepo(input)).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe('processImportQueue', () => {
|
||||
// Restoring here, not at the end of the test that spies on Date.now: an assertion throwing above
|
||||
// that line would leave the clock frozen for every later test in the file.
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockCreateMultipart.mockResolvedValue('upload-1');
|
||||
mockUploadPart.mockImplementation(
|
||||
async ({ partNumber }: { partNumber: number }) => `etag-${partNumber}`
|
||||
);
|
||||
mockComplete.mockResolvedValue(undefined);
|
||||
mockReadRange.mockImplementation(
|
||||
async ({ start, end }: { start: number; end: number }) => new Uint8Array(end - start + 1)
|
||||
);
|
||||
dbRead.huggingFaceImport.findUnique.mockResolvedValue({ status: 'Transferring' });
|
||||
// The pre-complete re-read goes to the PRIMARY: a cancel written there must be visible before
|
||||
// the upload is finalized, and replica lag would make the check decorative.
|
||||
dbWrite.huggingFaceImport.findUnique.mockResolvedValue({
|
||||
status: 'Transferring',
|
||||
claimedBy: 'test-worker',
|
||||
});
|
||||
dbWrite.huggingFaceImport.update.mockResolvedValue({});
|
||||
dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 });
|
||||
});
|
||||
|
||||
it('splits the file into parts on exact byte boundaries and completes the upload', async () => {
|
||||
claimOnce(baseRow());
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
const ranges = mockReadRange.mock.calls.map(([arg]) => [arg.start, arg.end]);
|
||||
expect(ranges).toEqual([
|
||||
[0, PART_SIZE - 1],
|
||||
[PART_SIZE, PART_SIZE * 2 - 1],
|
||||
// The tail part is short, and its end is the LAST byte — an off-by-one here silently truncates
|
||||
// or over-reads every import.
|
||||
[PART_SIZE * 2, PART_SIZE * 2 + 99],
|
||||
]);
|
||||
|
||||
expect(mockComplete).toHaveBeenCalledWith(
|
||||
'model-bucket',
|
||||
expect.stringMatching(/^model\/7\/model\./),
|
||||
'upload-1',
|
||||
[
|
||||
{ PartNumber: 1, ETag: 'etag-1' },
|
||||
{ PartNumber: 2, ETag: 'etag-2' },
|
||||
{ PartNumber: 3, ETag: 'etag-3' },
|
||||
],
|
||||
expect.anything()
|
||||
);
|
||||
});
|
||||
|
||||
it('heartbeats on every part so a live transfer is never re-claimed as stale', async () => {
|
||||
claimOnce(baseRow());
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
// Two runs holding their own parts arrays against one uploadId is what the heartbeat prevents;
|
||||
// no other assertion in this file fails if the write is deleted.
|
||||
const partWrites = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg)
|
||||
.filter((arg: { data: Record<string, unknown> }) => 'parts' in arg.data);
|
||||
expect(partWrites).toHaveLength(3);
|
||||
for (const write of partWrites) {
|
||||
expect(write.data.heartbeatAt).toBeInstanceOf(Date);
|
||||
expect(write.where.claimedBy).toBe('test-worker');
|
||||
}
|
||||
});
|
||||
|
||||
it('stores the object URL with the presigning query stripped', async () => {
|
||||
claimOnce(baseRow());
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
const completed = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg.data)
|
||||
.find((data: Record<string, unknown>) => data.status === 'Completed');
|
||||
expect(completed?.url).toMatch(/^https:\/\/s3\.example\/model-bucket\/model\/7\/model\./);
|
||||
expect(completed?.url).not.toContain('?');
|
||||
// Spent. A retained uploadId makes every later abort attempt fail against a finished upload,
|
||||
// which buries the one abort failure that means parts are still billed.
|
||||
expect(completed?.uploadId).toBeNull();
|
||||
|
||||
// The pre-complete re-read is one guard against a late cancel; this predicate is the other, and
|
||||
// it covers the window between that read and this write. Dropping it is invisible to the cancel
|
||||
// test, which the read alone already satisfies.
|
||||
const completedWhere = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg)
|
||||
.find((arg: { data: Record<string, unknown> }) => arg.data.status === 'Completed')?.where;
|
||||
expect(completedWhere).toMatchObject({
|
||||
claimedBy: 'test-worker',
|
||||
status: { not: 'Canceled' },
|
||||
});
|
||||
});
|
||||
|
||||
it('resumes an interrupted transfer at the next part instead of restarting the file', async () => {
|
||||
// A bucket that is NOT what `uploadTarget()` resolves to: a resume must address the bucket the
|
||||
// upload was created in, and with both mocked getters returning the same string this was
|
||||
// unobservable — parts went to the current-config bucket while complete named the row's.
|
||||
claimOnce(
|
||||
baseRow({
|
||||
uploadId: 'upload-1',
|
||||
key: 'model/7/resume-me.safetensors',
|
||||
bucket: 'other-bucket',
|
||||
partSize: PART_SIZE,
|
||||
parts: [{ PartNumber: 1, ETag: 'etag-1' }],
|
||||
})
|
||||
);
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
expect(mockCreateMultipart).not.toHaveBeenCalled();
|
||||
for (const [arg] of mockUploadPart.mock.calls) {
|
||||
expect(arg.bucket).toBe('other-bucket');
|
||||
expect(arg.key).toBe('model/7/resume-me.safetensors');
|
||||
}
|
||||
expect(mockComplete.mock.calls[0][0]).toBe('other-bucket');
|
||||
expect(mockComplete.mock.calls[0][1]).toBe('model/7/resume-me.safetensors');
|
||||
expect(mockReadRange.mock.calls.map(([arg]) => arg.start)).toEqual([PART_SIZE, PART_SIZE * 2]);
|
||||
expect(mockComplete.mock.calls[0][3]).toHaveLength(3);
|
||||
});
|
||||
|
||||
it('stops without completing when the deadline passes mid-file', async () => {
|
||||
claimOnce(baseRow({ sizeBytes: BigInt(PART_SIZE * 6) }));
|
||||
// The clock is moved past the deadline BY the first read rather than by a short real window: a
|
||||
// 1ms budget is missed outright on a loaded box, and the test then reads as "moved nothing".
|
||||
const deadline = Date.now() + 60_000;
|
||||
mockReadRange.mockImplementation(async ({ start, end }: { start: number; end: number }) => {
|
||||
vi.spyOn(Date, 'now').mockReturnValue(deadline + 1);
|
||||
return new Uint8Array(end - start + 1);
|
||||
});
|
||||
|
||||
await processImportQueue({
|
||||
deadline,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
expect(mockComplete).not.toHaveBeenCalled();
|
||||
expect(mockUploadPart.mock.calls.length).toBeGreaterThan(0);
|
||||
expect(mockUploadPart.mock.calls.length).toBeLessThan(6);
|
||||
});
|
||||
|
||||
it('resumes across a HOLE in the completed parts, not from their count', async () => {
|
||||
// Parts finish out of order, so part 2 can be missing while part 3 is done. Treating the array
|
||||
// length as "next part" would re-upload 2 as part 3 and silently corrupt the object.
|
||||
claimOnce(
|
||||
baseRow({
|
||||
uploadId: 'upload-1',
|
||||
key: 'model/7/model.abcd1234.safetensors',
|
||||
bucket: 'model-bucket',
|
||||
partSize: PART_SIZE,
|
||||
parts: [
|
||||
{ PartNumber: 1, ETag: 'etag-1' },
|
||||
{ PartNumber: 3, ETag: 'etag-3' },
|
||||
],
|
||||
})
|
||||
);
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
expect(mockUploadPart.mock.calls.map(([arg]) => arg.partNumber)).toEqual([2]);
|
||||
// Completion demands ascending order regardless of the order they finished in.
|
||||
expect(mockComplete.mock.calls[0][3].map((p: { PartNumber: number }) => p.PartNumber)).toEqual([
|
||||
1, 2, 3,
|
||||
]);
|
||||
});
|
||||
|
||||
it('moves exactly as many parts at a time as it is told to', async () => {
|
||||
claimOnce(baseRow({ sizeBytes: BigInt(PART_SIZE * 6) }));
|
||||
|
||||
// The barrier releases at PARTS_IN_FLIGHT and the assertion demands PARTS_IN_FLIGHT. Asserting
|
||||
// merely ">1" let a 3→2 change pass green: the hatch fired, everything unblocked, peak landed
|
||||
// on 2, and the only evidence was the test taking 250ms instead of 1ms — which nothing reads.
|
||||
let inFlight = 0;
|
||||
let peak = 0;
|
||||
let escaped = false;
|
||||
let release!: () => void;
|
||||
const gate = new Promise<void>((resolve) => {
|
||||
release = resolve;
|
||||
});
|
||||
const escapeHatch = setTimeout(() => {
|
||||
escaped = true;
|
||||
release();
|
||||
}, 250);
|
||||
|
||||
mockReadRange.mockImplementation(async ({ start, end }: { start: number; end: number }) => {
|
||||
inFlight++;
|
||||
peak = Math.max(peak, inFlight);
|
||||
if (inFlight >= TEST_PARTS_IN_FLIGHT) release();
|
||||
await gate;
|
||||
inFlight--;
|
||||
return new Uint8Array(end - start + 1);
|
||||
});
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
clearTimeout(escapeHatch);
|
||||
|
||||
// `escaped` is what makes this non-vacuous: a serial implementation never reaches the barrier,
|
||||
// falls through the hatch, and fails here in ~250ms. `peak === PARTS_IN_FLIGHT` pins that the
|
||||
// pool is as wide as configured — it cannot catch a deliberate change to the constant itself,
|
||||
// which is a sizing decision (see the memory note on PARTS_IN_FLIGHT), not a regression.
|
||||
expect(escaped).toBe(false);
|
||||
expect(peak).toBeGreaterThan(1);
|
||||
expect(peak).toBe(TEST_PARTS_IN_FLIGHT);
|
||||
expect(mockComplete).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('refuses a short range read rather than writing a truncated part', async () => {
|
||||
claimOnce(baseRow());
|
||||
mockReadRange.mockResolvedValue(new Uint8Array(10));
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
expect(mockUploadPart).not.toHaveBeenCalled();
|
||||
const failure = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg.data)
|
||||
.find((data: Record<string, unknown>) => typeof data.error === 'string');
|
||||
expect(failure?.error).toContain('returned 10 bytes');
|
||||
});
|
||||
|
||||
it('refuses to complete an upload that was canceled during the final part', async () => {
|
||||
claimOnce(baseRow());
|
||||
// The workers' probe runs BEFORE each takes its part, so the last worker never checks again.
|
||||
// Without the re-read here, a cancel landing in that window still finalized the object and
|
||||
// stamped Completed with a URL on a row the moderator had stopped.
|
||||
dbWrite.huggingFaceImport.findUnique.mockResolvedValue({
|
||||
status: 'Canceled',
|
||||
claimedBy: 'test-worker',
|
||||
});
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
expect(mockComplete).not.toHaveBeenCalled();
|
||||
const completed = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg.data)
|
||||
.find((data: Record<string, unknown>) => data.status === 'Completed');
|
||||
expect(completed).toBeUndefined();
|
||||
});
|
||||
|
||||
it('refuses to complete when the claim has been taken by another run', async () => {
|
||||
claimOnce(baseRow());
|
||||
dbWrite.huggingFaceImport.findUnique.mockResolvedValue({
|
||||
status: 'Transferring',
|
||||
claimedBy: 'a-different-worker',
|
||||
});
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
expect(mockComplete).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('stops a canceled import at the next part boundary', async () => {
|
||||
claimOnce(baseRow());
|
||||
dbRead.huggingFaceImport.findUnique.mockResolvedValue({ status: 'Canceled' });
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
expect(mockUploadPart).not.toHaveBeenCalled();
|
||||
expect(mockComplete).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('suggestFileType', () => {
|
||||
it.each([
|
||||
// The case that matters: a checkpoint that merely names its bundled VAE is still the checkpoint.
|
||||
['flux1-dev-vae-baked.safetensors', null],
|
||||
['Vaevictis-v1.safetensors', null],
|
||||
['flux1-dev-with-t5.safetensors', null],
|
||||
['ae.safetensors', 'VAE'],
|
||||
['vae/diffusion_pytorch_model.safetensors', 'VAE'],
|
||||
['text_encoder_2/model-00001-of-00002.safetensors', 'Text Encoder'],
|
||||
['t5xxl_fp16.safetensors', 'Text Encoder'],
|
||||
['clip_l.safetensors', 'Text Encoder'],
|
||||
['model_index.json', 'Config'],
|
||||
['flux1-dev.safetensors', null],
|
||||
])('%s -> %s', (path, expected) => {
|
||||
expect(suggestFileType(path)).toBe(expected);
|
||||
});
|
||||
});
|
||||
|
||||
describe('failOrRetry', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockCreateMultipart.mockResolvedValue('upload-1');
|
||||
dbRead.huggingFaceImport.findUnique.mockResolvedValue({ status: 'Transferring' });
|
||||
dbWrite.huggingFaceImport.findUnique.mockResolvedValue({
|
||||
status: 'Transferring',
|
||||
claimedBy: 'test-worker',
|
||||
});
|
||||
dbWrite.huggingFaceImport.update.mockResolvedValue({});
|
||||
dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 });
|
||||
});
|
||||
|
||||
const failingRead = () => mockReadRange.mockRejectedValue(new Error('HF said 503'));
|
||||
|
||||
it('fences the failure write by the claim, exactly as the success path does', async () => {
|
||||
claimOnce(baseRow());
|
||||
failingRead();
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
// Without this the failure path hands the row back without checking the claim is still ours,
|
||||
// and a superseded run can release a row a live run is still transferring — the two-runs-one-
|
||||
// uploadId state the heartbeat guards against on the success path only.
|
||||
const failure = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg)
|
||||
.find((arg: { data: Record<string, unknown> }) => typeof arg.data.error === 'string');
|
||||
expect(failure?.where).toMatchObject({
|
||||
claimedBy: 'test-worker',
|
||||
status: { not: 'Canceled' },
|
||||
});
|
||||
});
|
||||
|
||||
it('backs off rather than giving up while attempts remain', async () => {
|
||||
claimOnce(baseRow({ attempts: 0 }));
|
||||
failingRead();
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
const failure = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg.data)
|
||||
.find((data: Record<string, unknown>) => typeof data.error === 'string');
|
||||
expect(failure?.status).toBe('Transferring');
|
||||
expect(failure?.attempts).toBe(1);
|
||||
expect(failure?.nextAttemptAt).toBeInstanceOf(Date);
|
||||
expect((failure?.nextAttemptAt as Date).getTime()).toBeGreaterThan(Date.now());
|
||||
expect(mockAbort).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('gives up on the last attempt and aborts the multipart upload', async () => {
|
||||
claimOnce(
|
||||
baseRow({
|
||||
attempts: 4,
|
||||
uploadId: 'upload-1',
|
||||
key: 'model/7/x.safetensors',
|
||||
bucket: 'model-bucket',
|
||||
partSize: PART_SIZE,
|
||||
})
|
||||
);
|
||||
failingRead();
|
||||
mockAbort.mockResolvedValue(undefined);
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
// An upload abandoned without an abort keeps every part already written, billed, with nothing
|
||||
// left holding the id needed to free them.
|
||||
expect(mockAbort).toHaveBeenCalledTimes(1);
|
||||
const failure = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg.data)
|
||||
.find((data: Record<string, unknown>) => typeof data.error === 'string');
|
||||
expect(failure?.status).toBe('Failed');
|
||||
expect(failure?.nextAttemptAt).toBeNull();
|
||||
});
|
||||
|
||||
it('keeps the uploadId when the abort itself failed', async () => {
|
||||
claimOnce(
|
||||
baseRow({
|
||||
attempts: 4,
|
||||
uploadId: 'upload-1',
|
||||
key: 'model/7/x.safetensors',
|
||||
bucket: 'model-bucket',
|
||||
partSize: PART_SIZE,
|
||||
})
|
||||
);
|
||||
failingRead();
|
||||
mockAbort.mockRejectedValue(new Error('B2 unavailable'));
|
||||
|
||||
await processImportQueue({
|
||||
deadline: Date.now() + 60_000,
|
||||
worker: 'test',
|
||||
concurrency: 1,
|
||||
partsInFlight: TEST_PARTS_IN_FLIGHT,
|
||||
});
|
||||
|
||||
// Clearing it here would be the one thing that makes the orphaned parts unreclaimable.
|
||||
const failure = dbWrite.huggingFaceImport.updateMany.mock.calls
|
||||
.map(([arg]) => arg.data)
|
||||
.find((data: Record<string, unknown>) => typeof data.error === 'string');
|
||||
expect(failure).not.toHaveProperty('uploadId');
|
||||
});
|
||||
});
|
||||
|
||||
describe('getImports filtering', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
dbRead.huggingFaceImport.findMany.mockResolvedValue([]);
|
||||
});
|
||||
|
||||
const whereOf = () => dbRead.huggingFaceImport.findMany.mock.calls[0][0].where;
|
||||
|
||||
it('filters on the server, so results are not capped-then-filtered', async () => {
|
||||
await getImports({ userId: 7, isModerator: true, groupName: 'FLUX' });
|
||||
// The bug this replaces: both callers fetched `limit` rows and filtered in the client, so a
|
||||
// group older than that window returned nothing and looked like it had never been imported.
|
||||
expect(whereOf()).toMatchObject({
|
||||
groupName: { contains: 'FLUX', mode: 'insensitive' },
|
||||
});
|
||||
});
|
||||
|
||||
it('matches a repo exactly rather than by substring', async () => {
|
||||
await getImports({ userId: 7, isModerator: true, repo: 'owner/name' });
|
||||
expect(whereOf().repo).toBe('owner/name');
|
||||
});
|
||||
|
||||
it('adds no predicate when nothing is filtered', async () => {
|
||||
await getImports({ userId: 7, isModerator: true });
|
||||
expect(whereOf()).toEqual({});
|
||||
});
|
||||
|
||||
it('still scopes a non-moderator to their own rows while filtering', async () => {
|
||||
await getImports({ userId: 7, isModerator: false, groupName: 'FLUX' });
|
||||
expect(whereOf().userId).toBe(7);
|
||||
});
|
||||
});
|
||||
|
||||
describe('import config', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
sysRedisMock.packed.get.mockResolvedValue(null);
|
||||
});
|
||||
|
||||
it('falls open to the defaults when the store cannot be read', async () => {
|
||||
sysRedisMock.packed.get.mockRejectedValue(new Error('redis down'));
|
||||
// A config store that cannot be read must not stop transfers, and must never resolve
|
||||
// concurrency to zero — both are worse than running at the shipped shape.
|
||||
await expect(getHuggingFaceImportConfig()).resolves.toEqual(HUGGING_FACE_IMPORT_DEFAULTS);
|
||||
});
|
||||
|
||||
it('merges a partial stored value over the defaults', async () => {
|
||||
sysRedisMock.packed.get.mockResolvedValue({ partsInFlight: 1 });
|
||||
const config = await getHuggingFaceImportConfig();
|
||||
expect(config.partsInFlight).toBe(1);
|
||||
expect(config.filesInParallel).toBe(HUGGING_FACE_IMPORT_DEFAULTS.filesInParallel);
|
||||
});
|
||||
|
||||
it('ignores a stored value that is out of bounds rather than obeying it', async () => {
|
||||
// The bounds are what stop a text box setting pod memory to gigabytes.
|
||||
sysRedisMock.packed.get.mockResolvedValue({ partsInFlight: 500 });
|
||||
await expect(getHuggingFaceImportConfig()).resolves.toEqual(HUGGING_FACE_IMPORT_DEFAULTS);
|
||||
});
|
||||
|
||||
it('refuses to write a value outside the bounds', async () => {
|
||||
await expect(setHuggingFaceImportConfig({ filesInParallel: 99 })).rejects.toThrow();
|
||||
expect(sysRedisMock.packed.set).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('unattached and delete', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
dbRead.huggingFaceImport.findMany.mockResolvedValue([]);
|
||||
});
|
||||
|
||||
it('defines unattached as completed with no model file', async () => {
|
||||
await getImports({ userId: 7, isModerator: true, unattached: true });
|
||||
expect(dbRead.huggingFaceImport.findMany.mock.calls[0][0].where).toMatchObject({
|
||||
status: 'Completed',
|
||||
modelFileId: null,
|
||||
});
|
||||
});
|
||||
|
||||
it('counts the unattached tab with the same predicate the tab lists', async () => {
|
||||
dbRead.huggingFaceImport.count.mockResolvedValue(0);
|
||||
await getImportCounts({ userId: 7, isModerator: true, groupName: 'krea' });
|
||||
|
||||
const [unattachedWhere, totalWhere] = dbRead.huggingFaceImport.count.mock.calls.map(
|
||||
(call: [{ where: Record<string, unknown> }]) => call[0].where
|
||||
);
|
||||
// A label that counts a population the list is not showing is the bug this query exists to fix.
|
||||
expect(unattachedWhere).toMatchObject({
|
||||
status: 'Completed',
|
||||
modelFileId: null,
|
||||
groupName: { contains: 'krea', mode: 'insensitive' },
|
||||
});
|
||||
expect(totalWhere).toMatchObject({ groupName: { contains: 'krea', mode: 'insensitive' } });
|
||||
expect(totalWhere).not.toHaveProperty('modelFileId');
|
||||
});
|
||||
|
||||
it('scopes the lookup to the owner when the caller is not a moderator', async () => {
|
||||
dbRead.huggingFaceImport.findFirst.mockResolvedValue(null);
|
||||
await expect(deleteImport({ id: 1, userId: 7, isModerator: false })).rejects.toThrow();
|
||||
expect(dbRead.huggingFaceImport.findFirst.mock.calls[0][0].where).toMatchObject({
|
||||
id: 1,
|
||||
userId: 7,
|
||||
});
|
||||
});
|
||||
|
||||
it('refuses to delete an import that is still attached', async () => {
|
||||
// Deleting here would leave a model version pointing at bytes that no longer exist.
|
||||
dbRead.huggingFaceImport.findFirst.mockResolvedValue({
|
||||
id: 1,
|
||||
status: 'Completed',
|
||||
bucket: 'b',
|
||||
key: 'k',
|
||||
url: 'https://s3.example/b/k',
|
||||
uploadId: null,
|
||||
modelFileId: 99,
|
||||
});
|
||||
await expect(deleteImport({ id: 1, userId: 7, isModerator: true })).rejects.toThrow();
|
||||
expect(mockDeleteObject).not.toHaveBeenCalled();
|
||||
expect(dbWrite.huggingFaceImport.deleteMany).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('refuses to delete a DETACHED import a model file still points at', async () => {
|
||||
// The two-click data-loss path: detach leaves the ModelFile alive, so the row lands in the
|
||||
// unattached list while a published version is still serving those exact bytes. `modelFileId`
|
||||
// is a local pointer; the refcount over `ModelFile.url` is the global one.
|
||||
dbRead.huggingFaceImport.findFirst.mockResolvedValue({
|
||||
id: 1,
|
||||
status: 'Completed',
|
||||
bucket: 'b2-transfer-bucket',
|
||||
key: 'model/7/x.safetensors',
|
||||
url: 'https://s3.example/b2-transfer-bucket/model/7/x.safetensors',
|
||||
uploadId: null,
|
||||
modelFileId: null,
|
||||
});
|
||||
mockUrlsSafeToDelete.mockResolvedValue({ safe: [], skipped: 1 });
|
||||
|
||||
await expect(deleteImport({ id: 1, userId: 7, isModerator: true })).rejects.toThrow(
|
||||
/model file still points at/i
|
||||
);
|
||||
expect(mockDeleteObject).not.toHaveBeenCalled();
|
||||
expect(dbWrite.huggingFaceImport.deleteMany).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('refuses to delete a transfer that is still running', async () => {
|
||||
dbRead.huggingFaceImport.findFirst.mockResolvedValue({
|
||||
id: 1,
|
||||
status: 'Transferring',
|
||||
bucket: 'b',
|
||||
key: 'k',
|
||||
url: null,
|
||||
uploadId: 'u',
|
||||
modelFileId: null,
|
||||
});
|
||||
await expect(deleteImport({ id: 1, userId: 7, isModerator: true })).rejects.toThrow();
|
||||
expect(mockDeleteObject).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('frees the object before removing the row', async () => {
|
||||
dbRead.huggingFaceImport.findFirst.mockResolvedValue({
|
||||
id: 1,
|
||||
status: 'Completed',
|
||||
// Deliberately NOT what the env resolves to: these bytes are in the bucket the transfer used.
|
||||
bucket: 'b2-transfer-bucket',
|
||||
key: 'model/7/x.safetensors',
|
||||
url: 'https://s3.example/b2-transfer-bucket/model/7/x.safetensors',
|
||||
uploadId: null,
|
||||
modelFileId: null,
|
||||
});
|
||||
mockUrlsSafeToDelete.mockResolvedValue({ safe: ['https://s3.example/x'], skipped: 0 });
|
||||
mockDeleteObject.mockResolvedValue(undefined);
|
||||
dbWrite.huggingFaceImport.deleteMany.mockResolvedValue({ count: 1 });
|
||||
|
||||
await deleteImport({ id: 1, userId: 7, isModerator: true });
|
||||
|
||||
expect(mockDeleteObject).toHaveBeenCalledTimes(1);
|
||||
expect(mockDeleteObject).toHaveBeenCalledWith(
|
||||
'b2-transfer-bucket',
|
||||
'model/7/x.safetensors',
|
||||
expect.anything()
|
||||
);
|
||||
// The predicate rides into the write, because the read was against the replica.
|
||||
expect(dbWrite.huggingFaceImport.deleteMany).toHaveBeenCalledWith({
|
||||
where: { id: 1, modelFileId: null },
|
||||
});
|
||||
// Named for the ordering, so it asserts the ordering rather than leaning on the sibling test.
|
||||
expect(mockDeleteObject.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
dbWrite.huggingFaceImport.deleteMany.mock.invocationCallOrder[0]
|
||||
);
|
||||
});
|
||||
|
||||
it('aborts a live multipart before forgetting the row', async () => {
|
||||
// A Failed row can still hold an uploadId, and the row is the only handle that can free the
|
||||
// parts already uploaded — which are billed until something aborts them.
|
||||
dbRead.huggingFaceImport.findFirst.mockResolvedValue({
|
||||
id: 1,
|
||||
status: 'Failed',
|
||||
bucket: 'b2-transfer-bucket',
|
||||
key: 'model/7/x.safetensors',
|
||||
url: null,
|
||||
uploadId: 'upload-1',
|
||||
modelFileId: null,
|
||||
});
|
||||
mockAbort.mockResolvedValue(undefined);
|
||||
mockDeleteObject.mockResolvedValue(undefined);
|
||||
dbWrite.huggingFaceImport.deleteMany.mockResolvedValue({ count: 1 });
|
||||
|
||||
await deleteImport({ id: 1, userId: 7, isModerator: true });
|
||||
|
||||
expect(mockAbort).toHaveBeenCalledWith(
|
||||
'b2-transfer-bucket',
|
||||
'model/7/x.safetensors',
|
||||
'upload-1',
|
||||
expect.anything()
|
||||
);
|
||||
});
|
||||
|
||||
it('keeps the row when the multipart abort fails', async () => {
|
||||
dbRead.huggingFaceImport.findFirst.mockResolvedValue({
|
||||
id: 1,
|
||||
status: 'Failed',
|
||||
bucket: 'b2-transfer-bucket',
|
||||
key: 'model/7/x.safetensors',
|
||||
url: null,
|
||||
uploadId: 'upload-1',
|
||||
modelFileId: null,
|
||||
});
|
||||
mockAbort.mockRejectedValue(new Error('B2 unavailable'));
|
||||
|
||||
await expect(deleteImport({ id: 1, userId: 7, isModerator: true })).rejects.toThrow();
|
||||
expect(dbWrite.huggingFaceImport.deleteMany).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('keeps the row when the object could not be deleted', async () => {
|
||||
// Otherwise the bytes stay in the bucket with nothing left pointing at them.
|
||||
dbRead.huggingFaceImport.findFirst.mockResolvedValue({
|
||||
id: 1,
|
||||
status: 'Completed',
|
||||
bucket: 'b2-transfer-bucket',
|
||||
key: 'model/7/x.safetensors',
|
||||
url: 'https://s3.example/b2-transfer-bucket/model/7/x.safetensors',
|
||||
uploadId: null,
|
||||
modelFileId: null,
|
||||
});
|
||||
mockUrlsSafeToDelete.mockResolvedValue({ safe: ['https://s3.example/x'], skipped: 0 });
|
||||
mockDeleteObject.mockRejectedValue(new Error('B2 unavailable'));
|
||||
|
||||
await expect(deleteImport({ id: 1, userId: 7, isModerator: true })).rejects.toThrow();
|
||||
expect(dbWrite.huggingFaceImport.deleteMany).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('renameGroup', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
const input = {
|
||||
repo: 'owner/name',
|
||||
revision: 'abc123',
|
||||
from: 'flux-krea',
|
||||
groupName: 'FLUX Krea',
|
||||
userId: 7,
|
||||
isModerator: true,
|
||||
};
|
||||
|
||||
it('renames only the rows of the named group, whatever their status', async () => {
|
||||
dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 3 });
|
||||
|
||||
await expect(renameGroup(input)).resolves.toEqual({ renamed: 3, groupName: 'FLUX Krea' });
|
||||
|
||||
const { where, data } = dbWrite.huggingFaceImport.updateMany.mock.calls[0][0];
|
||||
// The current name is part of the scope: one repo at one revision can be two batches.
|
||||
expect(where).toEqual({ repo: 'owner/name', revision: 'abc123', groupName: 'flux-krea' });
|
||||
expect(data).toEqual({ groupName: 'FLUX Krea' });
|
||||
});
|
||||
|
||||
it('scopes to the owner when the caller is not a moderator', async () => {
|
||||
dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 });
|
||||
|
||||
await renameGroup({ ...input, isModerator: false });
|
||||
|
||||
expect(dbWrite.huggingFaceImport.updateMany.mock.calls[0][0].where).toMatchObject({
|
||||
userId: 7,
|
||||
});
|
||||
});
|
||||
|
||||
it('trims the new name', async () => {
|
||||
dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 1 });
|
||||
|
||||
await renameGroup({ ...input, groupName: ' FLUX Krea ' });
|
||||
|
||||
expect(dbWrite.huggingFaceImport.updateMany.mock.calls[0][0].data).toEqual({
|
||||
groupName: 'FLUX Krea',
|
||||
});
|
||||
});
|
||||
|
||||
it('refuses an empty name without writing', async () => {
|
||||
await expect(renameGroup({ ...input, groupName: ' ' })).rejects.toThrow();
|
||||
expect(dbWrite.huggingFaceImport.updateMany).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('writes nothing when the name is unchanged', async () => {
|
||||
await expect(renameGroup({ ...input, groupName: 'flux-krea' })).resolves.toEqual({
|
||||
renamed: 0,
|
||||
groupName: 'flux-krea',
|
||||
});
|
||||
expect(dbWrite.huggingFaceImport.updateMany).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('reports a group that no longer exists', async () => {
|
||||
dbWrite.huggingFaceImport.updateMany.mockResolvedValue({ count: 0 });
|
||||
|
||||
await expect(renameGroup(input)).rejects.toThrow(/No files found in group/);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,70 @@
|
||||
import * as z from 'zod';
|
||||
import { REDIS_SYS_KEYS, sysRedis, withSysReadDeadline } from '~/server/redis/client';
|
||||
import { logSysRedisFailOpen } from '~/server/redis/fail-open-log';
|
||||
|
||||
/**
|
||||
* Runtime knobs for the Hugging Face transfer, so the shape of a running import can be changed
|
||||
* without a deploy. Before this existed the only lever on a transfer hurting production was shipping
|
||||
* a new constant.
|
||||
*
|
||||
* 🔴 The bounds are the point, not the defaults. `partsInFlight × filesInParallel × PART_SIZE_BYTES`
|
||||
* is the RETAINED payload — 3 × 2 × 16MB = 96MB at the defaults — but measured RSS growth is ~200MB,
|
||||
* because `res.arrayBuffer()` leaves undici's concat buffer alive and that garbage is off-heap, where
|
||||
* it barely pressures V8's major-GC trigger. It oscillates rather than leaks, but a container limit is
|
||||
* a hard limit, so size headroom against the larger number. An unbounded value here would be an OOM on
|
||||
* a web pod set from a text box; the bounds live here rather than in the UI because the UI is not the
|
||||
* only caller.
|
||||
*/
|
||||
export const huggingFaceImportConfigSchema = z.object({
|
||||
/** The kill switch. Off leaves queued rows untouched; it stops claiming new work. */
|
||||
enabled: z.boolean(),
|
||||
/** Files transferred at once across the whole fleet — one job run holds the lock. */
|
||||
filesInParallel: z.number().int().min(1).max(4),
|
||||
/** Parts of one file in flight at once. */
|
||||
partsInFlight: z.number().int().min(1).max(6),
|
||||
/** How long one run transfers before yielding. Must stay well inside the job's 5-minute lock. */
|
||||
workBudgetSeconds: z.number().int().min(15).max(240),
|
||||
});
|
||||
|
||||
export type HuggingFaceImportConfig = z.infer<typeof huggingFaceImportConfigSchema>;
|
||||
|
||||
/**
|
||||
* Deliberately modest. 3 × 2 × 16MB is ~96MB retained and ~200MB resident, which a web pod can carry
|
||||
* while serving traffic; the budget leaves three minutes of the lock for the slowest part in flight.
|
||||
*/
|
||||
export const HUGGING_FACE_IMPORT_DEFAULTS: HuggingFaceImportConfig = {
|
||||
enabled: true,
|
||||
filesInParallel: 2,
|
||||
partsInFlight: 3,
|
||||
workBudgetSeconds: 120,
|
||||
};
|
||||
|
||||
/**
|
||||
* 🔴 Fails OPEN to the defaults. A config store that cannot be read must not stop transfers, and must
|
||||
* not silently resolve concurrency to zero — both are worse than running at the shipped shape.
|
||||
*/
|
||||
export async function getHuggingFaceImportConfig(): Promise<HuggingFaceImportConfig> {
|
||||
let stored: Partial<HuggingFaceImportConfig> = {};
|
||||
try {
|
||||
const raw = await withSysReadDeadline(
|
||||
sysRedis.packed.get<Partial<HuggingFaceImportConfig>>(
|
||||
REDIS_SYS_KEYS.HUGGING_FACE_IMPORT.CONFIG
|
||||
)
|
||||
);
|
||||
if (raw) stored = huggingFaceImportConfigSchema.partial().parse(raw);
|
||||
} catch (error) {
|
||||
logSysRedisFailOpen('read-degraded', 'getHuggingFaceImportConfig', error);
|
||||
}
|
||||
|
||||
return { ...HUGGING_FACE_IMPORT_DEFAULTS, ...stored };
|
||||
}
|
||||
|
||||
/** Merges over what is stored, so a caller may set one knob without restating the rest. */
|
||||
export async function setHuggingFaceImportConfig(input: Partial<HuggingFaceImportConfig>) {
|
||||
const next = huggingFaceImportConfigSchema.parse({
|
||||
...(await getHuggingFaceImportConfig()),
|
||||
...input,
|
||||
});
|
||||
await sysRedis.packed.set(REDIS_SYS_KEYS.HUGGING_FACE_IMPORT.CONFIG, next);
|
||||
return next;
|
||||
}
|
||||
@@ -0,0 +1,872 @@
|
||||
import { Prisma } from '@prisma/client';
|
||||
import { env } from '~/env/server';
|
||||
import { dbRead, dbWrite } from '~/server/db/client';
|
||||
import { logToAxiom } from '~/server/logging/client';
|
||||
import {
|
||||
getRepoFiles,
|
||||
headHuggingFaceFile,
|
||||
huggingFaceResolveUrl,
|
||||
defaultGroupName,
|
||||
readHuggingFaceRange,
|
||||
suggestFileType,
|
||||
type HuggingFaceRepo,
|
||||
} from '~/server/services/huggingface.service';
|
||||
import { UploadType } from '~/server/common/enums';
|
||||
import {
|
||||
abortMultipartUpload,
|
||||
completeMultipartUpload,
|
||||
createMultipartUpload,
|
||||
getGetUrlByKey,
|
||||
deleteObject,
|
||||
getBucket,
|
||||
getS3Client,
|
||||
urlsSafeToDelete,
|
||||
getUploadBucket,
|
||||
objectExists,
|
||||
getUploadS3Client,
|
||||
uploadPart,
|
||||
} from '~/utils/s3-utils';
|
||||
import { buildUploadKey } from '~/utils/upload-key';
|
||||
import { throwBadRequestError, throwNotFoundError } from '~/server/utils/errorHandling';
|
||||
import { bytesToKB } from '~/utils/number-helpers';
|
||||
import { getModelFileFormat } from '~/utils/file-helpers';
|
||||
import type { ModelFileCreateInput } from '~/server/schema/model-file.schema';
|
||||
import type { ModelFileType } from '~/server/common/constants';
|
||||
|
||||
/** A claim older than this belonged to a run that died mid-part. Sized well above one part's
|
||||
* transfer time so a slow part is never mistaken for an abandoned one. */
|
||||
const STALE_CLAIM_MINUTES = 20;
|
||||
/**
|
||||
* 🔴 A LITERAL, never a bind. `make_interval` takes int4, and a JS number bound through `$queryRaw`
|
||||
* arrives as int8 — Postgres then finds no matching function and fails the whole statement with
|
||||
* 42883. No unit test can see it: the SQL is mocked wholesale, so this only ever surfaces against a
|
||||
* real database. (`minor-hash.service.ts` carries the same note for the same reason.)
|
||||
*/
|
||||
const STALE_CLAIM_INTERVAL = Prisma.raw(`make_interval(mins => ${STALE_CLAIM_MINUTES})`);
|
||||
const MAX_ATTEMPTS = 5;
|
||||
const RETRY_BACKOFF_MINUTES = 5;
|
||||
/**
|
||||
* Fixed, and deliberately NOT `getUploadChunkSize`. That helper doubles the chunk to stay under
|
||||
* `MAX_UPLOAD_PARTS = 1000`, a bound that exists because the browser path presigns every part up
|
||||
* front — a cheap constraint there, and the wrong one to inherit here, where the part size IS the
|
||||
* pod's memory footprint. Under it a single 50GB file silently moved to 100MB parts and roughly
|
||||
* doubled resident memory. 16MB × 10,000 parts (what B2's S3 API allows) covers 160GB.
|
||||
*/
|
||||
export const PART_SIZE_BYTES = 16 * 1024 * 1024;
|
||||
const partSizeFor = (size: number) => Math.max(PART_SIZE_BYTES, Math.ceil(size / 10_000));
|
||||
|
||||
type MultipartPart = { PartNumber: number; ETag: string };
|
||||
|
||||
export type HuggingFaceImportView = {
|
||||
id: number;
|
||||
repo: string;
|
||||
revision: string;
|
||||
filename: string;
|
||||
sourceUrl: string;
|
||||
sizeBytes: number | null;
|
||||
groupName: string;
|
||||
sourceSha256: string | null;
|
||||
status: 'Queued' | 'Transferring' | 'Completed' | 'Failed' | 'Canceled';
|
||||
bytesTransferred: number;
|
||||
url: string | null;
|
||||
error: string | null;
|
||||
modelFileId: number | null;
|
||||
modelVersionId: number | null;
|
||||
suggestedType: ReturnType<typeof suggestFileType>;
|
||||
userId: number | null;
|
||||
createdAt: Date;
|
||||
startedAt: Date | null;
|
||||
completedAt: Date | null;
|
||||
};
|
||||
|
||||
const importSelect = {
|
||||
id: true,
|
||||
repo: true,
|
||||
revision: true,
|
||||
filename: true,
|
||||
sourceUrl: true,
|
||||
groupName: true,
|
||||
sizeBytes: true,
|
||||
sourceSha256: true,
|
||||
status: true,
|
||||
bytesTransferred: true,
|
||||
url: true,
|
||||
error: true,
|
||||
modelFileId: true,
|
||||
modelVersionId: true,
|
||||
userId: true,
|
||||
createdAt: true,
|
||||
startedAt: true,
|
||||
completedAt: true,
|
||||
} satisfies Prisma.HuggingFaceImportSelect;
|
||||
|
||||
/** BigInt columns arrive as `bigint`, which does not survive superjson to the client. */
|
||||
function toView(row: Prisma.HuggingFaceImportGetPayload<{ select: typeof importSelect }>) {
|
||||
return {
|
||||
...row,
|
||||
sizeBytes: row.sizeBytes === null ? null : Number(row.sizeBytes),
|
||||
bytesTransferred: Number(row.bytesTransferred),
|
||||
suggestedType: suggestFileType(row.filename),
|
||||
} as HuggingFaceImportView;
|
||||
}
|
||||
|
||||
/**
|
||||
* Transferred bytes no version has claimed. The only two ways off this list are attaching and
|
||||
* deleting — there is deliberately no dismissed-but-stored state, because a hidden row still costs
|
||||
* storage and would let the count understate what we hold.
|
||||
*
|
||||
* Shared so a tab's label and the rows under it cannot disagree about what unattached means.
|
||||
*/
|
||||
const UNATTACHED_WHERE = { status: 'Completed' as const, modelFileId: null };
|
||||
|
||||
/** The filters every list and count in this feature takes, including who may see a row. */
|
||||
function scopeWhere({
|
||||
userId,
|
||||
isModerator,
|
||||
groupName,
|
||||
repo,
|
||||
}: {
|
||||
userId: number;
|
||||
isModerator: boolean;
|
||||
groupName?: string;
|
||||
repo?: string;
|
||||
}) {
|
||||
return {
|
||||
...(isModerator ? {} : { userId }),
|
||||
...(repo ? { repo } : {}),
|
||||
...(groupName ? { groupName: { contains: groupName, mode: 'insensitive' as const } } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* 🔴 Filtering happens HERE, not in the caller. Both surfaces take `limit` rows and used to filter
|
||||
* them client-side, so once the table passed that limit an older group returned nothing — and "no
|
||||
* results" is indistinguishable from "never imported". The skill's docs even attributed that empty
|
||||
* result to a casing mismatch and recommended a re-run that would also have found nothing.
|
||||
*
|
||||
* `repo` is an equality match and is what `HuggingFaceImport_groupName_idx`'s sibling indexes serve.
|
||||
* `groupName` is a substring search, which a plain btree cannot serve — at this table's size that is
|
||||
* a few milliseconds of seq scan, and if it ever stops being one the answer is a trigram index, not
|
||||
* a narrower filter.
|
||||
*/
|
||||
export async function getImports({
|
||||
userId,
|
||||
isModerator,
|
||||
limit = 100,
|
||||
groupName,
|
||||
repo,
|
||||
unattached,
|
||||
}: {
|
||||
userId: number;
|
||||
isModerator: boolean;
|
||||
limit?: number;
|
||||
groupName?: string;
|
||||
repo?: string;
|
||||
unattached?: boolean;
|
||||
}) {
|
||||
const rows = await dbRead.huggingFaceImport.findMany({
|
||||
where: {
|
||||
...scopeWhere({ userId, isModerator, groupName, repo }),
|
||||
...(unattached ? UNATTACHED_WHERE : {}),
|
||||
},
|
||||
select: importSelect,
|
||||
orderBy: { createdAt: 'desc' },
|
||||
take: limit,
|
||||
});
|
||||
return rows.map(toView);
|
||||
}
|
||||
|
||||
/**
|
||||
* Counts for the queue's tabs. Separate from `getImports` because a page of rows cannot report how
|
||||
* many exist outside it — which is the mistake the client-side filter made. Takes the same filters
|
||||
* the rows do, or a label counts a population the list beneath it is not showing.
|
||||
*/
|
||||
export async function getImportCounts(input: {
|
||||
userId: number;
|
||||
isModerator: boolean;
|
||||
groupName?: string;
|
||||
repo?: string;
|
||||
}) {
|
||||
const scope = scopeWhere(input);
|
||||
const [unattached, total] = await Promise.all([
|
||||
dbRead.huggingFaceImport.count({ where: { ...scope, ...UNATTACHED_WHERE } }),
|
||||
dbRead.huggingFaceImport.count({ where: scope }),
|
||||
]);
|
||||
return { unattached, total };
|
||||
}
|
||||
|
||||
/**
|
||||
* Files we already store under a byte-identical sha256. HF reports each LFS file's content hash
|
||||
* before any bytes move, so a text encoder shared by a dozen repos is transferred once and the next
|
||||
* import points at what we have.
|
||||
*/
|
||||
async function findExistingByHash(sha256List: string[]) {
|
||||
const hashes = [...new Set(sha256List.filter(Boolean).map((h) => h.toUpperCase()))];
|
||||
if (!hashes.length) return new Map<string, { fileId: number; name: string; url: string }>();
|
||||
|
||||
const rows = await dbRead.modelFileHash.findMany({
|
||||
where: { type: 'SHA256', hash: { in: hashes } },
|
||||
select: { hash: true, file: { select: { id: true, name: true, url: true } } },
|
||||
});
|
||||
|
||||
return new Map(
|
||||
rows.map((row) => [
|
||||
row.hash.toUpperCase(),
|
||||
{ fileId: row.file.id, name: row.file.name, url: row.file.url },
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
export async function resolveRepoForImport(input: { repo: string; revision?: string }) {
|
||||
const repo = await getRepoFiles(input);
|
||||
const existing = await findExistingByHash(
|
||||
repo.files.map((f) => f.sha256).filter((s): s is string => !!s)
|
||||
);
|
||||
return {
|
||||
...repo,
|
||||
files: repo.files.map((file) => ({
|
||||
...file,
|
||||
existing: file.sha256 ? existing.get(file.sha256.toUpperCase()) ?? null : null,
|
||||
suggestedType: suggestFileType(file.path),
|
||||
})),
|
||||
};
|
||||
}
|
||||
|
||||
/** Model files go to B2 when it is configured, matching `/api/upload`. */
|
||||
async function uploadTarget() {
|
||||
const useB2 = !!env.S3_UPLOAD_B2_ENDPOINT;
|
||||
return {
|
||||
s3: useB2 ? getUploadS3Client('b2') : getS3Client(),
|
||||
bucket: useB2 ? getUploadBucket('b2') : await getBucket(),
|
||||
};
|
||||
}
|
||||
|
||||
export async function enqueueImports({
|
||||
repo,
|
||||
paths,
|
||||
userId,
|
||||
groupName,
|
||||
}: {
|
||||
repo: HuggingFaceRepo;
|
||||
paths: string[];
|
||||
userId: number;
|
||||
groupName?: string;
|
||||
}) {
|
||||
const wanted = repo.files.filter((file) => paths.includes(file.path));
|
||||
if (!wanted.length) return { queued: 0, skipped: 0 };
|
||||
|
||||
// `(repo, revision, filename)` is unique, so re-queueing a repo adds only what is new.
|
||||
const result = await dbWrite.huggingFaceImport.createMany({
|
||||
data: wanted.map((file) => ({
|
||||
repo: repo.repo,
|
||||
revision: repo.revision,
|
||||
filename: file.path,
|
||||
sourceUrl: huggingFaceResolveUrl(repo.repo, repo.revision, file.path),
|
||||
sizeBytes: file.size ? BigInt(file.size) : null,
|
||||
sourceSha256: file.sha256,
|
||||
groupName: groupName?.trim() || defaultGroupName(repo.repo),
|
||||
userId,
|
||||
})),
|
||||
skipDuplicates: true,
|
||||
});
|
||||
|
||||
return { queued: result.count, skipped: wanted.length - result.count };
|
||||
}
|
||||
|
||||
/** Scoped by owner as well as id — the page is moderator-only today, the service is not. */
|
||||
async function ownedImport({
|
||||
id,
|
||||
userId,
|
||||
isModerator,
|
||||
}: {
|
||||
id: number;
|
||||
userId: number;
|
||||
isModerator: boolean;
|
||||
}) {
|
||||
const row = await dbRead.huggingFaceImport.findFirst({
|
||||
where: { id, userId: isModerator ? undefined : userId },
|
||||
select: {
|
||||
id: true,
|
||||
status: true,
|
||||
bucket: true,
|
||||
key: true,
|
||||
url: true,
|
||||
uploadId: true,
|
||||
modelFileId: true,
|
||||
},
|
||||
});
|
||||
if (!row) throw throwNotFoundError('Import not found');
|
||||
return row;
|
||||
}
|
||||
|
||||
/**
|
||||
* Turns a finished import into the input `createFileHandler` takes. Going through that handler rather
|
||||
* than writing a `ModelFile` directly is what gets the storage-resolver registration and the inline
|
||||
* scan submission — an attached file has to be downloadable and scannable, not merely present.
|
||||
*/
|
||||
export async function buildAttachInput({
|
||||
id,
|
||||
modelVersionId,
|
||||
type,
|
||||
userId,
|
||||
isModerator,
|
||||
}: {
|
||||
id: number;
|
||||
modelVersionId: number;
|
||||
type: ModelFileType;
|
||||
userId: number;
|
||||
isModerator: boolean;
|
||||
}): Promise<ModelFileCreateInput & { importId: number }> {
|
||||
// Ownership is decided by `ownedImport` and nowhere else. This used to re-inline the same
|
||||
// predicate for the sake of a wider `select`, which left two copies of the rule — and this is the
|
||||
// copy that mints a `ModelFile` on a caller-supplied version, so it is the worst one to let drift.
|
||||
await ownedImport({ id, userId, isModerator });
|
||||
const row = await dbRead.huggingFaceImport.findUniqueOrThrow({
|
||||
where: { id },
|
||||
select: {
|
||||
id: true,
|
||||
filename: true,
|
||||
url: true,
|
||||
key: true,
|
||||
bucket: true,
|
||||
sizeBytes: true,
|
||||
status: true,
|
||||
modelFileId: true,
|
||||
},
|
||||
});
|
||||
if (row.status !== 'Completed' || !row.url || !row.key)
|
||||
throw throwBadRequestError('That import has not finished transferring.');
|
||||
if (row.modelFileId)
|
||||
throw throwBadRequestError(
|
||||
`Already attached as file ${row.modelFileId}. Detach or delete that file first.`
|
||||
);
|
||||
if (!row.sizeBytes) throw throwBadRequestError('That import has no recorded size.');
|
||||
|
||||
// 🔴 The row says Completed; the bytes may not be there. `deleteFile` refcounts live `ModelFile`
|
||||
// rows and knows nothing about imports, so deleting the one attached file — which this function's
|
||||
// own refusal message suggests — removes the object and leaves this row looking attachable.
|
||||
// Re-attaching would publish a model file pointing at nothing.
|
||||
//
|
||||
// `objectExists` is deliberately tri-state: `null` means the bucket could not be consulted, and a
|
||||
// guard that cannot ask must not block a legitimate attach.
|
||||
const { s3, bucket } = await uploadTarget();
|
||||
const present = await objectExists(row.bucket ?? bucket, row.key, s3);
|
||||
if (present === false)
|
||||
throw throwBadRequestError(
|
||||
'The stored object for that import is gone. Re-import it before attaching.'
|
||||
);
|
||||
|
||||
const name = row.filename.split('/').pop() ?? row.filename;
|
||||
return {
|
||||
importId: row.id,
|
||||
modelVersionId,
|
||||
type,
|
||||
name,
|
||||
url: row.url,
|
||||
sizeKB: bytesToKB(Number(row.sizeBytes)),
|
||||
// From the row, for the same reason the resume path reads `row.bucket`: these bytes are in the
|
||||
// bucket the transfer used, which is not necessarily what the env resolves to now.
|
||||
backend: row.bucket && row.bucket === getUploadBucket('b2') ? 'b2' : undefined,
|
||||
s3Path: row.key,
|
||||
metadata: { format: getModelFileFormat(name) },
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Claims the import for exactly one model file. The `modelFileId: null` predicate is the real
|
||||
* double-attach guard — `buildAttachInput`'s check reads the REPLICA, so two attaches issued close
|
||||
* together both see null there and both create a file. Zero rows updated means someone else won.
|
||||
*/
|
||||
export async function linkImportToFile({
|
||||
id,
|
||||
modelFileId,
|
||||
modelVersionId,
|
||||
}: {
|
||||
id: number;
|
||||
modelFileId: number;
|
||||
modelVersionId: number;
|
||||
}) {
|
||||
const { count } = await dbWrite.huggingFaceImport.updateMany({
|
||||
where: { id, modelFileId: null },
|
||||
data: { modelFileId, modelVersionId },
|
||||
});
|
||||
return count === 1;
|
||||
}
|
||||
|
||||
/**
|
||||
* Releases the import from its model file. Without this the refusal in `buildAttachInput` names a
|
||||
* recovery path that does not exist: attach to the wrong version, delete the file, and the row can
|
||||
* never be attached again — nor re-queued, since `(repo, revision, filename)` is unique.
|
||||
*/
|
||||
export async function detachImport(input: { id: number; userId: number; isModerator: boolean }) {
|
||||
const row = await ownedImport(input);
|
||||
const { count } = await dbWrite.huggingFaceImport.updateMany({
|
||||
where: { id: row.id, modelFileId: { not: null } },
|
||||
data: { modelFileId: null, modelVersionId: null },
|
||||
});
|
||||
return { ok: count === 1 };
|
||||
}
|
||||
|
||||
/**
|
||||
* Renames a batch at any status — the name never reaches a storage key, so nothing desynchronises.
|
||||
*
|
||||
* Scoped by the CURRENT name as well as repo and revision: one repo at one revision can be imported
|
||||
* as two batches, and renaming by repo alone would silently merge them.
|
||||
*/
|
||||
export async function renameGroup({
|
||||
repo,
|
||||
revision,
|
||||
from,
|
||||
groupName,
|
||||
userId,
|
||||
isModerator,
|
||||
}: {
|
||||
repo: string;
|
||||
revision: string;
|
||||
from: string;
|
||||
groupName: string;
|
||||
userId: number;
|
||||
isModerator: boolean;
|
||||
}) {
|
||||
const name = groupName.trim();
|
||||
if (!name) throw throwBadRequestError('A group name cannot be empty.');
|
||||
if (name === from) return { renamed: 0, groupName: name };
|
||||
|
||||
const { count } = await dbWrite.huggingFaceImport.updateMany({
|
||||
where: { repo, revision, groupName: from, ...(isModerator ? {} : { userId }) },
|
||||
data: { groupName: name },
|
||||
});
|
||||
if (!count) throw throwNotFoundError(`No files found in group "${from}".`);
|
||||
return { renamed: count, groupName: name };
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes an unattached import: the stored object first, then the row.
|
||||
*
|
||||
* 🔴 The row is hard-deleted rather than tombstoned. `(repo, revision, filename)` is unique, so a
|
||||
* tombstone would block ever re-importing that exact file — and a deliberate deletion is precisely
|
||||
* the case where you might want it back later. The cost is losing the record that we once held those
|
||||
* bytes, which is the cheaper of the two.
|
||||
*
|
||||
* Refuses while any `ModelFile` points at the object, which would leave a model version serving
|
||||
* nothing.
|
||||
*/
|
||||
export async function deleteImport(input: { id: number; userId: number; isModerator: boolean }) {
|
||||
const row = await ownedImport(input);
|
||||
if (row.modelFileId)
|
||||
throw throwBadRequestError(
|
||||
`That import is attached as file ${row.modelFileId}. Detach it first.`
|
||||
);
|
||||
if (row.status === 'Queued' || row.status === 'Transferring')
|
||||
throw throwBadRequestError('Cancel the transfer before deleting it.');
|
||||
|
||||
// 🔴 `modelFileId` is NOT a reference count — detach clears it and leaves the `ModelFile` alive,
|
||||
// so the check above passes on a row a published version is still serving. `ModelFile.url` is the
|
||||
// authoritative reference. Called directly rather than through `deleteModelFileObject`, which
|
||||
// skips silently when unsafe and would leave the row deleted while the bytes stayed.
|
||||
if (row.url) {
|
||||
const { safe } = await urlsSafeToDelete([row.url]);
|
||||
if (!safe.length)
|
||||
throw throwBadRequestError(
|
||||
'A model file still points at that object. Delete the model file first.'
|
||||
);
|
||||
}
|
||||
|
||||
// An in-flight multipart and a finished object are freed differently, and a row can carry either.
|
||||
// The row is the only handle that can free already-uploaded parts, so a failed abort must keep it.
|
||||
if (!(await abortIfInFlight(row)))
|
||||
throw throwBadRequestError(
|
||||
'Could not abort the in-flight upload, so its parts cannot be freed yet. Nothing was removed.'
|
||||
);
|
||||
|
||||
if (row.key) {
|
||||
// Bucket and client from one resolution: a row with no bucket falling back to `getBucket()`
|
||||
// would address the main bucket with the B2 client, and a delete against the wrong endpoint
|
||||
// returns 204 with the bytes still there.
|
||||
const { s3, bucket: fallbackBucket } = await uploadTarget();
|
||||
const bucket = row.bucket ?? fallbackBucket;
|
||||
await deleteObject(bucket, row.key, s3).catch((error) => {
|
||||
logToAxiom({
|
||||
type: 'error',
|
||||
name: 'huggingface-import',
|
||||
message: 'object delete failed; the row is kept so the bytes stay reachable',
|
||||
key: row.key,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
throw throwBadRequestError('Could not delete the stored object. Nothing was removed.');
|
||||
});
|
||||
}
|
||||
|
||||
// The predicate rides into the write: `ownedImport` read the REPLICA, so an attach that committed
|
||||
// during the lag would otherwise have its brand-new model file left pointing at deleted bytes.
|
||||
const { count } = await dbWrite.huggingFaceImport.deleteMany({
|
||||
where: { id: row.id, modelFileId: null },
|
||||
});
|
||||
if (!count) throw throwBadRequestError('That import was attached while you were deleting it.');
|
||||
return { ok: true as const };
|
||||
}
|
||||
|
||||
export async function retryImport(input: { id: number; userId: number; isModerator: boolean }) {
|
||||
const row = await ownedImport(input);
|
||||
if (row.status !== 'Failed' && row.status !== 'Canceled') return { ok: false as const };
|
||||
|
||||
// Restart from zero when the abort succeeded — a failure we could not finish is the case where the
|
||||
// bytes we did write are least worth trusting. A failed abort keeps the upload, so that row resumes.
|
||||
const aborted = await abortIfInFlight(row);
|
||||
await dbWrite.huggingFaceImport.update({
|
||||
where: { id: row.id },
|
||||
data: {
|
||||
status: 'Queued',
|
||||
error: null,
|
||||
attempts: 0,
|
||||
nextAttemptAt: null,
|
||||
bytesTransferred: BigInt(0),
|
||||
...(aborted ? { uploadId: null, parts: Prisma.DbNull } : {}),
|
||||
claimedBy: null,
|
||||
claimedAt: null,
|
||||
heartbeatAt: null,
|
||||
startedAt: null,
|
||||
},
|
||||
});
|
||||
return { ok: true as const };
|
||||
}
|
||||
|
||||
export async function cancelImport(input: { id: number; userId: number; isModerator: boolean }) {
|
||||
const row = await ownedImport(input);
|
||||
if (row.status !== 'Queued' && row.status !== 'Transferring') return { ok: false as const };
|
||||
|
||||
// 🔴 The status predicate is in the WHERE, not only in the check above: `ownedImport` reads the
|
||||
// REPLICA, so a transfer that completed during the lag would otherwise have `Completed` overwritten
|
||||
// with `Canceled` — stranding a finished multi-GB object on a row that can no longer be attached,
|
||||
// re-queued (the unique key) or retried (aborting a completed upload throws).
|
||||
const { count } = await dbWrite.huggingFaceImport.updateMany({
|
||||
where: { id: row.id, status: { in: ['Queued', 'Transferring'] } },
|
||||
data: { status: 'Canceled' },
|
||||
});
|
||||
if (!count) return { ok: false as const };
|
||||
|
||||
// 🔴 Re-read from the PRIMARY before aborting. `row` is the replica snapshot taken before the
|
||||
// cancel, and a run that created its multipart upload in between would not appear in it — aborting
|
||||
// with that stale row silently no-ops and leaves every transferred part in the bucket, billed,
|
||||
// with nothing holding the id needed to free them.
|
||||
const current = await dbWrite.huggingFaceImport.findUnique({
|
||||
where: { id: row.id },
|
||||
select: { bucket: true, key: true, uploadId: true },
|
||||
});
|
||||
if (current) await abortIfInFlight(current);
|
||||
return { ok: true as const };
|
||||
}
|
||||
|
||||
/** Returns whether the upload is known to be gone — callers may only forget an `uploadId` on true. */
|
||||
async function abortIfInFlight(row: {
|
||||
bucket: string | null;
|
||||
key: string | null;
|
||||
uploadId: string | null;
|
||||
}): Promise<boolean> {
|
||||
if (!row.bucket || !row.key || !row.uploadId) return true;
|
||||
const { s3 } = await uploadTarget();
|
||||
try {
|
||||
await abortMultipartUpload(row.bucket, row.key, row.uploadId, s3);
|
||||
return true;
|
||||
} catch (error) {
|
||||
logToAxiom({
|
||||
type: 'error',
|
||||
name: 'huggingface-import',
|
||||
message: 'multipart abort failed; upload id retained so it can be aborted later',
|
||||
key: row.key,
|
||||
uploadId: row.uploadId,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
});
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
type ClaimedRow = {
|
||||
id: number;
|
||||
repo: string;
|
||||
filename: string;
|
||||
sourceUrl: string;
|
||||
sizeBytes: bigint | null;
|
||||
status: string;
|
||||
uploadId: string | null;
|
||||
partSize: number | null;
|
||||
parts: unknown;
|
||||
bucket: string | null;
|
||||
key: string | null;
|
||||
attempts: number;
|
||||
userId: number | null;
|
||||
claimedBy: string | null;
|
||||
};
|
||||
|
||||
/**
|
||||
* Claims one row for this run. `FOR UPDATE SKIP LOCKED` keeps two runs off the same file, and a NULL
|
||||
* `claimedAt` is what a cleanly-yielded row leaves behind — so a transfer that ran out of budget is
|
||||
* eligible again on the very next tick, while one whose run died waits out the stale window.
|
||||
*/
|
||||
async function claimNext(worker: string) {
|
||||
const rows = await dbWrite.$queryRaw<ClaimedRow[]>`
|
||||
UPDATE "HuggingFaceImport" SET
|
||||
status = 'Transferring',
|
||||
"claimedBy" = ${worker},
|
||||
"claimedAt" = now(),
|
||||
"heartbeatAt" = now(),
|
||||
"startedAt" = coalesce("startedAt", now()),
|
||||
"updatedAt" = now()
|
||||
WHERE id = (
|
||||
SELECT id FROM "HuggingFaceImport"
|
||||
WHERE status IN ('Queued', 'Transferring')
|
||||
AND ("claimedAt" IS NULL OR "heartbeatAt" < now() - ${STALE_CLAIM_INTERVAL})
|
||||
AND ("nextAttemptAt" IS NULL OR "nextAttemptAt" <= now())
|
||||
ORDER BY "createdAt"
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT 1
|
||||
)
|
||||
RETURNING id, repo, filename, "sourceUrl", "sizeBytes", status, "uploadId", "partSize",
|
||||
parts, bucket, key, attempts, "userId", "claimedBy"
|
||||
`;
|
||||
return rows[0];
|
||||
}
|
||||
|
||||
async function yieldClaim(id: number) {
|
||||
await dbWrite.huggingFaceImport.update({
|
||||
where: { id },
|
||||
data: { claimedBy: null, claimedAt: null },
|
||||
});
|
||||
}
|
||||
|
||||
async function failOrRetry(row: ClaimedRow, error: unknown) {
|
||||
const message = error instanceof Error ? error.message : String(error);
|
||||
const attempts = row.attempts + 1;
|
||||
const giveUp = attempts >= MAX_ATTEMPTS;
|
||||
|
||||
const aborted = giveUp ? await abortIfInFlight(row) : false;
|
||||
await dbWrite.huggingFaceImport.updateMany({
|
||||
where: { id: row.id, claimedBy: row.claimedBy, status: { not: 'Canceled' } },
|
||||
data: {
|
||||
status: giveUp ? 'Failed' : 'Transferring',
|
||||
attempts,
|
||||
error: message.slice(0, 1000),
|
||||
claimedBy: null,
|
||||
claimedAt: null,
|
||||
nextAttemptAt: giveUp
|
||||
? null
|
||||
: new Date(Date.now() + RETRY_BACKOFF_MINUTES * attempts * 60_000),
|
||||
...(aborted ? { uploadId: null, parts: Prisma.DbNull } : {}),
|
||||
},
|
||||
});
|
||||
|
||||
logToAxiom({
|
||||
type: 'error',
|
||||
name: 'huggingface-import',
|
||||
message,
|
||||
importId: row.id,
|
||||
repo: row.repo,
|
||||
filename: row.filename,
|
||||
attempts,
|
||||
giveUp,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Moves one file as far as the deadline allows. Returns `true` when the file finished, `false` when
|
||||
* it yielded with work left — the caller reclaims either way.
|
||||
*/
|
||||
async function advanceImport(
|
||||
row: ClaimedRow,
|
||||
deadline: number,
|
||||
partsInFlight: number
|
||||
): Promise<number> {
|
||||
const size = row.sizeBytes
|
||||
? Number(row.sizeBytes)
|
||||
: (await headHuggingFaceFile(row.sourceUrl)) ?? 0;
|
||||
if (!size) throw new Error(`Hugging Face reported no size for ${row.repo}/${row.filename}`);
|
||||
|
||||
const target = await uploadTarget();
|
||||
const s3 = target.s3;
|
||||
let uploadId = row.uploadId;
|
||||
let key = row.key;
|
||||
// 🔴 A resume MUST address the bucket the multipart upload was created in, which is the one on the
|
||||
// row — not whatever the backend config resolves to now. Re-deriving it sends the remaining parts
|
||||
// to a different bucket than `completeMultipartUpload` names, and the transfer fails at the end
|
||||
// having moved every byte.
|
||||
let bucket = row.bucket ?? target.bucket;
|
||||
const partSize = row.partSize ?? partSizeFor(size);
|
||||
|
||||
if (!uploadId || !key) {
|
||||
bucket = target.bucket;
|
||||
// Refuse rather than invent an owner: `model/0/…` is a key no upload path could produce, and the
|
||||
// userId segment is what `/api/upload/sign-part` authorises against.
|
||||
if (!row.userId)
|
||||
throw new Error(`Import ${row.id} has no owner; refusing to build a key for it`);
|
||||
key = buildUploadKey(
|
||||
UploadType.Model,
|
||||
row.userId,
|
||||
row.filename.split('/').pop() ?? row.filename
|
||||
);
|
||||
uploadId = await createMultipartUpload({ bucket, key, s3 });
|
||||
await dbWrite.huggingFaceImport.update({
|
||||
where: { id: row.id },
|
||||
data: { uploadId, key, bucket, partSize, sizeBytes: BigInt(size) },
|
||||
});
|
||||
}
|
||||
|
||||
const parts: MultipartPart[] = Array.isArray(row.parts) ? (row.parts as MultipartPart[]) : [];
|
||||
const done = new Set(parts.map((part) => part.PartNumber));
|
||||
const totalParts = Math.ceil(size / partSize);
|
||||
|
||||
const sizeOfPart = (partNumber: number) =>
|
||||
Math.min(partNumber * partSize, size) - (partNumber - 1) * partSize;
|
||||
const transferred = () => parts.reduce((sum, part) => sum + sizeOfPart(part.PartNumber), 0);
|
||||
|
||||
// Parts complete out of order, so `done` is a set with holes and never a count. `parts.length + 1`
|
||||
// would re-upload a part already written the moment one finishes ahead of another.
|
||||
const pending: number[] = [];
|
||||
for (let partNumber = 1; partNumber <= totalParts; partNumber++)
|
||||
if (!done.has(partNumber)) pending.push(partNumber);
|
||||
|
||||
let stopped = false;
|
||||
let canceled = false;
|
||||
let movedBytes = 0;
|
||||
|
||||
const movePart = async (partNumber: number) => {
|
||||
const start = (partNumber - 1) * partSize;
|
||||
const end = Math.min(start + partSize, size) - 1;
|
||||
|
||||
const body = await readHuggingFaceRange({ url: row.sourceUrl, start, end });
|
||||
if (body.byteLength !== end - start + 1)
|
||||
throw new Error(
|
||||
`Range ${start}-${end} returned ${body.byteLength} bytes for ${row.repo}/${row.filename}`
|
||||
);
|
||||
|
||||
const etag = await uploadPart({ bucket, key, uploadId, partNumber, body, s3 });
|
||||
parts.push({ PartNumber: partNumber, ETag: etag });
|
||||
done.add(partNumber);
|
||||
movedBytes += body.byteLength;
|
||||
|
||||
// `heartbeatAt` is load-bearing, not telemetry: it is the only thing that stops another run
|
||||
// re-claiming this row once the stale window elapses, which would put two runs on one uploadId.
|
||||
await dbWrite.huggingFaceImport.updateMany({
|
||||
where: { id: row.id, claimedBy: row.claimedBy },
|
||||
data: {
|
||||
parts: parts as unknown as Prisma.InputJsonValue,
|
||||
bytesTransferred: BigInt(transferred()),
|
||||
heartbeatAt: new Date(),
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
const worker = async () => {
|
||||
for (;;) {
|
||||
if (stopped || canceled) return;
|
||||
if (Date.now() >= deadline) {
|
||||
stopped = true;
|
||||
return;
|
||||
}
|
||||
const current = await dbRead.huggingFaceImport.findUnique({
|
||||
where: { id: row.id },
|
||||
select: { status: true },
|
||||
});
|
||||
if (current?.status === 'Canceled') {
|
||||
canceled = true;
|
||||
return;
|
||||
}
|
||||
|
||||
const partNumber = pending.shift();
|
||||
if (partNumber === undefined) return;
|
||||
await movePart(partNumber);
|
||||
}
|
||||
};
|
||||
|
||||
await Promise.all(Array.from({ length: Math.min(partsInFlight, pending.length) }, worker));
|
||||
|
||||
if (canceled || done.size < totalParts) return movedBytes;
|
||||
|
||||
// Re-read the status immediately before finalizing. The workers' probe happens BEFORE each takes
|
||||
// its part, so the last one never checks again — without this, a cancel arriving during the final
|
||||
// part still completes the upload and stamps a URL on a row the moderator stopped.
|
||||
const beforeComplete = await dbWrite.huggingFaceImport.findUnique({
|
||||
where: { id: row.id },
|
||||
select: { status: true, claimedBy: true },
|
||||
});
|
||||
if (beforeComplete?.status === 'Canceled' || beforeComplete?.claimedBy !== row.claimedBy)
|
||||
return movedBytes;
|
||||
|
||||
parts.sort((a, b) => a.PartNumber - b.PartNumber);
|
||||
await completeMultipartUpload(bucket, key, uploadId, parts, s3);
|
||||
|
||||
// Presigned GET with the query stripped: the same `https://<endpoint>/<bucket>/<key>` shape the
|
||||
// browser path writes to `ModelFile.url`, without a second copy of the endpoint config here.
|
||||
// Deliberately not `getCustomPutUrl` — that bumps `recordB2PresignIssued`, a counter whose whole
|
||||
// purpose is measuring browser-direct uploads, and a server transfer is not one.
|
||||
const { url } = await getGetUrlByKey(key, { s3, bucket });
|
||||
await dbWrite.huggingFaceImport.updateMany({
|
||||
where: { id: row.id, claimedBy: row.claimedBy, status: { not: 'Canceled' } },
|
||||
data: {
|
||||
status: 'Completed',
|
||||
url: url.split('?')[0],
|
||||
bytesTransferred: BigInt(size),
|
||||
completedAt: new Date(),
|
||||
error: null,
|
||||
attempts: 0,
|
||||
// Spent. A retained id makes every later abort fail against a finished upload, burying the
|
||||
// one abort failure that means parts are still billed.
|
||||
uploadId: null,
|
||||
},
|
||||
});
|
||||
return movedBytes;
|
||||
}
|
||||
|
||||
/**
|
||||
* Drains the queue until `deadline`. Bounded work per call by design: a transfer is a sequence of
|
||||
* resumable parts, so the job never needs a run longer than its own lock.
|
||||
*/
|
||||
export async function processImportQueue({
|
||||
deadline,
|
||||
worker,
|
||||
concurrency,
|
||||
partsInFlight,
|
||||
}: {
|
||||
deadline: number;
|
||||
worker: string;
|
||||
concurrency: number;
|
||||
partsInFlight: number;
|
||||
}) {
|
||||
let moved = 0;
|
||||
let bytes = 0;
|
||||
|
||||
const drain = async () => {
|
||||
while (Date.now() < deadline) {
|
||||
const row = await claimNext(worker);
|
||||
if (!row) return;
|
||||
try {
|
||||
bytes += await advanceImport(row, deadline, partsInFlight);
|
||||
moved++;
|
||||
} catch (error) {
|
||||
await failOrRetry(row, error);
|
||||
continue;
|
||||
}
|
||||
await yieldClaim(row.id).catch(() => undefined);
|
||||
}
|
||||
};
|
||||
|
||||
const startedAt = Date.now();
|
||||
await Promise.all(Array.from({ length: concurrency }, drain));
|
||||
const seconds = (Date.now() - startedAt) / 1000;
|
||||
|
||||
// Throughput is the number every question about this feature turns on — whether it is too slow,
|
||||
// whether it is starving the pod, what any bandwidth limit should be set to — and nothing else
|
||||
// records it. Emitted per run rather than per part so a quiet run is one line, not none.
|
||||
if (bytes)
|
||||
logToAxiom({
|
||||
type: 'info',
|
||||
name: 'huggingface-import-throughput',
|
||||
bytes,
|
||||
seconds: Math.round(seconds),
|
||||
bytesPerSecond: Math.round(bytes / Math.max(seconds, 1)),
|
||||
filesTouched: moved,
|
||||
concurrency,
|
||||
partsInFlight,
|
||||
});
|
||||
|
||||
return { moved, bytes };
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
import { env } from '~/env/server';
|
||||
|
||||
const HF_API = 'https://huggingface.co/api';
|
||||
const HF_HOST = 'https://huggingface.co';
|
||||
|
||||
export type HuggingFaceFile = {
|
||||
path: string;
|
||||
size: number;
|
||||
/** Content sha256, present for LFS files — which is every weight file. Small non-LFS files carry a
|
||||
* git blob sha1 instead, which is not comparable to a `ModelFileHash`, so they report null. */
|
||||
sha256: string | null;
|
||||
};
|
||||
|
||||
export type HuggingFaceRepo = {
|
||||
repo: string;
|
||||
revision: string;
|
||||
license: string | null;
|
||||
gated: string | false;
|
||||
files: HuggingFaceFile[];
|
||||
};
|
||||
|
||||
export class HuggingFaceError extends Error {
|
||||
constructor(message: string, readonly status?: number) {
|
||||
super(message);
|
||||
this.name = 'HuggingFaceError';
|
||||
}
|
||||
}
|
||||
|
||||
const WEIGHT_EXTENSIONS = /\.(safetensors|sft|ckpt|pt|pth|bin|gguf|onnx)$/i;
|
||||
/** Per part, not per file — the run's budget is only checked BETWEEN parts, so one stalled read
|
||||
* otherwise defers that check indefinitely and the next tick starts a second run. */
|
||||
const RANGE_READ_TIMEOUT_MS = 90_000;
|
||||
|
||||
export function isWeightFile(path: string) {
|
||||
return WEIGHT_EXTENSIONS.test(path);
|
||||
}
|
||||
|
||||
/**
|
||||
* A `ModelFile.type` guess from the file's path, or null when the path doesn't say.
|
||||
*
|
||||
* 🔴 Advisory only — the attach call takes the type explicitly. Naming conventions are the whole
|
||||
* evidence here, and a wrong guess on the *primary* weights is the expensive one: it decides whether
|
||||
* a version is loadable at all. So the primary case is exactly the one this refuses to infer.
|
||||
*/
|
||||
export function suggestFileType(path: string): 'VAE' | 'Text Encoder' | 'Config' | null {
|
||||
const name = path.toLowerCase();
|
||||
const base = name.split('/').pop() ?? name;
|
||||
const stem = base.replace(/\.[^.]+$/, '');
|
||||
|
||||
if (/\.(json|yaml|yml|txt|md)$/.test(base)) return 'Config';
|
||||
// A containing directory is decisive; otherwise the WHOLE basename must be the accessory.
|
||||
// `flux1-dev-vae-baked.safetensors` is the primary weights and merely names its bundled VAE —
|
||||
// claiming it is the one mistake that produces a version nothing can load.
|
||||
if (/(^|\/)text_encoder/.test(name)) return 'Text Encoder';
|
||||
if (/^(t5|umt5|clip|clip_[lgh]|open_?clip)[\w.-]*$/.test(stem)) return 'Text Encoder';
|
||||
if (/(^|\/)vae(\/|$)/.test(name)) return 'VAE';
|
||||
if (/^(ae|vae)([_.-][\w.-]*)?$/.test(stem)) return 'VAE';
|
||||
return null;
|
||||
}
|
||||
|
||||
/** Accepts a repo URL, a `/tree/<rev>` or `/blob/<rev>/<file>` URL, or a bare `owner/name`. */
|
||||
export function parseHuggingFaceRepo(input: string): { repo: string; revision?: string } | null {
|
||||
const trimmed = input.trim();
|
||||
if (!trimmed) return null;
|
||||
|
||||
const withoutHost = trimmed
|
||||
.replace(/^https?:\/\/(www\.)?huggingface\.co\//i, '')
|
||||
.replace(/^\/+/, '');
|
||||
const segments = withoutHost.split('?')[0].split('#')[0].split('/').filter(Boolean);
|
||||
// A copied URL sometimes carries a `models/` prefix; datasets and spaces are a different API.
|
||||
const parts = segments[0] === 'models' ? segments.slice(1) : segments;
|
||||
if (parts.length < 2) return null;
|
||||
|
||||
const repo = `${parts[0]}/${parts[1]}`;
|
||||
const marker = parts[2];
|
||||
const revision = marker === 'tree' || marker === 'blob' ? parts[3] : undefined;
|
||||
return revision ? { repo, revision } : { repo };
|
||||
}
|
||||
|
||||
function authHeaders(): Record<string, string> {
|
||||
const token = env.HUGGING_FACE_TOKEN;
|
||||
return token ? { Authorization: `Bearer ${token}` } : {};
|
||||
}
|
||||
|
||||
async function hfFetch(url: string) {
|
||||
const res = await fetch(url, { headers: authHeaders() });
|
||||
if (!res.ok) {
|
||||
const detail = await res.text().catch(() => '');
|
||||
throw new HuggingFaceError(
|
||||
res.status === 401 || res.status === 403
|
||||
? `Hugging Face refused the request (${res.status}). The repo is gated or private; importing it needs a HUGGING_FACE_TOKEN whose account has accepted its terms.`
|
||||
: `Hugging Face returned ${res.status}${detail ? `: ${detail.slice(0, 200)}` : ''}`,
|
||||
res.status
|
||||
);
|
||||
}
|
||||
return res.json();
|
||||
}
|
||||
|
||||
type TreeEntry = {
|
||||
type: string;
|
||||
path: string;
|
||||
size?: number;
|
||||
oid?: string;
|
||||
lfs?: { oid?: string; size?: number };
|
||||
};
|
||||
|
||||
/**
|
||||
* Lists a repo's files at a pinned revision, resolving a branch to its commit sha — an import records
|
||||
* which bytes it took, and `main` moves.
|
||||
*/
|
||||
export async function getRepoFiles(input: {
|
||||
repo: string;
|
||||
revision?: string;
|
||||
}): Promise<HuggingFaceRepo> {
|
||||
const { repo } = input;
|
||||
const infoUrl = input.revision
|
||||
? `${HF_API}/models/${repo}/revision/${encodeURIComponent(input.revision)}`
|
||||
: `${HF_API}/models/${repo}`;
|
||||
const info = (await hfFetch(infoUrl)) as {
|
||||
id?: string;
|
||||
sha?: string;
|
||||
gated?: string | false;
|
||||
cardData?: { license?: string | string[] };
|
||||
};
|
||||
|
||||
// 🔴 The repo id comes from HF's response, not from what was pasted. It is the group every import
|
||||
// is filed under, and a URL typed with different casing would otherwise file the same repo under
|
||||
// two groups that nothing could merge.
|
||||
const canonicalRepo = info.id ?? repo;
|
||||
|
||||
const revision = info.sha ?? input.revision;
|
||||
if (!revision) throw new HuggingFaceError(`Could not resolve a commit sha for ${repo}`);
|
||||
|
||||
const tree = (await hfFetch(
|
||||
`${HF_API}/models/${repo}/tree/${encodeURIComponent(revision)}?recursive=true`
|
||||
)) as TreeEntry[];
|
||||
|
||||
const license = Array.isArray(info.cardData?.license)
|
||||
? info.cardData?.license[0] ?? null
|
||||
: info.cardData?.license ?? null;
|
||||
|
||||
return {
|
||||
repo: canonicalRepo,
|
||||
revision,
|
||||
license,
|
||||
gated: info.gated ?? false,
|
||||
files: tree
|
||||
.filter((entry) => entry.type === 'file')
|
||||
.map((entry) => ({
|
||||
path: entry.path,
|
||||
size: entry.lfs?.size ?? entry.size ?? 0,
|
||||
sha256: entry.lfs?.oid ?? null,
|
||||
}))
|
||||
.sort((a, b) => b.size - a.size),
|
||||
};
|
||||
}
|
||||
|
||||
/** The name a batch gets unless a moderator types a different one: the repo's own name, without the owner. */
|
||||
export function defaultGroupName(repo: string) {
|
||||
return repo.split('/').pop() ?? repo;
|
||||
}
|
||||
|
||||
export function huggingFaceResolveUrl(repo: string, revision: string, path: string) {
|
||||
const encodedPath = path.split('/').map(encodeURIComponent).join('/');
|
||||
return `${HF_HOST}/${repo}/resolve/${encodeURIComponent(revision)}/${encodedPath}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads one byte range of a file.
|
||||
*
|
||||
* A server that ignores `Range` answers 200 with the WHOLE file, so only a 206 is accepted.
|
||||
*/
|
||||
export async function readHuggingFaceRange({
|
||||
url,
|
||||
start,
|
||||
end,
|
||||
signal,
|
||||
timeoutMs = RANGE_READ_TIMEOUT_MS,
|
||||
}: {
|
||||
url: string;
|
||||
start: number;
|
||||
/** Inclusive, as HTTP ranges are. */
|
||||
end: number;
|
||||
signal?: AbortSignal;
|
||||
timeoutMs?: number;
|
||||
}): Promise<Uint8Array> {
|
||||
const timeout = AbortSignal.timeout(timeoutMs);
|
||||
const res = await fetch(url, {
|
||||
headers: { ...authHeaders(), Range: `bytes=${start}-${end}` },
|
||||
signal: signal ? AbortSignal.any([signal, timeout]) : timeout,
|
||||
});
|
||||
if (res.status !== 206) {
|
||||
throw new HuggingFaceError(
|
||||
`Expected 206 for range ${start}-${end} of ${url}, got ${res.status}`,
|
||||
res.status
|
||||
);
|
||||
}
|
||||
return new Uint8Array(await res.arrayBuffer());
|
||||
}
|
||||
|
||||
export async function headHuggingFaceFile(url: string) {
|
||||
const res = await fetch(url, { method: 'HEAD', headers: authHeaders(), redirect: 'follow' });
|
||||
if (!res.ok) throw new HuggingFaceError(`HEAD failed with ${res.status} for ${url}`, res.status);
|
||||
const length = Number(res.headers.get('content-length') ?? '0');
|
||||
return Number.isFinite(length) && length > 0 ? length : null;
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { getModelFileTypeOptions } from '~/utils/file-display-helpers';
|
||||
|
||||
const values = (options: { value: string }[]) => options.map((option) => option.value);
|
||||
const labelOf = (options: { value: string; label: string }[], value: string) =>
|
||||
options.find((option) => option.value === value)?.label;
|
||||
|
||||
describe('getModelFileTypeOptions', () => {
|
||||
it('offers only the types the extension can be', () => {
|
||||
const options = values(getModelFileTypeOptions('config.yaml'));
|
||||
expect(options).toContain('Config');
|
||||
expect(options).not.toContain('Model');
|
||||
});
|
||||
|
||||
it('keeps the current type selectable even when the extension would exclude it', () => {
|
||||
// A legacy file must not render as a blank Select.
|
||||
expect(values(getModelFileTypeOptions('config.yaml'))).not.toContain('VAE');
|
||||
expect(values(getModelFileTypeOptions('config.yaml', { currentType: 'VAE' }))).toContain('VAE');
|
||||
});
|
||||
|
||||
it('restricts to the types it is given', () => {
|
||||
expect(
|
||||
values(getModelFileTypeOptions('model.safetensors', { types: ['Model', 'VAE'] }))
|
||||
).toEqual(['Model', 'VAE']);
|
||||
});
|
||||
|
||||
it('labels the generic Model option by the model type', () => {
|
||||
const options = getModelFileTypeOptions('model.safetensors', {
|
||||
types: ['Model'],
|
||||
modelType: 'TextualInversion',
|
||||
});
|
||||
expect(labelOf(options, 'Model')).toBe('Embedding');
|
||||
});
|
||||
|
||||
it('prefers the ComfyUI label over the raw type name', () => {
|
||||
const options = getModelFileTypeOptions('upscaler.safetensors', { types: ['Upscaler'] });
|
||||
expect(labelOf(options, 'Upscaler')).toBe('Upscale Model');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,46 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import { buildUploadKey } from '~/utils/upload-key';
|
||||
|
||||
/**
|
||||
* `buildUploadKey` is the only sanitiser applied to a client-supplied filename on `/api/upload`, and
|
||||
* the shape it produces is an authorisation contract: `/api/upload/sign-part` authorises a part by
|
||||
* comparing `key.split('/')[1]` to the session's user id.
|
||||
*
|
||||
* The endpoint's own suite asserts the key as `expect.any(String)`, so before this file existed both
|
||||
* the collision token and the sanitiser could be deleted with every test still green — which in
|
||||
* production is one upload silently overwriting another's object.
|
||||
*/
|
||||
describe('buildUploadKey', () => {
|
||||
it('puts the userId in segment 1, where sign-part authorises against it', () => {
|
||||
const key = buildUploadKey('model', 42, 'thing.safetensors');
|
||||
expect(key.split('/')[0]).toBe('model');
|
||||
expect(key.split('/')[1]).toBe('42');
|
||||
});
|
||||
|
||||
it('sanitises the filename and appends a collision token', () => {
|
||||
expect(buildUploadKey('model', 42, 'My Model v1.0.safetensors')).toMatch(
|
||||
/^model\/42\/[A-Za-z0-9_]+\.[A-Za-z0-9]{4}\.safetensors$/
|
||||
);
|
||||
});
|
||||
|
||||
it('gives two uploads of the same filename different keys', () => {
|
||||
const a = buildUploadKey('model', 42, 'model.safetensors');
|
||||
const b = buildUploadKey('model', 42, 'model.safetensors');
|
||||
// Without the token these collide, and the second upload overwrites the first object while every
|
||||
// ModelFile.url already pointing at it silently starts serving different bytes.
|
||||
expect(a).not.toBe(b);
|
||||
});
|
||||
|
||||
it('cannot be made to add path segments from the filename', () => {
|
||||
const key = buildUploadKey('model', 42, '../../etc/pass wd.safetensors');
|
||||
expect(key.split('/')).toHaveLength(3);
|
||||
expect(key.split('/')[1]).toBe('42');
|
||||
});
|
||||
|
||||
it('keeps an extensionless filename in one segment', () => {
|
||||
const key = buildUploadKey('default', 7, 'README');
|
||||
expect(key.split('/')).toHaveLength(3);
|
||||
expect(key.startsWith('default/7/')).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -4,8 +4,14 @@
|
||||
*/
|
||||
|
||||
import type { ModelFileType } from '~/server/common/constants';
|
||||
import { constants } from '~/server/common/constants';
|
||||
import type { ModelType } from '~/shared/utils/prisma/enums';
|
||||
import { filenamize, getFileExtension, replaceInsensitive } from '~/utils/string-helpers';
|
||||
import {
|
||||
filenamize,
|
||||
getDisplayName,
|
||||
getFileExtension,
|
||||
replaceInsensitive,
|
||||
} from '~/utils/string-helpers';
|
||||
|
||||
/**
|
||||
* Metadata shape expected for file display functions
|
||||
@@ -80,6 +86,31 @@ export function filterFileTypeByExtension(value: ModelFileType, fileName: string
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* File-type options for a filename, with display labels. `currentType` stays selectable even
|
||||
* when not otherwise offered, so an existing file never renders as a blank Select.
|
||||
*/
|
||||
export function getModelFileTypeOptions(
|
||||
fileName: string,
|
||||
{
|
||||
types = constants.modelFileTypes,
|
||||
currentType,
|
||||
modelType,
|
||||
}: {
|
||||
types?: readonly ModelFileType[];
|
||||
currentType?: ModelFileType | null;
|
||||
modelType?: ModelType | null;
|
||||
} = {}
|
||||
) {
|
||||
return types
|
||||
.filter((type) => type === currentType || filterFileTypeByExtension(type, fileName))
|
||||
.map((type) => ({
|
||||
value: type,
|
||||
label:
|
||||
comfyFileTypeLabels[type] ?? getDisplayName(type === 'Model' ? modelType ?? type : type),
|
||||
}));
|
||||
}
|
||||
|
||||
/**
|
||||
* Minimal file shape for display functions
|
||||
*/
|
||||
|
||||
@@ -683,6 +683,61 @@ export async function getMultipartPutUrl(
|
||||
return { urls, bucket, key, uploadId: UploadId, chunkSize };
|
||||
}
|
||||
|
||||
/**
|
||||
* Start a multipart upload WITHOUT presigning its parts, for a server-side transfer that uploads
|
||||
* parts itself. {@link getMultipartPutUrl} presigns every part up front because the browser has no
|
||||
* credentials; a transfer that spans several job runs cannot use those URLs — they expire, and the
|
||||
* part count is only known once the source reports its size.
|
||||
*/
|
||||
export async function createMultipartUpload({
|
||||
bucket,
|
||||
key,
|
||||
mimeType,
|
||||
s3,
|
||||
}: {
|
||||
bucket: string;
|
||||
key: string;
|
||||
mimeType?: string;
|
||||
s3?: S3Client | null;
|
||||
}) {
|
||||
if (!s3) s3 = getS3Client();
|
||||
const { UploadId } = await s3.send(
|
||||
new CreateMultipartUploadCommand({ Bucket: bucket, Key: key, ContentType: mimeType })
|
||||
);
|
||||
if (!UploadId) throw new Error(`S3 returned no UploadId for ${key}`);
|
||||
return UploadId;
|
||||
}
|
||||
|
||||
/** Upload one part of an in-flight multipart upload from the pod. Returns the part's ETag. */
|
||||
export async function uploadPart({
|
||||
bucket,
|
||||
key,
|
||||
uploadId,
|
||||
partNumber,
|
||||
body,
|
||||
s3,
|
||||
}: {
|
||||
bucket: string;
|
||||
key: string;
|
||||
uploadId: string;
|
||||
partNumber: number;
|
||||
body: Uint8Array;
|
||||
s3?: S3Client | null;
|
||||
}) {
|
||||
if (!s3) s3 = getS3Client();
|
||||
const { ETag } = await s3.send(
|
||||
new UploadPartCommand({
|
||||
Bucket: bucket,
|
||||
Key: key,
|
||||
UploadId: uploadId,
|
||||
PartNumber: partNumber,
|
||||
Body: body,
|
||||
})
|
||||
);
|
||||
if (!ETag) throw new Error(`S3 returned no ETag for part ${partNumber} of ${key}`);
|
||||
return ETag;
|
||||
}
|
||||
|
||||
interface MultipartUploadPart {
|
||||
ETag: string;
|
||||
PartNumber: number;
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
import { extname } from 'node:path';
|
||||
import { filenamize, generateToken } from '~/utils/string-helpers';
|
||||
|
||||
/**
|
||||
* The one place an upload key is built. Every uploaded object — browser or server — is
|
||||
* `<type>/<userId>/<name>.<token><ext>`, and that shape is load-bearing beyond tidiness:
|
||||
* `/api/upload/sign-part` authorises a part by reading the userId out of segment 1, so nothing may be
|
||||
* inserted ahead of it.
|
||||
*
|
||||
* 🔴 Grouping does not belong in a key. A key is immutable once the object exists, so anything encoded
|
||||
* in it can never be corrected without copying the bytes. The `HuggingFaceImport` row is what makes an
|
||||
* imported object navigable — repo, revision and filename are columns, and columns can be fixed.
|
||||
*
|
||||
* Lives here rather than in `s3-utils` because that module builds S3 clients at load: a consumer that
|
||||
* only needs to name a file should not drag credentials into its module graph, and a test that mocks
|
||||
* it should not have to stand up an endpoint config.
|
||||
*/
|
||||
export function buildUploadKey(type: string, userId: number, fullFilename: string) {
|
||||
const ext = extname(fullFilename);
|
||||
const name = filenamize(fullFilename.replace(ext, ''));
|
||||
return `${type}/${userId}/${name}.${generateToken(4)}${ext}`;
|
||||
}
|
||||
Reference in New Issue
Block a user