

"""simple docstring"""




import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    UNetaDConditionModel,
    VideoToVideoSDPipeline,
)
from diffusers.utils import floats_tensor, is_xformers_available, skip_mps
from diffusers.utils.testing_utils import enable_full_determinism, slow, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





@skip_mps
class    lowerCamelCase						(_UpperCamelCase			,unittest.TestCase			):
			lowerCamelCase__   :		str							=							VideoToVideoSDPipeline
			lowerCamelCase__   :		Union[str, Any]							=							TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({'video'}			) - {'image', 'width', 'height'}
			lowerCamelCase__   :		Dict							=							TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'video'}			) - {'image'}
			lowerCamelCase__   :		Optional[int]							=							PipelineTesterMixin.required_optional_params - {'latents'}
			lowerCamelCase__   :		Dict							=							False

			# No `output_type`.
			lowerCamelCase__   :		str							=							frozenset(
			    [
			        'num_inference_steps',
			        'generator',
			        'latents',
			        'return_dict',
			        'callback',
			        'callback_steps',
			    ]			)

			def 					SCREAMING_SNAKE_CASE				(					self				:						Optional[int]       )   ->    int:
						torch.manual_seed(0       )
						SCREAMING_SNAKE_CASE__								=      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    ,				)
						SCREAMING_SNAKE_CASE__								=      DDIMScheduler(
						    beta_start=0.00_085    ,				beta_end=0.012    ,				beta_schedule="""scaled_linear"""    ,				clip_sample=_UpperCAmelCase    ,				set_alpha_to_one=_UpperCAmelCase    ,				)
						torch.manual_seed(0       )
						SCREAMING_SNAKE_CASE__								=      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       )
						SCREAMING_SNAKE_CASE__								=      CLIPTextConfig(
						    bos_token_id=0    ,				eos_token_id=2    ,				hidden_size=3_2    ,				intermediate_size=3_7    ,				layer_norm_eps=1e-05    ,				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    ,				)
						SCREAMING_SNAKE_CASE__								=      CLIPTextModel(_UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

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

			def 					SCREAMING_SNAKE_CASE				(					self				:						Optional[Any]    ,				__UpperCAmelCase				:						Optional[Any]    ,				__UpperCAmelCase				:						Any=0       )   ->    Optional[Any]:
						# 3 frames
						SCREAMING_SNAKE_CASE__								=      floats_tensor((1, 3, 3, 3_2, 3_2)    ,				rng=random.Random(_UpperCAmelCase       )       ).to(_UpperCAmelCase       )

						if str(_UpperCAmelCase       ).startswith("""mps"""       ):
									SCREAMING_SNAKE_CASE__								=      torch.manual_seed(_UpperCAmelCase       )
						else:
									SCREAMING_SNAKE_CASE__								=      torch.Generator(device=_UpperCAmelCase       ).manual_seed(_UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      {
						    '''prompt''': '''A painting of a squirrel eating a burger''',
						    '''video''': video,
						    '''generator''': generator,
						    '''num_inference_steps''': 2,
						    '''guidance_scale''': 6.0,
						    '''output_type''': '''pt''',
						}
						return inputs

			def 					SCREAMING_SNAKE_CASE				(					self				:						List[Any]       )   ->    List[str]:
						SCREAMING_SNAKE_CASE__								=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
						SCREAMING_SNAKE_CASE__								=      self.get_dummy_components()
						SCREAMING_SNAKE_CASE__								=      VideoToVideoSDPipeline(**_UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      sd_pipe.to(_UpperCAmelCase       )
						sd_pipe.set_progress_bar_config(disable=_UpperCAmelCase       )

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

						assert frames[0].shape == (3_2, 3_2, 3)
						SCREAMING_SNAKE_CASE__								=      np.array([1_0_6, 1_1_7, 1_1_3, 1_7_4, 1_3_7, 1_1_2, 1_4_8, 1_5_1, 1_3_1]       )

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

			@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				:						int       )   ->    List[str]:
						self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_UpperCAmelCase    ,				expected_max_diff=5e-3       )

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

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

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





			def 					SCREAMING_SNAKE_CASE				(					self				:						Tuple       )   ->    List[Any]:
						return super().test_progress_bar()





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

			def 					SCREAMING_SNAKE_CASE				(					self				:						Dict       )   ->    Any:
						SCREAMING_SNAKE_CASE__								=      VideoToVideoSDPipeline.from_pretrained("""cerspense/zeroscope_v2_XL"""    ,				torch_dtype=torch.floataa       )
						pipe.enable_model_cpu_offload()

						# 10 frames
						SCREAMING_SNAKE_CASE__								=      torch.Generator(device="""cpu"""       ).manual_seed(0       )
						SCREAMING_SNAKE_CASE__								=      torch.randn((1, 1_0, 3, 1_0_2_4, 5_7_6)    ,				generator=_UpperCAmelCase       )
						SCREAMING_SNAKE_CASE__								=      video.to("""cuda"""       )

						SCREAMING_SNAKE_CASE__								=      '''Spiderman is surfing'''

						SCREAMING_SNAKE_CASE__								=      pipe(_UpperCAmelCase    ,				video=_UpperCAmelCase    ,				generator=_UpperCAmelCase    ,				num_inference_steps=3    ,				output_type="""pt"""       ).frames

						SCREAMING_SNAKE_CASE__								=      np.array([-1.0_458_984, -1.1_279_297, -0.9_663_086, -0.91_503_906, -0.75_097_656]       )
						assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array       ).sum() < 1e-2





"""simple docstring"""




from .glue import GlueDataset, GlueDataTrainingArguments
from .language_modeling import (
    LineByLineTextDataset,
    LineByLineWithRefDataset,
    LineByLineWithSOPTextDataset,
    TextDataset,
    TextDatasetForNextSentencePrediction,
)
from .squad import SquadDataset, SquadDataTrainingArguments


