




'''simple docstring'''





import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class 							__UpperCamelCase			( __UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					unittest.TestCase  ):





				'''simple docstring'''



				__a							: List[Any]										=StableUnCLIPPipeline
				__a							: List[str]										=TEXT_TO_IMAGE_PARAMS
				__a							: Any										=TEXT_TO_IMAGE_BATCH_PARAMS
				__a							: Any										=TEXT_TO_IMAGE_IMAGE_PARAMS
				__a							: str										=TEXT_TO_IMAGE_IMAGE_PARAMS

				# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
				__a							: Optional[int]										=False



				def       __snake_case		( self						):
											lowerCAmelCase   =		32
											lowerCAmelCase   =		embedder_hidden_size

											# prior components

											torch.manual_seed(0						)
											lowerCAmelCase   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		CLIPTextModelWithProjection(
											    CLIPTextConfig(
											        bos_token_id=0       , eos_token_id=2       , hidden_size=UpperCAmelCase_       , projection_dim=UpperCAmelCase_       , intermediate_size=37       , layer_norm_eps=1E-0_5       , num_attention_heads=4       , num_hidden_layers=5       , pad_token_id=1       , vocab_size=10_00       , )						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		PriorTransformer(
											    num_attention_heads=2       , attention_head_dim=12       , embedding_dim=UpperCAmelCase_       , num_layers=1       , )

											torch.manual_seed(0						)
											lowerCAmelCase   =		DDPMScheduler(
											    variance_type='''fixed_small_log'''       , prediction_type='''sample'''       , num_train_timesteps=10_00       , clip_sample=UpperCAmelCase_       , clip_sample_range=5.0       , beta_schedule='''squaredcos_cap_v2'''       , )

											# regular denoising components

											torch.manual_seed(0						)
											lowerCAmelCase   =		StableUnCLIPImageNormalizer(embedding_dim=UpperCAmelCase_						)
											lowerCAmelCase   =		DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		CLIPTextModel(
											    CLIPTextConfig(
											        bos_token_id=0       , eos_token_id=2       , hidden_size=UpperCAmelCase_       , projection_dim=32       , intermediate_size=37       , layer_norm_eps=1E-0_5       , num_attention_heads=4       , num_hidden_layers=5       , pad_token_id=1       , vocab_size=10_00       , )						)

											torch.manual_seed(0						)
											lowerCAmelCase   =		UNetaDConditionModel(
											    sample_size=32       , in_channels=4       , out_channels=4       , down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')       , up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')       , block_out_channels=(32, 64)       , attention_head_dim=(2, 4)       , class_embed_type='''projection'''       , projection_class_embeddings_input_dim=embedder_projection_dim * 2       , cross_attention_dim=UpperCAmelCase_       , layers_per_block=1       , upcast_attention=UpperCAmelCase_       , use_linear_projection=UpperCAmelCase_       , )

											torch.manual_seed(0						)
											lowerCAmelCase   =		DDIMScheduler(
											    beta_schedule='''scaled_linear'''       , beta_start=0.00085       , beta_end=0.012       , prediction_type='''v_prediction'''       , set_alpha_to_one=UpperCAmelCase_       , steps_offset=1       , )

											torch.manual_seed(0						)
											lowerCAmelCase   =		AutoencoderKL()

											lowerCAmelCase   =		{
											    # prior components
											    '''prior_tokenizer''': prior_tokenizer,
											    '''prior_text_encoder''': prior_text_encoder,
											    '''prior''': prior,
											    '''prior_scheduler''': prior_scheduler,
											    # image noising components
											    '''image_normalizer''': image_normalizer,
											    '''image_noising_scheduler''': image_noising_scheduler,
											    # regular denoising components
											    '''tokenizer''': tokenizer,
											    '''text_encoder''': text_encoder,
											    '''unet''': unet,
											    '''scheduler''': scheduler,
											    '''vae''': vae,
											}

											return components



				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_=0						):
											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,
											    '''prior_num_inference_steps''': 2,
											    '''output_type''': '''numpy''',
											}
											return inputs



				def       __snake_case		( self						):
											lowerCAmelCase   =		torch_device == '''cpu'''

											self._test_attention_slicing_forward_pass(test_max_difference=UpperCAmelCase_						)






				def       __snake_case		( self						):
											lowerCAmelCase   =		torch_device in ['''cpu''', '''mps''']

											self._test_inference_batch_single_identical(test_max_difference=UpperCAmelCase_						)




@slow
@require_torch_gpu
class 							__UpperCamelCase			( unittest.TestCase  ):





				'''simple docstring'''



				def       __snake_case		( self						):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()



				def       __snake_case		( self						):
											lowerCAmelCase   =		load_numpy(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'''						)

											lowerCAmelCase   =		StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''       , torch_dtype=torch.floataa						)
											pipe.to(UpperCAmelCase_						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)
											# stable unclip will oom when integration tests are run on a V100,
											# so turn on memory savings
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											lowerCAmelCase   =		torch.Generator(device='''cpu'''						).manual_seed(0						)
											lowerCAmelCase   =		pipe('''anime turle'''       , generator=UpperCAmelCase_       , output_type='''np'''						)

											lowerCAmelCase   =		output.images[0]

											assert image.shape == (7_68, 7_68, 3)

											assert_mean_pixel_difference(UpperCAmelCase_       , UpperCAmelCase_						)






				def       __snake_case		( self						):
											torch.cuda.empty_cache()
											torch.cuda.reset_max_memory_allocated()
											torch.cuda.reset_peak_memory_stats()

											lowerCAmelCase   =		StableUnCLIPPipeline.from_pretrained('''fusing/stable-unclip-2-1-l'''       , torch_dtype=torch.floataa						)
											lowerCAmelCase   =		pipe.to(UpperCAmelCase_						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)
											pipe.enable_attention_slicing()
											pipe.enable_sequential_cpu_offload()

											lowerCAmelCase   =		pipe(
											    '''anime turtle'''       , prior_num_inference_steps=2       , num_inference_steps=2       , output_type='''np'''       , )

											lowerCAmelCase   =		torch.cuda.max_memory_allocated()
											# make sure that less than 7 GB is allocated
											assert mem_bytes < 7 * 10**9









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase_						={
    """configuration_jukebox""": [
        """JUKEBOX_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """JukeboxConfig""",
        """JukeboxPriorConfig""",
        """JukeboxVQVAEConfig""",
    ],
    """tokenization_jukebox""": ["""JukeboxTokenizer"""],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase_						=[
		    """JUKEBOX_PRETRAINED_MODEL_ARCHIVE_LIST""",
		    """JukeboxModel""",
		    """JukeboxPreTrainedModel""",
		    """JukeboxVQVAE""",
		    """JukeboxPrior""",
		]

if TYPE_CHECKING:
		from .configuration_jukebox import (
		    JUKEBOX_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    JukeboxConfig,
		    JukeboxPriorConfig,
		    JukeboxVQVAEConfig,
		)
		from .tokenization_jukebox import JukeboxTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_jukebox import (
				    JUKEBOX_PRETRAINED_MODEL_ARCHIVE_LIST,
				    JukeboxModel,
				    JukeboxPreTrainedModel,
				    JukeboxPrior,
				    JukeboxVQVAE,
				)

else:
		import sys

		UpperCAmelCase_						=_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



