Last active
February 20, 2024 22:20
-
-
Save ddrscott/14619e0459693aed1f2fde433b005c43 to your computer and use it in GitHub Desktop.
Generate Image from Text using Together API Generate Image
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment