# Open Source Model Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
# The below Model in this distribution may have been modified by THL A29 Limited ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.

# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved. 
# The below software and/or models in this distribution may have been 
# modified by THL A29 Limited ("Tencent Modifications"). 
# All Tencent Modifications are Copyright (C) THL A29 Limited.

# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT 
# except for the third-party components listed below. 
# Hunyuan 3D does not impose any additional limitations beyond what is outlined 
# in the repsective licenses of these third-party components. 
# Users must comply with all terms and conditions of original licenses of these third-party 
# components and must ensure that the usage of the third party components adheres to 
# all relevant laws and regulations. 

# For avoidance of doubts, Hunyuan 3D means the large language models and 
# their software and algorithms, including trained model weights, parameters (including 
# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, 
# fine-tuning enabling code and other elements of the foregoing made publicly available 
# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.

import math
import numpy
import torch
import inspect
import warnings
from PIL import Image
from einops import rearrange
import torch.nn.functional as F
from diffusers.utils.torch_utils import randn_tensor
from diffusers.configuration_utils import FrozenDict
from diffusers.image_processor import VaeImageProcessor
from typing import Any, Callable, Dict, List, Optional, Union
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers import DDPMScheduler, EulerAncestralDiscreteScheduler, ImagePipelineOutput
from diffusers.loaders import (
    FromSingleFileMixin, 
    LoraLoaderMixin, 
    TextualInversionLoaderMixin
)
from transformers import (
    CLIPImageProcessor, 
    CLIPTextModel, 
    CLIPTokenizer, 
    CLIPVisionModelWithProjection
)
from diffusers.models.attention_processor import (
    Attention, 
    AttnProcessor, 
    XFormersAttnProcessor, 
    AttnProcessor2_0
)

from .utils import to_rgb_image, white_out_background, recenter_img


EXAMPLE_DOC_STRING = """
    Examples:
        ```py
        >>> import torch
        >>> from here import Hunyuan3d_MVD_Qing_Pipeline

        >>> pipe = Hunyuan3d_MVD_Qing_Pipeline.from_pretrained(
        ...     "Tencent-Hunyuan-3D/MVD-Qing", torch_dtype=torch.float16
        ... )
        >>> pipe.to("cuda")

        >>> img = Image.open("demo.png")
        >>> res_img = pipe(img).images[0]
"""

def unscale_latents(latents): return latents / 0.75 + 0.22
def unscale_image  (image  ): return   image / 0.50 * 0.80


def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
    return noise_cfg



class ReferenceOnlyAttnProc(torch.nn.Module):
    # reference attention
    def __init__(self, chained_proc, enabled=False, name=None):
        super().__init__()
        self.enabled = enabled
        self.chained_proc = chained_proc
        self.name = name

    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, mode="w", ref_dict=None):
        if encoder_hidden_states is None: encoder_hidden_states = hidden_states
        if self.enabled:
            if mode == 'w': 
                ref_dict[self.name] = encoder_hidden_states
            elif mode == 'r': 
                encoder_hidden_states = torch.cat([encoder_hidden_states, ref_dict.pop(self.name)], dim=1)
        res = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask)
        return res


# class RowWiseAttnProcessor2_0:
#     def __call__(self, attn, 
#                  hidden_states, 
#                  encoder_hidden_states=None, 
#                  attention_mask=None, 
#                  temb=None, 
#                  num_views=6, 
#                  *args, 
#                  **kwargs):
#         residual = hidden_states
#         if attn.spatial_norm is not None: hidden_states = attn.spatial_norm(hidden_states, temb)

#         input_ndim = hidden_states.ndim
#         if input_ndim == 4:
#             batch_size, channel, height, width = hidden_states.shape
#             hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)

#         if encoder_hidden_states is None:
#             batch_size, sequence_length, _ = hidden_states.shape
#         else:
#             batch_size, sequence_length, _ = encoder_hidden_states.shape
        
#         if attention_mask is not None:
#             attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
#             attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
#         if attn.group_norm is not None: hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)

#         query = attn.to_q(hidden_states)
#         if encoder_hidden_states is None: encoder_hidden_states = hidden_states
#         elif attn.norm_cross: encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)

#         # encoder_hidden_states [B, 6hw+hw, C] if ref att
#         key = attn.to_k(encoder_hidden_states) # [B, Vhw+hw, C]
#         value = attn.to_v(encoder_hidden_states) # [B, Vhw+hw, C]

#         mv_flag = hidden_states.shape[1] < encoder_hidden_states.shape[1] and encoder_hidden_states.shape[1] != 77
#         if mv_flag:
#             target_size = int(math.sqrt(hidden_states.shape[1] // num_views))
#             assert target_size ** 2 * num_views == hidden_states.shape[1]

#             gen_key = key[:, :num_views*target_size*target_size, :]
#             ref_key = key[:, num_views*target_size*target_size:, :]
#             gen_value = value[:, :num_views*target_size*target_size, :]
#             ref_value = value[:, num_views*target_size*target_size:, :]

#             # rowwise attention
#             query, gen_key, gen_value = \
#                     rearrange(    query, "b (v1 h v2 w) c -> (b h) (v1 v2 w) c", 
#                               v1=num_views//2, v2=2, h=target_size, w=target_size), \
#                     rearrange(  gen_key, "b (v1 h v2 w) c -> (b h) (v1 v2 w) c", 
#                               v1=num_views//2, v2=2, h=target_size, w=target_size), \
#                     rearrange(gen_value, "b (v1 h v2 w) c -> (b h) (v1 v2 w) c", 
#                               v1=num_views//2, v2=2, h=target_size, w=target_size)

#             inner_dim = key.shape[-1]
#             ref_size = int(math.sqrt(ref_key.shape[1]))
#             ref_key_expanded = ref_key.view(batch_size, 1, ref_size * ref_size, inner_dim)
#             ref_key_expanded = ref_key_expanded.expand(-1, target_size, -1, -1).contiguous()
#             ref_key_expanded = ref_key_expanded.view(batch_size * target_size, ref_size * ref_size, inner_dim)
#             key   = torch.cat([  gen_key,   ref_key_expanded], dim=1)
            
#             ref_value_expanded = ref_value.view(batch_size, 1, ref_size * ref_size, inner_dim)
#             ref_value_expanded = ref_value_expanded.expand(-1, target_size, -1, -1).contiguous()
#             ref_value_expanded = ref_value_expanded.view(batch_size * target_size, ref_size * ref_size, inner_dim)
#             value = torch.cat([gen_value, ref_value_expanded], dim=1)
#             h = target_size
#         else:
#             target_size = int(math.sqrt(hidden_states.shape[1]))
#             h = 1
#             num_views = 1

#         inner_dim = key.shape[-1]
#         head_dim = inner_dim // attn.heads

#         query = query.view(batch_size * h, -1, attn.heads, head_dim).transpose(1, 2)
#         key   =   key.view(batch_size * h, -1, attn.heads, head_dim).transpose(1, 2)
#         value = value.view(batch_size * h, -1, attn.heads, head_dim).transpose(1, 2)

#         hidden_states = F.scaled_dot_product_attention(query, key, value, 
#                                                        attn_mask=attention_mask, 
#                                                        dropout_p=0.0, 
#                                                        is_causal=False)
#         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size * h, 
#                                                               -1, 
#                                                               attn.heads * head_dim).to(query.dtype)
#         hidden_states = attn.to_out[1](attn.to_out[0](hidden_states))

#         if mv_flag: hidden_states = rearrange(hidden_states, "(b h) (v1 v2 w) c -> b (v1 h v2 w) c", 
#                                               b=batch_size, v1=num_views//2,
#                                               v2=2, h=target_size, w=target_size)

#         if input_ndim == 4: 
#             hidden_states = hidden_states.transpose(-1, -2)
#             hidden_states = hidden_states.reshape(batch_size, 
#                                                   channel, 
#                                                   target_size, 
#                                                   target_size)
#         if attn.residual_connection: hidden_states = hidden_states + residual
#         hidden_states = hidden_states / attn.rescale_output_factor
#         return hidden_states


class RefOnlyNoisedUNet(torch.nn.Module):
    def __init__(self, unet, train_sched, val_sched):
        super().__init__()
        self.unet = unet
        self.train_sched = train_sched
        self.val_sched = val_sched

        unet_lora_attn_procs = dict()
        for name, _ in unet.attn_processors.items():
            unet_lora_attn_procs[name] = ReferenceOnlyAttnProc(AttnProcessor2_0(), 
                                                           enabled=name.endswith("attn1.processor"), 
                                                           name=name)
        unet.set_attn_processor(unet_lora_attn_procs)

    def __getattr__(self, name: str):
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.unet, name)

    def forward(self, sample, timestep, encoder_hidden_states, *args, cross_attention_kwargs, **kwargs):
        cond_lat = cross_attention_kwargs['cond_lat']
        noise = torch.randn_like(cond_lat)
        if self.training:
            noisy_cond_lat = self.train_sched.add_noise(cond_lat, noise, timestep)
            noisy_cond_lat = self.train_sched.scale_model_input(noisy_cond_lat, timestep)
        else:
            noisy_cond_lat = self.val_sched.add_noise(cond_lat, noise, timestep.reshape(-1))
            noisy_cond_lat = self.val_sched.scale_model_input(noisy_cond_lat, timestep.reshape(-1))

        ref_dict = {}
        self.unet(noisy_cond_lat, 
                  timestep, 
                  encoder_hidden_states, 
                  *args, 
                  cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict), 
                  **kwargs)
        return  self.unet(sample, 
                          timestep, 
                          encoder_hidden_states, 
                          *args, 
                          cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict), 
                          **kwargs)


class Hunyuan3d_MVD_Lite_Pipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin):
    def __init__(
        self,
        vae: AutoencoderKL,
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        unet: UNet2DConditionModel,
        scheduler: KarrasDiffusionSchedulers,
        vision_encoder: CLIPVisionModelWithProjection,
        feature_extractor_clip: CLIPImageProcessor, 
        feature_extractor_vae: CLIPImageProcessor,
        ramping_coefficients: Optional[list] = None,
        safety_checker=None,
    ):
        DiffusionPipeline.__init__(self)
        self.register_modules(
            vae=vae,
            unet=unet,
            tokenizer=tokenizer,
            scheduler=scheduler,
            text_encoder=text_encoder,
            vision_encoder=vision_encoder,
            feature_extractor_vae=feature_extractor_vae,
            feature_extractor_clip=feature_extractor_clip)
        '''
            rewrite the stable diffusion pipeline
            vae: vae
            unet: unet
            tokenizer: tokenizer
            scheduler: scheduler
            text_encoder: text_encoder
            vision_encoder: vision_encoder
            feature_extractor_vae: feature_extractor_vae
            feature_extractor_clip: feature_extractor_clip
        '''
        self.register_to_config(ramping_coefficients=ramping_coefficients)
        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)

    def prepare_extra_step_kwargs(self, generator, eta):
        extra_step_kwargs = {}
        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
        if accepts_eta: extra_step_kwargs["eta"] = eta

        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
        if accepts_generator: extra_step_kwargs["generator"] = generator
        return extra_step_kwargs

    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        latents = latents * self.scheduler.init_noise_sigma
        return latents

    @torch.no_grad()
    def _encode_prompt(
        self,
        prompt,
        device,
        num_images_per_prompt,
        do_classifier_free_guidance,
        negative_prompt=None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        lora_scale: Optional[float] = None,
    ):
        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
            self._lora_scale = lora_scale

        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]

        if prompt_embeds is None:
            if isinstance(self, TextualInversionLoaderMixin):
                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)

            text_inputs = self.tokenizer(
                prompt,
                padding="max_length",
                max_length=self.tokenizer.model_max_length,
                truncation=True,
                return_tensors="pt",
            )
            text_input_ids = text_inputs.input_ids

            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
                attention_mask = text_inputs.attention_mask.to(device)
            else:
                attention_mask = None

            prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)[0]

        if self.text_encoder is not None:
            prompt_embeds_dtype = self.text_encoder.dtype
        elif self.unet is not None:
            prompt_embeds_dtype = self.unet.dtype
        else:
            prompt_embeds_dtype = prompt_embeds.dtype

        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
        bs_embed, seq_len, _ = prompt_embeds.shape
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)

        if do_classifier_free_guidance and negative_prompt_embeds is None:
            uncond_tokens: List[str]
            if negative_prompt is None: uncond_tokens = [""] * batch_size
            elif prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError()
            elif isinstance(negative_prompt, str): uncond_tokens = [negative_prompt]
            elif batch_size != len(negative_prompt): raise ValueError()
            else: uncond_tokens = negative_prompt
            if isinstance(self, TextualInversionLoaderMixin):
                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)

            max_length = prompt_embeds.shape[1]
            uncond_input = self.tokenizer(uncond_tokens, 
                                          padding="max_length", 
                                          max_length=max_length, 
                                          truncation=True, 
                                          return_tensors="pt")

            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
                attention_mask = uncond_input.attention_mask.to(device)
            else:
                attention_mask = None

            negative_prompt_embeds = self.text_encoder(uncond_input.input_ids.to(device), attention_mask=attention_mask)
            negative_prompt_embeds = negative_prompt_embeds[0]

        if do_classifier_free_guidance:
            seq_len = negative_prompt_embeds.shape[1]
            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

        return prompt_embeds

    @torch.no_grad()
    def encode_condition_image(self, image: torch.Tensor): return self.vae.encode(image).latent_dist.sample()

    @torch.no_grad()
    def __call__(self, image=None, 
                 width=640, 
                 height=960, 
                 num_inference_steps=75, 
                 return_dict=True, 
                 generator=None, 
                 **kwargs):
        batch_size = 1
        num_images_per_prompt = 1
        output_type = 'pil'
        do_classifier_free_guidance = True
        guidance_rescale = 0.
        if isinstance(self.unet, UNet2DConditionModel): 
            self.unet = RefOnlyNoisedUNet(self.unet, None, self.scheduler).eval()

        cond_image = recenter_img(image)
        cond_image = to_rgb_image(image)
        image = cond_image
        image_1 = self.feature_extractor_vae(images=image, return_tensors="pt").pixel_values
        image_2 = self.feature_extractor_clip(images=image, return_tensors="pt").pixel_values
        image_1 = image_1.to(device=self.vae.device, dtype=self.vae.dtype)
        image_2 = image_2.to(device=self.vae.device, dtype=self.vae.dtype)

        cond_lat = self.encode_condition_image(image_1)
        negative_lat = self.encode_condition_image(torch.zeros_like(image_1))
        cond_lat = torch.cat([negative_lat, cond_lat])
        cross_attention_kwargs = dict(cond_lat=cond_lat)

        global_embeds = self.vision_encoder(image_2, output_hidden_states=False).image_embeds.unsqueeze(-2)
        encoder_hidden_states = self._encode_prompt('', self.device, num_images_per_prompt, False)
        ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
        prompt_embeds = torch.cat([encoder_hidden_states, encoder_hidden_states + global_embeds * ramp])

        device = self._execution_device
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(batch_size * num_images_per_prompt,
                                       num_channels_latents, 
                                       height, 
                                       width, 
                                       prompt_embeds.dtype, 
                                       device, 
                                       generator, 
                                       None)
        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, 0.0)
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order

        # set adaptive cfg
        # the image order is:
        #    [0, 60, 
        #     120, 180,
        #     240, 300]
        # the cfg is set as 3, 2.5, 2, 1.5
        
        tmp_guidance_scale = torch.ones_like(latents)
        tmp_guidance_scale[:, :, :40, :40] = 3
        tmp_guidance_scale[:, :, :40, 40:] =  2.5
        tmp_guidance_scale[:, :, 40:80, :40] =  2
        tmp_guidance_scale[:, :, 40:80, 40:] =  1.5
        tmp_guidance_scale[:, :, 80:120, :40] =  2
        tmp_guidance_scale[:, :, 80:120, 40:] =  2.5

        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                noise_pred = self.unet(latent_model_input, t,
                                encoder_hidden_states=prompt_embeds, 
                                cross_attention_kwargs=cross_attention_kwargs, 
                                return_dict=False)[0]

                adaptive_guidance_scale = (2 + 16 * (t / 1000) ** 5) / 3
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + \
                        tmp_guidance_scale * adaptive_guidance_scale * \
                        (noise_pred_text - noise_pred_uncond)

                if do_classifier_free_guidance and guidance_rescale > 0.0:
                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)

                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
                if i==len(timesteps)-1 or ((i+1)>num_warmup_steps and (i+1)%self.scheduler.order==0): 
                    progress_bar.update()

        latents = unscale_latents(latents)
        image = unscale_image(self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0])
        image = self.image_processor.postprocess(image, output_type='pil')[0]
        image = [image, cond_image]
        return ImagePipelineOutput(images=image) if return_dict else (image,)