
"""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  snake_case     (							UpperCamelCase_       ,			unittest.TestCase   ):
							UpperCAmelCase__						     =      VideoToVideoSDPipeline
							UpperCAmelCase__						     =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({'''video'''}   ) - {'''image''', '''width''', '''height'''}
							UpperCAmelCase__						     =      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''video'''}   ) - {'''image'''}
							UpperCAmelCase__						     =      PipelineTesterMixin.required_optional_params - {'''latents'''}
							UpperCAmelCase__						     =      False

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

							def _lowercase					(self    ):

									"""simple docstring"""







									torch.manual_seed(0    )
									SCREAMING_SNAKE_CASE_						 = 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       ,       )
									SCREAMING_SNAKE_CASE_						 = 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    )
									SCREAMING_SNAKE_CASE_						 = 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=1_28       ,       )
									torch.manual_seed(0    )
									SCREAMING_SNAKE_CASE_						 = 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=10_00       ,       hidden_act='''gelu'''       ,       projection_dim=5_12       ,       )
									SCREAMING_SNAKE_CASE_						 = CLIPTextModel(__a    )
									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 _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_=0    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = floats_tensor((1, 3, 3, 32, 32)       ,       rng=random.Random(__a    )    ).to(__a    )

									if str(__a    ).startswith('''mps'''    ):
											SCREAMING_SNAKE_CASE_						 = torch.manual_seed(__a    )
									else:
											SCREAMING_SNAKE_CASE_						 = torch.Generator(device=__a    ).manual_seed(__a    )
									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 _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 'cpu'  # ensure determinism for the device-dependent torch.Generator
									SCREAMING_SNAKE_CASE_						 = self.get_dummy_components()
									SCREAMING_SNAKE_CASE_						 = VideoToVideoSDPipeline(**__a    )
									SCREAMING_SNAKE_CASE_						 = sd_pipe.to(__a    )
									sd_pipe.set_progress_bar_config(disable=__a    )

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

									assert frames[0].shape == (32, 32, 3)
									SCREAMING_SNAKE_CASE_						 = np.array([1_06, 1_17, 1_13, 1_74, 1_37, 1_12, 1_48, 1_51, 1_31]    )

									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    ):

									"""simple docstring"""







									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    ):

									"""simple docstring"""







									pass

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

									"""simple docstring"""







									pass

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

									"""simple docstring"""







									pass







							def _lowercase					(self    ):

									"""simple docstring"""







									return super().test_progress_bar()







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

							def _lowercase					(self    ):

									"""simple docstring"""







									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, 10, 3, 10_24, 5_76)       ,       generator=__a    )
									SCREAMING_SNAKE_CASE_						 = video.to('''cuda'''    )

									SCREAMING_SNAKE_CASE_						 = 'Spiderman is surfing'

									SCREAMING_SNAKE_CASE_						 = pipe(__a       ,       video=__a       ,       generator=__a       ,       num_inference_steps=3       ,       output_type='''pt'''    ).frames

									SCREAMING_SNAKE_CASE_						 = np.array([-1.0_45_89_84, -1.1_27_92_97, -0.9_66_30_86, -0.91_50_39_06, -0.75_09_76_56]    )
									assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array    ).sum() < 1e-2

"""simple docstring"""




import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
						from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm



def        _lowerCamelCase ( __a    ):
		if is_torch_version('''<''',							'''2.0.0'''    ) or not hasattr(__a,							'''_dynamo'''    ):
				return False
		return isinstance(__a,							torch._dynamo.eval_frame.OptimizedModule    )



def        _lowerCamelCase ( __a,							__a = True    ):
		SCREAMING_SNAKE_CASE_						 = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

		SCREAMING_SNAKE_CASE_						 = is_compiled_module(__a    )
		if is_compiled:
				SCREAMING_SNAKE_CASE_						 = model
				SCREAMING_SNAKE_CASE_						 = model._orig_mod

		if is_deepspeed_available():
				options += (DeepSpeedEngine,)

		while isinstance(__a,							__a    ):
				SCREAMING_SNAKE_CASE_						 = model.module

		if not keep_fpaa_wrapper:
				SCREAMING_SNAKE_CASE_						 = getattr(__a,							'''forward'''    )
				SCREAMING_SNAKE_CASE_						 = model.__dict__.pop('''_original_forward''',							__a    )
				if original_forward is not None:
						while hasattr(__a,							'''__wrapped__'''    ):
								SCREAMING_SNAKE_CASE_						 = forward.__wrapped__
								if forward == original_forward:
										break
						SCREAMING_SNAKE_CASE_						 = forward
				if getattr(__a,							'''_converted_to_transformer_engine''',							__a    ):
						convert_model(__a,							to_transformer_engine=__a    )

		if is_compiled:
				SCREAMING_SNAKE_CASE_						 = model
				SCREAMING_SNAKE_CASE_						 = compiled_model

		return model



def        _lowerCamelCase ( ):
		PartialState().wait_for_everyone()



def        _lowerCamelCase ( __a,							__a    ):
		if PartialState().distributed_type == DistributedType.TPU:
				xm.save(__a,							__a    )
		elif PartialState().local_process_index == 0:
				torch.save(__a,							__a    )



@contextmanager
def        _lowerCamelCase ( **__a    ):
		for key, value in kwargs.items():
				SCREAMING_SNAKE_CASE_						 = str(__a    )

		yield

		for key in kwargs:
				if key.upper() in os.environ:
						del os.environ[key.upper()]



def        _lowerCamelCase ( __a    ):
		if not hasattr(__a,							'''__qualname__'''    ) and not hasattr(__a,							'''__name__'''    ):
				SCREAMING_SNAKE_CASE_						 = getattr(__a,							'''__class__''',							__a    )
		if hasattr(__a,							'''__qualname__'''    ):
				return obj.__qualname__
		if hasattr(__a,							'''__name__'''    ):
				return obj.__name__
		return str(__a    )



def        _lowerCamelCase ( __a,							__a    ):
		for key, value in source.items():
				if isinstance(__a,							__a    ):
						SCREAMING_SNAKE_CASE_						 = destination.setdefault(__a,							{}    )
						merge_dicts(__a,							__a    )
				else:
						SCREAMING_SNAKE_CASE_						 = value

		return destination



def        _lowerCamelCase ( __a = None    ):
		if port is None:
				SCREAMING_SNAKE_CASE_						 = 29_500
		with socket.socket(socket.AF_INET,							socket.SOCK_STREAM    ) as s:
				return s.connect_ex(('''localhost''', port)    ) == 0