import json
import argparse
from gradio_client import Client


def main():
    parser = argparse.ArgumentParser(description="x")
    parser.add_argument('--model', '-m', type=str, default="red")
    parser.add_argument('--prompt', '-p', type=str, default="a girl in beijing")
    parser.add_argument('--api-url', type=str, default="http://127.0.0.1:7860/")
    parser.add_argument('--api-name', type=str, default="/predict")

    args = parser.parse_args()
    client = Client(args.api_url)
    output = client.predict(args.model, args.prompt, api_name=args.api_name)    

    result = {
        "model": args.model, 
        "prompt": args.prompt, 
        "output": output
    }
    
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()