Stable LoRA

Stable LoRA

AI Исходный код

1.0.0

Скачать

Стабильная Лора (WIP)

Обучит Лору, используя официальную реализацию Microsoft со стабильными диффузионными моделями. Это самый эффективный и самый простой способ обучения лораса без дополнительной сложности, а также общаться между библиотеками и реализациями.

? Это репо все еще в активном развитии. Ожидайте ошибки, пока флаг WIP не будет удален. ?

Начиная

Клонировать этот репозиторий в свою учебную папку:

git clone https://github.com/ExponentialML/Stable-LoRA

Установите требования из этого хранилища:

pip install -r requirements.txt

Пример

Обучение проводится с помощью библиотеки диффузоров.

См. Пример avest_train_script.sh для примера сценария Dreambooth.

Запуск его так же просто, как и делать:

 accelerate launch train_lora . py 
  - - pretrained_model_name_or_path = "runwayml/stable-diffusion-v1-5" 
  - - instance_data_dir = "instance_images" 
  - - lora_rank = 64 
  - - output_dir = "./output" 
  - - instance_prompt = "a shld dog" 
  - - resolution = 512 
  - - mixed_precision = "fp16" 
  - - train_text_encoder 
  - - learning_rate = 1e-4 
  - - resize = True 
  - - save_steps = 200 
  - - preview_steps = 100 
  - - lr_scheduler = "constant_with_warmup" 
  - - lr_warmup_steps = 50 
  - - max_train_steps = 5000 
  - - save_preview 
  - - preview_prompt = "a shld dog as a super hero"

После обучения результаты будут сохранены в вашей выходной папке.

По умолчанию, вариант Automatic1111 Webui сохраняется вместе с ними (расширение WebUI находится в разработке).

Разработчики

Просто добавьте его в свою модель перед тренировочной петлей.

 # Import function to add the LoRA, as well as target modules.
from stable_lora . lora import add_lora_to , UNET_REPLACE , TEXT_ENCODER_REPLACE

from diffusers import StableDiffusionPipeline

# Load a Stable Diffusion Model
pipeline = StableDiffusionPipeline . from_pretrained ( pretrained_model_name_or_path )

# Freeze the models. Remember, only the LoRA weights get trained, not the model itself.
pipeline . unet . requires_grad_ ( False )
pipeline . text_encoder . requires_grad_ ( False )

# Add LoRA to the UNET
add_lora_to (
    pipeline . unet , 
    target_module = UNET_REPLACE , 
    search_class = [ torch . nn . Linear , torch . nn . Conv2d ], 
    r = 32
)

# Add LoRA to the Text Encoder
add_lora_to ( pipeline . text_encoder , target_module = TEXT_ENCODER_REPLACE , r = 32 )

# Your optimizers and training code...

После добавления LORA в вашу модель вы можете легко добавить параметры параметров оптимизатора, если это необходимо.

 # Example

unet_params = []
text_encoder_params = []

for n , p in unet . named_parameters ():
    if 'lora' in n :
        unet_params . append ({
            "params" : p , 
            "lr" : args . learning_rate
        })

for n , p in text_encoder . named_parameters ():
    if 'lora' in n :
        text_encoder_params . append ({
            "params" : p , 
            "lr" : args . learning_rate_text
        })

optimizer_params = unet_params + text_encoder_params

Сохранение может быть выполнено с помощью Safetensors или традиционного способа, используя файлы .pt :

 from stable_lora . lora import save_lora

save_lora ( unet , path = 'save_file_path.safetensors' )
save_lora ( text_encoder = text_encoder , use_safetensors = False , path = 'save_file_path.pt' )

Кроме того, вы можете установить режим для оценки или режима поезда.

 from stable_lora . lora import set_mode

# Train mode
set_mode ( unet , is_train = True )

# Evaluation mode
set_mode ( unet , is_train = False )

Советы

Это было протестировано на моделях на основе стабильной диффузии 1.5.

Рекомендуется попробовать 5e-6 , 3e-5 и 1e-4 .

Ранг по умолчанию составляет 32 , но может быть установлен, передавая его через параметр r Использование более низких рангов будет потреблять меньше памяти, причем высокие ряды потребляют больше. Это также способствует сохранению файлов.

Эти факторы будут зависеть исключительно от вашей задачи, но являются хорошей отправной точкой.

Вы также можете включить смещения для тренировки.

В зависимости от ваших данных и скорости обучения вы сможете выжать дополнительную производительность из обучения.

 # Example | Applicable to Unet or Text Encoder

lora_bias = 'lora_only'
add_lora_to ( pipeline . text_encoder , target_module = TEXT_ENCODER_REPLACE , r = 32 , lora_bias = lora_bias )

# Must be set here as well when saving.
save_lora ( text_encoder = text_encoder , use_safetensors = True , path = 'save_file_path.pt' , lora_bias = lora_bias )