


'''simple docstring'''



from __future__ import annotations

from collections import deque





class 				__UpperCamelCase							:




	def __init__(  self		,							__a				):

						'''simple docstring'''





						__a     :			list[dict]          =    []
						self.adlist.append(
						    {'value': '', 'next_states': [], 'fail_state': 0, 'output': []}				)

						for keyword in keywords:
											self.add_keyword(__lowerCamelCase				)
						self.set_fail_transitions()




	def  __UpperCAmelCase				(  self		,							__a		,							__a				):

						'''simple docstring'''





						for state in self.adlist[current_state]["next_states"]:
											if char == self.adlist[state]["value"]:
																return state
						return None




	def  __UpperCAmelCase				(  self		,							__a				):

						'''simple docstring'''





						__a     :			Tuple          =    0
						for character in keyword:
											__a     :			int          =    self.find_next_state(__lowerCamelCase		,							__lowerCamelCase				)
											if next_state is None:
																self.adlist.append(
																    {
																        'value': character,
																        'next_states': [],
																        'fail_state': 0,
																        'output': [],
																    }				)
																self.adlist[current_state]["next_states"].append(len(self.adlist				) - 1				)
																__a     :			List[str]          =    len(self.adlist				) - 1
											else:
																__a     :			int          =    next_state
						self.adlist[current_state]["output"].append(__lowerCamelCase				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			deque          =    deque()
						for node in self.adlist[0]["next_states"]:
											q.append(__lowerCamelCase				)
											__a     :			List[Any]          =    0
						while q:
											__a     :			Dict          =    q.popleft()
											for child in self.adlist[r]["next_states"]:
																q.append(__lowerCamelCase				)
																__a     :			List[Any]          =    self.adlist[r]['''fail_state''']
																while (
																    self.find_next_state(__lowerCamelCase		,							self.adlist[child]['value']				) is None
																    and state != 0
																):
																					__a     :			Tuple          =    self.adlist[state]['''fail_state''']
																__a     :			Dict          =    self.find_next_state(
																    __lowerCamelCase		,							self.adlist[child]['value']				)
																if self.adlist[child]["fail_state"] is None:
																					__a     :			Dict          =    0
																__a     :			Optional[int]          =    (
																    self.adlist[child]['''output''']
																    + self.adlist[self.adlist[child]['''fail_state''']]['''output''']
																)





	def  __UpperCAmelCase				(  self		,							__a				):

						'''simple docstring'''





						__a     :			dict          =    {}  # returns a dict with keywords and list of its occurrences
						__a     :			List[Any]          =    0
						for i in range(len(__lowerCamelCase				)				):
											while (
											    self.find_next_state(__lowerCamelCase		,							string[i]				) is None
											    and current_state != 0
											):
																__a     :			str          =    self.adlist[current_state]['''fail_state''']
											__a     :			Union[str, Any]          =    self.find_next_state(__lowerCamelCase		,							string[i]				)
											if next_state is None:
																__a     :			str          =    0
											else:
																__a     :			Any          =    next_state
																for key in self.adlist[current_state]["output"]:
																					if key not in result:
																										__a     :			List[Any]          =    []
																					result[key].append(i - len(__lowerCamelCase				) + 1				)
						return result


if __name__ == "__main__":
	import doctest

	doctest.testmod()







'''simple docstring'''





import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


__lowerCamelCase												=						False




class 				UpperCAmelCase			(			unittest.TestCase						):


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any] ):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    return 1_2


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    return 1_2


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    return 3_2


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    torch.manual_seed(0 )
    UpperCAmelCase__	:List[str]      =      VQModel(
        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=3			,       num_vq_embeddings=self.num_embed			,       vq_embed_dim=3			,       )
    return model


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    UpperCAmelCase__	:Tuple      =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )
    return tokenizer


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple ):
    torch.manual_seed(0 )
    UpperCAmelCase__	:Optional[Any]      =      CLIPTextConfig(
        bos_token_id=0			,       eos_token_id=2			,       hidden_size=self.text_embedder_hidden_size			,       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(__lowerCamelCase )


 @property
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    torch.manual_seed(0 )

    UpperCAmelCase__	:List[str]      =      1_2
    UpperCAmelCase__	:Optional[int]      =      1_2

    UpperCAmelCase__	:Optional[int]      =      {
        '''attention_bias''': True,
        '''cross_attention_dim''': 3_2,
        '''attention_head_dim''': height * width,
        '''num_attention_heads''': 1,
        '''num_vector_embeds''': self.num_embed,
        '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
        '''norm_num_groups''': 3_2,
        '''sample_size''': width,
        '''activation_fn''': '''geglu-approximate''',
    }

    UpperCAmelCase__	:List[str]      =      TransformeraDModel(**__lowerCamelCase )
    return model


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    UpperCAmelCase__	:Union[str, Any]      =      '''cpu'''

    UpperCAmelCase__	:Optional[Any]      =      self.dummy_vqvae
    UpperCAmelCase__	:Optional[int]      =      self.dummy_text_encoder
    UpperCAmelCase__	:Union[str, Any]      =      self.dummy_tokenizer
    UpperCAmelCase__	:List[Any]      =      self.dummy_transformer
    UpperCAmelCase__	:str      =      VQDiffusionScheduler(self.num_embed )
    UpperCAmelCase__	:Optional[Any]      =      LearnedClassifierFreeSamplingEmbeddings(learnable=__lowerCamelCase )

    UpperCAmelCase__	:Any      =      VQDiffusionPipeline(
        vqvae=__lowerCamelCase			,       text_encoder=__lowerCamelCase			,       tokenizer=__lowerCamelCase			,       transformer=__lowerCamelCase			,       scheduler=__lowerCamelCase			,       learned_classifier_free_sampling_embeddings=__lowerCamelCase			,       )
    UpperCAmelCase__	:Any      =      pipe.to(__lowerCamelCase )
    pipe.set_progress_bar_config(disable=__lowerCamelCase )

    UpperCAmelCase__	:str      =      '''teddy bear playing in the pool'''

    UpperCAmelCase__	:Optional[int]      =      torch.Generator(device=__lowerCamelCase ).manual_seed(0 )
    UpperCAmelCase__	:Any      =      pipe([prompt]			,       generator=__lowerCamelCase			,       num_inference_steps=2			,       output_type='''np''' )
    UpperCAmelCase__	:str      =      output.images

    UpperCAmelCase__	:List[Any]      =      torch.Generator(device=__lowerCamelCase ).manual_seed(0 )
    UpperCAmelCase__	:Dict      =      pipe(
        [prompt]			,       generator=__lowerCamelCase			,       output_type='''np'''			,       return_dict=__lowerCamelCase			,       num_inference_steps=2 )[0]

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

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

    UpperCAmelCase__	:int      =      np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92] )

    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__	:List[Any]      =      '''cpu'''

    UpperCAmelCase__	:int      =      self.dummy_vqvae
    UpperCAmelCase__	:List[str]      =      self.dummy_text_encoder
    UpperCAmelCase__	:Union[str, Any]      =      self.dummy_tokenizer
    UpperCAmelCase__	:Optional[Any]      =      self.dummy_transformer
    UpperCAmelCase__	:List[str]      =      VQDiffusionScheduler(self.num_embed )
    UpperCAmelCase__	:List[str]      =      LearnedClassifierFreeSamplingEmbeddings(
        learnable=__lowerCamelCase			,       hidden_size=self.text_embedder_hidden_size			,       length=tokenizer.model_max_length )

    UpperCAmelCase__	:List[Any]      =      VQDiffusionPipeline(
        vqvae=__lowerCamelCase			,       text_encoder=__lowerCamelCase			,       tokenizer=__lowerCamelCase			,       transformer=__lowerCamelCase			,       scheduler=__lowerCamelCase			,       learned_classifier_free_sampling_embeddings=__lowerCamelCase			,       )
    UpperCAmelCase__	:str      =      pipe.to(__lowerCamelCase )
    pipe.set_progress_bar_config(disable=__lowerCamelCase )

    UpperCAmelCase__	:str      =      '''teddy bear playing in the pool'''

    UpperCAmelCase__	:str      =      torch.Generator(device=__lowerCamelCase ).manual_seed(0 )
    UpperCAmelCase__	:Union[str, Any]      =      pipe([prompt]			,       generator=__lowerCamelCase			,       num_inference_steps=2			,       output_type='''np''' )
    UpperCAmelCase__	:Any      =      output.images

    UpperCAmelCase__	:Optional[int]      =      torch.Generator(device=__lowerCamelCase ).manual_seed(0 )
    UpperCAmelCase__	:int      =      pipe(
        [prompt]			,       generator=__lowerCamelCase			,       output_type='''np'''			,       return_dict=__lowerCamelCase			,       num_inference_steps=2 )[0]

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

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

    UpperCAmelCase__	:Union[str, Any]      =      np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88] )

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




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


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





 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    UpperCAmelCase__	:Optional[Any]      =      load_numpy(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy''' )

    UpperCAmelCase__	:Optional[Any]      =      VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq''' )
    UpperCAmelCase__	:List[str]      =      pipeline.to(__lowerCamelCase )
    pipeline.set_progress_bar_config(disable=__lowerCamelCase )

    # requires GPU generator for gumbel softmax
    # don't use GPU generator in tests though
    UpperCAmelCase__	:int      =      torch.Generator(device=__lowerCamelCase ).manual_seed(0 )
    UpperCAmelCase__	:Union[str, Any]      =      pipeline(
        '''teddy bear playing in the pool'''			,       num_images_per_prompt=1			,       generator=__lowerCamelCase			,       output_type='''np'''			,       )

    UpperCAmelCase__	:List[str]      =      output.images[0]

    assert image.shape == (2_5_6, 2_5_6, 3)
    assert np.abs(expected_image - image ).max() < 2.0



