



'''simple docstring'''






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class 				lowercase			(    _lowerCamelCase   ,   _lowerCamelCase   ,   _lowerCamelCase   ,   unittest.TestCase							):


 """simple docstring"""


 UpperCAmelCase		      =    StableDiffusionImageVariationPipeline
 UpperCAmelCase		      =    IMAGE_VARIATION_PARAMS
 UpperCAmelCase		      =    IMAGE_VARIATION_BATCH_PARAMS
 UpperCAmelCase		      =    frozenset([]							)
 # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 UpperCAmelCase		      =    frozenset([]							)







 def _snake_case					(				self				)						->  List[Any]:
        torch.manual_seed(0				)
        _UpperCAmelCase    :		Tuple     	=					UNetaDConditionModel(
            block_out_channels=(32, 64)    ,layers_per_block=2    ,sample_size=32    ,in_channels=4    ,out_channels=4    ,down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")    ,up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")    ,cross_attention_dim=32    ,)
        _UpperCAmelCase    :		List[str]     	=					PNDMScheduler(skip_prk_steps=a__				)
        torch.manual_seed(0				)
        _UpperCAmelCase    :		Union[str, Any]     	=					AutoencoderKL(
            block_out_channels=[32, 64]    ,in_channels=3    ,out_channels=3    ,down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]    ,up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]    ,latent_channels=4    ,)
        torch.manual_seed(0				)
        _UpperCAmelCase    :		List[Any]     	=					CLIPVisionConfig(
            hidden_size=32    ,projection_dim=32    ,intermediate_size=37    ,layer_norm_eps=1E-0_5    ,num_attention_heads=4    ,num_hidden_layers=5    ,image_size=32    ,patch_size=4    ,)
        _UpperCAmelCase    :		Tuple     	=					CLIPVisionModelWithProjection(a__				)
        _UpperCAmelCase    :		Optional[Any]     	=					CLIPImageProcessor(crop_size=32    ,size=32				)

        _UpperCAmelCase    :		List[Any]     	=					{
            """unet""": unet,
            """scheduler""": scheduler,
            """vae""": vae,
            """image_encoder""": image_encoder,
            """feature_extractor""": feature_extractor,
            """safety_checker""": None,
        }
        return components







 def _snake_case					(				self    ,a_    ,a_=0				)						->  List[Any]:
        _UpperCAmelCase    :		Optional[int]     	=					floats_tensor((1, 3, 32, 32)    ,rng=random.Random(a__				)				)
        _UpperCAmelCase    :		Any     	=					image.cpu().permute(0    ,2    ,3    ,1				)[0]
        _UpperCAmelCase    :		Optional[Any]     	=					Image.fromarray(np.uinta(a__				)				).convert("""RGB"""				).resize((32, 32)				)
        if str(a__				).startswith("""mps"""				):
               _UpperCAmelCase    :		List[str]     	=					torch.manual_seed(a__				)
        else:
               _UpperCAmelCase    :		Optional[Any]     	=					torch.Generator(device=a__				).manual_seed(a__				)
        _UpperCAmelCase    :		Optional[Any]     	=					{
            """image""": image,
            """generator""": generator,
            """num_inference_steps""": 2,
            """guidance_scale""": 6.0,
            """output_type""": """numpy""",
        }
        return inputs







 def _snake_case					(				self				)						->  Tuple:
        _UpperCAmelCase    :		int     	=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase    :		Optional[Any]     	=					self.get_dummy_components()
        _UpperCAmelCase    :		Any     	=					StableDiffusionImageVariationPipeline(**a__				)
        _UpperCAmelCase    :		Dict     	=					sd_pipe.to(a__				)
        sd_pipe.set_progress_bar_config(disable=a__				)

        _UpperCAmelCase    :		Any     	=					self.get_dummy_inputs(a__				)
        _UpperCAmelCase    :		Any     	=					sd_pipe(**a__				).images
        _UpperCAmelCase    :		Optional[int]     	=					image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)
        _UpperCAmelCase    :		List[str]     	=					np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921]				)

        assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-3







 def _snake_case					(				self				)						->  Optional[Any]:
        _UpperCAmelCase    :		str     	=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase    :		List[str]     	=					self.get_dummy_components()
        _UpperCAmelCase    :		Union[str, Any]     	=					StableDiffusionImageVariationPipeline(**a__				)
        _UpperCAmelCase    :		Tuple     	=					sd_pipe.to(a__				)
        sd_pipe.set_progress_bar_config(disable=a__				)

        _UpperCAmelCase    :		Optional[Any]     	=					self.get_dummy_inputs(a__				)
        _UpperCAmelCase    :		List[Any]     	=					2 * [inputs["""image"""]]
        _UpperCAmelCase    :		Dict     	=					sd_pipe(**a__				)

        _UpperCAmelCase    :		Tuple     	=					output.images

        _UpperCAmelCase    :		str     	=					image[-1, -3:, -3:, -1]

        assert image.shape == (2, 64, 64, 3)
        _UpperCAmelCase    :		Any     	=					np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289]				)

        assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-3







 def _snake_case					(				self				)						->  Dict:
        super().test_inference_batch_single_identical(expected_max_diff=3E-3				)







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


 """simple docstring"""


 def _snake_case					(				self				)						->  List[Any]:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







 def _snake_case					(				self    ,a_    ,a_="cpu"    ,a_=torch.floataa    ,a_=0				)						->  Optional[Any]:
        _UpperCAmelCase    :		List[str]     	=					torch.Generator(device=a__				).manual_seed(a__				)
        _UpperCAmelCase    :		int     	=					load_image(
            """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
            """/stable_diffusion_imgvar/input_image_vermeer.png"""				)
        _UpperCAmelCase    :		Dict     	=					np.random.RandomState(a__				).standard_normal((1, 4, 64, 64)				)
        _UpperCAmelCase    :		List[Any]     	=					torch.from_numpy(a__				).to(device=a__    ,dtype=a__				)
        _UpperCAmelCase    :		Optional[Any]     	=					{
            """image""": init_image,
            """latents""": latents,
            """generator""": generator,
            """num_inference_steps""": 3,
            """guidance_scale""": 7.5,
            """output_type""": """numpy""",
        }
        return inputs







 def _snake_case					(				self				)						->  Union[str, Any]:
        _UpperCAmelCase    :		int     	=					StableDiffusionImageVariationPipeline.from_pretrained(
            """lambdalabs/sd-image-variations-diffusers"""    ,safety_checker=a__				)
        _UpperCAmelCase    :		str     	=					sd_pipe.to(a__				)
        sd_pipe.set_progress_bar_config(disable=a__				)

        _UpperCAmelCase    :		Union[str, Any]     	=					self.get_inputs(a__				)
        _UpperCAmelCase    :		List[Any]     	=					sd_pipe(**a__				).images
        _UpperCAmelCase    :		Dict     	=					image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)
        _UpperCAmelCase    :		Tuple     	=					np.array([0.8_4491, 0.9_0789, 0.7_5708, 0.7_8734, 0.8_3485, 0.7_0099, 0.6_6938, 0.6_8727, 0.6_1379]				)
        assert np.abs(image_slice - expected_slice				).max() < 6E-3







 def _snake_case					(				self				)						->  Union[str, Any]:
        _UpperCAmelCase    :		List[str]     	=					0

        def callback_fn(a_    ,a_    ,a_				) -> None:
               _UpperCAmelCase    :		Union[str, Any]     	=					True
               nonlocal number_of_steps
               number_of_steps += 1
               if step == 1:
                      _UpperCAmelCase    :		Tuple     	=					latents.detach().cpu().numpy()
                      assert latents.shape == (1, 4, 64, 64)
                      _UpperCAmelCase    :		Tuple     	=					latents[0, -3:, -3:, -1]
                      _UpperCAmelCase    :		List[Any]     	=					np.array(
                          [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309]				)

                      assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5E-2
               elif step == 2:
                      _UpperCAmelCase    :		Union[str, Any]     	=					latents.detach().cpu().numpy()
                      assert latents.shape == (1, 4, 64, 64)
                      _UpperCAmelCase    :		Dict     	=					latents[0, -3:, -3:, -1]
                      _UpperCAmelCase    :		Union[str, Any]     	=					np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273]				)

                      assert np.abs(latents_slice.flatten() - expected_slice				).max() < 5E-2

        _UpperCAmelCase    :		int     	=					False

        _UpperCAmelCase    :		List[Any]     	=					StableDiffusionImageVariationPipeline.from_pretrained(
            """fusing/sd-image-variations-diffusers"""    ,safety_checker=a__    ,torch_dtype=torch.floataa    ,)
        pipe.to(a__				)
        pipe.set_progress_bar_config(disable=a__				)
        pipe.enable_attention_slicing()

        _UpperCAmelCase    :		Optional[int]     	=					self.get_inputs(a__    ,dtype=torch.floataa				)
        pipe(**a__    ,callback=a__    ,callback_steps=1				)
        assert callback_fn.has_been_called
        assert number_of_steps == inputs["num_inference_steps"]







 def _snake_case					(				self				)						->  Dict:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        _UpperCAmelCase    :		str     	=					"""fusing/sd-image-variations-diffusers"""
        _UpperCAmelCase    :		Optional[Any]     	=					StableDiffusionImageVariationPipeline.from_pretrained(
            a__    ,safety_checker=a__    ,torch_dtype=torch.floataa				)
        _UpperCAmelCase    :		Any     	=					pipe.to(a__				)
        pipe.set_progress_bar_config(disable=a__				)
        pipe.enable_attention_slicing(1				)
        pipe.enable_sequential_cpu_offload()

        _UpperCAmelCase    :		str     	=					self.get_inputs(a__    ,dtype=torch.floataa				)
        _UpperCAmelCase    :		Optional[Any]     	=					pipe(**a__				)

        _UpperCAmelCase    :		Dict     	=					torch.cuda.max_memory_allocated()
        # make sure that less than 2.6 GB is allocated
        assert mem_bytes < 2.6 * 10**9







@nightly
@require_torch_gpu
class 				lowercase			(    unittest.TestCase							):


 """simple docstring"""


 def _snake_case					(				self				)						->  Union[str, Any]:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







 def _snake_case					(				self    ,a_    ,a_="cpu"    ,a_=torch.floataa    ,a_=0				)						->  List[str]:
        _UpperCAmelCase    :		Dict     	=					torch.Generator(device=a__				).manual_seed(a__				)
        _UpperCAmelCase    :		Tuple     	=					load_image(
            """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
            """/stable_diffusion_imgvar/input_image_vermeer.png"""				)
        _UpperCAmelCase    :		Optional[Any]     	=					np.random.RandomState(a__				).standard_normal((1, 4, 64, 64)				)
        _UpperCAmelCase    :		Any     	=					torch.from_numpy(a__				).to(device=a__    ,dtype=a__				)
        _UpperCAmelCase    :		Optional[Any]     	=					{
            """image""": init_image,
            """latents""": latents,
            """generator""": generator,
            """num_inference_steps""": 50,
            """guidance_scale""": 7.5,
            """output_type""": """numpy""",
        }
        return inputs







 def _snake_case					(				self				)						->  Any:
        _UpperCAmelCase    :		Optional[Any]     	=					StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""				)
        sd_pipe.to(a__				)
        sd_pipe.set_progress_bar_config(disable=a__				)

        _UpperCAmelCase    :		Optional[Any]     	=					self.get_inputs(a__				)
        _UpperCAmelCase    :		Tuple     	=					sd_pipe(**a__				).images[0]

        _UpperCAmelCase    :		Any     	=					load_numpy(
            """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
            """/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy"""				)
        _UpperCAmelCase    :		Tuple     	=					np.abs(expected_image - image				).max()
        assert max_diff < 1E-3







 def _snake_case					(				self				)						->  Union[str, Any]:
        _UpperCAmelCase    :		Optional[Any]     	=					StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""				)
        _UpperCAmelCase    :		List[str]     	=					DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config				)
        sd_pipe.to(a__				)
        sd_pipe.set_progress_bar_config(disable=a__				)

        _UpperCAmelCase    :		Union[str, Any]     	=					self.get_inputs(a__				)
        _UpperCAmelCase    :		Any     	=					25
        _UpperCAmelCase    :		int     	=					sd_pipe(**a__				).images[0]

        _UpperCAmelCase    :		Optional[int]     	=					load_numpy(
            """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
            """/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy"""				)
        _UpperCAmelCase    :		Dict     	=					np.abs(expected_image - image				).max()
        assert max_diff < 1E-3










"""simple docstring"""






import copy
import inspect

from transformers.testing_utils import require_torch, torch_device
from transformers.utils.backbone_utils import BackboneType



@require_torch
class 						lowerCAmelCase       :
       lowerCAmelCase__    			=	()
       lowerCAmelCase__    			=	True






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.config_class

              # test default config
              _UpperCAmelCase			     =  config_class()
              self.assertIsNotNone(a__  )
              _UpperCAmelCase			     =  ['stem'] + [f"""stage{idx}""" for idx in range(1     ,	len(config.depths  ) + 1  )]
              self.assertEqual(config.stage_names     ,	a__  )
              self.assertTrue(set(config.out_features  ).issubset(set(config.stage_names  )  )  )

              # Test out_features and out_indices are correctly set
              # out_features and out_indices both None
              _UpperCAmelCase			     =  config_class(out_features=a__     ,	out_indices=a__  )
              self.assertEqual(config.out_features     ,	[config.stage_names[-1]]  )
              self.assertEqual(config.out_indices     ,	[len(config.stage_names  ) - 1]  )

              # out_features and out_indices both set
              _UpperCAmelCase			     =  config_class(out_features=['stem', 'stage1']     ,	out_indices=[0, 1]  )
              self.assertEqual(config.out_features     ,	['stem', 'stage1']  )
              self.assertEqual(config.out_indices     ,	[0, 1]  )

              # Only out_features set
              _UpperCAmelCase			     =  config_class(out_features=['stage1', 'stage3']  )
              self.assertEqual(config.out_features     ,	['stage1', 'stage3']  )
              self.assertEqual(config.out_indices     ,	[1, 3]  )

              # Only out_indices set
              _UpperCAmelCase			     =  config_class(out_indices=[0, 2]  )
              self.assertEqual(config.out_features     ,	[config.stage_names[0], config.stage_names[2]]  )
              self.assertEqual(config.out_indices     ,	[0, 2]  )

              # Error raised when out_indices do not correspond to out_features
              with self.assertRaises(a__  ):
                     _UpperCAmelCase			     =  config_class(out_features=['stage1', 'stage2']     ,	out_indices=[0, 2]  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     _UpperCAmelCase			     =  model_class(a__  )
                     _UpperCAmelCase			     =  inspect.signature(model.forward  )
                     # signature.parameters is an OrderedDict => so arg_names order is deterministic
                     _UpperCAmelCase			     =  [*signature.parameters.keys()]
                     _UpperCAmelCase			     =  ['pixel_values']
                     self.assertListEqual(arg_names[:1]     ,	a__  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     _UpperCAmelCase			     =  model_class(a__  )
                     self.assertEqual(len(model.channels  )     ,	len(config.out_features  )  )
                     _UpperCAmelCase			     =  model.num_features
                     _UpperCAmelCase			     =  [config.stage_names.index(a__  ) for feat in config.out_features]
                     _UpperCAmelCase			     =  [num_features[idx] for idx in out_indices]
                     self.assertListEqual(model.channels     ,	a__  )

                     _UpperCAmelCase			     =  copy.deepcopy(a__  )
                     _UpperCAmelCase			     =  None
                     _UpperCAmelCase			     =  model_class(a__  )
                     self.assertEqual(len(model.channels  )     ,	1  )
                     self.assertListEqual(model.channels     ,	[num_features[-1]]  )

                     _UpperCAmelCase			     =  copy.deepcopy(a__  )
                     _UpperCAmelCase			     =  None
                     _UpperCAmelCase			     =  model_class(a__  )
                     self.assertEqual(len(model.channels  )     ,	1  )
                     self.assertListEqual(model.channels     ,	[num_features[-1]]  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     _UpperCAmelCase			     =  model_class(a__  )
                     model.to(a__  )
                     model.eval()
                     _UpperCAmelCase			     =  model(**a__  )

                     self.assertEqual(len(result.feature_maps  )     ,	len(config.out_features  )  )
                     self.assertEqual(len(model.channels  )     ,	len(config.out_features  )  )
                     self.assertEqual(len(result.feature_maps  )     ,	len(config.out_indices  )  )
                     self.assertEqual(len(model.channels  )     ,	len(config.out_indices  )  )

                     # Check output of last stage is taken if out_features=None, out_indices=None
                     _UpperCAmelCase			     =  copy.deepcopy(a__  )
                     _UpperCAmelCase			     =  None
                     _UpperCAmelCase			     =  model_class(a__  )
                     model.to(a__  )
                     model.eval()
                     _UpperCAmelCase			     =  model(**a__  )

                     self.assertEqual(len(result.feature_maps  )     ,	1  )
                     self.assertEqual(len(model.channels  )     ,	1  )

                     _UpperCAmelCase			     =  copy.deepcopy(a__  )
                     _UpperCAmelCase			     =  None
                     _UpperCAmelCase			     =  model_class(a__  )
                     model.to(a__  )
                     model.eval()
                     _UpperCAmelCase			     =  model(**a__  )

                     self.assertEqual(len(result.feature_maps  )     ,	1  )
                     self.assertEqual(len(model.channels  )     ,	1  )

                     # Check backbone can be initialized with fresh weights
                     _UpperCAmelCase			     =  copy.deepcopy(a__  )
                     _UpperCAmelCase			     =  False
                     _UpperCAmelCase			     =  model_class(a__  )
                     model.to(a__  )
                     model.eval()
                     _UpperCAmelCase			     =  model(**a__  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs_for_common()

              for backbone_class in self.all_model_classes:
                     _UpperCAmelCase			     =  backbone_class(a__  )

                     self.assertTrue(hasattr(a__     ,	'backbone_type'  )  )
                     self.assertTrue(hasattr(a__     ,	'stage_names'  )  )
                     self.assertTrue(hasattr(a__     ,	'num_features'  )  )
                     self.assertTrue(hasattr(a__     ,	'out_indices'  )  )
                     self.assertTrue(hasattr(a__     ,	'out_features'  )  )
                     self.assertTrue(hasattr(a__     ,	'out_feature_channels'  )  )
                     self.assertTrue(hasattr(a__     ,	'channels'  )  )

                     self.assertIsInstance(backbone.backbone_type     ,	a__  )
                     # Verify num_features has been initialized in the backbone init
                     self.assertIsNotNone(backbone.num_features  )
                     self.assertTrue(len(backbone.channels  ) == len(backbone.out_indices  )  )
                     self.assertTrue(len(backbone.stage_names  ) == len(backbone.num_features  )  )
                     self.assertTrue(len(backbone.channels  ) <= len(backbone.num_features  )  )
                     self.assertTrue(len(backbone.out_feature_channels  ) == len(backbone.stage_names  )  )






       def 				__A			( self  ):
              _UpperCAmelCase     ,				_UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs_for_common()
              _UpperCAmelCase			     =  inputs_dict['pixel_values'].shape[0]

              for backbone_class in self.all_model_classes:
                     _UpperCAmelCase			     =  backbone_class(a__  )
                     backbone.to(a__  )
                     backbone.eval()

                     _UpperCAmelCase			     =  backbone(**a__  )

                     # Test default outputs and verify feature maps
                     self.assertIsInstance(outputs.feature_maps     ,	a__  )
                     self.assertTrue(len(outputs.feature_maps  ) == len(backbone.channels  )  )
                     for feature_map, n_channels in zip(outputs.feature_maps     ,	backbone.channels  ):
                            self.assertTrue(feature_map.shape[:2]     ,	(batch_size, n_channels)  )
                     self.assertIsNone(outputs.hidden_states  )
                     self.assertIsNone(outputs.attentions  )

                     # Test output_hidden_states=True
                     _UpperCAmelCase			     =  backbone(**a__     ,	output_hidden_states=a__  )
                     self.assertIsNotNone(outputs.hidden_states  )
                     self.assertTrue(len(outputs.hidden_states  )     ,	len(backbone.stage_names  )  )
                     for hidden_state, n_channels in zip(outputs.hidden_states     ,	backbone.channels  ):
                            self.assertTrue(hidden_state.shape[:2]     ,	(batch_size, n_channels)  )

                     # Test output_attentions=True
                     if self.has_attentions:
                            _UpperCAmelCase			     =  backbone(**a__     ,	output_attentions=a__  )
                            self.assertIsNotNone(outputs.attentions  )






