#
# SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
Model definitions for UNet models.
"""

import torch

from demo_diffusion.dynamic_import import import_from_diffusers
from demo_diffusion.model import base_model, load, optimizer
from diffusers import StableDiffusionXLControlNetPipeline

# List of models to import from diffusers.models
models_to_import = [
    "ControlNetModel",
    "UNet2DConditionModel",
    "UNetSpatioTemporalConditionModel",
    "StableCascadeUNet",
]
for model in models_to_import:
    globals()[model] = import_from_diffusers(model, "diffusers.models")


def get_unet_embedding_dim(version, pipeline):
    if version in ("1.4", "dreamshaper-7"):
        return 768
    elif version in ("xl-1.0", "xl-turbo") and pipeline.is_sd_xl_base():
        return 2048
    elif version in ("cascade"):
        return 1280
    elif version in ("xl-1.0", "xl-turbo") and pipeline.is_sd_xl_refiner():
        return 1280
    elif pipeline.is_img2vid():
        return 1024
    else:
        raise ValueError(f"Invalid version {version} + pipeline {pipeline}")


class UNet2DConditionControlNetModel(torch.nn.Module):
    def __init__(self, unet, controlnets) -> None:
        super().__init__()
        self.unet = unet
        self.controlnets = controlnets

    def forward(self, sample, timestep, encoder_hidden_states, images, controlnet_scales, added_cond_kwargs=None):
        for i, (image, conditioning_scale, controlnet) in enumerate(zip(images, controlnet_scales, self.controlnets)):
            down_samples, mid_sample = controlnet(
                sample,
                timestep,
                encoder_hidden_states=encoder_hidden_states,
                controlnet_cond=image,
                return_dict=False,
                added_cond_kwargs=added_cond_kwargs,
            )

            down_samples = [down_sample * conditioning_scale for down_sample in down_samples]
            mid_sample *= conditioning_scale

            # merge samples
            if i == 0:
                down_block_res_samples, mid_block_res_sample = down_samples, mid_sample
            else:
                down_block_res_samples = [
                    samples_prev + samples_curr
                    for samples_prev, samples_curr in zip(down_block_res_samples, down_samples)
                ]
                mid_block_res_sample += mid_sample

        noise_pred = self.unet(
            sample,
            timestep,
            encoder_hidden_states=encoder_hidden_states,
            down_block_additional_residuals=down_block_res_samples,
            mid_block_additional_residual=mid_block_res_sample,
            added_cond_kwargs=added_cond_kwargs,
        )
        return noise_pred


class UNetModel(base_model.BaseModel):
    def __init__(
        self,
        version,
        pipeline,
        device,
        hf_token,
        verbose,
        framework_model_dir,
        fp16=False,
        int8=False,
        fp8=False,
        max_batch_size=16,
        text_maxlen=77,
        controlnets=None,
        do_classifier_free_guidance=False,
    ):

        super(UNetModel, self).__init__(
            version,
            pipeline,
            device=device,
            hf_token=hf_token,
            verbose=verbose,
            framework_model_dir=framework_model_dir,
            fp16=fp16,
            int8=int8,
            fp8=fp8,
            max_batch_size=max_batch_size,
            text_maxlen=text_maxlen,
            embedding_dim=get_unet_embedding_dim(version, pipeline),
        )
        self.subfolder = "unet"
        self.controlnets = load.get_path(version, pipeline, controlnets) if controlnets else None
        self.unet_dim = 4
        self.xB = 2 if do_classifier_free_guidance else 1  # batch multiplier

    def get_model(self, torch_inference=""):
        model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {}
        if self.controlnets:
            unet_model = UNet2DConditionModel.from_pretrained(
                self.path,
                subfolder=self.subfolder,
                use_safetensors=self.hf_safetensor,
                token=self.hf_token,
                **model_opts,
            ).to(self.device)
            cnet_model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
            controlnets = torch.nn.ModuleList(
                [ControlNetModel.from_pretrained(path, **cnet_model_opts).to(self.device) for path in self.controlnets]
            )
            # FIXME - cache UNet2DConditionControlNetModel
            model = UNet2DConditionControlNetModel(unet_model, controlnets)
        else:
            unet_model_dir = load.get_checkpoint_dir(
                self.framework_model_dir, self.version, self.pipeline, self.subfolder
            )
            if not load.is_model_cached(unet_model_dir, model_opts, self.hf_safetensor):
                model = UNet2DConditionModel.from_pretrained(
                    self.path,
                    subfolder=self.subfolder,
                    use_safetensors=self.hf_safetensor,
                    token=self.hf_token,
                    **model_opts,
                ).to(self.device)
                model.save_pretrained(unet_model_dir, **model_opts)
            else:
                print(f"[I] Load UNet2DConditionModel  model from: {unet_model_dir}")
                model = UNet2DConditionModel.from_pretrained(unet_model_dir, **model_opts).to(self.device)
            if torch_inference:
                model.to(memory_format=torch.channels_last)
        model = optimizer.optimize_checkpoint(model, torch_inference)
        return model

    def get_input_names(self):
        if self.controlnets is None:
            return ["sample", "timestep", "encoder_hidden_states"]
        else:
            return ["sample", "timestep", "encoder_hidden_states", "images", "controlnet_scales"]

    def get_output_names(self):
        return ["latent"]

    def get_dynamic_axes(self):
        xB = "2B" if self.xB == 2 else "B"
        if self.controlnets is None:
            return {
                "sample": {0: xB, 2: "H", 3: "W"},
                "encoder_hidden_states": {0: xB},
                "latent": {0: xB, 2: "H", 3: "W"},
            }
        else:
            return {
                "sample": {0: xB, 2: "H", 3: "W"},
                "encoder_hidden_states": {0: xB},
                "images": {1: xB, 3: "8H", 4: "8W"},
                "latent": {0: xB, 2: "H", 3: "W"},
            }

    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
        # WAR to enable inference for H/W that are not multiples of 16
        # If building with Dynamic Shapes: ensure image height and width are not multiples of 16 for ONNX export and TensorRT engine build
        if not static_shape:
            image_height = image_height - 8 if image_height % 16 == 0 else image_height
            image_width = image_width - 8 if image_width % 16 == 0 else image_width
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        (
            min_batch,
            max_batch,
            min_image_height,
            max_image_height,
            min_image_width,
            max_image_width,
            min_latent_height,
            max_latent_height,
            min_latent_width,
            max_latent_width,
        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
        if self.controlnets is None:
            return {
                "sample": [
                    (self.xB * min_batch, self.unet_dim, min_latent_height, min_latent_width),
                    (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
                    (self.xB * max_batch, self.unet_dim, max_latent_height, max_latent_width),
                ],
                "encoder_hidden_states": [
                    (self.xB * min_batch, self.text_maxlen, self.embedding_dim),
                    (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
                    (self.xB * max_batch, self.text_maxlen, self.embedding_dim),
                ],
            }
        else:
            return {
                "sample": [
                    (self.xB * min_batch, self.unet_dim, min_latent_height, min_latent_width),
                    (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
                    (self.xB * max_batch, self.unet_dim, max_latent_height, max_latent_width),
                ],
                "encoder_hidden_states": [
                    (self.xB * min_batch, self.text_maxlen, self.embedding_dim),
                    (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
                    (self.xB * max_batch, self.text_maxlen, self.embedding_dim),
                ],
                "images": [
                    (len(self.controlnets), self.xB * min_batch, 3, min_image_height, min_image_width),
                    (len(self.controlnets), self.xB * batch_size, 3, image_height, image_width),
                    (len(self.controlnets), self.xB * max_batch, 3, max_image_height, max_image_width),
                ],
            }

    def get_shape_dict(self, batch_size, image_height, image_width):
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        if self.controlnets is None:
            return {
                "sample": (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
                "encoder_hidden_states": (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
                "latent": (self.xB * batch_size, 4, latent_height, latent_width),
            }
        else:
            return {
                "sample": (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
                "encoder_hidden_states": (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
                "images": (len(self.controlnets), self.xB * batch_size, 3, image_height, image_width),
                "latent": (self.xB * batch_size, 4, latent_height, latent_width),
            }

    def get_sample_input(self, batch_size, image_height, image_width, static_shape):
        # WAR to enable inference for H/W that are not multiples of 16
        # If building with Dynamic Shapes: ensure image height and width are not multiples of 16 for ONNX export and TensorRT engine build
        if not static_shape:
            image_height = image_height - 8 if image_height % 16 == 0 else image_height
            image_width = image_width - 8 if image_width % 16 == 0 else image_width
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        dtype = torch.float16 if self.fp16 else torch.float32
        if self.controlnets is None:
            return (
                torch.randn(batch_size, self.unet_dim, latent_height, latent_width, dtype=dtype, device=self.device),
                torch.tensor([1.0], dtype=dtype, device=self.device),
                torch.randn(batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),
            )
        else:
            return (
                torch.randn(batch_size, self.unet_dim, latent_height, latent_width, dtype=dtype, device=self.device),
                torch.tensor(999, dtype=dtype, device=self.device),
                torch.randn(batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),
                torch.randn(
                    len(self.controlnets), batch_size, 3, image_height, image_width, dtype=dtype, device=self.device
                ),
                torch.randn(len(self.controlnets), dtype=dtype, device=self.device),
            )

    def optimize(self, onnx_graph):
        if self.fp8:
            return super().optimize(onnx_graph, modify_fp8_graph=True)
        if self.int8:
            return super().optimize(onnx_graph, modify_int8_graph=True)
        return super().optimize(onnx_graph)


class UNetXLModel(base_model.BaseModel):
    def __init__(
        self,
        version,
        pipeline,
        device,
        hf_token,
        verbose,
        framework_model_dir,
        fp16=False,
        int8=False,
        fp8=False,
        max_batch_size=16,
        text_maxlen=77,
        do_classifier_free_guidance=False,
    ):
        super(UNetXLModel, self).__init__(
            version,
            pipeline,
            device=device,
            hf_token=hf_token,
            verbose=verbose,
            framework_model_dir=framework_model_dir,
            fp16=fp16,
            int8=int8,
            fp8=fp8,
            max_batch_size=max_batch_size,
            text_maxlen=text_maxlen,
            embedding_dim=get_unet_embedding_dim(version, pipeline),
        )
        self.subfolder = "unet"
        self.unet_dim = 4
        self.time_dim = 5 if pipeline.is_sd_xl_refiner() else 6
        self.xB = 2 if do_classifier_free_guidance else 1  # batch multiplier

    def get_model(self, torch_inference=""):
        model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {}
        unet_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
        if not load.is_model_cached(unet_model_dir, model_opts, self.hf_safetensor):
            model = UNet2DConditionModel.from_pretrained(
                self.path,
                subfolder=self.subfolder,
                use_safetensors=self.hf_safetensor,
                token=self.hf_token,
                **model_opts,
            ).to(self.device)
            # Use default attention processor for ONNX export
            if not torch_inference:
                model.set_default_attn_processor()
            model.save_pretrained(unet_model_dir, **model_opts)
        else:
            print(f"[I] Load UNet2DConditionModel model from: {unet_model_dir}")
            model = UNet2DConditionModel.from_pretrained(unet_model_dir, **model_opts).to(self.device)
        model = optimizer.optimize_checkpoint(model, torch_inference)
        return model

    def get_input_names(self):
        return ["sample", "timestep", "encoder_hidden_states", "text_embeds", "time_ids"]

    def get_output_names(self):
        return ["latent"]

    def get_dynamic_axes(self):
        xB = "2B" if self.xB == 2 else "B"
        return {
            "sample": {0: xB, 2: "H", 3: "W"},
            "encoder_hidden_states": {0: xB},
            "latent": {0: xB, 2: "H", 3: "W"},
            "text_embeds": {0: xB},
            "time_ids": {0: xB},
        }

    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
        # WAR to enable inference for H/W that are not multiples of 16
        # If building with Dynamic Shapes: ensure image height and width are not multiples of 16 for ONNX export and TensorRT engine build
        if not static_shape:
            image_height = image_height - 8 if image_height % 16 == 0 else image_height
            image_width = image_width - 8 if image_width % 16 == 0 else image_width
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
            self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
        )
        return {
            "sample": [
                (self.xB * min_batch, self.unet_dim, min_latent_height, min_latent_width),
                (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
                (self.xB * max_batch, self.unet_dim, max_latent_height, max_latent_width),
            ],
            "encoder_hidden_states": [
                (self.xB * min_batch, self.text_maxlen, self.embedding_dim),
                (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
                (self.xB * max_batch, self.text_maxlen, self.embedding_dim),
            ],
            "text_embeds": [(self.xB * min_batch, 1280), (self.xB * batch_size, 1280), (self.xB * max_batch, 1280)],
            "time_ids": [
                (self.xB * min_batch, self.time_dim),
                (self.xB * batch_size, self.time_dim),
                (self.xB * max_batch, self.time_dim),
            ],
        }

    def get_shape_dict(self, batch_size, image_height, image_width):
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        return {
            "sample": (self.xB * batch_size, self.unet_dim, latent_height, latent_width),
            "encoder_hidden_states": (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
            "latent": (self.xB * batch_size, 4, latent_height, latent_width),
            "text_embeds": (self.xB * batch_size, 1280),
            "time_ids": (self.xB * batch_size, self.time_dim),
        }

    def get_sample_input(self, batch_size, image_height, image_width, static_shape):
        # WAR to enable inference for H/W that are not multiples of 16
        # If building with Dynamic Shapes: ensure image height and width are not multiples of 16 for ONNX export and TensorRT engine build
        if not static_shape:
            image_height = image_height - 8 if image_height % 16 == 0 else image_height
            image_width = image_width - 8 if image_width % 16 == 0 else image_width
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        dtype = torch.float16 if self.fp16 else torch.float32
        return (
            torch.randn(
                self.xB * batch_size, self.unet_dim, latent_height, latent_width, dtype=dtype, device=self.device
            ),
            torch.tensor([1.0], dtype=dtype, device=self.device),
            torch.randn(self.xB * batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device),
            {
                "added_cond_kwargs": {
                    "text_embeds": torch.randn(self.xB * batch_size, 1280, dtype=dtype, device=self.device),
                    "time_ids": torch.randn(self.xB * batch_size, self.time_dim, dtype=dtype, device=self.device),
                }
            },
        )

    def optimize(self, onnx_graph):
        if self.fp8:
            return super().optimize(onnx_graph, modify_fp8_graph=True)
        if self.int8:
            return super().optimize(onnx_graph, modify_int8_graph=True)
        return super().optimize(onnx_graph)


class UNetXLModelControlNet(UNetXLModel):
    def __init__(
        self,
        version,
        pipeline,
        device,
        hf_token,
        verbose,
        framework_model_dir,
        fp16=False,
        int8=False,
        fp8=False,
        max_batch_size=16,
        text_maxlen=77,
        controlnets=None,
        do_classifier_free_guidance=False,
    ):
        super().__init__(
            version=version,
            pipeline=pipeline,
            device=device,
            hf_token=hf_token,
            verbose=verbose,
            framework_model_dir=framework_model_dir,
            fp16=fp16,
            int8=int8,
            fp8=fp8,
            max_batch_size=max_batch_size,
            text_maxlen=text_maxlen,
            do_classifier_free_guidance=do_classifier_free_guidance,
        )
        self.controlnets = load.get_path(version, pipeline, controlnets) if controlnets else None

    def get_pipeline(self):
        cnet_model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
        controlnets = [
            ControlNetModel.from_pretrained(path, **cnet_model_opts).to(self.device) for path in self.controlnets
        ]
        if self.bf16:
            model_opts = {"torch_dtype": torch.bfloat16}
        elif self.fp16:
            model_opts = {"variant": "fp16", "torch_dtype": torch.float16}
        else:
            model_opts = {}
        pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
            self.path,
            use_safetensors=self.hf_safetensor,
            token=self.hf_token,
            controlnet=controlnets,
            **model_opts,
        ).to(self.device)
        return pipeline

    def get_model(self, torch_inference=""):
        model_opts = {"variant": "fp16", "torch_dtype": torch.float16} if self.fp16 else {}
        unet_model = UNet2DConditionModel.from_pretrained(
            self.path,
            subfolder=self.subfolder,
            use_safetensors=self.hf_safetensor,
            token=self.hf_token,
            **model_opts,
        ).to(self.device)
        cnet_model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
        controlnets = torch.nn.ModuleList(
            [ControlNetModel.from_pretrained(path, **cnet_model_opts).to(self.device) for path in self.controlnets]
        )
        # FIXME - cache UNet2DConditionControlNetModel
        model = UNet2DConditionControlNetModel(unet_model, controlnets)
        model = optimizer.optimize_checkpoint(model, torch_inference)
        return model

    def get_input_names(self):
        return ["sample", "timestep", "encoder_hidden_states", "images", "controlnet_scales", "text_embeds", "time_ids"]

    def get_dynamic_axes(self):
        xB = "2B" if self.xB == 2 else "B"
        result = super().get_dynamic_axes()
        result["images"] = {1: xB, 3: "8H", 4: "8W"}
        return result

    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
        min_batch, max_batch, min_image_height, max_image_height, min_image_width, max_image_width, _, _, _, _ = (
            self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
        )
        result = super().get_input_profile(batch_size, image_height, image_width, static_batch, static_shape)
        result["images"] = [
            (len(self.controlnets), self.xB * min_batch, 3, min_image_height, min_image_width),
            (len(self.controlnets), self.xB * batch_size, 3, image_height, image_width),
            (len(self.controlnets), self.xB * max_batch, 3, max_image_height, max_image_width),
        ]
        return result

    def get_shape_dict(self, batch_size, image_height, image_width):
        result = super().get_shape_dict(batch_size, image_height, image_width)
        result["images"] = (len(self.controlnets), self.xB * batch_size, 3, image_height, image_width)
        return result

    def get_sample_input(self, batch_size, image_height, image_width, static_shape):
        dtype = torch.float16 if self.fp16 else torch.float32
        result = super().get_sample_input(batch_size, image_height, image_width, static_shape)
        result = (
            result[:-1]
            + (
                torch.randn(
                    len(self.controlnets),
                    self.xB * batch_size,
                    3,
                    image_height,
                    image_width,
                    dtype=dtype,
                    device=self.device,
                ),  # images
                torch.randn(len(self.controlnets), dtype=dtype, device=self.device),  # controlnet_scales
            )
            + result[-1:]
        )
        return result


class UNetTemporalModel(base_model.BaseModel):
    def __init__(
        self,
        version,
        pipeline,
        device,
        hf_token,
        verbose,
        framework_model_dir,
        fp16=False,
        fp8=False,
        max_batch_size=16,
        num_frames=14,
        do_classifier_free_guidance=True,
    ):
        super(UNetTemporalModel, self).__init__(
            version,
            pipeline,
            device=device,
            hf_token=hf_token,
            verbose=verbose,
            framework_model_dir=framework_model_dir,
            fp16=fp16,
            fp8=fp8,
            max_batch_size=max_batch_size,
            embedding_dim=get_unet_embedding_dim(version, pipeline),
        )
        self.subfolder = "unet"
        self.unet_dim = 4
        self.num_frames = num_frames
        self.out_channels = 4
        self.cross_attention_dim = 1024
        self.xB = 2 if do_classifier_free_guidance else 1  # batch multiplier

    def get_model(self, torch_inference=""):
        model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
        unet_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
        if not load.is_model_cached(unet_model_dir, model_opts, self.hf_safetensor):
            model = UNetSpatioTemporalConditionModel.from_pretrained(
                self.path,
                subfolder=self.subfolder,
                use_safetensors=self.hf_safetensor,
                token=self.hf_token,
                **model_opts,
            ).to(self.device)
            model.save_pretrained(unet_model_dir, **model_opts)
        else:
            print(f"[I] Load UNetSpatioTemporalConditionModel model from: {unet_model_dir}")
            model = UNetSpatioTemporalConditionModel.from_pretrained(unet_model_dir, **model_opts).to(self.device)
        model = optimizer.optimize_checkpoint(model, torch_inference)
        return model

    def get_input_names(self):
        return ["sample", "timestep", "encoder_hidden_states", "added_time_ids"]

    def get_output_names(self):
        return ["latent"]

    def get_dynamic_axes(self):
        xB = str(self.xB) + "B"
        return {
            "sample": {0: xB, 1: "num_frames", 3: "H", 4: "W"},
            "encoder_hidden_states": {0: xB},
            "added_time_ids": {0: xB},
        }

    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        (
            min_batch,
            max_batch,
            min_image_height,
            max_image_height,
            min_image_width,
            max_image_width,
            min_latent_height,
            max_latent_height,
            min_latent_width,
            max_latent_width,
        ) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
        return {
            "sample": [
                (self.xB * min_batch, self.num_frames, 2 * self.out_channels, min_latent_height, min_latent_width),
                (self.xB * batch_size, self.num_frames, 2 * self.out_channels, latent_height, latent_width),
                (self.xB * max_batch, self.num_frames, 2 * self.out_channels, max_latent_height, max_latent_width),
            ],
            "encoder_hidden_states": [
                (self.xB * min_batch, 1, self.cross_attention_dim),
                (self.xB * batch_size, 1, self.cross_attention_dim),
                (self.xB * max_batch, 1, self.cross_attention_dim),
            ],
            "added_time_ids": [(self.xB * min_batch, 3), (self.xB * batch_size, 3), (self.xB * max_batch, 3)],
        }

    def get_shape_dict(self, batch_size, image_height, image_width):
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        return {
            "sample": (self.xB * batch_size, self.num_frames, 2 * self.out_channels, latent_height, latent_width),
            "timestep": (1,),
            "encoder_hidden_states": (self.xB * batch_size, 1, self.cross_attention_dim),
            "added_time_ids": (self.xB * batch_size, 3),
        }

    def get_sample_input(self, batch_size, image_height, image_width, static_shape):
        # TODO chunk_size if forward_chunking is used
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)

        dtype = torch.float16 if self.fp16 else torch.float32
        return (
            torch.randn(
                self.xB * batch_size,
                self.num_frames,
                2 * self.out_channels,
                latent_height,
                latent_width,
                dtype=dtype,
                device=self.device,
            ),
            torch.tensor([1.0], dtype=torch.float32, device=self.device),
            torch.randn(self.xB * batch_size, 1, self.cross_attention_dim, dtype=dtype, device=self.device),
            torch.randn(self.xB * batch_size, 3, dtype=dtype, device=self.device),
        )

    def optimize(self, onnx_graph):
        return super().optimize(onnx_graph, modify_fp8_graph=self.fp8)


class UNetCascadeModel(base_model.BaseModel):
    def __init__(
        self,
        version,
        pipeline,
        device,
        hf_token,
        verbose,
        framework_model_dir,
        fp16=False,
        bf16=False,
        max_batch_size=16,
        text_maxlen=77,
        do_classifier_free_guidance=False,
        compression_factor=42,
        latent_dim_scale=10.67,
        image_embedding_dim=768,
        lite=False,
    ):
        super(UNetCascadeModel, self).__init__(
            version,
            pipeline,
            device=device,
            hf_token=hf_token,
            verbose=verbose,
            framework_model_dir=framework_model_dir,
            fp16=fp16,
            bf16=bf16,
            max_batch_size=max_batch_size,
            text_maxlen=text_maxlen,
            embedding_dim=get_unet_embedding_dim(version, pipeline),
            compression_factor=compression_factor,
        )
        self.is_prior = True if pipeline.is_cascade_prior() else False
        self.subfolder = "prior" if self.is_prior else "decoder"
        if lite:
            self.subfolder += "_lite"
        self.prior_dim = 16
        self.decoder_dim = 4
        self.xB = 2 if do_classifier_free_guidance else 1  # batch multiplier
        self.latent_dim_scale = latent_dim_scale
        self.min_latent_shape = self.min_image_shape // self.compression_factor
        self.max_latent_shape = self.max_image_shape // self.compression_factor
        self.do_constant_folding = False
        self.image_embedding_dim = image_embedding_dim

    def get_model(self, torch_inference=""):
        # FP16 variant doesn't exist
        model_opts = {"torch_dtype": torch.float16} if self.fp16 else {}
        model_opts = {"variant": "bf16", "torch_dtype": torch.bfloat16} if self.bf16 else model_opts
        unet_model_dir = load.get_checkpoint_dir(self.framework_model_dir, self.version, self.pipeline, self.subfolder)
        if not load.is_model_cached(unet_model_dir, model_opts, self.hf_safetensor):
            model = StableCascadeUNet.from_pretrained(
                self.path,
                subfolder=self.subfolder,
                use_safetensors=self.hf_safetensor,
                token=self.hf_token,
                **model_opts,
            ).to(self.device)
            model.save_pretrained(unet_model_dir, **model_opts)
        else:
            print(f"[I] Load Stable Cascade UNet pytorch model from: {unet_model_dir}")
            model = StableCascadeUNet.from_pretrained(unet_model_dir, **model_opts).to(self.device)
        model = optimizer.optimize_checkpoint(model, torch_inference)
        return model

    def get_input_names(self):
        if self.is_prior:
            return ["sample", "timestep_ratio", "clip_text_pooled", "clip_text", "clip_img"]
        else:
            return ["sample", "timestep_ratio", "clip_text_pooled", "effnet"]

    def get_output_names(self):
        return ["latent"]

    def get_dynamic_axes(self):
        xB = "2B" if self.xB == 2 else "B"
        if self.is_prior:
            return {
                "sample": {0: xB, 2: "H", 3: "W"},
                "timestep_ratio": {0: xB},
                "clip_text_pooled": {0: xB},
                "clip_text": {0: xB},
                "clip_img": {0: xB},
                "latent": {0: xB, 2: "H", 3: "W"},
            }
        else:
            return {
                "sample": {0: xB, 2: "H", 3: "W"},
                "timestep_ratio": {0: xB},
                "clip_text_pooled": {0: xB},
                "effnet": {0: xB, 2: "H_effnet", 3: "W_effnet"},
                "latent": {0: xB, 2: "H", 3: "W"},
            }

    def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
            self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
        )
        if self.is_prior:
            return {
                "sample": [
                    (self.xB * min_batch, self.prior_dim, min_latent_height, min_latent_width),
                    (self.xB * batch_size, self.prior_dim, latent_height, latent_width),
                    (self.xB * max_batch, self.prior_dim, max_latent_height, max_latent_width),
                ],
                "timestep_ratio": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)],
                "clip_text_pooled": [
                    (self.xB * min_batch, 1, self.embedding_dim),
                    (self.xB * batch_size, 1, self.embedding_dim),
                    (self.xB * max_batch, 1, self.embedding_dim),
                ],
                "clip_text": [
                    (self.xB * min_batch, self.text_maxlen, self.embedding_dim),
                    (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
                    (self.xB * max_batch, self.text_maxlen, self.embedding_dim),
                ],
                "clip_img": [
                    (self.xB * min_batch, 1, self.image_embedding_dim),
                    (self.xB * batch_size, 1, self.image_embedding_dim),
                    (self.xB * max_batch, 1, self.image_embedding_dim),
                ],
            }
        else:
            return {
                "sample": [
                    (
                        self.xB * min_batch,
                        self.decoder_dim,
                        int(min_latent_height * self.latent_dim_scale),
                        int(min_latent_width * self.latent_dim_scale),
                    ),
                    (
                        self.xB * batch_size,
                        self.decoder_dim,
                        int(latent_height * self.latent_dim_scale),
                        int(latent_width * self.latent_dim_scale),
                    ),
                    (
                        self.xB * max_batch,
                        self.decoder_dim,
                        int(max_latent_height * self.latent_dim_scale),
                        int(max_latent_width * self.latent_dim_scale),
                    ),
                ],
                "timestep_ratio": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)],
                "clip_text_pooled": [
                    (self.xB * min_batch, 1, self.embedding_dim),
                    (self.xB * batch_size, 1, self.embedding_dim),
                    (self.xB * max_batch, 1, self.embedding_dim),
                ],
                "effnet": [
                    (self.xB * min_batch, self.prior_dim, min_latent_height, min_latent_width),
                    (self.xB * batch_size, self.prior_dim, latent_height, latent_width),
                    (self.xB * max_batch, self.prior_dim, max_latent_height, max_latent_width),
                ],
            }

    def get_shape_dict(self, batch_size, image_height, image_width):
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        if self.is_prior:
            return {
                "sample": (self.xB * batch_size, self.prior_dim, latent_height, latent_width),
                "timestep_ratio": (self.xB * batch_size,),
                "clip_text_pooled": (self.xB * batch_size, 1, self.embedding_dim),
                "clip_text": (self.xB * batch_size, self.text_maxlen, self.embedding_dim),
                "clip_img": (self.xB * batch_size, 1, self.image_embedding_dim),
                "latent": (self.xB * batch_size, self.prior_dim, latent_height, latent_width),
            }
        else:
            return {
                "sample": (
                    self.xB * batch_size,
                    self.decoder_dim,
                    int(latent_height * self.latent_dim_scale),
                    int(latent_width * self.latent_dim_scale),
                ),
                "timestep_ratio": (self.xB * batch_size,),
                "clip_text_pooled": (self.xB * batch_size, 1, self.embedding_dim),
                "effnet": (self.xB * batch_size, self.prior_dim, latent_height, latent_width),
                "latent": (
                    self.xB * batch_size,
                    self.decoder_dim,
                    int(latent_height * self.latent_dim_scale),
                    int(latent_width * self.latent_dim_scale),
                ),
            }

    def get_sample_input(self, batch_size, image_height, image_width, static_shape):
        latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
        dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
        if self.is_prior:
            return (
                torch.randn(batch_size, self.prior_dim, latent_height, latent_width, dtype=dtype, device=self.device),
                torch.tensor([1.0] * batch_size, dtype=dtype, device=self.device),
                torch.randn(batch_size, 1, self.embedding_dim, dtype=dtype, device=self.device),
                {
                    "clip_text": torch.randn(
                        batch_size, self.text_maxlen, self.embedding_dim, dtype=dtype, device=self.device
                    ),
                    "clip_img": torch.randn(batch_size, 1, self.image_embedding_dim, dtype=dtype, device=self.device),
                },
            )
        else:
            return (
                torch.randn(
                    batch_size,
                    self.decoder_dim,
                    int(latent_height * self.latent_dim_scale),
                    int(latent_width * self.latent_dim_scale),
                    dtype=dtype,
                    device=self.device,
                ),
                torch.tensor([1.0] * batch_size, dtype=dtype, device=self.device),
                torch.randn(batch_size, 1, self.embedding_dim, dtype=dtype, device=self.device),
                {
                    "effnet": torch.randn(
                        batch_size, self.prior_dim, latent_height, latent_width, dtype=dtype, device=self.device
                    ),
                },
            )
