



import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from ...utils import logging
from ..auto import CONFIG_MAPPING


UpperCAmelCase							 =					logging.get_logger(__name__)

UpperCAmelCase							 =					{
    """salesforce/blip2-opt-2.7b""": """https://huggingface.co/salesforce/blip2-opt-2.7b/resolve/main/config.json""",
}

class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							UpperCamelCase__				:      List[str]         =					"""blip_2_vision_model"""


							def __init__(							self  ,  a_=1_408  ,  a_=6_144  ,  a_=39  ,  a_=16  ,  a_=224  ,  a_=14  ,  a_="gelu"  ,  a_=0.00_001  ,  a_=0.0  ,  a_=1E-1_0  ,  a_=True  ,  **a_  ,  ):
											super().__init__(**a_		)

											a__				      =		hidden_size
											a__				      =		intermediate_size
											a__				      =		num_hidden_layers
											a__				      =		num_attention_heads
											a__				      =		patch_size
											a__				      =		image_size
											a__				      =		initializer_range
											a__				      =		attention_dropout
											a__				      =		layer_norm_eps
											a__				      =		hidden_act
											a__				      =		qkv_bias







							@classmethod
							def _a      (							cls  ,  a_  ,  **a_		):
											cls._set_token_in_kwargs(a_		)

											a__      ,			a__				      =		cls.get_config_dict(a_  ,  **a_		)

											# get the vision config dict if we are loading from Blip2Config
											if config_dict.get("""model_type"""		) == "blip-2":
															a__				      =		config_dict["""vision_config"""]

											if "model_type" in config_dict and hasattr(cls  ,  """model_type"""		) and config_dict["model_type"] != cls.model_type:
															logger.warning(
															    F'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
															    F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''		)

											return cls.from_dict(a_  ,  **a_		)




class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							UpperCamelCase__				:      Optional[Any]         =					"""blip_2_qformer"""


							def __init__(							self  ,  a_=30_522  ,  a_=768  ,  a_=12  ,  a_=12  ,  a_=3_072  ,  a_="gelu"  ,  a_=0.1  ,  a_=0.1  ,  a_=512  ,  a_=0.02  ,  a_=1E-1_2  ,  a_=0  ,  a_="absolute"  ,  a_=2  ,  a_=1_408  ,  **a_  ,  ):
											super().__init__(pad_token_id=a_  ,  **a_		)

											a__				      =		vocab_size
											a__				      =		hidden_size
											a__				      =		num_hidden_layers
											a__				      =		num_attention_heads
											a__				      =		hidden_act
											a__				      =		intermediate_size
											a__				      =		hidden_dropout_prob
											a__				      =		attention_probs_dropout_prob
											a__				      =		max_position_embeddings
											a__				      =		initializer_range
											a__				      =		layer_norm_eps
											a__				      =		position_embedding_type
											a__				      =		cross_attention_frequency
											a__				      =		encoder_hidden_size







							@classmethod
							def _a      (							cls  ,  a_  ,  **a_		):
											cls._set_token_in_kwargs(a_		)

											a__      ,			a__				      =		cls.get_config_dict(a_  ,  **a_		)

											# get the qformer config dict if we are loading from Blip2Config
											if config_dict.get("""model_type"""		) == "blip-2":
															a__				      =		config_dict["""qformer_config"""]

											if "model_type" in config_dict and hasattr(cls  ,  """model_type"""		) and config_dict["model_type"] != cls.model_type:
															logger.warning(
															    F'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
															    F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''		)

											return cls.from_dict(a_  ,  **a_		)




class  __snake_case      (   SCREAMING_SNAKE_CASE):







							'''simple docstring'''






							UpperCamelCase__				:      Union[str, Any]         =					"""blip-2"""
							UpperCamelCase__				:      Optional[Any]         =					True


							def __init__(							self  ,  a_=None  ,  a_=None  ,  a_=None  ,  a_=32  ,  **a_		):
											super().__init__(**a_		)

											if vision_config is None:
															a__				      =		{}
															logger.info("""vision_config is None. initializing the Blip2VisionConfig with default values."""		)

											if qformer_config is None:
															a__				      =		{}
															logger.info("""qformer_config is None. Initializing the Blip2QFormerConfig with default values."""		)

											if text_config is None:
															a__				      =		{}
															logger.info("""text_config is None. Initializing the text config with default values (`OPTConfig`)."""		)

											a__				      =		BlipaVisionConfig(**a_		)
											a__				      =		BlipaQFormerConfig(**a_		)
											a__				      =		text_config["""model_type"""] if """model_type""" in text_config else """opt"""
											a__				      =		CONFIG_MAPPING[text_model_type](**a_		)

											a__				      =		self.text_config.tie_word_embeddings
											a__				      =		self.text_config.is_encoder_decoder

											a__				      =		num_query_tokens
											a__				      =		self.vision_config.hidden_size
											a__				      =		self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
											a__				      =		1.0
											a__				      =		0.02


							@classmethod
							def _a      (							cls  ,  a_  ,  a_  ,  a_  ,  **a_  ,  ):
											return cls(
											    vision_config=vision_config.to_dict()  ,  qformer_config=qformer_config.to_dict()  ,  text_config=text_config.to_dict()  ,  **a_  ,  )







							def _a      (							self		):
											a__				      =		copy.deepcopy(self.__dict__		)
											a__				      =		self.vision_config.to_dict()
											a__				      =		self.qformer_config.to_dict()
											a__				      =		self.text_config.to_dict()
											a__				      =		self.__class__.model_type
											return output






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  __snake_case      (   unittest.TestCase):







							'''simple docstring'''






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


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


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


							@property
							def _a      (							self		):
											torch.manual_seed(0		)
											a__				      =		AutoencoderKL(
											    sample_size=(128, 64)  ,  in_channels=1  ,  out_channels=1  ,  latent_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")  ,  up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")  ,  )
											a__				      =		UNetaDModel(
											    sample_size=(64, 32)  ,  in_channels=1  ,  out_channels=1  ,  layers_per_block=2  ,  block_out_channels=(128, 128)  ,  down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")  ,  up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")  ,  )
											return vqvae, unet







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

											a__				      =		DDPMScheduler()
											a__				      =		AudioDiffusionPipeline(vqvae=a_  ,  unet=self.dummy_unet  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4		)
											a__				      =		output.audios[0]
											a__				      =		output.images[0]

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_  ,  steps=4  ,  return_dict=a_		)
											a__				      =		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]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.frombuffer(image_from_tuple.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]		)

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

											a__				      =		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]  ,  )

											a__				      =		DDIMScheduler()
											a__				      =		self.dummy_vqvae_and_unet
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=dummy_vqvae_and_unet[1]  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		np.random.uniform(-1  ,  1  ,  ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)		)
											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(raw_audio=a_  ,  generator=a_  ,  start_step=5  ,  steps=10		)
											a__				      =		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]
											)
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]		)

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

											a__				      =		self.dummy_unet_condition
											a__				      =		AudioDiffusionPipeline(
											    vqvae=self.dummy_vqvae_and_unet[0]  ,  unet=a_  ,  mel=a_  ,  scheduler=a_		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											np.random.seed(0		)
											a__				      =		torch.rand((1, 1, 10)		)
											a__				      =		pipe(generator=a_  ,  encoding=a_		)
											a__				      =		output.images[0]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]		)

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




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







							'''simple docstring'''






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







							def _a      (							self		):
											a__				      =		torch_device

											a__				      =		DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""		)
											a__				      =		pipe.to(a_		)
											pipe.set_progress_bar_config(disable=a_		)

											a__				      =		torch.Generator(device=a_		).manual_seed(42		)
											a__				      =		pipe(generator=a_		)
											a__				      =		output.audios[0]
											a__				      =		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]
											a__				      =		np.frombuffer(image.tobytes()  ,  dtype="""uint8"""		)[:10]
											a__				      =		np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]		)

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

