import argparse
from model import models
from multit2i import (load_models, infer_fn, infer_rand_fn, save_gallery,
    change_model, warm_model, get_model_info_md, loaded_models,
    get_positive_prefix, get_positive_suffix, get_negative_prefix, get_negative_suffix,
    get_recom_prompt_type, set_recom_prompt_preset, randomize_seed, translate_to_en)

MAX_SEED = 2**32-1
load_models(models)

# Define a function to handle image generation
def generate_images(args):
    model_name = args.model
    prompt = args.prompt
    neg_prompt = args.neg_prompt
    width = args.width
    height = args.height
    steps = args.steps
    cfg = args.cfg
    seed = args.seed
    image_num = args.image_num

    positive_prefix = args.positive_prefix
    positive_suffix = args.positive_suffix
    negative_prefix = args.negative_prefix
    negative_suffix = args.negative_suffix

    if seed == -1:
        seed = randomize_seed()

    # Infer function
    for i in range(image_num):
        output_image = infer_fn(model_name, prompt, neg_prompt, height, width, steps, cfg, seed,
                                positive_prefix, positive_suffix, negative_prefix, negative_suffix)
        save_gallery(output_image, "gallery", f"output_image_{i}.png")

# CLI Arguments parsing
def main():
    parser = argparse.ArgumentParser(description="Custom Image Generation CLI")
    parser.add_argument("--model", type=str, required=True, help="Model name to use for generation.")
    parser.add_argument("--prompt", type=str, required=True, help="Text prompt for the image.")
    parser.add_argument("--neg_prompt", type=str, default="", help="Negative prompt for the image.")
    parser.add_argument("--width", type=int, default=0, help="Width of the image.")
    parser.add_argument("--height", type=int, default=0, help="Height of the image.")
    parser.add_argument("--steps", type=int, default=0, help="Number of inference steps.")
    parser.add_argument("--cfg", type=float, default=0, help="Guidance scale.")
    parser.add_argument("--seed", type=int, default=-1, help="Seed for randomization.")
    parser.add_argument("--image_num", type=int, default=1, help="Number of images to generate.")
    
    # Prefix and Suffix options
    parser.add_argument("--positive_prefix", nargs="*", default=[], help="Positive prompt prefix.")
    parser.add_argument("--positive_suffix", nargs="*", default=["Common"], help="Positive prompt suffix.")
    parser.add_argument("--negative_prefix", nargs="*", default=[], help="Negative prompt prefix.")
    parser.add_argument("--negative_suffix", nargs="*", default=["Common"], help="Negative prompt suffix.")
    
    args = parser.parse_args()
    generate_images(args)

if __name__ == "__main__":
    main()