'''simple docstring'''

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       UpperCAmelCase  (						lowerCamelCase__   ,						lowerCamelCase__   ,						lowerCamelCase__   ,						unittest.TestCase ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE 		=   AltDiffusionPipeline
       SCREAMING_SNAKE_CASE 		=   TEXT_TO_IMAGE_PARAMS
       SCREAMING_SNAKE_CASE 		=   TEXT_TO_IMAGE_BATCH_PARAMS
       SCREAMING_SNAKE_CASE 		=   TEXT_TO_IMAGE_IMAGE_PARAMS
       SCREAMING_SNAKE_CASE 		=   TEXT_TO_IMAGE_IMAGE_PARAMS

       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        torch.manual_seed(0			)
        lowercase__    :   int						= 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     ,       )
        lowercase__    :   List[str]						= DDIMScheduler(
            beta_start=0.0_0_0_8_5     ,       beta_end=0.0_1_2     ,       beta_schedule='''scaled_linear'''     ,       clip_sample=lowercase__     ,       set_alpha_to_one=lowercase__     ,       )
        torch.manual_seed(0			)
        lowercase__    :   str						= 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			)
        lowercase__    :   str						= 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=5002     ,       )
        lowercase__    :   Any						= CLIPTextModel(lowercase__			)

        lowercase__    :   Tuple						= XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''			)
        lowercase__    :   str						= 77

        lowercase__    :   Optional[int]						= {
            '''unet''': unet,
            '''scheduler''': scheduler,
            '''vae''': vae,
            '''text_encoder''': text_encoder,
            '''tokenizer''': tokenizer,
            '''safety_checker''': None,
            '''feature_extractor''': None,
        }
        return components

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase=0			)  ->  int:
        if str(lowercase__			).startswith('''mps'''			):
         lowercase__    :   Optional[Any]						= torch.manual_seed(lowercase__			)
        else:
         lowercase__    :   Optional[Any]						= torch.Generator(device=lowercase__			).manual_seed(lowercase__			)
        lowercase__    :   str						= {
            '''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 _lowerCAmelCase(						self			)  ->  Optional[Any]:
        super().test_attention_slicing_forward_pass(expected_max_diff=3E-3			)

       def _lowerCAmelCase(						self			)  ->  int:
        super().test_inference_batch_single_identical(expected_max_diff=3E-3			)

       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        lowercase__    :   int						= '''cpu'''  # ensure determinism for the device-dependent torch.Generator

        lowercase__    :   Any						= self.get_dummy_components()
        torch.manual_seed(0			)
        lowercase__    :   Optional[Any]						= 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     ,       )
        # TODO: remove after fixing the non-deterministic text encoder
        lowercase__    :   Optional[Any]						= RobertaSeriesModelWithTransformation(lowercase__			)
        lowercase__    :   str						= text_encoder

        lowercase__    :   Tuple						= AltDiffusionPipeline(**lowercase__			)
        lowercase__    :   Dict						= alt_pipe.to(lowercase__			)
        alt_pipe.set_progress_bar_config(disable=lowercase__			)

        lowercase__    :   List[str]						= self.get_dummy_inputs(lowercase__			)
        lowercase__    :   Optional[int]						= '''A photo of an astronaut'''
        lowercase__    :   int						= alt_pipe(**lowercase__			)
        lowercase__    :   int						= output.images
        lowercase__    :   int						= image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)
        lowercase__    :   Optional[int]						= 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 _lowerCAmelCase(						self			)  ->  str:
        lowercase__    :   Optional[int]						= '''cpu'''  # ensure determinism for the device-dependent torch.Generator

        lowercase__    :   Dict						= self.get_dummy_components()
        lowercase__    :   Union[str, Any]						= PNDMScheduler(skip_prk_steps=lowercase__			)
        torch.manual_seed(0			)
        lowercase__    :   Tuple						= 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     ,       )
        # TODO: remove after fixing the non-deterministic text encoder
        lowercase__    :   Dict						= RobertaSeriesModelWithTransformation(lowercase__			)
        lowercase__    :   Optional[Any]						= text_encoder
        lowercase__    :   Tuple						= AltDiffusionPipeline(**lowercase__			)
        lowercase__    :   List[str]						= alt_pipe.to(lowercase__			)
        alt_pipe.set_progress_bar_config(disable=lowercase__			)

        lowercase__    :   Tuple						= self.get_dummy_inputs(lowercase__			)
        lowercase__    :   List[Any]						= alt_pipe(**lowercase__			)
        lowercase__    :   Tuple						= output.images
        lowercase__    :   List[Any]						= image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)
        lowercase__    :   str						= 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       UpperCAmelCase  (						unittest.TestCase ):





       '''simple docstring'''



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

       def _lowerCAmelCase(						self			)  ->  int:
        # make sure here that pndm scheduler skips prk
        lowercase__    :   List[str]						= AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''     ,       safety_checker=lowercase__			)
        lowercase__    :   Dict						= alt_pipe.to(lowercase__			)
        alt_pipe.set_progress_bar_config(disable=lowercase__			)

        lowercase__    :   int						= '''A painting of a squirrel eating a burger'''
        lowercase__    :   Optional[int]						= torch.manual_seed(0			)
        lowercase__    :   List[str]						= alt_pipe([prompt]     ,       generator=lowercase__     ,       guidance_scale=6.0     ,       num_inference_steps=20     ,       output_type='''np'''			)

        lowercase__    :   int						= output.images

        lowercase__    :   Any						= image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        lowercase__    :   Tuple						= 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 _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        lowercase__    :   Optional[int]						= DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''     ,       subfolder='''scheduler'''			)

        lowercase__    :   Union[str, Any]						= AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''     ,       scheduler=lowercase__     ,       safety_checker=lowercase__			)
        lowercase__    :   List[Any]						= alt_pipe.to(lowercase__			)
        alt_pipe.set_progress_bar_config(disable=lowercase__			)

        lowercase__    :   Any						= '''A painting of a squirrel eating a burger'''
        lowercase__    :   Tuple						= torch.manual_seed(0			)

        lowercase__    :   List[Any]						= alt_pipe([prompt]     ,       generator=lowercase__     ,       num_inference_steps=2     ,       output_type='''numpy'''			)
        lowercase__    :   Dict						= output.images

        lowercase__    :   str						= image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        lowercase__    :   int						= 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


'''simple docstring'''

from manim import *



class       UpperCAmelCase  (						a__ ):





       '''simple docstring'''



       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   int						= Rectangle(height=0.5     ,       width=0.5			)
        lowercase__    :   Optional[int]						= Rectangle(height=0.2_5     ,       width=0.2_5			)
        lowercase__    :   Tuple						= Rectangle(height=0.4_6     ,       width=0.4_6			).set_stroke(width=0			)

        lowercase__    :   str						= [mem.copy() for i in range(6			)]
        lowercase__    :   Dict						= [mem.copy() for i in range(6			)]
        lowercase__    :   Tuple						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   List[str]						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   str						= VGroup(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   str						= Text('''CPU'''     ,       font_size=24			)
        lowercase__    :   List[Any]						= Group(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0.5     ,       aligned_edge=__lowerCAmelCase			)
        cpu.move_to([-2.5, -0.5, 0]			)
        self.add(__lowerCAmelCase			)

        lowercase__    :   Any						= [mem.copy() for i in range(4			)]
        lowercase__    :   int						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   Optional[int]						= Text('''GPU'''     ,       font_size=24			)
        lowercase__    :   Tuple						= Group(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0.5     ,       aligned_edge=__lowerCAmelCase			)
        gpu.move_to([-1, -1, 0]			)
        self.add(__lowerCAmelCase			)

        lowercase__    :   int						= [mem.copy() for i in range(6			)]
        lowercase__    :   Optional[int]						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)

        lowercase__    :   Any						= Text('''Model'''     ,       font_size=24			)
        lowercase__    :   Optional[Any]						= Group(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0.5     ,       aligned_edge=__lowerCAmelCase			)
        model.move_to([3, -1.0, 0]			)
        self.add(__lowerCAmelCase			)

        lowercase__    :   int						= []
        lowercase__    :   int						= []
        lowercase__    :   Any						= []

        for i, rect in enumerate(__lowerCAmelCase			):
         rect.set_stroke(__lowerCAmelCase			)

         lowercase__    :   Optional[int]						= Rectangle(height=0.4_6 / 4     ,       width=0.4_6 / 3			).set_stroke(width=0.0			).set_fill(__lowerCAmelCase     ,       opacity=0.7			)

         if i == 0:
          cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT			)     ,       buff=0.0_2     ,       direction=__lowerCAmelCase			)
          cpu_target.set_x(cpu_target.get_x() + 0.1			)
         elif i == 3:
          cpu_target.next_to(model_cpu_arr[0]     ,       direction=__lowerCAmelCase     ,       buff=0.0			)
         else:
          cpu_target.next_to(model_cpu_arr[i - 1]     ,       direction=__lowerCAmelCase     ,       buff=0.0			)
         self.add(__lowerCAmelCase			)
         model_cpu_arr.append(__lowerCAmelCase			)

        self.add(*__lowerCAmelCase     ,       *__lowerCAmelCase     ,       *__lowerCAmelCase			)

        lowercase__    :   Optional[int]						= [mem.copy() for i in range(6			)]
        lowercase__    :   List[Any]						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)

        lowercase__    :   Optional[Any]						= Text('''Loaded Checkpoint'''     ,       font_size=24			)
        lowercase__    :   Dict						= Group(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0.5     ,       aligned_edge=__lowerCAmelCase			)
        checkpoint.move_to([3, 0.5, 0]			)

        self.add(__lowerCAmelCase			)

        lowercase__    :   str						= []
        lowercase__    :   List[str]						= []

        for i, rect in enumerate(__lowerCAmelCase			):
         lowercase__    :   List[str]						= fill.copy().set_fill(__lowerCAmelCase     ,       opacity=0.7			)
         target.move_to(__lowerCAmelCase			)
         ckpt_arr.append(__lowerCAmelCase			)

         lowercase__    :   Any						= target.copy()
         if i < 5:
          cpu_target.move_to(cpu_left_col_base[i + 1]			)
         else:
          cpu_target.move_to(cpu_right_col_base[i - 5]			)
         ckpt_cpu_arr.append(__lowerCAmelCase			)
        self.add(*__lowerCAmelCase     ,       *__lowerCAmelCase			)

        lowercase__    :   Optional[int]						= Square(side_length=2.2			)
        key.move_to([-5, 2, 0]			)

        lowercase__    :   Optional[Any]						= MarkupText(
            F"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""     ,       font_size=18     ,       )

        key_text.move_to([-5, 2.4, 0]			)

        self.add(__lowerCAmelCase     ,       __lowerCAmelCase			)

        lowercase__    :   str						= MarkupText(
            F"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""     ,       font_size=18     ,       )

        blue_text.next_to(__lowerCAmelCase     ,       DOWN * 2.4     ,       aligned_edge=key_text.get_left()			)
        self.add(__lowerCAmelCase			)

        lowercase__    :   Union[str, Any]						= MarkupText(
            F"""Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device."""     ,       font_size=24     ,       )
        step_a.move_to([2, 2, 0]			)

        lowercase__    :   Tuple						= [meta_mem.copy() for i in range(6			)]
        lowercase__    :   Any						= [meta_mem.copy() for i in range(6			)]
        lowercase__    :   str						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   Any						= VGroup(*__lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   Union[str, Any]						= VGroup(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0			)
        lowercase__    :   Union[str, Any]						= Text('''Disk'''     ,       font_size=24			)
        lowercase__    :   Tuple						= Group(__lowerCAmelCase     ,       __lowerCAmelCase			).arrange(__lowerCAmelCase     ,       buff=0.5     ,       aligned_edge=__lowerCAmelCase			)
        disk.move_to([-4.0, -1.2_5, 0]			)
        self.play(Write(__lowerCAmelCase     ,       run_time=3			)     ,       Write(__lowerCAmelCase     ,       run_time=1			)     ,       Create(__lowerCAmelCase     ,       run_time=1			)			)

        lowercase__    :   Tuple						= []
        for i, rect in enumerate(__lowerCAmelCase			):
         lowercase__    :   Dict						= rect.copy()
         target.generate_target()
         target.target.move_to(disk_left_col_base[i]			).scale(0.5			)
         animations.append(MoveToTarget(__lowerCAmelCase     ,       run_time=1.5			)			)
        self.play(*__lowerCAmelCase			)

        self.play(FadeOut(__lowerCAmelCase			)			)

        lowercase__    :   Dict						= MarkupText(F"""Then, the checkpoint is removed from memory\nthrough garbage collection."""     ,       font_size=24			)
        step_a.move_to([2, 2, 0]			)

        self.play(Write(__lowerCAmelCase     ,       run_time=3			)			)

        self.play(
            FadeOut(__lowerCAmelCase     ,       __lowerCAmelCase     ,       *__lowerCAmelCase     ,       *__lowerCAmelCase			)     ,       )

        self.wait()

