import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class   snake_case_       (		__A  ,unittest.TestCase       ):
   __A	:		Optional[int]         =				LDMTextToImagePipeline
   __A	:		int         =				TEXT_TO_IMAGE_PARAMS - {
       "negative_prompt",
       "negative_prompt_embeds",
       "cross_attention_kwargs",
       "prompt_embeds",
   }
   __A	:		Optional[Any]         =				PipelineTesterMixin.required_optional_params - {
       "num_images_per_prompt",
       "callback",
       "callback_steps",
   }
   __A	:		Any         =				TEXT_TO_IMAGE_BATCH_PARAMS
   __A	:		str         =				False






   def       __UpperCamelCase     (      self				:     List[str]  )					->						Optional[int]:
       torch.manual_seed(0  )
       lowercase__ :   List[Any]    					=				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	,						)
       lowercase__ :   List[str]    					=				DDIMScheduler(
           beta_start=0.0_00_85	,						beta_end=0.0_12	,						beta_schedule="scaled_linear"	,						clip_sample=lowercase_	,						set_alpha_to_one=lowercase_	,						)
       torch.manual_seed(0  )
       lowercase__ :   str    					=				AutoencoderKL(
           block_out_channels=(32, 64)	,						in_channels=3	,						out_channels=3	,						down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D")	,						up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D")	,						latent_channels=4	,						)
       torch.manual_seed(0  )
       lowercase__ :   Optional[Any]    					=				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=10_00	,						)
       lowercase__ :   Optional[int]    					=				CLIPTextModel(lowercase_  )
       lowercase__ :   Union[str, Any]    					=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"  )

       lowercase__ :   List[Any]    					=				{
           "unet": unet,
           "scheduler": scheduler,
           "vqvae": vae,
           "bert": text_encoder,
           "tokenizer": tokenizer,
       }
       return components






   def       __UpperCamelCase     (      self				:     List[str]	,						lowercase_				:     Dict	,						lowercase_				:     Optional[int]=0  )					->						Dict:
       if str(lowercase_  ).startswith("mps"  ):
           lowercase__ :   int    					=				torch.manual_seed(lowercase_  )
       else:
           lowercase__ :   List[str]    					=				torch.Generator(device=lowercase_  ).manual_seed(lowercase_  )
       lowercase__ :   List[str]    					=				{
           "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       __UpperCamelCase     (      self				:     Union[str, Any]  )					->						Dict:
       lowercase__ :   List[Any]    					=				"cpu"  # ensure determinism for the device-dependent torch.Generator

       lowercase__ :   str    					=				self.get_dummy_components()
       lowercase__ :   Any    					=				LDMTextToImagePipeline(**lowercase_  )
       pipe.to(lowercase_  )
       pipe.set_progress_bar_config(disable=lowercase_  )

       lowercase__ :   str    					=				self.get_dummy_inputs(lowercase_  )
       lowercase__ :   Union[str, Any]    					=				pipe(**lowercase_  ).images
       lowercase__ :   List[Any]    					=				image[0, -3:, -3:, -1]

       assert image.shape == (1, 16, 16, 3)
       lowercase__ :   int    					=				np.array([0.61_01, 0.61_56, 0.56_22, 0.48_95, 0.66_61, 0.38_04, 0.57_48, 0.61_36, 0.50_14]  )

       assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-3



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






   def       __UpperCamelCase     (      self				:     str  )					->						Dict:
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






   def       __UpperCamelCase     (      self				:     List[Any]	,						lowercase_				:     Union[str, Any]	,						lowercase_				:     Union[str, Any]=torch.floataa	,						lowercase_				:     List[Any]=0  )					->						int:
       lowercase__ :   Dict    					=				torch.manual_seed(lowercase_  )
       lowercase__ :   Union[str, Any]    					=				np.random.RandomState(lowercase_  ).standard_normal((1, 4, 32, 32)  )
       lowercase__ :   Optional[Any]    					=				torch.from_numpy(lowercase_  ).to(device=lowercase_	,						dtype=lowercase_  )
       lowercase__ :   Optional[int]    					=				{
           "prompt": "A painting of a squirrel eating a burger",
           "latents": latents,
           "generator": generator,
           "num_inference_steps": 3,
           "guidance_scale": 6.0,
           "output_type": "numpy",
       }
       return inputs







   def       __UpperCamelCase     (      self				:     Dict  )					->						List[str]:
       lowercase__ :   str    					=				LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256"  ).to(lowercase_  )
       pipe.set_progress_bar_config(disable=lowercase_  )

       lowercase__ :   Tuple    					=				self.get_inputs(lowercase_  )
       lowercase__ :   Optional[Any]    					=				pipe(**lowercase_  ).images
       lowercase__ :   Optional[int]    					=				image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 2_56, 2_56, 3)
       lowercase__ :   List[Any]    					=				np.array([0.5_18_25, 0.5_28_50, 0.5_25_43, 0.5_42_58, 0.5_23_04, 0.5_25_69, 0.5_43_63, 0.5_52_76, 0.5_68_78]  )
       lowercase__ :   Optional[int]    					=				np.abs(expected_slice - image_slice  ).max()
       assert max_diff < 1E-3



@nightly
@require_torch_gpu
class   snake_case_       (		unittest.TestCase       ):






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






   def       __UpperCamelCase     (      self				:     Tuple	,						lowercase_				:     Optional[int]	,						lowercase_				:     Optional[Any]=torch.floataa	,						lowercase_				:     List[str]=0  )					->						Dict:
       lowercase__ :   Optional[Any]    					=				torch.manual_seed(lowercase_  )
       lowercase__ :   str    					=				np.random.RandomState(lowercase_  ).standard_normal((1, 4, 32, 32)  )
       lowercase__ :   Dict    					=				torch.from_numpy(lowercase_  ).to(device=lowercase_	,						dtype=lowercase_  )
       lowercase__ :   Tuple    					=				{
           "prompt": "A painting of a squirrel eating a burger",
           "latents": latents,
           "generator": generator,
           "num_inference_steps": 50,
           "guidance_scale": 6.0,
           "output_type": "numpy",
       }
       return inputs







   def       __UpperCamelCase     (      self				:     Tuple  )					->						str:
       lowercase__ :   Any    					=				LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256"  ).to(lowercase_  )
       pipe.set_progress_bar_config(disable=lowercase_  )

       lowercase__ :   Optional[Any]    					=				self.get_inputs(lowercase_  )
       lowercase__ :   Optional[int]    					=				pipe(**lowercase_  ).images[0]

       lowercase__ :   Union[str, Any]    					=				load_numpy(
           "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy"  )
       lowercase__ :   Dict    					=				np.abs(expected_image - image  ).max()
       assert max_diff < 1E-3









from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch





@require_torch
@require_tf
class     lowerCamelCase__(					__lowerCamelCase):
      def 			lowerCAmelCase__						(				self: List[Any]   ):
             __lowerCamelCase   =						SMALL_MODEL_IDENTIFIER
             __lowerCamelCase   =						"""pt"""
             __lowerCamelCase   =						"""tf"""
      def 			lowerCAmelCase__						(				self: Optional[Any]    ,       UpperCamelCase_: Optional[Any]   ):
             __lowerCamelCase   =						AutoModel.from_pretrained(self.test_model   )
             model_pt.save_pretrained(UpperCamelCase_   )
      def 			lowerCAmelCase__						(				self: Tuple    ,       UpperCamelCase_: Union[str, Any]   ):
             __lowerCamelCase   =						TFAutoModel.from_pretrained(self.test_model    ,       from_pt=UpperCamelCase_   )
             model_tf.save_pretrained(UpperCamelCase_   )
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             __lowerCamelCase   =						"""mock_framework"""

             # Framework provided - return whatever the user provides
             __lowerCamelCase   =						FeaturesManager.determine_framework(self.test_model    ,       UpperCamelCase_   )
             self.assertEqual(UpperCamelCase_    ,       UpperCamelCase_   )

             # Local checkpoint and framework provided - return provided framework
             # PyTorch checkpoint
             with TemporaryDirectory() as local_pt_ckpt:
                    self._setup_pt_ckpt(UpperCamelCase_   )
                    __lowerCamelCase   =						FeaturesManager.determine_framework(UpperCamelCase_    ,       UpperCamelCase_   )
                    self.assertEqual(UpperCamelCase_    ,       UpperCamelCase_   )

             # TensorFlow checkpoint
             with TemporaryDirectory() as local_tf_ckpt:
                    self._setup_tf_ckpt(UpperCamelCase_   )
                    __lowerCamelCase   =						FeaturesManager.determine_framework(UpperCamelCase_    ,       UpperCamelCase_   )
                    self.assertEqual(UpperCamelCase_    ,       UpperCamelCase_   )
      def 			lowerCAmelCase__						(				self: int   ):
             # PyTorch checkpoint
             with TemporaryDirectory() as local_pt_ckpt:
                    self._setup_pt_ckpt(UpperCamelCase_   )
                    __lowerCamelCase   =						FeaturesManager.determine_framework(UpperCamelCase_   )
                    self.assertEqual(UpperCamelCase_    ,       self.framework_pt   )

             # TensorFlow checkpoint
             with TemporaryDirectory() as local_tf_ckpt:
                    self._setup_tf_ckpt(UpperCamelCase_   )
                    __lowerCamelCase   =						FeaturesManager.determine_framework(UpperCamelCase_   )
                    self.assertEqual(UpperCamelCase_    ,       self.framework_tf   )

             # Invalid local checkpoint
             with TemporaryDirectory() as local_invalid_ckpt:
                    with self.assertRaises(UpperCamelCase_   ):
                           __lowerCamelCase   =						FeaturesManager.determine_framework(UpperCamelCase_   )




      def 			lowerCAmelCase__						(				self: Optional[int]   ):
             __lowerCamelCase   =						MagicMock(return_value=UpperCamelCase_   )
             with patch("""transformers.onnx.features.is_tf_available"""    ,       UpperCamelCase_   ):
                    __lowerCamelCase   =						FeaturesManager.determine_framework(self.test_model   )
                    self.assertEqual(UpperCamelCase_    ,       self.framework_pt   )

             # PyTorch not in environment -> use TensorFlow
             __lowerCamelCase   =						MagicMock(return_value=UpperCamelCase_   )
             with patch("""transformers.onnx.features.is_torch_available"""    ,       UpperCamelCase_   ):
                    __lowerCamelCase   =						FeaturesManager.determine_framework(self.test_model   )
                    self.assertEqual(UpperCamelCase_    ,       self.framework_tf   )

             # Both in environment -> use PyTorch
             __lowerCamelCase   =						MagicMock(return_value=UpperCamelCase_   )
             __lowerCamelCase   =						MagicMock(return_value=UpperCamelCase_   )
             with patch("""transformers.onnx.features.is_tf_available"""    ,       UpperCamelCase_   ), patch(
                 """transformers.onnx.features.is_torch_available"""    ,       UpperCamelCase_   ):
                    __lowerCamelCase   =						FeaturesManager.determine_framework(self.test_model   )
                    self.assertEqual(UpperCamelCase_    ,       self.framework_pt   )

             # Both not in environment -> raise error
             __lowerCamelCase   =						MagicMock(return_value=UpperCamelCase_   )
             __lowerCamelCase   =						MagicMock(return_value=UpperCamelCase_   )
             with patch("""transformers.onnx.features.is_tf_available"""    ,       UpperCamelCase_   ), patch(
                 """transformers.onnx.features.is_torch_available"""    ,       UpperCamelCase_   ):
                    with self.assertRaises(UpperCamelCase_   ):
                           __lowerCamelCase   =						FeaturesManager.determine_framework(self.test_model   )

