




import hashlib
import unittest
from typing import Dict

import numpy as np

from transformers import (
    MODEL_FOR_MASK_GENERATION_MAPPING,
    TF_MODEL_FOR_MASK_GENERATION_MAPPING,
    is_vision_available,
    pipeline,
)
from transformers.pipelines import MaskGenerationPipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)


if is_vision_available():
     from PIL import Image
else:




     class       __lowerCAmelCase				:






            @staticmethod
            def     SCREAMING_SNAKE_CASE			(      *_lowerCAmelCase:    Tuple				,    **_lowerCAmelCase:    List[Any] ):
              pass



def 					UpperCAmelCase__     (   lowerCamelCase      ):
  lowercase			:str							=  hashlib.mda(image.tobytes()      )
  return m.hexdigest()[:10]



def 					UpperCAmelCase__     (   lowerCamelCase      ):
  lowercase			:int							=  np.array(lowerCamelCase      )
  lowercase			:Any							=  npimg.shape
  return {"hash": hashimage(lowerCamelCase      ), "shape": shape}







@is_pipeline_test
@require_vision
@require_torch
class       __lowerCAmelCase				(							unittest.TestCase):
       _a           =  dict(
           (list(MODEL_FOR_MASK_GENERATION_MAPPING.items()) if MODEL_FOR_MASK_GENERATION_MAPPING else []))
       _a           =  dict(
           (list(TF_MODEL_FOR_MASK_GENERATION_MAPPING.items()) if TF_MODEL_FOR_MASK_GENERATION_MAPPING else []))






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any]				,    _lowerCAmelCase:    int				,    _lowerCAmelCase:    str				,    _lowerCAmelCase:    Tuple ):
         lowercase			:int							=  MaskGenerationPipeline(model=_lowerCAmelCase				,    image_processor=_lowerCAmelCase )
         return image_segmenter, [
             "./tests/fixtures/tests_samples/COCO/000000039769.png",
             "./tests/fixtures/tests_samples/COCO/000000039769.png",
         ]






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any]				,    _lowerCAmelCase:    List[str]				,    _lowerCAmelCase:    int ):
         pass






       @require_tf
       @unittest.skip("Image segmentation not implemented in TF" )
       def     SCREAMING_SNAKE_CASE			(      self:    Tuple ):
         pass






       @slow
       @require_torch
       def     SCREAMING_SNAKE_CASE			(      self:    Tuple ):
         lowercase			:Dict							=  pipeline("mask-generation"				,    model="facebook/sam-vit-huge" )

         lowercase			:Union[str, Any]							=  image_segmenter("http://images.cocodataset.org/val2017/000000039769.jpg"				,    points_per_batch=2_56 )

         # Shortening by hashing
         lowercase			:List[str]							=  []
         for i, o in enumerate(outputs["masks"] ):
           new_outupt += [{"mask": mask_to_test_readable(_lowerCAmelCase ), "scores": outputs["scores"][i]}]

         # fmt: off
         self.assertEqual(
             nested_simplify(_lowerCAmelCase				,    decimals=4 )				,    [
                 {"mask": {"hash": "115ad19f5f", "shape": (4_80, 6_40)}, "scores": 1.04_44},
                 {"mask": {"hash": "6affa964c6", "shape": (4_80, 6_40)}, "scores": 1.0_21},
                 {"mask": {"hash": "dfe28a0388", "shape": (4_80, 6_40)}, "scores": 1.01_67},
                 {"mask": {"hash": "c0a5f4a318", "shape": (4_80, 6_40)}, "scores": 1.01_32},
                 {"mask": {"hash": "fe8065c197", "shape": (4_80, 6_40)}, "scores": 1.00_53},
                 {"mask": {"hash": "e2d0b7a0b7", "shape": (4_80, 6_40)}, "scores": 0.99_67},
                 {"mask": {"hash": "453c7844bd", "shape": (4_80, 6_40)}, "scores": 0.9_93},
                 {"mask": {"hash": "3d44f2926d", "shape": (4_80, 6_40)}, "scores": 0.99_09},
                 {"mask": {"hash": "64033ddc3f", "shape": (4_80, 6_40)}, "scores": 0.98_79},
                 {"mask": {"hash": "801064ff79", "shape": (4_80, 6_40)}, "scores": 0.98_34},
                 {"mask": {"hash": "6172f276ef", "shape": (4_80, 6_40)}, "scores": 0.97_16},
                 {"mask": {"hash": "b49e60e084", "shape": (4_80, 6_40)}, "scores": 0.96_12},
                 {"mask": {"hash": "a811e775fd", "shape": (4_80, 6_40)}, "scores": 0.95_99},
                 {"mask": {"hash": "a6a8ebcf4b", "shape": (4_80, 6_40)}, "scores": 0.95_52},
                 {"mask": {"hash": "9d8257e080", "shape": (4_80, 6_40)}, "scores": 0.95_32},
                 {"mask": {"hash": "32de6454a8", "shape": (4_80, 6_40)}, "scores": 0.95_16},
                 {"mask": {"hash": "af3d4af2c8", "shape": (4_80, 6_40)}, "scores": 0.94_99},
                 {"mask": {"hash": "3c6db475fb", "shape": (4_80, 6_40)}, "scores": 0.94_83},
                 {"mask": {"hash": "c290813fb9", "shape": (4_80, 6_40)}, "scores": 0.94_64},
                 {"mask": {"hash": "b6f0b8f606", "shape": (4_80, 6_40)}, "scores": 0.9_43},
                 {"mask": {"hash": "92ce16bfdf", "shape": (4_80, 6_40)}, "scores": 0.9_43},
                 {"mask": {"hash": "c749b25868", "shape": (4_80, 6_40)}, "scores": 0.94_08},
                 {"mask": {"hash": "efb6cab859", "shape": (4_80, 6_40)}, "scores": 0.93_35},
                 {"mask": {"hash": "1ff2eafb30", "shape": (4_80, 6_40)}, "scores": 0.93_26},
                 {"mask": {"hash": "788b798e24", "shape": (4_80, 6_40)}, "scores": 0.92_62},
                 {"mask": {"hash": "abea804f0e", "shape": (4_80, 6_40)}, "scores": 0.89_99},
                 {"mask": {"hash": "7b9e8ddb73", "shape": (4_80, 6_40)}, "scores": 0.89_86},
                 {"mask": {"hash": "cd24047c8a", "shape": (4_80, 6_40)}, "scores": 0.89_84},
                 {"mask": {"hash": "6943e6bcbd", "shape": (4_80, 6_40)}, "scores": 0.88_73},
                 {"mask": {"hash": "b5f47c9191", "shape": (4_80, 6_40)}, "scores": 0.88_71}
             ]				,    )
         # fmt: on






       @require_torch
       @slow
       def     SCREAMING_SNAKE_CASE			(      self:    List[Any] ):
         lowercase			:int							=  "facebook/sam-vit-huge"
         lowercase			:List[Any]							=  pipeline("mask-generation"				,    model=_lowerCAmelCase )

         lowercase			:Any							=  image_segmenter(
             "http://images.cocodataset.org/val2017/000000039769.jpg"				,    pred_iou_thresh=1				,    points_per_batch=2_56 )

         # Shortening by hashing
         lowercase			:str							=  []
         for i, o in enumerate(outputs["masks"] ):
           new_outupt += [{"mask": mask_to_test_readable(_lowerCAmelCase ), "scores": outputs["scores"][i]}]

         self.assertEqual(
             nested_simplify(_lowerCAmelCase				,    decimals=4 )				,    [
                 {"mask": {"hash": "115ad19f5f", "shape": (4_80, 6_40)}, "scores": 1.04_44},
                 {"mask": {"hash": "6affa964c6", "shape": (4_80, 6_40)}, "scores": 1.02_10},
                 {"mask": {"hash": "dfe28a0388", "shape": (4_80, 6_40)}, "scores": 1.01_67},
                 {"mask": {"hash": "c0a5f4a318", "shape": (4_80, 6_40)}, "scores": 1.01_32},
                 {"mask": {"hash": "fe8065c197", "shape": (4_80, 6_40)}, "scores": 1.00_53},
             ]				,    )









import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class       __lowerCAmelCase				(							unittest.TestCase):






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    int ):
         torch.manual_seed(0 )
         lowercase			:str							=  UNetaDModel(
             sample_size=(32, 64)				,    in_channels=1				,    out_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("AttnDownBlock2D", "DownBlock2D")				,    up_block_types=("UpBlock2D", "AttnUpBlock2D")				,    )
         return model






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    Any ):
         torch.manual_seed(0 )
         lowercase			:Dict							=  UNetaDConditionModel(
             sample_size=(64, 32)				,    in_channels=1				,    out_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")				,    up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")				,    cross_attention_dim=10				,    )
         return model






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         torch.manual_seed(0 )
         lowercase			:List[str]							=  AutoencoderKL(
             sample_size=(1_28, 64)				,    in_channels=1				,    out_channels=1				,    latent_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D")				,    up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D")				,    )
         lowercase			:List[str]							=  UNetaDModel(
             sample_size=(64, 32)				,    in_channels=1				,    out_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("AttnDownBlock2D", "DownBlock2D")				,    up_block_types=("UpBlock2D", "AttnUpBlock2D")				,    )
         return vqvae, unet






       @slow
       def     SCREAMING_SNAKE_CASE			(      self:    str ):
         lowercase			:Optional[Any]							=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         lowercase			:Optional[int]							=  Mel(
             x_res=self.dummy_unet.config.sample_size[1]				,    y_res=self.dummy_unet.config.sample_size[0]				,    )

         lowercase			:List[str]							=  DDPMScheduler()
         lowercase			:Tuple							=  AudioDiffusionPipeline(vqvae=_lowerCAmelCase				,    unet=self.dummy_unet				,    mel=_lowerCAmelCase				,    scheduler=_lowerCAmelCase )
         lowercase			:Dict							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         lowercase			:Tuple							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Optional[int]							=  pipe(generator=_lowerCAmelCase				,    steps=4 )
         lowercase			:List[str]							=  output.audios[0]
         lowercase			:List[str]							=  output.images[0]

         lowercase			:List[str]							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Optional[int]							=  pipe(generator=_lowerCAmelCase				,    steps=4				,    return_dict=_lowerCAmelCase )
         lowercase			:int							=  output[0][0]

         assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
         assert (
             image.height == self.dummy_unet.config.sample_size[0]
             and image.width == self.dummy_unet.config.sample_size[1]
         )
         lowercase			:Any							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:Dict							=  np.frombuffer(image_from_tuple.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[Any]							=  np.array([69, 2_55, 2_55, 2_55, 0, 0, 77, 1_81, 12, 1_27] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() == 0
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() == 0

         lowercase			:Optional[Any]							=  Mel(
             x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]				,    y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]				,    )

         lowercase			:List[str]							=  DDIMScheduler()
         lowercase			:Tuple							=  self.dummy_vqvae_and_unet
         lowercase			:str							=  AudioDiffusionPipeline(
             vqvae=self.dummy_vqvae_and_unet[0]				,    unet=dummy_vqvae_and_unet[1]				,    mel=_lowerCAmelCase				,    scheduler=_lowerCAmelCase )
         lowercase			:Union[str, Any]							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         np.random.seed(0 )
         lowercase			:Dict							=  np.random.uniform(-1				,    1				,    ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,) )
         lowercase			:List[Any]							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Dict							=  pipe(raw_audio=_lowerCAmelCase				,    generator=_lowerCAmelCase				,    start_step=5				,    steps=10 )
         lowercase			:str							=  output.images[0]

         assert (
             image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
             and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
         )
         lowercase			:Optional[Any]							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[str]							=  np.array([1_20, 1_17, 1_10, 1_09, 1_38, 1_67, 1_38, 1_48, 1_32, 1_21] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() == 0

         lowercase			:Optional[Any]							=  self.dummy_unet_condition
         lowercase			:Optional[int]							=  AudioDiffusionPipeline(
             vqvae=self.dummy_vqvae_and_unet[0]				,    unet=_lowerCAmelCase				,    mel=_lowerCAmelCase				,    scheduler=_lowerCAmelCase )
         lowercase			:int							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         np.random.seed(0 )
         lowercase			:List[str]							=  torch.rand((1, 1, 10) )
         lowercase			:Union[str, Any]							=  pipe(generator=_lowerCAmelCase				,    encoding=_lowerCAmelCase )
         lowercase			:Tuple							=  output.images[0]
         lowercase			:Optional[Any]							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:Any							=  np.array([1_07, 1_03, 1_20, 1_27, 1_42, 1_22, 1_13, 1_22, 97, 1_11] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() == 0







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






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






       def     SCREAMING_SNAKE_CASE			(      self:    str ):
         lowercase			:Tuple							=  torch_device

         lowercase			:List[Any]							=  DiffusionPipeline.from_pretrained("teticio/audio-diffusion-ddim-256" )
         lowercase			:Dict							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         lowercase			:Any							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Optional[int]							=  pipe(generator=_lowerCAmelCase )
         lowercase			:List[str]							=  output.audios[0]
         lowercase			:Any							=  output.images[0]

         assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
         assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
         lowercase			:Dict							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[str]							=  np.array([1_51, 1_67, 1_54, 1_44, 1_22, 1_34, 1_21, 1_05, 70, 26] )

         assert np.abs(image_slice.flatten() - expected_slice ).max() == 0



