

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       __magic_name__       (     _lowerCamelCase    ,       _lowerCamelCase    ,       _lowerCamelCase    ,       unittest.TestCase ):
			SCREAMING_SNAKE_CASE				  =		StableUnCLIPPipeline
			SCREAMING_SNAKE_CASE				  =		TEXT_TO_IMAGE_PARAMS
			SCREAMING_SNAKE_CASE				  =		TEXT_TO_IMAGE_BATCH_PARAMS
			SCREAMING_SNAKE_CASE				  =		TEXT_TO_IMAGE_IMAGE_PARAMS
			SCREAMING_SNAKE_CASE				  =		TEXT_TO_IMAGE_IMAGE_PARAMS

			# TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
			SCREAMING_SNAKE_CASE				  =		False





			def 						__magic_name__       (						self   )  ->					str:




						'''simple docstring'''
						__a     =32
						__a     =embedder_hidden_size

						# prior components

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

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

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

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

						# regular denoising components

						torch.manual_seed(0   )
						__a     =StableUnCLIPImageNormalizer(embedding_dim=__UpperCamelCase   )
						__a     =DDPMScheduler(beta_schedule='squaredcos_cap_v2'   )

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

						torch.manual_seed(0   )
						__a     =CLIPTextModel(
						    CLIPTextConfig(
						        bos_token_id=0   ,    eos_token_id=2   ,    hidden_size=__UpperCamelCase   ,    projection_dim=32   ,    intermediate_size=37   ,    layer_norm_eps=1e-05   ,    num_attention_heads=4   ,    num_hidden_layers=5   ,    pad_token_id=1   ,    vocab_size=1000   ,    )   )

						torch.manual_seed(0   )
						__a     =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   )
						__a     =DDIMScheduler(
						    beta_schedule='scaled_linear'   ,    beta_start=0.0_0085   ,    beta_end=0.012   ,    prediction_type='v_prediction'   ,    set_alpha_to_one=__UpperCamelCase   ,    steps_offset=1   ,    )

						torch.manual_seed(0   )
						__a     =AutoencoderKL()

						__a     ={
						    # 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 						__magic_name__       (						self   ,    __snake_case   ,    __snake_case=0   )  ->					Optional[int]:




						'''simple docstring'''
						if str(__UpperCamelCase   ).startswith('mps'   ):
									__a     =torch.manual_seed(__UpperCamelCase   )
						else:
									__a     =torch.Generator(device=__UpperCamelCase   ).manual_seed(__UpperCamelCase   )
						__a     ={
						    '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 						__magic_name__       (						self   )  ->					Union[str, Any]:




						'''simple docstring'''
						__a     =torch_device == 'cpu'

						self._test_attention_slicing_forward_pass(test_max_difference=__UpperCamelCase   )





			def 						__magic_name__       (						self   )  ->					Tuple:




						'''simple docstring'''
						__a     =torch_device in ['cpu', 'mps']

						self._test_inference_batch_single_identical(test_max_difference=__UpperCamelCase   )






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





			def 						__magic_name__       (						self   )  ->					Union[str, Any]:




						'''simple docstring'''
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()





			def 						__magic_name__       (						self   )  ->					Tuple:




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

						__a     =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()

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

						__a     =output.images[0]

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

						assert_mean_pixel_difference(__UpperCamelCase   ,    __UpperCamelCase   )





			def 						__magic_name__       (						self   )  ->					Dict:




						'''simple docstring'''
						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()

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

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

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











import unittest
from queue import Empty
from threading import Thread

from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
from transformers.testing_utils import CaptureStdout, require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
				import torch

				from transformers import AutoModelForCausalLM



@require_torch
class        lowerCAmelCase__					(  unittest.TestCase       ):


				def __UpperCamelCase						(							self      :  List[Any]       )		->     Optional[Any]:
											A      						=    AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'       )
											A      						=    AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'       ).to(__UpperCamelCase       )
											A      						=    -1

											A      						=    ids_tensor((1, 5)						, vocab_size=model.config.vocab_size       ).to(__UpperCamelCase       )
											A      						=    model.generate(__UpperCamelCase						, max_new_tokens=10						, do_sample=__UpperCamelCase       )
											A      						=    tokenizer.decode(greedy_ids[0]       )

											with CaptureStdout() as cs:
																		A      						=    TextStreamer(__UpperCamelCase       )
																		model.generate(__UpperCamelCase						, max_new_tokens=10						, do_sample=__UpperCamelCase						, streamer=__UpperCamelCase       )
											# The greedy text should be printed to stdout, except for the final "\n" in the streamer
											A      						=    cs.out[:-1]

											self.assertEqual(__UpperCamelCase						, __UpperCamelCase       )


				def __UpperCamelCase						(							self      :  int       )		->     List[Any]:
											A      						=    AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'       )
											A      						=    AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'       ).to(__UpperCamelCase       )
											A      						=    -1

											A      						=    ids_tensor((1, 5)						, vocab_size=model.config.vocab_size       ).to(__UpperCamelCase       )
											A      						=    model.generate(__UpperCamelCase						, max_new_tokens=10						, do_sample=__UpperCamelCase       )
											A      						=    tokenizer.decode(greedy_ids[0]       )

											A      						=    TextIteratorStreamer(__UpperCamelCase       )
											A      						=    {'input_ids': input_ids, 'max_new_tokens': 10, 'do_sample': False, 'streamer': streamer}
											A      						=    Thread(target=model.generate						, kwargs=__UpperCamelCase       )
											thread.start()
											A      						=    ''
											for new_text in streamer:
																		streamer_text += new_text

											self.assertEqual(__UpperCamelCase						, __UpperCamelCase       )


				def __UpperCamelCase						(							self      :  int       )		->     Tuple:
											A      						=    AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'       )
											A      						=    AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'       ).to(__UpperCamelCase       )
											A      						=    -1

											A      						=    ids_tensor((1, 5)						, vocab_size=model.config.vocab_size       ).to(__UpperCamelCase       )
											A      						=    model.generate(__UpperCamelCase						, max_new_tokens=10						, do_sample=__UpperCamelCase       )
											A      						=    greedy_ids[:, input_ids.shape[1] :]
											A      						=    tokenizer.decode(new_greedy_ids[0]       )

											with CaptureStdout() as cs:
																		A      						=    TextStreamer(__UpperCamelCase						, skip_prompt=__UpperCamelCase       )
																		model.generate(__UpperCamelCase						, max_new_tokens=10						, do_sample=__UpperCamelCase						, streamer=__UpperCamelCase       )
											# The greedy text should be printed to stdout, except for the final "\n" in the streamer
											A      						=    cs.out[:-1]

											self.assertEqual(__UpperCamelCase						, __UpperCamelCase       )


				def __UpperCamelCase						(							self      :  List[Any]       )		->     Tuple:
											# Tests that we can pass `decode_kwargs` to the streamer to control how the tokens are decoded. Must be tested
											# with actual models -- the dummy models' tokenizers are not aligned with their models, and
											# `skip_special_tokens=True` has no effect on them
											A      						=    AutoTokenizer.from_pretrained('distilgpt2'       )
											A      						=    AutoModelForCausalLM.from_pretrained('distilgpt2'       ).to(__UpperCamelCase       )
											A      						=    -1

											A      						=    torch.ones((1, 5)						, device=__UpperCamelCase       ).long() * model.config.bos_token_id
											with CaptureStdout() as cs:
																		A      						=    TextStreamer(__UpperCamelCase						, skip_special_tokens=__UpperCamelCase       )
																		model.generate(__UpperCamelCase						, max_new_tokens=1						, do_sample=__UpperCamelCase						, streamer=__UpperCamelCase       )

											# The prompt contains a special token, so the streamer should not print it. As such, the output text, when
											# re-tokenized, must only contain one token
											A      						=    cs.out[:-1]  # Remove the final "\n"
											A      						=    tokenizer(__UpperCamelCase						, return_tensors='pt'       )
											self.assertEqual(streamer_text_tokenized.input_ids.shape						, (1, 1)       )


				def __UpperCamelCase						(							self      :  Any       )		->     Dict:
											A      						=    AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'       )
											A      						=    AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'       ).to(__UpperCamelCase       )
											A      						=    -1

											A      						=    ids_tensor((1, 5)						, vocab_size=model.config.vocab_size       ).to(__UpperCamelCase       )
											A      						=    TextIteratorStreamer(__UpperCamelCase						, timeout=0.0_0_1       )
											A      						=    {'input_ids': input_ids, 'max_new_tokens': 10, 'do_sample': False, 'streamer': streamer}
											A      						=    Thread(target=model.generate						, kwargs=__UpperCamelCase       )
											thread.start()

											# The streamer will timeout after 0.001 seconds, so an exception will be raised
											with self.assertRaises(__UpperCamelCase       ):
																		A      						=    ''
																		for new_text in streamer:
																									streamer_text += new_text