
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  lowerCamelCase__							(				snake_case_     ,						unittest.TestCase):




	'''simple docstring'''




	_A								=     TextToVideoSDPipeline
	_A								=     TEXT_TO_IMAGE_PARAMS
	_A								=     TEXT_TO_IMAGE_BATCH_PARAMS
	# No `output_type`.
	_A								=     frozenset(
	    [
	        'num_inference_steps',
	        'generator',
	        'latents',
	        'return_dict',
	        'callback',
	        'callback_steps',
	    ])







	def    _lowerCamelCase				(					self		:int						)		-> Any:
								torch.manual_seed(0						)
								__UpperCamelCase	:   int			   =					UNetaDConditionModel(
								    block_out_channels=(3_2, 6_4, 6_4, 6_4)     ,   layers_per_block=2     ,   sample_size=3_2     ,   in_channels=4     ,   out_channels=4     ,   down_block_types=("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")     ,   up_block_types=("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")     ,   cross_attention_dim=3_2     ,   attention_head_dim=4     ,   )
								__UpperCamelCase	:   Optional[Any]			   =					DDIMScheduler(
								    beta_start=0.00085     ,   beta_end=0.012     ,   beta_schedule="scaled_linear"     ,   clip_sample=a     ,   set_alpha_to_one=a     ,   )
								torch.manual_seed(0						)
								__UpperCamelCase	:   str			   =					AutoencoderKL(
								    block_out_channels=[3_2, 6_4]     ,   in_channels=3     ,   out_channels=3     ,   down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]     ,   up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]     ,   latent_channels=4     ,   sample_size=1_2_8     ,   )
								torch.manual_seed(0						)
								__UpperCamelCase	:   Union[str, Any]			   =					CLIPTextConfig(
								    bos_token_id=0     ,   eos_token_id=2     ,   hidden_size=3_2     ,   intermediate_size=3_7     ,   layer_norm_eps=1E-0_5     ,   num_attention_heads=4     ,   num_hidden_layers=5     ,   pad_token_id=1     ,   vocab_size=1_0_0_0     ,   hidden_act="gelu"     ,   projection_dim=5_1_2     ,   )
								__UpperCamelCase	:   str			   =					CLIPTextModel(a						)
								__UpperCamelCase	:   List[Any]			   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

								__UpperCamelCase	:   int			   =					{
								    'unet': unet,
								    'scheduler': scheduler,
								    'vae': vae,
								    'text_encoder': text_encoder,
								    'tokenizer': tokenizer,
								}
								return components







	def    _lowerCamelCase				(					self		:Dict     ,   a		:int     ,   a		:int=0						)		-> int:





								if str(a						).startswith("mps"						):
															__UpperCamelCase	:   int			   =					torch.manual_seed(a						)
								else:
															__UpperCamelCase	:   List[Any]			   =					torch.Generator(device=a						).manual_seed(a						)
								__UpperCamelCase	:   List[str]			   =					{
								    '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    _lowerCamelCase				(					self		:Tuple						)		-> str:
								__UpperCamelCase	:   Dict			   =					'cpu'  # ensure determinism for the device-dependent torch.Generator
								__UpperCamelCase	:   Tuple			   =					self.get_dummy_components()
								__UpperCamelCase	:   Union[str, Any]			   =					TextToVideoSDPipeline(**a						)
								__UpperCamelCase	:   List[Any]			   =					sd_pipe.to(a						)
								sd_pipe.set_progress_bar_config(disable=a						)

								__UpperCamelCase	:   Union[str, Any]			   =					self.get_dummy_inputs(a						)
								__UpperCamelCase	:   Union[str, Any]			   =					'np'
								__UpperCamelCase	:   List[Any]			   =					sd_pipe(**a						).frames
								__UpperCamelCase	:   int			   =					frames[0][-3:, -3:, -1]

								assert frames[0].shape == (6_4, 6_4, 3)
								__UpperCamelCase	:   List[str]			   =					np.array([158.0, 160.0, 153.0, 125.0, 100.0, 121.0, 111.0, 93.0, 113.0]						)

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







	def    _lowerCamelCase				(					self		:Optional[Any]						)		-> Optional[int]:
								self._test_attention_slicing_forward_pass(test_mean_pixel_difference=a     ,   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    _lowerCamelCase				(					self		:List[str]						)		-> Optional[Any]:
								self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=a     ,   expected_max_diff=1E-2						)







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







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







	@unittest.skip(reason="`num_images_per_prompt` argument is not supported for this pipeline."						)
	def    _lowerCamelCase				(					self		:int						)		-> Tuple:
								pass







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




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




	'''simple docstring'''




	def    _lowerCamelCase				(					self		:List[str]						)		-> Optional[int]:
								__UpperCamelCase	:   str			   =					load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy"						)

								__UpperCamelCase	:   str			   =					TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b"						)
								__UpperCamelCase	:   List[Any]			   =					DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
								__UpperCamelCase	:   Union[str, Any]			   =					pipe.to("cuda"						)

								__UpperCamelCase	:   List[str]			   =					'Spiderman is surfing'
								__UpperCamelCase	:   Tuple			   =					torch.Generator(device="cpu"						).manual_seed(0						)

								__UpperCamelCase	:   str			   =					pipe(a     ,   generator=a     ,   num_inference_steps=2_5     ,   output_type="pt"						).frames
								__UpperCamelCase	:   Optional[int]			   =					video_frames.cpu().numpy()

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







	def    _lowerCamelCase				(					self		:Optional[int]						)		-> Optional[int]:
								__UpperCamelCase	:   Any			   =					load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy"						)

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

								__UpperCamelCase	:   Optional[Any]			   =					'Spiderman is surfing'
								__UpperCamelCase	:   str			   =					torch.Generator(device="cpu"						).manual_seed(0						)

								__UpperCamelCase	:   List[Any]			   =					pipe(a     ,   generator=a     ,   num_inference_steps=2     ,   output_type="pt"						).frames
								__UpperCamelCase	:   List[Any]			   =					video_frames.cpu().numpy()

								assert np.abs(expected_video - video						).mean() < 5E-2
import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
   import jax
   import jax.numpy as jnp
   from flax.jax_utils import replicate
   from flax.training.common_utils import shard





@slow
@require_flax
class 		_lowercase						(					unittest.TestCase       ):

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						str:

  """simple docstring"""





  super().tearDown()
  gc.collect()

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Union[str, Any]			)   ->						Tuple:

  """simple docstring"""





  UpperCamelCase_, UpperCamelCase_					:						str =					FlaxStableDiffusionPipeline.from_pretrained(
      'stabilityai/stable-diffusion-2'   ,      revision='bf16'   ,      dtype=jnp.bfloataa   ,      )

  UpperCamelCase_					:						str =					'A painting of a squirrel eating a burger'
  UpperCamelCase_					:						Any =					jax.device_count()
  UpperCamelCase_					:						List[str] =					num_samples * [prompt]
  UpperCamelCase_					:						List[Any] =					sd_pipe.prepare_inputs(snake_case			)

  UpperCamelCase_					:						Dict =					replicate(snake_case			)
  UpperCamelCase_					:						Optional[Any] =					shard(snake_case			)

  UpperCamelCase_					:						Dict =					jax.random.PRNGKey(0			)
  UpperCamelCase_					:						Tuple =					jax.random.split(snake_case   ,      jax.device_count()			)

  UpperCamelCase_					:						Optional[Any] =					sd_pipe(snake_case   ,      snake_case   ,      snake_case   ,      num_inference_steps=2_5   ,      jit=snake_case			)[0]
  assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

  UpperCamelCase_					:						int =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]			)
  UpperCamelCase_					:						Tuple =					images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

  UpperCamelCase_					:						str =					jnp.asarray(jax.device_get(image_slice.flatten()			)			)
  UpperCamelCase_					:						List[str] =					jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512]			)
  print(f"output_slice: {output_slice}"			)
  assert jnp.abs(output_slice - expected_slice			).max() < 1e-2


 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any			)   ->						List[Any]:

  """simple docstring"""





  UpperCamelCase_					:						Tuple =					'stabilityai/stable-diffusion-2'
  UpperCamelCase_, UpperCamelCase_					:						Tuple =					FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case   ,      subfolder='scheduler'			)
  UpperCamelCase_, UpperCamelCase_					:						Union[str, Any] =					FlaxStableDiffusionPipeline.from_pretrained(
      snake_case   ,      scheduler=snake_case   ,      revision='bf16'   ,      dtype=jnp.bfloataa   ,      )
  UpperCamelCase_					:						Optional[int] =					scheduler_params

  UpperCamelCase_					:						Optional[Any] =					'A painting of a squirrel eating a burger'
  UpperCamelCase_					:						Union[str, Any] =					jax.device_count()
  UpperCamelCase_					:						Union[str, Any] =					num_samples * [prompt]
  UpperCamelCase_					:						Tuple =					sd_pipe.prepare_inputs(snake_case			)

  UpperCamelCase_					:						List[Any] =					replicate(snake_case			)
  UpperCamelCase_					:						Optional[Any] =					shard(snake_case			)

  UpperCamelCase_					:						Tuple =					jax.random.PRNGKey(0			)
  UpperCamelCase_					:						str =					jax.random.split(snake_case   ,      jax.device_count()			)

  UpperCamelCase_					:						str =					sd_pipe(snake_case   ,      snake_case   ,      snake_case   ,      num_inference_steps=2_5   ,      jit=snake_case			)[0]
  assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

  UpperCamelCase_					:						str =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]			)
  UpperCamelCase_					:						int =					images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

  UpperCamelCase_					:						int =					jnp.asarray(jax.device_get(image_slice.flatten()			)			)
  UpperCamelCase_					:						Union[str, Any] =					jnp.array([0.4336, 0.42969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]			)
  print(f"output_slice: {output_slice}"			)
  assert jnp.abs(output_slice - expected_slice			).max() < 1e-2



