Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,40 @@

All notable TrainKit changes are documented here. TrainKit follows semantic versioning.

## [1.3.0] - 2026-10-07

TrainKit 1.3.0 adds Claude and OpenAI captioning through your own API keys and improves batch safety and desktop reliability.

### Added

- API connections tab with masked key entry, session-only keys by default, optional OS-encrypted persistence, editable vision model IDs, removal, and an explicit connection test that sends no images.
- Caption provider selection with unavailable providers disabled.
- Anthropic Messages and OpenAI Responses image captioning with resized, metadata-free JPEGs, bounded responses, cancellable requests, and safe errors without automatic paid retries.

### Changed

- Local and cloud captioning reuse the same collision-safe batch loop, atomic output writes, and optional resumable manifests.
- Cloud batches stop on provider failures, refusals, or incomplete output; dry runs and fully skipped or completed batches make no provider requests.
- Updated Electron to 44.6.0, Electron Forge to 8, ESLint to 10, Vitest to 5, and Lucide to 1; refreshed both dependency lockfiles and raised the development Node.js minimum to 22.17.
- Removed unused comparison-slider and import-lint dependencies. TypeScript stays on 5.9 to match the supported range of the ESLint parser.
- Removed dry-run controls from the desktop panels and aligned the manifest checkbox with the collision policy selector.

### Fixed

- Failed model switches no longer retain a cleaned-up service in the model cache.
- Queued cancellation prevents the runner from starting, and missed job events are recovered after reconnecting.
- Caption model controls stay consistent during loading, unloading, and completed runs; duplicate submissions and cancellation errors are handled centrally.
- Backend spawn failures no longer hang shutdown, and empty error responses produce useful messages.
- New output directories and generated manifests resolve existing junctions; dangling folder links are rejected.
- Unreadable remembered API keys can be removed without entering a replacement key.
- Cloud captioning applies photo rotation and mirroring before removing metadata.
- Cloud captioning distinguishes local file-save failures from provider errors.
- Rename batches reject invalid images, restrict duplicate detection to supported images, and use deterministic tie-breakers for natural ordering.
- Single-channel NCNN output is converted to RGB and invalid tiled overlap is rejected before processing.
- Failed atomic text writes remove their temporary files, and damaged provider records do not discard another provider's remembered key.
- Development setup preserves test tools and excludes backend environments from frontend file watching.
- Setup and startup windows have small top-right minimize and quit buttons. Minimizing setup keeps the main window minimized when it opens; quitting stops setup.

## [1.2.1] - 2026-07-15

TrainKit 1.2.1 fixes NCNN graph compatibility, makes manifests opt-in, and adds individual-image workflows across every operation.
Expand Down
6 changes: 4 additions & 2 deletions CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ Thanks for improving TrainKit. Keep changes focused, add regression coverage for

## Development setup

Install Node.js 22.12+, uv, Python 3.12, and the Microsoft Visual C++ x64 Redistributable on Windows. Then run:
Install Node.js 22.17+, uv, Python 3.12, and the Microsoft Visual C++ x64 Redistributable on Windows. Then run:

```powershell
npm ci
Expand All @@ -30,6 +30,8 @@ npm run verify:package
## Project invariants

- The renderer never receives the backend token or its port.
- Configured API keys stay in the main process and per-job authenticated backend headers; never log them or include them in manifests. Persistence requires explicit opt-in and OS encryption.
- Cloud captioning must not automatically retry paid requests.
- Renderer filesystem access must flow through narrow IPC and a user-selected path grant.
- Backend HTTP and WebSocket access remains authenticated and loopback-only.
- Resume manifests are untrusted and must remain scoped to the requested input/output roots.
Expand All @@ -46,7 +48,7 @@ User-visible changes must update `CHANGELOG.md` under the target version. The re

## Maintainer release setup

Tagged releases must match `package.json` exactly (for example, version `1.2.1` requires tag `v1.2.1`). Releases are unsigned by default. To Authenticode-sign the packaged executable, configure both GitHub Actions secrets:
Tagged releases must match `package.json` exactly (for example, version `1.3.0` requires tag `v1.3.0`). Releases are unsigned by default. To Authenticode-sign the packaged executable, configure both GitHub Actions secrets:

- `WINDOWS_CERTIFICATE_BASE64`: Base64-encoded PFX code-signing certificate.
- `WINDOWS_CERTIFICATE_PASSWORD`: PFX password.
Expand Down
36 changes: 22 additions & 14 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,14 @@

<h1 align="center">TrainKit</h1>

<p align="center">A local, resumable dataset-preparation desktop app for AI image training.</p>
<p align="center">Prepare image datasets for AI training.</p>

<p align="center">
<a href="https://github.com/LightCyan01/TrainKit/actions/workflows/ci.yml"><img src="https://github.com/LightCyan01/TrainKit/actions/workflows/ci.yml/badge.svg" alt="CI"></a>
<a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue.svg" alt="MIT License"></a>
<a href="https://github.com/LightCyan01/TrainKit/releases/latest"><img src="https://img.shields.io/github/v/release/LightCyan01/TrainKit" alt="Latest release"></a>
<img src="https://img.shields.io/badge/platform-Windows-lightgrey" alt="Windows">
<img src="https://img.shields.io/badge/Electron-43-47848F?logo=electron" alt="Electron 43">
<img src="https://img.shields.io/badge/Electron-44-47848F?logo=electron" alt="Electron 44">
<img src="https://img.shields.io/badge/Python-3.12-3776AB?logo=python&logoColor=white" alt="Python 3.12">
<img src="https://img.shields.io/badge/PyTorch-2.11-EE4C2C?logo=pytorch&logoColor=white" alt="PyTorch 2.11">
</p>
Expand All @@ -23,21 +23,21 @@
<a href="docs/architecture.md">Architecture</a>
</p>

TrainKit prepares individual images or complete image folders without uploading them to a hosted service. Every run is deterministic, cancellable, and collision-aware; resumable manifest files are optional.
TrainKit prepares individual images or image folders with local captioning, upscaling, tagging, and renaming. Claude and OpenAI captioning use your own API keys. You can cancel batches and save progress in a manifest to resume later.

## What's new in 1.2.1
## What's new in 1.3.0

- Automatic NCNN input/output blob discovery instead of hard-coded `in0` and `out0` names.
- Individual-image selection for captioning, upscaling, tagging, and renaming.
- Manifests disabled by default, with explicit save, dry-run, and resume options.
- First processing errors remain visible even when no manifest is written.
- A single standard Windows ZIP distribution with no installer.
- Claude and OpenAI captioning through your own API keys, alongside local models.
- An API tab to save session keys, remember keys with OS encryption, choose vision models, and test key connections.
- Cloud requests can be cancelled. Failed requests stop the batch without automatic paid retries.
- Recovered progress after reconnecting, safer model switching and shutdown, and fixes for batch file handling.
- Updated Electron and packaging dependencies with a clean dependency audit.

See the [complete 1.2.1 changelog](CHANGELOG.md#121---2026-07-15) for details.
See the [complete 1.3.0 changelog](CHANGELOG.md#130---2026-10-07) for details.

## Features

- **Captioning:** local standard Transformers multimodal-chat, BLIP, and InstructBLIP models; auto-detection, explicit preload/unload, prompt control, and token-level cancellation.
- **Captioning:** local standard Transformers multimodal-chat, BLIP, and InstructBLIP models, or Claude/OpenAI vision APIs. Local models support auto-detection, preload/unload, prompt control, and token-level cancellation.
- **Upscaling:** safe `.safetensors` models through Spandrel, or NCNN `.param` + `.bin` models through the Python bindings with automatic blob discovery, CPU/Vulkan, and tiled inference.
- **Tagging:** local Hugging Face image-classification models with threshold/top-K controls and scored JSON, training-text, or paired sidecars.
- **Renaming:** natural input ordering, configurable zero padding, optional visual-duplicate filtering, and atomic copies.
Expand All @@ -61,7 +61,7 @@ An NVIDIA GPU is recommended for Transformers and Spandrel workloads. NCNN can r

## Build from source

Requirements: Git, Node.js 22.12 or newer, uv, Python 3.12 (uv can install it), and the Visual C++ x64 Redistributable on Windows.
Requirements: Git, Node.js 22.17 or newer, uv, Python 3.12 (uv can install it), and the Visual C++ x64 Redistributable on Windows.

```powershell
git clone https://github.com/LightCyan01/TrainKit.git
Expand Down Expand Up @@ -90,13 +90,21 @@ Normal runs do not write a manifest. Enable **Save resumable manifest** when you
<output>/.trainkit/manifests/<job-id>.json
```

Dry runs always write a manifest because the plan is their output. A resume uses the same operation and input/output paths, preserves completed and skipped items, retries failed items, and updates the selected manifest. Manifests are treated as untrusted input: paths outside the selected roots are rejected.
To resume, select the manifest and use the same operation and input/output paths. TrainKit preserves completed and skipped items, retries failed items, and updates that manifest. It rejects paths outside the selected folders.

## Cloud captioning

Open **API**, paste an Anthropic or OpenAI API key, choose a vision model available to your account, and save. Keys last for the session unless you select **Remember on this computer** to store them with OS encryption. Remembered keys are tied to your operating-system account. **Test key** checks authentication; it does not verify access to every model. API billing is separate from ChatGPT and Claude subscriptions.

In **Caption**, choose the configured provider, select your images and output folder, enter an instruction, and click **Start captioning**. Images are converted to JPEG, metadata is removed, and the longest edge is capped at 1568 pixels. Cloud captioning sends your instruction and resized images to the selected provider and uses API credits; the provider's data policies apply. OpenAI requests use `store: false`, which does not replace the provider's retention policy. Captions are saved as local UTF-8 text files.

Cloud requests run one at a time and are not automatically retried. The batch stops if the provider rejects your key, refuses a caption, or returns an incomplete response. Connection failures, exhausted quota, and rate limits also stop the batch. Save a manifest if you want to resume pending work. Cancellation stops pending requests and prevents saving their captions, but cannot undo provider work or charges already incurred. Completed resumes and skipped outputs make no provider requests. Local mode remains the default.

See [model support](docs/model-support.md) for model layouts and NCNN assumptions, [architecture](docs/architecture.md) for the desktop/backend trust boundaries, and the [changelog](CHANGELOG.md) for complete release history.

## Project status

Version 1.2.1 completes the original processing roadmap and improves NCNN compatibility, input selection, and output control. New feature proposals and model-compatibility reports are welcome through GitHub issues.
Version 1.3.0 adds cloud captioning and improves batch and desktop reliability. New feature proposals and model-compatibility reports are welcome through GitHub issues.

## Contributing and security

Expand Down
6 changes: 5 additions & 1 deletion backend/core/jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,8 +153,12 @@ async def reserve(self, operation: str):

async def _execute(self, record: JobRecord, cancelled: Event, runner: Runner):
context = JobContext(self, record, cancelled)
await self.update(record.job_id, status="running", message=f"Starting {record.operation}")
try:
context.raise_if_cancelled()
await self.update(
record.job_id, status="running", message=f"Starting {record.operation}"
)
context.raise_if_cancelled()
manifest_path = await runner(context)
context.raise_if_cancelled()
await self.update(
Expand Down
23 changes: 21 additions & 2 deletions backend/models/requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from datetime import datetime
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

CollisionPolicy = Literal["fail", "skip", "overwrite", "rename"]
JobOperation = Literal["rename", "caption", "upscale", "tag"]
Expand Down Expand Up @@ -38,13 +38,32 @@ class UpscaleRequest(BatchRequest):
ncnn_scale: int = Field(default=4, ge=1, le=16)
ncnn_use_vulkan: bool = True

@model_validator(mode="after")
def validate_tiling(self) -> UpscaleRequest:
if self.use_tiling and self.tile_overlap >= self.tile_size:
raise ValueError("Tile overlap must be smaller than tile size")
return self


class CaptionRequest(BatchRequest):
caption_model_path: str
caption_model_path: str = ""
provider: Literal["local", "anthropic", "openai"] = "local"
cloud_model: str = Field(default="", max_length=200)
prompt: str = Field(min_length=1, max_length=8000)
adapter: Literal["auto", "multimodal", "blip", "instructblip"] = "auto"
max_new_tokens: int = Field(default=256, ge=1, le=2048)

@model_validator(mode="after")
def validate_provider(self) -> CaptionRequest:
if self.provider == "local":
if not self.caption_model_path.strip():
raise ValueError("Local captioning requires a model folder")
else:
self.cloud_model = self.cloud_model.strip()
if not self.cloud_model:
raise ValueError("Cloud captioning requires a model ID")
return self


class TagRequest(BatchRequest):
tag_model_path: str
Expand Down
1 change: 1 addition & 0 deletions backend/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ dependencies = [
"accelerate>=1.11.0",
"difpy>=4.2.1",
"fastapi>=0.139.0,<0.140",
"httpx>=0.28.1,<1",
"ncnn>=1.0.20260526",
"numpy>=2.3.0",
"pillow>=12.3.0,<13",
Expand Down
17 changes: 16 additions & 1 deletion backend/routers/caption.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,14 @@

from pathlib import Path

from fastapi import APIRouter, Depends
from fastapi import APIRouter, Depends, Header

from core.dependencies import (
get_connection_manager,
get_job_manager,
get_service_manager,
)
from core.exceptions import TrainKitException
from core.jobs import JobManager
from core.websocket import ConnectionManager
from models import CaptionRequest, JobResponse, ModelStatusRequest, PreloadRequest
Expand All @@ -19,10 +20,24 @@
@router.post("/caption", response_model=JobResponse, status_code=202)
async def caption(
request: CaptionRequest,
provider_key: str | None = Header(default=None, alias="x-trainkit-provider-key"),
jobs: JobManager = Depends(get_job_manager),
services=Depends(get_service_manager),
):
if request.provider != "local" and not request.dry_run:
if not provider_key or not provider_key.strip():
raise TrainKitException(
"Configure the provider API key first", 400, "provider_key_missing"
)
provider_key = provider_key.strip()
if len(provider_key) > 8192 or any(not 33 <= ord(char) <= 126 for char in provider_key):
raise TrainKitException("Invalid provider API key", 400, "provider_key_invalid")

async def run(context):
if request.provider != "local":
from service.cloud_captioning import process_cloud_caption_batch

return await process_cloud_caption_batch(request, context, provider_key or "")
service = services.get_caption_service(
Path(request.caption_model_path), request.adapter, request.max_new_tokens
)
Expand Down
2 changes: 0 additions & 2 deletions backend/routers/rename.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,6 @@ async def rename(
jobs: JobManager = Depends(get_job_manager),
service=Depends(get_rename_service),
):
service.clear_cache()

async def run(context):
return await service.process(request, context)

Expand Down
Loading
Loading