






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	:       Dict   =					StableDiffusionLDMaDPipeline
						_SCREAMING_SNAKE_CASE	:       List[Any]   =					TEXT_TO_IMAGE_PARAMS
						_SCREAMING_SNAKE_CASE	:       int   =					TEXT_TO_IMAGE_BATCH_PARAMS
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					TEXT_TO_IMAGE_IMAGE_PARAMS






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												torch.manual_seed(0   )
												lowerCAmelCase__    						=						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     ,						)
												lowerCAmelCase__    						=						DDIMScheduler(
												    beta_start=0.0_00_85     ,						beta_end=0.0_12     ,						beta_schedule='scaled_linear'     ,						clip_sample=_UpperCamelCase     ,						set_alpha_to_one=_UpperCamelCase     ,						)
												torch.manual_seed(0   )
												lowerCAmelCase__    						=						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   )
												lowerCAmelCase__    						=						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=10_00     ,						)
												lowerCAmelCase__    						=						CLIPTextModel(_UpperCamelCase   )
												lowerCAmelCase__    						=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'   )

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






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase=0   ):






												"""simple docstring"""
												if str(_UpperCamelCase   ).startswith('mps'   ):
																		lowerCAmelCase__    						=						torch.manual_seed(_UpperCamelCase   )
												else:
																		lowerCAmelCase__    						=						torch.Generator(device=_UpperCamelCase   ).manual_seed(_UpperCamelCase   )
												lowerCAmelCase__    						=						{
												    '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   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						'cpu'  # ensure determinism for the device-dependent torch.Generator

												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						StableDiffusionLDMaDPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						ldmad_pipe.to(_UpperCamelCase   )
												ldmad_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						ldmad_pipe(**_UpperCamelCase   )
												lowerCAmelCase__						,       lowerCAmelCase__    						=						output.rgb, output.depth

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

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

												lowerCAmelCase__    						=						np.array(
												    [0.37_33_81_76, 0.7_02_47, 0.74_20_31_93, 0.51_64_36_04, 0.58_25_67_93, 0.60_93_21_36, 0.4_18_10_95, 0.48_35_58_77, 0.46_53_52_62]   )
												lowerCAmelCase__    						=						np.array([1_03.4_67_27, 85.81_20_04, 87.84_92_36]   )

												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   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						StableDiffusionLDMaDPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						ldmad_pipe.to(_UpperCamelCase   )
												ldmad_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						3 * [inputs['prompt']]

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

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						3 * [inputs.pop('prompt'   )]

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

												lowerCAmelCase__    						=						ldmad_pipe.text_encoder(_UpperCamelCase   )[0]

												lowerCAmelCase__    						=						prompt_embeds

												# forward
												lowerCAmelCase__    						=						ldmad_pipe(**_UpperCamelCase   )
												lowerCAmelCase__						,       lowerCAmelCase__    						=						output.rgb, output.depth
												lowerCAmelCase__    						=						rgb_slice_a[0, -3:, -3:, -1]
												lowerCAmelCase__    						=						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   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						'cpu'  # ensure determinism for the device-dependent torch.Generator
												lowerCAmelCase__    						=						self.get_dummy_components()
												lowerCAmelCase__    						=						PNDMScheduler(skip_prk_steps=_UpperCamelCase   )
												lowerCAmelCase__    						=						StableDiffusionLDMaDPipeline(**_UpperCamelCase   )
												lowerCAmelCase__    						=						ldmad_pipe.to(_UpperCamelCase   )
												ldmad_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_dummy_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						'french fries'
												lowerCAmelCase__    						=						ldmad_pipe(**_UpperCamelCase     ,						negative_prompt=_UpperCamelCase   )

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

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

												lowerCAmelCase__    						=						np.array(
												    [0.3_70_44, 0.71_81_15_03, 0.7_22_32_51, 0.48_60_36_75, 0.5_63_83_91, 0.6_36_49_48, 0.42_83_37_04, 0.4_90_13_15, 0.47_92_62_17]   )
												lowerCAmelCase__    						=						np.array([1_07.8_47_38, 84.6_28_02, 89.96_21_35]   )
												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   ):






												"""simple docstring"""
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase="cpu"     ,						_UpperCamelCase=torch.floataa     ,						_UpperCamelCase=0   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						torch.Generator(device=_UpperCamelCase   ).manual_seed(_UpperCamelCase   )
												lowerCAmelCase__    						=						np.random.RandomState(_UpperCamelCase   ).standard_normal((1, 4, 64, 64)   )
												lowerCAmelCase__    						=						torch.from_numpy(_UpperCamelCase   ).to(device=_UpperCamelCase     ,						dtype=_UpperCamelCase   )
												lowerCAmelCase__    						=						{
												    '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   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'   )
												lowerCAmelCase__    						=						ldmad_pipe.to(_UpperCamelCase   )
												ldmad_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

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

												assert rgb.shape == (1, 5_12, 5_12, 3)
												assert depth.shape == (1, 5_12, 5_12)

												lowerCAmelCase__    						=						np.array(
												    [0.53_80_54_65, 0.56_70_73_05, 0.5_48_65_15, 0.57_01_22_36, 0.5_81_45_11, 0.56_25_34_87, 0.54_84_30_14, 0.55_09_22_63, 0.6_45_97_06]   )
												lowerCAmelCase__    						=						np.array(
												    [0.9_26_37_81, 0.6_67_86_72, 0.5_48_65_15, 0.92_20_21_45, 0.67_83_11_35, 0.56_25_34_87, 0.9_24_16_94, 0.7_55_14_78, 0.6_45_97_06]   )
												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   ):






												"""simple docstring"""
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase="cpu"     ,						_UpperCamelCase=torch.floataa     ,						_UpperCamelCase=0   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						torch.Generator(device=_UpperCamelCase   ).manual_seed(_UpperCamelCase   )
												lowerCAmelCase__    						=						np.random.RandomState(_UpperCamelCase   ).standard_normal((1, 4, 64, 64)   )
												lowerCAmelCase__    						=						torch.from_numpy(_UpperCamelCase   ).to(device=_UpperCamelCase     ,						dtype=_UpperCamelCase   )
												lowerCAmelCase__    						=						{
												    '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   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'   ).to(_UpperCamelCase   )
												ldmad_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						ldmad_pipe(**_UpperCamelCase   )
												lowerCAmelCase__						,       lowerCAmelCase__    						=						output.rgb, output.depth

												lowerCAmelCase__    						=						0.49_55_86
												lowerCAmelCase__    						=						0.33_79_55_15
												lowerCAmelCase__    						=						1_12.4_85_18
												lowerCAmelCase__    						=						98.48_97_46
												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   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'   ).to(_UpperCamelCase   )
												ldmad_pipe.set_progress_bar_config(disable=_UpperCamelCase   )

												lowerCAmelCase__    						=						self.get_inputs(_UpperCamelCase   )
												lowerCAmelCase__    						=						ldmad_pipe(**_UpperCamelCase   )
												lowerCAmelCase__						,       lowerCAmelCase__    						=						output.rgb, output.depth

												lowerCAmelCase__    						=						0.4_19_41_27
												lowerCAmelCase__    						=						0.35_37_55_86
												lowerCAmelCase__    						=						0.5_63_85_02
												lowerCAmelCase__    						=						0.34_68_61_03

												assert rgb.shape == (1, 5_12, 5_12, 3)
												assert depth.shape == (1, 5_12, 5_12, 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 json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging


if TYPE_CHECKING:
					from transformers.pipelines.conversational import Conversation


__snake_case    :    Any         =     logging.get_logger(__name__)

__snake_case    :    Tuple         =     {"""tokenizer_file""": """tokenizer.json"""}

__snake_case    :    str         =     {
    """tokenizer_file""": {
        """bigscience/tokenizer""": """https://huggingface.co/bigscience/tokenizer/blob/main/tokenizer.json""",
        """bigscience/bloom-560m""": """https://huggingface.co/bigscience/bloom-560m/blob/main/tokenizer.json""",
        """bigscience/bloom-1b1""": """https://huggingface.co/bigscience/bloom-1b1/blob/main/tokenizer.json""",
        """bigscience/bloom-1b7""": """https://huggingface.co/bigscience/bloom-1b7/blob/main/tokenizer.json""",
        """bigscience/bloom-3b""": """https://huggingface.co/bigscience/bloom-3b/blob/main/tokenizer.json""",
        """bigscience/bloom-7b1""": """https://huggingface.co/bigscience/bloom-7b1/blob/main/tokenizer.json""",
        """bigscience/bloom""": """https://huggingface.co/bigscience/bloom/blob/main/tokenizer.json""",
    },
}


class     __SCREAMING_SNAKE_CASE					(   __lowercase):
						_SCREAMING_SNAKE_CASE	:       Dict   =					VOCAB_FILES_NAMES
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					PRETRAINED_VOCAB_FILES_MAP
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					['''input_ids''', '''attention_mask''']
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					None






						def __init__(	self     ,						_UpperCamelCase=None     ,						_UpperCamelCase=None     ,						_UpperCamelCase=None     ,						_UpperCamelCase="<unk>"     ,						_UpperCamelCase="<s>"     ,						_UpperCamelCase="</s>"     ,						_UpperCamelCase="<pad>"     ,						_UpperCamelCase=False     ,						_UpperCamelCase=False     ,						**_UpperCamelCase     ,						):






												"""simple docstring"""
												super().__init__(
												    _UpperCamelCase     ,						_UpperCamelCase     ,						tokenizer_file=_UpperCamelCase     ,						unk_token=_UpperCamelCase     ,						bos_token=_UpperCamelCase     ,						eos_token=_UpperCamelCase     ,						pad_token=_UpperCamelCase     ,						add_prefix_space=_UpperCamelCase     ,						clean_up_tokenization_spaces=_UpperCamelCase     ,						**_UpperCamelCase     ,						)
												lowerCAmelCase__    						=						json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()   )
												if pre_tok_state.get('add_prefix_space'     ,						_UpperCamelCase   ) != add_prefix_space:
																		lowerCAmelCase__    						=						getattr(_UpperCamelCase     ,						pre_tok_state.pop('type'   )   )
																		lowerCAmelCase__    						=						add_prefix_space
																		lowerCAmelCase__    						=						pre_tok_class(**_UpperCamelCase   )

												lowerCAmelCase__    						=						add_prefix_space






						def       UpperCamelCase__      (	self     ,						*_UpperCamelCase     ,						**_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						kwargs.get('is_split_into_words'     ,						_UpperCamelCase   )
												if not (self.add_prefix_space or not is_split_into_words):
																		raise Exception(
																		    F"You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"
																		    ' pretokenized inputs.'   )

												return super()._batch_encode_plus(*_UpperCamelCase     ,						**_UpperCamelCase   )






						def       UpperCamelCase__      (	self     ,						*_UpperCamelCase     ,						**_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						kwargs.get('is_split_into_words'     ,						_UpperCamelCase   )

												if not (self.add_prefix_space or not is_split_into_words):
																		raise Exception(
																		    F"You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"
																		    ' pretokenized inputs.'   )

												return super()._encode_plus(*_UpperCamelCase     ,						**_UpperCamelCase   )






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase = None   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self._tokenizer.model.save(_UpperCamelCase     ,						name=_UpperCamelCase   )
												return tuple(_UpperCamelCase   )







						def       UpperCamelCase__      (	self     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						[]
												for is_user, text in conversation.iter_texts():
																		input_ids.extend(self.encode(_UpperCamelCase     ,						add_special_tokens=_UpperCamelCase   ) + [self.eos_token_id]   )

												if len(_UpperCamelCase   ) > self.model_max_length:
																		lowerCAmelCase__    						=						input_ids[-self.model_max_length :]
												return input_ids





