| """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)) |
|
|
