



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        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ,							unittest.TestCase   ):







 '''simple docstring'''
 lowercase_		=      VideoToVideoSDPipeline
 lowercase_		=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({"video"}   ) - {"image", "width", "height"}
 lowercase_		=      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({"video"}   ) - {"image"}
 lowercase_		=      PipelineTesterMixin.required_optional_params - {"latents"}
 lowercase_		=      False

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

 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple)	->Dict:



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__:				Union[str, Any]        =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__:				int        =DDIMScheduler(
       beta_start=0.0_0085	,			beta_end=0.012	,			beta_schedule="scaled_linear"	,			clip_sample=__UpperCAmelCase	,			set_alpha_to_one=__UpperCAmelCase	,			)
   torch.manual_seed(0)
   lowerCamelCase__:				Tuple        =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__:				Dict        =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=1_000	,			hidden_act="gelu"	,			projection_dim=512	,			)
   lowerCamelCase__:				Tuple        =CLIPTextModel(__UpperCAmelCase)
   lowerCamelCase__:				int        =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

   lowerCamelCase__:				Optional[int]        ={
       "unet": unet,
       "scheduler": scheduler,
       "vae": vae,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
   }
   return components

 def 							SCREAMING_SNAKE_CASE_ (self    :       str	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Optional[Any]=0)	->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =floats_tensor((1, 3, 3, 32, 32)	,			rng=random.Random(__UpperCAmelCase)).to(__UpperCAmelCase)

   if str(__UpperCAmelCase).startswith("mps"):
     lowerCamelCase__:				Any        =torch.manual_seed(__UpperCAmelCase)
   else:
     lowerCamelCase__:				Tuple        =torch.Generator(device=__UpperCAmelCase).manual_seed(__UpperCAmelCase)
   lowerCamelCase__:				int        ={
       "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    :       Optional[int])	->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        ="cpu"  # ensure determinism for the device-dependent torch.Generator
   lowerCamelCase__:				List[str]        =self.get_dummy_components()
   lowerCamelCase__:				str        =VideoToVideoSDPipeline(**__UpperCAmelCase)
   lowerCamelCase__:				Optional[Any]        =sd_pipe.to(__UpperCAmelCase)
   sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase)

   lowerCamelCase__:				str        =self.get_dummy_inputs(__UpperCAmelCase)
   lowerCamelCase__:				Dict        ="np"
   lowerCamelCase__:				int        =sd_pipe(**__UpperCAmelCase).frames
   lowerCamelCase__:				Optional[int]        =frames[0][-3:, -3:, -1]

   assert frames[0].shape == (32, 32, 3)
   lowerCamelCase__:				Optional[int]        =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 							SCREAMING_SNAKE_CASE_ (self    :       List[Any])	->Dict:



   '''simple docstring'''



   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    :       List[str])	->Union[str, Any]:



   '''simple docstring'''



   pass

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



   '''simple docstring'''



   pass

 @unittest.skip(reason="`num_images_per_prompt` argument is not supported for this pipeline.")
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->List[Any]:



   '''simple docstring'''



   pass







 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict)	->str:



   '''simple docstring'''



   return super().test_progress_bar()



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







 '''simple docstring'''

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



   '''simple docstring'''



   lowerCamelCase__:				Dict        =VideoToVideoSDPipeline.from_pretrained("cerspense/zeroscope_v2_XL"	,			torch_dtype=torch.floataa)
   pipe.enable_model_cpu_offload()

   # 10 frames
   lowerCamelCase__:				int        =torch.Generator(device="cpu").manual_seed(0)
   lowerCamelCase__:				Dict        =torch.randn((1, 10, 3, 1_024, 576)	,			generator=__UpperCAmelCase)
   lowerCamelCase__:				Any        =video.to("cuda")

   lowerCamelCase__:				List[str]        ="Spiderman is surfing"

   lowerCamelCase__:				Dict        =pipe(__UpperCAmelCase	,			video=__UpperCAmelCase	,			generator=__UpperCAmelCase	,			num_inference_steps=3	,			output_type="pt").frames

   lowerCamelCase__:				List[Any]        =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








from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''

 def __init__(self    :       int	,			UpperCAmelCase_    :       Distribution	,			UpperCAmelCase_    :       Dict=None	,			UpperCAmelCase_    :       Tuple=None	,			UpperCAmelCase_    :       Tuple=0)	->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =1.0 if scale is None else scale
   lowerCamelCase__:				List[Any]        =0.0 if loc is None else loc

   super().__init__(UpperCAmelCase_	,			[AffineTransform(loc=self.loc	,			scale=self.scale	,			event_dim=UpperCAmelCase_)])

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any])	->int:



   '''simple docstring'''



   return self.base_dist.mean * self.scale + self.loc

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Dict:



   '''simple docstring'''



   return self.base_dist.variance * self.scale**2







 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Optional[Any]:



   '''simple docstring'''



   return self.variance.sqrt()



class        _SCREAMING_SNAKE_CASE						(     nn.Module   ):







 '''simple docstring'''

 def __init__(self    :       Tuple	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Dict[str, int]	,			UpperCAmelCase_    :       Callable[..., Tuple[torch.Tensor]]	,			**UpperCAmelCase_    :       Dict)	->None:



   '''simple docstring'''



   super().__init__(**UpperCAmelCase_)
   lowerCamelCase__:				Tuple        =args_dim
   lowerCamelCase__:				Any        =nn.ModuleList([nn.Linear(UpperCAmelCase_	,			UpperCAmelCase_) for dim in args_dim.values()])
   lowerCamelCase__:				Any        =domain_map







 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any]	,			UpperCAmelCase_    :       torch.Tensor)	->Tuple[torch.Tensor]:



   '''simple docstring'''



   lowerCamelCase__:				Any        =[proj(UpperCAmelCase_) for proj in self.proj]

   return self.domain_map(*UpperCAmelCase_)



class        _SCREAMING_SNAKE_CASE						(     nn.Module   ):







 '''simple docstring'''

 def __init__(self    :       str	,			UpperCAmelCase_    :       Optional[Any])	->Optional[int]:



   '''simple docstring'''



   super().__init__()
   lowerCamelCase__:				Any        =function







 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict	,			UpperCAmelCase_    :       Optional[Any]	,			*UpperCAmelCase_    :       Union[str, Any])	->Union[str, Any]:



   '''simple docstring'''



   return self.function(UpperCAmelCase_	,			*UpperCAmelCase_)



class        _SCREAMING_SNAKE_CASE						:







 '''simple docstring'''
 lowercase_		=      42
 lowercase_		=      42
 lowercase_		=      42

 def __init__(self    :       str	,			UpperCAmelCase_    :       int = 1)	->None:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =dim
   lowerCamelCase__:				int        ={k: dim * self.args_dim[k] for k in self.args_dim}

 def 							SCREAMING_SNAKE_CASE_ (self    :       int	,			UpperCAmelCase_    :       str)	->Optional[Any]:



   '''simple docstring'''



   if self.dim == 1:
     return self.distribution_class(*UpperCAmelCase_)
   else:
     return Independent(self.distribution_class(*UpperCAmelCase_)	,			1)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			UpperCAmelCase_    :       List[Any]	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =self._base_distribution(UpperCAmelCase_)
   if loc is None and scale is None:
     return distr
   else:
     return AffineTransformed(UpperCAmelCase_	,			loc=UpperCAmelCase_	,			scale=UpperCAmelCase_	,			event_dim=self.event_dim)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->Tuple:



   '''simple docstring'''



   return () if self.dim == 1 else (self.dim,)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->int:



   '''simple docstring'''



   return len(self.event_shape)

 @property
 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple)	->float:



   '''simple docstring'''



   return 0.0

 def 							SCREAMING_SNAKE_CASE_ (self    :       str	,			UpperCAmelCase_    :       int)	->nn.Module:



   '''simple docstring'''



   return ParameterProjection(
       in_features=UpperCAmelCase_	,			args_dim=self.args_dim	,			domain_map=LambdaLayer(self.domain_map)	,			)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Any	,			*UpperCAmelCase_    :       torch.Tensor)	->Any:



   '''simple docstring'''



   raise NotImplementedError()







 @staticmethod
 def 							SCREAMING_SNAKE_CASE_ (UpperCAmelCase_    :       torch.Tensor)	->torch.Tensor:



   '''simple docstring'''



   return (x + torch.sqrt(torch.square(UpperCAmelCase_) + 4.0)) / 2.0



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"df": 1, "loc": 1, "scale": 1}
 lowercase_		=      StudentT

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[int]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->Dict:



   '''simple docstring'''



   lowerCamelCase__:				str        =cls.squareplus(UpperCAmelCase_).clamp_min(torch.finfo(scale.dtype).eps)
   lowerCamelCase__:				Any        =2.0 + cls.squareplus(UpperCAmelCase_)
   return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"loc": 1, "scale": 1}
 lowercase_		=      Normal

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[Any]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:				List[Any]        =cls.squareplus(UpperCAmelCase_).clamp_min(torch.finfo(scale.dtype).eps)
   return loc.squeeze(-1), scale.squeeze(-1)



class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ):







 '''simple docstring'''
 lowercase_		=      {"total_count": 1, "logits": 1}
 lowercase_		=      NegativeBinomial

 @classmethod
 def 							SCREAMING_SNAKE_CASE_ (cls    :       Optional[int]	,			UpperCAmelCase_    :       torch.Tensor	,			UpperCAmelCase_    :       torch.Tensor)	->List[str]:



   '''simple docstring'''



   lowerCamelCase__:				Tuple        =cls.squareplus(UpperCAmelCase_)
   return total_count.squeeze(-1), logits.squeeze(-1)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple	,			UpperCAmelCase_    :       Tuple)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__  , lowerCamelCase__:				Any        =distr_args
   if self.dim == 1:
     return self.distribution_class(total_count=UpperCAmelCase_	,			logits=UpperCAmelCase_)
   else:
     return Independent(self.distribution_class(total_count=UpperCAmelCase_	,			logits=UpperCAmelCase_)	,			1)







 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any]	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None	,			UpperCAmelCase_    :       Optional[torch.Tensor] = None)	->Distribution:



   '''simple docstring'''



   lowerCamelCase__  , lowerCamelCase__:				Union[str, Any]        =distr_args

   if scale is not None:
     # See scaling property of Gamma.
     logits += scale.log()

   return self._base_distribution((total_count, logits))

