import torch
import os
import random
from diffusers import FluxTransformer2DModel
from diffusers.pipelines import FluxKontextPipeline
from diffusers.utils import load_image
from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint

from customization_kontext.evaluate import Evaluator
from customization_kontext.utils import parse_args_and_config
from customization_kontext.utils import * 



def main():
    # Load and overwrite config
    args, config = parse_args_and_config()
    test_config = config["test"]
    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))
    pipe = FluxKontextPipeline.from_pretrained(config["flux_path"], torch_dtype=torch.bfloat16)
    if test_config.get("ckpt_path"):
        print(f"Loading custom transformer weights from: {test_config['ckpt_path']}")
        state_dict = get_fp32_state_dict_from_zero_checkpoint(test_config['ckpt_path'])
        missing_keys, unexpected_keys = pipe.transformer.load_state_dict(state_dict, strict=True)
        print("Successfully loaded custom transformer weights with strict=True.")
        # 如果 strict=True 成功，missing_keys 和 unexpected_keys 应该都是空的
        if missing_keys:
            print(f"Warning: Missing keys in model: {missing_keys}")
        if unexpected_keys:
            print(f"Warning: Unexpected keys in checkpoint: {unexpected_keys}")
    
    pipe.to("cuda")
    if test_config.get("image"):
        os.makedirs(test_config['save_path'], exist_ok=True)
        input_image = load_image(test_config.get("image"))
        torch.manual_seed(42)
        seed = [42]
        seed.extend([random.randint(1, 10000) for _ in range(test_config['gen_num']-1)])
        for i in range(test_config['gen_num']):
            torch.manual_seed(seed[i])
            result_img = pipe(
                image=input_image,
                prompt=test_config.get("prompt"),
                num_inference_steps=test_config['num_inference_steps']
            ).images[0]
            output_name = generate_image_name(test_config.get("image"), test_config.get("prompt"), i)
            result_img.save(os.path.join(test_config['save_path'], output_name))
    else:
        # Evaluate and save 
        evaluator = Evaluator(
            test_config["json_file"],
            test_config['is_combine'],
            test_config['target_size'],
            gen_num=test_config['gen_num'],
            num_inference_steps=test_config['num_inference_steps']
        )
        evaluator.evaluate(
            pipe,
            test_config['save_path'], 
        )


if __name__ == "__main__":
    main()
