





import tempfile
import unittest

import numpy as np
import torch

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






class 							_a(     unittest.TestCase     ):







		def 							lowercase							(     self			,    __snake_case = False					)   ->   Tuple:


							'''simple docstring'''

							_snake_case							:       List[str]												= 1_0

							if only_cross_attention:
												_snake_case							:       Dict												= 1_2
							else:
												# when only cross attention is not set, the cross attention dim must be the same as the query dim
												_snake_case							:       int												= 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 							lowercase							(     self			,    __snake_case			,    __snake_case					)   ->   List[Any]:


							'''simple docstring'''

							_snake_case							:       List[Any]												= 2

							_snake_case							:       Any												= torch.rand(__snake_case			,    __snake_case			,    3			,    2					)
							_snake_case							:       Dict												= torch.rand(__snake_case			,    4			,    __snake_case					)
							_snake_case							:       Any												= None

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







		def 							lowercase							(     self					)   ->   Union[str, Any]:


							'''simple docstring'''

							torch.manual_seed(0					)

							_snake_case							:       Tuple												= self.get_constructor_arguments(only_cross_attention=__snake_case					)
							_snake_case							:       Dict												= Attention(**__snake_case					)

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

							_snake_case							:       Tuple												= self.get_forward_arguments(
							    query_dim=constructor_args["query_dim"]			,    added_kv_proj_dim=constructor_args["added_kv_proj_dim"]					)

							_snake_case							:       List[str]												= attn(**__snake_case					)

							# only self attention

							torch.manual_seed(0					)

							_snake_case							:       str												= self.get_constructor_arguments(only_cross_attention=__snake_case					)
							_snake_case							:       Tuple												= Attention(**__snake_case					)

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

							_snake_case							:       Any												= self.get_forward_arguments(
							    query_dim=constructor_args["query_dim"]			,    added_kv_proj_dim=constructor_args["added_kv_proj_dim"]					)

							_snake_case							:       str												= attn(**__snake_case					)

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






class 							_a(     unittest.TestCase     ):







		def 							lowercase							(     self					)   ->   Optional[int]:


							'''simple docstring'''

							_snake_case							:       Union[str, Any]												= DiffusionPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-pipe"			,    safety_checker=__snake_case					)

							_snake_case							:       Any												= pipe(
							    "foo"			,    num_inference_steps=2			,    generator=torch.Generator("cpu"					).manual_seed(0					)			,    output_type="np"			,    ).images

							# the initial conversion succeeds
							_snake_case							:       Any												= DiffusionPipeline.from_pretrained(
							    "hf-internal-testing/tiny-stable-diffusion-pipe"			,    device_map="sequential"			,    safety_checker=__snake_case					)

							_snake_case							:       Tuple												= 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(__snake_case					)

												# can also load the converted weights
												_snake_case							:       int												= DiffusionPipeline.from_pretrained(__snake_case			,    device_map="sequential"			,    safety_checker=__snake_case					)

							_snake_case							:       List[Any]												= pipe(
							    "foo"			,    num_inference_steps=2			,    generator=torch.Generator("cpu"					).manual_seed(0					)			,    output_type="np"			,    ).images

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






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__lowerCAmelCase			:Optional[int]	 =     {
    'configuration_mobilenet_v1': [
        'MOBILENET_V1_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'MobileNetV1Config',
        'MobileNetV1OnnxConfig',
    ],
}

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__lowerCAmelCase			:Dict	 =     ['MobileNetV1FeatureExtractor']
							__lowerCAmelCase			:List[Any]	 =     ['MobileNetV1ImageProcessor']

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__lowerCAmelCase			:Optional[int]	 =     [
							    'MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'MobileNetV1ForImageClassification',
							    'MobileNetV1Model',
							    'MobileNetV1PreTrainedModel',
							    'load_tf_weights_in_mobilenet_v1',
							]


if TYPE_CHECKING:
							from .configuration_mobilenet_va import (
							    MOBILENET_V1_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    MobileNetVaConfig,
							    MobileNetVaOnnxConfig,
							)

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .feature_extraction_mobilenet_va import MobileNetVaFeatureExtractor
														from .image_processing_mobilenet_va import MobileNetVaImageProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_mobilenet_va import (
														    MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST,
														    MobileNetVaForImageClassification,
														    MobileNetVaModel,
														    MobileNetVaPreTrainedModel,
														    load_tf_weights_in_mobilenet_va,
														)


else:
							import sys

							__lowerCAmelCase			:Optional[int]	 =     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)