import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference





class      lowerCAmelCase						(					_a			,		unittest.TestCase			):
					_SCREAMING_SNAKE_CASE     :     Union[str, Any]		=ShapEPipeline
					_SCREAMING_SNAKE_CASE     :     int		=["""prompt"""]
					_SCREAMING_SNAKE_CASE     :     List[str]		=["""prompt"""]
					_SCREAMING_SNAKE_CASE     :     Tuple		=[
					    """num_images_per_prompt""",
					    """num_inference_steps""",
					    """generator""",
					    """latents""",
					    """guidance_scale""",
					    """frame_size""",
					    """output_type""",
					    """return_dict""",
					]
					_SCREAMING_SNAKE_CASE     :     List[Any]		=False



					@property
					def 			a__  (					self       ):
												return 32



					@property
					def 			a__  (					self       ):
												return 32



					@property
					def 			a__  (					self       ):
												return self.time_input_dim * 4



					@property
					def 			a__  (					self       ):
												return 8



					@property
					def 			a__  (					self       ):
												_A=	CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'       )
												return tokenizer



					@property
					def 			a__  (					self       ):
												torch.manual_seed(0       )
												_A=	CLIPTextConfig(
												    bos_token_id=0				,      eos_token_id=2				,      hidden_size=self.text_embedder_hidden_size				,      projection_dim=self.text_embedder_hidden_size				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      pad_token_id=1				,      vocab_size=1000				,      )
												return CLIPTextModelWithProjection(lowerCAmelCase__       )



					@property
					def 			a__  (					self       ):
												torch.manual_seed(0       )

												_A=	{
												    'num_attention_heads': 2,
												    'attention_head_dim': 16,
												    'embedding_dim': self.time_input_dim,
												    'num_embeddings': 32,
												    'embedding_proj_dim': self.text_embedder_hidden_size,
												    'time_embed_dim': self.time_embed_dim,
												    'num_layers': 1,
												    'clip_embed_dim': self.time_input_dim * 2,
												    'additional_embeddings': 0,
												    'time_embed_act_fn': 'gelu',
												    'norm_in_type': 'layer',
												    'encoder_hid_proj_type': None,
												    'added_emb_type': None,
												}

												_A=	PriorTransformer(**lowerCAmelCase__       )
												return model



					@property
					def 			a__  (					self       ):
												torch.manual_seed(0       )

												_A=	{
												    'param_shapes': (
												        (self.renderer_dim, 93),
												        (self.renderer_dim, 8),
												        (self.renderer_dim, 8),
												        (self.renderer_dim, 8),
												    ),
												    'd_latent': self.time_input_dim,
												    'd_hidden': self.renderer_dim,
												    'n_output': 12,
												    'background': (
												        0.1,
												        0.1,
												        0.1,
												    ),
												}
												_A=	ShapERenderer(**lowerCAmelCase__       )
												return model



					def 			a__  (					self       ):
												_A=	self.dummy_prior
												_A=	self.dummy_text_encoder
												_A=	self.dummy_tokenizer
												_A=	self.dummy_renderer

												_A=	HeunDiscreteScheduler(
												    beta_schedule='exp'				,      num_train_timesteps=1024				,      prediction_type='sample'				,      use_karras_sigmas=lowerCAmelCase__				,      clip_sample=lowerCAmelCase__				,      clip_sample_range=1.0				,      )
												_A=	{
												    'prior': prior,
												    'text_encoder': text_encoder,
												    'tokenizer': tokenizer,
												    'renderer': renderer,
												    'scheduler': scheduler,
												}

												return components



					def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__=0       ):
												if str(lowerCAmelCase__       ).startswith('mps'       ):
																			_A=	torch.manual_seed(lowerCAmelCase__       )
												else:
																			_A=	torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
												_A=	{
												    'prompt': 'horse',
												    'generator': generator,
												    'num_inference_steps': 1,
												    'frame_size': 32,
												    'output_type': 'np',
												}
												return inputs



					def 			a__  (					self       ):
												_A=	'cpu'

												_A=	self.get_dummy_components()

												_A=	self.pipeline_class(**lowerCAmelCase__       )
												_A=	pipe.to(lowerCAmelCase__       )

												pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

												_A=	pipe(**self.get_dummy_inputs(lowerCAmelCase__       )       )
												_A=	output.images[0]
												_A=	image[0, -3:, -3:, -1]

												assert image.shape == (20, 32, 32, 3)

												_A=	np.array(
												    [
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												        0.00039216,
												    ]       )

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



					def 			a__  (					self       ):
												# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
												self._test_inference_batch_consistent(batch_sizes=[1, 2]       )



					def 			a__  (					self       ):
												_A=	torch_device == 'cpu'
												_A=	True

												self._test_inference_batch_single_identical(
												    batch_size=2				,      test_max_difference=lowerCAmelCase__				,      relax_max_difference=lowerCAmelCase__				,      )





					def 			a__  (					self       ):
												_A=	self.get_dummy_components()
												_A=	self.pipeline_class(**lowerCAmelCase__       )
												_A=	pipe.to(lowerCAmelCase__       )
												pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

												_A=	1
												_A=	2

												_A=	self.get_dummy_inputs(lowerCAmelCase__       )

												for key in inputs.keys():
																			if key in self.batch_params:
																										_A=	batch_size * [inputs[key]]

												_A=	pipe(**lowerCAmelCase__				,      num_images_per_prompt=lowerCAmelCase__       )[0]

												assert images.shape[0] == batch_size * num_images_per_prompt





@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'
												    '/shap_e/test_shap_e_np_out.npy'       )
												_A=	ShapEPipeline.from_pretrained('openai/shap-e'       )
												_A=	pipe.to(lowerCAmelCase__       )
												pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

												_A=	torch.Generator(device=lowerCAmelCase__       ).manual_seed(0       )

												_A=	pipe(
												    'a shark'				,      generator=lowerCAmelCase__				,      guidance_scale=15.0				,      num_inference_steps=64				,      frame_size=64				,      output_type='np'				,      ).images[0]

												assert images.shape == (20, 64, 64, 3)

												assert_mean_pixel_difference(lowerCAmelCase__				,      lowerCAmelCase__       )
import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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 PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





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



     def 			a__  (					self       ):
            torch.manual_seed(0       )
            _A=	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				,      )
            _A=	DDIMScheduler(
                beta_start=0.00085				,      beta_end=0.012				,      beta_schedule='scaled_linear'				,      clip_sample=lowerCAmelCase__				,      set_alpha_to_one=lowerCAmelCase__				,      )
            torch.manual_seed(0       )
            _A=	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				,      )

            # TODO: address the non-deterministic text encoder (fails for save-load tests)
            # torch.manual_seed(0)
            # text_encoder_config = RobertaSeriesConfig(
            #     hidden_size=32,
            #     project_dim=32,
            #     intermediate_size=37,
            #     layer_norm_eps=1e-05,
            #     num_attention_heads=4,
            #     num_hidden_layers=5,
            #     vocab_size=5002,
            # )
            # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

            torch.manual_seed(0       )
            _A=	CLIPTextConfig(
                bos_token_id=0				,      eos_token_id=2				,      hidden_size=32				,      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=5002				,      )
            _A=	CLIPTextModel(lowerCAmelCase__       )

            _A=	XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'       )
            _A=	77

            _A=	{
                'unet': unet,
                'scheduler': scheduler,
                'vae': vae,
                'text_encoder': text_encoder,
                'tokenizer': tokenizer,
                'safety_checker': None,
                'feature_extractor': None,
            }
            return components



     def 			a__  (					self				,      lowerCAmelCase__				,      lowerCAmelCase__=0       ):
            if str(lowerCAmelCase__       ).startswith('mps'       ):
                   _A=	torch.manual_seed(lowerCAmelCase__       )
            else:
                   _A=	torch.Generator(device=lowerCAmelCase__       ).manual_seed(lowerCAmelCase__       )
            _A=	{
                '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 			a__  (					self       ):
            super().test_attention_slicing_forward_pass(expected_max_diff=3E-3       )



     def 			a__  (					self       ):
            super().test_inference_batch_single_identical(expected_max_diff=3E-3       )



     def 			a__  (					self       ):
            _A=	'cpu'  # ensure determinism for the device-dependent torch.Generator

            _A=	self.get_dummy_components()
            torch.manual_seed(0       )
            _A=	RobertaSeriesConfig(
                hidden_size=32				,      project_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      vocab_size=5002				,      )
            # TODO: remove after fixing the non-deterministic text encoder
            _A=	RobertaSeriesModelWithTransformation(lowerCAmelCase__       )
            _A=	text_encoder

            _A=	AltDiffusionPipeline(**lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__       )
            _A=	'A photo of an astronaut'
            _A=	alt_pipe(**lowerCAmelCase__       )
            _A=	output.images
            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _A=	np.array(
                [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]       )

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





     def 			a__  (					self       ):
            _A=	'cpu'  # ensure determinism for the device-dependent torch.Generator

            _A=	self.get_dummy_components()
            _A=	PNDMScheduler(skip_prk_steps=lowerCAmelCase__       )
            torch.manual_seed(0       )
            _A=	RobertaSeriesConfig(
                hidden_size=32				,      project_dim=32				,      intermediate_size=37				,      layer_norm_eps=1E-05				,      num_attention_heads=4				,      num_hidden_layers=5				,      vocab_size=5002				,      )
            # TODO: remove after fixing the non-deterministic text encoder
            _A=	RobertaSeriesModelWithTransformation(lowerCAmelCase__       )
            _A=	text_encoder
            _A=	AltDiffusionPipeline(**lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	self.get_dummy_inputs(lowerCAmelCase__       )
            _A=	alt_pipe(**lowerCAmelCase__       )
            _A=	output.images
            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _A=	np.array(
                [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]       )

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





@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       ):
            # make sure here that pndm scheduler skips prk
            _A=	AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'				,      safety_checker=lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	'A painting of a squirrel eating a burger'
            _A=	torch.manual_seed(0       )
            _A=	alt_pipe([prompt]				,      generator=lowerCAmelCase__				,      guidance_scale=6.0				,      num_inference_steps=20				,      output_type='np'       )

            _A=	output.images

            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 512, 512, 3)
            _A=	np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]       )

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





     def 			a__  (					self       ):
            _A=	DDIMScheduler.from_pretrained('BAAI/AltDiffusion'				,      subfolder='scheduler'       )

            _A=	AltDiffusionPipeline.from_pretrained('BAAI/AltDiffusion'				,      scheduler=lowerCAmelCase__				,      safety_checker=lowerCAmelCase__       )
            _A=	alt_pipe.to(lowerCAmelCase__       )
            alt_pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

            _A=	'A painting of a squirrel eating a burger'
            _A=	torch.manual_seed(0       )

            _A=	alt_pipe([prompt]				,      generator=lowerCAmelCase__				,      num_inference_steps=2				,      output_type='numpy'       )
            _A=	output.images

            _A=	image[0, -3:, -3:, -1]

            assert image.shape == (1, 512, 512, 3)
            _A=	np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]       )

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