


"""simple docstring"""






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"""


					__a      			=				StableDiffusionLDMaDPipeline
					__a      			=				TEXT_TO_IMAGE_PARAMS
					__a      			=				TEXT_TO_IMAGE_BATCH_PARAMS
					__a      			=				TEXT_TO_IMAGE_IMAGE_PARAMS






					def lowerCamelCase__				(			self       :					Dict    ):
												'''simple docstring'''






												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  :		str 				=   DDIMScheduler(
												    beta_start=0.00085					,			beta_end=0.012					,			beta_schedule="""scaled_linear"""					,			clip_sample=UpperCamelCase_					,			set_alpha_to_one=UpperCamelCase_					,			)
												torch.manual_seed(0    )
												__UpperCAmelCase  :		Union[str, 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    )
												__UpperCAmelCase  :		Tuple 				=   CLIPTextConfig(
												    bos_token_id=0					,			eos_token_id=2					,			hidden_size=32					,			intermediate_size=37					,			layer_norm_eps=1e-0_5					,			num_attention_heads=4					,			num_hidden_layers=5					,			pad_token_id=1					,			vocab_size=1_000					,			)
												__UpperCAmelCase  :		Union[str, Any] 				=   CLIPTextModel(UpperCamelCase_    )
												__UpperCAmelCase  :		str 				=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    )

												__UpperCAmelCase  :		List[Any] 				=   {
												    """unet""": unet,
												    """scheduler""": scheduler,
												    """vae""": vae,
												    """text_encoder""": text_encoder,
												    """tokenizer""": tokenizer,
												    """safety_checker""": None,
												    """feature_extractor""": None,
												}
												return components






					def lowerCamelCase__				(			self       :					List[Any]					,			UpperCamelCase       :					List[Any]					,			UpperCamelCase       :					List[Any]=0    ):
												'''simple docstring'''






												if str(UpperCamelCase_    ).startswith("""mps"""    ):
																			__UpperCAmelCase  :		Optional[int] 				=   torch.manual_seed(UpperCamelCase_    )
												else:
																			__UpperCAmelCase  :		Optional[int] 				=   torch.Generator(device=UpperCamelCase_    ).manual_seed(UpperCamelCase_    )
												__UpperCAmelCase  :		int 				=   {
												    """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    ):
												'''simple docstring'''






												__UpperCAmelCase  :		Optional[int] 				=   """cpu"""  # ensure determinism for the device-dependent torch.Generator

												__UpperCAmelCase  :		Tuple 				=   self.get_dummy_components()
												__UpperCAmelCase  :		Optional[int] 				=   StableDiffusionLDMaDPipeline(**UpperCamelCase_    )
												__UpperCAmelCase  :		List[Any] 				=   ldmad_pipe.to(UpperCamelCase_    )
												ldmad_pipe.set_progress_bar_config(disable=UpperCamelCase_    )

												__UpperCAmelCase  :		int 				=   self.get_dummy_inputs(UpperCamelCase_    )
												__UpperCAmelCase  :		Tuple 				=   ldmad_pipe(**UpperCamelCase_    )
												__UpperCAmelCase				,__UpperCAmelCase  :		Optional[int] 				=   output.rgb, output.depth

												__UpperCAmelCase  :		str 				=   rgb[0, -3:, -3:, -1]
												__UpperCAmelCase  :		Dict 				=   depth[0, -3:, -1]

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

												__UpperCAmelCase  :		Any 				=   np.array(
												    [0.37338176, 0.70247, 0.74203193, 0.51643604, 0.58256793, 0.60932136, 0.4181095, 0.48355877, 0.46535262]    )
												__UpperCAmelCase  :		Optional[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]    ):
												'''simple docstring'''






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

												__UpperCAmelCase  :		Dict 				=   self.get_dummy_inputs(UpperCamelCase_    )
												__UpperCAmelCase  :		List[Any] 				=   3 * [inputs["""prompt"""]]

												# forward
												__UpperCAmelCase  :		str 				=   ldmad_pipe(**UpperCamelCase_    )
												__UpperCAmelCase				,__UpperCAmelCase  :		str 				=   output.rgb, output.depth
												__UpperCAmelCase  :		List[Any] 				=   rgb_slice_a[0, -3:, -3:, -1]
												__UpperCAmelCase  :		Any 				=   depth_slice_a[0, -3:, -1]

												__UpperCAmelCase  :		List[str] 				=   self.get_dummy_inputs(UpperCamelCase_    )
												__UpperCAmelCase  :		Tuple 				=   3 * [inputs.pop("""prompt"""    )]

												__UpperCAmelCase  :		Tuple 				=   ldmad_pipe.tokenizer(
												    UpperCamelCase_					,			padding="""max_length"""					,			max_length=ldmad_pipe.tokenizer.model_max_length					,			truncation=UpperCamelCase_					,			return_tensors="""pt"""					,			)
												__UpperCAmelCase  :		Tuple 				=   text_inputs["""input_ids"""].to(UpperCamelCase_    )

												__UpperCAmelCase  :		str 				=   ldmad_pipe.text_encoder(UpperCamelCase_    )[0]

												__UpperCAmelCase  :		Tuple 				=   prompt_embeds

												# forward
												__UpperCAmelCase  :		List[Any] 				=   ldmad_pipe(**UpperCamelCase_    )
												__UpperCAmelCase				,__UpperCAmelCase  :		Tuple 				=   output.rgb, output.depth
												__UpperCAmelCase  :		Tuple 				=   rgb_slice_a[0, -3:, -3:, -1]
												__UpperCAmelCase  :		Union[str, Any] 				=   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       :					List[Any]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		str 				=   """cpu"""  # ensure determinism for the device-dependent torch.Generator
												__UpperCAmelCase  :		List[str] 				=   self.get_dummy_components()
												__UpperCAmelCase  :		str 				=   PNDMScheduler(skip_prk_steps=UpperCamelCase_    )
												__UpperCAmelCase  :		str 				=   StableDiffusionLDMaDPipeline(**UpperCamelCase_    )
												__UpperCAmelCase  :		Dict 				=   ldmad_pipe.to(UpperCamelCase_    )
												ldmad_pipe.set_progress_bar_config(disable=UpperCamelCase_    )

												__UpperCAmelCase  :		str 				=   self.get_dummy_inputs(UpperCamelCase_    )
												__UpperCAmelCase  :		str 				=   """french fries"""
												__UpperCAmelCase  :		Dict 				=   ldmad_pipe(**UpperCamelCase_					,			negative_prompt=UpperCamelCase_    )

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

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

												__UpperCAmelCase  :		int 				=   np.array(
												    [0.37044, 0.71811503, 0.7223251, 0.48603675, 0.5638391, 0.6364948, 0.42833704, 0.4901315, 0.47926217]    )
												__UpperCAmelCase  :		Union[str, 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       :					List[Any]    ):
												'''simple docstring'''






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






					def lowerCamelCase__				(			self       :					List[str]					,			UpperCamelCase       :					List[str]					,			UpperCamelCase       :					int="cpu"					,			UpperCamelCase       :					Dict=torch.floataa					,			UpperCamelCase       :					Optional[int]=0    ):
												'''simple docstring'''






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






												__UpperCAmelCase  :		Dict 				=   StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d"""    )
												__UpperCAmelCase  :		Union[str, Any] 				=   ldmad_pipe.to(UpperCamelCase_    )
												ldmad_pipe.set_progress_bar_config(disable=UpperCamelCase_    )

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

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

												__UpperCAmelCase  :		List[Any] 				=   np.array(
												    [0.53805465, 0.56707305, 0.5486515, 0.57012236, 0.5814511, 0.56253487, 0.54843014, 0.55092263, 0.6459706]    )
												__UpperCAmelCase  :		Tuple 				=   np.array(
												    [0.9263781, 0.6678672, 0.5486515, 0.92202145, 0.67831135, 0.56253487, 0.9241694, 0.7551478, 0.6459706]    )
												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       :					Any    ):
												'''simple docstring'''






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






					def lowerCamelCase__				(			self       :					Tuple					,			UpperCamelCase       :					List[str]					,			UpperCamelCase       :					List[Any]="cpu"					,			UpperCamelCase       :					int=torch.floataa					,			UpperCamelCase       :					Union[str, Any]=0    ):
												'''simple docstring'''






												__UpperCAmelCase  :		int 				=   torch.Generator(device=UpperCamelCase_    ).manual_seed(UpperCamelCase_    )
												__UpperCAmelCase  :		int 				=   np.random.RandomState(UpperCamelCase_    ).standard_normal((1, 4, 64, 64)    )
												__UpperCAmelCase  :		int 				=   torch.from_numpy(UpperCamelCase_    ).to(device=UpperCamelCase_					,			dtype=UpperCamelCase_    )
												__UpperCAmelCase  :		Optional[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    ):
												'''simple docstring'''






												__UpperCAmelCase  :		Union[str, Any] 				=   StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d"""    ).to(UpperCamelCase_    )
												ldmad_pipe.set_progress_bar_config(disable=UpperCamelCase_    )

												__UpperCAmelCase  :		Union[str, Any] 				=   self.get_inputs(UpperCamelCase_    )
												__UpperCAmelCase  :		Tuple 				=   ldmad_pipe(**UpperCamelCase_    )
												__UpperCAmelCase				,__UpperCAmelCase  :		Optional[Any] 				=   output.rgb, output.depth

												__UpperCAmelCase  :		Tuple 				=   0.495586
												__UpperCAmelCase  :		List[str] 				=   0.33795515
												__UpperCAmelCase  :		Tuple 				=   112.48_518
												__UpperCAmelCase  :		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[str]    ):
												'''simple docstring'''






												__UpperCAmelCase  :		int 				=   StableDiffusionLDMaDPipeline.from_pretrained("""Intel/ldm3d-4c"""    ).to(UpperCamelCase_    )
												ldmad_pipe.set_progress_bar_config(disable=UpperCamelCase_    )

												__UpperCAmelCase  :		int 				=   self.get_inputs(UpperCamelCase_    )
												__UpperCAmelCase  :		Dict 				=   ldmad_pipe(**UpperCamelCase_    )
												__UpperCAmelCase				,__UpperCAmelCase  :		Any 				=   output.rgb, output.depth

												__UpperCAmelCase  :		Optional[Any] 				=   0.4194127
												__UpperCAmelCase  :		Any 				=   0.35375586
												__UpperCAmelCase  :		str 				=   0.5638502
												__UpperCAmelCase  :		str 				=   0.34686103

												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







'''simple docstring'''




from typing import List, Optional, Tuple, Union

import torch

from ...schedulers import DDIMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput






class 				__snake_case     (					__magic_name__							):





	def __init__(		self      ,       UpperCamelCase_      ,       UpperCamelCase_		)      ->		int:
		super().__init__()

		# make sure scheduler can always be converted to DDIM
		snake_case__	       =   DDIMScheduler.from_config(scheduler.config		)

		self.register_modules(unet=UpperCamelCase_      ,       scheduler=UpperCamelCase_		)





	@torch.no_grad()
	def __call__(		self      ,       UpperCamelCase_ = 1      ,       UpperCamelCase_ = None      ,       UpperCamelCase_ = 0.0      ,       UpperCamelCase_ = 50      ,       UpperCamelCase_ = None      ,       UpperCamelCase_ = "pil"      ,       UpperCamelCase_ = True      ,       )      ->		Union[ImagePipelineOutput, Tuple]:

		# Sample gaussian noise to begin loop
		if isinstance(self.unet.config.sample_size      ,       UpperCamelCase_		):
			snake_case__	       =   (
			    batch_size,
			    self.unet.config.in_channels,
			    self.unet.config.sample_size,
			    self.unet.config.sample_size,
			)
		else:
			snake_case__	       =   (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)

		if isinstance(UpperCamelCase_      ,       UpperCamelCase_		) and len(UpperCamelCase_		) != batch_size:
			raise ValueError(
			    F'''You have passed a list of generators of length {len(UpperCamelCase_		)}, but requested an effective batch'''
			    F''' size of {batch_size}. Make sure the batch size matches the length of the generators.'''		)

		snake_case__	       =   randn_tensor(UpperCamelCase_      ,       generator=UpperCamelCase_      ,       device=self.device      ,       dtype=self.unet.dtype		)

		# set step values
		self.scheduler.set_timesteps(UpperCamelCase_		)

		for t in self.progress_bar(self.scheduler.timesteps		):
			# 1. predict noise model_output
			snake_case__	       =   self.unet(UpperCamelCase_      ,       UpperCamelCase_		).sample

			# 2. predict previous mean of image x_t-1 and add variance depending on eta
			# eta corresponds to η in paper and should be between [0, 1]
			# do x_t -> x_t-1
			snake_case__	       =   self.scheduler.step(
			    UpperCamelCase_      ,       UpperCamelCase_      ,       UpperCamelCase_      ,       eta=UpperCamelCase_      ,       use_clipped_model_output=UpperCamelCase_      ,       generator=UpperCamelCase_		).prev_sample

		snake_case__	       =   (image / 2 + 0.5).clamp(0      ,       1		)
		snake_case__	       =   image.cpu().permute(0      ,       2      ,       3      ,       1		).numpy()
		if output_type == "pil":
			snake_case__	       =   self.numpy_to_pil(UpperCamelCase_		)

		if not return_dict:
			return (image,)

		return ImagePipelineOutput(images=UpperCamelCase_		)



