|
| 1 | +# CLAUDE.md |
| 2 | + |
| 3 | +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. Failure IS an option. |
| 4 | + |
| 5 | +IMPORTANT: ALWAYS RUN RUFF AND ESLINT ON STRICT MODE BEFORE FINALISING ANY TASK. NO FILES ARE EXEMPT FROM THIS RULE. |
| 6 | + |
| 7 | +## Project Overview |
| 8 | + |
| 9 | +OneTrainer is a mature diffusion model training framework supporting 21+ model architectures, 4 training methods (Fine-Tune, LoRA, Embedding, Fine-Tune VAE), and 40+ optimizers. The project is migrating from a customtkinter desktop GUI to an Electron + React web-based frontend (**OneTrainerWeb**), with a FastAPI bridge layer connecting to the existing Python backend. |
| 10 | + |
| 11 | +## Critical Constraints |
| 12 | + |
| 13 | +- **DO NOT modify existing backend code** in `modules/`, `scripts/`, or `training_presets/` unless absolutely necessary. The existing backend must remain 100% backwards compatible. |
| 14 | +- All new code lives under `web/gui/` (frontend) and `web/backend/` (FastAPI bridge). |
| 15 | +- The existing `TrainCallbacks`/`TrainCommands` interface is the only integration point between UI and training backend. |
| 16 | + |
| 17 | +## Tech Stack |
| 18 | + |
| 19 | +- **Frontend:** React 19.2.4, Vite 7.3.1, Tailwind CSS 4.2.0, Electron 40.6.0, TypeScript (strict mode) |
| 20 | +- **Backend bridge:** Python 3.12, FastAPI, Uvicorn |
| 21 | +- **Existing backend:** Python 3.10-3.12, PyTorch 2.8, built on Hugging Face diffusers |
| 22 | + |
| 23 | +## Build & Development Commands |
| 24 | + |
| 25 | +### Python (existing backend) |
| 26 | +```bash |
| 27 | +# Install (creates venv automatically) |
| 28 | +./install.bat # Windows |
| 29 | +./install.sh # Linux/Mac |
| 30 | + |
| 31 | +# Run existing GUI |
| 32 | +./start-ui.bat # Windows |
| 33 | +./start-ui.sh # Linux/Mac |
| 34 | + |
| 35 | +# CLI headless training |
| 36 | +python scripts/train.py --config-path <path-to-config.json> |
| 37 | + |
| 38 | +# Install dev dependencies (run OUTSIDE venv) |
| 39 | +pip install -r requirements-dev.txt |
| 40 | +pre-commit install |
| 41 | +``` |
| 42 | + |
| 43 | +### Linting (Python) |
| 44 | +```bash |
| 45 | +# Ruff linter with auto-fix (configured in pyproject.toml) |
| 46 | +ruff check --fix . |
| 47 | + |
| 48 | +# Pre-commit runs Ruff automatically on commit |
| 49 | +pre-commit run --all-files |
| 50 | +``` |
| 51 | + |
| 52 | +Ruff config: line-length 120, see `pyproject.toml` for full rule set. Import ordering uses custom sections (first-party `modules`, then `mgds`, `torch`, `hf`). |
| 53 | + |
| 54 | +### Frontend (web/gui/) |
| 55 | +```bash |
| 56 | +cd web/gui |
| 57 | +npm install |
| 58 | +npm run dev # Vite dev server |
| 59 | +npm run build # Production build |
| 60 | +npm run typecheck # TypeScript strict checking |
| 61 | +npm run lint # ESLint |
| 62 | +npm test # Vitest unit tests |
| 63 | +npm run test:e2e # Playwright E2E tests |
| 64 | +``` |
| 65 | + |
| 66 | +### Backend bridge (web/backend/) |
| 67 | +```bash |
| 68 | +cd web/backend |
| 69 | +pip install fastapi uvicorn |
| 70 | +uvicorn main:app --reload --port 8000 |
| 71 | +pytest # Backend tests |
| 72 | +``` |
| 73 | + |
| 74 | +## Architecture |
| 75 | + |
| 76 | +### Existing Backend (modules/) |
| 77 | + |
| 78 | +The codebase follows a modular architecture where each module type is interchangeable: |
| 79 | + |
| 80 | +- **`modules/trainer/`** — Training loop (`GenericTrainer`, `MultiTrainer`, `CloudTrainer`). Communicates with UI exclusively through `TrainCallbacks`/`TrainCommands`. |
| 81 | +- **`modules/util/config/`** — JSON-native config system. `TrainConfig` has 200+ parameters with a 10-generation version migration system. Config classes extend `BaseConfig` for serialization. |
| 82 | +- **`modules/util/enum/`** — 27 enum definitions (ModelType, TrainingMethod, Optimizer, DataType, etc.) that must be mirrored as TypeScript types. |
| 83 | +- **`modules/util/create.py`** — Central factory (86KB) that instantiates all modules based on config. |
| 84 | +- **`modules/util/callbacks/TrainCallbacks.py`** — Trainer-to-UI interface (progress, status, samples). |
| 85 | +- **`modules/util/commands/TrainCommands.py`** — UI-to-Trainer interface (stop, sample, backup, save). |
| 86 | +- **`modules/ui/`** — Legacy customtkinter GUI (27 files). Reference for feature parity but not modified. |
| 87 | +- **`modules/model/`, `modelLoader/`, `modelSampler/`, `modelSaver/`, `modelSetup/`** — Per-architecture implementations. |
| 88 | +- **`modules/dataLoader/`** — MGDS-based data loading per architecture. |
| 89 | + |
| 90 | +### New Frontend Architecture (web/) |
| 91 | + |
| 92 | +``` |
| 93 | +web/ |
| 94 | +├── gui/ # Electron + React + Vite + Tailwind |
| 95 | +│ └── src/ |
| 96 | +│ ├── main/ # Electron main process (spawn FastAPI, window mgmt) |
| 97 | +│ ├── renderer/ # React SPA |
| 98 | +│ │ ├── components/ # Shared UI components |
| 99 | +│ │ ├── pages/ # Tab pages (General, Model, Training, etc.) |
| 100 | +│ │ ├── hooks/ # Custom React hooks |
| 101 | +│ │ ├── store/ # State management |
| 102 | +│ │ ├── api/ # REST/WebSocket client layer |
| 103 | +│ │ └── types/ # TypeScript types (generated from Python enums) |
| 104 | +│ └── shared/ # Types shared between main/renderer |
| 105 | +├── backend/ # FastAPI bridge server |
| 106 | +│ ├── routers/ # REST endpoints (config, training, sampling, tools, system, presets) |
| 107 | +│ ├── services/ # Business logic wrappers |
| 108 | +│ ├── ws/ # WebSocket handlers (training progress, tensorboard, system metrics) |
| 109 | +│ └── models/ # Pydantic request/response models |
| 110 | +└── docs/ |
| 111 | + └── breakdown.md # Comprehensive change documentation (required) |
| 112 | +``` |
| 113 | + |
| 114 | +### Process Lifecycle |
| 115 | +1. Electron spawns FastAPI/Uvicorn as child process |
| 116 | +2. Waits for health check endpoint |
| 117 | +3. Loads React renderer pointing to FastAPI backend |
| 118 | +4. React communicates via REST (config CRUD) and WebSocket (real-time progress/metrics) |
| 119 | +5. FastAPI imports existing `modules/` directly — no backend code changes |
| 120 | + |
| 121 | +### Key Communication Pattern |
| 122 | +``` |
| 123 | +React UI → REST/WebSocket → FastAPI Bridge → TrainCallbacks/TrainCommands → GenericTrainer |
| 124 | +``` |
| 125 | + |
| 126 | +The FastAPI bridge imports `modules.util.config.TrainConfig`, `modules.util.create`, `modules.util.callbacks.TrainCallbacks`, and `modules.util.commands.TrainCommands` directly. |
| 127 | + |
| 128 | +## Dynamic UI Behavior |
| 129 | + |
| 130 | +The Model and Training tabs completely rebuild based on `model_type` (21 types) and `training_method` (4 methods). The LoRA and Embedding tabs appear/disappear based on training method. The React frontend should use schema-driven rendering (declarative config per model type) rather than large if/elif chains. |
| 131 | + |
| 132 | +## Configuration System |
| 133 | + |
| 134 | +- All configs serialize to/from JSON via `BaseConfig.to_dict()`/`from_dict()` |
| 135 | +- `TrainConfig` version is currently 10, with migrations from version 0-9 |
| 136 | +- Config presets live in `training_presets/*.json` (built-in prefixed with `#`) |
| 137 | +- Concepts in `training_concepts/concepts.json`, samples in `training_samples/samples.json` |
| 138 | +- Secrets in `secrets.json` (never commit) |
| 139 | +- TypeScript types must be generated from Python config/enum classes at build time |
| 140 | + |
| 141 | +## Testing |
| 142 | + |
| 143 | +The existing project has zero automated tests. The web frontend must build a comprehensive test suite: |
| 144 | +- **Config round-trip tests** (highest priority): load preset → serialize → deserialize → diff must be zero |
| 145 | +- **Parameter parity tests**: automated checks that React forms match `TrainConfig` fields |
| 146 | +- **Dynamic UI tests**: verify correct parameter sets per model type |
| 147 | +- Frontend: Vitest + React Testing Library (unit), Playwright (E2E) |
| 148 | +- Backend: pytest + httpx (API), websockets (WebSocket) |
| 149 | + |
| 150 | +## Code Style |
| 151 | + |
| 152 | +- Python: 4-space indent, 120 char line length, double quotes, Ruff linting |
| 153 | +- TypeScript: strict mode, all linting rules enabled |
| 154 | +- File encoding: UTF-8, LF line endings (CRLF for .bat files only) |
| 155 | + |
| 156 | +## Design Reference |
| 157 | + |
| 158 | +See `OneTrainerWeb.md` for complete brand spec (color palette, typography, elevation, motion, accessibility requirements). Key tokens: `--cobalt-600: #2563EB`, `--azure-500: #38BDF8`, dark surface `--surface-dark: #0B1120`, light surface `--surface-light: #F0F5FF`. |
| 159 | + |
| 160 | +## Reference Documents |
| 161 | + |
| 162 | +- `FeasibilityStudy.md` — Full technical analysis, UI migration map, risk assessment, implementation roadmap |
| 163 | +- `OneTrainerWeb.md` — Project requirements, design system, brand guidelines |
| 164 | +- `docs/ProjectStructure.md` — Module architecture explanation |
| 165 | +- `docs/Overview.md` — General OneTrainer overview |
0 commit comments