Compreendendo o LORA no Stable Diffusion

LORA (Low-Rank Adaptation) tornou-se uma ferramenta indispensável para quem deseja personalizar modelos de geração de imagem como o Stable Diffusion. Ele permite injetar características específicas em imagens geradas por IA sem a necessidade de retreinar o modelo base completo, o que seria uma tarefa computacionalmente intensiva. Este guia detalha o que é o LORA, como funciona e como utilizá-lo na prática, desde o treinamento até a aplicação.

Onde o LORA se Encaixa no Fluxo de Trabalho do Stable Diffusion?

A ppieline de geração de imagens do Stable Diffusion segue uma sequência de componentes que convertem texto em pixels visuais: um codificador de texto (Text Encoder) interpreta o prompt, uma rede UNet refina a representação latente e um VAE (Variational Autoencoder) decodifica a imagem final. O LORA age nas camadas de cross-attention dentro do UNet.

A grande inovação do LORA reside em sua abordagem: em vez de modificar os pesos originais do modelo base (representados por W), ele introduz e treina apenas um pequeno conjunto de pesos adicionais, denotados como ΔW. Para otimizar o uso de memória e a eficiência computacional, esse ΔW é decomposto em duas matrizes menores, A e B, de baixo rank, de modo que ΔW ≈ A × B. Isso permite ajustar o comportamento do modelo de forma eficaz, mantendo o modelo base intocado e com um custo computacional e de armazenamento muito menor.

Treinando Seu Próprio Modelo LORA

Criar seu próprio modelo LORA envolve algumas etapas essenciais. O exemplo a seguir utiliza o ambiente kohya-ss/sd-scripts, uma ferramenta popular para treinamento de LORAs.

1. Preparação do Conjunto de Dados

Para treinar um LORA eficaz, você precisará de um conjunto de imagens com anotações claras. Por exemplo, para ensinar o modelo a gerar "orelhas de gato":

# Crie uma estrutura de diretórios para as imagens de treinamento
mkdir -p ~/sd_treinamento/imagens/10_personagem_felina

# Adicione 15-20 imagens de 512x512 pixels de personagens com orelhas de gato.
# Nomeie-as sequencialmente, por exemplo, personagem_00001.png, personagem_00002.png, etc.

# Crie um arquivo de texto (.txt) correspondente para cada imagem, contendo prompts descritivos.
# Exemplo para 'personagem_00001.txt':
# 1girl, cat ears, solo, anime style, detailed eyes
# O número '10' no nome do diretório '10_personagem_felina' sugere que o token 'personagem_felina'
# será repetido 10 vezes durante o treinamento, o que pode ajudar na ênfase.

2. Configuração do Ambiente de Treinamento

Primeiro, clone o repositório kohya-ss/sd-scripts e configure um ambiente Python virtual.

git clone https://github.com/kohya-ss/sd-scripts.git
cd sd-scripts
python -m venv venv_lora
source venv_lora/bin/activate
pip install torch==2.0.1+cu118 torchvision xformers bitsandbytes

3. Executando o Script de Treinamento

Configure o accelerate e inicie o processo de treinamento com os parâmetros desejados. Os comentários no script explicam cada opção.

# Configuração do accelerate para otimização de recursos (responda NO,NO,NO,1,fp16 quando solicitado)
accelerate config

# Iniciar o treinamento do LORA
accelerate launch --num_cpu_threads_per_process 8 train_network.py \
  --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
  --train_data_dir="~/sd_treinamento/imagens" \
  --resolution=512,512 \
  --output_dir="~/sd_treinamento/saida_lora" \
  --output_name="orelhas_felinas_v1" \
  --network_module="networks.lora" \
  --network_dim=32 \          # Rank da matriz, define a capacidade do LORA. Valores maiores podem exigir mais VRAM.
  --network_alpha=16 \        # Coeficiente de suavização, ajuda a prevenir overfitting.
  --train_batch_size=2 \
  --max_train_epochs=10 \
  --save_every_n_epochs=2 \
  --optimizer_type="AdamW8bit" \
  --learning_rate=1e-4 \
  --lr_scheduler="cosine_with_restarts" \
  --lr_warmup_steps=200 \
  --mixed_precision="fp16" \
  --xformers \
  --cache_latents \
  --gradient_checkpointing

Após a conclusão do treinamento, você encontrará um arquivo orelhas_felinas_v1.safetensors no diretório de saída, geralmente com um tamanho de arquivo gerenciável (em torno de 30-40 MB).

Utilizando Seu LORA Gerado

Existem diversas maneiras de integrar seu LORA treinado ao Stable Diffusion para gerar imagens.

1. Através da Interface WebUI (Automatic1111/Fooocus)

1. Copie o arquivo orelhas_felinas_v1.safetensors para o diretório stable-diffusion-webui/models/Lora (ou o caminho equivalente em sua instalação).
2. Reinicie a interface WebUI.
3. Na interface, clique no botão para modelos extra (geralmente um ícone de "corações" ou "cartas"). Selecione a aba "LoRA" e clique no seu modelo orelhas_felinas_v1. Isso adicionará automaticamente o tag LORA ao seu prompt.

# Exemplo de prompt para adicionar à WebUI
positive: (obra-prima:1.2), uma garota, orelhas de gato, <lora:orelhas_felinas_v1:0.8>
negative: (qualidade baixa:2), borrado, orelhas de gato extras, <lora:orelhas_felinas_v1:-0.1>

O valor 0.8 ajusta a intensidade do LORA; um valor entre 0.6 e 0.9 é geralmente um bom ponto de partida. Usar um valor negativo (por exemplo, -0.1) no prompt negativo pode ajudar a suprimir características indesejadas que o LORA pode acidentalmente introduzir, como "orelhas extras".

2. Utilizando a API do Stable Diffusion WebUI

Para integração programática, a API do Stable Diffusion WebUI permite gerar imagens diretamente via código Python.

from webuiapi import WebUIApi

# Conecta ao servidor da API
sd_api_client = WebUIApi(host='127.0.0.1', port=7860)

# Define os parâmetros de geração
prompt_base = "uma garota, orelhas de gato, <lora:orelhas_felinas_v1:0.8>"
neg_prompt_base = "qualidade baixa, borrado"

# Gera a imagem
resultado_geracao = sd_api_client.txt2img(
    prompt=prompt_base,
    negative_prompt=neg_prompt_base,
    steps=28,
    cfg_scale=7,
    width=512,
    height=512,
    sampler_name="DPM++ 2M Karras"
)

# Salva a imagem resultante
imagem_gerada = resultado_geracao.image
imagem_gerada.save("personagem_com_orelhas.png")
print("Imagem com personagem felino gerada com sucesso!")

Diagnóstico e Solução de Problemas Comuns

Durante o uso e treinamento de LORAs, alguns problemas podem surgir. Aqui está uma tabela para auxiliar na resolução:

Técnicas Avançadas: Empilhando LORAs

É possível combinar múltiplos LORAs para criar resultados mais complexos. A ordem e a intensidade com que eles são aplicados no prompt podem influenciar significativamente a saída.

positive: uma garota, <lora:orelhas_felinas_v1:0.5>, <lora:estilo_cyberpunk:0.3>, <lora:iluminacao_suave:0.4>
negative: membros extras, (baixa qualidade:2)

Uma boa prática é colocar LORAs de personagens ou elementos específicos no início, seguidos por LORAs de estilo ou iluminação. Evite que a soma total das intensidades dos LORAs seja muito alta (geralmente acima de 1.2), pois isso pode levar a resultados caóticos.

Para exploração dinâmica, você pode variar as intensidades dos LORAs programaticamente:

import random
from webuiapi import WebUIApi

sd_api_client = WebUIApi(host='127.0.0.1', port=7860)

for i in range(5):
    peso_orelhas = round(random.uniform(0.4, 0.7), 2)
    peso_estilo = round(1.0 - peso_orelhas, 2) # Exemplo de pesos complementares
    
    prompt_combinado = f"uma garota, <lora:orelhas_felinas_v1:{peso_orelhas}>, <lora:estilo_cyberpunk:{peso_estilo}>"
    
    resultado_dinamico = sd_api_client.txt2img(
        prompt=prompt_combinado,
        negative_prompt="qualidade baixa",
        steps=28, cfg_scale=7, width=512, height=512, sampler_name="DPM++ 2M Karras"
    )
    resultado_dinamico.image.save(f"imagem_combinada_{i+1}.png")
    print(f"Imagem combinada {i+1} gerada com sucesso!")

Retreinando um LORA Existente (Fine-tuning de LORA)

Você pode usar um LORA já treinado como ponto de partida para um novo treinamento. Isso permite que você refine um estilo ou característica existente com seu próprio conjunto de dados, criando um "LORA do LORA". No kohya-ss, isso é feito adicionando um argumento:

# Adicione esta linha ao seu script de treinamento
--network_weights="caminho/para/seu/orelhas_felinas_v1.safetensors"
# O novo LORA (por exemplo, "orelhas_felinas_v2") herdará as características do v1
# e as adaptará ao seu novo conjunto de dados.

Otimização de Memória para Placas de Vídeo com Pouca VRAM

Mesmo com placas de vídeo com VRAM limitada (como 6GB), é possível treinar LORAs. As seguintes opções podem ser combinadas para reduzir o consumo de memória:

• --mixed_precision="fp16": Utiliza precisão de ponto flutuante de 16 bits, reduzindo a memória necessária para os tensores.
• --gradient_checkpointing: Troca computação por memória, armazenando menos valores intermediários e recomputando-os quando necessário.
• --cache_latents: Pré-calcula e armazena as representações latentes das imagens, economizando tempo de VAE durante o treinamento.
• --network_dim=16 e --network_alpha=8: Reduza o rank (network_dim) e o alpha para diminuir a complexidade do LORA e o uso de memória.
• --optimizer_type="AdaFactor": Otimizadores como AdaFactor geralmente consomem menos memória do que AdamW, embora possam ser mais lentos ou exigir ajuste de hiperparâmetros.

Com essas configurações, mesmo GPUs mais antigas podem ser capazes de realizar o treinamento, embora o tempo de processamento possa ser significativamente maior.

Checklist Essencial para o Treinamento de LORA

1. Dados de Treinamento: 15-20 imagens de 512x512 pixels, com o objeto/personagem principal bem centralizado e anotado com prompts claros.
2. Parâmetros de Treinamento: Experimente com network_dim=32, network_alpha=16, learning_rate=1e-4, max_train_epochs=10, train_batch_size=2.
3. Parâmetros de Geração: Para a imagem final, use Steps=28, Sampler=DPM++ 2M Karras, CFG=7, e uma intensidade de LORA entre 0.6 e 0.8.
4. Prompt Negativo: Inclua termos como orelhas extras, qualidade baixa, borrado, baixa resolução para melhorar a qualidade.
5. Dica Adicional: Incluir algumas imagens de "close-up frontal" ou "close-up do recurso desejado" no conjunto de treinamento pode melhorar a consistência.