






"""simple docstring"""


import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 						__UpperCamelCase     (     UpperCamelCase   ,     unittest.TestCase						):
					lowercase_    :				List[Any]      =  DiTPipeline
					lowercase_    :				Any      =  CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
					lowercase_    :				Union[str, Any]      =  PipelineTesterMixin.required_optional_params - {
					    """latents""",
					    """num_images_per_prompt""",
					    """callback""",
					    """callback_steps""",
					}
					lowercase_    :				Tuple      =  CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
					lowercase_    :				int      =  False
					def        UpperCAmelCase__ (							self							:	str   )	->       Union[str, Any]:
											torch.manual_seed(0   )
											lowerCAmelCase   :Dict						=			TransformeraDModel(
											    sample_size=16     , num_layers=2     , patch_size=4     , attention_head_dim=8     , num_attention_heads=2     , in_channels=4     , out_channels=8     , attention_bias=UpperCAmelCase     , activation_fn='gelu-approximate'     , num_embeds_ada_norm=1000     , norm_type='ada_norm_zero'     , norm_elementwise_affine=UpperCAmelCase     , )
											lowerCAmelCase   :Dict						=			AutoencoderKL()
											lowerCAmelCase   :Union[str, Any]						=			DDIMScheduler()
											lowerCAmelCase   :List[Any]						=			{'transformer': transformer.eval(), 'vae': vae.eval(), 'scheduler': scheduler}
											return components
					def        UpperCAmelCase__ (							self							:	List[Any]     , UpperCAmelCase							:	List[str]     , UpperCAmelCase							:	Dict=0   )	->       Optional[Any]:
											if str(UpperCAmelCase   ).startswith('mps'   ):
																	lowerCAmelCase   :List[Any]						=			torch.manual_seed(UpperCAmelCase   )
											else:
																	lowerCAmelCase   :Union[str, Any]						=			torch.Generator(device=UpperCAmelCase   ).manual_seed(UpperCAmelCase   )
											lowerCAmelCase   :List[str]						=			{
											    'class_labels': [1],
											    'generator': generator,
											    'num_inference_steps': 2,
											    'output_type': 'numpy',
											}
											return inputs
					def        UpperCAmelCase__ (							self							:	Any   )	->       List[Any]:
											lowerCAmelCase   :str						=			'cpu'

											lowerCAmelCase   :Optional[int]						=			self.get_dummy_components()
											lowerCAmelCase   :int						=			self.pipeline_class(**UpperCAmelCase   )
											pipe.to(UpperCAmelCase   )
											pipe.set_progress_bar_config(disable=UpperCAmelCase   )

											lowerCAmelCase   :Dict						=			self.get_dummy_inputs(UpperCAmelCase   )
											lowerCAmelCase   :int						=			pipe(**UpperCAmelCase   ).images
											lowerCAmelCase   :str						=			image[0, -3:, -3:, -1]

											self.assertEqual(image.shape     , (1, 16, 16, 3)   )
											lowerCAmelCase   :Tuple						=			np.array([0.2_9_4_6, 0.6_6_0_1, 0.4_3_2_9, 0.3_2_9_6, 0.4_1_4_4, 0.5_3_1_9, 0.7_2_7_3, 0.5_0_1_3, 0.4_4_5_7]   )
											lowerCAmelCase   :List[Any]						=			np.abs(image_slice.flatten() - expected_slice   ).max()
											self.assertLessEqual(UpperCAmelCase     , 1e-3   )
					def        UpperCAmelCase__ (							self							:	Optional[int]   )	->       int:
											self._test_inference_batch_single_identical(relax_max_difference=UpperCAmelCase     , expected_max_diff=1e-3   )
					@unittest.skipIf(
					    torch_device != 'cuda' or not is_xformers_available()     , reason='XFormers attention is only available with CUDA and `xformers` installed'     , )
					def        UpperCAmelCase__ (							self							:	str   )	->       List[str]:
											self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3   )




@require_torch_gpu
@slow
class 						__UpperCamelCase     (     unittest.TestCase						):
					def        UpperCAmelCase__ (							self							:	Optional[Any]   )	->       Optional[Any]:
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()
					def        UpperCAmelCase__ (							self							:	Dict   )	->       Union[str, Any]:
											lowerCAmelCase   :Union[str, Any]						=			torch.manual_seed(0   )

											lowerCAmelCase   :List[Any]						=			DiTPipeline.from_pretrained('facebook/DiT-XL-2-256'   )
											pipe.to('cuda'   )

											lowerCAmelCase   :int						=			['vase', 'umbrella', 'white shark', 'white wolf']
											lowerCAmelCase   :List[str]						=			pipe.get_label_ids(UpperCAmelCase   )

											lowerCAmelCase   :Optional[Any]						=			pipe(UpperCAmelCase     , generator=UpperCAmelCase     , num_inference_steps=40     , output_type='np'   ).images

											for word, image in zip(UpperCAmelCase     , UpperCAmelCase   ):
																	lowerCAmelCase   :Any						=			load_numpy(
																	    f"""https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy"""   )
																	assert np.abs((expected_image - image).max()   ) < 1e-2
					def        UpperCAmelCase__ (							self							:	List[Any]   )	->       int:
											lowerCAmelCase   :Tuple						=			DiTPipeline.from_pretrained('facebook/DiT-XL-2-512'   )
											lowerCAmelCase   :str						=			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config   )
											pipe.to('cuda'   )

											lowerCAmelCase   :Union[str, Any]						=			['vase', 'umbrella']
											lowerCAmelCase   :Optional[int]						=			pipe.get_label_ids(UpperCAmelCase   )

											lowerCAmelCase   :Optional[Any]						=			torch.manual_seed(0   )
											lowerCAmelCase   :Tuple						=			pipe(UpperCAmelCase     , generator=UpperCAmelCase     , num_inference_steps=25     , output_type='np'   ).images

											for word, image in zip(UpperCAmelCase     , UpperCAmelCase   ):
																	lowerCAmelCase   :Optional[Any]						=			load_numpy(
																	    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
																	    f"""/dit/{word}_512.npy"""   )

																	assert np.abs((expected_image - image).max()   ) < 1e-1







"""simple docstring"""


from math import sqrt


def        UpperCAmelCase      (					a__ = 1_00_00_00 ):





						'''simple docstring'''






						lowerCAmelCase   :int						=			0
						lowerCAmelCase   :int						=			0
						lowerCAmelCase   :int

						while num_cuboids <= limit:
												max_cuboid_size += 1
												for sum_shortest_sides in range(2     ,      2 * max_cuboid_size + 1 ):
																		if sqrt(sum_shortest_sides**2 + max_cuboid_size**2 ).is_integer():
																								num_cuboids += (
																								    min(a__     ,      sum_shortest_sides // 2 )
																								    - max(1     ,      sum_shortest_sides - max_cuboid_size )
																								    + 1
																								)

						return max_cuboid_size


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