Como treinei o Stable Diffusion para gerar imagens minhas?

Este tutorial mostrará como eu usei o DreamBooth do artigo do Google DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation. Para treinar um modelo de IA capaz de criar imagens minhas.

O artigo se apresenta como:

É como uma cabine de fotos, mas uma vez que o assunto é capturado, ele pode ser sintetizado onde quer que seus sonhos o levem…

Imagem mostrando os recursos do Dreambooth. Um cachorro é adicionado a diferentes imagens. Exemplos de imagens geradas pelo DreamBooth

A Hugging Face 🤗, conhecida por tornar os modelos de aprendizado de máquina acessíveis ao público, incorporou o DreamBooth em seu ecossistema em setembro deste ano. Isso foi um divisor de águas em termos de usabilidade do DreamBooth.

Passos para treinar o modelo

Todo o processo foi bastante simples depois que encontrei um notebook que funcionasse bem. O campo de IA generativa tem evoluído rapidamente nos últimos meses e, por causa disso, muitas implementações estão desatualizadas ou simplesmente quebradas devido a dependências que mudaram.

O notebook que usei originalmente e que funcionou para mim foi o implementado por ShivamShrirao. Você pode encontrá-lo aqui.

Para este tutorial, usaremos um fork que fiz do notebook do ShivamShrirao. Meu notebook é mais simples e se concentra na criação de imagens humanas. Se você deseja aprender como salvar seu modelo ou como gerar imagens com animais ou outros objetos, sugiro que você utilize o notebook do ShivamShrirao.

Dividi o treinamento e a inferência do modelo em cinco etapas. Elas são:

1. Crie uma conta na Hugging Face

Vamos usar modelos hospedados na Hugging Face. Para poder baixá-los, precisamos de uma conta. Você pode criar uma aqui.

Com a conta criada, vá para foto do perfil → Configurações → Tokens de acesso e gere um novo token. Vamos precisar deste token mais tarde para baixar o modelo do notebook. Um token de Leitura (Read) é suficiente para este tutorial.

2. Crie o conjunto de dados com suas imagens

Para o meu conjunto de dados, usei 16 imagens diferentes de mim mesmo. Segui as diretrizes do notebook de John Penna.

AVISO: Certifique-se de fazer o upload de um número par de imagens, caso contrário, o treinamento irá parar inexplicavelmente em 1500 passos.

  • 2–3 corpo inteiro
  • 3–5 parte superior do corpo
  • 5–12 close-up no rosto

As imagens precisam ser padronizadas. O modelo espera que a entrada seja 512×512px. Para padronizar minhas imagens, usei uma ferramenta online chamada Birme. Depois de fazer o upload das minhas fotos no Birme, certifiquei-me de que meu rosto estava no centro de cada imagem. Em seguida, baixei-as como um arquivo zip.

3. Abra o projeto no Google Colab

Para usar o Colab, você precisa de uma conta do Google. Depois de fazer login em sua conta, é hora de abrir o projeto.

Para abrir o notebook no Google Colab, vá para a página README do projeto e clique no botão Abrir no Colab. Ou clique aqui.

4. Execute as células do notebook até a seção de treinamento do modelo

Execute as células uma por uma pressionando shift+enter ou clicando no botão de execução à esquerda de cada célula. Há algumas células às quais você precisa prestar atenção extra, são elas:

  • Login no Hugging Face 🤗
    Antes de executar esta célula, você precisa fazer duas coisas:
    • Visite o repositório do modelo, leia a licença e marque a caixa de seleção da licença se concordar. Do contrário, você não poderá baixar o modelo.
    • Copie o token de acesso da Hugging Face gerado na Etapa 1 e cole-o no campo HUGGINGFACE_TOKEN. Execute a célula depois disso.
  • Faça o upload das imagens
    Depois de executar esta célula, um botão escrito Escolher arquivos (Choose files) aparecerá. Clique nele e selecione as imagens que deseja usar no treinamento.

Assim que o upload for concluído, você deve ver suas imagens na pasta data/ukj no painel de arquivos à esquerda.

Imagens carregadas na pasta data/ukj

  • Execute o treinamento
    Você pode ajustar os parâmetros do treinamento antes de executar esta célula. Com os parâmetros padrão, o treinamento leva cerca de 20 minutos.

5. Gerando novas imagens

A seção de inferência do notebook possui alguns parâmetros que afetarão os resultados de diferentes maneiras, os parâmetros são:

  • prompt: as imagens geradas serão baseadas nesta entrada de texto. Você pode encontrar exemplos de entradas na próxima seção deste tutorial. O que é importante para o prompt é que a pessoa aprendida pelo modelo é referida como ukj person. Por exemplo, se você deseja se inserir no filme Matrix, pode escrever o seguinte prompt:

> Cena do filme de ukj person como Morpheus de The Matrix (1999).

  • num_samples: quantas imagens devem ser geradas.
  • guidance_scale: o usuário da Hugging Face ghpkishore deu uma boa explicação do que este parâmetro significa aqui.

Quanto maior, mais de perto ele segue o prompt. No entanto, após um certo valor, ele se torna aleatório. Pense nisso como pesos para o prompt. Aumentar a guidance_scale faz com que a geração siga mais de perto o prompt.

  • num_inference_steps: da documentação do diffusers:

Em geral, os resultados são melhores quanto mais etapas você usa, no entanto, quanto mais etapas, mais tempo leva a geração.

  • height and width: Esses parâmetros definem o tamanho da imagem gerada.
  • seed: A semente garante a reprodutibilidade do seu prompt. Usando sementes diferentes para o mesmo prompt, você obterá resultados diferentes.

Exemplos de gerações

As imagens abaixo são os resultados que obtive após treinar meu modelo. Encontrei inspirações para os prompts no vídeo do Corridor Crew Is This The Death of VFX? e neste blog post.

Prompt: Movie still ukj person as Iron Man from The Avengers.

Uma imagem minha vestindo uma roupa do Homem de Ferro

Prompt: A digital painting of a warlord, Greg Rutkowski, ukj person

Eu como uma pintura digital de um senhor da guerra, Greg Rutkowski

Prompt: Movie still ukj person as Morpheus from The Matrix (1999).

Uma imagem minha como Morpheus de The Matrix (1999)

Prompt: Portrait of ukj person as a cartoon from the The Adventures of Tintin

Uma imagem minha como um desenho animado das Aventuras de Tintin

Prompt: Portrait of ukj person as a 3d model

Uma imagem minha como um modelo 3D

Conclusão

Devo dizer que não esperava que os resultados fossem tão bons. A quantidade de informações que esses modelos podem extrair de algumas imagens é surpreendente.

Ele aprendeu pequenas nuances do meu rosto que são quase imperceptíveis, como manchas onde minha barba não cresceu. Eu também vi imagens em que minha camiseta estava fora do lugar da mesma maneira que estava em duas imagens do conjunto de treinamento, o que é impressionante para mim.

Outros recursos e inspirações

Até o próximo post! 👋

Caso queira entrar em contato, você pode me encontrar no Twitter ou no LinkedIn.