





'''simple docstring'''






import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()




class  UpperCAmelCase     (			snake_case_							,		snake_case_							,		unittest.TestCase   ):
      _lowercase:			Optional[Any]						     =   AutoencoderKL
      _lowercase:			List[Any]						     =   '''sample'''
      _lowercase:			List[str]						     =   1E-2




      @property
      def 				lowercase__    (			self			: Union[str, Any]  )				->      str:
            _lowerCAmelCase							     =							4
            _lowerCAmelCase							     =							3
            _lowerCAmelCase							     =							(32, 32)

            _lowerCAmelCase							     =							floats_tensor((batch_size, num_channels) + sizes  ).to(__snake_case  )

            return {"sample": image}




      @property
      def 				lowercase__    (			self			: Tuple  )				->      Any:
            return (3, 32, 32)




      @property
      def 				lowercase__    (			self			: Union[str, Any]  )				->      Union[str, Any]:
            return (3, 32, 32)




      def 				lowercase__    (			self			: Tuple  )				->      Dict:
            _lowerCAmelCase							     =							{
                """block_out_channels""": [32, 64],
                """in_channels""": 3,
                """out_channels""": 3,
                """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
                """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
                """latent_channels""": 4,
            }
            _lowerCAmelCase							     =							self.dummy_input
            return init_dict, inputs_dict




      def 				lowercase__    (			self			: Union[str, Any]  )				->      int:
            pass




      def 				lowercase__    (			self			: List[Any]  )				->      Any:
            pass




      @unittest.skipIf(torch_device == """mps"""			,			"""Gradient checkpointing skipped on MPS"""  )
      def 				lowercase__    (			self			: int  )				->      Dict:
            # enable deterministic behavior for gradient checkpointing
            _lowerCAmelCase     ,							_lowerCAmelCase							     =							self.prepare_init_args_and_inputs_for_common()
            _lowerCAmelCase							     =							self.model_class(**__snake_case  )
            model.to(__snake_case  )

            assert not model.is_gradient_checkpointing and model.training

            _lowerCAmelCase							     =							model(**__snake_case  ).sample
            # run the backwards pass on the model. For backwards pass, for simplicity purpose,
            # we won't calculate the loss and rather backprop on out.sum()
            model.zero_grad()

            _lowerCAmelCase							     =							torch.randn_like(__snake_case  )
            _lowerCAmelCase							     =							(out - labels).mean()
            loss.backward()

            # re-instantiate the model now enabling gradient checkpointing
            _lowerCAmelCase							     =							self.model_class(**__snake_case  )
            # clone model
            model_a.load_state_dict(model.state_dict()  )
            model_a.to(__snake_case  )
            model_a.enable_gradient_checkpointing()

            assert model_a.is_gradient_checkpointing and model_a.training

            _lowerCAmelCase							     =							model_a(**__snake_case  ).sample
            # run the backwards pass on the model. For backwards pass, for simplicity purpose,
            # we won't calculate the loss and rather backprop on out.sum()
            model_a.zero_grad()
            _lowerCAmelCase							     =							(out_a - labels).mean()
            loss_a.backward()

            # compare the output and parameters gradients
            self.assertTrue((loss - loss_a).abs() < 1E-5  )
            _lowerCAmelCase							     =							dict(model.named_parameters()  )
            _lowerCAmelCase							     =							dict(model_a.named_parameters()  )
            for name, param in named_params.items():
                  self.assertTrue(torch_all_close(param.grad.data			,			named_params_a[name].grad.data			,			atol=5E-5  )  )




      def 				lowercase__    (			self			: int  )				->      Optional[Any]:
            _lowerCAmelCase     ,							_lowerCAmelCase							     =							AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""			,			output_loading_info=__snake_case  )
            self.assertIsNotNone(__snake_case  )
            self.assertEqual(len(loading_info["""missing_keys"""]  )			,			0  )

            model.to(__snake_case  )
            _lowerCAmelCase							     =							model(**self.dummy_input  )

            assert image is not None, "Make sure output is not None"




      def 				lowercase__    (			self			: Any  )				->      int:
            _lowerCAmelCase							     =							AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""  )
            _lowerCAmelCase							     =							model.to(__snake_case  )
            model.eval()

            if torch_device == "mps":
                  _lowerCAmelCase							     =							torch.manual_seed(0  )
            else:
                  _lowerCAmelCase							     =							torch.Generator(device=__snake_case  ).manual_seed(0  )

            _lowerCAmelCase							     =							torch.randn(
                1			,			model.config.in_channels			,			model.config.sample_size			,			model.config.sample_size			,			generator=torch.manual_seed(0  )			,			)
            _lowerCAmelCase							     =							image.to(__snake_case  )
            with torch.no_grad():
                  _lowerCAmelCase							     =							model(__snake_case			,			sample_posterior=__snake_case			,			generator=__snake_case  ).sample

            _lowerCAmelCase							     =							output[0, -1, -3:, -3:].flatten().cpu()

            # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
            # the expected output slices are not the same for CPU and GPU.
            if torch_device == "mps":
                  _lowerCAmelCase							     =							torch.tensor(
                      [
                          -4.0_0_7_8E-0_1,
                          -3.8_3_2_3E-0_4,
                          -1.2_6_8_1E-0_1,
                          -1.1_4_6_2E-0_1,
                          2.0_0_9_5E-0_1,
                          1.0_8_9_3E-0_1,
                          -8.8_2_4_7E-0_2,
                          -3.0_3_6_1E-0_1,
                          -9.8_6_4_4E-0_3,
                      ]  )
            elif torch_device == "cpu":
                  _lowerCAmelCase							     =							torch.tensor(
                      [-0.13_52, 0.08_78, 0.04_19, -0.08_18, -0.10_69, 0.06_88, -0.14_58, -0.44_46, -0.00_26]  )
            else:
                  _lowerCAmelCase							     =							torch.tensor(
                      [-0.24_21, 0.46_42, 0.25_07, -0.04_38, 0.06_82, 0.31_60, -0.20_18, -0.07_27, 0.24_85]  )

            self.assertTrue(torch_all_close(__snake_case			,			__snake_case			,			rtol=1E-2  )  )




@slow
class  UpperCAmelCase     (			unittest.TestCase   ):




      def 				lowercase__    (			self			: List[str]			,			__snake_case			: int			,			__snake_case			: Any  )				->      Union[str, Any]:
            return f"gaussian_noise_s={seed}_shape={'_'.join([str(__snake_case  ) for s in shape]  )}.npy"




      def 				lowercase__    (			self			: List[str]  )				->      Any:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def 				lowercase__    (			self			: Dict			,			__snake_case			: Optional[int]=0			,			__snake_case			: Optional[Any]=(4, 3, 5_12, 5_12)			,			__snake_case			: str=False  )				->      List[Any]:
            _lowerCAmelCase							     =							torch.floataa if fpaa else torch.floataa
            _lowerCAmelCase							     =							torch.from_numpy(load_hf_numpy(self.get_file_format(__snake_case			,			__snake_case  )  )  ).to(__snake_case  ).to(__snake_case  )
            return image




      def 				lowercase__    (			self			: str			,			__snake_case			: Union[str, Any]="CompVis/stable-diffusion-v1-4"			,			__snake_case			: Any=False  )				->      Union[str, Any]:
            _lowerCAmelCase							     =							"""fp16""" if fpaa else None
            _lowerCAmelCase							     =							torch.floataa if fpaa else torch.floataa

            _lowerCAmelCase							     =							AutoencoderKL.from_pretrained(
                __snake_case			,			subfolder="""vae"""			,			torch_dtype=__snake_case			,			revision=__snake_case			,			)
            model.to(__snake_case  ).eval()

            return model




      def 				lowercase__    (			self			: List[str]			,			__snake_case			: Optional[int]=0  )				->      List[str]:
            if torch_device == "mps":
                  return torch.manual_seed(__snake_case  )
            return torch.Generator(device=__snake_case  ).manual_seed(__snake_case  )




      @parameterized.expand(
          [
              # fmt: off
              [33, [-0.16_03, 0.98_78, -0.04_95, -0.07_90, -0.27_09, 0.83_75, -0.20_60, -0.08_24], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
              [47, [-0.23_76, 0.11_68, 0.13_32, -0.48_40, -0.25_08, -0.07_91, -0.04_93, -0.40_89], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
              # fmt: on
          ]  )
      def 				lowercase__    (			self			: Tuple			,			__snake_case			: Dict			,			__snake_case			: str			,			__snake_case			: List[str]  )				->      int:
            _lowerCAmelCase							     =							self.get_sd_vae_model()
            _lowerCAmelCase							     =							self.get_sd_image(__snake_case  )
            _lowerCAmelCase							     =							self.get_generator(__snake_case  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model(__snake_case			,			generator=__snake_case			,			sample_posterior=__snake_case  ).sample

            assert sample.shape == image.shape

            _lowerCAmelCase							     =							sample[-1, -2:, -2:, :2].flatten().float().cpu()
            _lowerCAmelCase							     =							torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice  )

            assert torch_all_close(__snake_case			,			__snake_case			,			atol=3E-3  )




      @parameterized.expand(
          [
              # fmt: off
              [33, [-0.05_13, 0.02_89, 1.37_99, 0.21_66, -0.25_73, -0.08_71, 0.51_03, -0.09_99]],
              [47, [-0.41_28, -0.13_20, -0.37_04, 0.19_65, -0.41_16, -0.23_32, -0.33_40, 0.22_47]],
              # fmt: on
          ]  )
      @require_torch_gpu
      def 				lowercase__    (			self			: List[Any]			,			__snake_case			: str			,			__snake_case			: List[Any]  )				->      Dict:
            _lowerCAmelCase							     =							self.get_sd_vae_model(fpaa=__snake_case  )
            _lowerCAmelCase							     =							self.get_sd_image(__snake_case			,			fpaa=__snake_case  )
            _lowerCAmelCase							     =							self.get_generator(__snake_case  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model(__snake_case			,			generator=__snake_case			,			sample_posterior=__snake_case  ).sample

            assert sample.shape == image.shape

            _lowerCAmelCase							     =							sample[-1, -2:, :2, -2:].flatten().float().cpu()
            _lowerCAmelCase							     =							torch.tensor(__snake_case  )

            assert torch_all_close(__snake_case			,			__snake_case			,			atol=1E-2  )




      @parameterized.expand(
          [
              # fmt: off
              [33, [-0.16_09, 0.98_66, -0.04_87, -0.07_77, -0.27_16, 0.83_68, -0.20_55, -0.08_14], [-0.23_95, 0.00_98, 0.01_02, -0.07_09, -0.28_40, -0.02_74, -0.07_18, -0.18_24]],
              [47, [-0.23_77, 0.11_47, 0.13_33, -0.48_41, -0.25_06, -0.08_05, -0.04_91, -0.40_85], [0.03_50, 0.08_47, 0.04_67, 0.03_44, -0.08_42, -0.05_47, -0.06_33, -0.11_31]],
              # fmt: on
          ]  )
      def 				lowercase__    (			self			: Tuple			,			__snake_case			: Union[str, Any]			,			__snake_case			: List[str]			,			__snake_case			: Dict  )				->      Dict:
            _lowerCAmelCase							     =							self.get_sd_vae_model()
            _lowerCAmelCase							     =							self.get_sd_image(__snake_case  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model(__snake_case  ).sample

            assert sample.shape == image.shape

            _lowerCAmelCase							     =							sample[-1, -2:, -2:, :2].flatten().float().cpu()
            _lowerCAmelCase							     =							torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice  )

            assert torch_all_close(__snake_case			,			__snake_case			,			atol=3E-3  )




      @parameterized.expand(
          [
              # fmt: off
              [13, [-0.20_51, -0.18_03, -0.23_11, -0.21_14, -0.32_92, -0.35_74, -0.29_53, -0.33_23]],
              [37, [-0.26_32, -0.26_25, -0.21_99, -0.27_41, -0.45_39, -0.49_90, -0.37_20, -0.49_25]],
              # fmt: on
          ]  )
      @require_torch_gpu
      def 				lowercase__    (			self			: Tuple			,			__snake_case			: Any			,			__snake_case			: Tuple  )				->      str:
            _lowerCAmelCase							     =							self.get_sd_vae_model()
            _lowerCAmelCase							     =							self.get_sd_image(__snake_case			,			shape=(3, 4, 64, 64)  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model.decode(__snake_case  ).sample

            assert list(sample.shape  ) == [3, 3, 5_12, 5_12]

            _lowerCAmelCase							     =							sample[-1, -2:, :2, -2:].flatten().cpu()
            _lowerCAmelCase							     =							torch.tensor(__snake_case  )

            assert torch_all_close(__snake_case			,			__snake_case			,			atol=1E-3  )




      @parameterized.expand(
          [
              # fmt: off
              [27, [-0.03_69, 0.02_07, -0.07_76, -0.06_82, -0.17_47, -0.19_30, -0.14_65, -0.20_39]],
              [16, [-0.16_28, -0.21_34, -0.27_47, -0.26_42, -0.37_74, -0.44_04, -0.36_87, -0.42_77]],
              # fmt: on
          ]  )
      @require_torch_gpu
      def 				lowercase__    (			self			: Any			,			__snake_case			: str			,			__snake_case			: str  )				->      int:
            _lowerCAmelCase							     =							self.get_sd_vae_model(fpaa=__snake_case  )
            _lowerCAmelCase							     =							self.get_sd_image(__snake_case			,			shape=(3, 4, 64, 64)			,			fpaa=__snake_case  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model.decode(__snake_case  ).sample

            assert list(sample.shape  ) == [3, 3, 5_12, 5_12]

            _lowerCAmelCase							     =							sample[-1, -2:, :2, -2:].flatten().float().cpu()
            _lowerCAmelCase							     =							torch.tensor(__snake_case  )

            assert torch_all_close(__snake_case			,			__snake_case			,			atol=5E-3  )




      @parameterized.expand([(13,), (16,), (27,)]  )
      @require_torch_gpu
      @unittest.skipIf(not is_xformers_available()			,			reason="""xformers is not required when using PyTorch 2.0."""  )
      def 				lowercase__    (			self			: Dict			,			__snake_case			: Tuple  )				->      List[str]:
            _lowerCAmelCase							     =							self.get_sd_vae_model(fpaa=__snake_case  )
            _lowerCAmelCase							     =							self.get_sd_image(__snake_case			,			shape=(3, 4, 64, 64)			,			fpaa=__snake_case  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model.decode(__snake_case  ).sample

            model.enable_xformers_memory_efficient_attention()
            with torch.no_grad():
                  _lowerCAmelCase							     =							model.decode(__snake_case  ).sample

            assert list(sample.shape  ) == [3, 3, 5_12, 5_12]

            assert torch_all_close(__snake_case			,			__snake_case			,			atol=1E-1  )




      @parameterized.expand([(13,), (16,), (37,)]  )
      @require_torch_gpu
      @unittest.skipIf(not is_xformers_available()			,			reason="""xformers is not required when using PyTorch 2.0."""  )
      def 				lowercase__    (			self			: Optional[int]			,			__snake_case			: Union[str, Any]  )				->      str:
            _lowerCAmelCase							     =							self.get_sd_vae_model()
            _lowerCAmelCase							     =							self.get_sd_image(__snake_case			,			shape=(3, 4, 64, 64)  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model.decode(__snake_case  ).sample

            model.enable_xformers_memory_efficient_attention()
            with torch.no_grad():
                  _lowerCAmelCase							     =							model.decode(__snake_case  ).sample

            assert list(sample.shape  ) == [3, 3, 5_12, 5_12]

            assert torch_all_close(__snake_case			,			__snake_case			,			atol=1E-2  )




      @parameterized.expand(
          [
              # fmt: off
              [33, [-0.30_01, 0.09_18, -2.69_84, -3.97_20, -3.20_99, -5.03_53, 1.73_38, -0.20_65, 3.42_67]],
              [47, [-1.50_30, -4.38_71, -6.03_55, -9.11_57, -1.66_61, -2.78_53, 2.16_07, -5.08_23, 2.56_33]],
              # fmt: on
          ]  )
      def 				lowercase__    (			self			: Tuple			,			__snake_case			: Tuple			,			__snake_case			: List[str]  )				->      Union[str, Any]:
            _lowerCAmelCase							     =							self.get_sd_vae_model()
            _lowerCAmelCase							     =							self.get_sd_image(__snake_case  )
            _lowerCAmelCase							     =							self.get_generator(__snake_case  )

            with torch.no_grad():
                  _lowerCAmelCase							     =							model.encode(__snake_case  ).latent_dist
                  _lowerCAmelCase							     =							dist.sample(generator=__snake_case  )

            assert list(sample.shape  ) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

            _lowerCAmelCase							     =							sample[0, -1, -3:, -3:].flatten().cpu()
            _lowerCAmelCase							     =							torch.tensor(__snake_case  )

            _lowerCAmelCase							     =							3E-3 if torch_device != """mps""" else 1E-2
            assert torch_all_close(__snake_case			,			__snake_case			,			atol=__snake_case  )










import argparse

import intel_extension_for_pytorch as ipex
import torch

from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline


_snake_case             =     argparse.ArgumentParser('''Stable Diffusion script with intel optimization''', add_help=False)
parser.add_argument('''--dpm''', action='''store_true''', help='''Enable DPMSolver or not''')
parser.add_argument('''--steps''', default=None, type=int, help='''Num inference steps''')
_snake_case             =     parser.parse_args()


_snake_case             =     '''cpu'''
_snake_case             =     '''a lovely <dicoo> in red dress and hat, in the snowly and brightly night, with many brighly buildings'''

_snake_case             =     '''path-to-your-trained-model'''
_snake_case             =     StableDiffusionPipeline.from_pretrained(model_id)
if args.dpm:
   _snake_case             =     DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
_snake_case             =     pipe.to(device)

# to channels last
_snake_case             =     pipe.unet.to(memory_format=torch.channels_last)
_snake_case             =     pipe.vae.to(memory_format=torch.channels_last)
_snake_case             =     pipe.text_encoder.to(memory_format=torch.channels_last)
if pipe.requires_safety_checker:
   _snake_case             =     pipe.safety_checker.to(memory_format=torch.channels_last)

# optimize with ipex
_snake_case             =     torch.randn(2, 4, 64, 64)
_snake_case             =     torch.rand(1) * 9_99
_snake_case             =     torch.randn(2, 77, 7_68)
_snake_case             =     (sample, timestep, encoder_hidden_status)
try:
   _snake_case             =     ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True, sample_input=input_example)
except Exception:
   _snake_case             =     ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True)
_snake_case             =     ipex.optimize(pipe.vae.eval(), dtype=torch.bfloataa, inplace=True)
_snake_case             =     ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloataa, inplace=True)
if pipe.requires_safety_checker:
   _snake_case             =     ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloataa, inplace=True)

# compute
_snake_case             =     6_66
_snake_case             =     torch.Generator(device).manual_seed(seed)
_snake_case             =     {'''generator''': generator}
if args.steps is not None:
   _snake_case             =     args.steps

with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloataa):
   _snake_case             =     pipe(prompt, **generate_kwargs).images[0]

# save image
image.save('''generated.png''')



