DALLE OpenAI API

main.py

import argparse
import os
from typing import List

import openai

secret_key = os.environ.get("OPENAI_SECRET_KEY")
if not secret_key:
    raise ValueError("Please set the OPENAI_SECRET_KEY environment variable")


class DALLE:
    def __init__(self, api_key: str = secret_key):
        openai.api_key = api_key

    def create_images(self, prompt: str, n_images: int = 1, size: str = "1024x1024") -> List[str]:
        response = openai.Image.create(
            prompt=prompt,
            n=n_images,
            size=size
        )
        urls = []
        for image_data in response['data']:
            urls.append(image_data['url'])
        return urls


if __name__ == '__main__':
    print('> Getting prices...')
    args = argparse.ArgumentParser()
    args.add_argument("-t", '--text', type=str, default=None, help="Description of the image", dest="text",
                      required=True)
    options = args.parse_args()
    text = options.text

    image_urls = DALLE().create_images(text)
    print(image_urls)