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


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

__magic_name__								='''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")
        >>> pipe_prior.to("cuda")
        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
        >>> pipe.to("cuda")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save("cat.png")
        ```
'''







def  __UpperCamelCase    (						A ,							A ,							A=8							):
     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   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				)		->      List[str]:





       '''simple docstring'''





       super().__init__()

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






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)		->      int:





       '''simple docstring'''





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

       UpperCamelCase__										=      latents * scheduler.init_noise_sigma
       return latents






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_=0	)		->      str:





       '''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(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_	)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_=0	)		->      Optional[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.'''	)

       UpperCamelCase__										=      torch.device(F"cuda:{gpu_id}"	)

       if self.device.type != "cpu":
            self.to('''cpu'''   ,				silence_dtype_warnings=SCREAMING_SNAKE_CASE_	)
            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(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				prev_module_hook=SCREAMING_SNAKE_CASE_	)

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






  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def 		_a					(self	)		->      List[str]:





       '''simple docstring'''





       if not hasattr(self.unet   ,				'''_hf_hook'''	):
            return self.device
       for module in self.unet.modules():
            if (
                hasattr(SCREAMING_SNAKE_CASE_   ,				'''_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(SCREAMING_SNAKE_CASE_	)
  def __call__(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_ = 512   ,				SCREAMING_SNAKE_CASE_ = 512   ,				SCREAMING_SNAKE_CASE_ = 100   ,				SCREAMING_SNAKE_CASE_ = 4.0   ,				SCREAMING_SNAKE_CASE_ = 1   ,				SCREAMING_SNAKE_CASE_ = None   ,				SCREAMING_SNAKE_CASE_ = None   ,				SCREAMING_SNAKE_CASE_ = "pil"   ,				SCREAMING_SNAKE_CASE_ = True   ,				)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      self._execution_device

       UpperCamelCase__										=      guidance_scale > 1.0

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

       if do_classifier_free_guidance:
            UpperCamelCase__										=      image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE_   ,				dim=0	)
            UpperCamelCase__										=      negative_image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE_   ,				dim=0	)

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

       self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE_   ,				device=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      self.scheduler.timesteps

       UpperCamelCase__										=      self.unet.config.in_channels

       UpperCamelCase__						,       UpperCamelCase__										=      downscale_height_and_width(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				self.movq_scale_factor	)

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

       for i, t in enumerate(self.progress_bar(SCREAMING_SNAKE_CASE_	)	):
            # 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=SCREAMING_SNAKE_CASE_   ,				timestep=SCREAMING_SNAKE_CASE_   ,				encoder_hidden_states=SCREAMING_SNAKE_CASE_   ,				added_cond_kwargs=SCREAMING_SNAKE_CASE_   ,				return_dict=SCREAMING_SNAKE_CASE_   ,				)[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(
                SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				generator=SCREAMING_SNAKE_CASE_   ,				)[0]
       # post-processing
       UpperCamelCase__										=      self.movq.decode(SCREAMING_SNAKE_CASE_   ,				force_not_quantize=SCREAMING_SNAKE_CASE_	)['''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(SCREAMING_SNAKE_CASE_	)

       if not return_dict:
            return (image,)

       return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE_	)




import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM



@require_torch
class     _A   :






  def __init__(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_=99   ,				SCREAMING_SNAKE_CASE_=13   ,				SCREAMING_SNAKE_CASE_=16   ,				SCREAMING_SNAKE_CASE_=7   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=False   ,				SCREAMING_SNAKE_CASE_=True   ,				SCREAMING_SNAKE_CASE_=2   ,				SCREAMING_SNAKE_CASE_=32   ,				SCREAMING_SNAKE_CASE_=4   ,				SCREAMING_SNAKE_CASE_=4   ,				SCREAMING_SNAKE_CASE_=30   ,				SCREAMING_SNAKE_CASE_=0   ,				SCREAMING_SNAKE_CASE_=1   ,				SCREAMING_SNAKE_CASE_=2   ,				SCREAMING_SNAKE_CASE_=None   ,				)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      parent
       UpperCamelCase__										=      batch_size
       UpperCamelCase__										=      decoder_seq_length
       # For common tests
       UpperCamelCase__										=      self.decoder_seq_length
       UpperCamelCase__										=      is_training
       UpperCamelCase__										=      use_attention_mask
       UpperCamelCase__										=      use_labels

       UpperCamelCase__										=      vocab_size
       UpperCamelCase__										=      d_model
       UpperCamelCase__										=      d_model
       UpperCamelCase__										=      decoder_layers
       UpperCamelCase__										=      decoder_layers
       UpperCamelCase__										=      decoder_ffn_dim
       UpperCamelCase__										=      decoder_attention_heads
       UpperCamelCase__										=      decoder_attention_heads
       UpperCamelCase__										=      eos_token_id
       UpperCamelCase__										=      bos_token_id
       UpperCamelCase__										=      pad_token_id
       UpperCamelCase__										=      decoder_start_token_id
       UpperCamelCase__										=      use_cache
       UpperCamelCase__										=      max_position_embeddings

       UpperCamelCase__										=      None
       UpperCamelCase__										=      decoder_seq_length
       UpperCamelCase__										=      2
       UpperCamelCase__										=      1






  def 		_a					(self	)		->      str:





       '''simple docstring'''





       UpperCamelCase__										=      ids_tensor([self.batch_size, self.decoder_seq_length]   ,				self.vocab_size	)

       UpperCamelCase__										=      None
       if self.use_attention_mask:
            UpperCamelCase__										=      ids_tensor([self.batch_size, self.decoder_seq_length]   ,				vocab_size=2	)

       UpperCamelCase__										=      None
       if self.use_labels:
            UpperCamelCase__										=      ids_tensor([self.batch_size, self.decoder_seq_length]   ,				self.vocab_size	)

       UpperCamelCase__										=      TrOCRConfig(
           vocab_size=self.vocab_size   ,				d_model=self.d_model   ,				decoder_layers=self.decoder_layers   ,				decoder_ffn_dim=self.decoder_ffn_dim   ,				decoder_attention_heads=self.decoder_attention_heads   ,				eos_token_id=self.eos_token_id   ,				bos_token_id=self.bos_token_id   ,				use_cache=self.use_cache   ,				pad_token_id=self.pad_token_id   ,				decoder_start_token_id=self.decoder_start_token_id   ,				max_position_embeddings=self.max_position_embeddings   ,				)

       return (config, input_ids, attention_mask, lm_labels)






  def 		_a					(self   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      True
       UpperCamelCase__										=      TrOCRDecoder(config=SCREAMING_SNAKE_CASE_	).to(SCREAMING_SNAKE_CASE_	).eval()
       UpperCamelCase__										=      input_ids[:2]

       input_ids[input_ids == 0] += 1
       # first forward pass
       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_   ,				use_cache=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_   ,				use_cache=SCREAMING_SNAKE_CASE_	)

       self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_	) == len(SCREAMING_SNAKE_CASE_	)	)
       self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_	) == len(SCREAMING_SNAKE_CASE_	) + 1	)

       UpperCamelCase__										=      outputs['''past_key_values''']

       # create hypothetical next token and extent to next_input_ids
       UpperCamelCase__										=      ids_tensor((2, 1)   ,				config.vocab_size - 1	) + 1

       # append to next input_ids and
       UpperCamelCase__										=      torch.cat([input_ids, next_tokens]   ,				dim=-1	)

       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_	)['''last_hidden_state''']
       UpperCamelCase__										=      model(SCREAMING_SNAKE_CASE_   ,				past_key_values=SCREAMING_SNAKE_CASE_	)['''last_hidden_state''']

       # select random slice
       UpperCamelCase__										=      ids_tensor((1,)   ,				output_from_past.shape[-1]	).item()
       UpperCamelCase__										=      output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
       UpperCamelCase__										=      output_from_past[:, 0, random_slice_idx].detach()

       # test that outputs are equal for slice
       assert torch.allclose(SCREAMING_SNAKE_CASE_   ,				SCREAMING_SNAKE_CASE_   ,				atol=1E-3	)






  def 		_a					(self	)		->      Dict:





       '''simple docstring'''





       UpperCamelCase__										=      self.prepare_config_and_inputs()
       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__						,       UpperCamelCase__										=      config_and_inputs

       UpperCamelCase__										=      {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
       return config, inputs_dict



@require_torch
class     _A   (    __UpperCamelCase	, __UpperCamelCase	, __UpperCamelCase	, unittest.TestCase							):
  SCREAMING_SNAKE_CASE_	:	Tuple						=(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE_	:	Any						=(TrOCRForCausalLM,) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE_	:	Optional[int]						={"text-generation": TrOCRForCausalLM} if is_torch_available() else {}
  SCREAMING_SNAKE_CASE_	:	Optional[int]						=True
  SCREAMING_SNAKE_CASE_	:	Dict						=False






  def 		_a					(self	)		->      int:





       '''simple docstring'''





       UpperCamelCase__										=      TrOCRStandaloneDecoderModelTester(self   ,				is_training=SCREAMING_SNAKE_CASE_	)
       UpperCamelCase__										=      ConfigTester(self   ,				config_class=SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      Dict:





       '''simple docstring'''





       pass






  def 		_a					(self	)		->      List[str]:





       '''simple docstring'''





       pass






  def 		_a					(self	)		->      int:





       '''simple docstring'''





       pass






  def 		_a					(self	)		->      Tuple:





       '''simple docstring'''





       self.config_tester.run_common_tests()






  def 		_a					(self	)		->      Any:





       '''simple docstring'''





       UpperCamelCase__										=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_decoder_model_past(*SCREAMING_SNAKE_CASE_	)






  def 		_a					(self	)		->      List[Any]:





       '''simple docstring'''





       return






  @unittest.skip('''The model doesn\'t support left padding'''	)  # and it's not used enough to be worth fixing :)
  def 		_a					(self	)		->      Optional[int]:





       '''simple docstring'''





       pass



