"""Examples: python inference.py /path/to/model --task t2i --prompt 'A red mug'."""
import argparse
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM

parser=argparse.ArgumentParser()
parser.add_argument('model')
parser.add_argument('--task',choices=['t2i','edit','understand'],default='t2i')
parser.add_argument('--prompt',required=True)
parser.add_argument('--image')
parser.add_argument('--output',default='output.png')
parser.add_argument('--seed',type=int,default=42)
args=parser.parse_args()
model=AutoModelForCausalLM.from_pretrained(args.model,trust_remote_code=True,torch_dtype=torch.bfloat16).to('cuda').eval()
if args.task=='t2i':
    model.text_to_image(args.prompt,seed=args.seed).save(args.output)
elif args.task=='edit':
    if not args.image:parser.error('--image is required for editing')
    model.image_edit(args.image,args.prompt,seed=args.seed).save(args.output)
else:
    print(model.generate(args.prompt,image=args.image))