OneTrainer

OneTrainer 项目介绍

OneTrainer 是一个专为稳定扩散（Stable Diffusion）训练需求提供的一站式解决方案。无论是模型支持、训练方法还是数据处理工具，OneTrainer 都具备多种功能，帮助用户高效开展和管理模型训练过程。以下是详细介绍。

功能特点

• 支持的模型：OneTrainer 支持包括 FLUX.1、Stable Diffusion 各个版本（1.5、2.0、2.1、3.0、3.5）、SDXL、Würstchen-v2、Stable Cascade、PixArt-Alpha、PixArt-Sigma 和修复（inpainting）模型在内的多种模型。
• 模型格式：支持 diffusers 和 ckpt 模型格式。
• 训练方法：包括全量微调训练、LoRA 训练及嵌入式训练。
• 遮罩训练：可以针对样本的特定部分进行训练，提高专注度和训练效果。
• 自动备份：在训练过程中，定期备份训练进度，确保可无缝继续训练。
• 图像增强：通过随机旋转、亮度、对比度和饱和度等转换，快速创造多样化的数据集。
• 可视化跟踪：内置简单的 tensorboard，方便跟踪训练进度。
• 多提示语训练：针对每个图像样本使用多个不同的提示语进行训练。
• 噪声调度重订：基于研究论文的噪声调度法进行重新校准。
• EMA 模型：支持训练自己的 EMA 模型，可选将 EMA 权重保存在 CPU 内存中以减少显存使用。
• 长宽比分割桶：自动在多种长宽设置下进行训练。
• 多分辨率训练：支持同时对多个分辨率进行训练。
• 数据工具：使用 BLIP、BLIP2 和 WD-1.4 自动为数据集添加标题，或使用 ClipSeg 或 Rembg 创建遮罩。
• 模型工具：通过简单的界面在不同模型格式间进行转换。
• 采样界面：在训练过程中实现模型采样而无需切换应用程序。
• AlignProp：一种文本到图像扩散模型的强化学习方法。

安装步骤

OneTrainer 的安装要求 Python 版本在 3.10 以上且 3.13 以下。推荐的安装步骤有自动安装和手动安装。

自动安装

1. 克隆代码库：git clone https://github.com/Nerogar/OneTrainer.git
2. 运行安装脚本：
   • Windows 用户执行 install.bat
   • Unix 系统用户执行 install.sh

手动安装

1. 克隆代码库：git clone https://github.com/Nerogar/OneTrainer.git
2. 进入目录：cd OneTrainer
3. 设置虚拟环境：python -m venv venv
4. 激活虚拟环境：
   • Windows 用户执行 venv\scripts\activate
   • Unix 系统用户执行 source venv/bin/activate
5. 安装依赖：pip install -r requirements.txt

对于部分 Linux 系统，例如 Ubuntu，需要安装 libGL：

sudo apt-get update
sudo apt-get install libgl1

更新方法

自动更新

执行 update.bat 或 update.sh

手动更新

1. 拉取更新：git pull
2. 激活虚拟环境：venv\scripts\activate
3. 重新安装所有依赖：pip install -r requirements.txt --force-reinstall

使用说明

启动用户界面，运行 start-ui.bat。用户可以选择通过命令行或用户界面（UI）运行 OneTrainer。所有的命令必须在安装时创建的活动虚拟环境中运行。

主要脚本包括：

• train.py：核心训练脚本
• train_ui.py：训练的 UI
• caption_ui.py：手动或者自动添加标题和创建遮罩的 UI
• convert_model_ui.py：模型格式转换的 UI
• convert_model.py：模型格式转换工具
• sample.py：模型采样工具
• create_train_files.py：创建训练所需文件的工具
• generate_captions.py：自动为数据集生成标题的工具
• generate_masks.py：自动为数据集生成遮罩的工具
• calculate_loss.py：计算每个样本训练损失的工具

贡献说明

欢迎所有形式的贡献。用户可以通过提交问题、参与讨论或者提出功能改进的 pull request 来参与项目。详细信息和项目结构可以参考相应的文档。

此外，贡献者需要安装开发者依赖项并启用 Git 提交钩子。这将有助于确保所有代码风格一致，从而提高代码审核的效率。

相关项目

• MGDS：一个基于节点图思想的自定义 Pytorch 数据集实现。
• StableTuner：另一个 Stable Diffusion 训练应用，OneTrainer 的灵感来源。
• Visions of Chaos：包含 OneTrainer 的一套机器学习工具。