Ostris的AI工具包

重要提示 - 请阅读

这是我的研究代码库。我在这里进行很多实验,可能会导致某些功能失效。如果遇到问题,请回退到之前的提交。这个代码库可以训练很多东西,很难跟上所有的更新。

支持我的工作

没有Glif团队的大力支持,我无法继续这个项目。如果你想支持我,请支持Glif。加入网站,加入我们的Discord,在Twitter上关注我们,一起来创造一些很酷的东西吧。

安装

要求:

Python >3.10
具有足够RAM的Nvidia GPU
Python虚拟环境
Git

Linux:

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
git submodule update --init --recursive
python3 -m venv venv
source venv/bin/activate
# 在Windows上使用 .\venv\Scripts\activate
# 先安装torch
pip3 install torch
pip3 install -r requirements.txt

Windows:

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
git submodule update --init --recursive
python -m venv venv
.\venv\Scripts\activate
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt

FLUX.1训练

正在开发中。我正在尽快更新文档并优化。如果有bug请提交issue。不知道如何使用不算bug。请耐心等待我继续开发。

要求

目前你需要一张至少24GB显存的GPU来训练FLUX.1。如果你正在用它来控制显示器,你可能需要在配置文件的model:下设置low_vram: true标志。这会在CPU上对模型进行量化,应该可以在连接显示器的情况下进行训练。用户已经在使用WSL的Windows上成功运行,但有报告称在Windows上原生运行时存在bug。目前我只在Linux上测试过。这仍然是非常实验性的,需要很多量化和技巧才能在24GB上运行。

模型许可

训练目前只适用于FLUX.1-dev。这意味着你训练的任何内容都将继承非商业许可。这也是一个受限模型,所以你需要在HF上接受许可才能使用它。否则,这将失败。以下是设置许可所需的步骤。

登录HF并在此接受模型访问权限black-forest-labs/FLUX.1-dev
在这个文件夹的根目录下创建一个名为.env的文件
从huggingface获取READ密钥并将其添加到.env文件中,如下所示:HF_TOKEN=your_key_here

训练

将位于config/examples/train_lora_flux_24gb.yaml的示例配置文件复制到config文件夹并重命名为whatever_you_want.yml
按照文件中的注释编辑文件
运行文件:python3 run.py config/whatever_you_want.yml

当你开始时,将创建一个与配置文件中的名称和训练文件夹相同的文件夹。它将包含所有检查点和图像。你可以随时使用ctrl+c停止训练,恢复时它将从最后一个检查点继续。

重要提示:如果你在保存时按下crtl+c,很可能会损坏该检查点。所以请等待保存完成。

需要帮助?

除非是代码中的bug,否则请不要提交bug报告。欢迎加入我的Discord寻求帮助。但是,请不要直接给我发私信询问一般问题或寻求支持。在discord上提问,我会在有空时回答。

在云端训练

即将推出。先发布基础版本,然后会提供一个可以让所有这些工作的notebook。

数据集准备

数据集通常需要是一个包含图像和相关文本文件的文件夹。目前,唯一支持的格式是jpg、jpeg和png。Webp目前存在问题。文本文件应与图像同名,但扩展名为.txt。例如image2.jpg和image2.txt。文本文件应只包含标题。你可以在标题文件中添加[trigger]一词,如果你在配置中有trigger_word,它将被自动替换。

图像永远不会被放大,但它们会被缩小并放入桶中进行批处理。你不需要裁剪/调整图像大小。加载器会自动调整它们的大小,并可以处理不同的宽高比。

以下内容已过时

它可能仍然可以这样工作,但我已经有一段时间没有测试过了。

批量图像生成

一个可以从配置文件或txt文件获取提示并生成到文件夹的图像生成器。我主要是为了SDXL测试而需要它,但为了通用的批量图像生成添加了一些改进。它完全基于配置文件运行,你可以在config/examples/generate.example.yaml中找到示例。更多信息在示例的注释中。

LoRA (lierla), LoCON (LyCORIS) 提取器

它基于LyCORIS工具中的提取器,但添加了一些生活质量功能和LoRA (lierla)支持。它可以在一次运行中进行多种类型的提取。它完全基于配置文件运行,你可以在config/examples/extract.example.yml中找到示例。只需将该文件复制到config文件夹,并将其重命名为whatever_you_want.yml。然后你可以编辑文件,并像这样调用它:

python3 run.py config/whatever_you_want.yml

你也可以使用配置文件的完整路径,如果你想把它放在其他地方。

python3 run.py "/home/user/whatever_you_want.yml"

示例配置文件本身中有更多关于它如何工作的说明。LoRA和LoCON都支持'fixed'、'threshold'、'ratio'、'quantile'提取。我稍后会更新这些是什么意思。大多数人使用fixed,这是传统的固定维度提取。

process是一个包含不同进程的数组。你可以添加几个并混合搭配。一个LoRA,一个LyCON等。

LoRA 重缩放

将<lora:my_lora:4.6>更改为<lora:my_lora:1.0>或你想要的任何值,效果相同。这是一个用于重缩放LoRA权重的工具。应该也适用于LoCON,但我还没有测试过。它完全基于配置文件运行,你可以在config/examples/mod_lora_scale.yml中找到示例。只需将该文件复制到config文件夹,并将其重命名为whatever_you_want.yml。然后你可以编辑文件,并像这样调用它:

python3 run.py config/whatever_you_want.yml

你也可以使用配置文件的完整路径,如果你想把它放在其他地方。

python3 run.py "/home/user/whatever_you_want.yml"

示例配置文件本身中有更多关于它如何工作的说明。这在制作所有LoRA时很有用,因为理想权重很少是1.0,但现在你可以修复这个问题。对于滑块,它们可能有从-2到2甚至-15到15的奇怪比例。这将允许你将它们调整到你想要的比例。

LoRA 滑块训练器

这是我最近在Civitai上训练大多数滑块的方法,你可以在我的Civitai个人资料中查看它们。它基于p1atdev/LECO和rohitgandikota/erasing的工作, 但经过大量修改,用于创建滑块而不是擦除概念。我对此还有很多计划,但目前已经非常实用了。它也非常易用。只需将config/examples/train_slider.example.yml中的示例配置文件复制到config文件夹,并将其重命名为whatever_you_want.yml。然后你可以编辑文件,并像这样调用它:

python3 run.py config/whatever_you_want.yml

该示例文件中有更多信息。你甚至可以不做任何修改直接运行示例来了解它是如何工作的。它将创建一个将所有动物变成狗(负)或猫(正)的滑块。只需这样运行:

python3 run.py config/examples/train_slider.example.yml

无需配置任何内容即可看到它的工作原理。这种方法不需要任何数据集。我很快会发布一个更好的教程。

扩展功能!!

现在你可以制作和分享自定义扩展。这些扩展在此框架内运行,并可使用所有内置工具。我可能会将其作为主要的开发方法,以避免不断向这个基础仓库添加越来越多的功能。我可能还会迁移很多现有功能,使一切都模块化。在 extensions 文件夹中有一个示例扩展,展示了如何制作模型合并扩展。所有代码都有详细注释,希望能帮助你入门。要制作扩展,只需复制该示例并替换所需的内容即可。

模型合并 - 示例扩展

位于 extensions 文件夹中。这是一个功能完整的模型合并器,可以合并任意多个模型。它不仅是制作扩展的好例子,也是一个非常有用的功能,因为大多数合并器一次只能处理一个模型,而这个可以处理任意多个模型。其中有一个示例配置文件,只需将其复制到 config 文件夹并重命名为 whatever_you_want.yml,然后像使用其他配置文件一样使用它。

正在开发的工具

VAE (变分自编码器) 训练器

这个功能可以工作,但还没有准备好供他人使用,因此没有示例配置。我还在改进中。完成后我会更新这部分。我正在添加许多用于我图像放大工作中的标准特性。包括评论家(判别器)、内容损失、风格损失等。如果你不知道,Stable Diffusion 的 VAE(即使是 MSE 和 SDXL 的)在较小的面部处理上很糟糕,这限制了 SD 的表现。我会解决这个问题。稍后我会发布更多相关内容和更好的示例,但这里有一个使用各种 VAE 的快速测试结果。只是简单地输入和输出。在较小的面部上效果比这里展示的要差得多。

待办事项

为滑块添加适当的正则化
添加 SDXL 支持(目前仅支持基础模型)
添加简单擦除功能
制作文本反转网络训练器(生成 TI 嵌入的网络)

更新日志

2023-08-05

内存使用和滑块功能大幅重构。滑块训练比以往更好,不再出现内存峰值。我还使滑块算法的 4 个部分在一个批次中运行,共享梯度累积。这使得训练更快更稳定。
更新了示例配置,使其更实用且符合当前方法。现在是一个细节滑块,展示了如何在没有主题的情况下训练。512x512 的 1.5 版滑块训练现在应该可以在 6GB GPU 上运行。稍后会测试验证。

2021-10-20

修复 Windows 支持的 bug
扩展功能!添加了制作和分享自定义训练、合并等扩展的功能。查看 extensions 文件夹中的示例。详情见上文。
模型合并功能,通过示例扩展提供。

2023-08-03

又一次大重构,使 SD 更加模块化。

制作了批量图像生成脚本

2023-08-01

重大更新。新增 LoRA 重缩放工具,详情见上文。添加了更好的元数据,使 Automatic1111 能识别基础模型。增加了一些实验功能和大量更新。目前这个版本仍不稳定,希望不会有太多破坏性变更。

很遗憾,我太懒了,没有写一个完整的更新日志来记录所有变化。

我为滑块训练添加了 SDXL 支持...但是...它无法正常工作。滑块训练依赖于模型理解无条件(负面提示)意味着你不希望输出中出现该概念。不知何故,SDXL 无法理解这一点, 这使得很难在模型内分离概念。我相信社区随着时间推移会找到解决方法, 但目前它无法正常工作。如果你们中有人在想"我们是否可以通过向模型添加 1 或 2 个文本编码器以及几个完全独立的扩散网络来修复它?"不。绝对不行。它只需要一点训练,而不是添加每一个实验性的新论文。KISS 原则。

2023-07-30

为滑块训练器添加了"锚点"。这允许你设置一个用作正则化器的提示。你可以设置网络乘数,以在高权重下强制保持一致性