"""Local or Hub inference for Image21-INT4."""
import argparse
from pathlib import Path

import torch

from scripts.editing_protocol import DEFAULT_SEEDS, edit_dimensions, load_input, validate_seeds
from scripts.runtime import load_pipeline


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', default='models/int4')
    parser.add_argument('--prompt', required=True)
    parser.add_argument('--input', nargs='*', default=[])
    parser.add_argument('--output', default='output.png')
    parser.add_argument('--seed', type=int)
    parser.add_argument('--source-seed', type=int)
    parser.add_argument('--width', type=int)
    parser.add_argument('--height', type=int)
    parser.add_argument('--resolution', type=int, default=1024)
    parser.add_argument('--steps', type=int, default=40)
    parser.add_argument('--offload', choices=['auto', 'model', 'group', 'resident'], default='auto')
    parser.add_argument('--device', default=None)
    args = parser.parse_args()
    if (args.width is None) != (args.height is None):
        parser.error('Specify both --width and --height, or neither')
    if args.steps <= 0 or any(value <= 0 or value % 32 for value in
                             (args.width, args.height, args.resolution) if value is not None):
        parser.error('Positive steps and dimensions divisible by 32 are required')
    if Path(args.output).suffix.lower() != '.png':
        parser.error('Use a .png output so RGBA channels are retained')
    images = [load_input(path) for path in args.input]
    seed = args.seed if args.seed is not None else (DEFAULT_SEEDS[0] if images else 42)
    case = {'source_seed': args.source_seed} if images else {}
    try:
        validate_seeds([case], [seed])
    except ValueError as exc:
        parser.error(str(exc))
    if args.width is None:
        args.width, args.height = (edit_dimensions(images[-1].size, args.resolution)
                                   if images else (args.resolution, args.resolution))
    pipe = load_pipeline(args.model, offload=args.offload, device=args.device,
                        local_files_only=Path(args.model).exists())
    with torch.inference_mode():
        image = pipe(prompt=args.prompt, image=images or None,
                     width=args.width, height=args.height, output_resolution=args.resolution,
                     num_inference_steps=args.steps, true_cfg_scale=1.0, use_kv_cache=True,
                     generator=torch.Generator('cpu').manual_seed(seed)).images[0]
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    image.save(args.output)
    policy = pipe.image21_runtime
    print(f'{args.output} (seed={seed}, size={image.size}, mode={image.mode}, '
          f'device={policy["device"]}, offload={policy["offload"]})')


if __name__ == '__main__':
    main()