






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 							__UpperCamelCase     (  unittest.TestCase ):
	__snake_case						:Union[str, Any] =							StableDiffusionLDMaDPipeline
	__snake_case						:List[str] =							TEXT_TO_IMAGE_PARAMS
	__snake_case						:List[str] =							TEXT_TO_IMAGE_BATCH_PARAMS
	__snake_case						:str =							TEXT_TO_IMAGE_IMAGE_PARAMS







	def 			_a (	self   :				Tuple      )		->			Optional[int]:
		"""simple docstring"""







		torch.manual_seed(0      )
		__lowercase   							=	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							,					)
		__lowercase   							=	DDIMScheduler(
		    beta_start=0.00_085							,					beta_end=0.012							,					beta_schedule="""scaled_linear"""							,					clip_sample=_lowerCAmelCase							,					set_alpha_to_one=_lowerCAmelCase							,					)
		torch.manual_seed(0      )
		__lowercase   							=	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      )
		__lowercase   							=	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=1000							,					)
		__lowercase   							=	CLIPTextModel(_lowerCAmelCase      )
		__lowercase   							=	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )

		__lowercase   							=	{
		    """unet""": unet,
		    """scheduler""": scheduler,
		    """vae""": vae,
		    """text_encoder""": text_encoder,
		    """tokenizer""": tokenizer,
		    """safety_checker""": None,
		    """feature_extractor""": None,
		}
		return components







	def 			_a (	self   :				str							,					_lowerCAmelCase   :				Any							,					_lowerCAmelCase   :				Dict=0      )		->			List[Any]:
		"""simple docstring"""







		if str(_lowerCAmelCase      ).startswith("""mps"""      ):
			__lowercase   							=	torch.manual_seed(_lowerCAmelCase      )
		else:
			__lowercase   							=	torch.Generator(device=_lowerCAmelCase      ).manual_seed(_lowerCAmelCase      )
		__lowercase   							=	{
		    """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 			_a (	self   :				List[Any]      )		->			List[Any]:
		"""simple docstring"""







		__lowercase   							=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator

		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	StableDiffusionLDMaDPipeline(**_lowerCAmelCase      )
		__lowercase   							=	ldmad_pipe.to(_lowerCAmelCase      )
		ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	ldmad_pipe(**_lowerCAmelCase      )
		__lowercase							,    __lowercase   							=	output.rgb, output.depth

		__lowercase   							=	rgb[0, -3:, -3:, -1]
		__lowercase   							=	depth[0, -3:, -1]

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

		__lowercase   							=	np.array(
		    [0.37_338_176, 0.70_247, 0.74_203_193, 0.51_643_604, 0.58_256_793, 0.60_932_136, 0.4_181_095, 0.48_355_877, 0.46_535_262]      )
		__lowercase   							=	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 			_a (	self   :				Tuple      )		->			Optional[Any]:
		"""simple docstring"""







		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	StableDiffusionLDMaDPipeline(**_lowerCAmelCase      )
		__lowercase   							=	ldmad_pipe.to(_lowerCAmelCase      )
		ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	3 * [inputs["""prompt"""]]

		# forward
		__lowercase   							=	ldmad_pipe(**_lowerCAmelCase      )
		__lowercase							,    __lowercase   							=	output.rgb, output.depth
		__lowercase   							=	rgb_slice_a[0, -3:, -3:, -1]
		__lowercase   							=	depth_slice_a[0, -3:, -1]

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	3 * [inputs.pop("""prompt"""      )]

		__lowercase   							=	ldmad_pipe.tokenizer(
		    _lowerCAmelCase							,					padding="""max_length"""							,					max_length=ldmad_pipe.tokenizer.model_max_length							,					truncation=_lowerCAmelCase							,					return_tensors="""pt"""							,					)
		__lowercase   							=	text_inputs["""input_ids"""].to(_lowerCAmelCase      )

		__lowercase   							=	ldmad_pipe.text_encoder(_lowerCAmelCase      )[0]

		__lowercase   							=	prompt_embeds

		# forward
		__lowercase   							=	ldmad_pipe(**_lowerCAmelCase      )
		__lowercase							,    __lowercase   							=	output.rgb, output.depth
		__lowercase   							=	rgb_slice_a[0, -3:, -3:, -1]
		__lowercase   							=	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 			_a (	self   :				Any      )		->			Optional[int]:
		"""simple docstring"""







		__lowercase   							=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	PNDMScheduler(skip_prk_steps=_lowerCAmelCase      )
		__lowercase   							=	StableDiffusionLDMaDPipeline(**_lowerCAmelCase      )
		__lowercase   							=	ldmad_pipe.to(_lowerCAmelCase      )
		ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	"""french fries"""
		__lowercase   							=	ldmad_pipe(**_lowerCAmelCase							,					negative_prompt=_lowerCAmelCase      )

		__lowercase							,    __lowercase   							=	output.rgb, output.depth
		__lowercase   							=	rgb[0, -3:, -3:, -1]
		__lowercase   							=	depth[0, -3:, -1]

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

		__lowercase   							=	np.array(
		    [0.37_044, 0.71_811_503, 0.7_223_251, 0.48_603_675, 0.5_638_391, 0.6_364_948, 0.42_833_704, 0.4_901_315, 0.47_926_217]      )
		__lowercase   							=	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 							__UpperCamelCase     (  unittest.TestCase ):







	def 			_a (	self   :				int      )		->			Union[str, Any]:
		"""simple docstring"""







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







	def 			_a (	self   :				Optional[int]							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				List[Any]="cpu"							,					_lowerCAmelCase   :				int=torch.floataa							,					_lowerCAmelCase   :				int=0      )		->			Optional[int]:
		"""simple docstring"""







		__lowercase   							=	torch.Generator(device=_lowerCAmelCase      ).manual_seed(_lowerCAmelCase      )
		__lowercase   							=	np.random.RandomState(_lowerCAmelCase      ).standard_normal((1, 4, 64, 64)      )
		__lowercase   							=	torch.from_numpy(_lowerCAmelCase      ).to(device=_lowerCAmelCase							,					dtype=_lowerCAmelCase      )
		__lowercase   							=	{
		    """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 			_a (	self   :				List[Any]      )		->			Dict:
		"""simple docstring"""







		__lowercase   							=	StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d"""      )
		__lowercase   							=	ldmad_pipe.to(_lowerCAmelCase      )
		ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_inputs(_lowerCAmelCase      )
		__lowercase   							=	ldmad_pipe(**_lowerCAmelCase      )
		__lowercase							,    __lowercase   							=	output.rgb, output.depth
		__lowercase   							=	rgb[0, -3:, -3:, -1].flatten()
		__lowercase   							=	rgb[0, -3:, -1].flatten()

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

		__lowercase   							=	np.array(
		    [0.53_805_465, 0.56_707_305, 0.5_486_515, 0.57_012_236, 0.5_814_511, 0.56_253_487, 0.54_843_014, 0.55_092_263, 0.6_459_706]      )
		__lowercase   							=	np.array(
		    [0.9_263_781, 0.6_678_672, 0.5_486_515, 0.92_202_145, 0.67_831_135, 0.56_253_487, 0.9_241_694, 0.7_551_478, 0.6_459_706]      )
		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 							__UpperCamelCase     (  unittest.TestCase ):







	def 			_a (	self   :				int      )		->			Tuple:
		"""simple docstring"""







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







	def 			_a (	self   :				Any							,					_lowerCAmelCase   :				Optional[Any]							,					_lowerCAmelCase   :				Optional[Any]="cpu"							,					_lowerCAmelCase   :				Dict=torch.floataa							,					_lowerCAmelCase   :				Optional[int]=0      )		->			Tuple:
		"""simple docstring"""







		__lowercase   							=	torch.Generator(device=_lowerCAmelCase      ).manual_seed(_lowerCAmelCase      )
		__lowercase   							=	np.random.RandomState(_lowerCAmelCase      ).standard_normal((1, 4, 64, 64)      )
		__lowercase   							=	torch.from_numpy(_lowerCAmelCase      ).to(device=_lowerCAmelCase							,					dtype=_lowerCAmelCase      )
		__lowercase   							=	{
		    """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 			_a (	self   :				Optional[Any]      )		->			int:
		"""simple docstring"""







		__lowercase   							=	StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d"""      ).to(_lowerCAmelCase      )
		ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_inputs(_lowerCAmelCase      )
		__lowercase   							=	ldmad_pipe(**_lowerCAmelCase      )
		__lowercase							,    __lowercase   							=	output.rgb, output.depth

		__lowercase   							=	0.495_586
		__lowercase   							=	0.33_795_515
		__lowercase   							=	112.48_518
		__lowercase   							=	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 			_a (	self   :				Tuple      )		->			Dict:
		"""simple docstring"""







		__lowercase   							=	StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d-4c"""      ).to(_lowerCAmelCase      )
		ldmad_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_inputs(_lowerCAmelCase      )
		__lowercase   							=	ldmad_pipe(**_lowerCAmelCase      )
		__lowercase							,    __lowercase   							=	output.rgb, output.depth

		__lowercase   							=	0.4_194_127
		__lowercase   							=	0.35_375_586
		__lowercase   							=	0.5_638_502
		__lowercase   							=	0.34_686_103

		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





from copy import deepcopy
class 		__magic_name__						:





     '''simple docstring'''



     def __init__(					self:int     ,		_a:list[int] | None = None     ,		_a:int | None = None   ):

       if arr is None and size is not None:
         snake_case__						= size
         snake_case__						= [0] * size
       elif arr is not None:
         self.init(_a   )
       else:
         raise ValueError('''Either arr or size must be specified'''   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Any     ,		_a:list[int]   ):
       snake_case__						= len(_a   )
       snake_case__						= deepcopy(_a   )
       for i in range(1     ,		self.size   ):
         snake_case__						= self.next_(_a   )
         if j < self.size:
           self.tree[j] += self.tree[i]

     def 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       snake_case__						= self.tree[:]
       for i in range(self.size - 1     ,		0     ,		-1   ):
         snake_case__						= self.next_(_a   )
         if j < self.size:
           arr[j] -= arr[i]
       return arr

     @staticmethod
     def 	SCREAMING_SNAKE_CASE__						(					_a:int   ):
       return index + (index & (-index))

     @staticmethod
     def 	SCREAMING_SNAKE_CASE__						(					_a:int   ):
       return index - (index & (-index))

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]     ,		_a:int     ,		_a:int   ):
       if index == 0:
         self.tree[0] += value
         return
       while index < self.size:
         self.tree[index] += value
         snake_case__						= self.next_(_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict     ,		_a:int     ,		_a:int   ):
       self.add(_a     ,		value - self.get(_a   )   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]     ,		_a:int   ):
       if right == 0:
         return 0
       snake_case__						= self.tree[0]
       right -= 1  # make right inclusive
       while right > 0:
         result += self.tree[right]
         snake_case__						= self.prev(_a   )
       return result

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict     ,		_a:int     ,		_a:int   ):
       return self.prefix(_a   ) - self.prefix(_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:str     ,		_a:int   ):
       return self.query(_a     ,		index + 1   )







     def 	SCREAMING_SNAKE_CASE__						(					self:str     ,		_a:int   ):
       value -= self.tree[0]
       if value < 0:
         return -1

       snake_case__						= 1  # Largest power of 2 <= size
       while j * 2 < self.size:
         j *= 2

       snake_case__						= 0

       while j > 0:
         if i + j < self.size and self.tree[i + j] <= value:
           value -= self.tree[i + j]
           i += j
         j //= 2
       return i


if __name__ == "__main__":
   import doctest

   doctest.testmod()

