




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS


enable_full_determinism()





class 	lowerCamelCase_							(				unittest.TestCase ):


 '''simple docstring'''







 lowercase_       =   StableDiffusionLDMaDPipeline
 lowercase_       =   TEXT_TO_IMAGE_PARAMS
 lowercase_       =   TEXT_TO_IMAGE_BATCH_PARAMS
 lowercase_       =   TEXT_TO_IMAGE_IMAGE_PARAMS



 def 				lowerCAmelCase_  (							self  :		Any						):
       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_ :    List[str]     						=      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    ,    )
       SCREAMING_SNAKE_CASE_ :    str     						=      DDIMScheduler(
           beta_start=0.0_0085    ,    beta_end=0.012    ,    beta_schedule='scaled_linear'    ,    clip_sample=_SCREAMING_SNAKE_CASE    ,    set_alpha_to_one=_SCREAMING_SNAKE_CASE    ,    )
       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      AutoencoderKL(
           block_out_channels=[32, 64]    ,    in_channels=6    ,    out_channels=6    ,    down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']    ,    up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']    ,    latent_channels=4    ,    )
       torch.manual_seed(0						)
       SCREAMING_SNAKE_CASE_ :    Tuple     						=      CLIPTextConfig(
           bos_token_id=0    ,    eos_token_id=2    ,    hidden_size=32    ,    intermediate_size=37    ,    layer_norm_eps=1E-05    ,    num_attention_heads=4    ,    num_hidden_layers=5    ,    pad_token_id=1    ,    vocab_size=1_000    ,    )
       SCREAMING_SNAKE_CASE_ :    str     						=      CLIPTextModel(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Dict     						=      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

       SCREAMING_SNAKE_CASE_ :    Optional[int]     						=      {
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       return components



 def 				lowerCAmelCase_  (							self  :		str    ,    _lowerCAmelCase  :		List[str]    ,    _lowerCAmelCase  :		Dict=0						):





       if str(_SCREAMING_SNAKE_CASE						).startswith('mps'						):
             SCREAMING_SNAKE_CASE_ :    Tuple     						=      torch.manual_seed(_SCREAMING_SNAKE_CASE						)
       else:
             SCREAMING_SNAKE_CASE_ :    List[Any]     						=      torch.Generator(device=_SCREAMING_SNAKE_CASE						).manual_seed(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }
       return inputs



 def 				lowerCAmelCase_  (							self  :		str						):
       SCREAMING_SNAKE_CASE_ :    int     						=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator

       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      self.get_dummy_components()
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      StableDiffusionLDMaDPipeline(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Dict     						=      ldmad_pipe.to(_SCREAMING_SNAKE_CASE						)
       ldmad_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

       SCREAMING_SNAKE_CASE_ :    Optional[int]     						=      self.get_dummy_inputs(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Tuple     						=      ldmad_pipe(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      output.rgb, output.depth

       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      rgb[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      depth[0, -3:, -1]

       assert rgb.shape == (1, 64, 64, 3)
       assert depth.shape == (1, 64, 64)

       SCREAMING_SNAKE_CASE_ :    Tuple     						=      np.array(
           [0.3733_8176, 0.7_0247, 0.7420_3193, 0.5164_3604, 0.5825_6793, 0.6093_2136, 0.418_1095, 0.4835_5877, 0.4653_5262]						)
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      np.array([103.46_727, 85.812_004, 87.849_236]						)

       assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb						).max() < 1E-2
       assert np.abs(image_slice_depth.flatten() - expected_slice_depth						).max() < 1E-2



 def 				lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      self.get_dummy_components()
       SCREAMING_SNAKE_CASE_ :    Any     						=      StableDiffusionLDMaDPipeline(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    List[str]     						=      ldmad_pipe.to(_SCREAMING_SNAKE_CASE						)
       ldmad_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      self.get_dummy_inputs(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    str     						=      3 * [inputs['''prompt''']]

       # forward
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      ldmad_pipe(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Optional[int]     						=      output.rgb, output.depth
       SCREAMING_SNAKE_CASE_ :    Tuple     						=      rgb_slice_a[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      depth_slice_a[0, -3:, -1]

       SCREAMING_SNAKE_CASE_ :    Tuple     						=      self.get_dummy_inputs(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Dict     						=      3 * [inputs.pop('prompt'						)]

       SCREAMING_SNAKE_CASE_ :    Tuple     						=      ldmad_pipe.tokenizer(
           _SCREAMING_SNAKE_CASE    ,    padding='max_length'    ,    max_length=ldmad_pipe.tokenizer.model_max_length    ,    truncation=_SCREAMING_SNAKE_CASE    ,    return_tensors='pt'    ,    )
       SCREAMING_SNAKE_CASE_ :    Dict     						=      text_inputs['''input_ids'''].to(_SCREAMING_SNAKE_CASE						)

       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      ldmad_pipe.text_encoder(_SCREAMING_SNAKE_CASE						)[0]

       SCREAMING_SNAKE_CASE_ :    Optional[int]     						=      prompt_embeds

       # forward
       SCREAMING_SNAKE_CASE_ :    Optional[int]     						=      ldmad_pipe(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Any     						=      output.rgb, output.depth
       SCREAMING_SNAKE_CASE_ :    Any     						=      rgb_slice_a[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE_ :    str     						=      depth_slice_a[0, -3:, -1]

       assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()						).max() < 1E-4
       assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()						).max() < 1E-4




 def 				lowerCAmelCase_  (							self  :		str						):
       SCREAMING_SNAKE_CASE_ :    Tuple     						=      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       SCREAMING_SNAKE_CASE_ :    List[str]     						=      self.get_dummy_components()
       SCREAMING_SNAKE_CASE_ :    int     						=      PNDMScheduler(skip_prk_steps=_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    int     						=      StableDiffusionLDMaDPipeline(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    str     						=      ldmad_pipe.to(_SCREAMING_SNAKE_CASE						)
       ldmad_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

       SCREAMING_SNAKE_CASE_ :    str     						=      self.get_dummy_inputs(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      '''french fries'''
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      ldmad_pipe(**_SCREAMING_SNAKE_CASE    ,    negative_prompt=_SCREAMING_SNAKE_CASE						)

       SCREAMING_SNAKE_CASE_ :    Optional[int]     						=      output.rgb, output.depth
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      rgb[0, -3:, -3:, -1]
       SCREAMING_SNAKE_CASE_ :    int     						=      depth[0, -3:, -1]

       assert rgb.shape == (1, 64, 64, 3)
       assert depth.shape == (1, 64, 64)

       SCREAMING_SNAKE_CASE_ :    int     						=      np.array(
           [0.3_7044, 0.7181_1503, 0.722_3251, 0.4860_3675, 0.563_8391, 0.636_4948, 0.4283_3704, 0.490_1315, 0.4792_6217]						)
       SCREAMING_SNAKE_CASE_ :    Any     						=      np.array([107.84_738, 84.62_802, 89.962_135]						)
       assert np.abs(rgb_slice.flatten() - expected_slice_rgb						).max() < 1E-2
       assert np.abs(depth_slice.flatten() - expected_slice_depth						).max() < 1E-2





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


 '''simple docstring'''







 def 				lowerCAmelCase_  (							self  :		Optional[Any]						):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 def 				lowerCAmelCase_  (							self  :		Union[str, Any]    ,    _lowerCAmelCase  :		List[Any]    ,    _lowerCAmelCase  :		List[str]="cpu"    ,    _lowerCAmelCase  :		Tuple=torch.floataa    ,    _lowerCAmelCase  :		Optional[Any]=0						):
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      torch.Generator(device=_SCREAMING_SNAKE_CASE						).manual_seed(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    str     						=      np.random.RandomState(_SCREAMING_SNAKE_CASE						).standard_normal((1, 4, 64, 64)						)
       SCREAMING_SNAKE_CASE_ :    int     						=      torch.from_numpy(_SCREAMING_SNAKE_CASE						).to(device=_SCREAMING_SNAKE_CASE    ,    dtype=_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    int     						=      {
           '''prompt''': '''a photograph of an astronaut riding a horse''',
           '''latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs




 def 				lowerCAmelCase_  (							self  :		str						):
       SCREAMING_SNAKE_CASE_ :    List[str]     						=      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'						)
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      ldmad_pipe.to(_SCREAMING_SNAKE_CASE						)
       ldmad_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

       SCREAMING_SNAKE_CASE_ :    int     						=      self.get_inputs(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Dict     						=      ldmad_pipe(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      output.rgb, output.depth
       SCREAMING_SNAKE_CASE_ :    int     						=      rgb[0, -3:, -3:, -1].flatten()
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      rgb[0, -3:, -1].flatten()

       assert rgb.shape == (1, 512, 512, 3)
       assert depth.shape == (1, 512, 512)

       SCREAMING_SNAKE_CASE_ :    Tuple     						=      np.array(
           [0.5380_5465, 0.5670_7305, 0.548_6515, 0.5701_2236, 0.581_4511, 0.5625_3487, 0.5484_3014, 0.5509_2263, 0.645_9706]						)
       SCREAMING_SNAKE_CASE_ :    Tuple     						=      np.array(
           [0.926_3781, 0.667_8672, 0.548_6515, 0.9220_2145, 0.6783_1135, 0.5625_3487, 0.924_1694, 0.755_1478, 0.645_9706]						)
       assert np.abs(rgb_slice - expected_slice_rgb						).max() < 3E-3
       assert np.abs(depth_slice - expected_slice_depth						).max() < 3E-3





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


 '''simple docstring'''







 def 				lowerCAmelCase_  (							self  :		Optional[int]						):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 def 				lowerCAmelCase_  (							self  :		List[Any]    ,    _lowerCAmelCase  :		str    ,    _lowerCAmelCase  :		Any="cpu"    ,    _lowerCAmelCase  :		Optional[Any]=torch.floataa    ,    _lowerCAmelCase  :		Tuple=0						):
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      torch.Generator(device=_SCREAMING_SNAKE_CASE						).manual_seed(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Any     						=      np.random.RandomState(_SCREAMING_SNAKE_CASE						).standard_normal((1, 4, 64, 64)						)
       SCREAMING_SNAKE_CASE_ :    str     						=      torch.from_numpy(_SCREAMING_SNAKE_CASE						).to(device=_SCREAMING_SNAKE_CASE    ,    dtype=_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      {
           '''prompt''': '''a photograph of an astronaut riding a horse''',
           '''latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 50,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs



 def 				lowerCAmelCase_  (							self  :		Dict						):
       SCREAMING_SNAKE_CASE_ :    Tuple     						=      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'						).to(_SCREAMING_SNAKE_CASE						)
       ldmad_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      self.get_inputs(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    str     						=      ldmad_pipe(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      output.rgb, output.depth

       SCREAMING_SNAKE_CASE_ :    int     						=      0.49_5586
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      0.3379_5515
       SCREAMING_SNAKE_CASE_ :    Optional[int]     						=      112.48_518
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      98.489_746
       assert np.abs(expected_rgb_mean - rgb.mean()						) < 1E-3
       assert np.abs(expected_rgb_std - rgb.std()						) < 1E-3
       assert np.abs(expected_depth_mean - depth.mean()						) < 1E-3
       assert np.abs(expected_depth_std - depth.std()						) < 1E-3




 def 				lowerCAmelCase_  (							self  :		List[Any]						):
       SCREAMING_SNAKE_CASE_ :    Any     						=      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'						).to(_SCREAMING_SNAKE_CASE						)
       ldmad_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)

       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      self.get_inputs(_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Tuple     						=      ldmad_pipe(**_SCREAMING_SNAKE_CASE						)
       SCREAMING_SNAKE_CASE_ :    Optional[Any]     						=      output.rgb, output.depth

       SCREAMING_SNAKE_CASE_ :    Tuple     						=      0.419_4127
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      0.3537_5586
       SCREAMING_SNAKE_CASE_ :    Union[str, Any]     						=      0.563_8502
       SCREAMING_SNAKE_CASE_ :    List[Any]     						=      0.3468_6103

       assert rgb.shape == (1, 512, 512, 3)
       assert depth.shape == (1, 512, 512, 1)
       assert np.abs(expected_rgb_mean - rgb.mean()						) < 1E-3
       assert np.abs(expected_rgb_std - rgb.std()						) < 1E-3
       assert np.abs(expected_depth_mean - depth.mean()						) < 1E-3
       assert np.abs(expected_depth_std - depth.std()						) < 1E-3





import gc
import threading
import time

import psutil
import torch





class 	lowerCamelCase_							:


 '''simple docstring'''







 def __init__(							self  :		Optional[Any]						):
       SCREAMING_SNAKE_CASE_     						=      psutil.Process()
       SCREAMING_SNAKE_CASE_     						=      False



 def 				lowerCAmelCase_  (							self  :		Dict						):
       SCREAMING_SNAKE_CASE_     						=      -1

       while True:
             SCREAMING_SNAKE_CASE_     						=      max(self.process.memory_info().rss    ,    self.cpu_memory_peak						)

             # can't sleep or will not catch the peak right (this comment is here on purpose)
             if not self.peak_monitoring:
                   break



 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      True
       SCREAMING_SNAKE_CASE_     						=      threading.Thread(target=self.peak_monitor						)
       SCREAMING_SNAKE_CASE_     						=      True
       self.thread.start()




 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      False
       self.thread.join()
       return self.cpu_memory_peak


lowerCamelCase__			:       List[str]         =							PeakCPUMemory()

def  UpperCAmelCase_							(				) ->			Tuple:
      # Time
      SCREAMING_SNAKE_CASE_     						=      {'time': time.time()}

      gc.collect()
      torch.cuda.empty_cache()

      # CPU mem
      SCREAMING_SNAKE_CASE_     						=      psutil.Process().memory_info().rss
      cpu_peak_tracker.start()

      # GPU mem
      for i in range(torch.cuda.device_count()							):
            SCREAMING_SNAKE_CASE_     						=      torch.cuda.memory_allocated(__UpperCAmelCase							)
      torch.cuda.reset_peak_memory_stats()

      return measures

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]							) ->			Optional[Any]:
      # Time
      SCREAMING_SNAKE_CASE_     						=      {'time': time.time() - start_measures['time']}

      gc.collect()
      torch.cuda.empty_cache()

      # CPU mem
      SCREAMING_SNAKE_CASE_     						=      (psutil.Process().memory_info().rss - start_measures['cpu']) / 2**20
      SCREAMING_SNAKE_CASE_     						=      (cpu_peak_tracker.stop() - start_measures['cpu']) / 2**20

      # GPU mem
      for i in range(torch.cuda.device_count()							):
            SCREAMING_SNAKE_CASE_     						=      (torch.cuda.memory_allocated(__UpperCAmelCase							) - start_measures[str(__UpperCAmelCase							)]) / 2**20
            SCREAMING_SNAKE_CASE_     						=      (torch.cuda.max_memory_allocated(__UpperCAmelCase							) - start_measures[str(__UpperCAmelCase							)]) / 2**20

      return measures





def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[int]      ,      __UpperCAmelCase						:     Any							) ->			Optional[Any]:
      print(f"{description}:"							)
      print(f"- Time: {measures['time']:.2f}s"							)
      for i in range(torch.cuda.device_count()							):
            print(f"- GPU {i} allocated: {measures[str(__UpperCAmelCase							)]:.2f}MiB"							)
            SCREAMING_SNAKE_CASE_     						=      measures[f"{i}-peak"]
            print(f"- GPU {i} peak: {peak:.2f}MiB"							)
      print(f"- CPU RAM allocated: {measures['cpu']:.2f}MiB"							)
      print(f"- CPU RAM peak: {measures['cpu-peak']:.2f}MiB"							)