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





	'''simple docstring'''



	_UpperCAmelCase	:		str          = StableDiffusionLDMaDPipeline
	_UpperCAmelCase	:		List[Any]          = TEXT_TO_IMAGE_PARAMS
	_UpperCAmelCase	:		Optional[Any]          = TEXT_TO_IMAGE_BATCH_PARAMS
	_UpperCAmelCase	:		Dict          = TEXT_TO_IMAGE_IMAGE_PARAMS



	def 			A							(  self     :   Tuple     ):


				'''simple docstring'''


				torch.manual_seed(0     )
				_snake_case          =      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      ,  )
				_snake_case          =      DDIMScheduler(
				    beta_start=0.00085      ,  beta_end=0.012      ,  beta_schedule='scaled_linear'      ,  clip_sample=lowercase      ,  set_alpha_to_one=lowercase      ,  )
				torch.manual_seed(0     )
				_snake_case          =      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     )
				_snake_case          =      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=1_000      ,  )
				_snake_case          =      CLIPTextModel(lowercase     )
				_snake_case          =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

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



	def 			A							(  self     :   Dict      ,  lowercase     :   int      ,  lowercase     :   Union[str, Any]=0     ):


				'''simple docstring'''


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


				'''simple docstring'''


				_snake_case          =      'cpu'  # ensure determinism for the device-dependent torch.Generator

				_snake_case          =      self.get_dummy_components()
				_snake_case          =      StableDiffusionLDMaDPipeline(**lowercase     )
				_snake_case          =      ldmad_pipe.to(lowercase     )
				ldmad_pipe.set_progress_bar_config(disable=lowercase     )

				_snake_case          =      self.get_dummy_inputs(lowercase     )
				_snake_case          =      ldmad_pipe(**lowercase     )
				_snake_case     , _snake_case          =      output.rgb, output.depth

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

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

				_snake_case          =      np.array(
				    [0.37338176, 0.70247, 0.74203193, 0.51643604, 0.58256793, 0.60932136, 0.4181095, 0.48355877, 0.46535262]     )
				_snake_case          =      np.array([103.46727, 85.812004, 87.849236]     )

				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     :   str     ):


				'''simple docstring'''


				_snake_case          =      self.get_dummy_components()
				_snake_case          =      StableDiffusionLDMaDPipeline(**lowercase     )
				_snake_case          =      ldmad_pipe.to(lowercase     )
				ldmad_pipe.set_progress_bar_config(disable=lowercase     )

				_snake_case          =      self.get_dummy_inputs(lowercase     )
				_snake_case          =      3 * [inputs['prompt']]

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

				_snake_case          =      self.get_dummy_inputs(lowercase     )
				_snake_case          =      3 * [inputs.pop('prompt'     )]

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

				_snake_case          =      ldmad_pipe.text_encoder(lowercase     )[0]

				_snake_case          =      prompt_embeds

				# forward
				_snake_case          =      ldmad_pipe(**lowercase     )
				_snake_case     , _snake_case          =      output.rgb, output.depth
				_snake_case          =      rgb_slice_a[0, -3:, -3:, -1]
				_snake_case          =      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]     ):


				'''simple docstring'''


				_snake_case          =      'cpu'  # ensure determinism for the device-dependent torch.Generator
				_snake_case          =      self.get_dummy_components()
				_snake_case          =      PNDMScheduler(skip_prk_steps=lowercase     )
				_snake_case          =      StableDiffusionLDMaDPipeline(**lowercase     )
				_snake_case          =      ldmad_pipe.to(lowercase     )
				ldmad_pipe.set_progress_bar_config(disable=lowercase     )

				_snake_case          =      self.get_dummy_inputs(lowercase     )
				_snake_case          =      'french fries'
				_snake_case          =      ldmad_pipe(**lowercase      ,  negative_prompt=lowercase     )

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

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

				_snake_case          =      np.array(
				    [0.37044, 0.71811503, 0.7223251, 0.48603675, 0.5638391, 0.6364948, 0.42833704, 0.4901315, 0.47926217]     )
				_snake_case          =      np.array([107.84738, 84.62802, 89.962135]     )
				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        SCREAMING_SNAKE_CASE__					(     unittest.TestCase ):





	'''simple docstring'''



	def 			A							(  self     :   Tuple     ):


				'''simple docstring'''


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



	def 			A							(  self     :   int      ,  lowercase     :   int      ,  lowercase     :   Tuple="cpu"      ,  lowercase     :   Optional[int]=torch.floataa      ,  lowercase     :   Union[str, Any]=0     ):


				'''simple docstring'''


				_snake_case          =      torch.Generator(device=lowercase     ).manual_seed(lowercase     )
				_snake_case          =      np.random.RandomState(lowercase     ).standard_normal((1, 4, 64, 64)     )
				_snake_case          =      torch.from_numpy(lowercase     ).to(device=lowercase      ,  dtype=lowercase     )
				_snake_case          =      {
				    '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     :   Any     ):


				'''simple docstring'''


				_snake_case          =      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'     )
				_snake_case          =      ldmad_pipe.to(lowercase     )
				ldmad_pipe.set_progress_bar_config(disable=lowercase     )

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

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

				_snake_case          =      np.array(
				    [0.53805465, 0.56707305, 0.5486515, 0.57012236, 0.5814511, 0.56253487, 0.54843014, 0.55092263, 0.6459706]     )
				_snake_case          =      np.array(
				    [0.9263781, 0.6678672, 0.5486515, 0.92202145, 0.67831135, 0.56253487, 0.9241694, 0.7551478, 0.6459706]     )
				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        SCREAMING_SNAKE_CASE__					(     unittest.TestCase ):





	'''simple docstring'''



	def 			A							(  self     :   Optional[Any]     ):


				'''simple docstring'''


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



	def 			A							(  self     :   Optional[int]      ,  lowercase     :   List[str]      ,  lowercase     :   Dict="cpu"      ,  lowercase     :   int=torch.floataa      ,  lowercase     :   int=0     ):


				'''simple docstring'''


				_snake_case          =      torch.Generator(device=lowercase     ).manual_seed(lowercase     )
				_snake_case          =      np.random.RandomState(lowercase     ).standard_normal((1, 4, 64, 64)     )
				_snake_case          =      torch.from_numpy(lowercase     ).to(device=lowercase      ,  dtype=lowercase     )
				_snake_case          =      {
				    '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     :   Optional[int]     ):


				'''simple docstring'''


				_snake_case          =      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'     ).to(lowercase     )
				ldmad_pipe.set_progress_bar_config(disable=lowercase     )

				_snake_case          =      self.get_inputs(lowercase     )
				_snake_case          =      ldmad_pipe(**lowercase     )
				_snake_case     , _snake_case          =      output.rgb, output.depth

				_snake_case          =      0.495586
				_snake_case          =      0.33795515
				_snake_case          =      112.48518
				_snake_case          =      98.489746
				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     :   Dict     ):


				'''simple docstring'''


				_snake_case          =      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'     ).to(lowercase     )
				ldmad_pipe.set_progress_bar_config(disable=lowercase     )

				_snake_case          =      self.get_inputs(lowercase     )
				_snake_case          =      ldmad_pipe(**lowercase     )
				_snake_case     , _snake_case          =      output.rgb, output.depth

				_snake_case          =      0.4194127
				_snake_case          =      0.35375586
				_snake_case          =      0.5638502
				_snake_case          =      0.34686103

				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

def 	a_		(	__lowercase       :     int = 50_000_000			)			->  int:
			_snake_case          =      set()
			_snake_case          =      int((limit - 24) ** (1 / 2)			)

			_snake_case          =      set(range(3	, prime_square_limit + 1	, 2			)			)
			primes.add(2			)
			for p in range(3	, prime_square_limit + 1	, 2			):
						if p not in primes:
									continue
						primes.difference_update(set(range(p * p	, prime_square_limit + 1	, __lowercase			)			)			)

			for primea in primes:
						_snake_case          =      primea * primea
						for primea in primes:
									_snake_case          =      primea * primea * primea
									if square + cube >= limit - 16:
												break
									for primea in primes:
												_snake_case          =      primea * primea * primea * primea
												_snake_case          =      square + cube + tetr
												if total >= limit:
															break
												ret.add(__lowercase			)

			return len(__lowercase			)


if __name__ == "__main__":
			print(F'{solution() = }')