


'''simple docstring'''





import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class 							UpperCAmelCase ( unittest.TestCase	):


 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[str]   ):
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()


 @property
 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[str]   ):
    UpperCAmelCase__  :List[Any]            =			1
    UpperCAmelCase__  :Optional[Any]            =			3
    UpperCAmelCase__  :Dict            =			(3_2, 3_2)

    UpperCAmelCase__  :Any            =			floats_tensor((batch_size, num_channels) + sizes ,					rng=random.Random(0   )   ).to(lowercase_   )
    return image


 @property
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int]   ):
    torch.manual_seed(0   )
    UpperCAmelCase__  :Any            =			UNetaDConditionModel(
        block_out_channels=(3_2, 6_4) ,					layers_per_block=2 ,					sample_size=3_2 ,					in_channels=4 ,					out_channels=4 ,					down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''') ,					up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''') ,					cross_attention_dim=3_2 ,					)
    return model


 @property
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Dict   ):
    torch.manual_seed(0   )
    UpperCAmelCase__  :Tuple            =			AutoencoderKL(
        block_out_channels=[3_2, 6_4] ,					in_channels=3 ,					out_channels=3 ,					down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''] ,					up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''] ,					latent_channels=4 ,					)
    return model


 @property
 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    torch.manual_seed(0   )
    UpperCAmelCase__  :Any            =			CLIPTextConfig(
        bos_token_id=0 ,					eos_token_id=2 ,					hidden_size=3_2 ,					intermediate_size=3_7 ,					layer_norm_eps=1e-05 ,					num_attention_heads=4 ,					num_hidden_layers=5 ,					pad_token_id=1 ,					vocab_size=1_0_0_0 ,					)
    return CLIPTextModel(lowercase_   )


 @property
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any]   ):






    def extract(*__lowerCamelCase       :  Dict ,					**__lowerCamelCase       :  List[str]   ):




       class 							UpperCAmelCase :


        def __init__(			self       :  Dict   ):
           UpperCAmelCase__  :int            =			torch.ones([0]   )





        def 					__SCREAMING_SNAKE_CASE				(			self       :  Dict ,					__lowerCamelCase       :  Dict   ):
           self.pixel_values.to(lowercase_   )
           return self

       return Out()

    return extract


 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    UpperCAmelCase__  :Any            =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
    UpperCAmelCase__  :Optional[Any]            =			self.dummy_cond_unet
    UpperCAmelCase__  :Union[str, Any]            =			DDIMScheduler(
        beta_start=0.0_00_85 ,					beta_end=0.0_12 ,					beta_schedule='''scaled_linear''' ,					clip_sample=lowercase_ ,					set_alpha_to_one=lowercase_ ,					)

    UpperCAmelCase__  :Optional[Any]            =			self.dummy_vae
    UpperCAmelCase__  :int            =			self.dummy_text_encoder
    UpperCAmelCase__  :Any            =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

    # make sure here that pndm scheduler skips prk
    UpperCAmelCase__  :str            =			StableDiffusionPipeline(
        unet=lowercase_ ,					scheduler=lowercase_ ,					vae=lowercase_ ,					text_encoder=lowercase_ ,					tokenizer=lowercase_ ,					safety_checker=lowercase_ ,					feature_extractor=self.dummy_extractor ,					)
    UpperCAmelCase__  :int            =			sd_pipe.to(lowercase_   )
    sd_pipe.set_progress_bar_config(disable=lowercase_   )

    UpperCAmelCase__  :Optional[int]            =			'''A painting of a squirrel eating a burger'''

    UpperCAmelCase__  :Tuple            =			torch.Generator(device=lowercase_   ).manual_seed(0   )
    UpperCAmelCase__  :Union[str, Any]            =			sd_pipe([prompt] ,					generator=lowercase_ ,					guidance_scale=6.0 ,					num_inference_steps=2 ,					output_type='''np'''   )
    UpperCAmelCase__  :Optional[int]            =			output.images

    UpperCAmelCase__  :int            =			torch.Generator(device=lowercase_   ).manual_seed(0   )
    UpperCAmelCase__  :str            =			sd_pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=6.0 ,					num_inference_steps=2 ,					output_type='''np''' ,					return_dict=lowercase_ ,					)[0]

    UpperCAmelCase__  :List[Any]            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :Union[str, Any]            =			image_from_tuple[0, -3:, -3:, -1]

    assert image.shape == (1, 6_4, 6_4, 3)
    UpperCAmelCase__  :Union[str, Any]            =			np.array([0.57_53, 0.61_14, 0.50_01, 0.50_34, 0.54_70, 0.47_29, 0.49_71, 0.48_67, 0.48_67]   )

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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Tuple   ):
    UpperCAmelCase__  :Optional[int]            =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
    UpperCAmelCase__  :Dict            =			self.dummy_cond_unet
    UpperCAmelCase__  :Optional[Any]            =			PNDMScheduler(skip_prk_steps=lowercase_   )
    UpperCAmelCase__  :Optional[int]            =			self.dummy_vae
    UpperCAmelCase__  :Tuple            =			self.dummy_text_encoder
    UpperCAmelCase__  :List[str]            =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

    # make sure here that pndm scheduler skips prk
    UpperCAmelCase__  :List[Any]            =			StableDiffusionPipeline(
        unet=lowercase_ ,					scheduler=lowercase_ ,					vae=lowercase_ ,					text_encoder=lowercase_ ,					tokenizer=lowercase_ ,					safety_checker=lowercase_ ,					feature_extractor=self.dummy_extractor ,					)
    UpperCAmelCase__  :Dict            =			sd_pipe.to(lowercase_   )
    sd_pipe.set_progress_bar_config(disable=lowercase_   )

    UpperCAmelCase__  :Optional[int]            =			'''A painting of a squirrel eating a burger'''
    UpperCAmelCase__  :Dict            =			torch.Generator(device=lowercase_   ).manual_seed(0   )
    UpperCAmelCase__  :Union[str, Any]            =			sd_pipe([prompt] ,					generator=lowercase_ ,					guidance_scale=6.0 ,					num_inference_steps=2 ,					output_type='''np'''   )

    UpperCAmelCase__  :Any            =			output.images

    UpperCAmelCase__  :List[Any]            =			torch.Generator(device=lowercase_   ).manual_seed(0   )
    UpperCAmelCase__  :int            =			sd_pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=6.0 ,					num_inference_steps=2 ,					output_type='''np''' ,					return_dict=lowercase_ ,					)[0]

    UpperCAmelCase__  :Dict            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :Dict            =			image_from_tuple[0, -3:, -3:, -1]

    assert image.shape == (1, 6_4, 6_4, 3)
    UpperCAmelCase__  :Optional[Any]            =			np.array([0.51_22, 0.57_12, 0.48_25, 0.50_53, 0.56_46, 0.47_69, 0.51_79, 0.48_94, 0.49_94]   )

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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[str]   ):
    UpperCAmelCase__  :int            =			StableDiffusionPipeline.from_pretrained(
        '''hf-internal-testing/tiny-stable-diffusion-lms-pipe''' ,					safety_checker=lowercase_   )
    assert isinstance(lowercase_ ,					lowercase_   )
    assert isinstance(pipe.scheduler ,					lowercase_   )
    assert pipe.safety_checker is None

    UpperCAmelCase__  :List[Any]            =			pipe('''example prompt''' ,					num_inference_steps=2   ).images[0]
    assert image is not None

    # check that there's no error when saving a pipeline with one of the models being None
    with tempfile.TemporaryDirectory() as tmpdirname:
       pipe.save_pretrained(lowercase_   )
       UpperCAmelCase__  :List[str]            =			StableDiffusionPipeline.from_pretrained(lowercase_   )

    # sanity check that the pipeline still works
    assert pipe.safety_checker is None
    UpperCAmelCase__  :str            =			pipe('''example prompt''' ,					num_inference_steps=2   ).images[0]
    assert image is not None





 @unittest.skipIf(torch_device != '''cuda''' ,					'''This test requires a GPU'''   )
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    UpperCAmelCase__  :Union[str, Any]            =			self.dummy_cond_unet
    UpperCAmelCase__  :int            =			PNDMScheduler(skip_prk_steps=lowercase_   )
    UpperCAmelCase__  :Dict            =			self.dummy_vae
    UpperCAmelCase__  :Optional[int]            =			self.dummy_text_encoder
    UpperCAmelCase__  :Optional[Any]            =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

    # put models in fp16
    UpperCAmelCase__  :List[Any]            =			unet.half()
    UpperCAmelCase__  :int            =			vae.half()
    UpperCAmelCase__  :int            =			bert.half()

    # make sure here that pndm scheduler skips prk
    UpperCAmelCase__  :Union[str, Any]            =			StableDiffusionPipeline(
        unet=lowercase_ ,					scheduler=lowercase_ ,					vae=lowercase_ ,					text_encoder=lowercase_ ,					tokenizer=lowercase_ ,					safety_checker=lowercase_ ,					feature_extractor=self.dummy_extractor ,					)
    UpperCAmelCase__  :Optional[int]            =			sd_pipe.to(lowercase_   )
    sd_pipe.set_progress_bar_config(disable=lowercase_   )

    UpperCAmelCase__  :Dict            =			'''A painting of a squirrel eating a burger'''
    UpperCAmelCase__  :List[str]            =			sd_pipe([prompt] ,					num_inference_steps=2 ,					output_type='''np'''   ).images

    assert image.shape == (1, 6_4, 6_4, 3)




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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any]   ):
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()


 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    UpperCAmelCase__  :Optional[Any]            =			'''cuda'''
    UpperCAmelCase__  :str            =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''   )
    UpperCAmelCase__  :int            =			pipe.to(lowercase_   )
    pipe.set_progress_bar_config(disable=lowercase_   )

    UpperCAmelCase__  :Optional[Any]            =			'''a photo of a cat'''
    UpperCAmelCase__  :Union[str, Any]            =			{
        '''editing_prompt''': ['''sunglasses'''],
        '''reverse_editing_direction''': [False],
        '''edit_warmup_steps''': 1_0,
        '''edit_guidance_scale''': 6,
        '''edit_threshold''': 0.95,
        '''edit_momentum_scale''': 0.5,
        '''edit_mom_beta''': 0.6,
    }

    UpperCAmelCase__  :Optional[int]            =			3
    UpperCAmelCase__  :Union[str, Any]            =			7

    # no sega enabled
    UpperCAmelCase__  :int            =			torch.Generator(lowercase_   )
    generator.manual_seed(lowercase_   )
    UpperCAmelCase__  :Optional[Any]            =			pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=lowercase_ ,					num_inference_steps=5_0 ,					output_type='''np''' ,					width=5_1_2 ,					height=5_1_2 ,					)

    UpperCAmelCase__  :List[Any]            =			output.images
    UpperCAmelCase__  :Dict            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :Optional[Any]            =			[
        0.34_67_31_13,
        0.38_49_27_33,
        0.37_59_73_52,
        0.34_08_63_35,
        0.35_65_07_48,
        0.35_57_92_05,
        0.3_38_47_63,
        0.34_34_02_36,
        0.3_57_32_71,
    ]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)

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

    # with sega enabled
    # generator = torch.manual_seed(seed)
    generator.manual_seed(lowercase_   )
    UpperCAmelCase__  :Tuple            =			pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=lowercase_ ,					num_inference_steps=5_0 ,					output_type='''np''' ,					width=5_1_2 ,					height=5_1_2 ,					**lowercase_ ,					)

    UpperCAmelCase__  :Union[str, Any]            =			output.images
    UpperCAmelCase__  :Union[str, Any]            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :List[str]            =			[
        0.41_88_78_26,
        0.37_72_87_66,
        0.30_13_82_72,
        0.41_41_63_35,
        0.41_66_49_85,
        0.36_28_33_92,
        0.36_19_12_46,
        0.43_36_44_65,
        0.43_00_17_32,
    ]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)

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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    UpperCAmelCase__  :Any            =			'''cuda'''
    UpperCAmelCase__  :Any            =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''   )
    UpperCAmelCase__  :List[Any]            =			pipe.to(lowercase_   )
    pipe.set_progress_bar_config(disable=lowercase_   )

    UpperCAmelCase__  :str            =			'''an image of a crowded boulevard, realistic, 4k'''
    UpperCAmelCase__  :int            =			{
        '''editing_prompt''': '''crowd, crowded, people''',
        '''reverse_editing_direction''': True,
        '''edit_warmup_steps''': 1_0,
        '''edit_guidance_scale''': 8.3,
        '''edit_threshold''': 0.9,
        '''edit_momentum_scale''': 0.5,
        '''edit_mom_beta''': 0.6,
    }

    UpperCAmelCase__  :Any            =			9
    UpperCAmelCase__  :int            =			7

    # no sega enabled
    UpperCAmelCase__  :int            =			torch.Generator(lowercase_   )
    generator.manual_seed(lowercase_   )
    UpperCAmelCase__  :Union[str, Any]            =			pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=lowercase_ ,					num_inference_steps=5_0 ,					output_type='''np''' ,					width=5_1_2 ,					height=5_1_2 ,					)

    UpperCAmelCase__  :Union[str, Any]            =			output.images
    UpperCAmelCase__  :List[Any]            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :Union[str, Any]            =			[
        0.43_49_79_98,
        0.91_81_40_65,
        0.7_54_07_39,
        0.55_58_02_05,
        0.8_46_72_65,
        0.5_38_96_91,
        0.62_57_45_06,
        0.58_89_77_63,
        0.50_92_67_57,
    ]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)

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

    # with sega enabled
    # generator = torch.manual_seed(seed)
    generator.manual_seed(lowercase_   )
    UpperCAmelCase__  :str            =			pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=lowercase_ ,					num_inference_steps=5_0 ,					output_type='''np''' ,					width=5_1_2 ,					height=5_1_2 ,					**lowercase_ ,					)

    UpperCAmelCase__  :int            =			output.images
    UpperCAmelCase__  :int            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :Tuple            =			[
        0.3_08_97_19,
        0.30_50_01_44,
        0.29_01_60_42,
        0.30_63_09_64,
        0.32_56_87,
        0.29_41_92_25,
        0.2_90_80_91,
        0.28_72_35_98,
        0.27_69_62_94,
    ]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)

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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    UpperCAmelCase__  :List[str]            =			'''cuda'''
    UpperCAmelCase__  :Optional[Any]            =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5'''   )
    UpperCAmelCase__  :Dict            =			pipe.to(lowercase_   )
    pipe.set_progress_bar_config(disable=lowercase_   )

    UpperCAmelCase__  :str            =			'''a castle next to a river'''
    UpperCAmelCase__  :str            =			{
        '''editing_prompt''': ['''boat on a river, boat''', '''monet, impression, sunrise'''],
        '''reverse_editing_direction''': False,
        '''edit_warmup_steps''': [1_5, 1_8],
        '''edit_guidance_scale''': 6,
        '''edit_threshold''': [0.9, 0.8],
        '''edit_momentum_scale''': 0.5,
        '''edit_mom_beta''': 0.6,
    }

    UpperCAmelCase__  :str            =			4_8
    UpperCAmelCase__  :Dict            =			7

    # no sega enabled
    UpperCAmelCase__  :str            =			torch.Generator(lowercase_   )
    generator.manual_seed(lowercase_   )
    UpperCAmelCase__  :Optional[Any]            =			pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=lowercase_ ,					num_inference_steps=5_0 ,					output_type='''np''' ,					width=5_1_2 ,					height=5_1_2 ,					)

    UpperCAmelCase__  :List[str]            =			output.images
    UpperCAmelCase__  :str            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :Union[str, Any]            =			[
        0.75_16_35_55,
        0.76_03_71_45,
        0.6_17_85,
        0.9_18_96_73,
        0.8_62_77_01,
        0.85_18_96_94,
        0.8_51_28_13,
        0.87_01_20_76,
        0.8_31_28_57,
    ]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)

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

    # with sega enabled
    # generator = torch.manual_seed(seed)
    generator.manual_seed(lowercase_   )
    UpperCAmelCase__  :Optional[Any]            =			pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=lowercase_ ,					num_inference_steps=5_0 ,					output_type='''np''' ,					width=5_1_2 ,					height=5_1_2 ,					**lowercase_ ,					)

    UpperCAmelCase__  :Tuple            =			output.images
    UpperCAmelCase__  :List[str]            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :str            =			[
        0.73_55_33_65,
        0.7_53_72_71,
        0.74_34_19_05,
        0.66_48_03_56,
        0.6_47_29_25,
        0.63_03_94_16,
        0.64_81_29_05,
        0.6_74_97_17,
        0.6_51_71_02,
    ]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)

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





 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    UpperCAmelCase__  :Tuple            =			'''cuda'''
    UpperCAmelCase__  :List[Any]            =			StableDiffusionPipeline.from_pretrained('''runwayml/stable-diffusion-v1-5''' ,					torch_dtype=torch.floataa   )
    UpperCAmelCase__  :str            =			pipe.to(lowercase_   )
    pipe.set_progress_bar_config(disable=lowercase_   )

    UpperCAmelCase__  :Any            =			'''a photo of a cat'''
    UpperCAmelCase__  :int            =			{
        '''editing_prompt''': ['''sunglasses'''],
        '''reverse_editing_direction''': [False],
        '''edit_warmup_steps''': 1_0,
        '''edit_guidance_scale''': 6,
        '''edit_threshold''': 0.95,
        '''edit_momentum_scale''': 0.5,
        '''edit_mom_beta''': 0.6,
    }

    UpperCAmelCase__  :Any            =			3
    UpperCAmelCase__  :str            =			7

    # no sega enabled
    UpperCAmelCase__  :Any            =			torch.Generator(lowercase_   )
    generator.manual_seed(lowercase_   )
    UpperCAmelCase__  :List[Any]            =			pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=lowercase_ ,					num_inference_steps=5_0 ,					output_type='''np''' ,					width=5_1_2 ,					height=5_1_2 ,					)

    UpperCAmelCase__  :Union[str, Any]            =			output.images
    UpperCAmelCase__  :Any            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :Tuple            =			[
        0.34_88_76_95,
        0.3_87_69_53,
        0.3_75,
        0.34_42_38_28,
        0.3_58_15_43,
        0.35_71_77_73,
        0.3_38_37_89,
        0.34_57_03_12,
        0.35_93_75,
    ]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)

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

    # with sega enabled
    # generator = torch.manual_seed(seed)
    generator.manual_seed(lowercase_   )
    UpperCAmelCase__  :str            =			pipe(
        [prompt] ,					generator=lowercase_ ,					guidance_scale=lowercase_ ,					num_inference_steps=5_0 ,					output_type='''np''' ,					width=5_1_2 ,					height=5_1_2 ,					**lowercase_ ,					)

    UpperCAmelCase__  :List[str]            =			output.images
    UpperCAmelCase__  :Any            =			image[0, -3:, -3:, -1]
    UpperCAmelCase__  :str            =			[
        0.42_28_51_56,
        0.36_91_40_62,
        0.29_07_71_48,
        0.42_04_10_16,
        0.41_91_89_45,
        0.35_49_80_47,
        0.3_61_81_64,
        0.4_42_38_28,
        0.43_11_52_34,
    ]

    assert image.shape == (1, 5_1_2, 5_1_2, 3)

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







'''simple docstring'''





import unittest
from unittest.util import safe_repr

from transformers import AutoTokenizer, RwkvConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
        RwkvForCausalLM,
        RwkvModel,
    )
    from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_0
else:
    __lowerCamelCase       				=	False




class 							UpperCAmelCase :


 def __init__(			self       :  int ,					__lowerCamelCase       :  str ,					__lowerCamelCase       :  int=1_4 ,					__lowerCamelCase       :  Tuple=7 ,					__lowerCamelCase       :  int=True ,					__lowerCamelCase       :  Tuple=False ,					__lowerCamelCase       :  Any=True ,					__lowerCamelCase       :  List[Any]=True ,					__lowerCamelCase       :  Dict=True ,					__lowerCamelCase       :  List[Any]=9_9 ,					__lowerCamelCase       :  Dict=3_2 ,					__lowerCamelCase       :  Optional[int]=5 ,					__lowerCamelCase       :  Union[str, Any]=3_7 ,					__lowerCamelCase       :  Dict="gelu" ,					__lowerCamelCase       :  Any=0.1 ,					__lowerCamelCase       :  Any=0.1 ,					__lowerCamelCase       :  Optional[int]=5_1_2 ,					__lowerCamelCase       :  Tuple=1_6 ,					__lowerCamelCase       :  str=2 ,					__lowerCamelCase       :  str=3 ,					__lowerCamelCase       :  Any=4 ,					__lowerCamelCase       :  Tuple=None ,					):
    UpperCAmelCase__  :Any            =			parent
    UpperCAmelCase__  :Optional[Any]            =			batch_size
    UpperCAmelCase__  :str            =			seq_length
    UpperCAmelCase__  :Optional[int]            =			is_training
    UpperCAmelCase__  :Optional[Any]            =			use_token_type_ids
    UpperCAmelCase__  :Dict            =			use_input_mask
    UpperCAmelCase__  :List[Any]            =			use_labels
    UpperCAmelCase__  :List[str]            =			use_mc_token_ids
    UpperCAmelCase__  :List[Any]            =			vocab_size
    UpperCAmelCase__  :Dict            =			hidden_size
    UpperCAmelCase__  :Union[str, Any]            =			num_hidden_layers
    UpperCAmelCase__  :List[Any]            =			intermediate_size
    UpperCAmelCase__  :int            =			hidden_act
    UpperCAmelCase__  :Optional[Any]            =			hidden_dropout_prob
    UpperCAmelCase__  :Optional[int]            =			attention_probs_dropout_prob
    UpperCAmelCase__  :str            =			max_position_embeddings
    UpperCAmelCase__  :int            =			type_vocab_size
    UpperCAmelCase__  :List[Any]            =			type_sequence_label_size
    UpperCAmelCase__  :int            =			num_labels
    UpperCAmelCase__  :str            =			num_choices
    UpperCAmelCase__  :Optional[int]            =			scope
    UpperCAmelCase__  :Optional[Any]            =			vocab_size - 1
    UpperCAmelCase__  :str            =			vocab_size - 1
    UpperCAmelCase__  :Tuple            =			vocab_size - 1


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any]   ):
    return RwkvConfig.from_pretrained('''sgugger/rwkv-4-pile-7b'''   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Tuple ,					__lowerCamelCase       :  Dict=False ,					__lowerCamelCase       :  Dict=False ,					__lowerCamelCase       :  List[str]=False   ):
    UpperCAmelCase__  :Union[str, Any]            =			ids_tensor([self.batch_size, self.seq_length] ,					self.vocab_size   )

    UpperCAmelCase__  :Dict            =			None
    if self.use_input_mask:
       UpperCAmelCase__  :Optional[int]            =			random_attention_mask([self.batch_size, self.seq_length]   )

    UpperCAmelCase__  :Tuple            =			None
    if self.use_token_type_ids:
       UpperCAmelCase__  :Any            =			ids_tensor([self.batch_size, self.seq_length] ,					self.type_vocab_size   )

    UpperCAmelCase__  :str            =			None
    if self.use_mc_token_ids:
       UpperCAmelCase__  :Optional[Any]            =			ids_tensor([self.batch_size, self.num_choices] ,					self.seq_length   )

    UpperCAmelCase__  :Optional[Any]            =			None
    UpperCAmelCase__  :List[Any]            =			None
    UpperCAmelCase__  :Any            =			None
    if self.use_labels:
       UpperCAmelCase__  :Optional[Any]            =			ids_tensor([self.batch_size] ,					self.type_sequence_label_size   )
       UpperCAmelCase__  :int            =			ids_tensor([self.batch_size, self.seq_length] ,					self.num_labels   )
       UpperCAmelCase__  :Tuple            =			ids_tensor([self.batch_size] ,					self.num_choices   )

    UpperCAmelCase__  :Optional[int]            =			self.get_config(
        gradient_checkpointing=__lowerCamelCase ,					scale_attn_by_inverse_layer_idx=__lowerCamelCase ,					reorder_and_upcast_attn=__lowerCamelCase ,					)

    return (
        config,
        input_ids,
        input_mask,
        None,
        token_type_ids,
        mc_token_ids,
        sequence_labels,
        token_labels,
        choice_labels,
    )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str ,					__lowerCamelCase       :  Union[str, Any]=False ,					__lowerCamelCase       :  int=False ,					__lowerCamelCase       :  int=False   ):
    return RwkvConfig(
        vocab_size=self.vocab_size ,					hidden_size=self.hidden_size ,					num_hidden_layers=self.num_hidden_layers ,					intermediate_size=self.intermediate_size ,					activation_function=self.hidden_act ,					resid_pdrop=self.hidden_dropout_prob ,					attn_pdrop=self.attention_probs_dropout_prob ,					n_positions=self.max_position_embeddings ,					type_vocab_size=self.type_vocab_size ,					use_cache=__lowerCamelCase ,					bos_token_id=self.bos_token_id ,					eos_token_id=self.eos_token_id ,					pad_token_id=self.pad_token_id ,					gradient_checkpointing=__lowerCamelCase ,					scale_attn_by_inverse_layer_idx=__lowerCamelCase ,					reorder_and_upcast_attn=__lowerCamelCase ,					)


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    UpperCAmelCase__  :int            =			self.get_config()
    UpperCAmelCase__  :Any            =			3_0_0
    return config


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Tuple   ):
    (
        (
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	
    )  :Dict            =			self.prepare_config_and_inputs()

    UpperCAmelCase__  :List[str]            =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
    UpperCAmelCase__  :Dict            =			ids_tensor([self.batch_size, self.seq_length] ,					vocab_size=2   )

    return (
        config,
        input_ids,
        input_mask,
        head_mask,
        token_type_ids,
        sequence_labels,
        token_labels,
        choice_labels,
        encoder_hidden_states,
        encoder_attention_mask,
    )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any] ,					__lowerCamelCase       :  Optional[Any] ,					__lowerCamelCase       :  Tuple ,					__lowerCamelCase       :  Dict ,					__lowerCamelCase       :  Union[str, Any] ,					__lowerCamelCase       :  Optional[int] ,					*__lowerCamelCase       :  Optional[Any]   ):
    UpperCAmelCase__  :str            =			True
    UpperCAmelCase__  :Optional[Any]            =			RwkvModel(config=__lowerCamelCase   )
    model.to(__lowerCamelCase   )
    model.eval()

    UpperCAmelCase__  :Optional[int]            =			model(__lowerCamelCase   )

    self.parent.assertEqual(result.last_hidden_state.shape ,					(self.batch_size, self.seq_length, self.hidden_size)   )
    self.parent.assertEqual(len(result.hidden_states   ) ,					config.num_hidden_layers + 1   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Dict ,					__lowerCamelCase       :  Tuple ,					__lowerCamelCase       :  Dict ,					__lowerCamelCase       :  Dict ,					__lowerCamelCase       :  Optional[int] ,					__lowerCamelCase       :  Union[str, Any] ,					*__lowerCamelCase       :  List[Any]   ):
    UpperCAmelCase__  :int            =			RwkvForCausalLM(__lowerCamelCase   )
    model.to(__lowerCamelCase   )
    model.eval()

    UpperCAmelCase__  :int            =			model(__lowerCamelCase ,					labels=__lowerCamelCase   )
    self.parent.assertEqual(result.loss.shape ,					()   )
    self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.seq_length, self.vocab_size)   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int] ,					__lowerCamelCase       :  Optional[int] ,					__lowerCamelCase       :  Tuple ,					__lowerCamelCase       :  Optional[int] ,					__lowerCamelCase       :  Union[str, Any] ,					__lowerCamelCase       :  Tuple ,					*__lowerCamelCase       :  Optional[Any]   ):
    UpperCAmelCase__  :Optional[Any]            =			RwkvModel(config=__lowerCamelCase   )
    model.to(__lowerCamelCase   )
    model.eval()

    UpperCAmelCase__  :Tuple            =			model(__lowerCamelCase   )
    UpperCAmelCase__  :int            =			outputs.last_hidden_state

    UpperCAmelCase__  :Tuple            =			model(input_ids[:, :2]   )
    UpperCAmelCase__  :List[Any]            =			outputs.last_hidden_state

    # Using the state computed on the first inputs, we will get the same output
    UpperCAmelCase__  :Any            =			model(input_ids[:, 2:] ,					state=outputs.state   )
    UpperCAmelCase__  :Any            =			outputs.last_hidden_state

    self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two] ,					dim=1   ) ,					__lowerCamelCase ,					atol=1e-5   )   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str ,					__lowerCamelCase       :  int ,					__lowerCamelCase       :  Optional[int] ,					__lowerCamelCase       :  List[Any] ,					__lowerCamelCase       :  List[str] ,					__lowerCamelCase       :  Any ,					*__lowerCamelCase       :  Any ,					__lowerCamelCase       :  Tuple=False   ):
    UpperCAmelCase__  :Any            =			RwkvForCausalLM(__lowerCamelCase   )
    model.to(__lowerCamelCase   )
    if gradient_checkpointing:
       model.gradient_checkpointing_enable()

    UpperCAmelCase__  :Union[str, Any]            =			model(__lowerCamelCase ,					labels=__lowerCamelCase   )
    self.parent.assertEqual(result.loss.shape ,					()   )
    self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.seq_length, self.vocab_size)   )
    result.loss.backward()





 def 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    UpperCAmelCase__  :Optional[Any]            =			self.prepare_config_and_inputs()

    (
        (
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	(
        UpperCAmelCase__
    )      ,	
    )  :Any            =			config_and_inputs

    UpperCAmelCase__  :int            =			{'''input_ids''': input_ids}

    return config, inputs_dict




@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0	,   reason="See https://github.com/huggingface/transformers/pull/24204"	)
@require_torch
class 							UpperCAmelCase ( _snake_case	,   _snake_case	,   _snake_case	,   unittest.TestCase	):
 UpperCAmelCase				=				(RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
 UpperCAmelCase				=				(
     {"feature-extraction": RwkvModel, "text-generation": RwkvForCausalLM} if is_torch_available() else {}
 )
 # all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
 UpperCAmelCase				=				False
 UpperCAmelCase				=				False
 UpperCAmelCase				=				False
 UpperCAmelCase				=				False
 UpperCAmelCase				=				False  # Rwkv does not support head masking


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int]   ):
    UpperCAmelCase__  :Dict            =			RwkvModelTester(self   )
    UpperCAmelCase__  :Union[str, Any]            =			ConfigTester(
        self ,					config_class=__lowerCamelCase ,					n_embd=3_7 ,					common_properties=['''hidden_size''', '''num_hidden_layers''']   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Dict ,					__lowerCamelCase       :  Optional[Any] ,					__lowerCamelCase       :  int ,					__lowerCamelCase       :  Optional[int]=None   ):
    if isinstance(__lowerCamelCase ,					torch.Tensor   ):
       UpperCAmelCase__      ,	UpperCAmelCase__  :Dict            =			member.max().item(), member.min().item()
    elif isinstance(__lowerCamelCase ,					__lowerCamelCase   ) or isinstance(__lowerCamelCase ,					__lowerCamelCase   ):
       UpperCAmelCase__      ,	UpperCAmelCase__  :List[Any]            =			max(__lowerCamelCase   ), min(__lowerCamelCase   )

    if not isinstance(__lowerCamelCase ,					__lowerCamelCase   ):
       raise TypeError('''container should be a list or tuple'''   )
    elif len(__lowerCamelCase   ) != 2:
       raise ValueError('''container should have 2 elements'''   )

    UpperCAmelCase__      ,	UpperCAmelCase__  :Any            =			container

    UpperCAmelCase__  :List[str]            =			(min_value >= expected_min) and (max_value <= expected_max)

    if not is_inside_interval:
       UpperCAmelCase__  :int            =			'''%s not found in %s''' % (safe_repr(__lowerCamelCase   ), safe_repr(__lowerCamelCase   ))
       self.fail(self._formatMessage(__lowerCamelCase ,					__lowerCamelCase   )   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    self.config_tester.run_common_tests()


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any]   ):
    UpperCAmelCase__  :Tuple            =			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_rwkv_model(*__lowerCamelCase   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  int   ):
    UpperCAmelCase__  :str            =			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_causl_lm(*__lowerCamelCase   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    UpperCAmelCase__  :Union[str, Any]            =			self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_state_equivalency(*__lowerCamelCase   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int]   ):
    UpperCAmelCase__      ,	UpperCAmelCase__  :Dict            =			self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
       UpperCAmelCase__  :Optional[int]            =			model_class(config=__lowerCamelCase   )
       for name, param in model.named_parameters():
          if "time_decay" in name:
             if param.requires_grad:
                self.assertTrue(param.data.max().item() == 3.0   )
                self.assertTrue(param.data.min().item() == -5.0   )
          elif "time_first" in name:
             if param.requires_grad:
                # check if it's a ones like
                self.assertTrue(torch.allclose(param.data ,					torch.ones_like(param.data   ) ,					atol=1e-5 ,					rtol=1e-5   )   )
          elif any(x in name for x in ['''time_mix_key''', '''time_mix_receptance''']   ):
             if param.requires_grad:
                self.assertInterval(
                    param.data ,					[0.0, 1.0] ,					msg=f'''Parameter {name} of model {model_class} seems not properly initialized''' ,					)
          elif "time_mix_value" in name:
             if param.requires_grad:
                self.assertInterval(
                    param.data ,					[0.0, 1.3] ,					msg=f'''Parameter {name} of model {model_class} seems not properly initialized''' ,					)


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    UpperCAmelCase__      ,	UpperCAmelCase__  :int            =			self.model_tester.prepare_config_and_inputs_for_common()
    UpperCAmelCase__  :Dict            =			True

    UpperCAmelCase__  :Dict            =			getattr(self.model_tester ,					'''seq_length''' ,					__lowerCamelCase   )

    for model_class in self.all_model_classes:
       UpperCAmelCase__  :Tuple            =			True
       UpperCAmelCase__  :Optional[int]            =			False
       UpperCAmelCase__  :List[Any]            =			True
       UpperCAmelCase__  :Optional[int]            =			model_class(__lowerCamelCase   )
       model.to(__lowerCamelCase   )
       model.eval()

       UpperCAmelCase__  :List[Any]            =			self._prepare_for_class(__lowerCamelCase ,					__lowerCamelCase   )
       UpperCAmelCase__  :Tuple            =			inputs['''input_ids'''].shape[0]
       with torch.no_grad():
          UpperCAmelCase__  :List[str]            =			model(**__lowerCamelCase   )
       UpperCAmelCase__  :int            =			outputs.attentions
       self.assertEqual(len(__lowerCamelCase   ) ,					self.model_tester.num_hidden_layers   )

       # check that output_attentions also work using config
       del inputs_dict["output_attentions"]
       UpperCAmelCase__  :Optional[Any]            =			True
       UpperCAmelCase__  :Union[str, Any]            =			model_class(__lowerCamelCase   )
       model.to(__lowerCamelCase   )
       model.eval()

       UpperCAmelCase__  :List[str]            =			self._prepare_for_class(__lowerCamelCase ,					__lowerCamelCase   )
       UpperCAmelCase__  :Optional[int]            =			inputs['''input_ids'''].shape[0]
       with torch.no_grad():
          UpperCAmelCase__  :Tuple            =			model(**__lowerCamelCase   )
       UpperCAmelCase__  :Union[str, Any]            =			outputs.attentions
       self.assertEqual(len(__lowerCamelCase   ) ,					self.model_tester.num_hidden_layers   )

       self.assertListEqual(
           list(attentions[0].shape[-3:]   ) ,					[batch_size, seq_len, config.hidden_size] ,					)
       UpperCAmelCase__  :Dict            =			len(__lowerCamelCase   )

       # Check attention is always last and order is fine
       UpperCAmelCase__  :int            =			True
       UpperCAmelCase__  :Tuple            =			True
       UpperCAmelCase__  :List[Any]            =			model_class(__lowerCamelCase   )
       model.to(__lowerCamelCase   )
       model.eval()

       UpperCAmelCase__  :Optional[int]            =			self._prepare_for_class(__lowerCamelCase ,					__lowerCamelCase   )
       UpperCAmelCase__  :int            =			inputs['''input_ids'''].shape[0]
       with torch.no_grad():
          UpperCAmelCase__  :Optional[int]            =			model(**__lowerCamelCase   )

       UpperCAmelCase__  :List[Any]            =			1
       self.assertEqual(out_len + added_hidden_states ,					len(__lowerCamelCase   )   )

       UpperCAmelCase__  :str            =			outputs.attentions

       self.assertEqual(len(__lowerCamelCase   ) ,					self.model_tester.num_hidden_layers   )
       self.assertListEqual(
           list(self_attentions[0].shape[-3:]   ) ,					[batch_size, seq_len, config.hidden_size] ,					)





 @slow
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Dict   ):
    for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       UpperCAmelCase__  :List[Any]            =			RwkvModel.from_pretrained(__lowerCamelCase   )
       self.assertIsNotNone(__lowerCamelCase   )




@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0	,   reason="See https://github.com/huggingface/transformers/pull/24204"	)
@slow
class 							UpperCAmelCase ( unittest.TestCase	):


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    UpperCAmelCase__  :List[Any]            =			'''RWKV/rwkv-4-169m-pile'''
    UpperCAmelCase__  :str            =			AutoTokenizer.from_pretrained(self.model_id   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Any   ):
    UpperCAmelCase__  :Optional[int]            =			'''Hello my name is Jasmine and I am a newbie to the'''
    UpperCAmelCase__  :Union[str, Any]            =			RwkvForCausalLM.from_pretrained(self.model_id   ).to(__lowerCamelCase   )

    UpperCAmelCase__  :List[str]            =			self.tokenizer('''Hello my name is''' ,					return_tensors='''pt'''   ).input_ids.to(__lowerCamelCase   )
    UpperCAmelCase__  :str            =			model.generate(__lowerCamelCase ,					max_new_tokens=1_0   )
    UpperCAmelCase__  :str            =			self.tokenizer.decode(output[0].tolist()   )

    self.assertEqual(__lowerCamelCase ,					__lowerCamelCase   )





 def 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    UpperCAmelCase__  :str            =			'''Hello my name is Jasmine and I am a newbie to the'''

    UpperCAmelCase__  :int            =			self.tokenizer('''Hello my name is''' ,					return_tensors='''pt'''   ).input_ids.to(__lowerCamelCase   )
    UpperCAmelCase__  :Dict            =			RwkvForCausalLM.from_pretrained(self.model_id ,					torch_dtype=torch.bfloataa   ).to(__lowerCamelCase   )

    UpperCAmelCase__  :str            =			model.generate(__lowerCamelCase ,					max_new_tokens=1_0   )
    UpperCAmelCase__  :int            =			self.tokenizer.decode(output[0].tolist()   )

    self.assertEqual(__lowerCamelCase ,					__lowerCamelCase   )



