

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 				UpperCAmelCase_		(			__lowercase				,       unittest.TestCase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						Any	    =    VideoToVideoSDPipeline
							UpperCAmelCase__       :						Dict	    =    TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"video"}   ) - {"image", "width", "height"}
							UpperCAmelCase__       :						Union[str, Any]	    =    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"video"}   ) - {"image"}
							UpperCAmelCase__       :						Dict	    =    PipelineTesterMixin.required_optional_params - {"latents"}
							UpperCAmelCase__       :						Any	    =    False

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


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

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


							def 				__lowercase				(			self       ,							_a       ,							_a=0       )    ->    int:
												# 3 frames
												_a	:    int         =		floats_tensor((1, 3, 3, 3_2, 3_2)       ,							rng=random.Random(_a       )       ).to(_a       )

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

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

												assert frames[0].shape == (3_2, 3_2, 3)
												_a	:    List[Any]         =		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 				__lowercase				(			self       )    ->    List[str]:
												self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_a       ,							expected_max_diff=5e-3       )


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


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


							@unittest.skip(reason='''`num_images_per_prompt` argument is not supported for this pipeline.'''       )
							def 				__lowercase				(			self       )    ->    Optional[int]:
												pass




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



@slow
@skip_mps
class 				UpperCAmelCase_		(			unittest.TestCase   ):
							"""simple docstring"""







							def 				__lowercase				(			self       )    ->    int:
												_a	:    Optional[Any]         =		VideoToVideoSDPipeline.from_pretrained('''cerspense/zeroscope_v2_XL'''       ,							torch_dtype=torch.floataa       )
												pipe.enable_model_cpu_offload()

												# 10 frames
												_a	:    Any         =		torch.Generator(device='''cpu'''       ).manual_seed(0       )
												_a	:    Tuple         =		torch.randn((1, 1_0, 3, 1_0_2_4, 5_7_6)       ,							generator=_a       )
												_a	:    str         =		video.to('''cuda'''       )

												_a	:    int         =		'''Spiderman is surfing'''

												_a	:    Optional[int]         =		pipe(_a       ,							video=_a       ,							generator=_a       ,							num_inference_steps=3       ,							output_type='''pt'''       ).frames

												_a	:    Optional[int]         =		np.array([-1.045_8984, -1.127_9297, -0.966_3086, -0.9150_3906, -0.7509_7656]       )
												assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array       ).sum() < 1e-2






import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


a__				=		logging.get_logger(__name__)



def    __UpperCAmelCase (	__a      :			Dict	)			->     Tuple:

					"""simple docstring"""

					_a	:    Any         =		R'''\w+[.]\d+'''
					_a	:    Union[str, Any]         =		re.findall(__a		,__a	)
					for pat in pats:
										_a	:    int         =		key.replace(__a		,'''_'''.join(pat.split('''.'''	)	)	)
					return key



def    __UpperCAmelCase (	__a      :			List[str]		,__a      :			Union[str, Any]		,__a      :			Optional[int]	)			->     Tuple:

					"""simple docstring"""

					_a	:    Dict         =		pt_tuple_key[:-1] + ('''scale''',)
					if (
					    any('''norm''' in str_ for str_ in pt_tuple_key	)
					    and (pt_tuple_key[-1] == "bias")
					    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
					    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
					):
										_a	:    Dict         =		pt_tuple_key[:-1] + ('''scale''',)
										return renamed_pt_tuple_key, pt_tensor
					elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
										_a	:    Optional[int]         =		pt_tuple_key[:-1] + ('''scale''',)
										return renamed_pt_tuple_key, pt_tensor

					# embedding
					if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
										_a	:    Union[str, Any]         =		pt_tuple_key[:-1] + ('''embedding''',)
										return renamed_pt_tuple_key, pt_tensor

					# conv layer
					_a	:    Tuple         =		pt_tuple_key[:-1] + ('''kernel''',)
					if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
										_a	:    List[str]         =		pt_tensor.transpose(2		,3		,1		,0	)
										return renamed_pt_tuple_key, pt_tensor

					# linear layer
					_a	:    Dict         =		pt_tuple_key[:-1] + ('''kernel''',)
					if pt_tuple_key[-1] == "weight":
										_a	:    Union[str, Any]         =		pt_tensor.T
										return renamed_pt_tuple_key, pt_tensor

					# old PyTorch layer norm weight
					_a	:    Dict         =		pt_tuple_key[:-1] + ('''weight''',)
					if pt_tuple_key[-1] == "gamma":
										return renamed_pt_tuple_key, pt_tensor

					# old PyTorch layer norm bias
					_a	:    Union[str, Any]         =		pt_tuple_key[:-1] + ('''bias''',)
					if pt_tuple_key[-1] == "beta":
										return renamed_pt_tuple_key, pt_tensor

					return pt_tuple_key, pt_tensor







def    __UpperCAmelCase (	__a      :			Dict		,__a      :			str		,__a      :			str=42	)			->     Optional[int]:

					"""simple docstring"""

					_a	:    Tuple         =		{k: v.numpy() for k, v in pt_state_dict.items()}

					# Step 2: Since the model is stateless, get random Flax params
					_a	:    List[Any]         =		flax_model.init_weights(PRNGKey(__a	)	)

					_a	:    Optional[int]         =		flatten_dict(__a	)
					_a	:    Tuple         =		{}

					# Need to change some parameters name to match Flax names
					for pt_key, pt_tensor in pt_state_dict.items():
										_a	:    List[str]         =		rename_key(__a	)
										_a	:    Optional[Any]         =		tuple(renamed_pt_key.split('''.'''	)	)

										# Correctly rename weight parameters
										_a      ,					_a	:    List[str]         =		rename_key_and_reshape_tensor(__a		,__a		,__a	)

										if flax_key in random_flax_state_dict:
															if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																				raise ValueError(
																				    F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
																				    F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""	)

        # also add unexpected weight so that warning is thrown
										_a	:    Dict         =		jnp.asarray(__a	)

					return unflatten_dict(__a	)



