File size: 2,175 Bytes
4f03424 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | """Generate an RGBA image from a local Image21-MLX or original BF16 model."""
import argparse
import time
from pathlib import Path
import mlx.core as mx
from scripts.runtime import load,generate
from scripts.common import write_json
def main():
ap=argparse.ArgumentParser()
ap.add_argument('--model',type=Path,required=True)
ap.add_argument('--prompt',required=True)
ap.add_argument('--output',type=Path,required=True)
ap.add_argument('--input',nargs='+')
ap.add_argument('--seed',type=int)
ap.add_argument('--source-seed',type=int)
ap.add_argument('--steps',type=int,default=40)
ap.add_argument('--width',type=int)
ap.add_argument('--height',type=int)
ap.add_argument('--resolution',type=int,default=1024)
args=ap.parse_args()
if args.output.exists(): raise FileExistsError(args.output)
if (args.width is None)!=(args.height is None): ap.error('Set width and height together')
if args.width is None:
if args.input:
from PIL import Image
from mlx_vlm.models.qwen_image.pipeline import _image_dimensions
with Image.open(args.input[-1]) as source:
args.width,args.height=_image_dimensions(args.resolution,source.width/source.height)
else: args.width=args.height=args.resolution
seed=args.seed if args.seed is not None else (1000042 if args.input else 42)
start=time.perf_counter(); pipe=load(args.model)
load_s=time.perf_counter()-start; mx.reset_peak_memory(); start=time.perf_counter()
img=generate(pipe,args.prompt,seed=seed,steps=args.steps,width=args.width,height=args.height,
inputs=args.input,source_seed=args.source_seed,resolution=args.resolution)
seconds=time.perf_counter()-start
args.output.parent.mkdir(parents=True,exist_ok=True); img.save(args.output)
row=dict(model=str(args.model),prompt=args.prompt,seed=seed,steps=args.steps,width=args.width,height=args.height,
mode=img.mode,pipeline_init_seconds=load_s,seconds=seconds,mlx_peak_gib=mx.get_peak_memory()/2**30)
write_json(args.output.with_suffix('.json'),row); print(row,flush=True)
if __name__=='__main__': main()
|