



'''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_			(			_a						,						unittest.TestCase      ):
    """simple docstring"""







    lowercase   =						DiTPipeline
    lowercase   =						CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
    lowercase   =						PipelineTesterMixin.required_optional_params - {
        "latents",
        "num_images_per_prompt",
        "callback",
        "callback_steps",
    }
    lowercase   =						CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
    lowercase   =						False




    def lowerCamelCase			(       self   :	Dict		):
          torch.manual_seed(0		)
          snake_case__    :       str       		= 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=snake_case_					,	activation_fn="""gelu-approximate"""					,	num_embeds_ada_norm=1_000					,	norm_type="""ada_norm_zero"""					,	norm_elementwise_affine=snake_case_					,	)
          snake_case__    :       Optional[int]       		= AutoencoderKL()
          snake_case__    :       str       		= DDIMScheduler()
          snake_case__    :       Optional[Any]       		= {"""transformer""": transformer.eval(), """vae""": vae.eval(), """scheduler""": scheduler}
          return components




    def lowerCamelCase			(       self   :	int					,	snake_case_   :	int					,	snake_case_   :	Union[str, Any]=0		):
          if str(snake_case_		).startswith("""mps"""		):
                snake_case__    :       Optional[Any]       		= torch.manual_seed(snake_case_		)
          else:
                snake_case__    :       int       		= torch.Generator(device=snake_case_		).manual_seed(snake_case_		)
          snake_case__    :       str       		= {
              """class_labels""": [1],
              """generator""": generator,
              """num_inference_steps""": 2,
              """output_type""": """numpy""",
          }
          return inputs




    def lowerCamelCase			(       self   :	Optional[int]		):
          snake_case__    :       Union[str, Any]       		= """cpu"""

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

          snake_case__    :       Tuple       		= self.get_dummy_inputs(snake_case_		)
          snake_case__    :       Union[str, Any]       		= pipe(**snake_case_		).images
          snake_case__    :       Tuple       		= image[0, -3:, -3:, -1]

          self.assertEqual(image.shape					,	(1, 16, 16, 3)		)
          snake_case__    :       List[str]       		= np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]		)
          snake_case__    :       List[Any]       		= np.abs(image_slice.flatten() - expected_slice		).max()
          self.assertLessEqual(snake_case_					,	1E-3		)




    def lowerCamelCase			(       self   :	Any		):
          self._test_inference_batch_single_identical(relax_max_difference=snake_case_					,	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 lowerCamelCase			(       self   :	Dict		):
          self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3		)







@require_torch_gpu
@slow
class        UpperCAmelCase_			(			unittest.TestCase      ):
    """simple docstring"""







    def lowerCamelCase			(       self   :	Optional[Any]		):
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()




    def lowerCamelCase			(       self   :	Tuple		):
          snake_case__    :       List[Any]       		= torch.manual_seed(0		)

          snake_case__    :       Any       		= DiTPipeline.from_pretrained("""facebook/DiT-XL-2-256"""		)
          pipe.to("""cuda"""		)

          snake_case__    :       List[str]       		= ["""vase""", """umbrella""", """white shark""", """white wolf"""]
          snake_case__    :       List[Any]       		= pipe.get_label_ids(snake_case_		)

          snake_case__    :       Any       		= pipe(snake_case_					,	generator=snake_case_					,	num_inference_steps=40					,	output_type="""np"""		).images

          for word, image in zip(snake_case_					,	snake_case_		):
                snake_case__    :       Tuple       		= 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 lowerCamelCase			(       self   :	Union[str, Any]		):
          snake_case__    :       Any       		= DiTPipeline.from_pretrained("""facebook/DiT-XL-2-512"""		)
          snake_case__    :       List[str]       		= DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
          pipe.to("""cuda"""		)

          snake_case__    :       int       		= ["""vase""", """umbrella"""]
          snake_case__    :       str       		= pipe.get_label_ids(snake_case_		)

          snake_case__    :       int       		= torch.manual_seed(0		)
          snake_case__    :       List[str]       		= pipe(snake_case_					,	generator=snake_case_					,	num_inference_steps=25					,	output_type="""np"""		).images

          for word, image in zip(snake_case_					,	snake_case_		):
                snake_case__    :       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'''





def 		__snake_case(							)						->					Optional[Any]:
      for n in range(1	,  1_000_000						):
            yield n * (n + 1) // 2





def 		__snake_case(							_lowerCAmelCase						)						->					str:
      snake_case__    :       Optional[int]       		= 1
      snake_case__    :       List[Any]       		= 2
      while i * i <= n:
            snake_case__    :       Dict       		= 0
            while n % i == 0:
                  n //= i
                  multiplicity += 1
            divisors_count *= multiplicity + 1
            i += 1
      if n > 1:
            divisors_count *= 2
      return divisors_count





def 		__snake_case(							)						->					List[str]:
      return next(i for i in triangle_number_generator() if count_divisors(_lowerCAmelCase						) > 500						)


if __name__ == "__main__":
       print(solution())
