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