
"""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 		a_    (			_UpperCAmelCase  ,    unittest.TestCase						):
    a   :					int			     =					DiTPipeline
    a   :					Union[str, Any]			     =					CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
    a   :					List[Any]			     =					PipelineTesterMixin.required_optional_params - {
        'latents',
        'num_images_per_prompt',
        'callback',
        'callback_steps',
    }
    a   :					List[Any]			     =					CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
    a   :					Dict			     =					False







    def   _snake_case       (      self      :							Union[str, Any]  )		->Union[str, Any]:






        '''simple docstring'''






        torch.manual_seed(0  )
        _UpperCAmelCase        = 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=10_00  ,       norm_type="""ada_norm_zero"""  ,       norm_elementwise_affine=__UpperCamelCase  ,       )
        _UpperCAmelCase        = AutoencoderKL()
        _UpperCAmelCase        = DDIMScheduler()
        _UpperCAmelCase        = {"""transformer""": transformer.eval(), """vae""": vae.eval(), """scheduler""": scheduler}
        return components







    def   _snake_case       (      self      :							Any  ,       __UpperCamelCase      :							Optional[Any]  ,       __UpperCamelCase      :							Any=0  )		->List[str]:






        '''simple docstring'''






        if str(__UpperCamelCase  ).startswith("""mps"""  ):
            _UpperCAmelCase        = torch.manual_seed(__UpperCamelCase  )
        else:
            _UpperCAmelCase        = torch.Generator(device=__UpperCamelCase  ).manual_seed(__UpperCamelCase  )
        _UpperCAmelCase        = {
            """class_labels""": [1],
            """generator""": generator,
            """num_inference_steps""": 2,
            """output_type""": """numpy""",
        }
        return inputs







    def   _snake_case       (      self      :							str  )		->Optional[int]:






        '''simple docstring'''






        _UpperCAmelCase        = """cpu"""

        _UpperCAmelCase        = self.get_dummy_components()
        _UpperCAmelCase        = self.pipeline_class(**__UpperCamelCase  )
        pipe.to(__UpperCamelCase  )
        pipe.set_progress_bar_config(disable=__UpperCamelCase  )

        _UpperCAmelCase        = self.get_dummy_inputs(__UpperCamelCase  )
        _UpperCAmelCase        = pipe(**__UpperCamelCase  ).images
        _UpperCAmelCase        = image[0, -3:, -3:, -1]

        self.assertEqual(image.shape  ,       (1, 16, 16, 3)  )
        _UpperCAmelCase        = 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]  )
        _UpperCAmelCase        = np.abs(image_slice.flatten() - expected_slice  ).max()
        self.assertLessEqual(__UpperCamelCase  ,       1e-3  )







    def   _snake_case       (      self      :							int  )		->Union[str, Any]:






        '''simple docstring'''






        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   _snake_case       (      self      :							Tuple  )		->Union[str, Any]:






        '''simple docstring'''






        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3  )







@require_torch_gpu
@slow
class 		a_    (			unittest.TestCase						):







    def   _snake_case       (      self      :							List[Any]  )		->Union[str, Any]:






        '''simple docstring'''






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







    def   _snake_case       (      self      :							Optional[Any]  )		->str:






        '''simple docstring'''






        _UpperCAmelCase        = torch.manual_seed(0  )

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

        _UpperCAmelCase        = ["""vase""", """umbrella""", """white shark""", """white wolf"""]
        _UpperCAmelCase        = pipe.get_label_ids(__UpperCamelCase  )

        _UpperCAmelCase        = pipe(__UpperCamelCase  ,       generator=__UpperCamelCase  ,       num_inference_steps=40  ,       output_type="""np"""  ).images

        for word, image in zip(__UpperCamelCase  ,       __UpperCamelCase  ):
            _UpperCAmelCase        = 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   _snake_case       (      self      :							Tuple  )		->Dict:






        '''simple docstring'''






        _UpperCAmelCase        = DiTPipeline.from_pretrained("""facebook/DiT-XL-2-512"""  )
        _UpperCAmelCase        = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
        pipe.to("""cuda"""  )

        _UpperCAmelCase        = ["""vase""", """umbrella"""]
        _UpperCAmelCase        = pipe.get_label_ids(__UpperCamelCase  )

        _UpperCAmelCase        = torch.manual_seed(0  )
        _UpperCAmelCase        = pipe(__UpperCamelCase  ,       generator=__UpperCamelCase  ,       num_inference_steps=25  ,       output_type="""np"""  ).images

        for word, image in zip(__UpperCamelCase  ,       __UpperCamelCase  ):
            _UpperCAmelCase        = 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"""



import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
 import torch

 from transformers.generation import DisjunctiveConstraint




@require_torch
class 		a_    (			unittest.TestCase						):







    def   _snake_case       (      self      :							Union[str, Any]  )		->Any:






        '''simple docstring'''






        _UpperCAmelCase        = [[1, 2, 4], [1, 2, 3, 4]]
        _UpperCAmelCase        = DisjunctiveConstraint(__UpperCamelCase  )
        self.assertTrue(isinstance(dc.token_ids  ,       __UpperCamelCase  )  )

        with self.assertRaises(__UpperCamelCase  ):
            DisjunctiveConstraint(torch.LongTensor([[1, 2, 4], [1, 2, 3]]  )  )

        with self.assertRaises(__UpperCamelCase  ):
            DisjunctiveConstraint([torch.LongTensor([1, 2, 4]  ), torch.LongTensor([1, 2, 3, 4, 5]  )]  )







    def   _snake_case       (      self      :							Optional[int]  )		->Dict:






        '''simple docstring'''






        _UpperCAmelCase        = [[1, 2], [1, 2, 3, 4]]

        with self.assertRaises(__UpperCamelCase  ):
            DisjunctiveConstraint(__UpperCamelCase  )  # fails here







    def   _snake_case       (      self      :							Optional[int]  )		->Dict:






        '''simple docstring'''






        _UpperCAmelCase        = [[1, 2, 3], [1, 2, 4]]

        _UpperCAmelCase        = DisjunctiveConstraint(__UpperCamelCase  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(1  )
        _UpperCAmelCase        = stepped is True and completed is False and reset is False
        self.assertTrue(__UpperCamelCase  )
        self.assertTrue(not dc.completed  )
        self.assertTrue(dc.current_seq == [1]  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(2  )
        _UpperCAmelCase        = stepped is True and completed is False and reset is False
        self.assertTrue(__UpperCamelCase  )
        self.assertTrue(not dc.completed  )
        self.assertTrue(dc.current_seq == [1, 2]  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(3  )
        _UpperCAmelCase        = stepped is True and completed is True and reset is False
        self.assertTrue(__UpperCamelCase  )
        self.assertTrue(dc.completed  )  # Completed!
        self.assertTrue(dc.current_seq == [1, 2, 3]  )







    def   _snake_case       (      self      :							List[Any]  )		->Any:






        '''simple docstring'''






        _UpperCAmelCase        = [[1, 2, 3], [1, 2, 4, 5], [1, 2, 5]]

        _UpperCAmelCase        = DisjunctiveConstraint(__UpperCamelCase  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(1  )
        self.assertTrue(not dc.completed  )
        self.assertTrue(dc.current_seq == [1]  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(2  )
        self.assertTrue(not dc.completed  )
        self.assertTrue(dc.current_seq == [1, 2]  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(4  )
        self.assertTrue(not dc.completed  )
        self.assertTrue(dc.current_seq == [1, 2, 4]  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(5  )
        self.assertTrue(dc.completed  )  # Completed!
        self.assertTrue(dc.current_seq == [1, 2, 4, 5]  )

        dc.reset()

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(1  )
        self.assertTrue(not dc.completed  )
        self.assertTrue(dc.remaining() == 3  )
        self.assertTrue(dc.current_seq == [1]  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(2  )
        self.assertTrue(not dc.completed  )
        self.assertTrue(dc.remaining() == 2  )
        self.assertTrue(dc.current_seq == [1, 2]  )

        _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase        = dc.update(5  )
        self.assertTrue(dc.completed  )  # Completed!
        self.assertTrue(dc.remaining() == 0  )
        self.assertTrue(dc.current_seq == [1, 2, 5]  )