import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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


enable_full_determinism()






class 		a__						(					unittest.TestCase					):

					def 				__UpperCamelCase   (    self      :       List[Any]    ):







							"""simple docstring"""






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

					@property
					def 				__UpperCamelCase   (    self      :       int    ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		str				      =    1
							SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    3
							SCREAMING_SNAKE_CASE_  :		Union[str, Any]				      =    (32, 32)

							SCREAMING_SNAKE_CASE_  :		List[Any]				      =    floats_tensor((batch_size, num_channels) + sizes,rng=random.Random(0    )    ).to(_A    )
							return image

					@property
					def 				__UpperCamelCase   (    self      :       str    ):







							"""simple docstring"""






							torch.manual_seed(0    )
							SCREAMING_SNAKE_CASE_  :		Dict				      =    UNetaDConditionModel(
							    block_out_channels=(32, 64),layers_per_block=2,sample_size=32,in_channels=4,out_channels=4,down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),cross_attention_dim=32,)
							return model

					@property
					def 				__UpperCamelCase   (    self      :       Any    ):







							"""simple docstring"""






							torch.manual_seed(0    )
							SCREAMING_SNAKE_CASE_  :		List[Any]				      =    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,)
							return model

					@property
					def 				__UpperCamelCase   (    self      :       str    ):







							"""simple docstring"""






							torch.manual_seed(0    )
							SCREAMING_SNAKE_CASE_  :		int				      =    RobertaSeriesConfig(
							    hidden_size=32,project_dim=32,intermediate_size=37,layer_norm_eps=1E-05,num_attention_heads=4,num_hidden_layers=5,pad_token_id=1,vocab_size=5006,)
							return RobertaSeriesModelWithTransformation(_A    )

					@property
					def 				__UpperCamelCase   (    self      :       List[str]    ):







							"""simple docstring"""






							def extract(*_A      :       str,**_A      :       Dict    ):






									class 		a__						:

														def __init__(    self      :       List[str]    ):







																"""simple docstring"""






																SCREAMING_SNAKE_CASE_  :		Any				      =    torch.ones([0]    )






														def 				__UpperCamelCase   (    self      :       Dict,_A      :       List[str]    ):







																"""simple docstring"""






																self.pixel_values.to(_A    )
																return self







									return Out()

							return extract

					def 				__UpperCamelCase   (    self      :       Optional[Any]    ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    "cpu"  # ensure determinism for the device-dependent torch.Generator
							SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    self.dummy_cond_unet
							SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    PNDMScheduler(skip_prk_steps=_A    )
							SCREAMING_SNAKE_CASE_  :		int				      =    self.dummy_vae
							SCREAMING_SNAKE_CASE_  :		Dict				      =    self.dummy_text_encoder
							SCREAMING_SNAKE_CASE_  :		str				      =    XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"    )
							SCREAMING_SNAKE_CASE_  :		Dict				      =    77

							SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    self.dummy_image.to(_A    )
							SCREAMING_SNAKE_CASE_  :		Union[str, Any]				      =    init_image / 2 + 0.5

							# make sure here that pndm scheduler skips prk
							SCREAMING_SNAKE_CASE_  :		Any				      =    AltDiffusionImgaImgPipeline(
							    unet=_A,scheduler=_A,vae=_A,text_encoder=_A,tokenizer=_A,safety_checker=_A,feature_extractor=self.dummy_extractor,)
							SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor,do_normalize=_A    )
							SCREAMING_SNAKE_CASE_  :		str				      =    alt_pipe.to(_A    )
							alt_pipe.set_progress_bar_config(disable=_A    )

							SCREAMING_SNAKE_CASE_  :		Union[str, Any]				      =    "A painting of a squirrel eating a burger"
							SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    torch.Generator(device=_A    ).manual_seed(0    )
							SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    alt_pipe(
							    [prompt],generator=_A,guidance_scale=6.0,num_inference_steps=2,output_type="np",image=_A,)

							SCREAMING_SNAKE_CASE_  :		List[str]				      =    output.images

							SCREAMING_SNAKE_CASE_  :		int				      =    torch.Generator(device=_A    ).manual_seed(0    )
							SCREAMING_SNAKE_CASE_  :		str				      =    alt_pipe(
							    [prompt],generator=_A,guidance_scale=6.0,num_inference_steps=2,output_type="np",image=_A,return_dict=_A,)[0]

							SCREAMING_SNAKE_CASE_  :		List[str]				      =    image[0, -3:, -3:, -1]
							SCREAMING_SNAKE_CASE_  :		Dict				      =    image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)
							SCREAMING_SNAKE_CASE_  :		Dict				      =    np.array([0.4427, 0.3731, 0.4249, 0.4941, 0.4546, 0.4148, 0.4193, 0.4666, 0.4499]    )

							assert np.abs(image_slice.flatten() - expected_slice    ).max() < 5E-3
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 5E-3

					@unittest.skipIf(torch_device != "cuda","This test requires a GPU"    )
					def 				__UpperCamelCase   (    self      :       Dict    ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		Union[str, Any]				      =    self.dummy_cond_unet
							SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    PNDMScheduler(skip_prk_steps=_A    )
							SCREAMING_SNAKE_CASE_  :		List[Any]				      =    self.dummy_vae
							SCREAMING_SNAKE_CASE_  :		str				      =    self.dummy_text_encoder
							SCREAMING_SNAKE_CASE_  :		str				      =    XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"    )
							SCREAMING_SNAKE_CASE_  :		Tuple				      =    77

							SCREAMING_SNAKE_CASE_  :		Tuple				      =    self.dummy_image.to(_A    )

							# put models in fp16
							SCREAMING_SNAKE_CASE_  :		Dict				      =    unet.half()
							SCREAMING_SNAKE_CASE_  :		str				      =    vae.half()
							SCREAMING_SNAKE_CASE_  :		int				      =    bert.half()

							# make sure here that pndm scheduler skips prk
							SCREAMING_SNAKE_CASE_  :		int				      =    AltDiffusionImgaImgPipeline(
							    unet=_A,scheduler=_A,vae=_A,text_encoder=_A,tokenizer=_A,safety_checker=_A,feature_extractor=self.dummy_extractor,)
							SCREAMING_SNAKE_CASE_  :		List[str]				      =    VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor,do_normalize=_A    )
							SCREAMING_SNAKE_CASE_  :		Any				      =    alt_pipe.to(_A    )
							alt_pipe.set_progress_bar_config(disable=_A    )

							SCREAMING_SNAKE_CASE_  :		Any				      =    "A painting of a squirrel eating a burger"
							SCREAMING_SNAKE_CASE_  :		Tuple				      =    torch.manual_seed(0    )
							SCREAMING_SNAKE_CASE_  :		int				      =    alt_pipe(
							    [prompt],generator=_A,num_inference_steps=2,output_type="np",image=_A,).images

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






					@unittest.skipIf(torch_device != "cuda","This test requires a GPU"    )
					def 				__UpperCamelCase   (    self      :       Any    ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/img2img/sketch-mountains-input.jpg"    )
							# resize to resolution that is divisible by 8 but not 16 or 32
							SCREAMING_SNAKE_CASE_  :		int				      =    init_image.resize((760, 504)    )

							SCREAMING_SNAKE_CASE_  :		Tuple				      =    "BAAI/AltDiffusion"
							SCREAMING_SNAKE_CASE_  :		Dict				      =    AltDiffusionImgaImgPipeline.from_pretrained(
							    _A,safety_checker=_A,)
							pipe.to(_A    )
							pipe.set_progress_bar_config(disable=_A    )
							pipe.enable_attention_slicing()

							SCREAMING_SNAKE_CASE_  :		int				      =    "A fantasy landscape, trending on artstation"

							SCREAMING_SNAKE_CASE_  :		Optional[Any]				      =    torch.manual_seed(0    )
							SCREAMING_SNAKE_CASE_  :		str				      =    pipe(
							    prompt=_A,image=_A,strength=0.75,guidance_scale=7.5,generator=_A,output_type="np",)
							SCREAMING_SNAKE_CASE_  :		List[str]				      =    output.images[0]

							SCREAMING_SNAKE_CASE_  :		List[Any]				      =    image[255:258, 383:386, -1]

							assert image.shape == (504, 760, 3)
							SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    np.array([0.9358, 0.9397, 0.9599, 0.9901, 1.0000, 1.0000, 0.9882, 1.0000, 1.0000]    )

							assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2







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

					def 				__UpperCamelCase   (    self      :       Any    ):







							"""simple docstring"""






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






					def 				__UpperCamelCase   (    self      :       Union[str, Any]    ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		List[str]				      =    load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/img2img/sketch-mountains-input.jpg"    )
							SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    init_image.resize((768, 512)    )
							SCREAMING_SNAKE_CASE_  :		int				      =    load_numpy(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"    )

							SCREAMING_SNAKE_CASE_  :		List[Any]				      =    "BAAI/AltDiffusion"
							SCREAMING_SNAKE_CASE_  :		List[Any]				      =    AltDiffusionImgaImgPipeline.from_pretrained(
							    _A,safety_checker=_A,)
							pipe.to(_A    )
							pipe.set_progress_bar_config(disable=_A    )
							pipe.enable_attention_slicing()

							SCREAMING_SNAKE_CASE_  :		Tuple				      =    "A fantasy landscape, trending on artstation"

							SCREAMING_SNAKE_CASE_  :		Optional[int]				      =    torch.manual_seed(0    )
							SCREAMING_SNAKE_CASE_  :		Any				      =    pipe(
							    prompt=_A,image=_A,strength=0.75,guidance_scale=7.5,generator=_A,output_type="np",)
							SCREAMING_SNAKE_CASE_  :		List[Any]				      =    output.images[0]

							assert image.shape == (512, 768, 3)
							# img2img is flaky across GPUs even in fp32, so using MAE here
							assert np.abs(expected_image - image    ).max() < 1E-2

from ...processing_utils import ProcessorMixin


class      lowerCamelCase    (SCREAMING_SNAKE_CASE__    ):
  """simple docstring"""







  lowerCamelCase__   =  ['''image_processor''', '''feature_extractor''']
  lowerCamelCase__   =  '''TvltImageProcessor'''
  lowerCamelCase__   =  '''TvltFeatureExtractor'''







  def __init__(							self       :				List[str]     ,   __magic_name__       :				Any     ,   __magic_name__       :				Any      )     ->		int:
        super().__init__(image_processor=__magic_name__     ,   feature_extractor=__magic_name__      )

        SCREAMING_SNAKE_CASE_      =     image_processor
        SCREAMING_SNAKE_CASE_      =     feature_extractor







  def __call__(							self       :				List[str]     ,   __magic_name__       :				Optional[Any]=None     ,   __magic_name__       :				Union[str, Any]=None     ,   __magic_name__       :				int=None     ,   __magic_name__       :				str=None     ,   __magic_name__       :				Any=False     ,   __magic_name__       :				int=False     ,   *__magic_name__       :				int     ,   **__magic_name__       :				Any     ,   )     ->		List[Any]:

        if images is None and audio is None:
              raise ValueError("You need to specify either an `images` or `audio` input to process."      )

        SCREAMING_SNAKE_CASE_      =     None
        if images is not None:
              SCREAMING_SNAKE_CASE_      =     self.image_processor(__magic_name__     ,   mask_pixel=__magic_name__     ,   *__magic_name__     ,   **__magic_name__      )
        if images_mixed is not None:
              SCREAMING_SNAKE_CASE_      =     self.image_processor(__magic_name__     ,   is_mixed=__magic_name__     ,   *__magic_name__     ,   **__magic_name__      )
        if audio is not None:
              SCREAMING_SNAKE_CASE_      =     self.feature_extractor(
                  __magic_name__     ,   *__magic_name__     ,   sampling_rate=__magic_name__     ,   mask_audio=__magic_name__     ,   **__magic_name__      )

        SCREAMING_SNAKE_CASE_      =     {}
        if audio is not None:
              output_dict.update(__magic_name__      )
        if images is not None:
              output_dict.update(__magic_name__      )
        if images_mixed_dict is not None:
              output_dict.update(__magic_name__      )
        return output_dict







  @property
  def       __A   (							self       :				Optional[Any]      )     ->		int:
        SCREAMING_SNAKE_CASE_      =     self.image_processor.model_input_names
        SCREAMING_SNAKE_CASE_      =     self.feature_extractor.model_input_names
        return list(dict.fromkeys(image_processor_input_names + feature_extractor_input_names      )      )


