code: 
  
  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
  
  
  lowerCamelCase     =      logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  
  
  
  
  
  @dataclass
  class 		__magic_name__    (							lowerCamelCase__							):
  
     '''simple docstring'''
  
  
  
  
     lowerCamelCase__				:			torch.FloatTensor
  
  
  
  
  
  
  
  class 		__magic_name__    (							lowerCamelCase__	, lowerCamelCase__							):
  
     '''simple docstring'''
  
  
  
  
     lowerCamelCase__				:			Optional[int]												=     1
  
  
  
     @register_to_config
     def __init__(       self,					lowercase_ = 40,					lowercase_ = 0.002,					lowercase_ = 80.0,					lowercase_ = 0.5,					lowercase_ = 1.0,					lowercase_ = 7.0,					lowercase_ = True,					)							->					Optional[Any]:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       a__						=sigma_max
  
       a__						=np.linspace(0,					1,					lowercase_ )
       a__						=self._convert_to_karras(lowercase_ )
       a__						=self.sigma_to_t(lowercase_ )
  
       # setable values
       a__						=None
       a__						=torch.from_numpy(lowercase_ )
       a__						=torch.from_numpy(lowercase_ )
       a__						=False
       a__						=False
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_=None )							->					str:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       if schedule_timesteps is None:
         a__						=self.timesteps
  
       a__						=(schedule_timesteps == timestep).nonzero()
       return indices.item()
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_ )							->					torch.FloatTensor:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       if isinstance(lowercase_,					torch.Tensor ):
         a__						=timestep.to(self.timesteps.device )
       a__						=self.index_for_timestep(lowercase_ )
       a__						=self.sigmas[step_idx]
  
       a__						=sample / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
  
       a__						=True
       return sample
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_ )							->					Optional[int]:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       if not isinstance(lowercase_,					np.ndarray ):
         a__						=np.array(lowercase_,					dtype=np.floataa )
  
       a__						=1000 * 0.25 * np.log(sigmas + 1E-44 )
  
       return timesteps
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_ = None,					lowercase_ = None,					lowercase_ = None,					)							->					int:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       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(lowercase_ ) ):
           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(lowercase_,					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,					lowercase_ ) * 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(lowercase_ )
       a__						=self.sigma_to_t(lowercase_ )
  
       a__						=np.concatenate([sigmas, [self.sigma_min]] ).astype(np.floataa )
       a__						=torch.from_numpy(lowercase_ ).to(device=lowercase_ )
  
       if str(lowercase_ ).startswith('''mps''' ):
         # mps does not support float64
         a__						=torch.from_numpy(lowercase_ ).to(lowercase_,					dtype=torch.floataa )
       else:
         a__						=torch.from_numpy(lowercase_ ).to(device=lowercase_ )
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_ )							->					int:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       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 				_UpperCAmelCase				(       self,					lowercase_ )							->					Optional[Any]:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       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 				_UpperCAmelCase				(       self,					lowercase_ )							->					List[Any]:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       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 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_ = None,					lowercase_ = True,					)							->					Union[CMStochasticIterativeSchedulerOutput, Tuple]:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       if (
           isinstance(lowercase_,					lowercase_ )
           or isinstance(lowercase_,					torch.IntTensor )
           or isinstance(lowercase_,					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(lowercase_,					torch.Tensor ):
         a__						=timestep.to(self.timesteps.device )
  
       a__						=self.config.sigma_min
       a__						=self.config.sigma_max
  
       a__						=self.index_for_timestep(lowercase_ )
  
       # 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(lowercase_ )
  
       # 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=lowercase_ )
       else:
         a__						=torch.zeros_like(lowercase_ )
       a__						=noise * self.config.s_noise
  
       a__						=sigma_next.clamp(min=lowercase_,					max=lowercase_ )
  
       # 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=lowercase_ )
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					)							->					torch.FloatTensor:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       a__						=self.sigmas.to(device=original_samples.device,					dtype=original_samples.dtype )
       if original_samples.device.type == "mps" and torch.is_floating_point(lowercase_ ):
         # 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 )							->					Tuple:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       return self.config.num_train_timesteps
  
style_context: 
  
  from dataclasses import dataclass
  from typing import List, Optional, Union
  
  import numpy as np
  import PIL
  from PIL import Image
  
  from ...utils import BaseOutput, is_torch_available, is_transformers_available
  
  
  
  
  
  
  
  @dataclass
  # Copied from diffusers.pipelines.stable_diffusion.__init__.StableDiffusionPipelineOutput with Stable->Alt
  class 		__magic_name__    (							lowerCamelCase__							):
  
     '''simple docstring'''
  
  
  
  
     lowerCamelCase__				:			Union[List[PIL.Image.Image], np.ndarray]
     lowerCamelCase__				:			Optional[List[bool]]
  
  
  if is_transformers_available() and is_torch_available():
         from .modeling_roberta_series import RobertaSeriesModelWithTransformation
         from .pipeline_alt_diffusion import AltDiffusionPipeline
         from .pipeline_alt_diffusion_imgaimg import AltDiffusionImgaImgPipeline
  