import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

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


enable_full_determinism()







class      SCREAMING_SNAKE_CASE__					(  __UpperCamelCase   ,							__UpperCamelCase   ,							unittest.TestCase    ):
       A							:	Union[str, Any]          =	StableDiffusionParadigmsPipeline
       A							:	int          =	TEXT_TO_IMAGE_PARAMS
       A							:	Any          =	TEXT_TO_IMAGE_BATCH_PARAMS
       A							:	Optional[Any]          =	TEXT_TO_IMAGE_IMAGE_PARAMS
       A							:	Union[str, Any]          =	TEXT_TO_IMAGE_IMAGE_PARAMS

       def 			snake_case__    (							self					:     Dict			):
           torch.manual_seed(0			)
           __snake_case       :				Optional[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			,			attention_head_dim=(2, 4)			,			use_linear_projection=_lowerCAmelCase			,			)
           __snake_case       :				Tuple 			=    DDIMParallelScheduler(
               beta_start=0.00085			,			beta_end=0.012			,			beta_schedule="""scaled_linear"""			,			clip_sample=_lowerCAmelCase			,			set_alpha_to_one=_lowerCAmelCase			,			)
           torch.manual_seed(0			)
           __snake_case       :				Tuple 			=    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			,			sample_size=1_28			,			)
           torch.manual_seed(0			)
           __snake_case       :				Tuple 			=    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			,			hidden_act="""gelu"""			,			projection_dim=5_12			,			)
           __snake_case       :				Dict 			=    CLIPTextModel(_lowerCAmelCase			)
           __snake_case       :				List[str] 			=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

           __snake_case       :				List[Any] 			=    {
               """unet""": unet,
               """scheduler""": scheduler,
               """vae""": vae,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """safety_checker""": None,
               """feature_extractor""": None,
           }
           return components

       def 			snake_case__    (							self					:     Dict			,			_lowerCAmelCase					:     Optional[int]			,			_lowerCAmelCase					:     Optional[Any]=0			):
           if str(_lowerCAmelCase			).startswith("""mps"""			):
               __snake_case       :				Any 			=    torch.manual_seed(_lowerCAmelCase			)
           else:
               __snake_case       :				Tuple 			=    torch.Generator(device=_lowerCAmelCase			).manual_seed(_lowerCAmelCase			)
           __snake_case       :				Any 			=    {
               """prompt""": """a photograph of an astronaut riding a horse""",
               """generator""": generator,
               """num_inference_steps""": 10,
               """guidance_scale""": 6.0,
               """output_type""": """numpy""",
               """parallel""": 3,
               """debug""": True,
           }
           return inputs

       def 			snake_case__    (							self					:     Tuple			):
           __snake_case       :				str 			=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
           __snake_case       :				List[str] 			=    self.get_dummy_components()
           __snake_case       :				Union[str, Any] 			=    StableDiffusionParadigmsPipeline(**_lowerCAmelCase			)
           __snake_case       :				Union[str, Any] 			=    sd_pipe.to(_lowerCAmelCase			)
           sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase			)

           __snake_case       :				List[str] 			=    self.get_dummy_inputs(_lowerCAmelCase			)
           __snake_case       :				Any 			=    sd_pipe(**_lowerCAmelCase			).images
           __snake_case       :				Optional[Any] 			=    image[0, -3:, -3:, -1]
           assert image.shape == (1, 64, 64, 3)

           __snake_case       :				Optional[Any] 			=    np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]			)

           assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2

       def 			snake_case__    (							self					:     Any			):
           __snake_case       :				List[Any] 			=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
           __snake_case       :				List[str] 			=    self.get_dummy_components()
           torch.manual_seed(0			)
           __snake_case       :				Any 			=    DDPMParallelScheduler()
           torch.manual_seed(0			)
           __snake_case       :				str 			=    StableDiffusionParadigmsPipeline(**_lowerCAmelCase			)
           __snake_case       :				Dict 			=    sd_pipe.to(_lowerCAmelCase			)
           sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase			)

           __snake_case       :				Any 			=    self.get_dummy_inputs(_lowerCAmelCase			)
           __snake_case       :				int 			=    sd_pipe(**_lowerCAmelCase			).images
           __snake_case       :				Optional[int] 			=    image[0, -3:, -3:, -1]
           assert image.shape == (1, 64, 64, 3)

           __snake_case       :				Dict 			=    np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]			)

           assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2

       def 			snake_case__    (							self					:     Any			):
           super().test_inference_batch_consistent(batch_sizes=[1, 2]			)

       def 			snake_case__    (							self					:     int			):
           super().test_inference_batch_single_identical(batch_size=2			,			expected_max_diff=3e-3			)





       def 			snake_case__    (							self					:     Tuple			):
           __snake_case       :				Tuple 			=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
           __snake_case       :				Optional[Any] 			=    self.get_dummy_components()
           __snake_case       :				List[str] 			=    StableDiffusionParadigmsPipeline(**_lowerCAmelCase			)
           __snake_case       :				List[str] 			=    sd_pipe.to(_lowerCAmelCase			)
           sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase			)

           __snake_case       :				Tuple 			=    self.get_dummy_inputs(_lowerCAmelCase			)
           __snake_case       :				Optional[int] 			=    """french fries"""
           __snake_case       :				Optional[int] 			=    sd_pipe(**_lowerCAmelCase			,			negative_prompt=_lowerCAmelCase			)
           __snake_case       :				List[str] 			=    output.images
           __snake_case       :				List[str] 			=    image[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)

           __snake_case       :				List[str] 			=    np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]			)

           assert np.abs(image_slice.flatten() - expected_slice			).max() < 1e-2







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

       def 			snake_case__    (							self					:     Union[str, Any]			):
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()

       def 			snake_case__    (							self					:     Optional[int]			,			_lowerCAmelCase					:     str=0			):
           __snake_case       :				List[str] 			=    torch.Generator(device=_lowerCAmelCase			).manual_seed(_lowerCAmelCase			)
           __snake_case       :				Any 			=    {
               """prompt""": """a photograph of an astronaut riding a horse""",
               """generator""": generator,
               """num_inference_steps""": 10,
               """guidance_scale""": 7.5,
               """output_type""": """numpy""",
               """parallel""": 3,
               """debug""": True,
           }
           return inputs





       def 			snake_case__    (							self					:     int			):
           __snake_case       :				str 			=    """stabilityai/stable-diffusion-2-base"""
           __snake_case       :				str 			=    DDIMParallelScheduler.from_pretrained(_lowerCAmelCase			,			subfolder="""scheduler"""			)
           __snake_case       :				str 			=    StableDiffusionParadigmsPipeline.from_pretrained(_lowerCAmelCase			,			scheduler=_lowerCAmelCase			,			safety_checker=_lowerCAmelCase			)
           pipe.to(_lowerCAmelCase			)
           pipe.set_progress_bar_config(disable=_lowerCAmelCase			)
           pipe.enable_attention_slicing()

           __snake_case       :				List[Any] 			=    self.get_inputs()
           __snake_case       :				str 			=    pipe(**_lowerCAmelCase			).images
           __snake_case       :				Dict 			=    image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 5_12, 5_12, 3)

           __snake_case       :				Union[str, Any] 			=    np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]			)

           assert np.abs(expected_slice - image_slice			).max() < 1e-2




import warnings

from .generation import FlaxGenerationMixin







class      SCREAMING_SNAKE_CASE__					(  __UpperCamelCase    ):
       # warning at import time
       warnings.warn(
           "Importing `FlaxGenerationMixin` from `src/transformers/generation_flax_utils.py` is deprecated and will "
           "be removed in Transformers v5. Import as `from transformers import FlaxGenerationMixin` instead."   ,							__UpperCamelCase   ,							)



