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


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

__SCREAMING_SNAKE_CASE							=     '\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       A_  (				__lowercase						,	__lowercase						,	__lowercase=8	):
 UpperCamelCase_     :					str   	=height // scale_factor**2
 if height % scale_factor**2 != 0:
  new_height += 1
 UpperCamelCase_     :					Tuple   	=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      ,  _lowerCamelCase :UNetaDConditionModel      ,  _lowerCamelCase :DDPMScheduler      ,  _lowerCamelCase :VQModel      ,  ):







      '''simple docstring'''



      super().__init__()

      self.register_modules(
          unet=_lowerCamelCase      ,  scheduler=_lowerCamelCase      ,  movq=_lowerCamelCase      ,  )
      UpperCamelCase_     :					Any   	=2 ** (len(self.movq.config.block_out_channels   ) - 1)



     def  lowerCamelCase_						(							self :int      ,  _lowerCamelCase :List[str]      ,  _lowerCamelCase :List[Any]      ,  _lowerCamelCase :Any      ,  _lowerCamelCase :int      ,  _lowerCamelCase :List[Any]      ,  _lowerCamelCase :Dict   ):







      '''simple docstring'''



      if latents is None:
       UpperCamelCase_     :					Union[str, Any]   	=randn_tensor(_lowerCamelCase      ,  generator=_lowerCamelCase      ,  device=_lowerCamelCase      ,  dtype=_lowerCamelCase   )
      else:
       if latents.shape != shape:
        raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {shape}'''   )
       UpperCamelCase_     :					Any   	=latents.to(_lowerCamelCase   )

      UpperCamelCase_     :					List[Any]   	=latents * scheduler.init_noise_sigma
      return latents



     def  lowerCamelCase_						(							self :Union[str, Any]      ,  _lowerCamelCase :List[Any]=0   ):







      '''simple docstring'''



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

      UpperCamelCase_     :					Tuple   	=torch.device(f'''cuda:{gpu_id}'''   )

      UpperCamelCase_     :					List[str]   	=[
          self.unet,
          self.movq,
      ]
      for cpu_offloaded_model in models:
       if cpu_offloaded_model is not None:
        cpu_offload(_lowerCamelCase      ,  _lowerCamelCase   )



     def  lowerCamelCase_						(							self :Tuple      ,  _lowerCamelCase :Union[str, Any]=0   ):







      '''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_     :					Any   	=torch.device(f'''cuda:{gpu_id}'''   )

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

      UpperCamelCase_     :					Optional[Any]   	=None
      for cpu_offloaded_model in [self.unet, self.movq]:
       UpperCamelCase_	,     UpperCamelCase_     :					List[str]   	=cpu_offload_with_hook(_lowerCamelCase      ,  _lowerCamelCase      ,  prev_module_hook=_lowerCamelCase   )

      # We'll offload the last model manually.
      UpperCamelCase_     :					Optional[int]   	=hook



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







      '''simple docstring'''



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







      '''simple docstring'''



      UpperCamelCase_     :					Optional[Any]   	=self._execution_device

      UpperCamelCase_     :					Any   	=guidance_scale > 1.0

      if isinstance(_lowerCamelCase      ,  _lowerCamelCase   ):
       UpperCamelCase_     :					str   	=torch.cat(_lowerCamelCase      ,  dim=0   )
      UpperCamelCase_     :					str   	=image_embeds.shape[0] * num_images_per_prompt
      if isinstance(_lowerCamelCase      ,  _lowerCamelCase   ):
       UpperCamelCase_     :					Any   	=torch.cat(_lowerCamelCase      ,  dim=0   )

      if do_classifier_free_guidance:
       UpperCamelCase_     :					List[str]   	=image_embeds.repeat_interleave(_lowerCamelCase      ,  dim=0   )
       UpperCamelCase_     :					List[str]   	=negative_image_embeds.repeat_interleave(_lowerCamelCase      ,  dim=0   )

      UpperCamelCase_     :					List[str]   	=torch.cat([negative_image_embeds, image_embeds]      ,  dim=0   ).to(dtype=self.unet.dtype      ,  device=_lowerCamelCase   )

      self.scheduler.set_timesteps(_lowerCamelCase      ,  device=_lowerCamelCase   )
      UpperCamelCase_     :					Any   	=self.scheduler.timesteps

      UpperCamelCase_     :					List[Any]   	=self.unet.config.in_channels

      UpperCamelCase_	,     UpperCamelCase_     :					Any   	=downscale_height_and_width(_lowerCamelCase      ,  _lowerCamelCase      ,  self.movq_scale_factor   )

      # create initial latent
      UpperCamelCase_     :					str   	=self.prepare_latents(
          (batch_size, num_channels_latents, height, width)      ,  image_embeds.dtype      ,  _lowerCamelCase      ,  _lowerCamelCase      ,  _lowerCamelCase      ,  self.scheduler      ,  )

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

       UpperCamelCase_     :					Optional[int]   	={'image_embeds': image_embeds}
       UpperCamelCase_     :					str   	=self.unet(
           sample=_lowerCamelCase      ,  timestep=_lowerCamelCase      ,  encoder_hidden_states=_lowerCamelCase      ,  added_cond_kwargs=_lowerCamelCase      ,  return_dict=_lowerCamelCase      ,  )[0]

       if do_classifier_free_guidance:
        UpperCamelCase_	,     UpperCamelCase_     :					Optional[Any]   	=noise_pred.split(latents.shape[1]      ,  dim=1   )
        UpperCamelCase_	,     UpperCamelCase_     :					List[str]   	=noise_pred.chunk(2   )
        UpperCamelCase_	,     UpperCamelCase_     :					List[Any]   	=variance_pred.chunk(2   )
        UpperCamelCase_     :					List[Any]   	=noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
        UpperCamelCase_     :					Optional[int]   	=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_     :					Optional[Any]   	=noise_pred.split(latents.shape[1]      ,  dim=1   )

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

      if output_type == "pil":
       UpperCamelCase_     :					List[str]   	=self.numpy_to_pil(_lowerCamelCase   )

      if not return_dict:
       return (image,)

      return ImagePipelineOutput(images=_lowerCamelCase   )


"""simple docstring"""






import copy
import tempfile
import unittest

from huggingface_hub import HfFolder, delete_repo
from parameterized import parameterized
from requests.exceptions import HTTPError

from transformers import AutoConfig, GenerationConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test





class 				a__					(						unittest.TestCase       ):



     @parameterized.expand([(None,), ('foo.json',)]   )
     def  lowerCamelCase_						(							self :Union[str, Any]      ,  _lowerCamelCase :Any   ):







      '''simple docstring'''



      UpperCamelCase_     :					str   	=GenerationConfig(
          do_sample=_lowerCamelCase      ,  temperature=0.7      ,  length_penalty=1.0      ,  bad_words_ids=[[1, 2, 3], [4, 5]]      ,  )
      with tempfile.TemporaryDirectory() as tmp_dir:
       config.save_pretrained(_lowerCamelCase      ,  config_name=_lowerCamelCase   )
       UpperCamelCase_     :					Union[str, Any]   	=GenerationConfig.from_pretrained(_lowerCamelCase      ,  config_name=_lowerCamelCase   )

      # Checks parameters that were specified
      self.assertEqual(loaded_config.do_sample      ,  _lowerCamelCase   )
      self.assertEqual(loaded_config.temperature      ,  0.7   )
      self.assertEqual(loaded_config.length_penalty      ,  1.0   )
      self.assertEqual(loaded_config.bad_words_ids      ,  [[1, 2, 3], [4, 5]]   )

      # Checks parameters that were not specified (defaults)
      self.assertEqual(loaded_config.top_k      ,  50   )
      self.assertEqual(loaded_config.max_length      ,  20   )
      self.assertEqual(loaded_config.max_time      ,  _lowerCamelCase   )



     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Optional[int]   	=AutoConfig.from_pretrained('gpt2'   )
      UpperCamelCase_     :					Tuple   	=GenerationConfig.from_model_config(_lowerCamelCase   )
      UpperCamelCase_     :					Optional[int]   	=GenerationConfig()

      # The generation config has loaded a few non-default parameters from the model config
      self.assertNotEqual(_lowerCamelCase      ,  _lowerCamelCase   )

      # One of those parameters is eos_token_id -- check if it matches
      self.assertNotEqual(generation_config_from_model.eos_token_id      ,  default_generation_config.eos_token_id   )
      self.assertEqual(generation_config_from_model.eos_token_id      ,  model_config.eos_token_id   )



     def  lowerCamelCase_						(							self :Tuple   ):







      '''simple docstring'''



      UpperCamelCase_     :					str   	=GenerationConfig()
      UpperCamelCase_     :					Optional[Any]   	={
          'max_new_tokens': 1_024,
          'foo': 'bar',
      }
      UpperCamelCase_     :					Optional[Any]   	=copy.deepcopy(_lowerCamelCase   )
      UpperCamelCase_     :					Optional[int]   	=generation_config.update(**_lowerCamelCase   )

      # update_kwargs was not modified (no side effects)
      self.assertEqual(_lowerCamelCase      ,  _lowerCamelCase   )

      # update_kwargs was used to update the config on valid attributes
      self.assertEqual(generation_config.max_new_tokens      ,  1_024   )

      # `.update()` returns a dictionary of unused kwargs
      self.assertEqual(_lowerCamelCase      ,  {'foo': 'bar'}   )



     def  lowerCamelCase_						(							self :Union[str, Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Union[str, Any]   	=GenerationConfig()
      UpperCamelCase_     :					Union[str, Any]   	='bar'

      with tempfile.TemporaryDirectory('test-generation-config'   ) as tmp_dir:
       generation_config.save_pretrained(_lowerCamelCase   )

       UpperCamelCase_     :					Dict   	=GenerationConfig.from_pretrained(_lowerCamelCase   )
      # update_kwargs was used to update the config on valid attributes
      self.assertEqual(new_config.foo      ,  'bar'   )

      UpperCamelCase_     :					List[str]   	=GenerationConfig.from_model_config(_lowerCamelCase   )
      assert not hasattr(_lowerCamelCase      ,  'foo'   )  # no new kwargs should be initialized if from config




     def  lowerCamelCase_						(							self :Optional[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[str]   	=GenerationConfig()
      self.assertEqual(default_config.temperature      ,  1.0   )
      self.assertEqual(default_config.do_sample      ,  _lowerCamelCase   )
      self.assertEqual(default_config.num_beams      ,  1   )

      UpperCamelCase_     :					Tuple   	=GenerationConfig(
          do_sample=_lowerCamelCase      ,  temperature=0.7      ,  length_penalty=1.0      ,  bad_words_ids=[[1, 2, 3], [4, 5]]      ,  )
      self.assertEqual(config.temperature      ,  0.7   )
      self.assertEqual(config.do_sample      ,  _lowerCamelCase   )
      self.assertEqual(config.num_beams      ,  1   )

      with tempfile.TemporaryDirectory() as tmp_dir:
       config.save_pretrained(_lowerCamelCase   )
       UpperCamelCase_     :					str   	=GenerationConfig.from_pretrained(_lowerCamelCase      ,  temperature=1.0   )

      self.assertEqual(loaded_config.temperature      ,  1.0   )
      self.assertEqual(loaded_config.do_sample      ,  _lowerCamelCase   )
      self.assertEqual(loaded_config.num_beams      ,  1   )  # default value





@is_staging_test
class 				a__					(						unittest.TestCase       ):



     @classmethod
     def  lowerCamelCase_						(							cls :Optional[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Dict   	=TOKEN
      HfFolder.save_token(_lowerCamelCase   )



     @classmethod
     def  lowerCamelCase_						(							cls :List[Any]   ):







      '''simple docstring'''



      try:
       delete_repo(token=cls._token      ,  repo_id='test-generation-config'   )
      except HTTPError:
       pass

      try:
       delete_repo(token=cls._token      ,  repo_id='valid_org/test-generation-config-org'   )
      except HTTPError:
       pass



     def  lowerCamelCase_						(							self :List[Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[str]   	=GenerationConfig(
          do_sample=_lowerCamelCase      ,  temperature=0.7      ,  length_penalty=1.0      ,  )
      config.push_to_hub('test-generation-config'      ,  use_auth_token=self._token   )

      UpperCamelCase_     :					List[Any]   	=GenerationConfig.from_pretrained(f'''{USER}/test-generation-config'''   )
      for k, v in config.to_dict().items():
       if k != "transformers_version":
        self.assertEqual(_lowerCamelCase      ,  getattr(_lowerCamelCase      ,  _lowerCamelCase   )   )

        # Reset repo
      delete_repo(token=self._token      ,  repo_id='test-generation-config'   )

      # Push to hub via save_pretrained
      with tempfile.TemporaryDirectory() as tmp_dir:
       config.save_pretrained(
           _lowerCamelCase      ,  repo_id='test-generation-config'      ,  push_to_hub=_lowerCamelCase      ,  use_auth_token=self._token   )

      UpperCamelCase_     :					Any   	=GenerationConfig.from_pretrained(f'''{USER}/test-generation-config'''   )
      for k, v in config.to_dict().items():
       if k != "transformers_version":
        self.assertEqual(_lowerCamelCase      ,  getattr(_lowerCamelCase      ,  _lowerCamelCase   )   )




     def  lowerCamelCase_						(							self :Union[str, Any]   ):







      '''simple docstring'''



      UpperCamelCase_     :					Union[str, Any]   	=GenerationConfig(
          do_sample=_lowerCamelCase      ,  temperature=0.7      ,  length_penalty=1.0      ,  )
      config.push_to_hub('valid_org/test-generation-config-org'      ,  use_auth_token=self._token   )

      UpperCamelCase_     :					int   	=GenerationConfig.from_pretrained('valid_org/test-generation-config-org'   )
      for k, v in config.to_dict().items():
       if k != "transformers_version":
        self.assertEqual(_lowerCamelCase      ,  getattr(_lowerCamelCase      ,  _lowerCamelCase   )   )

        # Reset repo
      delete_repo(token=self._token      ,  repo_id='valid_org/test-generation-config-org'   )

      # Push to hub via save_pretrained
      with tempfile.TemporaryDirectory() as tmp_dir:
       config.save_pretrained(
           _lowerCamelCase      ,  repo_id='valid_org/test-generation-config-org'      ,  push_to_hub=_lowerCamelCase      ,  use_auth_token=self._token   )

      UpperCamelCase_     :					Optional[Any]   	=GenerationConfig.from_pretrained('valid_org/test-generation-config-org'   )
      for k, v in config.to_dict().items():
       if k != "transformers_version":
        self.assertEqual(_lowerCamelCase      ,  getattr(_lowerCamelCase      ,  _lowerCamelCase   )   )
