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,
)


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

UpperCamelCase__    							= '\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_			(					__A,  __A,  __A=8							)					->    Any:







      '''simple docstring'''


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







class 					A							(    UpperCAmelCase_					):


   def __init__(self					:						Union[str, Any]	,  __UpperCAmelCase					:						UNetaDConditionModel	,  __UpperCAmelCase					:						DDPMScheduler	,  __UpperCAmelCase					:						VQModel	,  )						->    Any:
         """simple docstring"""



         super().__init__()

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


   def       lowercase_				(self					:						Tuple	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Optional[Any]		)						->    Optional[int]:
         """simple docstring"""



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

         UpperCAmelCase__								=  latents * scheduler.init_noise_sigma
         return latents


   def       lowercase_				(self					:						Any	,  __UpperCAmelCase					:						Any=0		)						->    Optional[int]:
         """simple docstring"""



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

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

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


   def       lowercase_				(self					:						Dict	,  __UpperCAmelCase					:						List[str]=0		)						->    int:
         """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."		)

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

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

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

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


   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def       lowercase_				(self					:						Dict		)						->    Tuple:
         """simple docstring"""



         if not hasattr(self.unet	,  "_hf_hook"		):
               return self.device
         for module in self.unet.modules():
               if (
                   hasattr(__UpperCAmelCase	,  "_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(__UpperCAmelCase		)
   def __call__(self					:						Optional[int]	,  __UpperCAmelCase					:						Union[torch.FloatTensor, List[torch.FloatTensor]]	,  __UpperCAmelCase					:						Union[torch.FloatTensor, List[torch.FloatTensor]]	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 5_1_2	,  __UpperCAmelCase					:						int = 1_0_0	,  __UpperCAmelCase					:						float = 4.0	,  __UpperCAmelCase					:						int = 1	,  __UpperCAmelCase					:						Optional[Union[torch.Generator, List[torch.Generator]]] = None	,  __UpperCAmelCase					:						Optional[torch.FloatTensor] = None	,  __UpperCAmelCase					:						Optional[str] = "pil"	,  __UpperCAmelCase					:						bool = True	,  )						->    Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self._execution_device

         UpperCAmelCase__								=  guidance_scale > 1.0

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

         if do_classifier_free_guidance:
               UpperCAmelCase__								=  image_embeds.repeat_interleave(__UpperCAmelCase	,  dim=0		)
               UpperCAmelCase__								=  negative_image_embeds.repeat_interleave(__UpperCAmelCase	,  dim=0		)

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

         self.scheduler.set_timesteps(__UpperCAmelCase	,  device=__UpperCAmelCase		)
         UpperCAmelCase__								=  self.scheduler.timesteps

         UpperCAmelCase__								=  self.unet.config.in_channels

         UpperCAmelCase__		,				UpperCAmelCase__								=  downscale_height_and_width(__UpperCAmelCase	,  __UpperCAmelCase	,  self.movq_scale_factor		)

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

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

               UpperCAmelCase__								=  {"image_embeds": image_embeds}
               UpperCAmelCase__								=  self.unet(
                   sample=__UpperCAmelCase	,  timestep=__UpperCAmelCase	,  encoder_hidden_states=__UpperCAmelCase	,  added_cond_kwargs=__UpperCAmelCase	,  return_dict=__UpperCAmelCase	,  )[0]

               if do_classifier_free_guidance:
                     UpperCAmelCase__		,				UpperCAmelCase__								=  noise_pred.split(latents.shape[1]	,  dim=1		)
                     UpperCAmelCase__		,				UpperCAmelCase__								=  noise_pred.chunk(2		)
                     UpperCAmelCase__		,				UpperCAmelCase__								=  variance_pred.chunk(2		)
                     UpperCAmelCase__								=  noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                     UpperCAmelCase__								=  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"]
               ):
                     UpperCAmelCase__		,				UpperCAmelCase__								=  noise_pred.split(latents.shape[1]	,  dim=1		)

               # compute the previous noisy sample x_t -> x_t-1
               UpperCAmelCase__								=  self.scheduler.step(
                   __UpperCAmelCase	,  __UpperCAmelCase	,  __UpperCAmelCase	,  generator=__UpperCAmelCase	,  )[0]
         # post-processing
         UpperCAmelCase__								=  self.movq.decode(__UpperCAmelCase	,  force_not_quantize=__UpperCAmelCase		)["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"]:
               UpperCAmelCase__								=  image * 0.5 + 0.5
               UpperCAmelCase__								=  image.clamp(0	,  1		)
               UpperCAmelCase__								=  image.cpu().permute(0	,  2	,  3	,  1		).float().numpy()

         if output_type == "pil":
               UpperCAmelCase__								=  self.numpy_to_pil(__UpperCAmelCase		)

         if not return_dict:
               return (image,)

         return ImagePipelineOutput(images=__UpperCAmelCase		)




from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers.models.distilbert.modeling_tf_distilbert import (
          TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFDistilBertForMaskedLM,
          TFDistilBertForMultipleChoice,
          TFDistilBertForQuestionAnswering,
          TFDistilBertForSequenceClassification,
          TFDistilBertForTokenClassification,
          TFDistilBertModel,
      )





class 					A							:


   def __init__(self					:						Tuple	,  __UpperCAmelCase					:						str	,  )						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  parent
         UpperCAmelCase__								=  1_3
         UpperCAmelCase__								=  7
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  False
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  9_9
         UpperCAmelCase__								=  3_2
         UpperCAmelCase__								=  2
         UpperCAmelCase__								=  4
         UpperCAmelCase__								=  3_7
         UpperCAmelCase__								=  "gelu"
         UpperCAmelCase__								=  0.1
         UpperCAmelCase__								=  0.1
         UpperCAmelCase__								=  5_1_2
         UpperCAmelCase__								=  1_6
         UpperCAmelCase__								=  2
         UpperCAmelCase__								=  0.02
         UpperCAmelCase__								=  3
         UpperCAmelCase__								=  4
         UpperCAmelCase__								=  None


   def       lowercase_				(self					:						int		)						->    Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	,  self.vocab_size		)

         UpperCAmelCase__								=  None
         if self.use_input_mask:
               UpperCAmelCase__								=  random_attention_mask([self.batch_size, self.seq_length]		)

         UpperCAmelCase__								=  None
         UpperCAmelCase__								=  None
         UpperCAmelCase__								=  None
         if self.use_labels:
               UpperCAmelCase__								=  ids_tensor([self.batch_size]	,  self.type_sequence_label_size		)
               UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]	,  self.num_labels		)
               UpperCAmelCase__								=  ids_tensor([self.batch_size]	,  self.num_choices		)

         UpperCAmelCase__								=  DistilBertConfig(
             vocab_size=self.vocab_size	,  dim=self.hidden_size	,  n_layers=self.num_hidden_layers	,  n_heads=self.num_attention_heads	,  hidden_dim=self.intermediate_size	,  hidden_act=self.hidden_act	,  dropout=self.hidden_dropout_prob	,  attention_dropout=self.attention_probs_dropout_prob	,  max_position_embeddings=self.max_position_embeddings	,  initializer_range=self.initializer_range	,  )

         return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Tuple		)						->    Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModel(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}

         UpperCAmelCase__								=  model(__UpperCAmelCase		)

         UpperCAmelCase__								=  [input_ids, input_mask]

         UpperCAmelCase__								=  model(__UpperCAmelCase		)

         self.parent.assertEqual(result.last_hidden_state.shape	,  (self.batch_size, self.seq_length, self.hidden_size)		)


   def       lowercase_				(self					:						str	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Dict		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertForMaskedLM(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.vocab_size)		)


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						List[str]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertForQuestionAnswering(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "input_ids": input_ids,
             "attention_mask": input_mask,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.start_logits.shape	,  (self.batch_size, self.seq_length)		)
         self.parent.assertEqual(result.end_logits.shape	,  (self.batch_size, self.seq_length)		)


   def       lowercase_				(self					:						Optional[int]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						List[str]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  TFDistilBertForSequenceClassification(__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_labels)		)


   def       lowercase_				(self					:						Optional[int]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Any		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_choices
         UpperCAmelCase__								=  TFDistilBertForMultipleChoice(__UpperCAmelCase		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  {
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_choices)		)


   def       lowercase_				(self					:						Dict	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Optional[int]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  TFDistilBertForTokenClassification(__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.num_labels)		)





   def       lowercase_				(self					:						Any		)						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  self.prepare_config_and_inputs()
         ((UpperCAmelCase__)		,				(UpperCAmelCase__)		,				(UpperCAmelCase__)		,				(UpperCAmelCase__)		,				(UpperCAmelCase__)		,				(UpperCAmelCase__))								=  config_and_inputs
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict





@require_tf
class 					A							(    UpperCAmelCase_		,    UpperCAmelCase_		,    unittest.TestCase					):
   __UpperCAmelCase			:			Union[str, Any]              =      (
       (
           TFDistilBertModel,
           TFDistilBertForMaskedLM,
           TFDistilBertForQuestionAnswering,
           TFDistilBertForSequenceClassification,
           TFDistilBertForTokenClassification,
           TFDistilBertForMultipleChoice,
       )
       if is_tf_available()
       else None
   )
   __UpperCAmelCase			:			Optional[int]              =      (
       {
           'feature-extraction': TFDistilBertModel,
           'fill-mask': TFDistilBertForMaskedLM,
           'question-answering': TFDistilBertForQuestionAnswering,
           'text-classification': TFDistilBertForSequenceClassification,
           'token-classification': TFDistilBertForTokenClassification,
           'zero-shot': TFDistilBertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   __UpperCAmelCase			:			Tuple              =      False
   __UpperCAmelCase			:			str              =      False


   def       lowercase_				(self					:						Union[str, Any]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModelTester(self		)
         UpperCAmelCase__								=  ConfigTester(self	,  config_class=__UpperCAmelCase	,  dim=3_7		)


   def       lowercase_				(self					:						Any		)						->    int:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def       lowercase_				(self					:						int		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_model(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Optional[int]		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_masked_lm(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Any		)						->    Optional[int]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_question_answering(*__UpperCAmelCase		)


   def       lowercase_				(self					:						List[str]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_sequence_classification(*__UpperCAmelCase		)


   def       lowercase_				(self					:						Dict		)						->    List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_multiple_choice(*__UpperCAmelCase		)


   def       lowercase_				(self					:						str		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_token_classification(*__UpperCAmelCase		)





   @slow
   def       lowercase_				(self					:						Optional[int]		)						->    Optional[int]:
         """simple docstring"""



         for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]		):
               UpperCAmelCase__								=  TFDistilBertModel.from_pretrained(__UpperCAmelCase		)
               self.assertIsNotNone(__UpperCAmelCase		)





@require_tf
class 					A							(    unittest.TestCase					):


   @slow
   def       lowercase_				(self					:						List[Any]		)						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModel.from_pretrained("distilbert-base-uncased"		)
         UpperCAmelCase__								=  tf.constant([[0, 1, 2, 3, 4, 5]]		)
         UpperCAmelCase__								=  model(__UpperCAmelCase		)[0]

         UpperCAmelCase__								=  [1, 6, 7_6_8]
         self.assertEqual(output.shape	,  __UpperCAmelCase		)

         UpperCAmelCase__								=  tf.constant(
             [
                 [
                     [0.19261885, -0.13732955, 0.4119799],
                     [0.22150156, -0.07422661, 0.39037204],
                     [0.22756018, -0.0896414, 0.3701467],
                 ]
             ]		)
         tf.debugging.assert_near(output[:, :3, :3]	,  __UpperCAmelCase	,  atol=1E-4		)



