Custom Diffusion

Custom Diffusion é uma técnica de treinamento para personalizar modelos de geração de imagens. Assim como Textual Inversion, DreamBooth e LoRA, o Custom Diffusion requer apenas algumas (~4-5) imagens de exemplo. Essa técnica funciona treinando apenas os pesos nas camadas de cross-attention e utiliza uma palavra especial para representar o conceito recém-aprendido. O Custom Diffusion é único porque também pode aprender múltiplos conceitos ao mesmo tempo.

Se você estiver treinando em uma GPU com vRAM limitado, você deve tentar habilitar o xFormers com --enable_xformers_memory_efficient_attention para um treinamento mais rápido com requisitos de vRAM menores (16GB). Para economizar ainda mais memória, adicione --set_grads_to_none no argumento de treinamento para definir os gradientes como None em vez de zero (essa opção pode causar alguns problemas, então, se você tiver algum, tente remover esse parâmetro).

Este guia explorará o script train_custom_diffusion.py para ajudá-lo a se familiarizar melhor com ele e mostrar como você pode adaptá-lo para o seu próprio caso de uso.

Antes de executar o script, certifique-se de instalar a biblioteca a partir do código-fonte:

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

Navegue até a pasta de exemplo com o script de treinamento e instale as dependências necessárias:

cd examples/custom_diffusion
pip install -r requirements.txt
pip install clip-retrieval

🤗 Accelerate é uma biblioteca para ajudar você a treinar em múltiplos GPUs/TPUs ou com precisão mista. Ela configura automaticamente seu ambiente de treinamento com base no seu hardware e ambiente. Dê uma olhada no Quick tour do 🤗 Accelerate para saber mais.

Inicialize um ambiente do 🤗 Accelerate:

accelerate config

Para configurar um ambiente padrão do 🤗 Accelerate sem escolher nenhuma configuração:

accelerate config default

Ou, se o seu ambiente não suporta um shell interativo, como um notebook, você pode usar:

from accelerate.utils import write_basic_config

write_basic_config()

Por fim, se você deseja treinar um modelo em seu próprio conjunto de dados, consulte o guia Criar um conjunto de dados para treinamento para aprender a criar um conjunto de dados que funcione com o script de treinamento.

As seções a seguir destacam partes do script de treinamento que são importantes para entender como modificá-lo, mas não cobrem todos os aspectos do script em detalhes. Se você tiver interesse em aprender mais, sinta-se à vontade para ler o script e nos informar caso tenha alguma dúvida ou preocupação.

Parâmetros do script

O script de treinamento contém todos os parâmetros para ajudar você a personalizar sua execução de treinamento. Eles estão na função parse_args(). A função vem com valores padrão, mas você também pode definir seus próprios valores no comando de treinamento, se desejar.

Por exemplo, para alterar a resolução da imagem de entrada:

accelerate launch train_custom_diffusion.py \
  --resolution=256

Muitos dos parâmetros básicos são descritos no guia de treinamento do DreamBooth, portanto, este guia se concentra nos parâmetros exclusivos do Custom Diffusion:

--freeze_model: congela os parâmetros de chave e valor na camada de atenção cruzada; o padrão é crossattn_kv, mas você pode configurá-lo como crossattn para treinar todos os parâmetros na camada de atenção cruzada
--concepts_list: para aprender múltiplos conceitos, forneça um caminho para um arquivo JSON contendo os conceitos
--modifier_token: uma palavra especial usada para representar o conceito aprendido
--initializer_token: uma palavra especial usada para inicializar os embeddings do modifier_token

Perda de preservação de prioridade

A perda de preservação de prioridade é um método que utiliza os próprios amostras geradas por um modelo para ajudá-lo a aprender a gerar imagens mais diversas. Como essas imagens de amostra geradas pertencem à mesma classe das imagens fornecidas, elas ajudam o modelo a reter o que ele aprendeu sobre a classe e como pode usar o que já sabe sobre a classe para criar novas composições.

Muitos dos parâmetros para a perda de preservação anterior são descritos no guia de treinamento do DreamBooth.

Regularização

Custom Diffusion inclui o treinamento das imagens de destino com um pequeno conjunto de imagens reais para evitar overfitting. Como você pode imaginar, isso pode ser fácil de fazer quando você está treinando apenas com algumas imagens! Baixe 200 imagens reais com clip_retrieval. O class_prompt deve ser a mesma categoria das imagens de destino. Essas imagens são armazenadas em class_data_dir.

python retrieve.py --class_prompt cat --class_data_dir real_reg/samples_cat --num_class_images 200

Para habilitar a regularização, adicione os seguintes parâmetros:

--with_prior_preservation: se deve usar a perda de preservação de prior
--prior_loss_weight: controla a influência da perda de preservação do prior no modelo
--real_prior: se deve usar um pequeno conjunto de imagens reais para evitar overfitting

accelerate launch train_custom_diffusion.py \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --class_data_dir="./real_reg/samples_cat" \
  --class_prompt="cat" \
  --real_prior=True \

Script de treinamento

Muito do código no script de treinamento do Custom Diffusion é semelhante ao script do DreamBooth. Este guia, em vez disso, foca no código que é relevante para o Custom Diffusion.

O script de treinamento do Custom Diffusion possui duas classes de conjunto de dados:

CustomDiffusionDataset: pré-processa as imagens, imagens de classe e prompts para o treinamento
PromptDataset: prepara os prompts para gerar imagens de classes

Em seguida, o modifier_token é adicionado ao tokenizer, convertido em ids de tokens e os embeddings de tokens são redimensionados para levar em conta o novo modifier_token. Depois, os embeddings do modifier_token são inicializados com os embeddings do initializer_token. Todos os parâmetros no codificador de texto são congelados, exceto os embeddings de tokens, pois é isso que o modelo está tentando aprender a associar aos conceitos.

params_to_freeze = itertools.chain(
    text_encoder.text_model.encoder.parameters(),
    text_encoder.text_model.final_layer_norm.parameters(),
    text_encoder.text_model.embeddings.position_embedding.parameters(),
)
freeze_params(params_to_freeze)

Agora você precisará adicionar os pesos do Custom Diffusion às camadas de atenção. Este é um passo muito importante para obter a forma e o tamanho corretos dos pesos de atenção e para definir o número apropriado de processadores de atenção em cada bloco do UNet.

st = unet.state_dict()
for name, _ in unet.attn_processors.items():
    cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
    if name.startswith("mid_block"):
        hidden_size = unet.config.block_out_channels[-1]
    elif name.startswith("up_blocks"):
        block_id = int(name[len("up_blocks.")])
        hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
    elif name.startswith("down_blocks"):
        block_id = int(name[len("down_blocks.")])
        hidden_size = unet.config.block_out_channels[block_id]
    layer_name = name.split(".processor")[0]
    weights = {
        "to_k_custom_diffusion.weight": st[layer_name + ".to_k.weight"],
        "to_v_custom_diffusion.weight": st[layer_name + ".to_v.weight"],
    }
    if train_q_out:
        weights["to_q_custom_diffusion.weight"] = st[layer_name + ".to_q.weight"]
        weights["to_out_custom_diffusion.0.weight"] = st[layer_name + ".to_out.0.weight"]
        weights["to_out_custom_diffusion.0.bias"] = st[layer_name + ".to_out.0.bias"]
    if cross_attention_dim is not None:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=train_kv,
            train_q_out=train_q_out,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        ).to(unet.device)
        custom_diffusion_attn_procs[name].load_state_dict(weights)
    else:
        custom_diffusion_attn_procs[name] = attention_class(
            train_kv=False,
            train_q_out=False,
            hidden_size=hidden_size,
            cross_attention_dim=cross_attention_dim,
        )
del st
unet.set_attn_processor(custom_diffusion_attn_procs)
custom_diffusion_layers = AttnProcsLayers(unet.attn_processors)

O optimizer é inicializado para atualizar os parâmetros da camada de cross-attention:

optimizer = optimizer_class(
    itertools.chain(text_encoder.get_input_embeddings().parameters(), custom_diffusion_layers.parameters())
    if args.modifier_token is not None
    else custom_diffusion_layers.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

No loop de treinamento, é importante atualizar apenas os embeddings do conceito que você está tentando aprender. Isso significa definir os gradientes de todos os outros embeddings de tokens como zero:

if args.modifier_token is not None:
    if accelerator.num_processes > 1:
        grads_text_encoder = text_encoder.module.get_input_embeddings().weight.grad
    else:
        grads_text_encoder = text_encoder.get_input_embeddings().weight.grad
    index_grads_to_zero = torch.arange(len(tokenizer)) != modifier_token_id[0]
    for i in range(len(modifier_token_id[1:])):
        index_grads_to_zero = index_grads_to_zero & (
            torch.arange(len(tokenizer)) != modifier_token_id[i]
        )
    grads_text_encoder.data[index_grads_to_zero, :] = grads_text_encoder.data[
        index_grads_to_zero, :
    ].fill_(0)

Inicie o script

Depois de fazer todas as alterações ou se você estiver satisfeito com a configuração padrão, está pronto para executar o script de treinamento! 🚀

Neste guia, você baixará e usará estas imagens de gatos de exemplo. Você também pode criar e usar seu próprio conjunto de dados, se desejar (veja o guia Criar um conjunto de dados para treinamento).

Defina a variável de ambiente MODEL_NAME para um id de modelo no Hub ou um caminho para um modelo local, INSTANCE_DIR para o caminho onde você acabou de baixar as imagens de gatos e OUTPUT_DIR para onde deseja salvar o modelo. Você usará <new1> como a palavra especial para vincular os embeddings recém-aprendidos. O script cria e salva checkpoints do modelo e um arquivo pytorch_custom_diffusion_weights.bin no seu repositório.

Para monitorar o progresso do treinamento com Weights and Biases, adicione o parâmetro --report_to=wandb ao comando de treinamento e especifique um prompt de validação com --validation_prompt. Isso é útil para depuração e para salvar resultados intermediários.

Se você estiver treinando com rostos humanos, a equipe do Custom Diffusion descobriu que os seguintes parâmetros funcionam bem:

--learning_rate=5e-6

--max_train_steps pode estar em qualquer lugar entre 1000 e 2000

--freeze_model=crossattn

use pelo menos 15-20 imagens para treinar

conceito único

vários conceitos

Quando o treinamento for concluído, você pode usar seu novo modelo Custom Diffusion para inferência.

conceito único

vários conceitos

Próximas etapas

Parabéns por treinar um modelo com Custom Diffusion! 🎉 Para saber mais:

Leia o post no blog Multi-Concept Customization of Text-to-Image Diffusion para conhecer mais detalhes sobre os resultados experimentais da equipe Custom Diffusion.

Update on GitHub

Diffusers

Custom Diffusion

Parâmetros do script

Perda de preservação de prioridade

Regularização

Script de treinamento

Inicie o script

Próximas etapas