




'''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 							_a				(			SCREAMING_SNAKE_CASE ,      unittest.TestCase ):





       '''simple docstring'''






       A						:				Optional[Any]        =						VideoToVideoSDPipeline
       A						:				Optional[Any]        =						TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({'''video'''} ) - {'''image''', '''width''', '''height'''}
       A						:				Tuple        =						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''video'''} ) - {'''image'''}
       A						:				Optional[Any]        =						PipelineTesterMixin.required_optional_params - {'''latents'''}
       A						:				Optional[Any]        =						False

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







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             torch.manual_seed(0  )
             SCREAMING_SNAKE_CASE   :      int    					=						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   :      Tuple    					=						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   :      str    					=						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   :      int    					=						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=1_000,							hidden_act='gelu',							projection_dim=512,							)
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						CLIPTextModel(A  )
             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

             SCREAMING_SNAKE_CASE   :      List[Any]    					=						{
                 'unet': unet,
                 'scheduler': scheduler,
                 'vae': vae,
                 'text_encoder': text_encoder,
                 'tokenizer': tokenizer,
             }
             return components







       def        UpperCamelCase_							(      self,							A,							A=0  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						floats_tensor((1, 3, 3, 32, 32),							rng=random.Random(A  )  ).to(A  )

             if str(A  ).startswith('mps'  ):
                   SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.manual_seed(A  )
             else:
                   SCREAMING_SNAKE_CASE   :      Dict    					=						torch.Generator(device=A  ).manual_seed(A  )
             SCREAMING_SNAKE_CASE   :      Union[str, 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        UpperCamelCase_							(      self  ):







             '''simple docstring'''







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

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

             assert frames[0].shape == (32, 32, 3)
             SCREAMING_SNAKE_CASE   :      Dict    					=						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        UpperCamelCase_							(      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        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             pass







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







             '''simple docstring'''







             pass







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







             '''simple docstring'''







             pass







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             return super().test_progress_bar()






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





       '''simple docstring'''






       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						VideoToVideoSDPipeline.from_pretrained('cerspense/zeroscope_v2_XL',							torch_dtype=torch.floataa  )
             pipe.enable_model_cpu_offload()

             # 10 frames
             SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.Generator(device='cpu'  ).manual_seed(0  )
             SCREAMING_SNAKE_CASE   :      Dict    					=						torch.randn((1, 10, 3, 1_024, 576),							generator=A  )
             SCREAMING_SNAKE_CASE   :      Dict    					=						video.to('cuda'  )

             SCREAMING_SNAKE_CASE   :      Tuple    					=						'Spiderman is surfing'

             SCREAMING_SNAKE_CASE   :      int    					=						pipe(A,							video=A,							generator=A,							num_inference_steps=3,							output_type='pt'  ).frames

             SCREAMING_SNAKE_CASE   :      Dict    					=						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'''



from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


UpperCamelCase_    	=				logging.get_logger(__name__)


UpperCamelCase_    	=				TypeVar("DatasetType", Dataset, IterableDataset)




def   lowercase__(			__UpperCamelCase:					List[DatasetType]      ,__UpperCamelCase:					Optional[List[float]] = None      ,__UpperCamelCase:					Optional[int] = None      ,__UpperCamelCase:					Optional[DatasetInfo] = None      ,__UpperCamelCase:					Optional[NamedSplit] = None      ,__UpperCamelCase:					Literal["first_exhausted", "all_exhausted"] = "first_exhausted"      ,):

      """simple docstring"""






      from .arrow_dataset import Dataset
      from .iterable_dataset import IterableDataset

      if not datasets:
            raise ValueError('Unable to interleave an empty list of datasets.'    )
      for i, dataset in enumerate(__UpperCamelCase    ):
            if not isinstance(__UpperCamelCase      ,(Dataset, IterableDataset)    ):
                  if isinstance(__UpperCamelCase      ,(DatasetDict, IterableDatasetDict)    ):
                        if not dataset:
                              raise ValueError(
                                  f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
                                  'is an empty dataset dictionary.'    )
                        raise ValueError(
                            f"Dataset at position {i} has at least one split: {list(__UpperCamelCase    )}\n"
                            f"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase    )    )}']"    )
                  raise ValueError(
                      f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase    ).__name__}."    )
            if i == 0:
                  SCREAMING_SNAKE_CASE					,    SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						(
                      (Dataset, IterableDataset) if isinstance(__UpperCamelCase      ,__UpperCamelCase    ) else (IterableDataset, Dataset)
                  )
            elif not isinstance(__UpperCamelCase      ,__UpperCamelCase    ):
                  raise ValueError(
                      f"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."    )
      if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
            raise ValueError(f"{stopping_strategy} is not supported. Please enter a valid stopping_strategy."    )
      if dataset_type is Dataset:
            return _interleave_map_style_datasets(
                __UpperCamelCase      ,__UpperCamelCase      ,__UpperCamelCase      ,info=__UpperCamelCase      ,split=__UpperCamelCase      ,stopping_strategy=__UpperCamelCase    )
      else:
            return _interleave_iterable_datasets(
                __UpperCamelCase      ,__UpperCamelCase      ,__UpperCamelCase      ,info=__UpperCamelCase      ,split=__UpperCamelCase      ,stopping_strategy=__UpperCamelCase    )





def   lowercase__(			__UpperCamelCase:					List[DatasetType]      ,__UpperCamelCase:					Optional[DatasetInfo] = None      ,__UpperCamelCase:					Optional[NamedSplit] = None      ,__UpperCamelCase:					int = 0      ,):

      """simple docstring"""






      if not dsets:
            raise ValueError('Unable to concatenate an empty list of datasets.'    )
      for i, dataset in enumerate(__UpperCamelCase    ):
            if not isinstance(__UpperCamelCase      ,(Dataset, IterableDataset)    ):
                  if isinstance(__UpperCamelCase      ,(DatasetDict, IterableDatasetDict)    ):
                        if not dataset:
                              raise ValueError(
                                  f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
                                  'is an empty dataset dictionary.'    )
                        raise ValueError(
                            f"Dataset at position {i} has at least one split: {list(__UpperCamelCase    )}\n"
                            f"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase    )    )}']"    )
                  raise ValueError(
                      f"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase    ).__name__}."    )
            if i == 0:
                  SCREAMING_SNAKE_CASE					,    SCREAMING_SNAKE_CASE   :      int    					=						(
                      (Dataset, IterableDataset) if isinstance(__UpperCamelCase      ,__UpperCamelCase    ) else (IterableDataset, Dataset)
                  )
            elif not isinstance(__UpperCamelCase      ,__UpperCamelCase    ):
                  raise ValueError(
                      f"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."    )
      if dataset_type is Dataset:
            return _concatenate_map_style_datasets(__UpperCamelCase      ,info=__UpperCamelCase      ,split=__UpperCamelCase      ,axis=__UpperCamelCase    )
      else:
            return _concatenate_iterable_datasets(__UpperCamelCase      ,info=__UpperCamelCase      ,split=__UpperCamelCase      ,axis=__UpperCamelCase    )


