"""Generate with Image 2.1 Calibrated FP8. Built with Qwen."""
import argparse, json, time
from pathlib import Path
import torch
from fp8_runtime import load_pipeline

@torch.inference_mode()
def main():
    ap=argparse.ArgumentParser()
    inputs=ap.add_mutually_exclusive_group(required=True)
    inputs.add_argument('--prompt')
    inputs.add_argument('--prompts-json', help='JSON list of prompt strings; one loaded pipeline serves the whole list')
    ap.add_argument('--base',default='Qwen/Qwen-Image-2.1')
    ap.add_argument('--quant',default=str(Path(__file__).parent/'transformer'))
    ap.add_argument('--width',type=int,default=2048)
    ap.add_argument('--height',type=int,default=2048)
    ap.add_argument('--steps',type=int,default=40)
    ap.add_argument('--seed',type=int,default=42)
    ap.add_argument('--output',default='output.png')
    ap.add_argument('--image', nargs='+', action='extend', metavar='PATH',
                    help='One or more local reference images for editing, in image 1, image 2, ... order. '
                         'Pass multiple paths or repeat --image. Omit for text-to-image generation.')
    ap.add_argument('--transparent', action='store_true',
                    help='Request native transparent RGBA output and require a PNG filename')
    ap.add_argument('--eager', action='store_true', help='Use the original uncompiled runtime')
    ap.add_argument('--warmup', action='store_true', help='Run a full untimed warmup; report its cost separately')
    ap.add_argument('--cache-prompts', action='store_true',
                    help='Reuse exact encoder outputs for repeated text-only prompts in this process')
    args=ap.parse_args()
    prompts=json.loads(Path(args.prompts_json).read_text()) if args.prompts_json else [args.prompt]
    if not isinstance(prompts,list) or not prompts or not all(isinstance(p,str) and p for p in prompts):
        ap.error('prompts-json must contain a nonempty list of prompt strings')
    if args.transparent:
        if Path(args.output).suffix.lower() != '.png':
            ap.error('--transparent requires --output ending in .png to preserve the alpha channel')
        from transparency import transparent_prompt
        prompts = [transparent_prompt(p) for p in prompts]
    kw={}
    if args.image:
        from PIL import Image
        reference_images=[]
        for filename in args.image:
            try:
                with Image.open(filename) as reference:
                    reference_images.append(reference.copy())
            except (OSError, ValueError) as exc:
                ap.error(f'cannot read --image {filename!r}: {exc}')
        kw['image']=reference_images[0] if len(reference_images)==1 else reference_images
    start=time.perf_counter();pipe=load_pipeline(args.base,args.quant)
    if not args.eager:
        from acceleration import accelerate_pipeline
        accelerate_pipeline(pipe)
    if args.cache_prompts:
        from prompt_cache import cache_prompt_embeddings
        cache_prompt_embeddings(pipe)
    torch.cuda.synchronize();load_seconds=time.perf_counter()-start
    def run(prompt,seed):
        return pipe(prompt=prompt,width=args.width,height=args.height,num_inference_steps=args.steps,
                    generator=torch.Generator('cuda').manual_seed(seed),**kw).images[0]
    warmup_seconds=0
    if args.warmup:
        start=time.perf_counter();run(prompts[0],args.seed);torch.cuda.synchronize()
        warmup_seconds=time.perf_counter()-start
    path=Path(args.output);path.parent.mkdir(parents=True,exist_ok=True)
    for i,prompt in enumerate(prompts):
        torch.cuda.synchronize();start=time.perf_counter();result=run(prompt,args.seed+i)
        torch.cuda.synchronize();seconds=time.perf_counter()-start
        destination=path if len(prompts)==1 else path.with_name(f'{path.stem}-{i:03d}{path.suffix or ".png"}')
        result.save(destination)
        from transparency import alpha_statistics
        print(json.dumps({'seconds':seconds,'output':str(destination),'width':result.width,'height':result.height,
                          'steps':args.steps,'eager':args.eager,'load_seconds':load_seconds,
                          'input_image_count':len(args.image or []),
                          'transparent_requested':args.transparent,**alpha_statistics(result),
                          'prompt_cache':getattr(pipe,'_image21_prompt_cache_stats',None),
                          'warmup_seconds':warmup_seconds,
                          'timing_note':'Generation includes any compilation on this call; excludes model load and PNG writing.'}),flush=True)

if __name__=='__main__':main()