Usando Python e OpenAI para gerar e redimensionar imagens (#dev #python #openai #ai #aiart)

Já que geral está usando a OpenAI (API ou ChatGPT) para criar conteúdo, decidi entrar na festa e escrever um tutorial mostrando uma maneira simples de usar a OpenAI (não ChatGPT) para automatizar o processo de gerar imagens a partir de um prompt, e depois criar variantes dessa imagem em diferentes tamanhos.

Antes de começar - Uma observação importante

Este tutorial é apenas para fins educacionais e de aprendizado. Eu não sou responsável por nada que você faça. Se você criar algo que seja ofensivo, ilegal ou prejudicial, isso é com você. Também se você postar acidentalmente sua chave de API no GitHub ou em algum lugar e depois receber uma fatura de milhares de dólares, isso também é com você. Tenha cuidado com sua chave de API e use-a de forma responsável.

O que você vai precisar

Como você pode imaginar, a OpenAI não é gratuita, então você precisará criar uma conta lá e obter uma chave de API. Do site deles:

Primeiro, crie uma conta na OpenAI ou faça login. Em seguida, navegue até a página da chave da API e “Crie uma nova chave secreta”, nomeando opcionalmente a chave. Certifique-se de salvar isso em um local seguro e não compartilhá-lo com ninguém.

Também recomendo que você acesse a página Usage: Cost e defina um limite de orçamento para sua conta. Dessa forma, você não terá uma surpresa na fatura no final do mês.

Para minhas necessidades , um limite de orçamento de $20 dólares (USD) é mais do que suficiente. O código neste script consome $0,08 do meu orçamento.

O que vamos fazer

O script que vamos criar chamará a API da OpenAI para gerar uma imagem com base em um prompt + configuração do modelo e, em seguida, criará diferentes versões dessa imagem com tamanhos diferentes, mas isso será feito localmente, então você não será cobrado por isso.

Configuração do ambiente

Para isso, precisaremos das seguintes bibliotecas:

• Biblioteca Python da OpenAI
• Pillow (PIL)
• Requests
• Io
• Types
• PathLib
• PPrint

Felizmente, a maioria dessas bibliotecas já está instalada em seu ambiente Python, então você só precisará instalar o seguinte:

pip install --upgrade openai
pip install --upgrade Pillow

Configuração inicial

Isso é algo que pode mudar no futuro, mas por enquanto podemos usar 2 modelos para isso: dall-e-2 e dall-e-3. Então, se quisermos permitir mais flexibilidade, devemos criar uma maneira de definir os valores permitidos e os parâmetros para cada um desses modelos.

Neste caso, usaremos um dicionário para armazenar os parâmetros de cada modelo.

Explique o que cada parâmetro significa:

• max_prompt_size: O número máximo de caracteres permitidos para o prompt.
• n: O número de imagens a serem geradas.
• quality: A qualidade da imagem gerada.
• response_format: O formato da resposta.
• size: O tamanho da imagem gerada.
• style: O estilo da imagem gerada.

Após ler a documentação, para o modelo dall-e-2 temos:

{
    "max_prompt_size": 1000,
    "n": {
        "min": 1,
        "max": 10
    },
    "quality": None,  # Não suportado.
    "response_format": [
        "url",  # Disponível apenas por 60 minutos após a solicitação
        "b64_json"
    ],
    "size": [
        "256x256",   # Quadrado
        "512x512",   # Quadrado maior
        "1024x1024"  # Quadrado ainda maior
    ],
    "style": None  # Não suportado.
}

E para o modelo dall-e-3 temos:

{
    "max_prompt_size": 4000,
    "n": {
        "min": 1,
        "max": 1
    },
    "quality": ["standard", "hd"],
    "response_format": [
        "url",  # Disponível apenas por 60 minutos após a solicitação
        "b64_json"
    ],
    "size": [
        "1024x1792",  # Retrato
        "1792x1024",  # Paisagem
        "1024x1024"   # Quadrado
    ],
    "style": ["vivid", "natural"]
}

Sinta-se à vontade para alterar esses valores posteriormente, mas usarei dall-e-3 como o modelo padrão para este script e a seguinte configuração:

• n: 1
• quality: standard
• response_format: url
• size: 1024x1792
• style: vivid

Como esses valores são fixos, podemos criar um dicionário com eles:

from types import MappingProxyType

default_config = MappingProxyType({  # Immutable dictionary
    "n": params[default_model]["n"]["min"],
    "quality": params[default_model]["quality"][0],
    "response_format": params[default_model]["response_format"][0],
    "size": params[default_model]["size"][0],
    "style": params[default_model]["style"][0]
})

A razão pela qual estou usando MappingProxyType é para tornar o dicionário imutável, para que não possamos alterar os valores posteriormente e possamos usá-lo como o valor padrão de um argumento sem que o linter reclame.

• Referência: https://platform.openai.com/docs/api-reference/images/create

Gerando a imagem

Agora que temos a configuração padrão, podemos criar uma função que chamará a API da OpenAI para gerar a imagem.

Esta função retornará a URL da imagem gerada e receberá os seguintes argumentos:

• Prompt - str;
• model - str (padrão: dall-e-3);
• config - dict (padrão: default_config).

Como estamos recebendo argumentos, precisamos garantir que eles sejam válidos.

Validando o modelo

Para validar o modelo, verificamos se temos um nome de modelo e se ele é válido, e podemos fazer isso com esta função:

def _ensure_model_name_is_valid(model: str):
    if model is None:
        raise ValueError("Model name cannot be None.")

    if model in params:
        return

    raise ValueError(f"Model name must be one of the following: {', '.join(params.keys())}")

Validando o prompt

Isso é ainda mais simples de verificar. Só precisamos garantir que temos um prompt e que ele está dentro do tamanho permitido. Se isso fosse um código de nível de produção, poderíamos verificar se este era um prompt malicioso, etc., mas por enquanto a vida é simples.

def _ensure_prompt_is_valid(prompt: str, model: str):
    max_prompt_size = params[model]["max_prompt_size"]

    if prompt is None:
        raise ValueError("Prompt cannot be None.")

    if len(prompt) <= max_prompt_size:
        return

    raise ValueError(f"Prompt must be less than or equal to {max_prompt_size} characters.")

Validando a configuração

Isso é um pouco mais complexo, pois precisamos verificar se os valores são válidos para o modelo que estamos usando. Não vou me preocupar muito com a otimização aqui, então vou apenas verificar cada valor de configuração em relação às definições do modelo.

Nota: Eu poderia criar uma função para sanitizar os valores de configuração, mas este post já é grande o suficiente.

def _ensure_config_is_valid(config: dict, model: str):
    for key, value in config.items():
        if key in params[model]:
            if value is None:
                raise ValueError(f"Value for key '{key}' cannot be None.")

            if key == "n":
                if value < params[model][key]["min"] or value > params[model][key]["max"]:
                    raise ValueError(f"Value for key '{key}' must be between {params[model][key]['min']} and {params[model][key]['max']}.")

            if key == "quality":
                if value not in params[model][key]:
                    raise ValueError(f"Value for key '{key}' must be one of the following: {', '.join(params[model][key])}")

            if key == "response_format":
                if value not in params[model][key]:
                    raise ValueError(f"Value for key '{key}' must be one of the following: {', '.join(params[model][key])}")

            if key == "size":
                if value not in params[model][key]:
                    raise ValueError(f"Value for key '{key}' must be one of the following: {', '.join(params[model][key])}")

            if key == "style":
                if value not in params[model][key]:
                    raise ValueError(f"Value for key '{key}' must be one of the following: {', '.join(params[model][key])}")

Gerando a imagem

Agora que temos as funções de validação, vamos passar para a função que realmente chamará a API da OpenAI.

Primeiro, chamamos as funções de validação que acabamos de criar:

_ensure_model_name_is_valid(model)
_ensure_prompt_is_valid(prompt, model)
_ensure_config_is_valid(config, model)

Em seguida, criamos o cliente OpenAI:

import openai

client = openai.OpenAI(api_key="ab-proj-1234567890")

E então fazemos a solicitação para gerar a imagem:

response = client.images.generate(
    prompt=prompt,
    model=model,
    **config
)

Neste código, passamos explicitamente o prompt e o modelo. A configuração que criamos, estou usando um desempacotamento de dicionário para passar os valores como argumentos de palavra-chave (como fazemos no TypeScript…)

Esta é uma chamada de bloqueio e lançará uma exceção se algo der errado, então podemos simplesmente retornar a URL da imagem gerada:

return response.data[0].url

E é isso! Até agora, temos uma função que gerará uma imagem com base em um prompt e retornará a URL da imagem gerada de uma maneira flexível o suficiente para que possamos alterar facilmente os parâmetros da solicitação. Se isso é tudo o que você queria, você pode parar por aqui e usar esta função em seus projetos.

Baixando a imagem gerada

Agora que temos a URL da imagem gerada, podemos baixá-la e salvá-la em um arquivo. Como estamos gerando variações dessa imagem, adicionaremos a palavra original a este nome de arquivo.

Esta função receberá dois argumentos:

• url - str;
• nome do arquivo sem a extensão - str.

Vamos confiar no processo e não validaremos essas entradas, mas você poderia.

Primeiro, fazemos uma solicitação para a URL e obtemos os dados da imagem:

import requests
response = requests.get(image_url)

Para garantir que nada deu errado, e como já estamos gerando exceptions em caso de falha, podemos fazer isso:

response.raise_for_status()

Se nenhuma exceção foi gerada, podemos carregar com segurança a imagem para um objeto Pillow:

from io import BytesIO
from PIL import Image

original_image = Image.open(BytesIO(response.content))

Agora vamos ajustar o nome do arquivo:

original_filename = f"{filename_without_ext}_original.jpg"

E salvar a imagem:

original_image.save(original_filename)

E, por fim, retornar o nome do arquivo e o objeto da imagem:

# Get the image width and height
original_width, original_height = original_image.size

# Calculate the ratio
ratio = original_height / original_width

# Return the image details
return {
    "filename_without_ext": filename_without_ext,
    "original_filename": original_filename,
    "image_object": original_image,
    "path": Path(original_filename),
    "width": original_width,
    "height": original_height,
    "ratio": ratio
}

Criando variações da imagem

Agora que temos a imagem original, podemos criar variações dela com tamanhos diferentes. Esta última função receberá os seguintes argumentos:

• detalhes da imagem - dict;
• larguras alvo - lista (padrão: [256, 512, 1024]);

Primeiro, criamos um dicionário para armazenar os detalhes de cada imagem, começando com a original:

images = {
    "original": {
        "filename": image_details["original_filename"],
        "path": image_details["path"],
        "width": image_details["width"],
        "height": image_details["height"]
    }
}

Em seguida, vamos percorrer as larguras alvo e criar as variações da imagem. Para cada loop, faremos o seguinte:

Definir a altura da imagem, com base na largura:

height = int(width * image_details["ratio"])

Para facilitar, criaremos uma variável com o tamanho da imagem:

size = f"{width}x{height}"

Como temos o objeto de imagem, podemos simplesmente redimensioná-lo várias vezes:

resized_image = image_details["image_object"].resize((width, height))

Agora definimos um nome de arquivo para esta imagem:

filename = f"{image_details['filename_without_ext']}_{size}.jpg"

Então salvamos o arquivo:

resized_image.save(filename)

E, por fim, salvamos os detalhes desta imagem:

from pathlib import Path

images[size] = {
    "filename": filename,
    "path": Path(filename),
    "width": width,
    "height": height
}

Depois do loop for, retornamos as imagens:

return images

Organizando tudo

Agora que temos todas as funções de que precisamos, podemos organizá as chamadas em um script que chamará essas funções na ordem correta e gerará as imagens.

if __name__ == '__main__':
    # Prompt for the image
    img_prompt = "A cute corgi dog in a space suit, floating in space, and trying to reach a tasty treat."

    print("Generating image...")
    generated_image_url = generate_image(img_prompt)

    print("Downloading generated image...")
    original_image_details = download_image(generated_image_url, "doggo_in_space")

    # Define the desired target widths
    desired_target_widths = [780, 500, 342, 185, 154, 92]

    print("Creating image variants...")
    generated_images_details = create_variants(
        image_details=original_image_details,
        target_widths=desired_target_widths
    )

    print("Image processing completed.")
    print("Generated image details:")
    pprint(generated_images_details)

Você pode alterar o script e pedir ao usuário o prompt, ou alterar algumas configurações, etc.

Exemplo

No script, usei o seguinte prompt: A cute corgi dog in a space suit, floating in space, and trying to reach a tasty treat.

Este foi o resultado (que me custou USD$0,08) 🐶:

Se você quiser o script completo, você pode obter uma versão completa (incluindo as imagens resultantes) funcionando (sem a chave da API) aqui. Espero ter ajudado! 🙂