grimdarkinator

grimdarkinate.py

import base64
from typing import Union, Literal, Optional
import openai
import requests
import os, io
from PIL import Image
from openai import OpenAI
import uuid
import argparse

def resize_and_encode(image_path_or_url: str) -> str:
    # Check if it is a URL
    if image_path_or_url.startswith('http://') or image_path_or_url.startswith('https://'):
        # Get the image from the URL
        response = requests.get(image_path_or_url)
        image_data = response.content
        img = Image.open(io.BytesIO(image_data))
    else:
        # Read image from local file
        img = Image.open(open(image_path_or_url, 'rb'))
    # Resize the image to fit within 1024x1024, keeping the aspect ratio
    img.thumbnail((1024, 1024))
    # Encode the image to base64
    buffer = io.BytesIO()
    img.save(buffer, format="JPEG")
    img_str = base64.b64encode(buffer.getvalue()).decode()
    return img_str

def describe_image(img: str, max_tokens: int = 512) -> str:
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {os.environ['OPENAI_API_KEY']}"
        }

    payload = {
        "model": "gpt-4-vision-preview",
        "messages": [{
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "Please describe this image in as much detail as possible. "
                            "Focus on describing it in a way that would make it easy for a generative "
                            "model to replicate a similar image. If the image is of a person or people, "
                            "describe and characterise their appearance to facilitate a reproduction. "
                            "Do not try to identify the individual. If the image is of a place, describe "
                            "the place in as much detail as possible."
                    },
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{img}"}
                    }
                ]
            }],
        "max_tokens": max_tokens,
        }

    response = requests.post("https://api.openai.com/v1/chat/completions",
                             headers=headers,
                             json=payload)

    try:
        return response.json()["choices"][0]["message"]["content"]
    except openai.OpenAIError as e:
        print(f"OpenAI API error {e.http_status}: {e.error}")



def transform_description(description: str) -> str:
    client = OpenAI()
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": "Your job is to analyze descriptions of images and change them to transpose the image into "
                           "the dark, futuristic aesthetic of the Warhammer 40k universe. You are given a description "
                           "of a scene, and your job is to change that scene according to the instructions, and give "
                           "a detailed description of the changed scene that would allow a generative model to create "
                           "an image.\n\n1. Replace clothing items with a mixture of futuristic and mediaeval looking "
                           "armour.\n2. Replace every-day objects held by characters with power swords, chainswords, "
                           "futuristic looking guns, staffs or magic items.\n3. Replace architecture with the "
                           "monumental, dark architecture common to the Warhammer 40k universe. \n4. Include things "
                           "like monumental gothic spaceships in the image description."
                },
            {
                "role": "user",
                "content": description}
            ],
        temperature=0.66,
        max_tokens=2048,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0
        )

    try:
        return response.choices[0].message.content
    except openai.OpenAIError as e:
        print(f"OpenAI API error {e.http_status}: {e.error}")


def generate_grimdark(description: str,
                     quality: Literal["hd", "standard"] = "standard",
                     return_url: bool = False,
                     filename: Optional[str] = None) -> Union[str, bytes]:
    client = OpenAI()
    response = client.images.generate(model="dall-e-3",
                                      prompt=f"Generate an image in the style of science fiction book cover "
                                             f"illustrations of the Warhammer 40k science fiction universe, in a dark,"
                                             f"moody fantasy science fiction style, based on the description provided "
                                             f"at the end of the prompt. The image should reflect the dystopian, "
                                             f"grimdark, gothic and militaristic aesthetics of the Warhammer 40k "
                                             f"universe. This universe melds science fiction and fantasy, and the "
                                             f"image should reflect this. "
                                             f"The background should be atmospheric, with monumental structures "
                                             f"hinting at a vast and advanced civilization. Use a palette that "
                                             f"emphasizes dark tones with strategic use of bright contrasts for "
                                             f"dramatic effect. The visual language should be similar to that of a "
                                             f"science fiction paperback cover illustration. "
                                             f"Here is the description of the subject of the image to be generated: "
                                             f"\n\n{description}\n\n",
                                      size="1024x1024",
                                      quality=quality,
                                      response_format="url" if return_url else "b64_json",
                                      n=1)

    if return_url:
        return response.data[0].url
    else:
        img_data = base64.b64decode(response.data[0].b64_json)

        if filename is None:
            return img_data
        else:
            with open(filename, 'wb') as f:
                f.write(img_data)
            return filename


def main():
    # Instantiate the parser
    parser = argparse.ArgumentParser(description='Just grimdark up my shit, bro.')

    # Add arguments
    parser.add_argument('image_path', type=str,
                        help='A path (local or URL) to the image file to process.')

    parser.add_argument('--max_tokens', type=int, default=512,
                        help='The maximum number of tokens for the descriptive step. Default is 512.')

    parser.add_argument('--quality', type=str, choices=["hd", "standard"], default='standard',
                        help='The quality of the generated image. Choices are "hd" and "standard". Default is "standard".')

    parser.add_argument('--return_url', action='store_true',
                        help='Whether to return the URL of the generated image. Default is False.')

    parser.add_argument('--filename', '-o', type=str,
                        help='The name of the file to save the generated image. If not set, the image will be saved to a random filename.')

    # Parse arguments
    args = parser.parse_args()

    # Process the image
    img_str = resize_and_encode(args.image_path)

    # Generate the description
    description = describe_image(img_str, args.max_tokens)
    print("IR: ", description, "\n\n")

    # Transform the description
    description = transform_description(description)
    print("Transformed IR: ", description, "\n\n")

    # Generate the image
    if args.return_url is True and args.filename is not None:
        raise ValueError("You cannot specify both --return_url and --filename.")

    if args.return_url is False and args.filename is None:
        # We generate a random uuid filename
        filename = f"{uuid.uuid4()}.png"
    else:
        filename = args.filename

    res = generate_grimdark(description, args.quality, args.return_url, filename)

    if args.return_url is True:
        print("Generated image: ", res)
    else:
        print(f"Generated image: saved to {filename}.")

if __name__ == "__main__":
    main()