




'''simple docstring'''
from typing import Dict, Optional, Tuple

import torch
def 				snake_case_  (   a__		:	torch.Tensor			):



    """simple docstring"""




    __lowercase							=						boundaries[1] - boundaries[0]
    __lowercase							=						boundaries + step / 2
    __lowercase							=						torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1			)]    ,dim=0			)
    return bin_centers
def 				snake_case_  (   a__		:	torch.Tensor    ,a__		:	torch.Tensor    ,):



    """simple docstring"""




    __lowercase							=						_calculate_bin_centers(a__			)
    return (
        torch.sum(aligned_distance_error_probs * bin_centers    ,dim=-1			),
        bin_centers[-1],
    )
def 				snake_case_  (   a__		:	torch.Tensor    ,a__		:	int = 31    ,a__		:	int = 64    ,**a__		:	Tuple    ,):



    """simple docstring"""




    __lowercase							=						torch.linspace(0    ,a__    ,steps=(no_bins - 1)    ,device=logits.device			)

    __lowercase							=						torch.nn.functional.softmax(a__    ,dim=-1			)
    __lowercase				,__lowercase							=						_calculate_expected_aligned_error(
        alignment_confidence_breaks=a__    ,aligned_distance_error_probs=a__    ,)

    return {
        "aligned_confidence_probs": aligned_confidence_probs,
        "predicted_aligned_error": predicted_aligned_error,
        "max_predicted_aligned_error": max_predicted_aligned_error,
    }



def 				snake_case_  (   a__		:	torch.Tensor    ,a__		:	Optional[torch.Tensor] = None    ,a__		:	int = 31    ,a__		:	int = 64    ,a__		:	float = 1e-8    ,**a__		:	List[Any]    ,):



    """simple docstring"""




    if residue_weights is None:
        __lowercase							=						logits.new_ones(logits.shape[-2]			)

    __lowercase							=						torch.linspace(0    ,a__    ,steps=(no_bins - 1)    ,device=logits.device			)

    __lowercase							=						_calculate_bin_centers(a__			)
    torch.sum(a__			)
    __lowercase							=						logits.shape[-2]
    __lowercase							=						max(a__    ,19			)

    __lowercase							=						1.2_4 * (clipped_n - 15) ** (1.0 / 3) - 1.8

    __lowercase							=						torch.nn.functional.softmax(a__    ,dim=-1			)

    __lowercase							=						1.0 / (1 + (bin_centers**2) / (da**2))
    __lowercase							=						torch.sum(probs * tm_per_bin    ,dim=-1			)

    __lowercase							=						residue_weights / (eps + residue_weights.sum())
    __lowercase							=						torch.sum(predicted_tm_term * normed_residue_mask    ,dim=-1			)

    __lowercase							=						per_alignment * residue_weights

    __lowercase							=						(weighted == torch.max(a__			)).nonzero()[0]
    return per_alignment[tuple(a__			)]










'''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        SCREAMING_SNAKE_CASE(				__A   ,      __A   ,      unittest.TestCase     ):
       snake_case_ :     str    =					StableDiffusionParadigmsPipeline
       snake_case_ :     Optional[Any]    =					TEXT_TO_IMAGE_PARAMS
       snake_case_ :     Tuple    =					TEXT_TO_IMAGE_BATCH_PARAMS
       snake_case_ :     Optional[Any]    =					TEXT_TO_IMAGE_IMAGE_PARAMS
       snake_case_ :     str    =					TEXT_TO_IMAGE_IMAGE_PARAMS

       def     snake_case__			(    self     )    ->					int:

           """simple docstring"""



           torch.manual_seed(0     )
           __lowercase							=						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=lowerCamelCase__					,		)
           __lowercase							=						DDIMParallelScheduler(
               beta_start=0.0_00_85					,		beta_end=0.0_12					,		beta_schedule="""scaled_linear"""					,		clip_sample=lowerCamelCase__					,		set_alpha_to_one=lowerCamelCase__					,		)
           torch.manual_seed(0     )
           __lowercase							=						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     )
           __lowercase							=						CLIPTextConfig(
               bos_token_id=0					,		eos_token_id=2					,		hidden_size=32					,		intermediate_size=37					,		layer_norm_eps=1E-0_5					,		num_attention_heads=4					,		num_hidden_layers=5					,		pad_token_id=1					,		vocab_size=1000					,		hidden_act="""gelu"""					,		projection_dim=512					,		)
           __lowercase							=						CLIPTextModel(lowerCamelCase__     )
           __lowercase							=						CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

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

       def     snake_case__			(    self					,		lowerCamelCase__					,		lowerCamelCase__=0     )    ->					Any:

           """simple docstring"""



           if str(lowerCamelCase__     ).startswith("""mps"""     ):
               __lowercase							=						torch.manual_seed(lowerCamelCase__     )
           else:
               __lowercase							=						torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
           __lowercase							=						{
               """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     snake_case__			(    self     )    ->					List[str]:

           """simple docstring"""



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

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

           __lowercase							=						np.array([0.47_73, 0.54_17, 0.47_23, 0.49_25, 0.56_31, 0.47_52, 0.52_40, 0.49_35, 0.50_23]     )

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

       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



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

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

           __lowercase							=						np.array([0.35_73, 0.44_20, 0.49_60, 0.47_99, 0.37_96, 0.38_79, 0.48_19, 0.43_65, 0.44_68]     )

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

       def     snake_case__			(    self     )    ->					Any:

           """simple docstring"""



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

       def     snake_case__			(    self     )    ->					Optional[Any]:

           """simple docstring"""



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




       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



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

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

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

           __lowercase							=						np.array([0.47_71, 0.54_20, 0.46_83, 0.49_18, 0.56_36, 0.47_25, 0.52_30, 0.49_23, 0.50_15]     )

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






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

       def     snake_case__			(    self     )    ->					str:

           """simple docstring"""



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

       def     snake_case__			(    self					,		lowerCamelCase__=0     )    ->					int:

           """simple docstring"""



           __lowercase							=						torch.Generator(device=lowerCamelCase__     ).manual_seed(lowerCamelCase__     )
           __lowercase							=						{
               """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     snake_case__			(    self     )    ->					Optional[Any]:

           """simple docstring"""



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

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

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

           __lowercase							=						np.array([0.96_22, 0.96_02, 0.97_48, 0.95_91, 0.96_30, 0.96_91, 0.96_61, 0.96_31, 0.97_41]     )

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




