





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        SCREAMING_SNAKE_CASE__					(		unittest.TestCase ):
 __SCREAMING_SNAKE_CASE							=      StableDiffusionLDMaDPipeline
 __SCREAMING_SNAKE_CASE							=      TEXT_TO_IMAGE_PARAMS
 __SCREAMING_SNAKE_CASE							=      TEXT_TO_IMAGE_BATCH_PARAMS
 __SCREAMING_SNAKE_CASE							=      TEXT_TO_IMAGE_IMAGE_PARAMS



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

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



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

        A__		   =		self.get_dummy_components()
        A__		   =		StableDiffusionLDMaDPipeline(**__lowerCamelCase					)
        A__		   =		ldmad_pipe.to(__lowerCamelCase					)
        ldmad_pipe.set_progress_bar_config(disable=__lowerCamelCase					)

        A__		   =		self.get_dummy_inputs(__lowerCamelCase					)
        A__		   =		ldmad_pipe(**__lowerCamelCase					)
        A__   ,      A__		   =		output.rgb, output.depth

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

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

        A__		   =		np.array(
            [0.37338176, 0.70247, 0.74203193, 0.51643604, 0.58256793, 0.60932136, 0.4181095, 0.48355877, 0.46535262]					)
        A__		   =		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 							UpperCamelCase     (				self					):
        A__		   =		self.get_dummy_components()
        A__		   =		StableDiffusionLDMaDPipeline(**__lowerCamelCase					)
        A__		   =		ldmad_pipe.to(__lowerCamelCase					)
        ldmad_pipe.set_progress_bar_config(disable=__lowerCamelCase					)

        A__		   =		self.get_dummy_inputs(__lowerCamelCase					)
        A__		   =		3 * [inputs['''prompt''']]

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

        A__		   =		self.get_dummy_inputs(__lowerCamelCase					)
        A__		   =		3 * [inputs.pop('''prompt'''					)]

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

        A__		   =		ldmad_pipe.text_encoder(__lowerCamelCase					)[0]

        A__		   =		prompt_embeds

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

        A__		   =		self.get_dummy_inputs(__lowerCamelCase					)
        A__		   =		'''french fries'''
        A__		   =		ldmad_pipe(**__lowerCamelCase,negative_prompt=__lowerCamelCase					)

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

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

        A__		   =		np.array(
            [0.37044, 0.71811503, 0.7223251, 0.48603675, 0.5638391, 0.6364948, 0.42833704, 0.4901315, 0.47926217]					)
        A__		   =		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        SCREAMING_SNAKE_CASE__					(		unittest.TestCase ):



 def 							UpperCamelCase     (				self					):
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



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

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

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

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



 def 							UpperCamelCase     (				self					):
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



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

        A__		   =		self.get_inputs(__lowerCamelCase					)
        A__		   =		ldmad_pipe(**__lowerCamelCase					)
        A__   ,      A__		   =		output.rgb, output.depth

        A__		   =		0.495586
        A__		   =		0.33795515
        A__		   =		112.48518
        A__		   =		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 							UpperCamelCase     (				self					):
        A__		   =		StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d-4c'''					).to(__lowerCamelCase					)
        ldmad_pipe.set_progress_bar_config(disable=__lowerCamelCase					)

        A__		   =		self.get_inputs(__lowerCamelCase					)
        A__		   =		ldmad_pipe(**__lowerCamelCase					)
        A__   ,      A__		   =		output.rgb, output.depth

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












a__: str							   =						'0.21.0'

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
       from .utils import rich





