





import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


UpperCamelCase__           =			False

class 							__SCREAMING_SNAKE_CASE       (  unittest.TestCase							):
							pass






@nightly
@require_torch_gpu
class 							__SCREAMING_SNAKE_CASE       (  unittest.TestCase							):






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






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      VersatileDiffusionTextToImagePipeline.from_pretrained("""shi-labs/versatile-diffusion"""	)
														# remove text_unet
														pipe.remove_unused_weights()
														pipe.to(__lowerCAmelCase	)
														pipe.set_progress_bar_config(disable=__lowerCAmelCase	)

														UpperCamelCase__					  =      """A painting of a squirrel eating a burger """
														UpperCamelCase__					  =      torch.manual_seed(0	)
														UpperCamelCase__					  =      pipe(
														    prompt=__lowerCAmelCase	,			generator=__lowerCAmelCase	,			guidance_scale=7.5	,			num_inference_steps=2	,			output_type="""numpy"""	).images

														with tempfile.TemporaryDirectory() as tmpdirname:
																					pipe.save_pretrained(__lowerCAmelCase	)
																					UpperCamelCase__					  =      VersatileDiffusionTextToImagePipeline.from_pretrained(__lowerCAmelCase	)
														pipe.to(__lowerCAmelCase	)
														pipe.set_progress_bar_config(disable=__lowerCAmelCase	)

														UpperCamelCase__					  =      generator.manual_seed(0	)
														UpperCamelCase__					  =      pipe(
														    prompt=__lowerCAmelCase	,			generator=__lowerCAmelCase	,			guidance_scale=7.5	,			num_inference_steps=2	,			output_type="""numpy"""	).images

														assert np.abs(image - new_image	).sum() < 1E-5, "Models don't have the same forward pass"







							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      VersatileDiffusionTextToImagePipeline.from_pretrained(
														    """shi-labs/versatile-diffusion"""	,			torch_dtype=torch.floataa	)
														pipe.to(__lowerCAmelCase	)
														pipe.set_progress_bar_config(disable=__lowerCAmelCase	)

														UpperCamelCase__					  =      """A painting of a squirrel eating a burger """
														UpperCamelCase__					  =      torch.manual_seed(0	)
														UpperCamelCase__					  =      pipe(
														    prompt=__lowerCAmelCase	,			generator=__lowerCAmelCase	,			guidance_scale=7.5	,			num_inference_steps=50	,			output_type="""numpy"""	).images

														UpperCamelCase__					  =      image[0, 253:256, 253:256, -1]

														assert image.shape == (1, 512, 512, 3)
														UpperCamelCase__					  =      np.array([0.3367, 0.3169, 0.2656, 0.3870, 0.4790, 0.3796, 0.4009, 0.4878, 0.4778]	)

														assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2











from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__           =			logging.get_logger(__name__)

UpperCamelCase__           =			{
    "tanreinama/GPTSAN-2.8B-spout_is_uniform": (
        "https://huggingface.co/tanreinama/GPTSAN-2.8B-spout_is_uniform/resolve/main/config.json"
    ),
}

class 							__SCREAMING_SNAKE_CASE       (  _a							):
							snake_case  :   Optional[Any]      =							"""gptsan-japanese"""
							snake_case  :   Any      =							[
							    """past_key_values""",
							]
							snake_case  :   str      =							{
							    """hidden_size""": """d_model""",
							    """num_attention_heads""": """num_heads""",
							    """num_hidden_layers""": """num_layers""",
							}






							def __init__(    self	,			__lowerCAmelCase=36000	,			__lowerCAmelCase=1280	,			__lowerCAmelCase=1024	,			__lowerCAmelCase=8192	,			__lowerCAmelCase=4096	,			__lowerCAmelCase=128	,			__lowerCAmelCase=10	,			__lowerCAmelCase=0	,			__lowerCAmelCase=16	,			__lowerCAmelCase=16	,			__lowerCAmelCase=128	,			__lowerCAmelCase=0.0	,			__lowerCAmelCase=1E-5	,			__lowerCAmelCase=False	,			__lowerCAmelCase=0.0	,			__lowerCAmelCase="float32"	,			__lowerCAmelCase=False	,			__lowerCAmelCase=False	,			__lowerCAmelCase=False	,			__lowerCAmelCase=0.002	,			__lowerCAmelCase=False	,			__lowerCAmelCase=True	,			__lowerCAmelCase=35998	,			__lowerCAmelCase=35995	,			__lowerCAmelCase=35999	,			**__lowerCAmelCase	,			):
														UpperCamelCase__					  =      vocab_size
														UpperCamelCase__					  =      max_position_embeddings
														UpperCamelCase__					  =      d_model
														UpperCamelCase__					  =      d_ff
														UpperCamelCase__					  =      d_ext
														UpperCamelCase__					  =      d_spout
														UpperCamelCase__					  =      num_switch_layers
														UpperCamelCase__					  =      num_ext_layers
														UpperCamelCase__					  =      num_switch_layers + num_ext_layers
														UpperCamelCase__					  =      num_heads
														UpperCamelCase__					  =      num_experts
														UpperCamelCase__					  =      expert_capacity
														UpperCamelCase__					  =      dropout_rate
														UpperCamelCase__					  =      layer_norm_epsilon
														UpperCamelCase__					  =      router_bias
														UpperCamelCase__					  =      router_jitter_noise
														UpperCamelCase__					  =      router_dtype
														UpperCamelCase__					  =      router_ignore_padding_tokens
														UpperCamelCase__					  =      output_hidden_states
														UpperCamelCase__					  =      output_attentions
														UpperCamelCase__					  =      initializer_factor
														UpperCamelCase__					  =      output_router_logits
														UpperCamelCase__					  =      use_cache

														super().__init__(
														    separator_token_id=__lowerCAmelCase	,			pad_token_id=__lowerCAmelCase	,			eos_token_id=__lowerCAmelCase	,			**__lowerCAmelCase	,			)




