import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class      lowerCAmelCase						(					_a			,		_a			,		_a			,		unittest.TestCase			):
     _SCREAMING_SNAKE_CASE     :     Dict		=StableUnCLIPPipeline
     _SCREAMING_SNAKE_CASE     :     Any		=TEXT_TO_IMAGE_PARAMS
     _SCREAMING_SNAKE_CASE     :     List[str]		=TEXT_TO_IMAGE_BATCH_PARAMS
     _SCREAMING_SNAKE_CASE     :     Dict		=TEXT_TO_IMAGE_IMAGE_PARAMS
     _SCREAMING_SNAKE_CASE     :     str		=TEXT_TO_IMAGE_IMAGE_PARAMS

     # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
     _SCREAMING_SNAKE_CASE     :     Union[str, Any]		=False



     def 			a__  (					self       ):
            _A= 32
            _A= embedder_hidden_size

            # prior components

            torch.manual_seed(0       )
            _A= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

            torch.manual_seed(0       )
            _A= CLIPTextModelWithProjection(
                CLIPTextConfig(
                    bos_token_id=0				,      eos_token_id=2				,      hidden_size=A__				,      projection_dim=A__				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      pad_token_id=1				,      vocab_size=1000				,      )       )

            torch.manual_seed(0       )
            _A= PriorTransformer(
                num_attention_heads=2				,      attention_head_dim=12				,      embedding_dim=A__				,      num_layers=1				,      )

            torch.manual_seed(0       )
            _A= DDPMScheduler(
                variance_type='fixed_small_log'				,      prediction_type='sample'				,      num_train_timesteps=1000				,      clip_sample=A__				,      clip_sample_range=5.0				,      beta_schedule='squaredcos_cap_v2'				,      )

            # regular denoising components

            torch.manual_seed(0       )
            _A= StableUnCLIPImageNormalizer(embedding_dim=A__       )
            _A= DDPMScheduler(beta_schedule='squaredcos_cap_v2'       )

            torch.manual_seed(0       )
            _A= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )

            torch.manual_seed(0       )
            _A= CLIPTextModel(
                CLIPTextConfig(
                    bos_token_id=0				,      eos_token_id=2				,      hidden_size=A__				,      projection_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      pad_token_id=1				,      vocab_size=1000				,      )       )

            torch.manual_seed(0       )
            _A= UNetaDConditionModel(
                sample_size=32				,      in_channels=4				,      out_channels=4				,      down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')				,      up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')				,      block_out_channels=(32, 64)				,      attention_head_dim=(2, 4)				,      class_embed_type='projection'				,      projection_class_embeddings_input_dim=embedder_projection_dim * 2				,      cross_attention_dim=A__				,      layers_per_block=1				,      upcast_attention=A__				,      use_linear_projection=A__				,      )

            torch.manual_seed(0       )
            _A= DDIMScheduler(
                beta_schedule='scaled_linear'				,      beta_start=0.00085				,      beta_end=0.012				,      prediction_type='v_prediction'				,      set_alpha_to_one=A__				,      steps_offset=1				,      )

            torch.manual_seed(0       )
            _A= AutoencoderKL()

            _A= {
                # prior components
                """prior_tokenizer""": prior_tokenizer,
                """prior_text_encoder""": prior_text_encoder,
                """prior""": prior,
                """prior_scheduler""": prior_scheduler,
                # image noising components
                """image_normalizer""": image_normalizer,
                """image_noising_scheduler""": image_noising_scheduler,
                # regular denoising components
                """tokenizer""": tokenizer,
                """text_encoder""": text_encoder,
                """unet""": unet,
                """scheduler""": scheduler,
                """vae""": vae,
            }

            return components



     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__=0       ):
            if str(A__       ).startswith('mps'       ):
                   _A= torch.manual_seed(A__       )
            else:
                   _A= torch.Generator(device=A__       ).manual_seed(A__       )
            _A= {
                """prompt""": """A painting of a squirrel eating a burger""",
                """generator""": generator,
                """num_inference_steps""": 2,
                """prior_num_inference_steps""": 2,
                """output_type""": """numpy""",
            }
            return inputs



     def 			a__  (					self       ):
            _A= torch_device == """cpu"""

            self._test_attention_slicing_forward_pass(test_max_difference=A__       )





     def 			a__  (					self       ):
            _A= torch_device in ["""cpu""", """mps"""]

            self._test_inference_batch_single_identical(test_max_difference=A__       )





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



     def 			a__  (					self       ):
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()



     def 			a__  (					self       ):
            _A= load_numpy(
                'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'       )

            _A= StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'				,      torch_dtype=torch.floataa       )
            pipe.to(A__       )
            pipe.set_progress_bar_config(disable=A__       )
            # stable unclip will oom when integration tests are run on a V100,
            # so turn on memory savings
            pipe.enable_attention_slicing()
            pipe.enable_sequential_cpu_offload()

            _A= torch.Generator(device='cpu'       ).manual_seed(0       )
            _A= pipe('anime turle'				,      generator=A__				,      output_type='np'       )

            _A= output.images[0]

            assert image.shape == (768, 768, 3)

            assert_mean_pixel_difference(A__				,      A__       )





     def 			a__  (					self       ):
            torch.cuda.empty_cache()
            torch.cuda.reset_max_memory_allocated()
            torch.cuda.reset_peak_memory_stats()

            _A= StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'				,      torch_dtype=torch.floataa       )
            _A= pipe.to(A__       )
            pipe.set_progress_bar_config(disable=A__       )
            pipe.enable_attention_slicing()
            pipe.enable_sequential_cpu_offload()

            _A= pipe(
                'anime turtle'				,      prior_num_inference_steps=2				,      num_inference_steps=2				,      output_type='np'				,      )

            _A= torch.cuda.max_memory_allocated()
            # make sure that less than 7 GB is allocated
            assert mem_bytes < 7 * 10**9
import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
      import torch

      from transformers.generation import (
          MaxLengthCriteria,
          MaxNewTokensCriteria,
          MaxTimeCriteria,
          StoppingCriteriaList,
          validate_stopping_criteria,
      )





@require_torch
class      lowerCAmelCase						(					unittest.TestCase			):



     def 			a__  (					self				,      lowerCAmelCase__       ):
            _A=	3
            _A=	250

            _A=	ids_tensor((batch_size, length)				,      lowerCAmelCase__       )
            _A=	torch.ones((batch_size, length)				,      device=lowerCAmelCase__				,      dtype=torch.float       ) / length
            return input_ids, scores



     def 			a__  (					self       ):
            _A,			_A=	self._get_tensors(5       )

            _A=	StoppingCriteriaList(
                [
                    MaxLengthCriteria(max_length=10       ),
                    MaxTimeCriteria(max_time=0.1       ),
                ]       )

            self.assertFalse(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )

            _A,			_A=	self._get_tensors(9       )
            self.assertFalse(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )

            _A,			_A=	self._get_tensors(10       )
            self.assertTrue(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )



     def 			a__  (					self       ):
            _A=	MaxLengthCriteria(max_length=10       )

            _A,			_A=	self._get_tensors(5       )
            self.assertFalse(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )

            _A,			_A=	self._get_tensors(9       )
            self.assertFalse(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )

            _A,			_A=	self._get_tensors(10       )
            self.assertTrue(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )



     def 			a__  (					self       ):
            _A=	MaxNewTokensCriteria(start_length=5				,      max_new_tokens=5       )

            _A,			_A=	self._get_tensors(5       )
            self.assertFalse(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )

            _A,			_A=	self._get_tensors(9       )
            self.assertFalse(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )

            _A,			_A=	self._get_tensors(10       )
            self.assertTrue(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )

            _A=	StoppingCriteriaList([criteria]       )
            self.assertEqual(criteria_list.max_length				,      10       )



     def 			a__  (					self       ):
            _A,			_A=	self._get_tensors(5       )

            _A=	MaxTimeCriteria(max_time=0.1       )
            self.assertFalse(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )

            _A=	MaxTimeCriteria(max_time=0.1				,      initial_timestamp=time.time() - 0.2       )
            self.assertTrue(criteria(lowerCAmelCase__				,      lowerCAmelCase__       )       )





     def 			a__  (					self       ):
            validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10       )]       )				,      10       )

            with self.assertWarns(lowerCAmelCase__       ):
                   validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10       )]       )				,      11       )

            _A=	validate_stopping_criteria(StoppingCriteriaList()				,      11       )

            self.assertEqual(len(lowerCAmelCase__       )				,      1       )