


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      =						StableDiffusionLDMaDPipeline
							_UpperCAmelCase      =						TEXT_TO_IMAGE_PARAMS
							_UpperCAmelCase      =						TEXT_TO_IMAGE_BATCH_PARAMS
							_UpperCAmelCase      =						TEXT_TO_IMAGE_IMAGE_PARAMS






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Any					=					UNetaDConditionModel(
										    block_out_channels=(3_2, 6_4)			,						layers_per_block=2			,						sample_size=3_2			,						in_channels=4			,						out_channels=4			,						down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')			,						up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')			,						cross_attention_dim=3_2			,						)
										SCREAMING_SNAKE_CASE_		:     int					=					DDIMScheduler(
										    beta_start=0.00_085			,						beta_end=0.012			,						beta_schedule='scaled_linear'			,						clip_sample=lowerCAmelCase__			,						set_alpha_to_one=lowerCAmelCase__			,						)
										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     int					=					AutoencoderKL(
										    block_out_channels=[3_2, 6_4]			,						in_channels=6			,						out_channels=6			,						down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']			,						up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']			,						latent_channels=4			,						)
										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					CLIPTextConfig(
										    bos_token_id=0			,						eos_token_id=2			,						hidden_size=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			,						)
										SCREAMING_SNAKE_CASE_		:     Any					=					CLIPTextModel(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					{
										    'unet': unet,
										    'scheduler': scheduler,
										    'vae': vae,
										    'text_encoder': text_encoder,
										    'tokenizer': tokenizer,
										    'safety_checker': None,
										    'feature_extractor': None,
										}
										return components






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__=0       ):





										"""simple docstring"""




										if str(lowerCAmelCase__       ).startswith('mps'       ):
													SCREAMING_SNAKE_CASE_		:     List[str]					=					torch.manual_seed(lowerCAmelCase__       )
										else:
													SCREAMING_SNAKE_CASE_		:     str					=					torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					{
										    '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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					'cpu'  # ensure determinism for the device-dependent torch.Generator

										SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_dummy_components()
										SCREAMING_SNAKE_CASE_		:     Tuple					=					StableDiffusionLDMaDPipeline(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					ldmad_pipe.to(lowerCAmelCase__       )
										ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Any					=					self.get_dummy_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					ldmad_pipe(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					output.rgb, output.depth

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					rgb[0, -3:, -3:, -1]
										SCREAMING_SNAKE_CASE_		:     Any					=					depth[0, -3:, -1]

										assert rgb.shape == (1, 6_4, 6_4, 3)
										assert depth.shape == (1, 6_4, 6_4)

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					np.array(
										    [0.37_338_176, 0.70_247, 0.74_203_193, 0.51_643_604, 0.58_256_793, 0.60_932_136, 0.4_181_095, 0.48_355_877, 0.46_535_262]       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					np.array([103.46_727, 85.812_004, 87.849_236]       )

										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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_dummy_components()
										SCREAMING_SNAKE_CASE_		:     Any					=					StableDiffusionLDMaDPipeline(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Any					=					ldmad_pipe.to(lowerCAmelCase__       )
										ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.get_dummy_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     int					=					3 * [inputs['prompt']]

										# forward
										SCREAMING_SNAKE_CASE_		:     str					=					ldmad_pipe(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     List[str]					=					output.rgb, output.depth
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					rgb_slice_a[0, -3:, -3:, -1]
										SCREAMING_SNAKE_CASE_		:     List[str]					=					depth_slice_a[0, -3:, -1]

										SCREAMING_SNAKE_CASE_		:     int					=					self.get_dummy_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					3 * [inputs.pop('prompt'       )]

										SCREAMING_SNAKE_CASE_		:     List[str]					=					ldmad_pipe.tokenizer(
										    lowerCAmelCase__			,						padding='max_length'			,						max_length=ldmad_pipe.tokenizer.model_max_length			,						truncation=lowerCAmelCase__			,						return_tensors='pt'			,						)
										SCREAMING_SNAKE_CASE_		:     str					=					text_inputs['input_ids'].to(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[Any]					=					ldmad_pipe.text_encoder(lowerCAmelCase__       )[0]

										SCREAMING_SNAKE_CASE_		:     List[str]					=					prompt_embeds

										# forward
										SCREAMING_SNAKE_CASE_		:     Tuple					=					ldmad_pipe(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					output.rgb, output.depth
										SCREAMING_SNAKE_CASE_		:     str					=					rgb_slice_a[0, -3:, -3:, -1]
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					'cpu'  # ensure determinism for the device-dependent torch.Generator
										SCREAMING_SNAKE_CASE_		:     List[str]					=					self.get_dummy_components()
										SCREAMING_SNAKE_CASE_		:     Tuple					=					PNDMScheduler(skip_prk_steps=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					StableDiffusionLDMaDPipeline(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					ldmad_pipe.to(lowerCAmelCase__       )
										ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_dummy_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Any					=					'french fries'
										SCREAMING_SNAKE_CASE_		:     str					=					ldmad_pipe(**lowerCAmelCase__			,						negative_prompt=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     List[str]					=					output.rgb, output.depth
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					rgb[0, -3:, -3:, -1]
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					depth[0, -3:, -1]

										assert rgb.shape == (1, 6_4, 6_4, 3)
										assert depth.shape == (1, 6_4, 6_4)

										SCREAMING_SNAKE_CASE_		:     Tuple					=					np.array(
										    [0.37_044, 0.71_811_503, 0.7_223_251, 0.48_603_675, 0.5_638_391, 0.6_364_948, 0.42_833_704, 0.4_901_315, 0.47_926_217]       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					np.array([107.84_738, 84.62_802, 89.962_135]       )
										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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__="cpu"			,						lowerCAmelCase__=torch.floataa			,						lowerCAmelCase__=0       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Dict					=					np.random.RandomState(lowerCAmelCase__       ).standard_normal((1, 4, 6_4, 6_4)       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					torch.from_numpy(lowerCAmelCase__       ).to(device=lowerCAmelCase__			,						dtype=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					{
										    '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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					ldmad_pipe.to(lowerCAmelCase__       )
										ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.get_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					ldmad_pipe(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Dict					=					output.rgb, output.depth
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					rgb[0, -3:, -3:, -1].flatten()
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					rgb[0, -3:, -1].flatten()

										assert rgb.shape == (1, 5_1_2, 5_1_2, 3)
										assert depth.shape == (1, 5_1_2, 5_1_2)

										SCREAMING_SNAKE_CASE_		:     int					=					np.array(
										    [0.53_805_465, 0.56_707_305, 0.5_486_515, 0.57_012_236, 0.5_814_511, 0.56_253_487, 0.54_843_014, 0.55_092_263, 0.6_459_706]       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					np.array(
										    [0.9_263_781, 0.6_678_672, 0.5_486_515, 0.92_202_145, 0.67_831_135, 0.56_253_487, 0.9_241_694, 0.7_551_478, 0.6_459_706]       )
										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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__="cpu"			,						lowerCAmelCase__=torch.floataa			,						lowerCAmelCase__=0       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					np.random.RandomState(lowerCAmelCase__       ).standard_normal((1, 4, 6_4, 6_4)       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					torch.from_numpy(lowerCAmelCase__       ).to(device=lowerCAmelCase__			,						dtype=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					{
										    'prompt': 'a photograph of an astronaut riding a horse',
										    'latents': latents,
										    'generator': generator,
										    'num_inference_steps': 5_0,
										    'guidance_scale': 7.5,
										    'output_type': 'numpy',
										}
										return inputs






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'       ).to(lowerCAmelCase__       )
										ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Any					=					self.get_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					ldmad_pipe(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Tuple					=					output.rgb, output.depth

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					0.495_586
										SCREAMING_SNAKE_CASE_		:     str					=					0.33_795_515
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					112.48_518
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					98.489_746
										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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[str]					=					StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'       ).to(lowerCAmelCase__       )
										ldmad_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Any					=					ldmad_pipe(**lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     str					=					output.rgb, output.depth

										SCREAMING_SNAKE_CASE_		:     Dict					=					0.4_194_127
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					0.35_375_586
										SCREAMING_SNAKE_CASE_		:     int					=					0.5_638_502
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					0.34_686_103

										assert rgb.shape == (1, 5_1_2, 5_1_2, 3)
										assert depth.shape == (1, 5_1_2, 5_1_2, 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 __future__ import annotations




def 							lowerCamelCase	(					_snake_case	:    int  ,_snake_case	:    int							):





    '''simple docstring'''


    if b == 0:
        return (1, 0)
    ((lowercase__)							,							(lowercase__))									=    extended_euclid(_snake_case  ,a % b							)
    lowercase__									=    a // b
    return (y, x - k * y)




def 							lowerCamelCase	(					_snake_case	:    int  ,_snake_case	:    int  ,_snake_case	:    int  ,_snake_case	:    int							):





    '''simple docstring'''


    ((lowercase__)							,							(lowercase__))									=    extended_euclid(_snake_case  ,_snake_case							)
    lowercase__									=    na * na
    lowercase__									=    ra * x * na + ra * y * na
    return (n % m + m) % m




def 							lowerCamelCase	(					_snake_case	:    int  ,_snake_case	:    int							):





    '''simple docstring'''


    ((lowercase__)							,							(lowercase__))									=    extended_euclid(_snake_case  ,_snake_case							)
    if b < 0:
        lowercase__									=    (b % n + n) % n
    return b




def 							lowerCamelCase	(					_snake_case	:    int  ,_snake_case	:    int  ,_snake_case	:    int  ,_snake_case	:    int							):





    '''simple docstring'''


    lowercase__							,							lowercase__									=    invert_modulo(_snake_case  ,_snake_case							), invert_modulo(_snake_case  ,_snake_case							)
    lowercase__									=    na * na
    lowercase__									=    ra * x * na + ra * y * na
    return (n % m + m) % m


if __name__ == "__main__":
 from doctest import testmod

 testmod(name="chinese_remainder_theorem", verbose=True)
 testmod(name="chinese_remainder_theorem2", verbose=True)
 testmod(name="invert_modulo", verbose=True)
 testmod(name="extended_euclid", verbose=True)




