A dataset preparation toolkit for AI image training
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.
- 🖼️ 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.
- Windows 10/11
- NVIDIA GPU (8GB+ VRAM recommended for captioning)
- Python 3.12
- Node.js 18+
Download the latest release from the Releases page.
# 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- Launch TrainKit — The backend server starts automatically
- Select a panel — Choose Caption, Upscale, or Rename from the sidebar
- Configure paths — Set your input folder, output folder, and model path
- Start processing — Click the action button and monitor progress in the logs panel
TrainKit supports LLaVA-based captioning models.
TrainKit uses Spandrel for upscaling. See the full list of supported architectures.
# Start in development mode
npm start
# Build for production
npm run make- NCNN support
- Additional captioning model architectures
- Image tagging
Contributions are welcome! Please open an issue first to discuss what you would like to change.
TrainKit is licensed under the MIT license. See the LICENSE file for details.