Modelos do tipo stable diffusion são uma das técnicas mais populares atualmente em inteligência artificial generativa. Porém, eles são de implementação complexa, com requerimentos computacionais de alto custo. Mas, felizmente, existem bibliotecas e ferramentas que podem ser empregadas para facilitar imensamente a implementação ou o uso desses modelos. Neste post, vamos usar a plataforma Hugging Face para explorar um pouco o que stable diffusion models podem fazer na geração de imagens.
O que é a plataforma Hugging Face?
Hugging Face é uma comunidade de inteligência artifical e plataforma criada em 2016 por Julien Chaumond, Clément Delangue e Thomas Wolf. Ela foi desenvolvida para democratizar principalmente o acesso a ferramentas de natural language processing (NLP). Porém, ela também conta com recursos que podem ser usados em aplicações de outras áreas.
Para realizar seus objetivos, a Hugging Face desenvolve e mantém ferramentas e bibliotecas populares relacionadas a NLP, machine learning e deep learning.
A plataforma Hugging Face também disponibiliza inúmeras pipelines. Elas possibilitam o acesso fácil a APIs com modelos pré-treinados da biblioteca Hugging Face Transformers. Isso permite que desenvolvedores usem modelos de última geração sem precisar escrever códigos extensos.
Neste post, iremos usar pipelines com modelos pré-treinados de stable diffusion para explorar a criação de imagens.
Por que usar modelos pré-treinados?
As principais razões para usar modelos pré-treinado se relacionam com custo e viabilidade computacional. O treinamento de modelos de deep learning do zero é uma tarefa muitas vezes impossível em termos de custos computacionais para a maioria dos mortais. Por causa disso, uma alternativa bastante comum é o uso de modelos pré-treinados. Eles podem ser facilmente customizados e adaptados para tarefas específicas com um custo computacional muito menor.
Modelos Stable diffusion pré-treinados para a criação de imagens com IA generativa
Diffusers é a biblioteca da Hugging Face ideal para modelos de difusão pré-treinados de última geração. Ela pode ser usada para gerar imagens, áudios e até estruturas 3D de moléculas. Além disso, ela serve para tarefas de inferências simples ou para o treinamento de modelos customizados. Segundo a Hugging Face:
“Nossa biblioteca foi projetada com foco na usabilidade em vez do desempenho, na simplicidade em vez da facilidade e na personalização em vez de abstrações.”
Portanto, neste post iremos explorar um pouco as capacidades da biblioteca Diffusers da Hugging Face para gerar imagens.
Pré-requisitos
Para criar imagens com a biblioteca Diffusers, é preciso instalar várias bibliotecas de suporte. Algumas dessas instalações podem ser problemáticas, por isso classificamos esse post como avançado.
Embora o uso de recursos da Hugging Face facilite muito a manipulação de modelos stable diffusion, isso não significa que essa tarefa é simples e apropriada para iniciantes. Ela requer o uso de GPU e pode envolver vários problemas de compatibilidade e memória.
Para esse post, é necessário instalar:
- numpy
- requests
- pytorch
- pillow
- transformers
- diffusers
- accelerate
O PyTorch precisa ser instalado com capacidade para GPU. Essa instalação pode ser problemática. Usando o conda, conseguimos instalar depois de várias tentativas com o seguinte comando:
conda install pytorch torchvision torchaudio pytorch-cuda=X.X -c pytorch -c nvidia
Onde X.X é a sua versão de GPU.
Para testar se o PyTorch está instalado corretamente, abra um terminal do Python e digite:
import torch
torch.cuda.is_available()
Se a execução dos comandos acima retornar True, a instalação foi bem sucedida.
Modelos Stable Diffusion: Text-to-img
A primeira coisa que testamos com a biblioteca Diffusers foi a geração de imagens a partir de textos (prompts), conforme feito em aplicações como DALL-E-2.
O código que usamos é mostrado abaixo. Ele é bem simples e consiste em comandos para chamar a pipeline que iremos usar (linha 1), definir o modelo pré-treinado (linhas 3-4), garantir o uso de GPU (linha 5), definir um prompt (linha 6), gerar a imagem com o modelo (linha 7) e salvar a imagem gerada (linha 8). O prompt mostrado no código abaixo foi usado para gerar a imagem de abertura do post.
from diffusers import StableDiffusionPipeline
model_id ="runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe = pipe.to("cuda")
prompt = "portrait photo of an old warrior chief"
image = pipe(prompt).images[0]
image.save("warrior_chief.png")
Caso esse código não funcione para você, uma alternativa é mostrada abaixo. Ela foi modificada daqui. Nessa versão, definimos safety_checker=None para acelerar os processos de inferência. Além disso, nas linhas em destaque, usamos dois comandos para reduzir os requerimentos de memória para execuções do código. Isso torna as execuções mais lentas, mas com consumo menor de memória. Por fim, também limitamos o número de inferências em 20. O número de inferências se relaciona às etapas de remoção de ruído do modelo.
from diffusers import StableDiffusionPipeline
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, safety_checker=None)
pipe.enable_sequential_cpu_offload()
pipe.enable_attention_slicing("max")
prompt = "portrait photo of an old warrior chief"
image = pipe(prompt, num_inference_steps=20).images[0]
image.save("img.png")
As imagens abaixo foram criadas com os seguintes prompts: “an astronaut in the beach”, “a realistic 3D rabbit”, “an astronaut riding a horse in Mars”, “an abstract representation of a human mind”. A geração de cada imagem levou cerca de 1 minuto no computador que testamos.
Modelos Stable Diffusion: Text-Guided Image-to-Image Generation
Em seguida, testamos uma outra pipeline super legal que gera imagens a partir de imagens usando texto como guia. O processo é simples. Nós apresentamos uma imagem e um prompt para o modelo. O modelo parte da imagem inicial para construir a solicitação do prompt.
Nosso código é mostrado abaixo. As mesmas alterações mostradas acima para redução de custos computacionais podem ser usadas aqui caso seja necessário.
Suas primeiras linhas logo após a importação das bibliotecas são parecidas com os comandos do código anterior. A principal diferença é que é preciso abrir uma imagem com PIL. Na linha 8, é substitua path_to_img pela localização da sua imagem. Se você não tiver imagens disponíveis, use o pexels. Selecione imagens pequenas e, preferencialmente, quadradas.
Nas linhas 9 e 10 do código, a imagem é convertida nos canais de cores e tamanho correto. Depois, definimos o prompt, geramos e salvamos a imagem.
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")
img = "path_to_img"
init_image = Image.open(img).convert("RGB")
init_image.thumbnail((512, 512))
prompt = "A realistic red rose on a landscape"
image = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5, num_inference_steps=50).images[0]
image.save("img.png")
O prompt mostrado no código acima foi usado para criar a 2ª e a 4ª imagens mostradas abaixo. As imagens originais utilizadas em suas gerações são a 1ª e a 3ª abaixo.
Depois testamos um prompt diferente partindo de outra imagem. No exemplo abaixo, a imagem original é a primeira. Nosso prompt foi “a magical landscape”. Executamos o código 3 vezes e esses foram os resultados obtidos.
Modelos Stable Diffusion: Text-Guided Image-Inpainting
A última pipeline que exploramos foi a mais assustadora. Ela se chama InpaintPipeline. Basicamente, ela parte de uma imagem e uma máscara. A máscara é usada para fazer uma substituição na imagem original a partir de um prompt.
Veja o código abaixo, ele segue a mesma lógica dos anteriores. Mas note que as imagens estão sendo buscadas com a biblioteca requests.
from diffusers import StableDiffusionInpaintPipeline
from io import BytesIO
from PIL import Image
import requests
import torch
pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
response = requests.get(img_url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
response = requests.get(mask_url)
mask = Image.open(BytesIO(response.content)).convert("RGB")
init_image = init_image.resize((512, 512))
mask_image = mask.resize((512, 512))
#prompt = "red rose, high resolution"
#prompt = "white rabbit, high resolution"
prompt = "small white dog"
image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
image.save("dog.png")
A imagem original e a máscara que usamos foram obtidas no site da Hugging Face e são mostradas abaixo.
Partindo da imagem original e de sua máscara, usamos os três prompts mostrados no código acima. Esses são os resultados obtidos. Na nossa opinião, eles são impressionantes.
Conclusões
Nesse post, exploramos um pouco as potencialidades dos modelos de stable diffusion da plataforma Hugging Face. Eles são relativamente fáceis de usar, embora seus custos computacionais ainda sejam altos para muitos computadores. Mas seus resultados são impressionantes. Se você quiser explorar mais, consulte o site da plataforma e as outras pipelines disponíveis.
