



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


__snake_case		:List[Any]   	=  logging.get_logger(__name__)  # pylint: disable=invalid-name

__snake_case		:List[str]   	=  '''
    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   __snake_case			(    _UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase=8       ):
    __a		 =       height // scale_factor**2
    if height % scale_factor**2 != 0:
        new_height += 1
    __a		 =       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							:  Dict  ,	__SCREAMING_SNAKE_CASE							:  UNetaDConditionModel  ,	__SCREAMING_SNAKE_CASE							:  DDPMScheduler  ,	__SCREAMING_SNAKE_CASE							:  VQModel  ,	):


           '''simple docstring'''

           super().__init__()

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



       def 	_lowerCamelCase    (	self							:  Any  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  Optional[Any]  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  List[str]  ,	__SCREAMING_SNAKE_CASE							:  List[str]):


           '''simple docstring'''

           if latents is None:
               __a		 =       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}')
               __a		 =       latents.to(__SCREAMING_SNAKE_CASE)

           __a		 =       latents * scheduler.init_noise_sigma
           return latents



       def 	_lowerCamelCase    (	self							:  List[str]  ,	__SCREAMING_SNAKE_CASE							:  Optional[Any]=0):


           '''simple docstring'''

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

           __a		 =       torch.device(F'cuda:{gpu_id}')

           __a		 =       [
               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 	_lowerCamelCase    (	self							:  int  ,	__SCREAMING_SNAKE_CASE							:  str=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.''')

           __a		 =       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)

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

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



       @property
       # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
       def 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''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							:  str  ,	__SCREAMING_SNAKE_CASE							:  Union[torch.FloatTensor, List[torch.FloatTensor]]  ,	__SCREAMING_SNAKE_CASE							:  Union[torch.FloatTensor, List[torch.FloatTensor]]  ,	__SCREAMING_SNAKE_CASE							:  int = 512  ,	__SCREAMING_SNAKE_CASE							:  int = 512  ,	__SCREAMING_SNAKE_CASE							:  int = 100  ,	__SCREAMING_SNAKE_CASE							:  float = 4.0  ,	__SCREAMING_SNAKE_CASE							:  int = 1  ,	__SCREAMING_SNAKE_CASE							:  Optional[Union[torch.Generator, List[torch.Generator]]] = None  ,	__SCREAMING_SNAKE_CASE							:  Optional[torch.FloatTensor] = None  ,	__SCREAMING_SNAKE_CASE							:  Optional[str] = "pil"  ,	__SCREAMING_SNAKE_CASE							:  bool = True  ,	):


           '''simple docstring'''

           __a		 =       self._execution_device

           __a		 =       guidance_scale > 1.0

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

           if do_classifier_free_guidance:
               __a		 =       image_embeds.repeat_interleave(__SCREAMING_SNAKE_CASE  ,	dim=0)
               __a		 =       negative_image_embeds.repeat_interleave(__SCREAMING_SNAKE_CASE  ,	dim=0)

           __a		 =       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)
           __a		 =       self.scheduler.timesteps

           __a		 =       self.unet.config.in_channels

           __a ,     __a		 =       downscale_height_and_width(__SCREAMING_SNAKE_CASE  ,	__SCREAMING_SNAKE_CASE  ,	self.movq_scale_factor)

           # create initial latent
           __a		 =       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
               __a		 =       torch.cat([latents] * 2) if do_classifier_free_guidance else latents

               __a		 =       {'''image_embeds''': image_embeds}
               __a		 =       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:
                   __a ,     __a		 =       noise_pred.split(latents.shape[1]  ,	dim=1)
                   __a ,     __a		 =       noise_pred.chunk(2)
                   __a ,     __a		 =       variance_pred.chunk(2)
                   __a		 =       noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                   __a		 =       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"]
               ):
                   __a ,     __a		 =       noise_pred.split(latents.shape[1]  ,	dim=1)

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

           if output_type == "pil":
               __a		 =       self.numpy_to_pil(__SCREAMING_SNAKE_CASE)

           if not return_dict:
               return (image,)

           return ImagePipelineOutput(images=__SCREAMING_SNAKE_CASE)










import os
import sys
import unittest


__snake_case		:Union[str, Any]   	=  os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, '''utils'''))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
__snake_case		:List[str]   	=  os.path.join(git_repo_path, '''src''', '''transformers''')

__snake_case		:Any   	=  '''
{0} = None
'''

__snake_case		:Dict   	=  '''
class {0}(metaclass=DummyObject):
    _backends = {1}

    def __init__(self, *args, **kwargs):
        requires_backends(self, {1})
'''


__snake_case		:str   	=  '''
def {0}(*args, **kwargs):
    requires_backends({0}, {1})
'''

class 					_A (			unittest.TestCase					):



       def 	_lowerCamelCase    (	self							:  List[Any]):


           '''simple docstring'''

           __a		 =       find_backend('''    _import_structure["models.albert"].append("AlbertTokenizerFast")''')
           self.assertIsNone(__SCREAMING_SNAKE_CASE)

           __a		 =       find_backend('''    if not is_tokenizers_available():''')
           self.assertEqual(__SCREAMING_SNAKE_CASE  ,	'''tokenizers''')

           __a		 =       find_backend('''    if not is_tensorflow_text_available():''')
           self.assertEqual(__SCREAMING_SNAKE_CASE  ,	'''tensorflow_text''')

           __a		 =       find_backend('''    if not (is_sentencepiece_available() and is_tokenizers_available()):''')
           self.assertEqual(__SCREAMING_SNAKE_CASE  ,	'''sentencepiece_and_tokenizers''')

           __a		 =       find_backend(
               '''    if not (is_sentencepiece_available() and is_tensorflow_text_available()):''')
           self.assertEqual(__SCREAMING_SNAKE_CASE  ,	'''sentencepiece_and_tensorflow_text''')

           __a		 =       find_backend(
               '''    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):''')
           self.assertEqual(__SCREAMING_SNAKE_CASE  ,	'''sentencepiece_and_tokenizers_and_vision''')



       def 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''simple docstring'''

           __a		 =       read_init()
           # We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
           self.assertIn('''torch'''  ,	__SCREAMING_SNAKE_CASE)
           self.assertIn('''tensorflow_text'''  ,	__SCREAMING_SNAKE_CASE)
           self.assertIn('''sentencepiece_and_tokenizers'''  ,	__SCREAMING_SNAKE_CASE)

           # Likewise, we can't assert on the exact content of a key
           self.assertIn('''BertModel'''  ,	objects['''torch'''])
           self.assertIn('''TFBertModel'''  ,	objects['''tf'''])
           self.assertIn('''FlaxBertModel'''  ,	objects['''flax'''])
           self.assertIn('''BertModel'''  ,	objects['''torch'''])
           self.assertIn('''TFBertTokenizer'''  ,	objects['''tensorflow_text'''])
           self.assertIn('''convert_slow_tokenizer'''  ,	objects['''sentencepiece_and_tokenizers'''])



       def 	_lowerCamelCase    (	self							:  Any):


           '''simple docstring'''

           __a		 =       create_dummy_object('''CONSTANT'''  ,	'''\'torch\'''')
           self.assertEqual(__SCREAMING_SNAKE_CASE  ,	'''\nCONSTANT = None\n''')

           __a		 =       create_dummy_object('''function'''  ,	'''\'torch\'''')
           self.assertEqual(
               __SCREAMING_SNAKE_CASE  ,	'''\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n''')

           __a		 =       '''
class FakeClass(metaclass=DummyObject):
    _backends = \'torch\'

    def __init__(self, *args, **kwargs):
        requires_backends(self, \'torch\')
'''
           __a		 =       create_dummy_object('''FakeClass'''  ,	'''\'torch\'''')
           self.assertEqual(__SCREAMING_SNAKE_CASE  ,	__SCREAMING_SNAKE_CASE)





       def 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''simple docstring'''

           __a		 =       '''# This file is autogenerated by the command `make fix-copies`, do not edit.
from ..utils import DummyObject, requires_backends


CONSTANT = None


def function(*args, **kwargs):
    requires_backends(function, ["torch"])


class FakeClass(metaclass=DummyObject):
    _backends = ["torch"]

    def __init__(self, *args, **kwargs):
        requires_backends(self, ["torch"])
'''
           __a		 =       create_dummy_files({'''torch''': ['''CONSTANT''', '''function''', '''FakeClass''']})
           self.assertEqual(dummy_files['''torch''']  ,	__SCREAMING_SNAKE_CASE)





