




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

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





	def 						A_(							self      )  ->					Tuple:




				"""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_0_0_8_5      ,     beta_end=0.0_1_2      ,     beta_schedule='scaled_linear'      ,     clip_sample=SCREAMING_SNAKE_CASE      ,     set_alpha_to_one=SCREAMING_SNAKE_CASE      ,     )
				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=128      ,     )
				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-05      ,     num_attention_heads=4      ,     num_hidden_layers=5      ,     pad_token_id=1      ,     vocab_size=1000      ,     hidden_act='gelu'      ,     projection_dim=512      ,     )
				SCREAMING_SNAKE_CASE_								=						CLIPTextModel(SCREAMING_SNAKE_CASE      )
				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 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE=0      )  ->					Union[str, Any]:




				"""simple docstring"""



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

				if str(SCREAMING_SNAKE_CASE      ).startswith('mps'      ):
							SCREAMING_SNAKE_CASE_								=						torch.manual_seed(SCREAMING_SNAKE_CASE      )
				else:
							SCREAMING_SNAKE_CASE_								=						torch.Generator(device=SCREAMING_SNAKE_CASE      ).manual_seed(SCREAMING_SNAKE_CASE      )
				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 						A_(							self      )  ->					Union[str, Any]:




				"""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(**SCREAMING_SNAKE_CASE      )
				SCREAMING_SNAKE_CASE_								=						sd_pipe.to(SCREAMING_SNAKE_CASE      )
				sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

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

				assert frames[0].shape == (32, 32, 3)
				SCREAMING_SNAKE_CASE_								=						np.array([106, 117, 113, 174, 137, 112, 148, 151, 131]      )

				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 						A_(							self      )  ->					Optional[Any]:




				"""simple docstring"""



				self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=SCREAMING_SNAKE_CASE      ,     expected_max_diff=5e-3      )





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




				"""simple docstring"""



				pass





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




				"""simple docstring"""



				pass





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




				"""simple docstring"""



				pass





	def 						A_(							self      )  ->					Optional[int]:




				"""simple docstring"""



				return super().test_progress_bar()




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





	def 						A_(							self      )  ->					str:




				"""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, 1024, 576)      ,     generator=SCREAMING_SNAKE_CASE      )
				SCREAMING_SNAKE_CASE_								=						video.to('cuda'      )

				SCREAMING_SNAKE_CASE_								=						'Spiderman is surfing'

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

				SCREAMING_SNAKE_CASE_								=						np.array([-1.0_4_5_8_9_8_4, -1.1_2_7_9_2_9_7, -0.9_6_6_3_0_8_6, -0.9_1_5_0_3_9_0_6, -0.7_5_0_9_7_6_5_6]      )
				assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array      ).sum() < 1e-2






import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn

from ...models.controlnet import ControlNetModel, ControlNetOutput
from ...models.modeling_utils import ModelMixin
from ...utils import logging


SCREAMING_SNAKE_CASE__     :     Union[str, Any]		    =       logging.get_logger(__name__)



class       a_       (     SCREAMING_SNAKE_CASE__     ):





	def __init__(							self      ,     SCREAMING_SNAKE_CASE      )  ->					Union[str, Any]:




				"""simple docstring"""



				super().__init__()
				SCREAMING_SNAKE_CASE_								=						nn.ModuleList(SCREAMING_SNAKE_CASE      )





	def 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = False      ,     SCREAMING_SNAKE_CASE = True      ,     )  ->					Union[ControlNetOutput, Tuple]:




				"""simple docstring"""



				for i, (image, scale, controlnet) in enumerate(zip(SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     self.nets      )      ):
							SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						controlnet(
							    SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      ,     )

							# merge samples
							if i == 0:
										SCREAMING_SNAKE_CASE_						,		SCREAMING_SNAKE_CASE_								=						down_samples, mid_sample
							else:
										SCREAMING_SNAKE_CASE_								=						[
										    samples_prev + samples_curr
										    for samples_prev, samples_curr in zip(SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      )
										]
										mid_block_res_sample += mid_sample

				return down_block_res_samples, mid_block_res_sample





	def 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE = True      ,     SCREAMING_SNAKE_CASE = None      ,     SCREAMING_SNAKE_CASE = False      ,     SCREAMING_SNAKE_CASE = None      ,     )  ->					int:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						0
				SCREAMING_SNAKE_CASE_								=						save_directory
				for controlnet in self.nets:
							controlnet.save_pretrained(
							    SCREAMING_SNAKE_CASE      ,     is_main_process=SCREAMING_SNAKE_CASE      ,     save_function=SCREAMING_SNAKE_CASE      ,     safe_serialization=SCREAMING_SNAKE_CASE      ,     variant=SCREAMING_SNAKE_CASE      ,     )

							idx += 1
							SCREAMING_SNAKE_CASE_								=						model_path_to_save + f'_{idx}'





	@classmethod
	def 						A_(							cls      ,     SCREAMING_SNAKE_CASE      ,     **SCREAMING_SNAKE_CASE      )  ->					Tuple:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						0
				SCREAMING_SNAKE_CASE_								=						[]

				# load controlnet and append to list until no controlnet directory exists anymore
				# first controlnet has to be saved under `./mydirectory/controlnet` to be compliant with `DiffusionPipeline.from_prertained`
				# second, third, ... controlnets have to be saved under `./mydirectory/controlnet_1`, `./mydirectory/controlnet_2`, ...
				SCREAMING_SNAKE_CASE_								=						pretrained_model_path
				while os.path.isdir(SCREAMING_SNAKE_CASE      ):
							SCREAMING_SNAKE_CASE_								=						ControlNetModel.from_pretrained(SCREAMING_SNAKE_CASE      ,     **SCREAMING_SNAKE_CASE      )
							controlnets.append(SCREAMING_SNAKE_CASE      )

							idx += 1
							SCREAMING_SNAKE_CASE_								=						pretrained_model_path + f'_{idx}'

				logger.info(f'{len(SCREAMING_SNAKE_CASE      )} controlnets loaded from {pretrained_model_path}.'      )

				if len(SCREAMING_SNAKE_CASE      ) == 0:
							raise ValueError(
							    f'No ControlNets found under {os.path.dirname(SCREAMING_SNAKE_CASE      )}. Expected at least {pretrained_model_path + "_0"}.'      )

				return cls(SCREAMING_SNAKE_CASE      )
