Flux Fill LoRA Train

Этот репозиторий содержит продвинутый скрипт для обучения модели FluxFill, основанный на архитектуре T2ITrainer.

Проект адаптирован для решения задачи заполнения изображений (image inpainting) с использованием передовых техник глубокого обучения.

Также наконец-то стабильное обучение задач Outpaint (дорисовки изображений)

Ключевые особенности

• 🚀 Эффективное обучение с применением LoRA (Low-Rank Adaptation)
• 💾 Автоматическое кэширование эмбеддингов и латентных представлений
• 🌈 Поддержка Flux Transformer для генерации изображений высокого разрешения
• ⚙️ Гибкая конфигурация через параметры в начале скрипта
• 📊 Интеграция с TensorBoard для мониторинга обучения
• 🐇 Оптимизация памяти через gradient checkpointing

Быстрый старт

1. Установка зависимостей:

PyTorch устанавливается отдельно для вашей версии CUDA.

SageAttention устанавливается также отдельно.

Мог что пропустить в requirements.txt - устанавливается через pip install.

Настоятельно рекомендуется установить OpenCV с CUDA для скорости.

pip install -r requirements.txt

2. Подготовка датасета:

• Разместите изображения в структуре:

  dataset/train/
      ├── image1_F.jpg    # Фактическое изображение
      ├── image1_G.jpg    # Эталонное изображение
      ├── image1_M.png    # Маска
      └── image1.txt      # Описание (опционально)
  

3. Настройка параметров: Отредактируйте параметры в верхней части скрипта:

# Общие настройки
PRETRAINED_MODEL_NAME = "black-forest-labs/FLUX.1-Fill-dev"
SEED = 41
TRAIN_DATA_DIR = "dataset/train"
# ... другие параметры ...

4. Запуск обучения:

accelerate launch train_flux_fill.py

Конфигурация обучения

Основные настраиваемые параметры:

Параметр	Описание	По умолчанию
LEARNING_RATE	Скорость обучения	1e-4
TRAIN_BATCH_SIZE	Размер батча	2
NUM_TRAIN_EPOCHS	Количество эпох	5
LORA_RANK	Размерность LoRA-матриц	64
MIXED_PRECISION	Смешанная точность (fp16/bf16)	"bf16"
OUTPUT_DIR	Директория для сохранения чекпоинтов	"train"

Системные требования

• Python 3.12+
• PyTorch 2.0+
• NVIDIA GPU с 12+ GB VRAM
• CUDA 12.*

Мониторинг обучения

Для просмотра метрик в реальном времени:

tensorboard --logdir=logs

Особенности реализации

1. Эффективная работа с памятью:

   • Gradient checkpointing
   • Пакетная обработка схожих изображений
   • Оптимизированные латентные представления

2. Продвинутые техники обучения:

   • LoRA-адаптация
   • Несколько схем взвешивания потерь
   • Смешанная точность
   • Регуляризация временных шагов

3. Гибкий пайплайн:

   • Автоматическая обработка датасета
   • Поддержка возобновления обучения
   • Валидация во время обучения

Лицензия

Проект распространяется под лицензией Apache 2.0. Полный текст лицензии доступен в файле LICENSE.