Generate an image using OpenAI's Image generator.

oai-gen-img.py

#!/usr/bin/env python3
'''
Image generation test using OpenAI's Image model and code.
Usage:

LINE1=Sample LINE2=Video ./oai-gen-img.py

Ouputs `./thumbnail.png` based on `build/*.txt` and input lines.
'''

from dotenv import load_dotenv
load_dotenv()

import os, base64
from kizano import getLogger
from openai import OpenAI
from glob import glob
from subprocess import Popen, PIPE

log = getLogger(__name__)

TEMPLATE_SVG = '''<svg width="585" height="1024" viewBox="0 0 585 1024" xmlns="http://www.w3.org/2000/svg">
  <style>
    .outlined {
      fill: #04547a;
      stroke: #e9e43c;
      stroke-width: 18;
      stroke-linejoin: round;
      stroke-linecap: round;
      paint-order: stroke fill;
      text-anchor: middle;
      dominant-baseline: middle;
      font-family: DejaVu Sans;
      font-weight: 900;
    }
  </style>

  <!-- Center the group, then position lines above and below center -->
  <g transform="translate(292.5, 512)">
    <text class="outlined" font-size="100" y="-84">%(line1)s</text>
    <text class="outlined" font-size="100" y="84">%(line2)s</text>
  </g>
</svg>
'''

IMAGE_PROMPT = """I'm trying to create an interesting thumbnail for my TikTok video.
Here's the content of the video:

---

%(content)s

---
Can you take the input image and the content and help me to craft an interesting thumbnail based on the context provided, please?
The aspect ratio needs to be 9:16. The size can be no greater than 585x1024.
Don't include the full context of the video in the thumbnail. It needs to have elements of the seed image and use the content of the video to help you craft an interesting background.
Thanks!
"""

def genTemplate(template: str) -> None:
    line1 = os.environ['LINE1']
    line2 = os.environ['LINE2']
    log.info(f'Generating SVG->PNG from "{line1} / {line2}"')
    svg = TEMPLATE_SVG % {'line1': line1, 'line2': line2}
    # Create the thumbnail from SVG using ImageMagick
    p = Popen(['convert', '-transparent', '#FFFFFF', '-', template], stdin=PIPE)
    p.communicate(input=svg.encode('utf-8'))

    if p.returncode != 0:
        log.error(f"Error: convert command failed with return code {p.returncode}")
        return 1
    log.info(f'Converted SVG written to {template} as PNG')


def genThumbnail(imgpath: str, thumbnail: str, prompt: str) -> None:
    'Generate the thumbnail for the video based on inputs.'
    image = OpenAI().images.edit(
        model='gpt-image-1',
        image=[open(imgpath, "rb")],
        prompt=prompt,
        n=1,
    )
    open(thumbnail, 'wb').write(base64.b64decode(image.data[0].b64_json))

def main():
    log.info('Creating thumbnail...')
    template = 'build/thumbnail.png'
    thumbnail = 'thumbnail.png'
    genTemplate(template)

    content_txt = glob('build/*.txt').pop(0)
    content = open(content_txt).read()

    # Generate output filename based on input
    prompt = IMAGE_PROMPT % {'content': content}
    genThumbnail(template, thumbnail, prompt)
    log.info(f"Image saved to: {thumbnail}")
    return 0

if __name__ == "__main__":
    main()