
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()

class 		lowercase_							(							unittest.TestCase					):

 @property
 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_											=							1
        UpperCamelCase_											=							3
        UpperCamelCase_											=							(3_2, 3_2)

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

 @property
 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        torch.manual_seed(0  )
        UpperCamelCase_											=							UNetaDModel(
            block_out_channels=(3_2, 6_4)  ,						layers_per_block=2  ,						sample_size=3_2  ,						in_channels=6  ,						out_channels=3  ,						down_block_types=("""DownBlock2D""", """AttnDownBlock2D""")  ,						up_block_types=("""AttnUpBlock2D""", """UpBlock2D""")  ,						)
        return model

 @property
 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        torch.manual_seed(0  )
        UpperCamelCase_											=							VQModel(
            block_out_channels=[3_2, 6_4]  ,						in_channels=3  ,						out_channels=3  ,						down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]  ,						up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]  ,						latent_channels=3  ,						)
        return model

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_											=							"""cpu"""
        UpperCamelCase_											=							self.dummy_uncond_unet
        UpperCamelCase_											=							DDIMScheduler()
        UpperCamelCase_											=							self.dummy_vq_model

        UpperCamelCase_											=							LDMSuperResolutionPipeline(unet=lowercase_  ,						vqvae=lowercase_  ,						scheduler=lowercase_  )
        ldm.to(lowercase_  )
        ldm.set_progress_bar_config(disable=lowercase_  )

        UpperCamelCase_											=							self.dummy_image.to(lowercase_  )

        UpperCamelCase_											=							torch.Generator(device=lowercase_  ).manual_seed(0  )
        UpperCamelCase_											=							ldm(image=lowercase_  ,						generator=lowercase_  ,						num_inference_steps=2  ,						output_type="""numpy"""  ).images

        UpperCamelCase_											=							image[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)
        UpperCamelCase_											=							np.array([0.8_678, 0.8_245, 0.6_381, 0.6_830, 0.4_385, 0.5_599, 0.4_641, 0.6_201, 0.5_150]  )

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

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




        """simple docstring"""


        UpperCamelCase_											=							self.dummy_uncond_unet
        UpperCamelCase_											=							DDIMScheduler()
        UpperCamelCase_											=							self.dummy_vq_model

        # put models in fp16
        UpperCamelCase_											=							unet.half()
        UpperCamelCase_											=							vqvae.half()

        UpperCamelCase_											=							LDMSuperResolutionPipeline(unet=lowercase_  ,						vqvae=lowercase_  ,						scheduler=lowercase_  )
        ldm.to(lowercase_  )
        ldm.set_progress_bar_config(disable=lowercase_  )

        UpperCamelCase_											=							self.dummy_image.to(lowercase_  )

        UpperCamelCase_											=							ldm(lowercase_  ,						num_inference_steps=2  ,						output_type="""numpy"""  ).images

        assert image.shape == (1, 6_4, 6_4, 3)



@slow
@require_torch
class 		lowercase_							(							unittest.TestCase					):

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_											=							load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/vq_diffusion/teddy_bear_pool.png"""  )
        UpperCamelCase_											=							init_image.resize((6_4, 6_4)  ,						resample=PIL_INTERPOLATION["""lanczos"""]  )

        UpperCamelCase_											=							LDMSuperResolutionPipeline.from_pretrained("""duongna/ldm-super-resolution"""  ,						device_map="""auto"""  )
        ldm.set_progress_bar_config(disable=lowercase_  )

        UpperCamelCase_											=							torch.manual_seed(0  )
        UpperCamelCase_											=							ldm(image=lowercase_  ,						generator=lowercase_  ,						num_inference_steps=2_0  ,						output_type="""numpy"""  ).images

        UpperCamelCase_											=							image[0, -3:, -3:, -1]

        assert image.shape == (1, 2_5_6, 2_5_6, 3)
        UpperCamelCase_											=							np.array([0.7_644, 0.7_679, 0.7_642, 0.7_633, 0.7_666, 0.7_560, 0.7_425, 0.7_257, 0.6_907]  )

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








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


lowerCAmelCase					:				Optional[int]   	={
    'configuration_speech_to_text_2': ['SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP', 'Speech2Text2Config'],
    'processing_speech_to_text_2': ['Speech2Text2Processor'],
    'tokenization_speech_to_text_2': ['Speech2Text2Tokenizer'],
}


try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowerCAmelCase					:				int   	=[
		    'SPEECH_TO_TEXT_2_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'Speech2Text2ForCausalLM',
		    'Speech2Text2PreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_speech_to_text_a import SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextaConfig
		from .processing_speech_to_text_a import SpeechaTextaProcessor
		from .tokenization_speech_to_text_a import SpeechaTextaTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_speech_to_text_a import (
				    SPEECH_TO_TEXT_2_PRETRAINED_MODEL_ARCHIVE_LIST,
				    SpeechaTextaForCausalLM,
				    SpeechaTextaPreTrainedModel,
				)

else:
		import sys

		lowerCAmelCase					:				Union[str, Any]   	=_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)

