AI Toolkit de Ostris: Difusión estable con FLUX.1 Kit de herramientas de formación de modelos

Introducción general

AI Toolkit de Ostris es un conjunto de herramientas de IA de código abierto centrado en los modelos Stable Diffusion y FLUX.1 para tareas de entrenamiento y generación de imágenes. Creado y mantenido por el desarrollador Ostris y alojado en GitHub, el kit de herramientas tiene como objetivo proporcionar una plataforma flexible para que investigadores y desarrolladores ajusten y experimenten con modelos. Contiene una serie de scripts de IA que admiten funciones como la extracción de LoRA, la generación de imágenes por lotes y el entrenamiento específico de capas. El proyecto se encuentra actualmente en fase de desarrollo y es posible que algunas de las funciones no sean lo suficientemente estables, pero su gran capacidad de personalización lo hace adecuado para usuarios avanzados en el campo del aprendizaje profundo. El conjunto de herramientas es compatible con sistemas Linux y Windows, y se requiere una GPU Nvidia con al menos 24 GB de memoria de vídeo para ejecutar el entrenamiento del modelo FLUX.1.

Lista de funciones

• formación de modelos:: Admite el ajuste fino de los modelos Stable Diffusion y FLUX.1 para el entrenamiento de los modelos LoRA y LoKr.
• Generación de imágenesGeneración de imágenes por lotes a partir de perfiles o mensajes de texto.
• Extracción y optimización de LoRA: Proporcionar herramientas de extracción LoRA y LoCON para optimizar la extracción de características del modelo.
• Entrenamiento específico por capasCapas de la red neuronal: se pueden seleccionar capas específicas de la red neuronal para el entrenamiento y los pesos se pueden ajustar con flexibilidad.
• Interfaz de usuarioInterfaz de usuario: Proporciona la interfaz de usuario de AI Toolkit y Gradio para simplificar la gestión de tareas y las operaciones de formación de modelos.
• Tratamiento del conjunto de datos: Ajusta automáticamente la resolución de las imágenes y las agrupa por cubos, admitiendo una amplia gama de formatos de imagen.
• Formación en la nube: Soporte para ejecutar tareas de entrenamiento en las plataformas RunPod y Modal.

Utilizar la ayuda

Proceso de instalación

Instalación del sistema Linux

1. almacén de clonesEjecute el siguiente comando en el terminal para descargar el código:

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

2. Actualización de submódulos: Asegúrese de que todas las bibliotecas dependientes están completas:

git submodule update --init --recursive

3. Creación de un entorno virtualUtilice Python 3.10 o posterior:

python3 -m venv venv
source venv/bin/activate

4. Instalación de dependenciasInstala PyTorch primero, luego las otras dependencias:

pip3 install torch
pip3 install -r requirements.txt

Instalación del sistema Windows

1. almacén de clonesEjecutar en el símbolo del sistema:

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

2. Actualización de submódulos:

git submodule update --init --recursive

3. Creación de un entorno virtual:

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

4. Instalación de dependenciasInstale la versión de PyTorch compatible con CUDA 12.4 y, a continuación, instale las demás dependencias:

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

Instalación de la interfaz de usuario

1. Instalación de Node.jsAsegúrese de que Node.js 18 o posterior está instalado en su sistema.
2. Construir la interfaz de usuario: Entra en el directorio ui e instala las dependencias:

cd ui
npm install
npm run build
npm run update_db

3. Ejecutar la interfaz de usuarioPantalla de inicio:

npm run start

4. Acceso a la IUEscriba en su navegadorhttp://localhost:8675.

Funciones principales

Entrenamiento del modelo FLUX.1

1. Preparar el entornoAsegúrese de que la GPU tiene al menos 24 GB de memoria de vídeo, si se utiliza para la salida de pantalla, configúrelo en el archivo de configuración.low_vram: truepara cuantificar el modelo en la CPU.
2. Configuración de FLUX.1-dev:

• Conéctate a Hugging Face y visitablack-forest-labs/FLUX.1-devy acepta la licencia.
• En el directorio raíz del proyecto, cree el archivo.envañada el campoHF_TOKEN=你的读取密钥.

3. Configurar FLUX.1-schnell:

• Edite el archivo de configuración (por ejemplotrain_lora_flux_schnell_24gb.yaml), añade:

  model:
  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. Preparación del conjunto de datosCrear en el directorio raízdataseten la carpeta.jpgy.jpegtal vez.pngimagen y la correspondiente.txtDescribe el expediente.
5. Editar archivo de configuración: Reproducciónconfig/examples/train_lora_flux_24gb.yamlhastaconfigy cámbiele el nombre amy_config.ymlModificacionesfolder_pathes la ruta del conjunto de datos.
6. entrenamiento de running:: Implementación:

python run.py config/my_config.yml

Los resultados del entrenamiento se guardan en la carpeta de salida especificada y pueden pausarse con Ctrl+C y reanudarse desde el punto de control más reciente.

Formación con Gradio UI

1. Conectarse a Hugging Face: Correrhuggingface-cli loginLas entradas tienenwriteLa llave del privilegio.
2. Iniciar interfaz de usuario:: Implementación:

python flux_train_ui.py

3. Operación UI: Cargue imágenes, rellene descripciones, establezca parámetros y haga clic en formación en la interfaz, y después podrá publicar el modelo LoRA una vez finalizado.

Entrenamiento en la nube (RunPod)

1. Creación de una instancia de RunPodUso de plantillasrunpod/pytorch:2.2.0-py3.10-cuda12.1.1-devel-ubuntu22.04Elige A40 (48 GB de memoria de vídeo).
2. Juego de herramientas de instalaciónConéctese a Jupyter Notebook y ejecute el comando de instalación de Linux en el terminal.
3. Cargar conjuntos de datosCrear en el directorio raízdatasetcarpeta, arrastre las imágenes y los archivos de descripción.
4. Configurar y ejecutarModifica el archivo de configuración delfolder_pathAplicaciónpython run.py config/my_config.yml.

Preparación del conjunto de datos

• Requisitos de formato: Apoyo.jpgy.jpegy.pngel archivo de descripción es.txtLos nombres de los archivos deben ser coherentes (p. ej.image1.jpgcontrapartidasimage1.txt).
• Descripción: .txtEscriba la descripción en el[trigger]tal y como se define en el archivo de configuracióntrigger_wordSustitución.
• ajuste automáticoLa herramienta reduce y agrupa automáticamente las imágenes en función de la resolución configurada.

Entrenamiento específico por capas

1. Editar archivo de configuraciónEnnetworkParcialmente añadido:

network:
type: "lora"
linear: 128
linear_alpha: 128
network_kwargs:
only_if_contains:
- "transformer.single_transformer_blocks.7.proj_out"
- "transformer.single_transformer_blocks.20.proj_out"

2. entrenamiento de running: Se inicia con un archivo de configuración modificado, entrenando sólo la capa especificada.

advertencia

• Interrupción de la formaciónEvita pulsar Ctrl+C mientras guardas los puntos de control para evitar corromper el archivo.
• Seguridad de la interfaz de usuarioActualmente, la interfaz de usuario sólo se ha probado en Linux, que es menos seguro y no se recomienda exponerlo a la red pública.
• Obtener ayudaPuede unirse a la comunidad de Discord de Ostris para hacer preguntas y evitar los mensajes privados directos a los desarrolladores.