
import gc
import unittest

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

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

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


enable_full_determinism()




class        _SCREAMING_SNAKE_CASE	(   __UpperCamelCase							,       __UpperCamelCase							,       __UpperCamelCase							,       unittest.TestCase			):
      _A    :		List[Any]     						=  AltDiffusionPipeline
      _A    :		List[Any]     						=  TEXT_TO_IMAGE_PARAMS
      _A    :		Any     						=  TEXT_TO_IMAGE_BATCH_PARAMS
      _A    :		Any     						=  TEXT_TO_IMAGE_IMAGE_PARAMS
      _A    :		Tuple     						=  TEXT_TO_IMAGE_IMAGE_PARAMS
      def A_      (      self			):
          torch.manual_seed(0			)
          snake_case__    =					UNetaDConditionModel(
              block_out_channels=(32, 64)		,						layers_per_block=2		,						sample_size=32		,						in_channels=4		,						out_channels=4		,						down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")		,						up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")		,						cross_attention_dim=32		,						)
          snake_case__    =					DDIMScheduler(
              beta_start=0.0_0_0_8_5		,						beta_end=0.0_1_2		,						beta_schedule="scaled_linear"		,						clip_sample=lowerCamelCase		,						set_alpha_to_one=lowerCamelCase		,						)
          torch.manual_seed(0			)
          snake_case__    =					AutoencoderKL(
              block_out_channels=[32, 64]		,						in_channels=3		,						out_channels=3		,						down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]		,						up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]		,						latent_channels=4		,						)

          # TODO: address the non-deterministic text encoder (fails for save-load tests)
          # torch.manual_seed(0)
          # text_encoder_config = RobertaSeriesConfig(
          #     hidden_size=32,
          #     project_dim=32,
          #     intermediate_size=37,
          #     layer_norm_eps=1e-05,
          #     num_attention_heads=4,
          #     num_hidden_layers=5,
          #     vocab_size=5002,
          # )
          # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

          torch.manual_seed(0			)
          snake_case__    =					CLIPTextConfig(
              bos_token_id=0		,						eos_token_id=2		,						hidden_size=32		,						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=50_02		,						)
          snake_case__    =					CLIPTextModel(lowerCamelCase			)

          snake_case__    =					XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"			)
          snake_case__    =					77

          snake_case__    =					{
              "unet": unet,
              "scheduler": scheduler,
              "vae": vae,
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "safety_checker": None,
              "feature_extractor": None,
          }
          return components
      def A_      (      self		,						lowerCamelCase		,						lowerCamelCase=0			):
          if str(lowerCamelCase			).startswith("mps"			):
              snake_case__    =					torch.manual_seed(lowerCamelCase			)
          else:
              snake_case__    =					torch.Generator(device=lowerCamelCase			).manual_seed(lowerCamelCase			)
          snake_case__    =					{
              "prompt": "A painting of a squirrel eating a burger",
              "generator": generator,
              "num_inference_steps": 2,
              "guidance_scale": 6.0,
              "output_type": "numpy",
          }
          return inputs
      def A_      (      self			):
          super().test_attention_slicing_forward_pass(expected_max_diff=3e-3			)
      def A_      (      self			):
          super().test_inference_batch_single_identical(expected_max_diff=3e-3			)
      def A_      (      self			):
          snake_case__    =					"cpu"  # ensure determinism for the device-dependent torch.Generator

          snake_case__    =					self.get_dummy_components()
          torch.manual_seed(0			)
          snake_case__    =					RobertaSeriesConfig(
              hidden_size=32		,						project_dim=32		,						intermediate_size=37		,						layer_norm_eps=1e-05		,						num_attention_heads=4		,						num_hidden_layers=5		,						vocab_size=50_02		,						)
          # TODO: remove after fixing the non-deterministic text encoder
          snake_case__    =					RobertaSeriesModelWithTransformation(lowerCamelCase			)
          snake_case__    =					text_encoder

          snake_case__    =					AltDiffusionPipeline(**lowerCamelCase			)
          snake_case__    =					alt_pipe.to(lowerCamelCase			)
          alt_pipe.set_progress_bar_config(disable=lowerCamelCase			)

          snake_case__    =					self.get_dummy_inputs(lowerCamelCase			)
          snake_case__    =					"A photo of an astronaut"
          snake_case__    =					alt_pipe(**lowerCamelCase			)
          snake_case__    =					output.images
          snake_case__    =					image[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          snake_case__    =					np.array(
              [0.5_7_4_8_1_6_2, 0.6_0_4_4_7_1_4_5, 0.4_8_8_2_1_2_1_7, 0.5_0_1_0_0_6_3_6, 0.5_4_3_1_1_8_5, 0.4_5_7_6_3_6_8_3, 0.4_9_6_5_7_6_9_6, 0.4_8_1_3_2_7_3_3, 0.4_7_5_7_3_0_9_3]			)

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



      def A_      (      self			):
          snake_case__    =					"cpu"  # ensure determinism for the device-dependent torch.Generator

          snake_case__    =					self.get_dummy_components()
          snake_case__    =					PNDMScheduler(skip_prk_steps=lowerCamelCase			)
          torch.manual_seed(0			)
          snake_case__    =					RobertaSeriesConfig(
              hidden_size=32		,						project_dim=32		,						intermediate_size=37		,						layer_norm_eps=1e-05		,						num_attention_heads=4		,						num_hidden_layers=5		,						vocab_size=50_02		,						)
          # TODO: remove after fixing the non-deterministic text encoder
          snake_case__    =					RobertaSeriesModelWithTransformation(lowerCamelCase			)
          snake_case__    =					text_encoder
          snake_case__    =					AltDiffusionPipeline(**lowerCamelCase			)
          snake_case__    =					alt_pipe.to(lowerCamelCase			)
          alt_pipe.set_progress_bar_config(disable=lowerCamelCase			)

          snake_case__    =					self.get_dummy_inputs(lowerCamelCase			)
          snake_case__    =					alt_pipe(**lowerCamelCase			)
          snake_case__    =					output.images
          snake_case__    =					image[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          snake_case__    =					np.array(
              [0.5_1_6_0_5_0_9_3, 0.5_7_0_7_2_4_1, 0.4_7_3_6_5_5_0_7, 0.5_0_5_7_8_8_8_6, 0.5_6_3_3_8_7_7, 0.4_6_4_2_5_0_3, 0.5_1_8_2_0_8_1, 0.4_8_7_6_3_4_8_4, 0.4_9_0_8_4_2_3_7]			)

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





@slow
@require_torch_gpu
class        _SCREAMING_SNAKE_CASE	(   unittest.TestCase			):
      def A_      (      self			):
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()
      def A_      (      self			):
          # make sure here that pndm scheduler skips prk
          snake_case__    =					AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"		,						safety_checker=lowerCamelCase			)
          snake_case__    =					alt_pipe.to(lowerCamelCase			)
          alt_pipe.set_progress_bar_config(disable=lowerCamelCase			)

          snake_case__    =					"A painting of a squirrel eating a burger"
          snake_case__    =					torch.manual_seed(0			)
          snake_case__    =					alt_pipe([prompt]		,						generator=lowerCamelCase		,						guidance_scale=6.0		,						num_inference_steps=20		,						output_type="np"			)

          snake_case__    =					output.images

          snake_case__    =					image[0, -3:, -3:, -1]

          assert image.shape == (1, 5_12, 5_12, 3)
          snake_case__    =					np.array([0.1_0_1_0, 0.0_8_0_0, 0.0_7_9_4, 0.0_8_8_5, 0.0_8_4_3, 0.0_7_6_2, 0.0_7_6_9, 0.0_7_2_9, 0.0_5_8_6]			)

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



      def A_      (      self			):
          snake_case__    =					DDIMScheduler.from_pretrained("BAAI/AltDiffusion"		,						subfolder="scheduler"			)

          snake_case__    =					AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion"		,						scheduler=lowerCamelCase		,						safety_checker=lowerCamelCase			)
          snake_case__    =					alt_pipe.to(lowerCamelCase			)
          alt_pipe.set_progress_bar_config(disable=lowerCamelCase			)

          snake_case__    =					"A painting of a squirrel eating a burger"
          snake_case__    =					torch.manual_seed(0			)

          snake_case__    =					alt_pipe([prompt]		,						generator=lowerCamelCase		,						num_inference_steps=2		,						output_type="numpy"			)
          snake_case__    =					output.images

          snake_case__    =					image[0, -3:, -3:, -1]

          assert image.shape == (1, 5_12, 5_12, 3)
          snake_case__    =					np.array([0.4_0_1_9, 0.4_0_5_2, 0.3_8_1_0, 0.4_1_1_9, 0.3_9_1_6, 0.3_9_8_2, 0.4_6_5_1, 0.4_1_9_5, 0.5_3_2_3]			)

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





from __future__ import annotations

import math





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase				):

    if 1 < number < 4:
        # 2 and 3 are primes
        return True
    elif number < 2 or number % 2 == 0 or number % 3 == 0:
        # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
        return False

    # All primes number are in format of 6k +/- 1
    for i in range(5     ,		int(math.sqrt(__lowerCAmelCase				) + 1				)     ,		6				):
        if number % i == 0 or number % (i + 2) == 0:
            return False
    return True





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase				):
    snake_case__    =					str(__lowerCAmelCase				)
    snake_case__    =					[n]
    for i in range(1     ,		len(__lowerCAmelCase				)				):
        list_nums.append(int(str_num[i:]				)				)
        list_nums.append(int(str_num[:-i]				)				)
    return list_nums





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase				):
    if len(str(__lowerCAmelCase				)				) > 3:
        if not is_prime(int(str(__lowerCAmelCase				)[-3:]				)				) or not is_prime(int(str(__lowerCAmelCase				)[:3]				)				):
            return False
    return True





def 						SCREAMING_SNAKE_CASE__     (   __lowerCAmelCase = 11				):
    snake_case__    =					[]
    snake_case__    =					13
    while len(__lowerCAmelCase				) != count:
        if validate(__lowerCAmelCase				):
            snake_case__    =					list_truncated_nums(__lowerCAmelCase				)
            if all(is_prime(__lowerCAmelCase				) for i in list_nums				):
                list_truncated_primes.append(__lowerCAmelCase				)
        num += 2
    return list_truncated_primes





def 						SCREAMING_SNAKE_CASE__     (   ):
    return sum(compute_truncated_primes(11				)				)


if __name__ == "__main__":
     print(F'''{sum(compute_truncated_primes(11)) = }''')



