
import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class   lowerCamelCase	(						unittest.TestCase						):





     def 	UpperCAmelCase(self   :					str     )		->		Optional[int]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





     @property
     def 	UpperCAmelCase(self   :					Optional[Any]     )		->		int:
         snake_case				  =      1
         snake_case				  =      3
         snake_case				  =      (3_2, 3_2)

         snake_case				  =      floats_tensor((batch_size, num_channels) + sizes       ,					rng=random.Random(0     )     ).to(_A     )
         return image





     @property
     def 	UpperCAmelCase(self   :					List[str]     )		->		Tuple:
         torch.manual_seed(0     )
         snake_case				  =      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 	UpperCAmelCase(self   :					List[str]     )		->		Union[str, Any]:
         torch.manual_seed(0     )
         snake_case				  =      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 	UpperCAmelCase(self   :					int     )		->		List[Any]:
         torch.manual_seed(0     )
         snake_case				  =      RobertaSeriesConfig(
             hidden_size=3_2       ,					project_dim=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=5_0_0_6       ,					)
         return RobertaSeriesModelWithTransformation(_A     )





     @property
     def 	UpperCAmelCase(self   :					Union[str, Any]     )		->		List[Any]:
         def extract(*_A   :					int       ,					**_A   :					int     ):







             class   lowerCamelCase	:





                  def __init__(self   :					List[str]     )		->		Any:
                      snake_case				  =      torch.ones([0]     )





                  def 	UpperCAmelCase(self   :					List[str]       ,					_A   :					List[Any]     )		->		List[Any]:
                      self.pixel_values.to(_A     )
                      return self


             return Out()

         return extract





     def 	UpperCAmelCase(self   :					Dict     )		->		List[Any]:
         snake_case				  =      "cpu"  # ensure determinism for the device-dependent torch.Generator
         snake_case				  =      self.dummy_cond_unet
         snake_case				  =      PNDMScheduler(skip_prk_steps=_A     )
         snake_case				  =      self.dummy_vae
         snake_case				  =      self.dummy_text_encoder
         snake_case				  =      XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"     )
         snake_case				  =      7_7

         snake_case				  =      self.dummy_image.to(_A     )
         snake_case				  =      init_image / 2 + 0.5

         # make sure here that pndm scheduler skips prk
         snake_case				  =      AltDiffusionImgaImgPipeline(
             unet=_A       ,					scheduler=_A       ,					vae=_A       ,					text_encoder=_A       ,					tokenizer=_A       ,					safety_checker=_A       ,					feature_extractor=self.dummy_extractor       ,					)
         snake_case				  =      VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor       ,					do_normalize=_A     )
         snake_case				  =      alt_pipe.to(_A     )
         alt_pipe.set_progress_bar_config(disable=_A     )

         snake_case				  =      "A painting of a squirrel eating a burger"
         snake_case				  =      torch.Generator(device=_A     ).manual_seed(0     )
         snake_case				  =      alt_pipe(
             [prompt]       ,					generator=_A       ,					guidance_scale=6.0       ,					num_inference_steps=2       ,					output_type="np"       ,					image=_A       ,					)

         snake_case				  =      output.images

         snake_case				  =      torch.Generator(device=_A     ).manual_seed(0     )
         snake_case				  =      alt_pipe(
             [prompt]       ,					generator=_A       ,					guidance_scale=6.0       ,					num_inference_steps=2       ,					output_type="np"       ,					image=_A       ,					return_dict=_A       ,					)[0]

         snake_case				  =      image[0, -3:, -3:, -1]
         snake_case				  =      image_from_tuple[0, -3:, -3:, -1]

         assert image.shape == (1, 3_2, 3_2, 3)
         snake_case				  =      np.array([0.44_27, 0.37_31, 0.42_49, 0.49_41, 0.45_46, 0.41_48, 0.41_93, 0.46_66, 0.44_99]     )

         assert np.abs(image_slice.flatten() - expected_slice     ).max() < 5E-3
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 5E-3





     @unittest.skipIf(torch_device != "cuda"       ,					"This test requires a GPU"     )
     def 	UpperCAmelCase(self   :					str     )		->		Any:
         snake_case				  =      self.dummy_cond_unet
         snake_case				  =      PNDMScheduler(skip_prk_steps=_A     )
         snake_case				  =      self.dummy_vae
         snake_case				  =      self.dummy_text_encoder
         snake_case				  =      XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"     )
         snake_case				  =      7_7

         snake_case				  =      self.dummy_image.to(_A     )

         # put models in fp16
         snake_case				  =      unet.half()
         snake_case				  =      vae.half()
         snake_case				  =      bert.half()

         # make sure here that pndm scheduler skips prk
         snake_case				  =      AltDiffusionImgaImgPipeline(
             unet=_A       ,					scheduler=_A       ,					vae=_A       ,					text_encoder=_A       ,					tokenizer=_A       ,					safety_checker=_A       ,					feature_extractor=self.dummy_extractor       ,					)
         snake_case				  =      VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor       ,					do_normalize=_A     )
         snake_case				  =      alt_pipe.to(_A     )
         alt_pipe.set_progress_bar_config(disable=_A     )

         snake_case				  =      "A painting of a squirrel eating a burger"
         snake_case				  =      torch.manual_seed(0     )
         snake_case				  =      alt_pipe(
             [prompt]       ,					generator=_A       ,					num_inference_steps=2       ,					output_type="np"       ,					image=_A       ,					).images

         assert image.shape == (1, 3_2, 3_2, 3)





     @unittest.skipIf(torch_device != "cuda"       ,					"This test requires a GPU"     )
     def 	UpperCAmelCase(self   :					int     )		->		Optional[int]:
         snake_case				  =      load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/img2img/sketch-mountains-input.jpg"     )
         # resize to resolution that is divisible by 8 but not 16 or 32
         snake_case				  =      init_image.resize((7_6_0, 5_0_4)     )

         snake_case				  =      "BAAI/AltDiffusion"
         snake_case				  =      AltDiffusionImgaImgPipeline.from_pretrained(
             _A       ,					safety_checker=_A       ,					)
         pipe.to(_A     )
         pipe.set_progress_bar_config(disable=_A     )
         pipe.enable_attention_slicing()

         snake_case				  =      "A fantasy landscape, trending on artstation"

         snake_case				  =      torch.manual_seed(0     )
         snake_case				  =      pipe(
             prompt=_A       ,					image=_A       ,					strength=0.75       ,					guidance_scale=7.5       ,					generator=_A       ,					output_type="np"       ,					)
         snake_case				  =      output.images[0]

         snake_case				  =      image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

         assert image.shape == (5_0_4, 7_6_0, 3)
         snake_case				  =      np.array([0.93_58, 0.93_97, 0.95_99, 0.99_01, 1.00_00, 1.00_00, 0.98_82, 1.00_00, 1.00_00]     )

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







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





     def 	UpperCAmelCase(self   :					List[Any]     )		->		List[str]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()





     def 	UpperCAmelCase(self   :					str     )		->		str:
         snake_case				  =      load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/img2img/sketch-mountains-input.jpg"     )
         snake_case				  =      init_image.resize((7_6_8, 5_1_2)     )
         snake_case				  =      load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"     )

         snake_case				  =      "BAAI/AltDiffusion"
         snake_case				  =      AltDiffusionImgaImgPipeline.from_pretrained(
             _A       ,					safety_checker=_A       ,					)
         pipe.to(_A     )
         pipe.set_progress_bar_config(disable=_A     )
         pipe.enable_attention_slicing()

         snake_case				  =      "A fantasy landscape, trending on artstation"

         snake_case				  =      torch.manual_seed(0     )
         snake_case				  =      pipe(
             prompt=_A       ,					image=_A       ,					strength=0.75       ,					guidance_scale=7.5       ,					generator=_A       ,					output_type="np"       ,					)
         snake_case				  =      output.images[0]

         assert image.shape == (5_1_2, 7_6_8, 3)
         # img2img is flaky across GPUs even in fp32, so using MAE here
         assert np.abs(expected_image - image     ).max() < 1E-2








import os

# Precomputes a list of the 100 first triangular numbers
_A						      =				[int(0.5 * n * (n + 1)) for n in range(1, 1_01)]



def 		lowercase_    (     )							->							Tuple:
    """simple docstring"""

    snake_case				  =      os.path.dirname(os.path.realpath(A__	)	)
    snake_case				  =      os.path.join(A__				,	"words.txt"	)

    snake_case				  =      ""
    with open(A__	) as f:
        snake_case				  =      f.readline()

    snake_case				  =      [word.strip("\""	) for word in words.strip("\r\n"	).split(","	)]
    snake_case				  =      [
        word
        for word in [sum(ord(A__	) - 64 for x in word	) for word in words]
        if word in TRIANGULAR_NUMBERS
    ]
    return len(A__	)


if __name__ == "__main__":
   print(solution())

