Skip to content

Instantly share code, notes, and snippets.

@ddrscott
Last active February 20, 2024 22:20
Show Gist options
  • Save ddrscott/14619e0459693aed1f2fde433b005c43 to your computer and use it in GitHub Desktop.
Save ddrscott/14619e0459693aed1f2fde433b005c43 to your computer and use it in GitHub Desktop.
Generate Image from Text using Together API Generate Image
#!/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()
@ddrscott
Copy link
Author

squirrel-50

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment