Generate Image from Text using Together API Generate Image

txt2img.py

#!/usr/bin/env python3
"""
Generate Image from Text using Together API Generate Image 

## requirements.txt
click
requests
pillow
piexif
"""
import click
import requests
import base64
import sys
import os
from PIL import Image, PngImagePlugin
from io import BytesIO
import json
import random
import piexif
import piexif.helper

TOGETHER_API_KEY = os.environ.get('TOGETHER_API_KEY', None)
if not TOGETHER_API_KEY:
    print("TOGETHER_API_KEY environment variable is required.", file=sys.stderr)
    raise SystemExit(1)

def save_image_with_metadata(image, output, metadata):
    """
    Save the image with metadata. Handles both PNG and JPG.
    For JPG, metadata is embedded into the Exif.
    """
    if output.lower().endswith('.png'):
        pnginfo = PngImagePlugin.PngInfo()
        for key, value in metadata.items():
            pnginfo.add_text(key, json.dumps(value, indent=2))
        image.save(output, "PNG", pnginfo=pnginfo)
    elif output.lower().endswith('.jpg') or output.lower().endswith('.jpeg'):
        exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}, "thumbnail": None}
        # Embedding metadata as a JSON string in the UserComment field
        exif_dict["Exif"][piexif.ExifIFD.UserComment] = piexif.helper.UserComment.dump(json.dumps(metadata, indent=2), encoding="unicode")
        exif_bytes = piexif.dump(exif_dict)
        image.save(output, "JPEG", exif=exif_bytes)
    else:
        raise ValueError("Unsupported file format. Please use PNG or JPG.")
    click.echo(f"Image saved to {output} with embedded JSON metadata.", err=True)

@click.command()
@click.option('--output', help='Save as jpg or png file')
@click.option('--model', default="SG161222/Realistic_Vision_V3.0_VAE", help='Model name')
@click.option('--negative_prompt', default=None, help='Negative prompt')
@click.option('--width', default=512, help='Width')
@click.option('--height', default=512, help='Height')
@click.option('--steps', default=20, help='Steps')
@click.option('--n', default=1, help='N')
@click.option('--seed', default=random.randint(1, 1024 * 64), help='Seed')
@click.argument('prompt', nargs=-1, required=True)
def main(**kwargs):
    """
    Make a POST request to an endpoint with individual data attributes, with prompt as optional remaining args.
    """

    output = kwargs.pop('output', None)
    if not output or (not output.lower().endswith('.png') and not output.lower().endswith('.jpg') and not output.lower().endswith('.jpeg')):
        print("Output file must be specified with a .png or .jpg extension.", file=sys.stderr)
        raise SystemExit(1)

    # Extract prompt from kwargs and join it into a single string
    prompt_str = ' '.join(kwargs.pop('prompt'))
    kwargs['prompt'] = prompt_str  # Update kwargs with the joined prompt string

    # strip empty values
    params = {k: v for k, v in kwargs.items() if v is not None}

    endpoint = 'https://api.together.xyz/v1/completions'
    
    res = requests.post(endpoint, json=params, headers={
        "Authorization": f"Bearer {TOGETHER_API_KEY}",
    })
    
    click.echo(f"JSON: {params}", err=True)
    response_json = res.json()
    
    if response_json.get("choices") and len(response_json["choices"]) > 0:
        image_base64 = response_json["choices"][0].get("image_base64")
        if image_base64:
            image_data = base64.b64decode(image_base64)
            image = Image.open(BytesIO(image_data))
            
            # Embed JSON payload as metadata
            metadata = {"params": params}
            save_image_with_metadata(image, output, metadata)
        else:
            click.echo("No image data found.", err=True)
    else:
        click.echo("No choices found in the response.", err=True)

if __name__ == "__main__":
    main()