# TextSynth text to image example
import sys
import argparse
import requests
import json
import time
import base64

parser = argparse.ArgumentParser(description='TextSynth text to image example')
parser.add_argument('--url', type=str, default = "https://api.textsynth.com", help="server URL")
parser.add_argument('--api_key', type=str, default = "")
parser.add_argument('--seed', type=int, default = None, help="seed (0 = random)")
parser.add_argument('--model', type=str, default = "stable_diffusion", help = "model name")
parser.add_argument('-o', type=str, default = "/tmp/out.jpg", help="output image")
parser.add_argument('--width', type=int, default = 512, help="generated image width")
parser.add_argument('--height', type=int, default = 512, help="generated image height")
parser.add_argument('--image', type=str, default = None, help="optional initial image filename (JPEG)")
parser.add_argument('--strength', type=float, default=0.5, help="strength, only meaningful when --image is used")
parser.add_argument('--image_count', type=int, default = 1, help="number of generate images")
parser.add_argument('prompt', type=str, help="image generation prompt")

args = parser.parse_args()

def make_request(path, json):
    response = requests.post(args.url + path, headers = { "Authorization": "Bearer " + args.api_key }, json = json)
    return response

def main():
    req = { "prompt": args.prompt, "image_count": args.image_count, "width": args.width, "height": args.height,
            "strength": args.strength }
    if args.seed != None:
        req["seed"] = args.seed
        
    if args.image != None:
        f = open(args.image, "rb")
        image_data = f.read()
        f.close()
        req["image"] = base64.b64encode(image_data).decode('ascii')
        
    result = make_request("/v1/engines/" + args.model + "/text_to_image", req);
    if result.status_code != 200:
        print("Request error:", result.text)
        sys.exit(1)

    resp = result.json()
    #    print(json.dumps(resp, indent=4, ensure_ascii=False))
    images = resp["images"]
    for i in range(len(images)):
        data = images[i]["data"]
        image = base64.b64decode(data)
        filename = args.o
        filename = filename.replace("%d", str(i))
        f = open(filename, "wb")
        f.write(image)
        f.close()
    
main()