

"""simple docstring"""


import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class     snake_case	(				_snake_case,							_snake_case,							_snake_case,							unittest.TestCase  ):
				SCREAMING_SNAKE_CASE_   :      List[str]         =  StableUnCLIPImgaImgPipeline
				SCREAMING_SNAKE_CASE_   :      Tuple         =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS
				SCREAMING_SNAKE_CASE_   :      Any         =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				SCREAMING_SNAKE_CASE_   :      Union[str, Any]         =  frozenset(
				    []  )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
				SCREAMING_SNAKE_CASE_   :      Optional[Any]         =  frozenset([]  )

				def 							lowercase_     (	self							:				Optional[Any])->   Dict:



						'''simple docstring'''


						__lowerCAmelCase:       List[Any]          =			3_2
						__lowerCAmelCase:       Tuple          =			embedder_hidden_size

						# image encoding components

						__lowerCAmelCase:       Optional[Any]          =			CLIPImageProcessor(crop_size=3_2     ,  size=3_2)

						torch.manual_seed(0)
						__lowerCAmelCase:       Optional[int]          =			CLIPVisionModelWithProjection(
						    CLIPVisionConfig(
						        hidden_size=_lowerCamelCase     ,  projection_dim=_lowerCamelCase     ,  num_hidden_layers=5     ,  num_attention_heads=4     ,  image_size=3_2     ,  intermediate_size=3_7     ,  patch_size=1     ,  ))

						# regular denoising components

						torch.manual_seed(0)
						__lowerCAmelCase:       int          =			StableUnCLIPImageNormalizer(embedding_dim=_lowerCamelCase)
						__lowerCAmelCase:       List[str]          =			DDPMScheduler(beta_schedule="squaredcos_cap_v2")

						torch.manual_seed(0)
						__lowerCAmelCase:       Tuple          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						torch.manual_seed(0)
						__lowerCAmelCase:       str          =			CLIPTextModel(
						    CLIPTextConfig(
						        bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=_lowerCamelCase     ,  projection_dim=3_2     ,  intermediate_size=3_7     ,  layer_norm_eps=1e-05     ,  num_attention_heads=4     ,  num_hidden_layers=5     ,  pad_token_id=1     ,  vocab_size=1_0_0_0     ,  ))

						torch.manual_seed(0)
						__lowerCAmelCase:       int          =			UNetaDConditionModel(
						    sample_size=3_2     ,  in_channels=4     ,  out_channels=4     ,  down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")     ,  up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")     ,  block_out_channels=(3_2, 6_4)     ,  attention_head_dim=(2, 4)     ,  class_embed_type="projection"     ,  projection_class_embeddings_input_dim=embedder_projection_dim * 2     ,  cross_attention_dim=_lowerCamelCase     ,  layers_per_block=1     ,  upcast_attention=_lowerCamelCase     ,  use_linear_projection=_lowerCamelCase     ,  )

						torch.manual_seed(0)
						__lowerCAmelCase:       Tuple          =			DDIMScheduler(
						    beta_schedule="scaled_linear"     ,  beta_start=0.00085     ,  beta_end=0.012     ,  prediction_type="v_prediction"     ,  set_alpha_to_one=_lowerCamelCase     ,  steps_offset=1     ,  )

						torch.manual_seed(0)
						__lowerCAmelCase:       Union[str, Any]          =			AutoencoderKL()

						__lowerCAmelCase:       Optional[Any]          =			{
						    # image encoding components
						    '''feature_extractor''': feature_extractor,
						    '''image_encoder''': image_encoder.eval(),
						    # image noising components
						    '''image_normalizer''': image_normalizer.eval(),
						    '''image_noising_scheduler''': image_noising_scheduler,
						    # regular denoising components
						    '''tokenizer''': tokenizer,
						    '''text_encoder''': text_encoder.eval(),
						    '''unet''': unet.eval(),
						    '''scheduler''': scheduler,
						    '''vae''': vae.eval(),
						}

						return components

				def 							lowercase_     (	self							:				Any     ,  UpperCamelCase__							:				Optional[int]     ,  UpperCamelCase__							:				Dict=0     ,  UpperCamelCase__							:				Dict=True)->   List[Any]:



						'''simple docstring'''


						if str(_lowerCamelCase).startswith("mps"):
								__lowerCAmelCase:       Optional[int]          =			torch.manual_seed(_lowerCamelCase)
						else:
								__lowerCAmelCase:       List[Any]          =			torch.Generator(device=_lowerCamelCase).manual_seed(_lowerCamelCase)

						__lowerCAmelCase:       List[str]          =			floats_tensor((1, 3, 3_2, 3_2)     ,  rng=random.Random(_lowerCamelCase)).to(_lowerCamelCase)

						if pil_image:
								__lowerCAmelCase:       Any          =			input_image * 0.5 + 0.5
								__lowerCAmelCase:       Optional[int]          =			input_image.clamp(0     ,  1)
								__lowerCAmelCase:       List[str]          =			input_image.cpu().permute(0     ,  2     ,  3     ,  1).float().numpy()
								__lowerCAmelCase:       List[str]          =			DiffusionPipeline.numpy_to_pil(_lowerCamelCase)[0]

						return {
						    "prompt": "An anime racoon running a marathon",
						    "image": input_image,
						    "generator": generator,
						    "num_inference_steps": 2,
						    "output_type": "np",
						}

				@skip_mps
				def 							lowercase_     (	self							:				Optional[Any])->   Optional[int]:



						'''simple docstring'''


						__lowerCAmelCase:       int          =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:       Any          =			self.get_dummy_components()
						__lowerCAmelCase:       Tuple          =			StableUnCLIPImgaImgPipeline(**_lowerCamelCase)
						__lowerCAmelCase:       List[Any]          =			sd_pipe.to(_lowerCamelCase)
						sd_pipe.set_progress_bar_config(disable=_lowerCamelCase)

						__lowerCAmelCase:       int          =			self.get_dummy_inputs(_lowerCamelCase)
						inputs.update({"image_embeds": None})
						__lowerCAmelCase:       Dict          =			sd_pipe(**_lowerCamelCase).images
						__lowerCAmelCase:       Union[str, Any]          =			image[0, -3:, -3:, -1]

						assert image.shape == (1, 3_2, 3_2, 3)
						__lowerCAmelCase:       List[str]          =			np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078])

						assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3

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



						'''simple docstring'''


						__lowerCAmelCase:       Dict          =			torch_device in ['''cpu''', '''mps''']

						self._test_attention_slicing_forward_pass(test_max_difference=_lowerCamelCase)

				def 							lowercase_     (	self							:				Optional[Any])->   Dict:



						'''simple docstring'''


						__lowerCAmelCase:       List[str]          =			torch_device in ['''cpu''', '''mps''']

						self._test_inference_batch_single_identical(test_max_difference=_lowerCamelCase)




				@unittest.skipIf(
				    torch_device != "cuda" or not is_xformers_available()     ,  reason="XFormers attention is only available with CUDA and `xformers` installed"     ,  )
				def 							lowercase_     (	self							:				Optional[Any])->   Optional[Any]:



						'''simple docstring'''


						self._test_xformers_attention_forwardGenerator_pass(test_max_difference=_lowerCamelCase)





@slow
@require_torch_gpu
class     snake_case	(				unittest.TestCase  ):

				def 							lowercase_     (	self							:				Tuple)->   Optional[Any]:



						'''simple docstring'''


						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()

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



						'''simple docstring'''


						__lowerCAmelCase:       Any          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png")

						__lowerCAmelCase:       Any          =			load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy")

						__lowerCAmelCase:       Optional[int]          =			StableUnCLIPImgaImgPipeline.from_pretrained(
						    "fusing/stable-unclip-2-1-l-img2img"     ,  torch_dtype=torch.floataa)
						pipe.to(_lowerCamelCase)
						pipe.set_progress_bar_config(disable=_lowerCamelCase)
						# stable unclip will oom when integration tests are run on a V100,
						# so turn on memory savings
						pipe.enable_attention_slicing()
						pipe.enable_sequential_cpu_offload()

						__lowerCAmelCase:       Tuple          =			torch.Generator(device="cpu").manual_seed(0)
						__lowerCAmelCase:       List[str]          =			pipe(_lowerCamelCase     ,  "anime turle"     ,  generator=_lowerCamelCase     ,  output_type="np")

						__lowerCAmelCase:       List[Any]          =			output.images[0]

						assert image.shape == (7_6_8, 7_6_8, 3)

						assert_mean_pixel_difference(_lowerCamelCase     ,  _lowerCamelCase)

				def 							lowercase_     (	self							:				Union[str, Any])->   Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       Optional[Any]          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png")

						__lowerCAmelCase:       Optional[Any]          =			load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy")

						__lowerCAmelCase:       List[str]          =			StableUnCLIPImgaImgPipeline.from_pretrained(
						    "fusing/stable-unclip-2-1-h-img2img"     ,  torch_dtype=torch.floataa)
						pipe.to(_lowerCamelCase)
						pipe.set_progress_bar_config(disable=_lowerCamelCase)
						# stable unclip will oom when integration tests are run on a V100,
						# so turn on memory savings
						pipe.enable_attention_slicing()
						pipe.enable_sequential_cpu_offload()

						__lowerCAmelCase:       Tuple          =			torch.Generator(device="cpu").manual_seed(0)
						__lowerCAmelCase:       Optional[int]          =			pipe(_lowerCamelCase     ,  "anime turle"     ,  generator=_lowerCamelCase     ,  output_type="np")

						__lowerCAmelCase:       int          =			output.images[0]

						assert image.shape == (7_6_8, 7_6_8, 3)

						assert_mean_pixel_difference(_lowerCamelCase     ,  _lowerCamelCase)




				def 							lowercase_     (	self							:				Any)->   Tuple:



						'''simple docstring'''


						__lowerCAmelCase:       List[str]          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png")

						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()

						__lowerCAmelCase:       int          =			StableUnCLIPImgaImgPipeline.from_pretrained(
						    "fusing/stable-unclip-2-1-h-img2img"     ,  torch_dtype=torch.floataa)
						__lowerCAmelCase:       Any          =			pipe.to(_lowerCamelCase)
						pipe.set_progress_bar_config(disable=_lowerCamelCase)
						pipe.enable_attention_slicing()
						pipe.enable_sequential_cpu_offload()

						__lowerCAmelCase:       str          =			pipe(
						    _lowerCamelCase     ,  "anime turtle"     ,  num_inference_steps=2     ,  output_type="np"     ,  )

						__lowerCAmelCase:       List[str]          =			torch.cuda.max_memory_allocated()
						# make sure that less than 7 GB is allocated
						assert mem_bytes < 7 * 1_0**9






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				:    Any       				=     logging.get_logger(__name__)  # pylint: disable=invalid-name

snake_case				:    Union[str, Any]       				=     '''
    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 				__lowerCamelCase						( UpperCAmelCase_						:		List[Any] ,  UpperCAmelCase_						:		Optional[int] ,  UpperCAmelCase_						:		List[str]=8						):




      """simple docstring"""





      a						:List[str] 							=       height // scale_factor**2
      if height % scale_factor**2 != 0:
            new_height += 1
      a						:int 							=       width // scale_factor**2
      if width % scale_factor**2 != 0:
            new_width += 1
      return new_height * scale_factor, new_width * scale_factor







class 		_snake_case  (      _snake_case							):




 def __init__(		self					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  ):
       super().__init__()

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




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase    ):
       if latents is None:
             a						:str 							=       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}'''    )
             a						:Any 							=       latents.to(_lowerCamelCase    )

       a						:Dict 							=       latents * scheduler.init_noise_sigma
       return latents




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase=0    ):
       if is_accelerate_available():
             from accelerate import cpu_offload
       else:
             raise ImportError('''Please install accelerate via `pip install accelerate`'''    )

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

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




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase=0    ):
       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						: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)

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

       # We'll offload the last model manually.
       a						:str 							=       hook




 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def        SCREAMING_SNAKE_CASE__				(		self    ):
       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					,  _lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase = 512					,  _lowerCamelCase = 512					,  _lowerCamelCase = 100					,  _lowerCamelCase = 4.0					,  _lowerCamelCase = 1					,  _lowerCamelCase = None					,  _lowerCamelCase = None					,  _lowerCamelCase = "pil"					,  _lowerCamelCase = True					,  ):
       a						:int 							=       self._execution_device

       a						:Optional[Any] 							=       guidance_scale > 1.0

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

       if do_classifier_free_guidance:
             a						:Union[str, Any] 							=       image_embeds.repeat_interleave(_lowerCamelCase					,  dim=0    )
             a						:Optional[int] 							=       negative_image_embeds.repeat_interleave(_lowerCamelCase					,  dim=0    )

       a						:Optional[int] 							=       torch.cat([negative_image_embeds, image_embeds]					,  dim=0    ).to(dtype=self.unet.dtype					,  device=_lowerCamelCase    )

       self.scheduler.set_timesteps(_lowerCamelCase					,  device=_lowerCamelCase    )
       a						:Optional[Any] 							=       self.scheduler.timesteps

       a						:List[str] 							=       self.unet.config.in_channels

       a							,				a						:str 							=       downscale_height_and_width(_lowerCamelCase					,  _lowerCamelCase					,  self.movq_scale_factor    )

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

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

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

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

       if output_type == "pil":
             a						:str 							=       self.numpy_to_pil(_lowerCamelCase    )

       if not return_dict:
             return (image,)

       return ImagePipelineOutput(images=_lowerCamelCase    )





