Skip to content

A dataset preparation toolkit for AI image training

Notifications You must be signed in to change notification settings

LightCyan01/TrainKit

Repository files navigation

TrainKit Logo

TrainKit

A dataset preparation toolkit for AI image training

License Release Platform Electron Python PyTorch


About

TrainKit streamlines the tedious process of preparing image datasets for AI model training. It provides GPU-accelerated tools for captioning, upscaling, and organizing your images wrapped in a modern desktop interface.

Features

  • 🖼️ Image Captioning — Generate captions using local LLaVA-based models (e.g., JoyCaption) with customizable prompts
  • ⬆️ Image Upscaling — Upscale images using Spandrel supported architectures
  • 📁 Batch Renaming — Rename and organize image files with flexible naming patterns.

Requirements

  • Windows 10/11
  • NVIDIA GPU (8GB+ VRAM recommended for captioning)
  • Python 3.12
  • Node.js 18+

Installation

From Release

Download the latest release from the Releases page.

From Source

# Clone the repository
git clone https://github.com/LightCyan01/TrainKit.git
cd TrainKit

# Install frontend dependencies
npm install

# Install backend dependencies
cd backend
uv sync

# Start the app
npm start

Usage

  1. Launch TrainKit — The backend server starts automatically
  2. Select a panel — Choose Caption, Upscale, or Rename from the sidebar
  3. Configure paths — Set your input folder, output folder, and model path
  4. Start processing — Click the action button and monitor progress in the logs panel

Captioning

TrainKit supports LLaVA-based captioning models.

Upscaling

TrainKit uses Spandrel for upscaling. See the full list of supported architectures.

Development

# Start in development mode
npm start

# Build for production
npm run make

Roadmap

  • NCNN support
  • Additional captioning model architectures
  • Image tagging

Contributing

Contributions are welcome! Please open an issue first to discuss what you would like to change.

License

TrainKit is licensed under the MIT license. See the LICENSE file for details.

About

A dataset preparation toolkit for AI image training

Resources

Stars

Watchers

Forks

Packages

No packages published