




"""simple docstring"""



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class      UpperCAmelCase_  (						__snake_case,       unittest.TestCase							):
    lowercase__						 =				KandinskyVaaImgaImgPipeline
    lowercase__						 =				['image_embeds', 'negative_image_embeds', 'image']
    lowercase__						 =				[
        'image_embeds',
        'negative_image_embeds',
        'image',
    ]
    lowercase__						 =				[
        'generator',
        'height',
        'width',
        'strength',
        'guidance_scale',
        'num_inference_steps',
        'return_dict',
        'guidance_scale',
        'num_images_per_prompt',
        'output_type',
        'return_dict',
    ]
    lowercase__						 =				False




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




     '''simple docstring'''

     return 32




    @property
    def 							__magic_name__							(   self						:	Dict				)				->							Any:




     '''simple docstring'''

     return 32




    @property
    def 							__magic_name__							(   self						:	str				)				->							Dict:




     '''simple docstring'''

     return self.time_input_dim




    @property
    def 							__magic_name__							(   self						:	str				)				->							int:




     '''simple docstring'''

     return self.time_input_dim * 4




    @property
    def 							__magic_name__							(   self						:	int				)				->							List[Any]:




     '''simple docstring'''

     return 100




    @property
    def 							__magic_name__							(   self						:	List[str]				)				->							Tuple:




     '''simple docstring'''

     torch.manual_seed(0				)

     A__	=    {
         "in_channels": 4,
         # Out channels is double in channels because predicts mean and variance
         "out_channels": 8,
         "addition_embed_type": "image",
         "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
         "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
         "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
         "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
         "layers_per_block": 1,
         "encoder_hid_dim": self.text_embedder_hidden_size,
         "encoder_hid_dim_type": "image_proj",
         "cross_attention_dim": self.cross_attention_dim,
         "attention_head_dim": 4,
         "resnet_time_scale_shift": "scale_shift",
         "class_embed_type": None,
     }

     A__	=    UNetaDConditionModel(**snake_case_				)
     return model




    @property
    def 							__magic_name__							(   self						:	Dict				)				->							Union[str, Any]:




     '''simple docstring'''

     return {
         "block_out_channels": [32, 64],
         "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
         "in_channels": 3,
         "latent_channels": 4,
         "layers_per_block": 1,
         "norm_num_groups": 8,
         "norm_type": "spatial",
         "num_vq_embeddings": 12,
         "out_channels": 3,
         "up_block_types": [
             "AttnUpDecoderBlock2D",
             "UpDecoderBlock2D",
         ],
         "vq_embed_dim": 4,
     }




    @property
    def 							__magic_name__							(   self						:	List[str]				)				->							List[Any]:




     '''simple docstring'''

     torch.manual_seed(0				)
     A__	=    VQModel(**self.dummy_movq_kwargs				)
     return model




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




     '''simple docstring'''

     A__	=    self.dummy_unet
     A__	=    self.dummy_movq

     A__	=    {
         "num_train_timesteps": 1_000,
         "beta_schedule": "linear",
         "beta_start": 0.0_0085,
         "beta_end": 0.012,
         "clip_sample": False,
         "set_alpha_to_one": False,
         "steps_offset": 0,
         "prediction_type": "epsilon",
         "thresholding": False,
     }

     A__	=    DDIMScheduler(**snake_case_				)

     A__	=    {
         "unet": unet,
         "scheduler": scheduler,
         "movq": movq,
     }

     return components




    def 							__magic_name__							(   self						:	Tuple     ,			snake_case_						:	List[Any]     ,			snake_case_						:	List[Any]=0				)				->							Any:




     '''simple docstring'''

     A__	=    floats_tensor((1, self.text_embedder_hidden_size)     ,			rng=random.Random(snake_case_				)				).to(snake_case_				)
     A__	=    floats_tensor((1, self.text_embedder_hidden_size)     ,			rng=random.Random(seed + 1				)				).to(
         snake_case_				)
     # create init_image
     A__	=    floats_tensor((1, 3, 64, 64)     ,			rng=random.Random(snake_case_				)				).to(snake_case_				)
     A__	=    image.cpu().permute(0     ,			2     ,			3     ,			1				)[0]
     A__	=    Image.fromarray(np.uinta(snake_case_				)				).convert("RGB"				).resize((256, 256)				)

     if str(snake_case_				).startswith("mps"				):
      A__	=    torch.manual_seed(snake_case_				)
     else:
      A__	=    torch.Generator(device=snake_case_				).manual_seed(snake_case_				)
     A__	=    {
         "image": init_image,
         "image_embeds": image_embeds,
         "negative_image_embeds": negative_image_embeds,
         "generator": generator,
         "height": 64,
         "width": 64,
         "num_inference_steps": 10,
         "guidance_scale": 7.0,
         "strength": 0.2,
         "output_type": "np",
     }
     return inputs




    def 							__magic_name__							(   self						:	Optional[Any]				)				->							List[str]:




     '''simple docstring'''

     A__	=    "cpu"

     A__	=    self.get_dummy_components()

     A__	=    self.pipeline_class(**snake_case_				)
     A__	=    pipe.to(snake_case_				)

     pipe.set_progress_bar_config(disable=snake_case_				)

     A__	=    pipe(**self.get_dummy_inputs(snake_case_				)				)
     A__	=    output.images

     A__	=    pipe(
         **self.get_dummy_inputs(snake_case_				)     ,			return_dict=snake_case_     ,			)[0]

     A__	=    image[0, -3:, -3:, -1]
     A__	=    image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)

     A__	=    np.array(
         [0.619_9778, 0.6398_4406, 0.4614_5785, 0.6294_4984, 0.562_2215, 0.4730_6132, 0.4744_1456, 0.460_7606, 0.4871_9263]				)
     assert (
         np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2
     ), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
     assert (
         np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1e-2
     ), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""






@slow
@require_torch_gpu
class      UpperCAmelCase_  (						unittest.TestCase							):




    def 							__magic_name__							(   self						:	int				)				->							Tuple:




     '''simple docstring'''

     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()




    def 							__magic_name__							(   self						:	Optional[int]				)				->							Any:




     '''simple docstring'''

     A__	=    load_numpy(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
         "/kandinskyv22/kandinskyv22_img2img_frog.npy"				)

     A__	=    load_image(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"				)
     A__	=    "A red cartoon frog, 4k"

     A__	=    KandinskyVaaPriorPipeline.from_pretrained(
         "kandinsky-community/kandinsky-2-2-prior"     ,			torch_dtype=torch.floataa				)
     pipe_prior.to(snake_case_				)

     A__	=    KandinskyVaaImgaImgPipeline.from_pretrained(
         "kandinsky-community/kandinsky-2-2-decoder"     ,			torch_dtype=torch.floataa				)
     A__	=    pipeline.to(snake_case_				)

     pipeline.set_progress_bar_config(disable=snake_case_				)

     A__	=    torch.Generator(device="cpu"				).manual_seed(0				)
     A__,				A__	=    pipe_prior(
         snake_case_     ,			generator=snake_case_     ,			num_inference_steps=5     ,			negative_prompt=""     ,			).to_tuple()

     A__	=    pipeline(
         image=snake_case_     ,			image_embeds=snake_case_     ,			negative_image_embeds=snake_case_     ,			generator=snake_case_     ,			num_inference_steps=100     ,			height=768     ,			width=768     ,			strength=0.2     ,			output_type="np"     ,			)

     A__	=    output.images[0]

     assert image.shape == (768, 768, 3)

     assert_mean_pixel_difference(snake_case_     ,			snake_case_				)











"""simple docstring"""



import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import is_speech_available, is_vision_available
from transformers.testing_utils import require_torch


if is_vision_available():
      from transformers import TvltImageProcessor

if is_speech_available():
      from transformers import TvltFeatureExtractor

from transformers import TvltProcessor




@require_torch
class      UpperCAmelCase_  (						unittest.TestCase							):




    def 							__magic_name__							(   self						:	str				)				->							Dict:




     '''simple docstring'''

     A__	=    "ZinengTang/tvlt-base"
     A__	=    tempfile.mkdtemp()




    def 							__magic_name__							(   self						:	int     ,			**snake_case_						:	Dict				)				->							Union[str, Any]:




     '''simple docstring'''

     return TvltImageProcessor.from_pretrained(self.checkpoint     ,			**snake_case_				)




    def 							__magic_name__							(   self						:	Optional[int]     ,			**snake_case_						:	str				)				->							List[str]:




     '''simple docstring'''

     return TvltFeatureExtractor.from_pretrained(self.checkpoint     ,			**snake_case_				)




    def 							__magic_name__							(   self						:	Optional[int]				)				->							str:




     '''simple docstring'''

     shutil.rmtree(self.tmpdirname				)




    def 							__magic_name__							(   self						:	Tuple				)				->							Any:




     '''simple docstring'''

     A__	=    self.get_image_processor()
     A__	=    self.get_feature_extractor()

     A__	=    TvltProcessor(image_processor=snake_case_     ,			feature_extractor=snake_case_				)
     processor.save_pretrained(self.tmpdirname				)
     A__	=    TvltProcessor.from_pretrained(self.tmpdirname				)

     self.assertIsInstance(processor.feature_extractor     ,			snake_case_				)
     self.assertIsInstance(processor.image_processor     ,			snake_case_				)




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




     '''simple docstring'''

     A__	=    self.get_image_processor()
     A__	=    self.get_feature_extractor()

     A__	=    TvltProcessor(image_processor=snake_case_     ,			feature_extractor=snake_case_				)

     A__	=    np.ones([12_000]				)

     A__	=    feature_extractor(snake_case_     ,			return_tensors="np"				)
     A__	=    processor(audio=snake_case_     ,			return_tensors="np"				)

     for key in audio_dict.keys():
      self.assertAlmostEqual(audio_dict[key].sum()     ,			input_processor[key].sum()     ,			delta=1e-2				)




    def 							__magic_name__							(   self						:	Tuple				)				->							Tuple:




     '''simple docstring'''

     A__	=    self.get_image_processor()
     A__	=    self.get_feature_extractor()

     A__	=    TvltProcessor(image_processor=snake_case_     ,			feature_extractor=snake_case_				)

     A__	=    np.ones([3, 224, 224]				)

     A__	=    image_processor(snake_case_     ,			return_tensors="np"				)
     A__	=    processor(images=snake_case_     ,			return_tensors="np"				)

     for key in image_dict.keys():
      self.assertAlmostEqual(image_dict[key].sum()     ,			input_processor[key].sum()     ,			delta=1e-2				)




    def 							__magic_name__							(   self						:	Tuple				)				->							Optional[int]:




     '''simple docstring'''

     A__	=    self.get_image_processor()
     A__	=    self.get_feature_extractor()

     A__	=    TvltProcessor(image_processor=snake_case_     ,			feature_extractor=snake_case_				)

     A__	=    np.ones([12_000]				)
     A__	=    np.ones([3, 224, 224]				)

     A__	=    processor(audio=snake_case_     ,			images=snake_case_				)

     self.assertListEqual(list(inputs.keys()				)     ,			["audio_values", "audio_mask", "pixel_values", "pixel_mask"]				)

     # test if it raises when no input is passed
     with pytest.raises(snake_case_				):
      processor()




    def 							__magic_name__							(   self						:	Optional[Any]				)				->							Tuple:




     '''simple docstring'''

     A__	=    self.get_image_processor()
     A__	=    self.get_feature_extractor()

     A__	=    TvltProcessor(image_processor=snake_case_     ,			feature_extractor=snake_case_				)

     self.assertListEqual(
         processor.model_input_names     ,			image_processor.model_input_names + feature_extractor.model_input_names     ,			msg="`processor` and `image_processor`+`feature_extractor` model input names do not match"     ,			)





