#!/usr/bin/env python3
# Copyright (C) 2025-2026 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import argparse

import numpy as np
import openvino as ov
import openvino_genai

from pathlib import Path

from PIL import Image


def pipeline_export_import(root_dir: Path):
    pipe = openvino_genai.Text2ImagePipeline(root_dir, "CPU")
    pipe.export_model(root_dir / "exported")
    # pipeline models are exported to dedicated subfolders
    # for stable diffusion xl:
    # exported/
    # ├── text_encoder/
    # │   └── openvino_model.blob
    # ├── text_encoder_2/
    # │   └── openvino_model.blob
    # ├── unet/
    # │   └── openvino_model.blob
    # └── vae_decoder/
    #     └── openvino_model.blob

    # during import, specify blob_path property to point to the exported model location
    imported_pipe = openvino_genai.Text2ImagePipeline(root_dir, "CPU", blob_path=root_dir / "exported")


def dedicated_models_export_import(root_dir: Path):
    blob_path = root_dir / "blobs"
    device = "CPU"

    text_encoder = openvino_genai.CLIPTextModel(root_dir / "text_encoder", device)
    text_encoder.export_model(blob_path / "text_encoder")

    text_encoder_2 = openvino_genai.CLIPTextModelWithProjection(root_dir / "text_encoder_2", device)
    text_encoder_2.export_model(blob_path / "text_encoder_2")

    unet = openvino_genai.UNet2DConditionModel(root_dir / "unet", device)
    unet.export_model(blob_path / "unet")

    vae = openvino_genai.AutoencoderKL(root_dir / "vae_decoder")
    vae.compile(device)
    vae.export_model(blob_path)
    # AutoencoderKL can be composed with decoder and encoder models
    # exported/
    # └── vae_decoder/
    #     └── openvino_model.blob
    # └── vae_encoder/
    #     └── openvino_model.blob

    pipe = openvino_genai.Text2ImagePipeline.stable_diffusion_xl(
        scheduler=openvino_genai.Scheduler.from_config(root_dir / "scheduler" / "scheduler_config.json"),
        clip_text_model=openvino_genai.CLIPTextModel(
            root_dir / "text_encoder", device, blob_path=blob_path / "text_encoder"
        ),
        clip_text_model_with_projection=openvino_genai.CLIPTextModelWithProjection(
            root_dir / "text_encoder_2", device, blob_path=blob_path / "text_encoder_2"
        ),
        unet=openvino_genai.UNet2DConditionModel(root_dir / "unet", device, blob_path=blob_path / "unet"),
        vae=openvino_genai.AutoencoderKL(root_dir / "vae_decoder", device, blob_path=blob_path),
    )


def tensor_models_export_import(root_dir: Path):
    blob_path = root_dir / "exported_tensors"
    device = "CPU"

    pipe = openvino_genai.Text2ImagePipeline(root_dir, device)
    pipe.export_model(blob_path)

    def read_blob_tensor(model_folder: str):
        blob = np.fromfile(blob_path / model_folder / "openvino_model.blob", dtype=np.uint8)
        return ov.Tensor(blob)

    text_encoder = openvino_genai.CLIPTextModel(
        read_blob_tensor("text_encoder"),
        openvino_genai.CLIPTextModel.Config(root_dir / "text_encoder" / "config.json"),
        openvino_genai.Tokenizer(root_dir / "tokenizer"),
        device,
    )
    text_encoder_2 = openvino_genai.CLIPTextModelWithProjection(
        read_blob_tensor("text_encoder_2"),
        openvino_genai.CLIPTextModelWithProjection.Config(root_dir / "text_encoder_2" / "config.json"),
        openvino_genai.Tokenizer(root_dir / "tokenizer_2"),
        device,
    )
    vae = openvino_genai.AutoencoderKL(
        read_blob_tensor("vae_decoder"),
        openvino_genai.AutoencoderKL.Config(root_dir / "vae_decoder" / "config.json"),
        device,
    )
    unet = openvino_genai.UNet2DConditionModel(
        read_blob_tensor("unet"),
        openvino_genai.UNet2DConditionModel.Config(root_dir / "unet" / "config.json"),
        vae.get_vae_scale_factor(),
        device,
    )

    imported_pipe = openvino_genai.Text2ImagePipeline.stable_diffusion_xl(
        scheduler=openvino_genai.Scheduler.from_config(root_dir / "scheduler" / "scheduler_config.json"),
        clip_text_model=text_encoder,
        clip_text_model_with_projection=text_encoder_2,
        unet=unet,
        vae=vae,
    )


def export_import_with_reshape(root_dir: Path, prompt: str):
    device = "CPU"

    width = 512
    height = 512
    number_of_images_to_generate = 1
    number_of_inference_steps_per_image = 20

    pipe = openvino_genai.Text2ImagePipeline(root_dir)
    pipe.reshape(1, height, width, pipe.get_generation_config().guidance_scale)
    pipe.compile(device)
    pipe.export_model(root_dir / "exported")

    imported_pipe = openvino_genai.Text2ImagePipeline(root_dir, device, blob_path=root_dir / "exported")

    # update generation config according to the new shape parameters
    config = imported_pipe.get_generation_config()
    config.height = height
    config.width = width
    config.num_images_per_prompt = number_of_images_to_generate
    imported_pipe.set_generation_config(config)

    for imagei in range(0, number_of_images_to_generate):
        image_tensor = imported_pipe.generate(
            prompt,
            num_inference_steps=number_of_inference_steps_per_image,
        )

        image = Image.fromarray(image_tensor.data[0])
        image.save("image_" + str(imagei) + ".bmp")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("model_dir")
    parser.add_argument("prompt")

    args = parser.parse_args()

    root_dir = Path(args.model_dir)

    pipeline_export_import(root_dir)
    dedicated_models_export_import(root_dir)
    tensor_models_export_import(root_dir)
    export_import_with_reshape(root_dir, args.prompt)


if "__main__" == __name__:
    main()