


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    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		Union[str, Any]              =  StableDiffusionLDMaDPipeline
    UpperCAmelCase__						:		Tuple              =  TEXT_TO_IMAGE_PARAMS
    UpperCAmelCase__						:		Tuple              =  TEXT_TO_IMAGE_BATCH_PARAMS
    UpperCAmelCase__						:		Any              =  TEXT_TO_IMAGE_IMAGE_PARAMS






    def _a (      self   )  ->	str:
         torch.manual_seed(0   )
         __UpperCamelCase   							=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 ,       )
         __UpperCamelCase   							=DDIMScheduler(
             beta_start=0.0_0085 ,       beta_end=0.012 ,       beta_schedule='scaled_linear' ,       clip_sample=A_ ,       set_alpha_to_one=A_ ,       )
         torch.manual_seed(0   )
         __UpperCamelCase   							=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   )
         __UpperCamelCase   							=CLIPTextConfig(
             bos_token_id=0 ,       eos_token_id=2 ,       hidden_size=32 ,       intermediate_size=37 ,       layer_norm_eps=1E-05 ,       num_attention_heads=4 ,       num_hidden_layers=5 ,       pad_token_id=1 ,       vocab_size=1000 ,       )
         __UpperCamelCase   							=CLIPTextModel(A_   )
         __UpperCamelCase   							=CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

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






    def _a (      self ,       A_ ,       A_=0   )  ->	Optional[Any]:
         if str(A_   ).startswith('mps'   ):
              __UpperCamelCase   							=torch.manual_seed(A_   )
         else:
              __UpperCamelCase   							=torch.Generator(device=A_   ).manual_seed(A_   )
         __UpperCamelCase   							={
             '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 _a (      self   )  ->	str:
         __UpperCamelCase   							='cpu'  # ensure determinism for the device-dependent torch.Generator

         __UpperCamelCase   							=self.get_dummy_components()
         __UpperCamelCase   							=StableDiffusionLDMaDPipeline(**A_   )
         __UpperCamelCase   							=ldmad_pipe.to(A_   )
         ldmad_pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=ldmad_pipe(**A_   )
         __UpperCamelCase							,				__UpperCamelCase   							=output.rgb, output.depth

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

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

         __UpperCamelCase   							=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]   )
         __UpperCamelCase   							=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 _a (      self   )  ->	Optional[Any]:
         __UpperCamelCase   							=self.get_dummy_components()
         __UpperCamelCase   							=StableDiffusionLDMaDPipeline(**A_   )
         __UpperCamelCase   							=ldmad_pipe.to(A_   )
         ldmad_pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=3 * [inputs['prompt']]

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

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							=3 * [inputs.pop('prompt'   )]

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

         __UpperCamelCase   							=ldmad_pipe.text_encoder(A_   )[0]

         __UpperCamelCase   							=prompt_embeds

         # forward
         __UpperCamelCase   							=ldmad_pipe(**A_   )
         __UpperCamelCase							,				__UpperCamelCase   							=output.rgb, output.depth
         __UpperCamelCase   							=rgb_slice_a[0, -3:, -3:, -1]
         __UpperCamelCase   							=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 _a (      self   )  ->	Optional[Any]:
         __UpperCamelCase   							='cpu'  # ensure determinism for the device-dependent torch.Generator
         __UpperCamelCase   							=self.get_dummy_components()
         __UpperCamelCase   							=PNDMScheduler(skip_prk_steps=A_   )
         __UpperCamelCase   							=StableDiffusionLDMaDPipeline(**A_   )
         __UpperCamelCase   							=ldmad_pipe.to(A_   )
         ldmad_pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_dummy_inputs(A_   )
         __UpperCamelCase   							='french fries'
         __UpperCamelCase   							=ldmad_pipe(**A_ ,       negative_prompt=A_   )

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

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

         __UpperCamelCase   							=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]   )
         __UpperCamelCase   							=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    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    def _a (      self   )  ->	Optional[int]:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    def _a (      self ,       A_ ,       A_="cpu" ,       A_=torch.floataa ,       A_=0   )  ->	Optional[int]:
         __UpperCamelCase   							=torch.Generator(device=A_   ).manual_seed(A_   )
         __UpperCamelCase   							=np.random.RandomState(A_   ).standard_normal((1, 4, 64, 64)   )
         __UpperCamelCase   							=torch.from_numpy(A_   ).to(device=A_ ,       dtype=A_   )
         __UpperCamelCase   							={
             '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 _a (      self   )  ->	int:
         __UpperCamelCase   							=StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'   )
         __UpperCamelCase   							=ldmad_pipe.to(A_   )
         ldmad_pipe.set_progress_bar_config(disable=A_   )

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

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

         __UpperCamelCase   							=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]   )
         __UpperCamelCase   							=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    UpperCAmelCase__	( unittest.TestCase		):



    """simple docstring"""






    def _a (      self   )  ->	int:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    def _a (      self ,       A_ ,       A_="cpu" ,       A_=torch.floataa ,       A_=0   )  ->	int:
         __UpperCamelCase   							=torch.Generator(device=A_   ).manual_seed(A_   )
         __UpperCamelCase   							=np.random.RandomState(A_   ).standard_normal((1, 4, 64, 64)   )
         __UpperCamelCase   							=torch.from_numpy(A_   ).to(device=A_ ,       dtype=A_   )
         __UpperCamelCase   							={
             '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 _a (      self   )  ->	Union[str, Any]:
         __UpperCamelCase   							=StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'   ).to(A_   )
         ldmad_pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_inputs(A_   )
         __UpperCamelCase   							=ldmad_pipe(**A_   )
         __UpperCamelCase							,				__UpperCamelCase   							=output.rgb, output.depth

         __UpperCamelCase   							=0.49_5586
         __UpperCamelCase   							=0.3379_5515
         __UpperCamelCase   							=112.4_8518
         __UpperCamelCase   							=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 _a (      self   )  ->	List[Any]:
         __UpperCamelCase   							=StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'   ).to(A_   )
         ldmad_pipe.set_progress_bar_config(disable=A_   )

         __UpperCamelCase   							=self.get_inputs(A_   )
         __UpperCamelCase   							=ldmad_pipe(**A_   )
         __UpperCamelCase							,				__UpperCamelCase   							=output.rgb, output.depth

         __UpperCamelCase   							=0.419_4127
         __UpperCamelCase   							=0.3537_5586
         __UpperCamelCase   							=0.563_8502
         __UpperCamelCase   							=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









import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    UpperCAmelCase__						:		Optional[int]              =  "MCTCTFeatureExtractor"
    UpperCAmelCase__						:		str              =  "AutoTokenizer"






    def __init__(      self ,       A_ ,       A_   )  ->	Dict:
         super().__init__(A_ ,       A_   )
         __UpperCamelCase   							=self.feature_extractor
         __UpperCamelCase   							=False






    def __call__(      self ,       *A_ ,       **A_   )  ->	Optional[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor(*A_ ,       **A_   )

         if "raw_speech" in kwargs:
              warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'   )
              __UpperCamelCase   							=kwargs.pop('raw_speech'   )
         else:
              __UpperCamelCase   							=kwargs.pop('audio' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('sampling_rate' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('text' ,       A_   )
         if len(A_   ) > 0:
              __UpperCamelCase   							=args[0]
              __UpperCamelCase   							=args[1:]

         if audio is None and text is None:
              raise ValueError('You need to specify either an `audio` or `text` input to process.'   )

         if audio is not None:
              __UpperCamelCase   							=self.feature_extractor(A_ ,       *A_ ,       sampling_rate=A_ ,       **A_   )
         if text is not None:
              __UpperCamelCase   							=self.tokenizer(A_ ,       **A_   )

         if text is None:
              return inputs
         elif audio is None:
              return encodings
         else:
              __UpperCamelCase   							=encodings['input_ids']
              return inputs






    def _a (      self ,       *A_ ,       **A_   )  ->	str:
         return self.tokenizer.batch_decode(*A_ ,       **A_   )






    def _a (      self ,       *A_ ,       **A_   )  ->	List[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor.pad(*A_ ,       **A_   )

         __UpperCamelCase   							=kwargs.pop('input_features' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('labels' ,       A_   )
         if len(A_   ) > 0:
              __UpperCamelCase   							=args[0]
              __UpperCamelCase   							=args[1:]

         if input_features is not None:
              __UpperCamelCase   							=self.feature_extractor.pad(A_ ,       *A_ ,       **A_   )
         if labels is not None:
              __UpperCamelCase   							=self.tokenizer.pad(A_ ,       **A_   )

         if labels is None:
              return input_features
         elif input_features is None:
              return labels
         else:
              __UpperCamelCase   							=labels['input_ids']
              return input_features






    def _a (      self ,       *A_ ,       **A_   )  ->	Optional[int]:
         return self.tokenizer.decode(*A_ ,       **A_   )






    @contextmanager
    def _a (      self   )  ->	str:
         warnings.warn(
             '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
             'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
             'your audio inputs, or in a separate call.'   )
         __UpperCamelCase   							=True
         __UpperCamelCase   							=self.tokenizer
         yield
         __UpperCamelCase   							=self.feature_extractor
         __UpperCamelCase   							=False





