

import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
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 ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


def        __lowerCAmelCase    (       _UpperCamelCase		:  Dict	)		->							int:
 '''simple docstring'''


 SCREAMING_SNAKE_CASE								=			[tensor.shape for tensor in tensor_list]
 return all(shape == shapes[0] for shape in shapes[1:]	)




class      UpperCamelCase (     SCREAMING_SNAKE_CASE	,			SCREAMING_SNAKE_CASE	,			SCREAMING_SNAKE_CASE	,			unittest.TestCase   ):
    __UpperCamelCase         =StableDiffusionLatentUpscalePipeline
    __UpperCamelCase         =TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
        "height",
        "width",
        "cross_attention_kwargs",
        "negative_prompt_embeds",
        "prompt_embeds",
    }
    __UpperCamelCase         =PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
    __UpperCamelCase         =TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    __UpperCamelCase         =frozenset(
        []   )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    __UpperCamelCase         =frozenset([]   )

    __UpperCamelCase         =True

    @property
    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			1
     SCREAMING_SNAKE_CASE								=			4
     SCREAMING_SNAKE_CASE								=			(1_6, 1_6)

     SCREAMING_SNAKE_CASE								=			floats_tensor((batch_size, num_channels) + sizes							,    rng=random.Random(0    )    ).to(snake_case__    )
     return image

    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     torch.manual_seed(0    )
     SCREAMING_SNAKE_CASE								=			UNetaDConditionModel(
         act_fn='gelu'							,    attention_head_dim=8							,    norm_num_groups=snake_case__							,    block_out_channels=[3_2, 3_2, 6_4, 6_4]							,    time_cond_proj_dim=1_6_0							,    conv_in_kernel=1							,    conv_out_kernel=1							,    cross_attention_dim=3_2							,    down_block_types=(
             'KDownBlock2D',
             'KCrossAttnDownBlock2D',
             'KCrossAttnDownBlock2D',
             'KCrossAttnDownBlock2D',
         )							,    in_channels=8							,    mid_block_type=snake_case__							,    only_cross_attention=snake_case__							,    out_channels=5							,    resnet_time_scale_shift='scale_shift'							,    time_embedding_type='fourier'							,    timestep_post_act='gelu'							,    up_block_types=('KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KUpBlock2D')							,    )
     SCREAMING_SNAKE_CASE								=			AutoencoderKL(
         block_out_channels=[3_2, 3_2, 6_4, 6_4]							,    in_channels=3							,    out_channels=3							,    down_block_types=[
             'DownEncoderBlock2D',
             'DownEncoderBlock2D',
             'DownEncoderBlock2D',
             'DownEncoderBlock2D',
         ]							,    up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D']							,    latent_channels=4							,    )
     SCREAMING_SNAKE_CASE								=			EulerDiscreteScheduler(prediction_type='sample'    )
     SCREAMING_SNAKE_CASE								=			CLIPTextConfig(
         bos_token_id=0							,    eos_token_id=2							,    hidden_size=3_2							,    intermediate_size=3_7							,    layer_norm_eps=1E-05							,    num_attention_heads=4							,    num_hidden_layers=5							,    pad_token_id=1							,    vocab_size=1_0_0_0							,    hidden_act='quick_gelu'							,    projection_dim=5_1_2							,    )
     SCREAMING_SNAKE_CASE								=			CLIPTextModel(snake_case__    )
     SCREAMING_SNAKE_CASE								=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

     SCREAMING_SNAKE_CASE								=			{
         'unet': model.eval(),
         'vae': vae.eval(),
         'scheduler': scheduler,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
     }

     return components

    def UpperCamelCase			(			self :      List[str]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Union[str, Any]=0    ):

     """simple docstring"""






     if str(snake_case__    ).startswith('mps'    ):
      SCREAMING_SNAKE_CASE								=			torch.manual_seed(snake_case__    )
     else:
      SCREAMING_SNAKE_CASE								=			torch.Generator(device=snake_case__    ).manual_seed(snake_case__    )
     SCREAMING_SNAKE_CASE								=			{
         'prompt': 'A painting of a squirrel eating a burger',
         'image': self.dummy_image.cpu(),
         'generator': generator,
         'num_inference_steps': 2,
         'output_type': 'numpy',
     }
     return inputs

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'cpu'

     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			self.pipeline_class(**snake_case__    )
     pipe.to(snake_case__    )
     pipe.set_progress_bar_config(disable=snake_case__    )

     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			pipe(**snake_case__    ).images
     SCREAMING_SNAKE_CASE								=			image[0, -3:, -3:, -1]

     self.assertEqual(image.shape							,    (1, 2_5_6, 2_5_6, 3)    )
     SCREAMING_SNAKE_CASE								=			np.array(
         [0.47_222_412, 0.41_921_633, 0.44_717_434, 0.46_874_192, 0.42_588_258, 0.46_150_726, 0.4_677_534, 0.45_583_832, 0.48_579_055]    )
     SCREAMING_SNAKE_CASE								=			np.abs(image_slice.flatten() - expected_slice    ).max()
     self.assertLessEqual(snake_case__							,    1E-3    )

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     super().test_attention_slicing_forward_pass(expected_max_diff=7E-3    )

    def UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     super().test_cpu_offload_forward_pass(expected_max_diff=3E-3    )

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3    )

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     super().test_inference_batch_single_identical(expected_max_diff=7E-3    )

    def UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3E-3    )

    def UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     super().test_save_load_local(expected_max_difference=3E-3    )

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

     """simple docstring"""






     super().test_save_load_optional_components(expected_max_difference=3E-3    )




    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			[
         'DDIMScheduler',
         'DDPMScheduler',
         'PNDMScheduler',
         'HeunDiscreteScheduler',
         'EulerAncestralDiscreteScheduler',
         'KDPM2DiscreteScheduler',
         'KDPM2AncestralDiscreteScheduler',
         'DPMSolverSDEScheduler',
     ]
     SCREAMING_SNAKE_CASE								=			self.get_dummy_components()
     SCREAMING_SNAKE_CASE								=			self.pipeline_class(**snake_case__    )

     # make sure that PNDM does not need warm-up
     pipe.scheduler.register_to_config(skip_prk_steps=snake_case__    )

     pipe.to(snake_case__    )
     pipe.set_progress_bar_config(disable=snake_case__    )
     SCREAMING_SNAKE_CASE								=			self.get_dummy_inputs(snake_case__    )
     SCREAMING_SNAKE_CASE								=			2

     SCREAMING_SNAKE_CASE								=			[]
     for scheduler_enum in KarrasDiffusionSchedulers:
      if scheduler_enum.name in skip_schedulers:
       # no sigma schedulers are not supported
       # no schedulers
       continue

      SCREAMING_SNAKE_CASE								=			getattr(snake_case__							,    scheduler_enum.name    )
      SCREAMING_SNAKE_CASE								=			scheduler_cls.from_config(pipe.scheduler.config    )
      SCREAMING_SNAKE_CASE								=			pipe(**snake_case__    )[0]
      outputs.append(snake_case__    )

     assert check_same_shape(snake_case__    )




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

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






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

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			torch.manual_seed(3_3    )

     SCREAMING_SNAKE_CASE								=			StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'							,    torch_dtype=torch.floataa    )
     pipe.to('cuda'    )

     SCREAMING_SNAKE_CASE								=			StableDiffusionLatentUpscalePipeline.from_pretrained(
         'stabilityai/sd-x2-latent-upscaler'							,    torch_dtype=torch.floataa    )
     upscaler.to('cuda'    )

     SCREAMING_SNAKE_CASE								=			'a photo of an astronaut high resolution, unreal engine, ultra realistic'

     SCREAMING_SNAKE_CASE								=			pipe(snake_case__							,    generator=snake_case__							,    output_type='latent'    ).images

     SCREAMING_SNAKE_CASE								=			upscaler(
         prompt=snake_case__							,    image=snake_case__							,    num_inference_steps=2_0							,    guidance_scale=0							,    generator=snake_case__							,    output_type='np'							,    ).images[0]

     SCREAMING_SNAKE_CASE								=			load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy'    )
     assert np.abs((expected_image - image).mean()    ) < 5E-2




    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			torch.manual_seed(3_3    )

     SCREAMING_SNAKE_CASE								=			StableDiffusionLatentUpscalePipeline.from_pretrained(
         'stabilityai/sd-x2-latent-upscaler'							,    torch_dtype=torch.floataa    )
     upscaler.to('cuda'    )

     SCREAMING_SNAKE_CASE								=			'the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'

     SCREAMING_SNAKE_CASE								=			load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png'    )

     SCREAMING_SNAKE_CASE								=			upscaler(
         prompt=snake_case__							,    image=snake_case__							,    num_inference_steps=2_0							,    guidance_scale=0							,    generator=snake_case__							,    output_type='np'							,    ).images[0]

     SCREAMING_SNAKE_CASE								=			load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy'    )
     assert np.abs((expected_image - image).max()    ) < 5E-2




# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ..models.whisper import WhisperForConditionalGeneration, WhisperProcessor
from .base import PipelineTool




class      UpperCamelCase (     SCREAMING_SNAKE_CASE   ):
    __UpperCamelCase         ="openai/whisper-base"
    __UpperCamelCase         =(
        "This is a tool that transcribes an audio into text. It takes an input named `audio` and returns the "
        "transcribed text."
    )
    __UpperCamelCase         ="transcriber"
    __UpperCamelCase         =WhisperProcessor
    __UpperCamelCase         =WhisperForConditionalGeneration

    __UpperCamelCase         =["audio"]
    __UpperCamelCase         =["text"]

    def UpperCamelCase			(			self :      Dict							,    snake_case__ :      Tuple    ):

     """simple docstring"""






     return self.pre_processor(snake_case__							,    return_tensors='pt'    ).input_features

    def UpperCamelCase			(			self :      Optional[int]							,    snake_case__ :      Tuple    ):

     """simple docstring"""






     return self.model.generate(inputs=snake_case__    )




    def UpperCamelCase			(			self :      str							,    snake_case__ :      Union[str, Any]    ):

     """simple docstring"""






     return self.pre_processor.batch_decode(snake_case__							,    skip_special_tokens=snake_case__    )[0]

