Implementando o Stable Diffusion no Vertex AI

O Stable Diffusion é um modelo de geração de imagens. Seu código-fonte foi aberto em 2022 e ganhou popularidade devido à sua capacidade de gerar imagens de alta qualidade a partir de descrições de texto. Como outros modelos de geração de imagens, como o Dall-E, o Stable Diffusion usa técnicas de aprendizado de máquina para gerar imagens com base em uma determinada entrada.

A Moët Hennessy, a divisão de vinhos e destilados do conglomerado de luxo LVMH, administra um portfólio de mais de 26 marcas icônicas, como Moët & Chandon, Hennessy e Veuve Clicquot. A Moët Hennessy colaborou com a Artefact para investigar os possíveis usos da tecnologia de ponta na geração de conteúdo de marketing. Com foco na privacidade e na segurança, a equipe decidiu explorar a implantação do Stable Diffusion no Google Cloud Platform (GCP) para permitir que a Moët Hennessy fizesse o ajuste fino e executasse o modelo em sua própria infraestrutura, proporcionando uma experiência perfeita desde o ajuste fino do modelo até a exposição da API.

Antes de começar, é importante observar que este artigo pressupõe que você tenha conhecimento prévio do Google Cloud Platform (GCP) e, especificamente, do Vertex AI. Isso inclui conceitos como registro de modelos e pontos de extremidade do Vertex. Além disso, você precisa ter experiência prévia com o Docker para seguir algumas das etapas. Se não estiver familiarizado com esses conceitos, é recomendável que você se familiarize com eles antes de prosseguir.
Além disso, para fazer o download do Stable Diffusion Weights, você precisa ter uma conta na huggingface; se ainda não tiver uma, poderá criá-la facilmente no site da huggingface.
Dito isso, vamos começar!

Recomendo clonar orepositório do github que preparei para seguir as etapas do artigo.

É sempre importante criar um ambiente virtual para instalar os pacotes. Pessoalmente, usarei o anaconda e instalarei o arquivo requirements.txt com cada dependência:

Agora você está pronto para fazer o download dos pesos do Stable Diffusion. Neste artigo, usaremos o Stable Diffusion 1.5. Você precisa aceitar a licença na página do modelo, caso contrário, haverá erros ao fazer o download dos pesos do modelo.
Vá para sua conta → configurações → token de acesso → novo token (acesso de leitura)

Você pode adicionar o token como uma variável de ambiente. Pessoalmente, recomendo usar um arquivo .env e carregar a variável de ambiente usando a biblioteca python-dotenv.
Você precisa navegar até `src/stable_diffusion` e executar:

Isso fará o download dos pesos internos

O Torchserve é uma estrutura para servir modelos PyTorch. Ele permite implantar modelos do PyTorch em um ambiente de produção e fornece recursos como versão de modelo e veiculação de vários modelos. Ele foi projetado para ser fácil de usar e permite que você se concentre na criação e na implantação de seus modelos, em vez de se preocupar com a infraestrutura.

Um manipulador personalizado é uma classe Python que define como pré-processar a entrada data, como executar o modelo e como pós-processar a saída. Para criar um manipulador personalizado para o seu modelo, você precisará criar uma classe Python que siga o formato do TorchServe.
Um manipulador personalizado para o Stable Diffusion já é fornecido no repositório do TorchServe. Mas o ponto de extremidade Vertex espera um formato específico para as solicitações, portanto, precisamos adaptar o método preprocess() do manipulador para levar em conta o formato Vertex. Você pode usar a versão modificada do manipulador chamada `stable_diffusion_handler.py` fornecida no repositório do github deste artigo.

Depois de criar o manipulador personalizado, você precisará empacotá-lo junto com as dependências e o próprio modelo em um arquivo .mar usando a ferramenta model-archiver. A ferramenta model-archiver é uma ferramenta de linha de comando que permite empacotar o modelo, o manipulador e as dependências em um único arquivo.

Isso criará um arquivo .mar chamado output.mar que contém o modelo, o manipulador e as dependências.

Talvez seja necessário editar o caminho com base na localização dos arquivos em seu computador:

Depois de criar o arquivo .mar, você pode iniciar o servidor TorchServe usando o comando torchserve. Para fazer isso, você precisará executar o seguinte comando:

O config.properties permite especificar as configurações do seu servidor TorchServe, como porta para inferência, verificações de integridade, número de trabalhadores etc.
AVISO: Todos os scripts precisam ser executados onde o arquivo está localizado para evitar erros de caminho.

É importante testar o código localmente antes de começar a dockerizar nossa implantação. Já preparei um script bash que criará o arquivo .mar e iniciará o servidor TorchServe.

Você pode executá-lo com :

É necessário aguardar alguns minutos para que o servidor seja inicializado e o trabalhador carregue o modelo. Você poderá começar a executar a inferência se vir o seguinte registro:

Em seguida, você pode usar o código a seguir para fazer solicitações de inferência em seu modelo:

Você pode verificar se o servidor recebeu sua solicitação examinando os logs do servidor:

O Stable Diffusion requer uma GPU para ser executado sem problemas, portanto, você não terá nenhuma saída no momento. Você pode parar o servidor torchserve com `torhcserve -stop`.

Seu servidor TorchServe está funcionando localmente. Para implantar o Stable Diffusion no Vertex AI, você precisará colocá-lo no Docker. Isso significa criar uma imagem do Docker que contenha o modelo, o manipulador personalizado e todas as dependências necessárias. Isso é simplesmente todas as etapas que fizemos acima em um Dockerfile.
Felizmente, ele já está preparado e pronto para uso aqui : Dockerfile.

É importante executar o contêiner localmente para verificar se ele está funcionando corretamente. Vou construí-lo localmente, mas você pode construí-lo com o cloud build e colocá-lo em sua máquina.

Crie a imagem localmente (você precisa ter o daemon do Docker em execução):

A construção da imagem levará de 20 a 30 minutos; a fase de construção é longa porque os pesos do modelo precisam ser copiados dentro da imagem antes de serem empacotados pelo arquivador de modelos.

Você pode executar um contêiner do docker e escutar na porta 7080 com:

Para verificar se tudo está funcionando corretamente, você pode esperar até que o trabalhador tenha carregado o modelo e executar o mesmo código de inferência de antes, já que estamos usando a mesma porta 7080.

Agora que o Dockerfile está pronto e funcionando, precisamos:

Isso é exatamente o que o script bash deploy.sh fará, e você pode executá-lo com:

A implantação leva cerca de 1 hora com uma boa velocidade de Internet:

Depois que o modelo for anexado com êxito ao endpoint, você poderá consultar o endpoint com o seguinte código:

PROJECT_NAME = ""
REGION = ""
ENDPOINT_ID = ""
aip.init(project=PROJECT_NAME, location=REGION)
endpoint = aip.Endpoint(endpoint_name=ENDPOINT_ID)
text_input = """Uma garrafa de conhaque envelhecido e exclusivo
fica em uma superfície reflexiva, em frente a um bar vibrante,
hiperdetalhado, 4K, bokeh"""

def query_endpoint(endpoint, text_input):
payload = {“data”: text_input}
response = endpoint.predict(instances=[payload])
return response

Uma garrafa de conhaque envelhecido e exclusivo em uma superfície reflexiva, em frente a um bar vibrante, hiperdetalhado, 4K, bokeh - Stable Diffusion 1.5

A inferência deve levar entre 10 e 15 segundos na GPU T4. Você pode aumentar a velocidade escolhendo uma GPU melhor e pode alterar a variável ACCELERATOR_TYPE no arquivo deploy.sh.

Talvez você queira manter o controle do histórico das imagens ou esteja tendo erros devido ao limite de tamanho de 1,5 MB para a resposta do endpoint. Nesse caso, recomendo usar o método de pós-processamento do manipulador para salvar a imagem dentro do GCS e retornar somente o caminho do GCS da imagem.

Felizmente, já preparei um manipulador que faz exatamente isso para você. Mudar de stable_diffusion_handler.py para stable_diffusion_handler_gcs.py deve resolver o problema.

AVISO: Antes de executar a implantação com o novo manipulador, você precisa:

Neste artigo, foi fornecido um guia abrangente para a implantação do modelo Stable Diffusion no Google Cloud Platform usando o Vertex AI.

O guia abordou etapas essenciais, como:

É importante lembrar que, embora o ponto de extremidade Vertex possa ser uma solução eficaz, ele não oferece suporte à redução para 0 instâncias, o que pode levar a um aumento nos custos, pois a GPU continua em uso.

Além disso, com a implantação do modelo de difusão estável concluída, estamos explorando outras possibilidades, como o ajuste fino do modelo em produtos específicos da Moet Hennessy para aprimorar ainda mais os recursos do modelo.