Implantação do Stable Diffusion no Vertex AI

O Stable Diffusion é um modelo de geração de imagens. Foi disponibilizado como código aberto em 2022 e ganhou popularidade devido à sua capacidade de gerar imagens de alta qualidade a partir de descrições textuais. Assim como outros modelos de geração de imagens, tais como o Dall-E, o Stable Diffusion utiliza técnicas de aprendizado de máquina para gerar imagens com base em uma entrada fornecida.

A Moët Hennessy, divisão de vinhos e bebidas destiladas 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 já colaborou com Artefact para investigar os usos potenciais de tecnologia de ponta na geração de conteúdo de marketing. Com foco na privacidade e segurança, a equipe decidiu explorar a implantação do Stable Diffusion na Google Cloud (GCP) para permitir que a Moët Hennessy ajustasse e executasse o modelo dentro de sua própria infraestrutura, proporcionando uma experiência integrada desde o ajuste 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 (GCP) e, especificamente, do Vertex AI. Isso inclui conceitos como o registro de modelos e os endpoints do Vertex. Além disso, é necessário ter experiência prévia com o Docker para seguir algumas das etapas. Se você não estiver familiarizado com esses conceitos, recomenda-se que se familiarize com eles antes de prosseguir.
Além disso, para baixar os pesos do Stable Diffusion, você precisa ter uma conta no Hugging Face; caso ainda não tenha uma, pode criá-la facilmente nosite do Hugging Face.
Dito isso, vamos começar!

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

É sempre importante criar um ambiente virtual para instalar os pacotes. Pessoalmente, vou usar o Anaconda e instalar o arquivo requirements.txt com todas as dependências:

Agora você está pronto para baixar os pesos do Stable Diffusion; neste artigo, usaremoso Stable Diffusion 1.5. É necessário aceitar a licença na página do modelo; caso contrário, ocorrerão erros ao baixar os pesos do modelo.
Acesse 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 com abiblioteca python-dotenv.
Você precisa acessar a pasta `src/stable_diffusion` e executar:

Isso fará o download dos pesos

O Torchserveé uma estrutura para a implementação de modelos PyTorch. Ele permite que você implemente modelos PyTorch em um ambiente de produção e oferece recursos como controle de versões de modelos e implementação de múltiplos modelos. Foi projetado para ser fácil de usar e permite que você se concentre na criação e na implementação de seus modelos, em vez de se preocupar com a infraestrutura.

Um manipulador personalizado é uma classe Python que define como pré-processar data de entrada, 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á está disponível no repositório do TorchServe. No entanto, o endpoint do Vertex espera um formato específico para as solicitações, portanto, precisamos adaptar o método preprocess() do manipulador para atender ao formato do Vertex. Você pode usar a versão modificada do manipulador chamada `stable_diffusion_handler.py`, disponível no repositório do GitHub deste artigo.

Depois de criar seu manipulador personalizado, você precisará empacotá-lo, juntamente com todas 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 seu modelo, manipulador e dependências em um único arquivo.

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

Talvez seja necessário editar o caminho de acordo com a localização dos arquivos no seu computador:

Depois de criar seu arquivo .mar, você pode iniciar o servidor TorchServe usando o comando torchserve. Para isso, execute o seguinte comando:

O arquivo config.properties permite especificar as configurações do seu servidor TorchServe, como a porta para inferência, verificações de integridade, número de workers, etc.
AVISO: Todos os scripts devem ser executados no mesmo diretório em que o arquivo está localizado para evitar erros de caminho.

É importante testar o código localmente antes de começar a implementar o Docker na nossa implantação; já preparei um script em Bash que criará o arquivo .mar e iniciará o servidor TorchServe.

Você pode executá-lo com:

É preciso aguardar alguns minutos até que o servidor seja inicializado e o worker carregue o modelo; você poderá iniciar a inferência assim que vir o seguinte registro:

Você pode então usar o código a seguir para fazer solicitações de inferência ao seu modelo:

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

O Stable Diffusion requer uma GPU para funcionar corretamente, por isso você não terá nenhum resultado por enquanto. Você pode parar o servidor torchserve com o comando `torchserve –stop`.

Seu servidor TorchServe está funcionando localmente; para implantar o Stable Diffusion no Vertex AI, você precisará criá-lo como um contêiner Docker. Isso significa criar uma imagem Docker que contenha o modelo, o manipulador personalizado e todas as dependências necessárias. Trata-se simplesmente de todas as etapas que realizamos acima dentro de 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 compilá-lo localmente, mas você pode compilá-lo usando cloud e baixá-lo para o seu computador.

Compile a imagem localmente (é necessário que o daemon do Docker esteja em execução):

A compilação da imagem levará de 20 a 30 minutos; a fase de compilação é demorada porque os pesos do modelo precisam ser copiados para dentro da imagem antes de serem empacotados pelo model-archiver.

Você pode executar um contêiner Docker e ouvir na porta 7080 com:

Para verificar se tudo está funcionando corretamente, você pode esperar até que o worker 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:

É exatamente isso que o script bash `deploy.sh` vai fazer; você pode executá-lo com:

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

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

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

def query_endpoint(endpoint, text_input):
payload =
response = endpoint.predict(instances=[payload])
return response

Uma garrafa de conhaque envelhecido e exclusivo repousa sobre uma superfície reflexiva, em frente a um bar animado, com detalhes impressionantes, em 4K e efeito bokeh — Stable Diffusion 1.5

A inferência deve levar entre 10 e 15 segundos em uma GPU T4. Você pode aumentar a velocidade escolhendo uma GPU mais potente; para isso, altere a variável ACCELERATOR_TYPE no arquivo deploy.sh.

Talvez você queira manter um registro do histórico das imagens, ou esteja enfrentando 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 no GCS e retornar apenas o caminho da imagem no GCS.

Felizmente, já preparei um handler que faz exatamente isso para você; basta mudar do stable_diffusion_handler.py parao stable_diffusion_handler_gcs.pye o problema estará resolvido.

AVISO: Antes de executar a implantação com o novo manipulador, é necessário:

Neste artigo, é apresentado um guia completo para a implantação do modelo Stable Diffusion na Google Cloud usando o Vertex AI.

O guia abordou etapas essenciais, tais como:

É importante lembrar que, embora o endpoint do Vertex possa ser uma solução eficaz, ele não permite reduzir a escala para 0 instâncias, o que pode resultar em aumento de custos, uma vez que a GPU continua em uso.

Além disso, com a implantação do modelo Stable Diffusion concluída, estamos atualmente explorando novas possibilidades, como o ajuste fino do modelo para produtos específicos da Moët Hennessy, a fim de aprimorar ainda mais as capacidades do modelo.