


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 		__a				(    UpperCAmelCase   ,				UpperCAmelCase   ,				UpperCAmelCase   ,				unittest.TestCase      ):
		_a     :  int           =     StableUnCLIPPipeline
		_a     :  Any           =     TEXT_TO_IMAGE_PARAMS
		_a     :  Tuple           =     TEXT_TO_IMAGE_BATCH_PARAMS
		_a     :  List[str]           =     TEXT_TO_IMAGE_IMAGE_PARAMS
		_a     :  Dict           =     TEXT_TO_IMAGE_IMAGE_PARAMS

		# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
		_a     :  List[str]           =     False
		def      UpperCAmelCase__    (      self						)			->	Dict:




									"""simple docstring"""





									_UpperCAmelCase					=      32
									_UpperCAmelCase					=      embedder_hidden_size

									# prior components

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

									torch.manual_seed(0						)
									_UpperCAmelCase					=      CLIPTextModelWithProjection(
									    CLIPTextConfig(
									        bos_token_id=0      ,   eos_token_id=2      ,   hidden_size=_SCREAMING_SNAKE_CASE      ,   projection_dim=_SCREAMING_SNAKE_CASE      ,   intermediate_size=37      ,   layer_norm_eps=1e-0_5      ,   num_attention_heads=4      ,   num_hidden_layers=5      ,   pad_token_id=1      ,   vocab_size=1000      ,   )						)

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

									torch.manual_seed(0						)
									_UpperCAmelCase					=      DDPMScheduler(
									    variance_type='fixed_small_log'      ,   prediction_type='sample'      ,   num_train_timesteps=1000      ,   clip_sample=_SCREAMING_SNAKE_CASE      ,   clip_sample_range=5.0      ,   beta_schedule='squaredcos_cap_v2'      ,   )

									# regular denoising components

									torch.manual_seed(0						)
									_UpperCAmelCase					=      StableUnCLIPImageNormalizer(embedding_dim=_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      DDPMScheduler(beta_schedule='squaredcos_cap_v2'						)

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

									torch.manual_seed(0						)
									_UpperCAmelCase					=      CLIPTextModel(
									    CLIPTextConfig(
									        bos_token_id=0      ,   eos_token_id=2      ,   hidden_size=_SCREAMING_SNAKE_CASE      ,   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=1000      ,   )						)

									torch.manual_seed(0						)
									_UpperCAmelCase					=      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=_SCREAMING_SNAKE_CASE      ,   layers_per_block=1      ,   upcast_attention=_SCREAMING_SNAKE_CASE      ,   use_linear_projection=_SCREAMING_SNAKE_CASE      ,   )

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

									torch.manual_seed(0						)
									_UpperCAmelCase					=      AutoencoderKL()

									_UpperCAmelCase					=      {
									    # 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      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=0						)			->	Union[str, Any]:




									"""simple docstring"""





									if str(_SCREAMING_SNAKE_CASE						).startswith('mps'						):
																_UpperCAmelCase					=      torch.manual_seed(_SCREAMING_SNAKE_CASE						)
									else:
																_UpperCAmelCase					=      torch.Generator(device=_SCREAMING_SNAKE_CASE						).manual_seed(_SCREAMING_SNAKE_CASE						)
									_UpperCAmelCase					=      {
									    '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      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





									_UpperCAmelCase					=      torch_device == 'cpu'

									self._test_attention_slicing_forward_pass(test_max_difference=_SCREAMING_SNAKE_CASE						)






		def      UpperCAmelCase__    (      self						)			->	Dict:




									"""simple docstring"""





									_UpperCAmelCase					=      torch_device in ['cpu', 'mps']

									self._test_inference_batch_single_identical(test_max_difference=_SCREAMING_SNAKE_CASE						)



@slow
@require_torch_gpu
class 		__a				(    unittest.TestCase      ):
		def      UpperCAmelCase__    (      self						)			->	List[Any]:




									"""simple docstring"""





									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()
		def      UpperCAmelCase__    (      self						)			->	Any:




									"""simple docstring"""





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

									_UpperCAmelCase					=      StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'      ,   torch_dtype=torch.floataa						)
									pipe.to(_SCREAMING_SNAKE_CASE						)
									pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE						)
									# 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()

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

									_UpperCAmelCase					=      output.images[0]

									assert image.shape == (768, 768, 3)

									assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)






		def      UpperCAmelCase__    (      self						)			->	int:




									"""simple docstring"""





									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()
									torch.cuda.reset_peak_memory_stats()

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

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

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




from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import torch

from ...utils import BaseOutput, OptionalDependencyNotAvailable, is_torch_available, is_transformers_available

@dataclass
class 		__a				(    UpperCAmelCase      ):
		_a     :  Union[List[np.ndarray], torch.FloatTensor]


try:
		if not (is_transformers_available() and is_torch_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
		from .pipeline_text_to_video_synth import TextToVideoSDPipeline
		from .pipeline_text_to_video_synth_imgaimg import VideoToVideoSDPipeline  # noqa: F401
		from .pipeline_text_to_video_zero import TextToVideoZeroPipeline
