import os
import torch
import numpy as np
import cv2
from PIL import Image
from argparse import ArgumentParser
from model import DiffMorpherPipeline

parser = ArgumentParser()
parser.add_argument(
    "--model_path", type=str, default="stabilityai/stable-diffusion-2-1-base",
    help="Pretrained model to use (default: %(default)s)"
)
parser.add_argument(
    "--image_path_0", type=str, default="",
    help="Path of the first image (default: %(default)s)")
parser.add_argument(
    "--prompt_0", type=str, default="",
    help="Prompt of the second image (default: %(default)s)")
parser.add_argument(
    "--image_path_1", type=str, default="",
    help="Path of the first image (default: %(default)s)")
parser.add_argument(
    "--prompt_1", type=str, default="",
    help="Prompt of the second image (default: %(default)s)")
parser.add_argument(
    "--output_path", type=str, default="./results",
    help="Path of the output image (default: %(default)s)"
)
parser.add_argument(
    "--save_lora_dir", type=str, default="./lora",
    help="Path of the output lora directory (default: %(default)s)"
)
parser.add_argument(
    "--load_lora_path_0", type=str, default="",
    help="Path of the lora directory of the first image (default: %(default)s)"
)
parser.add_argument(
    "--load_lora_path_1", type=str, default="",
    help="Path of the lora directory of the second image (default: %(default)s)"
)
parser.add_argument(
    "--use_adain", action="store_true",
    help="Use AdaIN (default: %(default)s)"
)
parser.add_argument(
    "--use_reschedule",  action="store_true",
    help="Use reschedule sampling (default: %(default)s)"
)
parser.add_argument(
    "--lamb",  type=float, default=0.6,
    help="Lambda for self-attention replacement (default: %(default)s)"
)
parser.add_argument(
    "--fix_lora_value", type=float, default=None,
    help="Fix lora value (default: LoRA Interp., not fixed)"
)
parser.add_argument(
    "--save_inter", action="store_true",
    help="Save intermediate results (default: %(default)s)"
)
parser.add_argument(
    "--num_frames", type=int, default=16,
    help="Number of frames to generate (default: %(default)s)"
)
parser.add_argument(
    "--duration", type=int, default=100,
    help="Duration of each frame (default: %(default)s ms)"
)
parser.add_argument(
    "--no_lora", action="store_true"
)

args = parser.parse_args()

os.makedirs(args.output_path, exist_ok=True)
pipeline = DiffMorpherPipeline.from_pretrained(
    args.model_path, torch_dtype=torch.float32)
pipeline.to("cuda")
images = pipeline(
    img_path_0=args.image_path_0,
    img_path_1=args.image_path_1,
    prompt_0=args.prompt_0,
    prompt_1=args.prompt_1,
    save_lora_dir=args.save_lora_dir,
    load_lora_path_0=args.load_lora_path_0,
    load_lora_path_1=args.load_lora_path_1,
    use_adain=args.use_adain,
    use_reschedule=args.use_reschedule,
    lamd=args.lamb,
    output_path=args.output_path,
    num_frames=args.num_frames,
    fix_lora=args.fix_lora_value,
    save_intermediates=args.save_inter,
    use_lora=not args.no_lora
)
images[0].save(f"{args.output_path}/output.gif", save_all=True,
               append_images=images[1:], duration=args.duration, loop=0)