


import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin







@skip_mps
class 						_snake_case    ( _snake_case		,       _snake_case		,       unittest.TestCase				):
 SCREAMING_SNAKE_CASE__ =			IFPipeline
 SCREAMING_SNAKE_CASE__ =			TEXT_TO_IMAGE_PARAMS - {'width', 'height', 'latents'}
 SCREAMING_SNAKE_CASE__ =			TEXT_TO_IMAGE_BATCH_PARAMS
 SCREAMING_SNAKE_CASE__ =			PipelineTesterMixin.required_optional_params - {'latents'}




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       return self._get_dummy_components()




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase=0       ):
       if str(_lowerCamelCase       ).startswith('''mps'''       ):
             a		:str									=			torch.manual_seed(_lowerCamelCase       )
       else:
             a		:Dict									=			torch.Generator(device=_lowerCamelCase       ).manual_seed(_lowerCamelCase       )

       a		:str									=			{
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''output_type''': '''numpy''',
       }

       return inputs




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       self._test_save_load_optional_components()




 @unittest.skipIf(torch_device != '''cuda'''   ,    reason='''float16 requires CUDA'''       )
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       # Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
       super().test_save_load_floataa(expected_max_diff=1e-1       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       self._test_attention_slicing_forward_pass(expected_max_diff=1e-2       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       self._test_save_load_local()




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       self._test_inference_batch_single_identical(
           expected_max_diff=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       ):
       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3       )







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




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       # if

       a		:Optional[int]									=			IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''   ,    variant='''fp16'''   ,    torch_dtype=torch.floataa       )

       a		:List[str]									=			IFSuperResolutionPipeline.from_pretrained(
           '''DeepFloyd/IF-II-L-v1.0'''   ,    variant='''fp16'''   ,    torch_dtype=torch.floataa   ,    text_encoder=_lowerCamelCase   ,    tokenizer=_lowerCamelCase       )

       # pre compute text embeddings and remove T5 to save memory

       pipe_a.text_encoder.to('''cuda'''       )

       a				,							a		:Tuple									=			pipe_a.encode_prompt('''anime turtle'''   ,    device='''cuda'''       )

       del pipe_a.tokenizer
       del pipe_a.text_encoder
       gc.collect()

       a		:List[str]									=			None
       a		:Tuple									=			None

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()       )
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()       )

       self._test_if(_lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       )

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # img2img

       a		:Optional[Any]									=			IFImgaImgPipeline(**pipe_a.components       )
       a		:Any									=			IFImgaImgSuperResolutionPipeline(**pipe_a.components       )

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()       )
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()       )

       self._test_if_imgaimg(_lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       )

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # inpainting

       a		:List[str]									=			IFInpaintingPipeline(**pipe_a.components       )
       a		:str									=			IFInpaintingSuperResolutionPipeline(**pipe_a.components       )

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()       )
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()       )

       self._test_if_inpainting(_lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       # pipeline 1

       _start_torch_memory_measurement()

       a		:int									=			torch.Generator(device='''cpu'''       ).manual_seed(0       )
       a		:str									=			pipe_a(
           prompt_embeds=_lowerCamelCase   ,    negative_prompt_embeds=_lowerCamelCase   ,    num_inference_steps=2   ,    generator=_lowerCamelCase   ,    output_type='''np'''   ,    )

       a		:Optional[int]									=			output.images[0]

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

       a		:int									=			torch.cuda.max_memory_allocated()
       assert mem_bytes < 13 * 10**9

       a		:Any									=			load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'''       )
       assert_mean_pixel_difference(_lowerCamelCase   ,    _lowerCamelCase       )

       # pipeline 2

       _start_torch_memory_measurement()

       a		:str									=			torch.Generator(device='''cpu'''       ).manual_seed(0       )

       a		:Dict									=			floats_tensor((1, 3, 64, 64)   ,    rng=random.Random(0       )       ).to(_lowerCamelCase       )

       a		:List[Any]									=			pipe_a(
           prompt_embeds=_lowerCamelCase   ,    negative_prompt_embeds=_lowerCamelCase   ,    image=_lowerCamelCase   ,    generator=_lowerCamelCase   ,    num_inference_steps=2   ,    output_type='''np'''   ,    )

       a		:Optional[int]									=			output.images[0]

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

       a		:Optional[Any]									=			torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       a		:Tuple									=			load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'''       )
       assert_mean_pixel_difference(_lowerCamelCase   ,    _lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       # pipeline 1

       _start_torch_memory_measurement()

       a		:int									=			floats_tensor((1, 3, 64, 64)   ,    rng=random.Random(0       )       ).to(_lowerCamelCase       )

       a		:List[str]									=			torch.Generator(device='''cpu'''       ).manual_seed(0       )

       a		:str									=			pipe_a(
           prompt_embeds=_lowerCamelCase   ,    negative_prompt_embeds=_lowerCamelCase   ,    image=_lowerCamelCase   ,    num_inference_steps=2   ,    generator=_lowerCamelCase   ,    output_type='''np'''   ,    )

       a		:Optional[int]									=			output.images[0]

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

       a		:Tuple									=			torch.cuda.max_memory_allocated()
       assert mem_bytes < 10 * 10**9

       a		:str									=			load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'''       )
       assert_mean_pixel_difference(_lowerCamelCase   ,    _lowerCamelCase       )

       # pipeline 2

       _start_torch_memory_measurement()

       a		:Any									=			torch.Generator(device='''cpu'''       ).manual_seed(0       )

       a		:Tuple									=			floats_tensor((1, 3, 256, 256)   ,    rng=random.Random(0       )       ).to(_lowerCamelCase       )
       a		:str									=			floats_tensor((1, 3, 64, 64)   ,    rng=random.Random(0       )       ).to(_lowerCamelCase       )

       a		:Any									=			pipe_a(
           prompt_embeds=_lowerCamelCase   ,    negative_prompt_embeds=_lowerCamelCase   ,    image=_lowerCamelCase   ,    original_image=_lowerCamelCase   ,    generator=_lowerCamelCase   ,    num_inference_steps=2   ,    output_type='''np'''   ,    )

       a		:List[str]									=			output.images[0]

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

       a		:Optional[Any]									=			torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       a		:Any									=			load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'''       )
       assert_mean_pixel_difference(_lowerCamelCase   ,    _lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       # pipeline 1

       _start_torch_memory_measurement()

       a		:Any									=			floats_tensor((1, 3, 64, 64)   ,    rng=random.Random(0       )       ).to(_lowerCamelCase       )
       a		:Optional[Any]									=			floats_tensor((1, 3, 64, 64)   ,    rng=random.Random(1       )       ).to(_lowerCamelCase       )

       a		:str									=			torch.Generator(device='''cpu'''       ).manual_seed(0       )
       a		:Optional[int]									=			pipe_a(
           prompt_embeds=_lowerCamelCase   ,    negative_prompt_embeds=_lowerCamelCase   ,    image=_lowerCamelCase   ,    mask_image=_lowerCamelCase   ,    num_inference_steps=2   ,    generator=_lowerCamelCase   ,    output_type='''np'''   ,    )

       a		:str									=			output.images[0]

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

       a		:Any									=			torch.cuda.max_memory_allocated()
       assert mem_bytes < 10 * 10**9

       a		:Tuple									=			load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'''       )
       assert_mean_pixel_difference(_lowerCamelCase   ,    _lowerCamelCase       )

       # pipeline 2

       _start_torch_memory_measurement()

       a		:Any									=			torch.Generator(device='''cpu'''       ).manual_seed(0       )

       a		:Optional[Any]									=			floats_tensor((1, 3, 64, 64)   ,    rng=random.Random(0       )       ).to(_lowerCamelCase       )
       a		:str									=			floats_tensor((1, 3, 256, 256)   ,    rng=random.Random(0       )       ).to(_lowerCamelCase       )
       a		:Dict									=			floats_tensor((1, 3, 256, 256)   ,    rng=random.Random(1       )       ).to(_lowerCamelCase       )

       a		:str									=			pipe_a(
           prompt_embeds=_lowerCamelCase   ,    negative_prompt_embeds=_lowerCamelCase   ,    image=_lowerCamelCase   ,    mask_image=_lowerCamelCase   ,    original_image=_lowerCamelCase   ,    generator=_lowerCamelCase   ,    num_inference_steps=2   ,    output_type='''np'''   ,    )

       a		:Optional[int]									=			output.images[0]

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

       a		:List[Any]									=			torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       a		:Any									=			load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'''       )
       assert_mean_pixel_difference(_lowerCamelCase   ,    _lowerCamelCase       )







def 							__lowerCamelCase		(					):




      """simple docstring"""





      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()









import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case							:       Optional[int]										=	logging.get_logger(__name__)

snake_case							:       Optional[Any]										=	{
    '''microsoft/wavlm-base''': '''https://huggingface.co/microsoft/wavlm-base/resolve/main/config.json''',
    # See all WavLM models at https://huggingface.co/models?filter=wavlm
}







class 						_snake_case    ( _snake_case				):
 SCREAMING_SNAKE_CASE__ =			'wavlm'




 def __init__(						self   ,    _lowerCamelCase=32   ,    _lowerCamelCase=768   ,    _lowerCamelCase=12   ,    _lowerCamelCase=12   ,    _lowerCamelCase=3072   ,    _lowerCamelCase="gelu"   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=0.0   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=0.02   ,    _lowerCamelCase=1e-5   ,    _lowerCamelCase="group"   ,    _lowerCamelCase="gelu"   ,    _lowerCamelCase=(512, 512, 512, 512, 512, 512, 512)   ,    _lowerCamelCase=(5, 2, 2, 2, 2, 2, 2)   ,    _lowerCamelCase=(10, 3, 3, 3, 3, 2, 2)   ,    _lowerCamelCase=False   ,    _lowerCamelCase=128   ,    _lowerCamelCase=16   ,    _lowerCamelCase=320   ,    _lowerCamelCase=800   ,    _lowerCamelCase=False   ,    _lowerCamelCase=True   ,    _lowerCamelCase=0.05   ,    _lowerCamelCase=10   ,    _lowerCamelCase=2   ,    _lowerCamelCase=0.0   ,    _lowerCamelCase=10   ,    _lowerCamelCase=320   ,    _lowerCamelCase=2   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=100   ,    _lowerCamelCase=256   ,    _lowerCamelCase=256   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase="mean"   ,    _lowerCamelCase=False   ,    _lowerCamelCase=False   ,    _lowerCamelCase=256   ,    _lowerCamelCase=(512, 512, 512, 512, 1500)   ,    _lowerCamelCase=(5, 3, 3, 1, 1)   ,    _lowerCamelCase=(1, 2, 3, 1, 1)   ,    _lowerCamelCase=512   ,    _lowerCamelCase=80   ,    _lowerCamelCase=0   ,    _lowerCamelCase=1   ,    _lowerCamelCase=2   ,    _lowerCamelCase=False   ,    _lowerCamelCase=3   ,    _lowerCamelCase=2   ,    _lowerCamelCase=3   ,    _lowerCamelCase=None   ,    **_lowerCamelCase   ,    ):
       super().__init__(**_lowerCamelCase   ,    pad_token_id=_lowerCamelCase   ,    bos_token_id=_lowerCamelCase   ,    eos_token_id=_lowerCamelCase       )
       a		:Union[str, Any]									=			hidden_size
       a		:List[Any]									=			feat_extract_norm
       a		:Tuple									=			feat_extract_activation
       a		:Tuple									=			list(_lowerCamelCase       )
       a		:Optional[Any]									=			list(_lowerCamelCase       )
       a		:Optional[Any]									=			list(_lowerCamelCase       )
       a		:int									=			conv_bias
       a		:Tuple									=			num_buckets
       a		:Tuple									=			max_bucket_distance
       a		:int									=			num_conv_pos_embeddings
       a		:List[str]									=			num_conv_pos_embedding_groups
       a		:Tuple									=			len(self.conv_dim       )
       a		:Any									=			num_hidden_layers
       a		:Tuple									=			intermediate_size
       a		:Any									=			hidden_act
       a		:List[str]									=			num_attention_heads
       a		:int									=			hidden_dropout
       a		:Optional[Any]									=			attention_dropout
       a		:Dict									=			activation_dropout
       a		:Any									=			feat_proj_dropout
       a		:Dict									=			final_dropout
       a		:int									=			layerdrop
       a		:Tuple									=			layer_norm_eps
       a		:Union[str, Any]									=			initializer_range
       a		:Union[str, Any]									=			num_ctc_classes
       a		:Optional[Any]									=			vocab_size
       a		:str									=			do_stable_layer_norm
       a		:Dict									=			use_weighted_layer_sum
       a		:Any									=			classifier_proj_size

       if (
           (len(self.conv_stride       ) != self.num_feat_extract_layers)
           or (len(self.conv_kernel       ) != self.num_feat_extract_layers)
           or (len(self.conv_dim       ) != self.num_feat_extract_layers)
       ):
             raise ValueError(
                 '''Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =='''
                 ''' `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ='''
                 F''' {len(self.conv_dim       )}`, `len(config.conv_stride) = {len(self.conv_stride       )}`,'''
                 F''' `len(config.conv_kernel) = {len(self.conv_kernel       )}`.'''       )

       # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
       a		:List[Any]									=			apply_spec_augment
       a		:Optional[int]									=			mask_time_prob
       a		:Optional[Any]									=			mask_time_length
       a		:Tuple									=			mask_time_min_masks
       a		:Dict									=			mask_feature_prob
       a		:List[Any]									=			mask_feature_length

       # parameters for pretraining with codevector quantized representations
       a		:Tuple									=			num_codevectors_per_group
       a		:List[str]									=			num_codevector_groups
       a		:List[Any]									=			contrastive_logits_temperature
       a		:Union[str, Any]									=			num_negatives
       a		:Union[str, Any]									=			codevector_dim
       a		:Dict									=			proj_codevector_dim
       a		:Union[str, Any]									=			diversity_loss_weight

       # ctc loss
       a		:Dict									=			ctc_loss_reduction
       a		:Optional[int]									=			ctc_zero_infinity

       # adapter
       a		:Any									=			add_adapter
       a		:str									=			adapter_kernel_size
       a		:Tuple									=			adapter_stride
       a		:Dict									=			num_adapter_layers
       a		:Any									=			output_hidden_size or hidden_size

       # SequenceClassification-specific parameter. Feel free to ignore for other classes.
       a		:Tuple									=			classifier_proj_size

       # XVector-specific parameters. Feel free to ignore for other classes.
       a		:List[Any]									=			list(_lowerCamelCase       )
       a		:Optional[int]									=			list(_lowerCamelCase       )
       a		:List[str]									=			list(_lowerCamelCase       )
       a		:Optional[int]									=			xvector_output_dim




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       return functools.reduce(operator.mul   ,    self.conv_stride   ,    1       )





