



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 	__snake_case   ( SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		SCREAMING_SNAKE_CASE__								=	TextToVideoSDPipeline
		SCREAMING_SNAKE_CASE__								=	TEXT_TO_IMAGE_PARAMS
		SCREAMING_SNAKE_CASE__								=	TEXT_TO_IMAGE_BATCH_PARAMS
		# No `output_type`.
		SCREAMING_SNAKE_CASE__								=	frozenset(
		    [
		        'num_inference_steps',
		        'generator',
		        'latents',
		        'return_dict',
		        'callback',
		        'callback_steps',
		    ] )
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						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.00085    ,beta_end=0.012    ,beta_schedule='scaled_linear'    ,clip_sample=a_    ,set_alpha_to_one=a_    ,)
						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-0_5    ,num_attention_heads=4    ,num_hidden_layers=5    ,pad_token_id=1    ,vocab_size=1000    ,hidden_act='gelu'    ,projection_dim=512    ,)
						lowerCAmelCase__		  =  CLIPTextModel(a_						)
						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    ,a_    ,a_=0						):







						"""simple docstring"""
						if str(a_						).startswith('mps'						):
										lowerCAmelCase__		  =  torch.manual_seed(a_						)
						else:
										lowerCAmelCase__		  =  torch.Generator(device=a_						).manual_seed(a_						)
						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						):







						"""simple docstring"""
						lowerCAmelCase__		  =  'cpu'  # ensure determinism for the device-dependent torch.Generator
						lowerCAmelCase__		  =  self.get_dummy_components()
						lowerCAmelCase__		  =  TextToVideoSDPipeline(**a_						)
						lowerCAmelCase__		  =  sd_pipe.to(a_						)
						sd_pipe.set_progress_bar_config(disable=a_						)

						lowerCAmelCase__		  =  self.get_dummy_inputs(a_						)
						lowerCAmelCase__		  =  'np'
						lowerCAmelCase__		  =  sd_pipe(**a_						).frames
						lowerCAmelCase__		  =  frames[0][-3:, -3:, -1]

						assert frames[0].shape == (64, 64, 3)
						lowerCAmelCase__		  =  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 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						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 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						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 						SCREAMING_SNAKE_CASE_					(							self						):







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







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







						"""simple docstring"""
						pass



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						return super().test_progress_bar()



@slow
@skip_mps
class 	__snake_case   ( unittest.TestCase ):
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						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(a_    ,generator=a_    ,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						):







						"""simple docstring"""
						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(a_    ,generator=a_    ,num_inference_steps=2    ,output_type='pt'						).frames
						lowerCAmelCase__		  =  video_frames.cpu().numpy()

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





import os
from collections import deque

import torch
from torch.utils.data import Dataset


class 	__snake_case   ( SCREAMING_SNAKE_CASE ):
		def __init__(							self    ,a_=""    ,a_="train"						):







						"""simple docstring"""
						assert os.path.isdir(a_						)

						lowerCAmelCase__		  =  []
						lowerCAmelCase__		  =  os.listdir(a_						)
						for story_filename in story_filenames_list:
										if "summary" in story_filename:
														continue
										lowerCAmelCase__		  =  os.path.join(a_    ,a_						)
										if not os.path.isfile(a_						):
														continue
										self.documents.append(a_						)
		def __len__(							self						):







						"""simple docstring"""
						return len(self.documents						)



		def __getitem__(							self    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.documents[idx]
						lowerCAmelCase__		  =  document_path.split('/'						)[-1]
						with open(a_    ,encoding='utf-8'						) as source:
										lowerCAmelCase__		  =  source.read()
										lowerCAmelCase__  ,   lowerCAmelCase__		  =  process_story(a_						)
						return document_name, story_lines, summary_lines







def     UpperCAmelCase_		(					snake_case__					)	->     str:



				"""simple docstring"""





				lowerCAmelCase__		  =  list(filter(lambda snake_case__					: len(snake_case__					) != 0    ,			[line.strip() for line in raw_story.split('\n'					)]					)					)

				# for some unknown reason some lines miss a period, add it
				lowerCAmelCase__		  =  [_add_missing_period(snake_case__					) for line in nonempty_lines]

				# gather article lines
				lowerCAmelCase__		  =  []
				lowerCAmelCase__		  =  deque(snake_case__					)
				while True:
								try:
												lowerCAmelCase__		  =  lines.popleft()
												if element.startswith('@highlight'					):
																break
												story_lines.append(snake_case__					)
								except IndexError:
												# if "@highlight" is absent from the file we pop
												# all elements until there is None, raising an exception.
												return story_lines, []

    # gather summary lines
				lowerCAmelCase__		  =  list(filter(lambda snake_case__					: not t.startswith('@highlight'					)    ,			snake_case__					)					)

				return story_lines, summary_lines







def     UpperCAmelCase_		(					snake_case__					)	->     Optional[Any]:



				"""simple docstring"""





				lowerCAmelCase__		  =  ['.', '!', '?', '...', '\'', '`', '"', '\u2019', '\u2019', ')']
				if line.startswith('@highlight'					):
								return line
				if line[-1] in END_TOKENS:
								return line
				return line + "."







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__    ,			snake_case__					)	->     str:



				"""simple docstring"""





				if len(snake_case__					) > block_size:
								return sequence[:block_size]
				else:
								sequence.extend([pad_token_id] * (block_size - len(snake_case__					))					)
								return sequence







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__					)	->     Optional[int]:



				"""simple docstring"""





				lowerCAmelCase__		  =  torch.ones_like(snake_case__					)
				lowerCAmelCase__		  =  sequence == pad_token_id
				lowerCAmelCase__		  =  0
				return mask







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__    ,			snake_case__					)	->     Dict:



				"""simple docstring"""





				lowerCAmelCase__		  =  [tokenizer.encode(snake_case__					) for line in story_lines]
				lowerCAmelCase__		  =  [token for sentence in story_lines_token_ids for token in sentence]
				lowerCAmelCase__		  =  [tokenizer.encode(snake_case__					) for line in summary_lines]
				lowerCAmelCase__		  =  [token for sentence in summary_lines_token_ids for token in sentence]

				return story_token_ids, summary_token_ids







def     UpperCAmelCase_		(					snake_case__    ,			snake_case__					)	->     Union[str, Any]:



				"""simple docstring"""





				lowerCAmelCase__		  =  []
				for sequence in batch:
								lowerCAmelCase__		  =  -1
								lowerCAmelCase__		  =  []
								for s in sequence:
												if s == separator_token_id:
																sentence_num += 1
												embeddings.append(sentence_num % 2					)
								batch_embeddings.append(snake_case__					)
				return torch.tensor(snake_case__					)
