OpenAI DALL-E Image Generator Script: Dynamic Model Selection, Custom Prompts, and Image Download

openai_dalle_generator.py

"""
Usage Guidelines and Command Line Arguments:
    This script generates images using OpenAI's DALL-E model based on a given prompt.
    To use this script, provide the model version (dall-e-2 or dall-e-3), the prompt for the image, 
    and optionally the size, quality, number of images, and a flag to download the images.

1. model: The version of DALL-E model to use. 
   - Possible values: 'dall-e-2', 'dall-e-3'
   - Example: --model dall-e-2

2. prompt: The prompt based on which the image will be generated. 
   - Format: A string describing the image.
   - Example: "a sunny beach"

3. --size: The size of the generated image (optional). 
   - Format: WIDTHxHEIGHT (in pixels)
   - Default: "1024x1024"
   - Example: --size "800x800"

4. --quality: The quality of the generated image (optional).
   - Possible values: "standard", "hd"
   - Default: "standard"
   - Example: --quality "hd"

5. --n: The number of images to generate (optional).
   - Format: An integer (For DALL-E v2, the maximum is 10; For DALL-E v3, it's always 1)
   - Default: 1
   - Example: --n 5

6. --download: Flag to download the generated images (optional).
   - If included, images will be downloaded to a 'generated_images' directory.
   - Example: --download

Example Commands:
python3 openai_dalle_generator.py dall-e-2 "a sunny beach"
python3 openai_dalle_generator.py dall-e-3 "a futuristic city" --size "800x800"
python3 openai_dalle_generator.py dall-e-2 "a bowl of fruit" --n 5
python3 openai_dalle_generator.py dall-e-3 "an old library" --quality "hd"
python3 openai_dalle_generator.py dall-e-2 "a space station" --download
python3 openai_dalle_generator.py dall-e-3 "a robot playing guitar" --size "1024x1024" --download
python3 openai_dalle_generator.py dall-e-2 "mountain landscape" --n 3 --quality "standard"
python3 openai_dalle_generator.py dall-e-3 "a cat in a superhero costume" --size "600x600" --quality "hd" --download
python3 openai_dalle_generator.py dall-e-2 "an underwater city" --n 10 --download
python3 openai_dalle_generator.py dall-e-3 "a fantasy dragon" --quality "standard"

Success Response Structure:

When the script successfully generates images, the output will be a JSON-formatted string containing a list of objects. Each object in the list represents one generated image and may contain the following keys:

1. 'image_url': The URL of the generated image.
   - Type: String
   - Description: Direct link to the image generated by the DALL-E model based on the provided prompt.

2. 'file_path' (optional): The absolute file path of the downloaded image.
   - Type: String
   - Description: Present only if the '--download' flag is used. It provides the absolute path where the generated image is saved locally.

Example of a Success Response:
[
    {
        "image_url": "https://example.com/generated_image_1.jpg",
        "file_path": "/absolute/path/to/generated_images/image_1.jpg"
    },
    {
        "image_url": "https://example.com/generated_image_2.jpg"
        // 'file_path' will be absent if the '--download' flag was not used
    }
    // Additional objects for each generated image
]

Note: The 'file_path' key appears only when the image is successfully downloaded to the local system. If the download fails, the 'file_path' key will either be absent or will indicate a failure message.
"""


import os
import json
import requests
import random
import string
import argparse
from openai import OpenAI

# Function to generate images using OpenAI's DALL-E model.
# Parameters:
# model: The version of DALL-E model to use ('dall-e-2' or 'dall-e-3').
# prompt: The prompt based on which the image will be generated.
# size: The size of the generated image.
# quality: The quality of the generated image.
# n: The number of images to generate.
def generate_image(model, prompt, size, quality, n):
    client = OpenAI(api_key="api_key") # Replace with your actual API key
    response = client.images.generate(
        model=model,
        prompt=prompt,
        size=size,
        quality=quality,
        n=n
    )

    if response.data:
        return [{'image_url': img.url} for img in response.data]
    else:
        return []

# Function to download the image from the given URL.
# Parameters:
# image_url: The URL of the image to download.
# prompt: The prompt used to generate the image (used for naming the file).
def download_image(image_url, prompt):
    # Convert prompt to snake case for file naming
    snake_case_prompt = "_".join(prompt.lower().split())
    # Generate a random string for unique file naming
    random_string = ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))
    filename = f"{snake_case_prompt}_{random_string}.jpg"

    # Create directory if it doesn't exist
    if not os.path.exists('generated_images'):
        os.makedirs('generated_images')

    # Generate absolute path for the image file
    filepath = os.path.abspath(os.path.join('generated_images', filename))

    # Download and save the image
    response = requests.get(image_url)
    if response.status_code == 200:
        with open(filepath, 'wb') as file:
            file.write(response.content)
        return filepath
    else:
        return None

# Setting up command-line argument parsing
parser = argparse.ArgumentParser(description='Generate images using DALL-E')
parser.add_argument('model', choices=['dall-e-2', 'dall-e-3'], help='Model to use (dall-e-2 or dall-e-3)')
parser.add_argument('prompt', type=str, help='Prompt for generating the image')
parser.add_argument('--size', type=str, default="1024x1024", help='Size of the image (default: 1024x1024)')
parser.add_argument('--quality', type=str, default="standard", help='Quality of the image (default: standard)')
parser.add_argument('--n', type=int, default=1, help='Number of images to generate (default: 1)')
parser.add_argument('--download', action='store_true', help='Download images to local directory')
args = parser.parse_args()

# Validate number of images based on the model version
if args.model == 'dall-e-2' and args.n > 10:
    print("For DALL-E v2, the maximum number of images is 10.")
    exit()
elif args.model == 'dall-e-3':
    args.n = 1  # For DALL-E v3, number of images is always 1

# Generate images based on the provided arguments
image_data = generate_image(args.model, args.prompt, args.size, args.quality, args.n)

# Download images if the download flag is set
if args.download:
    for data in image_data:
        path = download_image(data['image_url'], args.prompt)
        # Add file path to the image data
        data['file_path'] = path if path else "Failed to download"

# Print the image data in JSON format
print(json.dumps(image_data))