Article image
Carlos Lima
Carlos Lima16/04/2024 20:13
Compartilhe

Como gerar imagens de textos: Código em Python

  • #Machine Learning
  • #Python

Requerimentos

O objetivo do código é gera imagem de texto (text2image) usando modelos pré-treinados.

Para isto vamos usar o The Stable Diffusion cuja documentação se encontra no site hugginface.

Para usarmos precisamos instar a classe diffusers.

Código

!pip install -qq diffusers

import torch
from diffusers import StableDiffusionPipeline

def generate_image(prompt):
  '''
  Essa função entrega um prompt (descrição do user) e 
  o modelo pré-treinado retorna uma imagem.
  '''

  pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-   diffusion-v1-5",  # <- Modelo usado
                                                 torch_dtype=torch.float16) # <- Tipo do tensor
  pipe = pipe.to("cuda")
  image = pipe(prompt).images[0]

  return image

prompt = "a cute dog"
generate_image(prompt)

Saída

image

Obs: Todo o código foi feito usando o Google Colabs, em caso de tentar localmente considere maior atenção quantos aos pré-requisitos.

Compartilhe
Comentários (0)