
"""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 			lowerCAmelCase_			( _lowercase						,   _lowercase						,   unittest.TestCase      ):



						'''simple docstring'''


						_lowerCamelCase:      Any =      StableDiffusionParadigmsPipeline
						_lowerCamelCase:      Any =      TEXT_TO_IMAGE_PARAMS
						_lowerCamelCase:      Any =      TEXT_TO_IMAGE_BATCH_PARAMS
						_lowerCamelCase:      str =      TEXT_TO_IMAGE_IMAGE_PARAMS
						_lowerCamelCase:      Dict =      TEXT_TO_IMAGE_IMAGE_PARAMS


						def   _SCREAMING_SNAKE_CASE (	self       : Tuple )       ->				Optional[Any]:
							torch.manual_seed(0 )
							A  				=				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=A_					,)
							A  				=				DDIMParallelScheduler(
							    beta_start=0.0_00_85					,beta_end=0.0_12					,beta_schedule='scaled_linear'					,clip_sample=A_					,set_alpha_to_one=A_					,)
							torch.manual_seed(0 )
							A  				=				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 )
							A  				=				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					,)
							A  				=				CLIPTextModel(A_ )
							A  				=				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

							A  				=				{
							    '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       : Dict					,A_       : List[Any]					,A_       : Any=0 )       ->				int:
							if str(A_ ).startswith('mps' ):
								A  				=				torch.manual_seed(A_ )
							else:
								A  				=				torch.Generator(device=A_ ).manual_seed(A_ )
							A  				=				{
							    '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       : List[str] )       ->				List[str]:
							A  				=				'cpu'  # ensure determinism for the device-dependent torch.Generator
							A  				=				self.get_dummy_components()
							A  				=				StableDiffusionParadigmsPipeline(**A_ )
							A  				=				sd_pipe.to(A_ )
							sd_pipe.set_progress_bar_config(disable=A_ )

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

							A  				=				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   _SCREAMING_SNAKE_CASE (	self       : Tuple )       ->				Optional[int]:
							A  				=				'cpu'  # ensure determinism for the device-dependent torch.Generator
							A  				=				self.get_dummy_components()
							torch.manual_seed(0 )
							A  				=				DDPMParallelScheduler()
							torch.manual_seed(0 )
							A  				=				StableDiffusionParadigmsPipeline(**A_ )
							A  				=				sd_pipe.to(A_ )
							sd_pipe.set_progress_bar_config(disable=A_ )

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

							A  				=				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   _SCREAMING_SNAKE_CASE (	self       : List[str] )       ->				int:
							super().test_inference_batch_consistent(batch_sizes=[1, 2] )


						def   _SCREAMING_SNAKE_CASE (	self       : str )       ->				Any:
							super().test_inference_batch_single_identical(batch_size=2					,expected_max_diff=3e-3 )




						def   _SCREAMING_SNAKE_CASE (	self       : Optional[Any] )       ->				List[str]:
							A  				=				'cpu'  # ensure determinism for the device-dependent torch.Generator
							A  				=				self.get_dummy_components()
							A  				=				StableDiffusionParadigmsPipeline(**A_ )
							A  				=				sd_pipe.to(A_ )
							sd_pipe.set_progress_bar_config(disable=A_ )

							A  				=				self.get_dummy_inputs(A_ )
							A  				=				'french fries'
							A  				=				sd_pipe(**A_					,negative_prompt=A_ )
							A  				=				output.images
							A  				=				image[0, -3:, -3:, -1]

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

							A  				=				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 			lowerCAmelCase_			( unittest.TestCase      ):



						'''simple docstring'''


						def   _SCREAMING_SNAKE_CASE (	self       : Dict )       ->				Dict:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


						def   _SCREAMING_SNAKE_CASE (	self       : Union[str, Any]					,A_       : Union[str, Any]=0 )       ->				Tuple:
							A  				=				torch.Generator(device=A_ ).manual_seed(A_ )
							A  				=				{
							    '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       : Optional[Any] )       ->				Dict:
							A  				=				'stabilityai/stable-diffusion-2-base'
							A  				=				DDIMParallelScheduler.from_pretrained(A_					,subfolder='scheduler' )
							A  				=				StableDiffusionParadigmsPipeline.from_pretrained(A_					,scheduler=A_					,safety_checker=A_ )
							pipe.to(A_ )
							pipe.set_progress_bar_config(disable=A_ )
							pipe.enable_attention_slicing()

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

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

							A  				=				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

"""simple docstring"""


from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, logging, randn_tensor
from .scheduling_utils import SchedulerMixin


_lowercase											=			logging.get_logger(__name__)  # pylint: disable=invalid-name







@dataclass
class 			lowerCAmelCase_			( _lowercase      ):



						'''simple docstring'''


						_lowerCamelCase:      torch.FloatTensor







class 			lowerCAmelCase_			( _lowercase						,   _lowercase      ):



						'''simple docstring'''


						_lowerCamelCase:      Dict =      1


						@register_to_config
						def __init__(	self       : List[Any]					,A_       : int = 40					,A_       : float = 0.0_02					,A_       : float = 80.0					,A_       : float = 0.5					,A_       : float = 1.0					,A_       : float = 7.0					,A_       : bool = True					,)       ->				List[str]:
							# standard deviation of the initial noise distribution
							A  				=				sigma_max

							A  				=				np.linspace(0					,1					,A_ )
							A  				=				self._convert_to_karras(A_ )
							A  				=				self.sigma_to_t(A_ )

							# setable values
							A  				=				None
							A  				=				torch.from_numpy(A_ )
							A  				=				torch.from_numpy(A_ )
							A  				=				False
							A  				=				False


						def   _SCREAMING_SNAKE_CASE (	self       : Union[str, Any]					,A_       : List[Any]					,A_       : List[Any]=None )       ->				List[Any]:
							if schedule_timesteps is None:
								A  				=				self.timesteps

							A  				=				(schedule_timesteps == timestep).nonzero()
							return indices.item()


						def   _SCREAMING_SNAKE_CASE (	self       : List[str]					,A_       : torch.FloatTensor					,A_       : Union[float, torch.FloatTensor] )       ->				torch.FloatTensor:
							# Get sigma corresponding to timestep
							if isinstance(A_					,torch.Tensor ):
								A  				=				timestep.to(self.timesteps.device )
							A  				=				self.index_for_timestep(A_ )
							A  				=				self.sigmas[step_idx]

							A  				=				sample / ((sigma**2 + self.config.sigma_data**2) ** 0.5)

							A  				=				True
							return sample


						def   _SCREAMING_SNAKE_CASE (	self       : Union[str, Any]					,A_       : Union[float, np.ndarray] )       ->				Tuple:
							if not isinstance(A_					,np.ndarray ):
								A  				=				np.array(A_					,dtype=np.floataa )

							A  				=				1000 * 0.25 * np.log(sigmas + 1e-44 )

							return timesteps


						def   _SCREAMING_SNAKE_CASE (	self       : List[Any]					,A_       : Optional[int] = None					,A_       : Union[str, torch.device] = None					,A_       : Optional[List[int]] = None					,)       ->				Optional[Any]:
							if num_inference_steps is None and timesteps is None:
								raise ValueError('Exactly one of `num_inference_steps` or `timesteps` must be supplied.' )

							if num_inference_steps is not None and timesteps is not None:
								raise ValueError('Can only pass one of `num_inference_steps` or `timesteps`.' )

							# Follow DDPMScheduler custom timesteps logic
							if timesteps is not None:
								for i in range(1					,len(A_ ) ):
									if timesteps[i] >= timesteps[i - 1]:
										raise ValueError('`timesteps` must be in descending order.' )

								if timesteps[0] >= self.config.num_train_timesteps:
									raise ValueError(
									    F'`timesteps` must start before `self.config.train_timesteps`:'
									    F' {self.config.num_train_timesteps}.' )

								A  				=				np.array(A_					,dtype=np.intaa )
								A  				=				True
							else:
								if num_inference_steps > self.config.num_train_timesteps:
									raise ValueError(
									    F'`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:'
									    F' {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle'
									    F' maximal {self.config.num_train_timesteps} timesteps.' )

								A  				=				num_inference_steps

								A  				=				self.config.num_train_timesteps // self.num_inference_steps
								A  				=				(np.arange(0					,A_ ) * step_ratio).round()[::-1].copy().astype(np.intaa )
								A  				=				False

							# Map timesteps to Karras sigmas directly for multistep sampling
							# See https://github.com/openai/consistency_models/blob/main/cm/karras_diffusion.py#L675
							A  				=				self.config.num_train_timesteps
							A  				=				timesteps[::-1].copy()
							A  				=				ramp / (num_train_timesteps - 1)
							A  				=				self._convert_to_karras(A_ )
							A  				=				self.sigma_to_t(A_ )

							A  				=				np.concatenate([sigmas, [self.sigma_min]] ).astype(np.floataa )
							A  				=				torch.from_numpy(A_ ).to(device=A_ )

							if str(A_ ).startswith('mps' ):
								# mps does not support float64
								A  				=				torch.from_numpy(A_ ).to(A_					,dtype=torch.floataa )
							else:
								A  				=				torch.from_numpy(A_ ).to(device=A_ )


						def   _SCREAMING_SNAKE_CASE (	self       : Tuple					,A_       : Optional[Any] )       ->				Any:
							A  				=				self.config.sigma_min
							A  				=				self.config.sigma_max

							A  				=				self.config.rho
							A  				=				sigma_min ** (1 / rho)
							A  				=				sigma_max ** (1 / rho)
							A  				=				(max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
							return sigmas


						def   _SCREAMING_SNAKE_CASE (	self       : Optional[int]					,A_       : Dict )       ->				Any:
							A  				=				self.config.sigma_data

							A  				=				sigma_data**2 / (sigma**2 + sigma_data**2)
							A  				=				sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
							return c_skip, c_out


						def   _SCREAMING_SNAKE_CASE (	self       : Tuple					,A_       : Optional[Any] )       ->				Dict:
							A  				=				self.config.sigma_min
							A  				=				self.config.sigma_data

							A  				=				sigma_data**2 / ((sigma - sigma_min) ** 2 + sigma_data**2)
							A  				=				(sigma - sigma_min) * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
							return c_skip, c_out


						def   _SCREAMING_SNAKE_CASE (	self       : List[Any]					,A_       : torch.FloatTensor					,A_       : Union[float, torch.FloatTensor]					,A_       : torch.FloatTensor					,A_       : Optional[torch.Generator] = None					,A_       : bool = True					,)       ->				Union[CMStochasticIterativeSchedulerOutput, Tuple]:

							if (
							    isinstance(A_					,A_ )
							    or isinstance(A_					,torch.IntTensor )
							    or isinstance(A_					,torch.LongTensor )
							):
								raise ValueError(
								    (
								        'Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to'
								        F' `{self.__class__}.step()` is not supported. Make sure to pass'
								        ' one of the `scheduler.timesteps` as a timestep.'
								    )					,)

							if not self.is_scale_input_called:
								logger.warning(
								    'The `scale_model_input` function should be called before `step` to ensure correct denoising. '
								    'See `StableDiffusionPipeline` for a usage example.' )

							if isinstance(A_					,torch.Tensor ):
								A  				=				timestep.to(self.timesteps.device )

							A  				=				self.config.sigma_min
							A  				=				self.config.sigma_max

							A  				=				self.index_for_timestep(A_ )

							# sigma_next corresponds to next_t in original implementation
							A  				=				self.sigmas[step_index]
							if step_index + 1 < self.config.num_train_timesteps:
								A  				=				self.sigmas[step_index + 1]
							else:
								# Set sigma_next to sigma_min
								A  				=				self.sigmas[-1]

							# Get scalings for boundary conditions
							A				, A  				=				self.get_scalings_for_boundary_condition(A_ )

							# 1. Denoise model output using boundary conditions
							A  				=				c_out * model_output + c_skip * sample
							if self.config.clip_denoised:
								A  				=				denoised.clamp(-1					,1 )

							# 2. Sample z ~ N(0, s_noise^2 * I)
							# Noise is not used for onestep sampling.
							if len(self.timesteps ) > 1:
								A  				=				randn_tensor(
								    model_output.shape					,dtype=model_output.dtype					,device=model_output.device					,generator=A_ )
							else:
								A  				=				torch.zeros_like(A_ )
							A  				=				noise * self.config.s_noise

							A  				=				sigma_next.clamp(min=A_					,max=A_ )

							# 3. Return noisy sample
							# tau = sigma_hat, eps = sigma_min
							A  				=				denoised + z * (sigma_hat**2 - sigma_min**2) ** 0.5

							if not return_dict:
								return (prev_sample,)

							return CMStochasticIterativeSchedulerOutput(prev_sample=A_ )


						def   _SCREAMING_SNAKE_CASE (	self       : int					,A_       : torch.FloatTensor					,A_       : torch.FloatTensor					,A_       : torch.FloatTensor					,)       ->				torch.FloatTensor:
							# Make sure sigmas and timesteps have the same device and dtype as original_samples
							A  				=				self.sigmas.to(device=original_samples.device					,dtype=original_samples.dtype )
							if original_samples.device.type == "mps" and torch.is_floating_point(A_ ):
								# mps does not support float64
								A  				=				self.timesteps.to(original_samples.device					,dtype=torch.floataa )
								A  				=				timesteps.to(original_samples.device					,dtype=torch.floataa )
							else:
								A  				=				self.timesteps.to(original_samples.device )
								A  				=				timesteps.to(original_samples.device )

							A  				=				[(schedule_timesteps == t).nonzero().item() for t in timesteps]

							A  				=				sigmas[step_indices].flatten()
							while len(sigma.shape ) < len(original_samples.shape ):
								A  				=				sigma.unsqueeze(-1 )

							A  				=				original_samples + noise * sigma
							return noisy_samples




						def __len__(	self       : List[str] )       ->				List[Any]:
							return self.config.num_train_timesteps