



"""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 			lowercase__			(      unittest.TestCase    ):
       '''simple docstring'''







       _UpperCAmelCase      				=					StableDiffusionLDMaDPipeline
       _UpperCAmelCase      				=					TEXT_TO_IMAGE_PARAMS
       _UpperCAmelCase      				=					TEXT_TO_IMAGE_BATCH_PARAMS
       _UpperCAmelCase      				=					TEXT_TO_IMAGE_IMAGE_PARAMS






       def        lowerCamelCase_   (  self						)			->							str:
         torch.manual_seed(0						)
         _UpperCAmelCase							 =		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							 =		DDIMScheduler(
             beta_start=0.00085	,       beta_end=0.012	,       beta_schedule='scaled_linear'	,       clip_sample=snake_case	,       set_alpha_to_one=snake_case	,       )
         torch.manual_seed(0						)
         _UpperCAmelCase							 =		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							 =		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	,       )
         _UpperCAmelCase							 =		CLIPTextModel(snake_case						)
         _UpperCAmelCase							 =		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

         _UpperCAmelCase							 =		{
             'unet': unet,
             'scheduler': scheduler,
             'vae': vae,
             'text_encoder': text_encoder,
             'tokenizer': tokenizer,
             'safety_checker': None,
             'feature_extractor': None,
         }
         return components






       def        lowerCamelCase_   (  self	,       snake_case	,       snake_case=0						)			->							Dict:
         if str(snake_case						).startswith('mps'						):
           _UpperCAmelCase							 =		torch.manual_seed(snake_case						)
         else:
           _UpperCAmelCase							 =		torch.Generator(device=snake_case						).manual_seed(snake_case						)
         _UpperCAmelCase							 =		{
             '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						)			->							Dict:
         _UpperCAmelCase							 =		'cpu'  # ensure determinism for the device-dependent torch.Generator

         _UpperCAmelCase							 =		self.get_dummy_components()
         _UpperCAmelCase							 =		StableDiffusionLDMaDPipeline(**snake_case						)
         _UpperCAmelCase							 =		ldmad_pipe.to(snake_case						)
         ldmad_pipe.set_progress_bar_config(disable=snake_case						)

         _UpperCAmelCase							 =		self.get_dummy_inputs(snake_case						)
         _UpperCAmelCase							 =		ldmad_pipe(**snake_case						)
         _UpperCAmelCase ,					_UpperCAmelCase							 =		output.rgb, output.depth

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

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

         _UpperCAmelCase							 =		np.array(
             [0.37338176, 0.70247, 0.74203193, 0.51643604, 0.58256793, 0.60932136, 0.4181095, 0.48355877, 0.46535262]						)
         _UpperCAmelCase							 =		np.array([103.46727, 85.812004, 87.849236]						)

         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						)			->							Dict:
         _UpperCAmelCase							 =		self.get_dummy_components()
         _UpperCAmelCase							 =		StableDiffusionLDMaDPipeline(**snake_case						)
         _UpperCAmelCase							 =		ldmad_pipe.to(snake_case						)
         ldmad_pipe.set_progress_bar_config(disable=snake_case						)

         _UpperCAmelCase							 =		self.get_dummy_inputs(snake_case						)
         _UpperCAmelCase							 =		3 * [inputs['prompt']]

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

         _UpperCAmelCase							 =		self.get_dummy_inputs(snake_case						)
         _UpperCAmelCase							 =		3 * [inputs.pop('prompt'						)]

         _UpperCAmelCase							 =		ldmad_pipe.tokenizer(
             snake_case	,       padding='max_length'	,       max_length=ldmad_pipe.tokenizer.model_max_length	,       truncation=snake_case	,       return_tensors='pt'	,       )
         _UpperCAmelCase							 =		text_inputs['input_ids'].to(snake_case						)

         _UpperCAmelCase							 =		ldmad_pipe.text_encoder(snake_case						)[0]

         _UpperCAmelCase							 =		prompt_embeds

         # forward
         _UpperCAmelCase							 =		ldmad_pipe(**snake_case						)
         _UpperCAmelCase ,					_UpperCAmelCase							 =		output.rgb, output.depth
         _UpperCAmelCase							 =		rgb_slice_a[0, -3:, -3:, -1]
         _UpperCAmelCase							 =		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						)			->							int:
         _UpperCAmelCase							 =		'cpu'  # ensure determinism for the device-dependent torch.Generator
         _UpperCAmelCase							 =		self.get_dummy_components()
         _UpperCAmelCase							 =		PNDMScheduler(skip_prk_steps=snake_case						)
         _UpperCAmelCase							 =		StableDiffusionLDMaDPipeline(**snake_case						)
         _UpperCAmelCase							 =		ldmad_pipe.to(snake_case						)
         ldmad_pipe.set_progress_bar_config(disable=snake_case						)

         _UpperCAmelCase							 =		self.get_dummy_inputs(snake_case						)
         _UpperCAmelCase							 =		'french fries'
         _UpperCAmelCase							 =		ldmad_pipe(**snake_case	,       negative_prompt=snake_case						)

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

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

         _UpperCAmelCase							 =		np.array(
             [0.37044, 0.71811503, 0.7223251, 0.48603675, 0.5638391, 0.6364948, 0.42833704, 0.4901315, 0.47926217]						)
         _UpperCAmelCase							 =		np.array([107.84738, 84.62802, 89.962135]						)
         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 			lowercase__			(      unittest.TestCase    ):
       '''simple docstring'''







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






       def        lowerCamelCase_   (  self	,       snake_case	,       snake_case="cpu"	,       snake_case=torch.floataa	,       snake_case=0						)			->							Union[str, Any]:
         _UpperCAmelCase							 =		torch.Generator(device=snake_case						).manual_seed(snake_case						)
         _UpperCAmelCase							 =		np.random.RandomState(snake_case						).standard_normal((1, 4, 64, 64)						)
         _UpperCAmelCase							 =		torch.from_numpy(snake_case						).to(device=snake_case	,       dtype=snake_case						)
         _UpperCAmelCase							 =		{
             '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:
         _UpperCAmelCase							 =		StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'						)
         _UpperCAmelCase							 =		ldmad_pipe.to(snake_case						)
         ldmad_pipe.set_progress_bar_config(disable=snake_case						)

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

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

         _UpperCAmelCase							 =		np.array(
             [0.53805465, 0.56707305, 0.5486515, 0.57012236, 0.5814511, 0.56253487, 0.54843014, 0.55092263, 0.6459706]						)
         _UpperCAmelCase							 =		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 			lowercase__			(      unittest.TestCase    ):
       '''simple docstring'''







       def        lowerCamelCase_   (  self						)			->							int:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






       def        lowerCamelCase_   (  self	,       snake_case	,       snake_case="cpu"	,       snake_case=torch.floataa	,       snake_case=0						)			->							List[str]:
         _UpperCAmelCase							 =		torch.Generator(device=snake_case						).manual_seed(snake_case						)
         _UpperCAmelCase							 =		np.random.RandomState(snake_case						).standard_normal((1, 4, 64, 64)						)
         _UpperCAmelCase							 =		torch.from_numpy(snake_case						).to(device=snake_case	,       dtype=snake_case						)
         _UpperCAmelCase							 =		{
             '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						)			->							Optional[int]:
         _UpperCAmelCase							 =		StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'						).to(snake_case						)
         ldmad_pipe.set_progress_bar_config(disable=snake_case						)

         _UpperCAmelCase							 =		self.get_inputs(snake_case						)
         _UpperCAmelCase							 =		ldmad_pipe(**snake_case						)
         _UpperCAmelCase ,					_UpperCAmelCase							 =		output.rgb, output.depth

         _UpperCAmelCase							 =		0.495586
         _UpperCAmelCase							 =		0.33795515
         _UpperCAmelCase							 =		112.48518
         _UpperCAmelCase							 =		98.489746
         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						)			->							str:
         _UpperCAmelCase							 =		StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'						).to(snake_case						)
         ldmad_pipe.set_progress_bar_config(disable=snake_case						)

         _UpperCAmelCase							 =		self.get_inputs(snake_case						)
         _UpperCAmelCase							 =		ldmad_pipe(**snake_case						)
         _UpperCAmelCase ,					_UpperCAmelCase							 =		output.rgb, output.depth

         _UpperCAmelCase							 =		0.4194127
         _UpperCAmelCase							 =		0.35375586
         _UpperCAmelCase							 =		0.5638502
         _UpperCAmelCase							 =		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 Optional

from torch import nn

from .transformer_ad import TransformeraDModel, TransformeraDModelOutput




class 			lowercase__			(      nn.Module    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case = 16	,       snake_case = 88	,       snake_case = None	,       snake_case = 1	,       snake_case = 0.0	,       snake_case = 32	,       snake_case = None	,       snake_case = False	,       snake_case = None	,       snake_case = None	,       snake_case = "geglu"	,       snake_case = None	,       )			->							str:
									super().__init__()
									_UpperCAmelCase							 =		nn.ModuleList(
									    [
									        TransformeraDModel(
									            num_attention_heads=snake_case	,       attention_head_dim=snake_case	,       in_channels=snake_case	,       num_layers=snake_case	,       dropout=snake_case	,       norm_num_groups=snake_case	,       cross_attention_dim=snake_case	,       attention_bias=snake_case	,       sample_size=snake_case	,       num_vector_embeds=snake_case	,       activation_fn=snake_case	,       num_embeds_ada_norm=snake_case	,       )
									        for _ in range(2						)
									    ]						)

									# Variables that can be set by a pipeline:

									# The ratio of transformer1 to transformer2's output states to be combined during inference
									_UpperCAmelCase							 =		0.5

									# The shape of `encoder_hidden_states` is expected to be
									# `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)`
									_UpperCAmelCase							 =		[77, 257]

									# Which transformer to use to encode which condition.
									# E.g. `(1, 0)` means that we'll use `transformers[1](conditions[0])` and `transformers[0](conditions[1])`
									_UpperCAmelCase							 =		[1, 0]






							def        lowerCamelCase_   (  self	,       snake_case	,       snake_case	,       snake_case=None	,       snake_case=None	,       snake_case=None	,       snake_case = True	,       )			->							Any:
									_UpperCAmelCase							 =		hidden_states

									_UpperCAmelCase							 =		[]
									_UpperCAmelCase							 =		0
									# attention_mask is not used yet
									for i in range(2						):
											# for each of the two transformers, pass the corresponding condition tokens
											_UpperCAmelCase							 =		encoder_hidden_states[:, tokens_start : tokens_start + self.condition_lengths[i]]
											_UpperCAmelCase							 =		self.transformer_index_for_condition[i]
											_UpperCAmelCase							 =		self.transformers[transformer_index](
											    snake_case	,       encoder_hidden_states=snake_case	,       timestep=snake_case	,       cross_attention_kwargs=snake_case	,       return_dict=snake_case	,       )[0]
											encoded_states.append(encoded_state - input_states						)
											tokens_start += self.condition_lengths[i]

									_UpperCAmelCase							 =		encoded_states[0] * self.mix_ratio + encoded_states[1] * (1 - self.mix_ratio)
									_UpperCAmelCase							 =		output_states + input_states

									if not return_dict:
											return (output_states,)

									return TransformeraDModelOutput(sample=snake_case						)


