




import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor

class      UpperCAmelCase						(			unittest.TestCase   ):






   def  _a						(	self   :       int						,					UpperCamelCase_   :       bool = False     )->			Optional[Any]:
     _lowercase: Optional[int]						     =    10

     if only_cross_attention:
       _lowercase: List[Any]						     =    12
     else:
       # when only cross attention is not set, the cross attention dim must be the same as the query dim
       _lowercase: Dict						     =    query_dim

     return {
         "query_dim": query_dim,
         "cross_attention_dim": cross_attention_dim,
         "heads": 2,
         "dim_head": 4,
         "added_kv_proj_dim": 6,
         "norm_num_groups": 1,
         "only_cross_attention": only_cross_attention,
         "processor": AttnAddedKVProcessor(),
     }






   def  _a						(	self   :       Tuple						,					UpperCamelCase_   :       Union[str, Any]						,					UpperCamelCase_   :       Optional[int]     )->			Dict:
     _lowercase: Optional[Any]						     =    2

     _lowercase: List[str]						     =    torch.rand(UpperCamelCase_						,					UpperCamelCase_						,					3						,					2     )
     _lowercase: List[Any]						     =    torch.rand(UpperCamelCase_						,					4						,					UpperCamelCase_     )
     _lowercase: List[str]						     =    None

     return {
         "hidden_states": hidden_states,
         "encoder_hidden_states": encoder_hidden_states,
         "attention_mask": attention_mask,
     }






   def  _a						(	self   :       Dict     )->			List[str]:
     # self and cross attention

     torch.manual_seed(0     )

     _lowercase: Dict						     =    self.get_constructor_arguments(only_cross_attention=UpperCamelCase_     )
     _lowercase: List[Any]						     =    Attention(**UpperCamelCase_     )

     self.assertTrue(attn.to_k is not None     )
     self.assertTrue(attn.to_v is not None     )

     _lowercase: Union[str, Any]						     =    self.get_forward_arguments(
         query_dim=constructor_args['query_dim']						,					added_kv_proj_dim=constructor_args['added_kv_proj_dim']     )

     _lowercase: List[Any]						     =    attn(**UpperCamelCase_     )

     # only self attention

     torch.manual_seed(0     )

     _lowercase: Tuple						     =    self.get_constructor_arguments(only_cross_attention=UpperCamelCase_     )
     _lowercase: str						     =    Attention(**UpperCamelCase_     )

     self.assertTrue(attn.to_k is None     )
     self.assertTrue(attn.to_v is None     )

     _lowercase: Optional[Any]						     =    self.get_forward_arguments(
         query_dim=constructor_args['query_dim']						,					added_kv_proj_dim=constructor_args['added_kv_proj_dim']     )

     _lowercase: Dict						     =    attn(**UpperCamelCase_     )

     self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()     )

class      UpperCAmelCase						(			unittest.TestCase   ):






   def  _a						(	self   :       Tuple     )->			str:
     _lowercase: List[Any]						     =    DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe'						,					safety_checker=UpperCamelCase_     )

     _lowercase: Optional[Any]						     =    pipe(
         'foo'						,					num_inference_steps=2						,					generator=torch.Generator('cpu'     ).manual_seed(0     )						,					output_type='np'						,					).images

     # the initial conversion succeeds
     _lowercase: Dict						     =    DiffusionPipeline.from_pretrained(
         'hf-internal-testing/tiny-stable-diffusion-pipe'						,					device_map='sequential'						,					safety_checker=UpperCamelCase_     )

     _lowercase: int						     =    pipe(
         'foo'						,					num_inference_steps=2						,					generator=torch.Generator('cpu'     ).manual_seed(0     )						,					output_type='np'						,					).images

     with tempfile.TemporaryDirectory() as tmpdir:
       # save the converted model
       pipe.save_pretrained(UpperCamelCase_     )

       # can also load the converted weights
       _lowercase: List[str]						     =    DiffusionPipeline.from_pretrained(UpperCamelCase_						,					device_map='sequential'						,					safety_checker=UpperCamelCase_     )

     _lowercase: str						     =    pipe(
         'foo'						,					num_inference_steps=2						,					generator=torch.Generator('cpu'     ).manual_seed(0     )						,					output_type='np'						,					).images

     self.assertTrue(np.allclose(UpperCamelCase_						,					UpperCamelCase_						,					atol=1E-5     )     )
     self.assertTrue(np.allclose(UpperCamelCase_						,					UpperCamelCase_						,					atol=1E-5     )     )






import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin

class      UpperCAmelCase						(			A   ):
   __SCREAMING_SNAKE_CASE		:							Optional[Any]							=					"AutoFeatureExtractor"
   __SCREAMING_SNAKE_CASE		:							Optional[Any]							=					"Speech2Text2Tokenizer"






   def __init__(	self   :       Optional[Any]						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       List[Any]     )->			Dict:
     super().__init__(UpperCamelCase_						,					UpperCamelCase_     )
     _lowercase: Optional[int]						     =    self.feature_extractor
     _lowercase: int						     =    False






   def __call__(	self   :       List[Any]						,					*UpperCamelCase_   :       Dict						,					**UpperCamelCase_   :       Optional[int]     )->			str:
     # For backward compatibility
     if self._in_target_context_manager:
       return self.current_processor(*UpperCamelCase_						,					**UpperCamelCase_     )

     if "raw_speech" in kwargs:
       warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'     )
       _lowercase: Tuple						     =    kwargs.pop('raw_speech'     )
     else:
       _lowercase: Optional[Any]						     =    kwargs.pop('audio'						,					UpperCamelCase_     )
     _lowercase: int						     =    kwargs.pop('sampling_rate'						,					UpperCamelCase_     )
     _lowercase: Dict						     =    kwargs.pop('text'						,					UpperCamelCase_     )
     if len(UpperCamelCase_     ) > 0:
       _lowercase: List[Any]						     =    args[0]
       _lowercase: List[Any]						     =    args[1:]

     if audio is None and text is None:
       raise ValueError('You need to specify either an `audio` or `text` input to process.'     )

     if audio is not None:
       _lowercase: Optional[int]						     =    self.feature_extractor(UpperCamelCase_						,					*UpperCamelCase_						,					sampling_rate=UpperCamelCase_						,					**UpperCamelCase_     )
     if text is not None:
       _lowercase: str						     =    self.tokenizer(UpperCamelCase_						,					**UpperCamelCase_     )

     if text is None:
       return inputs
     elif audio is None:
       return encodings
     else:
       _lowercase: int						     =    encodings['input_ids']
       return inputs






   def  _a						(	self   :       Optional[int]						,					*UpperCamelCase_   :       Tuple						,					**UpperCamelCase_   :       Dict     )->			List[str]:
     return self.tokenizer.batch_decode(*UpperCamelCase_						,					**UpperCamelCase_     )






   def  _a						(	self   :       Tuple						,					*UpperCamelCase_   :       List[Any]						,					**UpperCamelCase_   :       Union[str, Any]     )->			Optional[int]:
     return self.tokenizer.decode(*UpperCamelCase_						,					**UpperCamelCase_     )






   @contextmanager
   def  _a						(	self   :       Tuple     )->			Union[str, Any]:
     warnings.warn(
         '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
         'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
         'your audio inputs, or in a separate call.'     )
     _lowercase: Any						     =    True
     _lowercase: Optional[Any]						     =    self.tokenizer
     yield
     _lowercase: Union[str, Any]						     =    self.feature_extractor
     _lowercase: str						     =    False
