





"""simple docstring"""


import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
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 PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class        lowercase_			(       __lowerCAmelCase			,__lowerCAmelCase			,unittest.TestCase							):
       snake_case__			:							int										=       StableDiffusionParadigmsPipeline
       snake_case__			:							Dict										=       TEXT_TO_IMAGE_PARAMS
       snake_case__			:							List[str]										=       TEXT_TO_IMAGE_BATCH_PARAMS
       snake_case__			:							Dict										=       TEXT_TO_IMAGE_IMAGE_PARAMS
       snake_case__			:							Tuple										=       TEXT_TO_IMAGE_IMAGE_PARAMS

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        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				,   attention_head_dim=(2, 4)				,   use_linear_projection=__SCREAMING_SNAKE_CASE				,   )
        snake_case__      							=			DDIMParallelScheduler(
            beta_start=0.00085				,   beta_end=0.012				,   beta_schedule="scaled_linear"				,   clip_sample=__SCREAMING_SNAKE_CASE				,   set_alpha_to_one=__SCREAMING_SNAKE_CASE				,   )
        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				,   sample_size=128				,   )
        torch.manual_seed(0       )
        snake_case__      							=			CLIPTextConfig(
            bos_token_id=0				,   eos_token_id=2				,   hidden_size=32				,   intermediate_size=37				,   layer_norm_eps=1e-05				,   num_attention_heads=4				,   num_hidden_layers=5				,   pad_token_id=1				,   vocab_size=1000				,   hidden_act="gelu"				,   projection_dim=512				,   )
        snake_case__      							=			CLIPTextModel(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

        snake_case__      							=			{
            "unet": unet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "safety_checker": None,
            "feature_extractor": None,
        }
        return components

       def 						_SCREAMING_SNAKE_CASE(     self				,   __SCREAMING_SNAKE_CASE				,   __SCREAMING_SNAKE_CASE=0       ):






        '''simple docstring'''

        if str(__SCREAMING_SNAKE_CASE       ).startswith("mps"       ):
         snake_case__      							=			torch.manual_seed(__SCREAMING_SNAKE_CASE       )
        else:
         snake_case__      							=			torch.Generator(device=__SCREAMING_SNAKE_CASE       ).manual_seed(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			{
            "prompt": "a photograph of an astronaut riding a horse",
            "generator": generator,
            "num_inference_steps": 10,
            "guidance_scale": 6.0,
            "output_type": "numpy",
            "parallel": 3,
            "debug": True,
        }
        return inputs

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__      							=			self.get_dummy_components()
        snake_case__      							=			StableDiffusionParadigmsPipeline(**__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			sd_pipe.to(__SCREAMING_SNAKE_CASE       )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )

        snake_case__      							=			self.get_dummy_inputs(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			sd_pipe(**__SCREAMING_SNAKE_CASE       ).images
        snake_case__      							=			image[0, -3:, -3:, -1]
        assert image.shape == (1, 64, 64, 3)

        snake_case__      							=			np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]       )

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

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__      							=			self.get_dummy_components()
        torch.manual_seed(0       )
        snake_case__      							=			DDPMParallelScheduler()
        torch.manual_seed(0       )
        snake_case__      							=			StableDiffusionParadigmsPipeline(**__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			sd_pipe.to(__SCREAMING_SNAKE_CASE       )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )

        snake_case__      							=			self.get_dummy_inputs(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			sd_pipe(**__SCREAMING_SNAKE_CASE       ).images
        snake_case__      							=			image[0, -3:, -3:, -1]
        assert image.shape == (1, 64, 64, 3)

        snake_case__      							=			np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]       )

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

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        super().test_inference_batch_consistent(batch_sizes=[1, 2]       )

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        super().test_inference_batch_single_identical(batch_size=2				,   expected_max_diff=3e-3       )

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			"cpu"  # ensure determinism for the device-dependent torch.Generator
        snake_case__      							=			self.get_dummy_components()
        snake_case__      							=			StableDiffusionParadigmsPipeline(**__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			sd_pipe.to(__SCREAMING_SNAKE_CASE       )
        sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )

        snake_case__      							=			self.get_dummy_inputs(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			"french fries"
        snake_case__      							=			sd_pipe(**__SCREAMING_SNAKE_CASE				,   negative_prompt=__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			output.images
        snake_case__      							=			image[0, -3:, -3:, -1]

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

        snake_case__      							=			np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]       )

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



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

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

       def 						_SCREAMING_SNAKE_CASE(     self				,   __SCREAMING_SNAKE_CASE=0       ):






        '''simple docstring'''

        snake_case__      							=			torch.Generator(device=__SCREAMING_SNAKE_CASE       ).manual_seed(__SCREAMING_SNAKE_CASE       )
        snake_case__      							=			{
            "prompt": "a photograph of an astronaut riding a horse",
            "generator": generator,
            "num_inference_steps": 10,
            "guidance_scale": 7.5,
            "output_type": "numpy",
            "parallel": 3,
            "debug": True,
        }
        return inputs

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			"stabilityai/stable-diffusion-2-base"
        snake_case__      							=			DDIMParallelScheduler.from_pretrained(__SCREAMING_SNAKE_CASE				,   subfolder="scheduler"       )
        snake_case__      							=			StableDiffusionParadigmsPipeline.from_pretrained(__SCREAMING_SNAKE_CASE				,   scheduler=__SCREAMING_SNAKE_CASE				,   safety_checker=__SCREAMING_SNAKE_CASE       )
        pipe.to(__SCREAMING_SNAKE_CASE       )
        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE       )
        pipe.enable_attention_slicing()

        snake_case__      							=			self.get_inputs()
        snake_case__      							=			pipe(**__SCREAMING_SNAKE_CASE       ).images
        snake_case__      							=			image[0, -3:, -3:, -1].flatten()

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

        snake_case__      							=			np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]       )

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








"""simple docstring"""


import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


a__   :       Optional[Any]   						=			logging.get_logger(__name__)

a__   :       List[str]   						=			{}



class        lowercase_			(       __lowerCAmelCase							):
       snake_case__			:							List[str]										=       """fsmt_decoder"""

       def __init__(     self				,   __SCREAMING_SNAKE_CASE=0				,   __SCREAMING_SNAKE_CASE=0       ):






        '''simple docstring'''

        super().__init__()
        snake_case__      							=			vocab_size
        snake_case__      							=			bos_token_id



class        lowercase_			(       __lowerCAmelCase							):
       snake_case__			:							str										=       """fsmt"""
       snake_case__			:							List[Any]										=       {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}

       def __init__(     self				,   __SCREAMING_SNAKE_CASE=["en", "de"]				,   __SCREAMING_SNAKE_CASE=4_2024				,   __SCREAMING_SNAKE_CASE=4_2024				,   __SCREAMING_SNAKE_CASE="relu"				,   __SCREAMING_SNAKE_CASE=1024				,   __SCREAMING_SNAKE_CASE=200				,   __SCREAMING_SNAKE_CASE=1024				,   __SCREAMING_SNAKE_CASE=4096				,   __SCREAMING_SNAKE_CASE=12				,   __SCREAMING_SNAKE_CASE=16				,   __SCREAMING_SNAKE_CASE=0.0				,   __SCREAMING_SNAKE_CASE=4096				,   __SCREAMING_SNAKE_CASE=12				,   __SCREAMING_SNAKE_CASE=16				,   __SCREAMING_SNAKE_CASE=0.0				,   __SCREAMING_SNAKE_CASE=0.0				,   __SCREAMING_SNAKE_CASE=0.1				,   __SCREAMING_SNAKE_CASE=0.0				,   __SCREAMING_SNAKE_CASE=0.02				,   __SCREAMING_SNAKE_CASE=2				,   __SCREAMING_SNAKE_CASE=True				,   __SCREAMING_SNAKE_CASE=True				,   __SCREAMING_SNAKE_CASE=False				,   __SCREAMING_SNAKE_CASE=5				,   __SCREAMING_SNAKE_CASE=1.0				,   __SCREAMING_SNAKE_CASE=False				,   __SCREAMING_SNAKE_CASE=True				,   __SCREAMING_SNAKE_CASE=1				,   __SCREAMING_SNAKE_CASE=0				,   __SCREAMING_SNAKE_CASE=2				,   __SCREAMING_SNAKE_CASE=2				,   **__SCREAMING_SNAKE_CASE				,   ):






        '''simple docstring'''

        snake_case__      							=			langs
        snake_case__      							=			src_vocab_size
        snake_case__      							=			tgt_vocab_size
        snake_case__      							=			d_model  # encoder_embed_dim and decoder_embed_dim

        snake_case__      							=			encoder_ffn_dim
        snake_case__      							=			snake_case__      							=			encoder_layers
        snake_case__      							=			encoder_attention_heads
        snake_case__      							=			encoder_layerdrop
        snake_case__      							=			decoder_layerdrop
        snake_case__      							=			decoder_ffn_dim
        snake_case__      							=			decoder_layers
        snake_case__      							=			decoder_attention_heads
        snake_case__      							=			max_position_embeddings
        snake_case__      							=			init_std  # Normal(0, this parameter)
        snake_case__      							=			activation_function

        snake_case__      							=			DecoderConfig(vocab_size=__SCREAMING_SNAKE_CASE				,   bos_token_id=__SCREAMING_SNAKE_CASE       )
        if "decoder" in common_kwargs:
         del common_kwargs["decoder"]

        snake_case__      							=			scale_embedding  # scale factor will be sqrt(d_model) if True

        # 3 Types of Dropout
        snake_case__      							=			attention_dropout
        snake_case__      							=			activation_dropout
        snake_case__      							=			dropout

        snake_case__      							=			use_cache
        super().__init__(
            pad_token_id=__SCREAMING_SNAKE_CASE				,   bos_token_id=__SCREAMING_SNAKE_CASE				,   eos_token_id=__SCREAMING_SNAKE_CASE				,   decoder_start_token_id=__SCREAMING_SNAKE_CASE				,   is_encoder_decoder=__SCREAMING_SNAKE_CASE				,   tie_word_embeddings=__SCREAMING_SNAKE_CASE				,   forced_eos_token_id=__SCREAMING_SNAKE_CASE				,   max_length=__SCREAMING_SNAKE_CASE				,   num_beams=__SCREAMING_SNAKE_CASE				,   length_penalty=__SCREAMING_SNAKE_CASE				,   early_stopping=__SCREAMING_SNAKE_CASE				,   **__SCREAMING_SNAKE_CASE				,   )

       def 						_SCREAMING_SNAKE_CASE(     self       ):






        '''simple docstring'''

        snake_case__      							=			copy.deepcopy(self.__dict__       )
        snake_case__      							=			self.decoder.to_dict()
        snake_case__      							=			self.__class__.model_type
        return output

