Generate an image based on a seed image and prompt. TODO Implement this in ytffmpeg.

gen-img.py

#!/usr/bin/env python3
'''
Run me to generate a thumbnail.
'''

import os, sys, io
from glob import glob
from dotenv import load_dotenv

from google import genai
from PIL import Image
from langchain_core.runnables import RunnableLambda

from subprocess import Popen, PIPE
from kizano import getLogger
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>
'''

load_dotenv()

# Choose the Google image model. You can override with GEMINI_IMAGE_MODEL in env.
# MODEL_ID = os.getenv("GEMINI_IMAGE_MODEL", "imagen-3.0-generate-002")  # or "imagen-3.0-generate-001"
MODEL_ID = os.getenv("GEMINI_IMAGE_MODEL", "gemini-2.5-flash-image")
client = genai.Client()

def edit_image_to_file(params: dict) -> str:
    """
    params:
      - prompt: str
      - image_path: str, path to seed image
      - output_path: str, path to save the edited image
      - aspect_ratio: optional, like "9:16", "1:1", "16:9"
      - mask_path: optional, supply a mask for inpainting
    returns:
      - path to the saved image file
    """
    log.info('Editing image...')
    prompt = params["prompt"]
    image_path = params["image_path"]
    output_path = params["output_path"]

    src_img = Image.open(image_path).convert("RGBA")
    log.info('Image loaded.')

    result = client.models.generate_content(
        model=MODEL_ID,
        contents=[prompt, src_img],
    )

    image_parts = [
        part.inline_data.data
        for part in result.candidates[0].content.parts
        if part.inline_data
    ]

    # Most SDK builds return a PIL.Image with .save available
    if image_parts:
        image = Image.open(io.BytesIO(image_parts[0]))
        image.save(output_path)
    return output_path

def main() -> int:
    log.info('hi')
    line1 = os.environ['LINE1']
    line2 = os.environ['LINE2']
    svg = TEMPLATE_SVG % {'line1': line1, 'line2': line2}
    template = 'build/thumbnail.png'
    # Create the thumbnail from SVG using ImageMagick
    log.info('Creating thumbnail...')
    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}", file=sys.stderr)
        return 1

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

    # Generate output filename based on input
    thumbnail = 'thumbnail.png'
    log.info(f'Converted SVG written to {thumbnail} as PNG')

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

---

{content}

---
Can you take the input 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
Thanks!
"""
    log.info('Invoking image-gen chain...')
    # LangChain runnable
    chain = RunnableLambda(edit_image_to_file)
    output = chain.invoke({
        "prompt": prompt,
        "image_path": template,
        "output_path": thumbnail,
        "aspect_ratio": "9:16"
    })
    log.info(f"Image saved to: {output}")
    return 0

if __name__ == "__main__":
    sys.exit(main())