

'''simple docstring'''
from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


lowerCamelCase			:				Dict 			=					logging.get_logger(__name__)  # pylint: disable=invalid-name

lowerCamelCase			:				int 			=					"\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-prior\")\n        >>> pipe_prior.to(\"cuda\")\n        >>> prompt = \"red cat, 4k photo\"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-decoder\")\n        >>> pipe.to(\"cuda\")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save(\"cat.png\")\n        ```\n"


def 	_lowerCAmelCase (						_UpperCamelCase						:					Any							,			_UpperCamelCase						:					Optional[Any]							,			_UpperCamelCase						:					Optional[Any]=8    )				->      Optional[int]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   =height // scale_factor**2
     if height % scale_factor**2 != 0:
          new_height += 1
     _SCREAMING_SNAKE_CASE							   =width // scale_factor**2
     if width % scale_factor**2 != 0:
          new_width += 1
     return new_height * scale_factor, new_width * scale_factor







class 							A__   (	A__  ):





  def __init__(		self				:     int							, _a				:     UNetaDConditionModel							, _a				:     DDPMScheduler							, _a				:     VQModel							, )     ->   int:





       '''simple docstring'''






       super().__init__()

       self.register_modules(
           unet=_a							, scheduler=_a							, movq=_a							, )
       _SCREAMING_SNAKE_CASE							   =2 ** (len(self.movq.config.block_out_channels			) - 1)





  def 						A  (		self				:     str							, _a				:     Any							, _a				:     Dict							, _a				:     List[str]							, _a				:     Optional[Any]							, _a				:     Any							, _a				:     List[Any]			)     ->   Dict:





       '''simple docstring'''






       if latents is None:
            _SCREAMING_SNAKE_CASE							   =randn_tensor(_a							, generator=_a							, device=_a							, dtype=_a			)
       else:
            if latents.shape != shape:
                 raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}"			)
            _SCREAMING_SNAKE_CASE							   =latents.to(_a			)

       _SCREAMING_SNAKE_CASE							   =latents * scheduler.init_noise_sigma
       return latents





  def 						A  (		self				:     Dict							, _a				:     Optional[Any]=0			)     ->   str:





       '''simple docstring'''






       if is_accelerate_available():
            from accelerate import cpu_offload
       else:
            raise ImportError('Please install accelerate via `pip install accelerate`'			)

       _SCREAMING_SNAKE_CASE							   =torch.device(f"cuda:{gpu_id}"			)

       _SCREAMING_SNAKE_CASE							   =[
           self.unet,
           self.movq,
       ]
       for cpu_offloaded_model in models:
            if cpu_offloaded_model is not None:
                 cpu_offload(_a							, _a			)





  def 						A  (		self				:     str							, _a				:     Dict=0			)     ->   List[Any]:





       '''simple docstring'''






       if is_accelerate_available() and is_accelerate_version('>='							, '0.17.0.dev0'			):
            from accelerate import cpu_offload_with_hook
       else:
            raise ImportError('`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'			)

       _SCREAMING_SNAKE_CASE							   =torch.device(f"cuda:{gpu_id}"			)

       if self.device.type != "cpu":
            self.to('cpu'							, silence_dtype_warnings=_a			)
            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

       _SCREAMING_SNAKE_CASE							   =None
       for cpu_offloaded_model in [self.unet, self.movq]:
            _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE							   =cpu_offload_with_hook(_a							, _a							, prev_module_hook=_a			)

       # We'll offload the last model manually.
       _SCREAMING_SNAKE_CASE							   =hook





  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def 						A  (		self				:     Tuple			)     ->   Union[str, Any]:





       '''simple docstring'''






       if not hasattr(self.unet							, '_hf_hook'			):
            return self.device
       for module in self.unet.modules():
            if (
                hasattr(_a							, '_hf_hook'			)
                and hasattr(module._hf_hook							, 'execution_device'			)
                and module._hf_hook.execution_device is not None
            ):
                 return torch.device(module._hf_hook.execution_device			)
       return self.device





  @torch.no_grad()
  @replace_example_docstring(_a			)
  def __call__(		self				:     Optional[Any]							, _a				:     Union[torch.FloatTensor, List[torch.FloatTensor]]							, _a				:     Union[torch.FloatTensor, List[torch.FloatTensor]]							, _a				:     int = 512							, _a				:     int = 512							, _a				:     int = 100							, _a				:     float = 4.0							, _a				:     int = 1							, _a				:     Optional[Union[torch.Generator, List[torch.Generator]]] = None							, _a				:     Optional[torch.FloatTensor] = None							, _a				:     Optional[str] = "pil"							, _a				:     bool = True							, )     ->   Dict:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =self._execution_device

       _SCREAMING_SNAKE_CASE							   =guidance_scale > 1.0

       if isinstance(_a							, _a			):
            _SCREAMING_SNAKE_CASE							   =torch.cat(_a							, dim=0			)
       _SCREAMING_SNAKE_CASE							   =image_embeds.shape[0] * num_images_per_prompt
       if isinstance(_a							, _a			):
            _SCREAMING_SNAKE_CASE							   =torch.cat(_a							, dim=0			)

       if do_classifier_free_guidance:
            _SCREAMING_SNAKE_CASE							   =image_embeds.repeat_interleave(_a							, dim=0			)
            _SCREAMING_SNAKE_CASE							   =negative_image_embeds.repeat_interleave(_a							, dim=0			)

       _SCREAMING_SNAKE_CASE							   =torch.cat([negative_image_embeds, image_embeds]							, dim=0			).to(dtype=self.unet.dtype							, device=_a			)

       self.scheduler.set_timesteps(_a							, device=_a			)
       _SCREAMING_SNAKE_CASE							   =self.scheduler.timesteps

       _SCREAMING_SNAKE_CASE							   =self.unet.config.in_channels

       _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE							   =downscale_height_and_width(_a							, _a							, self.movq_scale_factor			)

       # create initial latent
       _SCREAMING_SNAKE_CASE							   =self.prepare_latents(
           (batch_size, num_channels_latents, height, width)							, image_embeds.dtype							, _a							, _a							, _a							, self.scheduler							, )

       for i, t in enumerate(self.progress_bar(_a			)			):
            # expand the latents if we are doing classifier free guidance
            _SCREAMING_SNAKE_CASE							   =torch.cat([latents] * 2			) if do_classifier_free_guidance else latents

            _SCREAMING_SNAKE_CASE							   ={'image_embeds': image_embeds}
            _SCREAMING_SNAKE_CASE							   =self.unet(
                sample=_a							, timestep=_a							, encoder_hidden_states=_a							, added_cond_kwargs=_a							, return_dict=_a							, )[0]

            if do_classifier_free_guidance:
                 _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE							   =noise_pred.split(latents.shape[1]							, dim=1			)
                 _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE							   =noise_pred.chunk(2			)
                 _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE							   =variance_pred.chunk(2			)
                 _SCREAMING_SNAKE_CASE							   =noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                 _SCREAMING_SNAKE_CASE							   =torch.cat([noise_pred, variance_pred_text]							, dim=1			)

            if not (
                hasattr(self.scheduler.config							, 'variance_type'			)
                and self.scheduler.config.variance_type in ["learned", "learned_range"]
            ):
                 _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE							   =noise_pred.split(latents.shape[1]							, dim=1			)

            # compute the previous noisy sample x_t -> x_t-1
            _SCREAMING_SNAKE_CASE							   =self.scheduler.step(
                _a							, _a							, _a							, generator=_a							, )[0]
       # post-processing
       _SCREAMING_SNAKE_CASE							   =self.movq.decode(_a							, force_not_quantize=_a			)['sample']

       if output_type not in ["pt", "np", "pil"]:
            raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"			)

       if output_type in ["np", "pil"]:
            _SCREAMING_SNAKE_CASE							   =image * 0.5 + 0.5
            _SCREAMING_SNAKE_CASE							   =image.clamp(0							, 1			)
            _SCREAMING_SNAKE_CASE							   =image.cpu().permute(0							, 2							, 3							, 1			).float().numpy()

       if output_type == "pil":
            _SCREAMING_SNAKE_CASE							   =self.numpy_to_pil(_a			)

       if not return_dict:
            return (image,)

       return ImagePipelineOutput(images=_a			)



'''simple docstring'''
import argparse
import io

import requests
import torch
from omegaconf import OmegaConf

from diffusers import AutoencoderKL
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
    assign_to_checkpoint,
    conv_attn_to_linear,
    create_vae_diffusers_config,
    renew_vae_attention_paths,
    renew_vae_resnet_paths,
)


def 	_lowerCAmelCase (						_UpperCamelCase						:					Tuple							,			_UpperCamelCase						:					List[Any]    )				->      Dict:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   =checkpoint

     _SCREAMING_SNAKE_CASE							   ={}

     _SCREAMING_SNAKE_CASE							   =vae_state_dict['encoder.conv_in.weight']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['encoder.conv_in.bias']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['encoder.conv_out.weight']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['encoder.conv_out.bias']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['encoder.norm_out.weight']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['encoder.norm_out.bias']

     _SCREAMING_SNAKE_CASE							   =vae_state_dict['decoder.conv_in.weight']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['decoder.conv_in.bias']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['decoder.conv_out.weight']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['decoder.conv_out.bias']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['decoder.norm_out.weight']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['decoder.norm_out.bias']

     _SCREAMING_SNAKE_CASE							   =vae_state_dict['quant_conv.weight']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['quant_conv.bias']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['post_quant_conv.weight']
     _SCREAMING_SNAKE_CASE							   =vae_state_dict['post_quant_conv.bias']

     # Retrieves the keys for the encoder down blocks only
     _SCREAMING_SNAKE_CASE							   =len({'.'.join(layer.split('.'    )[:3]    ) for layer in vae_state_dict if 'encoder.down' in layer}    )
     _SCREAMING_SNAKE_CASE							   ={
         layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(_UpperCamelCase    )
     }

     # Retrieves the keys for the decoder up blocks only
     _SCREAMING_SNAKE_CASE							   =len({'.'.join(layer.split('.'    )[:3]    ) for layer in vae_state_dict if 'decoder.up' in layer}    )
     _SCREAMING_SNAKE_CASE							   ={
         layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(_UpperCamelCase    )
     }

     for i in range(_UpperCamelCase    ):
          _SCREAMING_SNAKE_CASE							   =[key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]

          if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
               _SCREAMING_SNAKE_CASE							   =vae_state_dict.pop(
                   f"encoder.down.{i}.downsample.conv.weight"    )
               _SCREAMING_SNAKE_CASE							   =vae_state_dict.pop(
                   f"encoder.down.{i}.downsample.conv.bias"    )

          _SCREAMING_SNAKE_CASE							   =renew_vae_resnet_paths(_UpperCamelCase    )
          _SCREAMING_SNAKE_CASE							   ={'old': f"down.{i}.block", 'new': f"down_blocks.{i}.resnets"}
          assign_to_checkpoint(_UpperCamelCase							,			_UpperCamelCase							,			_UpperCamelCase							,			additional_replacements=[meta_path]							,			config=_UpperCamelCase    )

     _SCREAMING_SNAKE_CASE							   =[key for key in vae_state_dict if 'encoder.mid.block' in key]
     _SCREAMING_SNAKE_CASE							   =2
     for i in range(1							,			num_mid_res_blocks + 1    ):
          _SCREAMING_SNAKE_CASE							   =[key for key in mid_resnets if f"encoder.mid.block_{i}" in key]

          _SCREAMING_SNAKE_CASE							   =renew_vae_resnet_paths(_UpperCamelCase    )
          _SCREAMING_SNAKE_CASE							   ={'old': f"mid.block_{i}", 'new': f"mid_block.resnets.{i - 1}"}
          assign_to_checkpoint(_UpperCamelCase							,			_UpperCamelCase							,			_UpperCamelCase							,			additional_replacements=[meta_path]							,			config=_UpperCamelCase    )

     _SCREAMING_SNAKE_CASE							   =[key for key in vae_state_dict if 'encoder.mid.attn' in key]
     _SCREAMING_SNAKE_CASE							   =renew_vae_attention_paths(_UpperCamelCase    )
     _SCREAMING_SNAKE_CASE							   ={'old': 'mid.attn_1', 'new': 'mid_block.attentions.0'}
     assign_to_checkpoint(_UpperCamelCase							,			_UpperCamelCase							,			_UpperCamelCase							,			additional_replacements=[meta_path]							,			config=_UpperCamelCase    )
     conv_attn_to_linear(_UpperCamelCase    )

     for i in range(_UpperCamelCase    ):
          _SCREAMING_SNAKE_CASE							   =num_up_blocks - 1 - i
          _SCREAMING_SNAKE_CASE							   =[
              key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
          ]

          if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
               _SCREAMING_SNAKE_CASE							   =vae_state_dict[
                   f"decoder.up.{block_id}.upsample.conv.weight"
               ]
               _SCREAMING_SNAKE_CASE							   =vae_state_dict[
                   f"decoder.up.{block_id}.upsample.conv.bias"
               ]

          _SCREAMING_SNAKE_CASE							   =renew_vae_resnet_paths(_UpperCamelCase    )
          _SCREAMING_SNAKE_CASE							   ={'old': f"up.{block_id}.block", 'new': f"up_blocks.{i}.resnets"}
          assign_to_checkpoint(_UpperCamelCase							,			_UpperCamelCase							,			_UpperCamelCase							,			additional_replacements=[meta_path]							,			config=_UpperCamelCase    )

     _SCREAMING_SNAKE_CASE							   =[key for key in vae_state_dict if 'decoder.mid.block' in key]
     _SCREAMING_SNAKE_CASE							   =2
     for i in range(1							,			num_mid_res_blocks + 1    ):
          _SCREAMING_SNAKE_CASE							   =[key for key in mid_resnets if f"decoder.mid.block_{i}" in key]

          _SCREAMING_SNAKE_CASE							   =renew_vae_resnet_paths(_UpperCamelCase    )
          _SCREAMING_SNAKE_CASE							   ={'old': f"mid.block_{i}", 'new': f"mid_block.resnets.{i - 1}"}
          assign_to_checkpoint(_UpperCamelCase							,			_UpperCamelCase							,			_UpperCamelCase							,			additional_replacements=[meta_path]							,			config=_UpperCamelCase    )

     _SCREAMING_SNAKE_CASE							   =[key for key in vae_state_dict if 'decoder.mid.attn' in key]
     _SCREAMING_SNAKE_CASE							   =renew_vae_attention_paths(_UpperCamelCase    )
     _SCREAMING_SNAKE_CASE							   ={'old': 'mid.attn_1', 'new': 'mid_block.attentions.0'}
     assign_to_checkpoint(_UpperCamelCase							,			_UpperCamelCase							,			_UpperCamelCase							,			additional_replacements=[meta_path]							,			config=_UpperCamelCase    )
     conv_attn_to_linear(_UpperCamelCase    )
     return new_checkpoint


def 	_lowerCAmelCase (						_UpperCamelCase						:					str							,			_UpperCamelCase						:					str							,			)				->      List[Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   =requests.get(
         ' https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml'    )
     _SCREAMING_SNAKE_CASE							   =io.BytesIO(r.content    )

     _SCREAMING_SNAKE_CASE							   =OmegaConf.load(_UpperCamelCase    )
     _SCREAMING_SNAKE_CASE							   =5_12
     _SCREAMING_SNAKE_CASE							   ='cuda' if torch.cuda.is_available() else 'cpu'
     if checkpoint_path.endswith('safetensors'    ):
          from safetensors import safe_open

          _SCREAMING_SNAKE_CASE							   ={}
          with safe_open(_UpperCamelCase							,			framework='pt'							,			device='cpu'    ) as f:
               for key in f.keys():
                    _SCREAMING_SNAKE_CASE							   =f.get_tensor(_UpperCamelCase    )
     else:
          _SCREAMING_SNAKE_CASE							   =torch.load(_UpperCamelCase							,			map_location=_UpperCamelCase    )['state_dict']

     # Convert the VAE model.
     _SCREAMING_SNAKE_CASE							   =create_vae_diffusers_config(_UpperCamelCase							,			image_size=_UpperCamelCase    )
     _SCREAMING_SNAKE_CASE							   =custom_convert_ldm_vae_checkpoint(_UpperCamelCase							,			_UpperCamelCase    )

     _SCREAMING_SNAKE_CASE							   =AutoencoderKL(**_UpperCamelCase    )
     vae.load_state_dict(_UpperCamelCase    )
     vae.save_pretrained(_UpperCamelCase    )


if __name__ == "__main__":
     lowerCamelCase			:				Union[str, Any] 			=					argparse.ArgumentParser()

     parser.add_argument("--vae_pt_path", default=None, type=str, required=True, help="Path to the VAE.pt to convert.")
     parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the VAE.pt to convert.")

     lowerCamelCase			:				List[str] 			=					parser.parse_args()

     vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)
