
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    TextToVideoSDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, load_numpy, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

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


enable_full_determinism()




@skip_mps
class     _SCREAMING_SNAKE_CASE     (      snake_case_			,       unittest.TestCase						):
	lowerCAmelCase__         = TextToVideoSDPipeline
	lowerCAmelCase__         = TEXT_TO_IMAGE_PARAMS
	lowerCAmelCase__         = TEXT_TO_IMAGE_BATCH_PARAMS
	# No `output_type`.
	lowerCAmelCase__         = frozenset(
	    [
	        'num_inference_steps',
	        'generator',
	        'latents',
	        'return_dict',
	        'callback',
	        'callback_steps',
	    ]						)



	def     SCREAMING_SNAKE_CASE_(    self )				->       Tuple:
					torch.manual_seed(0 )
					lowerCamelCase_						=			UNetaDConditionModel(
					    block_out_channels=(32, 64, 64, 64)						,  layers_per_block=2						,  sample_size=32						,  in_channels=4						,  out_channels=4						,  down_block_types=("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")						,  up_block_types=("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")						,  cross_attention_dim=32						,  attention_head_dim=4						,  )
					lowerCamelCase_						=			DDIMScheduler(
					    beta_start=0.0_0_0_8_5						,  beta_end=0.0_1_2						,  beta_schedule="scaled_linear"						,  clip_sample=lowercase						,  set_alpha_to_one=lowercase						,  )
					torch.manual_seed(0 )
					lowerCamelCase_						=			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=128						,  )
					torch.manual_seed(0 )
					lowerCamelCase_						=			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=1000						,  hidden_act="gelu"						,  projection_dim=512						,  )
					lowerCamelCase_						=			CLIPTextModel(lowercase )
					lowerCamelCase_						=			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

					lowerCamelCase_						=			{
					    "unet": unet,
					    "scheduler": scheduler,
					    "vae": vae,
					    "text_encoder": text_encoder,
					    "tokenizer": tokenizer,
					}
					return components



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase=0 )				->       Optional[int]:
					if str(lowercase ).startswith("mps" ):
									lowerCamelCase_						=			torch.manual_seed(lowercase )
					else:
									lowerCamelCase_						=			torch.Generator(device=lowercase ).manual_seed(lowercase )
					lowerCamelCase_						=			{
					    "prompt": "A painting of a squirrel eating a burger",
					    "generator": generator,
					    "num_inference_steps": 2,
					    "guidance_scale": 6.0,
					    "output_type": "pt",
					}
					return inputs



	def     SCREAMING_SNAKE_CASE_(    self )				->       int:
					lowerCamelCase_						=			"cpu"  # ensure determinism for the device-dependent torch.Generator
					lowerCamelCase_						=			self.get_dummy_components()
					lowerCamelCase_						=			TextToVideoSDPipeline(**lowercase )
					lowerCamelCase_						=			sd_pipe.to(lowercase )
					sd_pipe.set_progress_bar_config(disable=lowercase )

					lowerCamelCase_						=			self.get_dummy_inputs(lowercase )
					lowerCamelCase_						=			"np"
					lowerCamelCase_						=			sd_pipe(**lowercase ).frames
					lowerCamelCase_						=			frames[0][-3:, -3:, -1]

					assert frames[0].shape == (64, 64, 3)
					lowerCamelCase_						=			np.array([1_5_8.0, 1_6_0.0, 1_5_3.0, 1_2_5.0, 1_0_0.0, 1_2_1.0, 1_1_1.0, 9_3.0, 1_1_3.0] )

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



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowercase						,  expected_max_diff=3e-3 )



	@unittest.skipIf(
	    torch_device != "cuda" or not is_xformers_available()						,  reason="XFormers attention is only available with CUDA and `xformers` installed"						,  )
	def     SCREAMING_SNAKE_CASE_(    self )				->       Tuple:
					self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowercase						,  expected_max_diff=1e-2 )



	@unittest.skip(reason="Batching needs to be properly figured out first for this pipeline." )
	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[Any]:
					pass



	@unittest.skip(reason="Batching needs to be properly figured out first for this pipeline." )
	def     SCREAMING_SNAKE_CASE_(    self )				->       str:
					pass



	@unittest.skip(reason="`num_images_per_prompt` argument is not supported for this pipeline." )
	def     SCREAMING_SNAKE_CASE_(    self )				->       List[str]:
					pass



	def     SCREAMING_SNAKE_CASE_(    self )				->       Any:
					return super().test_progress_bar()




@slow
@skip_mps
class     _SCREAMING_SNAKE_CASE     (      unittest.TestCase						):



	def     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					lowerCamelCase_						=			load_numpy(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy" )

					lowerCamelCase_						=			TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b" )
					lowerCamelCase_						=			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config )
					lowerCamelCase_						=			pipe.to("cuda" )

					lowerCamelCase_						=			"Spiderman is surfing"
					lowerCamelCase_						=			torch.Generator(device="cpu" ).manual_seed(0 )

					lowerCamelCase_						=			pipe(lowercase						,  generator=lowercase						,  num_inference_steps=25						,  output_type="pt" ).frames
					lowerCamelCase_						=			video_frames.cpu().numpy()

					assert np.abs(expected_video - video ).mean() < 5e-2



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					lowerCamelCase_						=			load_numpy(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy" )

					lowerCamelCase_						=			TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b" )
					lowerCamelCase_						=			pipe.to("cuda" )

					lowerCamelCase_						=			"Spiderman is surfing"
					lowerCamelCase_						=			torch.Generator(device="cpu" ).manual_seed(0 )

					lowerCamelCase_						=			pipe(lowercase						,  generator=lowercase						,  num_inference_steps=2						,  output_type="pt" ).frames
					lowerCamelCase_						=			video_frames.cpu().numpy()

					assert np.abs(expected_video - video ).mean() < 5e-2




from abc import ABC, abstractmethod
from argparse import ArgumentParser




class     _SCREAMING_SNAKE_CASE     (      snake_case_						):



	@staticmethod
	@abstractmethod
	def     SCREAMING_SNAKE_CASE_(    lowercase )				->       int:
					raise NotImplementedError()



	@abstractmethod
	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					raise NotImplementedError()


