






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 						UpperCAmelCase__ (							__snake_case     ,				unittest.TestCase							):
		__snake_case      :       int     =   TextToVideoSDPipeline
		__snake_case      :       Optional[Any]     =   TEXT_TO_IMAGE_PARAMS
		__snake_case      :       Union[str, Any]     =   TEXT_TO_IMAGE_BATCH_PARAMS
		# No `output_type`.
		__snake_case      :       str     =   frozenset(
		    [
		        "num_inference_steps",
		        "generator",
		        "latents",
		        "return_dict",
		        "callback",
		        "callback_steps",
		    ]							)







		def  A__      (					self		):
						torch.manual_seed(0		)
						_A	:				int   =      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       ,)
						_A	:				str   =      DDIMScheduler(
						    beta_start=0.0_00_85       ,beta_end=0.0_12       ,beta_schedule='''scaled_linear'''       ,clip_sample=A__       ,set_alpha_to_one=A__       ,)
						torch.manual_seed(0		)
						_A	:				Union[str, Any]   =      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		)
						_A	:				int   =      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       ,)
						_A	:				int   =      CLIPTextModel(A__		)
						_A	:				Optional[Any]   =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

						_A	:				Any   =      {
						    '''unet''': unet,
						    '''scheduler''': scheduler,
						    '''vae''': vae,
						    '''text_encoder''': text_encoder,
						    '''tokenizer''': tokenizer,
						}
						return components







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

						_A	:				Tuple   =      self.get_dummy_inputs(A__		)
						_A	:				Union[str, Any]   =      '''np'''
						_A	:				Optional[int]   =      sd_pipe(**A__		).frames
						_A	:				str   =      frames[0][-3:, -3:, -1]

						assert frames[0].shape == (64, 64, 3)
						_A	:				Dict   =      np.array([1_58.0, 1_60.0, 1_53.0, 1_25.0, 1_00.0, 1_21.0, 1_11.0, 93.0, 1_13.0]		)

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







		def  A__      (					self		):
						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  A__      (					self		):
						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  A__      (					self		):
						pass







		@unittest.skip(reason='''Batching needs to be properly figured out first for this pipeline.'''		)
		def  A__      (					self		):
						pass







		@unittest.skip(reason='''`num_images_per_prompt` argument is not supported for this pipeline.'''		)
		def  A__      (					self		):
						pass







		def  A__      (					self		):
						return super().test_progress_bar()






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







		def  A__      (					self		):
						_A	:				Dict   =      load_numpy(
						    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy'''		)

						_A	:				Any   =      TextToVideoSDPipeline.from_pretrained('''damo-vilab/text-to-video-ms-1.7b'''		)
						_A	:				Optional[int]   =      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
						_A	:				List[Any]   =      pipe.to('''cuda'''		)

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

						_A	:				str   =      pipe(A__       ,generator=A__       ,num_inference_steps=25       ,output_type='''pt'''		).frames
						_A	:				Optional[Any]   =      video_frames.cpu().numpy()

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







		def  A__      (					self		):
						_A	:				Optional[int]   =      load_numpy(
						    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy'''		)

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

						_A	:				List[str]   =      '''Spiderman is surfing'''
						_A	:				Union[str, Any]   =      torch.Generator(device='''cpu'''		).manual_seed(0		)

						_A	:				Tuple   =      pipe(A__       ,generator=A__       ,num_inference_steps=2       ,output_type='''pt'''		).frames
						_A	:				Optional[int]   =      video_frames.cpu().numpy()

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











import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []
						_A	:				str   =      []

						for i in range(self.num_layers		):
										_A	:				Union[str, Any]   =      self.in_channels if i == 0 else self.out_channels

										_A	:				str   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

										_A	:				List[str]   =      FlaxTransformeraDModel(
										    in_channels=self.out_channels       ,n_heads=self.num_attention_heads       ,d_head=self.out_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,only_cross_attention=self.only_cross_attention       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

						_A	:				List[str]   =      resnets
						_A	:				Any   =      attentions

						if self.add_downsample:
										_A	:				Any   =      FlaxDownsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				List[str]   =      ()

						for resnet, attn in zip(self.resnets       ,self.attentions		):
										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Tuple   =      attn(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				Optional[int]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				List[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				int   =      self.in_channels if i == 0 else self.out_channels

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)
						_A	:				Tuple   =      resnets

						if self.add_downsample:
										_A	:				Optional[Any]   =      FlaxDownsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__=True		):
						_A	:				List[Any]   =      ()

						for resnet in self.resnets:
										_A	:				Optional[int]   =      resnet(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				List[Any]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				List[str]   =      []
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Dict   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				List[str]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				Optional[int]   =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

										_A	:				Union[str, Any]   =      FlaxTransformeraDModel(
										    in_channels=self.out_channels       ,n_heads=self.num_attention_heads       ,d_head=self.out_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,only_cross_attention=self.only_cross_attention       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

						_A	:				Dict   =      resnets
						_A	:				int   =      attentions

						if self.add_upsample:
										_A	:				str   =      FlaxUpsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet, attn in zip(self.resnets       ,self.attentions		):
										# pop res hidden states
										_A	:				List[Any]   =      res_hidden_states_tuple[-1]
										_A	:				int   =      res_hidden_states_tuple[:-1]
										_A	:				List[Any]   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				str   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Optional[int]   =      attn(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Union[str, Any]   =      self.upsamplers_a(A__		)

						return hidden_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Any   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				Optional[Any]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				int   =      resnets

						if self.add_upsample:
										_A	:				List[str]   =      FlaxUpsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet in self.resnets:
										# pop res hidden states
										_A	:				Tuple   =      res_hidden_states_tuple[-1]
										_A	:				Optional[int]   =      res_hidden_states_tuple[:-1]
										_A	:				Any   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Tuple   =      self.upsamplers_a(A__		)

						return hidden_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						# there is always at least one resnet
						_A	:				List[str]   =      [
						    FlaxResnetBlockaD(
						        in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
						]

						_A	:				Union[str, Any]   =      []

						for _ in range(self.num_layers		):
										_A	:				Any   =      FlaxTransformeraDModel(
										    in_channels=self.in_channels       ,n_heads=self.num_attention_heads       ,d_head=self.in_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

										_A	:				Any   =      FlaxResnetBlockaD(
										    in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				List[str]   =      resnets
						_A	:				str   =      attentions







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				Optional[int]   =      self.resnets[0](A__       ,A__		)
						for attn, resnet in zip(self.attentions       ,self.resnets[1:]		):
										_A	:				Dict   =      attn(A__       ,A__       ,deterministic=A__		)
										_A	:				Any   =      resnet(A__       ,A__       ,deterministic=A__		)

						return hidden_states



