Kit de ferramentas de IA da Ostris: Difusão estável com o kit de ferramentas de treinamento de modelos FLUX.1

Introdução geral

O AI Toolkit by Ostris é um conjunto de ferramentas de IA de código aberto focado no suporte aos modelos Stable Diffusion e FLUX.1 para tarefas de treinamento e geração de imagens. Criado e mantido pelo desenvolvedor Ostris e hospedado no GitHub, o kit de ferramentas tem como objetivo fornecer uma plataforma flexível para que pesquisadores e desenvolvedores ajustem e experimentem modelos. Ele contém uma variedade de scripts de IA que oferecem suporte a recursos como extração de LoRA, geração de imagens em lote e treinamento específico de camadas. O projeto está atualmente em fase de desenvolvimento e alguns dos recursos podem não ser suficientemente estáveis, mas sua alta capacidade de personalização o torna adequado para usuários avançados no campo da aprendizagem profunda. O conjunto de ferramentas é compatível com os sistemas Linux e Windows, e é necessária uma GPU Nvidia com pelo menos 24 GB de memória de vídeo para executar o treinamento do modelo FLUX.1.

Lista de funções

• treinamento de modelosSuporte a difusão estável e ajuste fino do modelo FLUX.1 para treinamento dos modelos LoRA e LoKr.
• Geração de imagensGeração de imagens em lote com base em perfis ou prompts de texto.
• Extração e otimização de LoRAFornecer ferramentas de extração LoRA e LoCON para otimizar a extração de recursos do modelo.
• Treinamento específico da camadaCamadas específicas da rede neural podem ser selecionadas para treinamento e os pesos podem ser ajustados com flexibilidade.
• Suporte à interface do usuárioInterface de usuário do AI Toolkit e do Gradio para simplificar o gerenciamento de tarefas e as operações de treinamento de modelos.
• Processamento do conjunto de dadosAjuste automático da resolução da imagem e agrupe as imagens por bloco, suportando uma ampla variedade de formatos de imagem.
• Treinamento na nuvemSuporte para execução de tarefas de treinamento nas plataformas RunPod e Modal.

Usando a Ajuda

Processo de instalação

Instalação do sistema Linux

1. armazém de clonesExecute o seguinte comando no terminal para fazer o download do código:

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit

2. Atualização de submódulosVerifique se todas as bibliotecas dependentes estão completas:

git submodule update --init --recursive

3. Criação de um ambiente virtualUse o Python 3.10 ou posterior:

python3 -m venv venv
fonte venv/bin/activate

4. Instalação de dependênciasInstalar o PyTorch primeiro e depois as outras dependências:

pip3 install torch
pip3 install -r requirements.txt

Instalação do sistema Windows

1. armazém de clonesExecutar no prompt de comando:

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit

2. Atualização de submódulos:

git submodule update --init --recursive

3. Criação de um ambiente virtual:

python -m venv venv
. \venv\Scripts\activate

4. Instalação de dependênciasInstale a versão do PyTorch compatível com CUDA 12.4 e, em seguida, instale as outras dependências:

pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt

Instalação da interface UI

1. Instalação do Node.jsVerifique se o Node.js 18 ou posterior está instalado em seu sistema.
2. Criação da interface do usuárioDigite o diretório ui e instale as dependências:

cd ui
npm install
npm run build
npm run update_db

3. Executando a interface do usuário: Tela de inicialização:

npm run start

4. Acesso à interface do usuárioDigite em seu navegadorhttp://localhost:8675.

Funções principais

Treinamento do modelo FLUX.1

1. Preparação do ambienteVerifique se a GPU tem pelo menos 24 GB de memória de vídeo; se for usada para saída de vídeo, defina-a no arquivo de configuração.low_vram: truepara quantificar o modelo na CPU.
2. Configuração do FLUX.1-dev:

• Faça login no Hugging Face e visiteblack-forest-labs/FLUX.1-deve aceitar a licença.
• No diretório raiz do projeto, crie o arquivo.envadicione o arquivoHF_TOKEN=sua chave de leitura.

3. Configuração do FLUX.1-schnell:

• Edite o arquivo de configuração (por exemplotrain_lora_flux_schnell_24gb.yaml), adicione:

  modelo.
  name_or_path: "black-forest-labs/FLUX.1-schnell"
  assistant_lora_path: "ostris/FLUX.1-schnell-training-adapter"
  is_flux: true
  quantize: true
  sample.
  guidance_scale: 1
  sample_steps: 4
  

4. Preparação do conjunto de dadosCriar no diretório raizconjunto de dadosna pasta.jpge.jpegtalvez.pnge a imagem correspondente.txtDescreva o arquivo.
5. Editar arquivo de configuração: Reproduçãoconfig/examples/train_lora_flux_24gb.yamlatéconfiguraçãorenomeie-o paramy_config.ymlModificaçõescaminho_da_pastaé o caminho do conjunto de dados.
6. treinamento de corrida:: Implementação:

python run.py config/my_config.yml

Os resultados do treinamento são salvos na pasta de saída especificada e podem ser pausados com Ctrl+C e retomados a partir do ponto de verificação mais próximo.

Treinamento com o Gradio UI

1. Faça login no Hugging Face: Corridahuggingface-cli loginAs entradas têmescreverA chave para o privilégio.
2. Iniciar a interface do usuário:: Implementação:

python flux_train_ui.py

3. Operação da interface do usuárioFaça upload de imagens, preencha as descrições, defina os parâmetros e clique em treinamento na interface e, em seguida, você poderá publicar o modelo LoRA após a conclusão.

Treinamento na nuvem (RunPod)

1. Criação de uma instância do RunPodUso de modelosrunpod/pytorch:2.2.0-py3.10-cuda12.1.1-devel-ubuntu22.04escolha A40 (48 GB de memória de vídeo).
2. Conjunto de ferramentas de instalaçãoConecte-se ao Jupyter Notebook e execute o comando de instalação do Linux no terminal.
3. Upload de conjuntos de dadosCriar no diretório raizconjunto de dadosarraste as imagens e os arquivos de descrição.
4. Configurar e executarModificar o arquivo de configuração docaminho_da_pastaImplementaçãopython run.py config/my_config.yml.

Preparação do conjunto de dados

• Requisitos de formatação: Suporte.jpge.jpege.pngo arquivo de descrição é.txtOs nomes dos arquivos precisam ser consistentes (por exemploimage1.jpgcontrapartesimage1.txt).
• Descrição: .txtpara escrever uma descrição no[acionador].Os espaços reservados, definidos pelopalavra-gatilhoSubstituição.
• ajuste automáticoA ferramenta encolhe e agrupa automaticamente as imagens de acordo com a resolução configurada; não há suporte para zoom.

Treinamento específico da camada

1. Editar arquivo de configuração: EmredeParcialmente adicionado:

rede.
tipo: "lora"
linear: 128
linear_alpha: 128
network_kwargs.
only_if_contains.
- "transformador.single_transformer_blocks.7.proj_out"
- "transformer.single_transformer_blocks.20.proj_out"

2. treinamento de corridaInicializa com um arquivo de configuração modificado, treinando apenas a camada especificada.

advertência

• Interrupção do treinamentoEvite pressionar Ctrl+C ao salvar pontos de controle para não corromper o arquivo.
• Segurança da interface do usuárioA interface do usuário é atualmente testada apenas no Linux, que é menos seguro e não é recomendado para exposição à rede pública.
• Obtendo ajudaPode participar da comunidade Discord do Ostris para fazer perguntas e evitar mensagens privadas diretas com os desenvolvedores.