




"""simple docstring"""

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS


enable_full_determinism()





class     lowercase_	(					unittest.TestCase							):



 '''simple docstring'''
 UpperCAmelCase   :					Optional[int]				  =     StableDiffusionLDMaDPipeline
 UpperCAmelCase   :					List[Any]				  =     TEXT_TO_IMAGE_PARAMS
 UpperCAmelCase   :					Tuple				  =     TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCAmelCase   :					Optional[Any]				  =     TEXT_TO_IMAGE_IMAGE_PARAMS

 def       lowerCAmelCase_							(			self	:	Optional[Any]		):
       torch.manual_seed(0		)
       _A        =       UNetaDConditionModel(
           block_out_channels=(32, 64)   ,							layers_per_block=2   ,							sample_size=32   ,							in_channels=4   ,							out_channels=4   ,							down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')   ,							up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')   ,							cross_attention_dim=32   ,							)
       _A        =       DDIMScheduler(
           beta_start=0.0_0085   ,							beta_end=0.012   ,							beta_schedule='scaled_linear'   ,							clip_sample=_UpperCAmelCase   ,							set_alpha_to_one=_UpperCAmelCase   ,							)
       torch.manual_seed(0		)
       _A        =       AutoencoderKL(
           block_out_channels=[32, 64]   ,							in_channels=6   ,							out_channels=6   ,							down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']   ,							up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']   ,							latent_channels=4   ,							)
       torch.manual_seed(0		)
       _A        =       CLIPTextConfig(
           bos_token_id=0   ,							eos_token_id=2   ,							hidden_size=32   ,							intermediate_size=37   ,							layer_norm_eps=1E-0_5   ,							num_attention_heads=4   ,							num_hidden_layers=5   ,							pad_token_id=1   ,							vocab_size=1_000   ,							)
       _A        =       CLIPTextModel(_UpperCAmelCase		)
       _A        =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'		)

       _A        =       {
           'unet': unet,
           'scheduler': scheduler,
           'vae': vae,
           'text_encoder': text_encoder,
           'tokenizer': tokenizer,
           'safety_checker': None,
           'feature_extractor': None,
       }
       return components

 def       lowerCAmelCase_							(			self	:	Optional[int]   ,							_UpperCAmelCase	:	List[str]   ,							_UpperCAmelCase	:	int=0		):
       if str(_UpperCAmelCase		).startswith('mps'		):
             _A        =       torch.manual_seed(_UpperCAmelCase		)
       else:
             _A        =       torch.Generator(device=_UpperCAmelCase		).manual_seed(_UpperCAmelCase		)
       _A        =       {
           'prompt': 'A painting of a squirrel eating a burger',
           'generator': generator,
           'num_inference_steps': 2,
           'guidance_scale': 6.0,
           'output_type': 'numpy',
       }
       return inputs

 def       lowerCAmelCase_							(			self	:	int		):
       _A        =       'cpu'  # ensure determinism for the device-dependent torch.Generator

       _A        =       self.get_dummy_components()
       _A        =       StableDiffusionLDMaDPipeline(**_UpperCAmelCase		)
       _A        =       ldmad_pipe.to(_UpperCAmelCase		)
       ldmad_pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       ldmad_pipe(**_UpperCAmelCase		)
       _A       ,   _A        =       output.rgb, output.depth

       _A        =       rgb[0, -3:, -3:, -1]
       _A        =       depth[0, -3:, -1]

       assert rgb.shape == (1, 64, 64, 3)
       assert depth.shape == (1, 64, 64)

       _A        =       np.array(
           [0.3733_8176, 0.7_0247, 0.7420_3193, 0.5164_3604, 0.5825_6793, 0.6093_2136, 0.418_1095, 0.4835_5877, 0.4653_5262]		)
       _A        =       np.array([103.4_6727, 85.81_2004, 87.84_9236]		)

       assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb		).max() < 1E-2
       assert np.abs(image_slice_depth.flatten() - expected_slice_depth		).max() < 1E-2

 def       lowerCAmelCase_							(			self	:	Optional[int]		):
       _A        =       self.get_dummy_components()
       _A        =       StableDiffusionLDMaDPipeline(**_UpperCAmelCase		)
       _A        =       ldmad_pipe.to(_UpperCAmelCase		)
       ldmad_pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       3 * [inputs['prompt']]

       # forward
       _A        =       ldmad_pipe(**_UpperCAmelCase		)
       _A       ,   _A        =       output.rgb, output.depth
       _A        =       rgb_slice_a[0, -3:, -3:, -1]
       _A        =       depth_slice_a[0, -3:, -1]

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       3 * [inputs.pop('prompt'		)]

       _A        =       ldmad_pipe.tokenizer(
           _UpperCAmelCase   ,							padding='max_length'   ,							max_length=ldmad_pipe.tokenizer.model_max_length   ,							truncation=_UpperCAmelCase   ,							return_tensors='pt'   ,							)
       _A        =       text_inputs['input_ids'].to(_UpperCAmelCase		)

       _A        =       ldmad_pipe.text_encoder(_UpperCAmelCase		)[0]

       _A        =       prompt_embeds

       # forward
       _A        =       ldmad_pipe(**_UpperCAmelCase		)
       _A       ,   _A        =       output.rgb, output.depth
       _A        =       rgb_slice_a[0, -3:, -3:, -1]
       _A        =       depth_slice_a[0, -3:, -1]

       assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()		).max() < 1E-4
       assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()		).max() < 1E-4



 def       lowerCAmelCase_							(			self	:	Dict		):
       _A        =       'cpu'  # ensure determinism for the device-dependent torch.Generator
       _A        =       self.get_dummy_components()
       _A        =       PNDMScheduler(skip_prk_steps=_UpperCAmelCase		)
       _A        =       StableDiffusionLDMaDPipeline(**_UpperCAmelCase		)
       _A        =       ldmad_pipe.to(_UpperCAmelCase		)
       ldmad_pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       'french fries'
       _A        =       ldmad_pipe(**_UpperCAmelCase   ,							negative_prompt=_UpperCAmelCase		)

       _A       ,   _A        =       output.rgb, output.depth
       _A        =       rgb[0, -3:, -3:, -1]
       _A        =       depth[0, -3:, -1]

       assert rgb.shape == (1, 64, 64, 3)
       assert depth.shape == (1, 64, 64)

       _A        =       np.array(
           [0.3_7044, 0.7181_1503, 0.722_3251, 0.4860_3675, 0.563_8391, 0.636_4948, 0.4283_3704, 0.490_1315, 0.4792_6217]		)
       _A        =       np.array([107.8_4738, 84.6_2802, 89.96_2135]		)
       assert np.abs(rgb_slice.flatten() - expected_slice_rgb		).max() < 1E-2
       assert np.abs(depth_slice.flatten() - expected_slice_depth		).max() < 1E-2





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



 '''simple docstring'''

 def       lowerCAmelCase_							(			self	:	Dict		):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()

 def       lowerCAmelCase_							(			self	:	List[Any]   ,							_UpperCAmelCase	:	Union[str, Any]   ,							_UpperCAmelCase	:	str="cpu"   ,							_UpperCAmelCase	:	Any=torch.floataa   ,							_UpperCAmelCase	:	str=0		):
       _A        =       torch.Generator(device=_UpperCAmelCase		).manual_seed(_UpperCAmelCase		)
       _A        =       np.random.RandomState(_UpperCAmelCase		).standard_normal((1, 4, 64, 64)		)
       _A        =       torch.from_numpy(_UpperCAmelCase		).to(device=_UpperCAmelCase   ,							dtype=_UpperCAmelCase		)
       _A        =       {
           'prompt': 'a photograph of an astronaut riding a horse',
           'latents': latents,
           'generator': generator,
           'num_inference_steps': 3,
           'guidance_scale': 7.5,
           'output_type': 'numpy',
       }
       return inputs



 def       lowerCAmelCase_							(			self	:	Optional[int]		):
       _A        =       StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'		)
       _A        =       ldmad_pipe.to(_UpperCAmelCase		)
       ldmad_pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       self.get_inputs(_UpperCAmelCase		)
       _A        =       ldmad_pipe(**_UpperCAmelCase		)
       _A       ,   _A        =       output.rgb, output.depth
       _A        =       rgb[0, -3:, -3:, -1].flatten()
       _A        =       rgb[0, -3:, -1].flatten()

       assert rgb.shape == (1, 512, 512, 3)
       assert depth.shape == (1, 512, 512)

       _A        =       np.array(
           [0.5380_5465, 0.5670_7305, 0.548_6515, 0.5701_2236, 0.581_4511, 0.5625_3487, 0.5484_3014, 0.5509_2263, 0.645_9706]		)
       _A        =       np.array(
           [0.926_3781, 0.667_8672, 0.548_6515, 0.9220_2145, 0.6783_1135, 0.5625_3487, 0.924_1694, 0.755_1478, 0.645_9706]		)
       assert np.abs(rgb_slice - expected_slice_rgb		).max() < 3E-3
       assert np.abs(depth_slice - expected_slice_depth		).max() < 3E-3





@nightly
@require_torch_gpu
class     lowercase_	(					unittest.TestCase							):



 '''simple docstring'''

 def       lowerCAmelCase_							(			self	:	str		):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()

 def       lowerCAmelCase_							(			self	:	Dict   ,							_UpperCAmelCase	:	List[str]   ,							_UpperCAmelCase	:	Union[str, Any]="cpu"   ,							_UpperCAmelCase	:	Dict=torch.floataa   ,							_UpperCAmelCase	:	List[str]=0		):
       _A        =       torch.Generator(device=_UpperCAmelCase		).manual_seed(_UpperCAmelCase		)
       _A        =       np.random.RandomState(_UpperCAmelCase		).standard_normal((1, 4, 64, 64)		)
       _A        =       torch.from_numpy(_UpperCAmelCase		).to(device=_UpperCAmelCase   ,							dtype=_UpperCAmelCase		)
       _A        =       {
           'prompt': 'a photograph of an astronaut riding a horse',
           'latents': latents,
           'generator': generator,
           'num_inference_steps': 50,
           'guidance_scale': 7.5,
           'output_type': 'numpy',
       }
       return inputs

 def       lowerCAmelCase_							(			self	:	List[Any]		):
       _A        =       StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'		).to(_UpperCAmelCase		)
       ldmad_pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       self.get_inputs(_UpperCAmelCase		)
       _A        =       ldmad_pipe(**_UpperCAmelCase		)
       _A       ,   _A        =       output.rgb, output.depth

       _A        =       0.49_5586
       _A        =       0.3379_5515
       _A        =       112.4_8518
       _A        =       98.48_9746
       assert np.abs(expected_rgb_mean - rgb.mean()		) < 1E-3
       assert np.abs(expected_rgb_std - rgb.std()		) < 1E-3
       assert np.abs(expected_depth_mean - depth.mean()		) < 1E-3
       assert np.abs(expected_depth_std - depth.std()		) < 1E-3



 def       lowerCAmelCase_							(			self	:	List[Any]		):
       _A        =       StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'		).to(_UpperCAmelCase		)
       ldmad_pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       self.get_inputs(_UpperCAmelCase		)
       _A        =       ldmad_pipe(**_UpperCAmelCase		)
       _A       ,   _A        =       output.rgb, output.depth

       _A        =       0.419_4127
       _A        =       0.3537_5586
       _A        =       0.563_8502
       _A        =       0.3468_6103

       assert rgb.shape == (1, 512, 512, 3)
       assert depth.shape == (1, 512, 512, 1)
       assert np.abs(expected_rgb_mean - rgb.mean()		) < 1E-3
       assert np.abs(expected_rgb_std - rgb.std()		) < 1E-3
       assert np.abs(expected_depth_mean - depth.mean()		) < 1E-3
       assert np.abs(expected_depth_std - depth.std()		) < 1E-3







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


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

a         =      '''
    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					( _snake_case							: Dict    ,  _snake_case							: List[Any]    ,  _snake_case							: List[Any]=8							)   -> Optional[int]:






      '''simple docstring'''


      _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     lowercase_	(					__lowerCAmelCase							):



 '''simple docstring'''

 def __init__(			self	:	Dict   ,							_UpperCAmelCase	:	UNetaDConditionModel   ,							_UpperCAmelCase	:	DDPMScheduler   ,							_UpperCAmelCase	:	VQModel   ,							):
       super().__init__()

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

 def       lowerCAmelCase_							(			self	:	Any   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	Tuple   ,							_UpperCAmelCase	:	Optional[Any]   ,							_UpperCAmelCase	:	Dict   ,							_UpperCAmelCase	:	List[str]   ,							_UpperCAmelCase	:	int		):
       if latents is None:
             _A        =       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}'''		)
             _A        =       latents.to(_UpperCAmelCase		)

       _A        =       latents * scheduler.init_noise_sigma
       return latents

 def       lowerCAmelCase_							(			self	:	Tuple   ,							_UpperCAmelCase	:	Tuple=0		):
       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(_UpperCAmelCase   ,							_UpperCAmelCase		)

 def       lowerCAmelCase_							(			self	:	Union[str, Any]   ,							_UpperCAmelCase	:	Union[str, Any]=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        =       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)

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

       # 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	:	int		):
       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	:	Union[str, Any]   ,							_UpperCAmelCase	:	Union[torch.FloatTensor, List[torch.FloatTensor]]   ,							_UpperCAmelCase	:	Union[torch.FloatTensor, List[torch.FloatTensor]]   ,							_UpperCAmelCase	:	int = 512   ,							_UpperCAmelCase	:	int = 512   ,							_UpperCAmelCase	:	int = 100   ,							_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   ,							):
       _A        =       self._execution_device

       _A        =       guidance_scale > 1.0

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

       if do_classifier_free_guidance:
             _A        =       image_embeds.repeat_interleave(_UpperCAmelCase   ,							dim=0		)
             _A        =       negative_image_embeds.repeat_interleave(_UpperCAmelCase   ,							dim=0		)

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

       self.scheduler.set_timesteps(_UpperCAmelCase   ,							device=_UpperCAmelCase		)
       _A        =       self.scheduler.timesteps

       _A        =       self.unet.config.in_channels

       _A       ,   _A        =       downscale_height_and_width(_UpperCAmelCase   ,							_UpperCAmelCase   ,							self.movq_scale_factor		)

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

             _A        =       {'image_embeds': image_embeds}
             _A        =       self.unet(
                 sample=_UpperCAmelCase   ,							timestep=_UpperCAmelCase   ,							encoder_hidden_states=_UpperCAmelCase   ,							added_cond_kwargs=_UpperCAmelCase   ,							return_dict=_UpperCAmelCase   ,							)[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(
                 _UpperCAmelCase   ,							_UpperCAmelCase   ,							_UpperCAmelCase   ,							generator=_UpperCAmelCase   ,							)[0]
       # post-processing
       _A        =       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"]:
             _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(_UpperCAmelCase		)

       if not return_dict:
             return (image,)

       return ImagePipelineOutput(images=_UpperCAmelCase		)

