




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 							__SCREAMING_SNAKE_CASE(  self						,     _A = False     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			1_0

											if only_cross_attention:
																__lowerCAmelCase							  =			1_2
											else:
																# when only cross attention is not set, the cross attention dim must be the same as the query dim
																__lowerCAmelCase							  =			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 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			2

											__lowerCAmelCase							  =			torch.rand(_A						,     _A						,     3						,     2     )
											__lowerCAmelCase							  =			torch.rand(_A						,     4						,     _A     )
											__lowerCAmelCase							  =			None

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






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											torch.manual_seed(0     )

											__lowerCAmelCase							  =			self.get_constructor_arguments(only_cross_attention=_A     )
											__lowerCAmelCase							  =			Attention(**_A     )

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

											__lowerCAmelCase							  =			self.get_forward_arguments(
											    query_dim=constructor_args["query_dim"]						,     added_kv_proj_dim=constructor_args["added_kv_proj_dim"]     )

											__lowerCAmelCase							  =			attn(**_A     )

											# only self attention

											torch.manual_seed(0     )

											__lowerCAmelCase							  =			self.get_constructor_arguments(only_cross_attention=_A     )
											__lowerCAmelCase							  =			Attention(**_A     )

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

											__lowerCAmelCase							  =			self.get_forward_arguments(
											    query_dim=constructor_args["query_dim"]						,     added_kv_proj_dim=constructor_args["added_kv_proj_dim"]     )

											__lowerCAmelCase							  =			attn(**_A     )

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






class      a__       (   unittest.TestCase ):






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			DiffusionPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-pipe"						,     safety_checker=_A     )

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

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

											__lowerCAmelCase							  =			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(_A     )

																# can also load the converted weights
																__lowerCAmelCase							  =			DiffusionPipeline.from_pretrained(_A						,     device_map="sequential"						,     safety_checker=_A     )

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

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










from typing import TYPE_CHECKING

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


UpperCamelCase__            =     {"""configuration_vit_hybrid""": ["""VIT_HYBRID_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ViTHybridConfig"""]}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCamelCase__            =     [
					    """VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """ViTHybridForImageClassification""",
					    """ViTHybridModel""",
					    """ViTHybridPreTrainedModel""",
					]

try:
					if not is_vision_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCamelCase__            =     ["""ViTHybridImageProcessor"""]


if TYPE_CHECKING:
					from .configuration_vit_hybrid import VIT_HYBRID_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTHybridConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_vit_hybrid import (
										    VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST,
										    ViTHybridForImageClassification,
										    ViTHybridModel,
										    ViTHybridPreTrainedModel,
										)

					try:
										if not is_vision_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .image_processing_vit_hybrid import ViTHybridImageProcessor


else:
					import sys

					UpperCamelCase__            =     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




