code: 
  
  # DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch
  
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  
  import flax
  import jax.numpy as jnp
  from jax import random
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import FlaxSchedulerMixin, FlaxSchedulerOutput, broadcast_to_shape_from_left
  
  
  
  
  
  
  
  @flax.struct.dataclass
  class 		__magic_name__    :
  
     '''simple docstring'''
  
  
  
  
     lowerCamelCase__				:			Optional[jnp.ndarray]												=     None
     lowerCamelCase__				:			Optional[jnp.ndarray]												=     None
     lowerCamelCase__				:			Optional[jnp.ndarray]												=     None
  
  
  
     @classmethod
     def 				_UpperCAmelCase				(       cls )							->					Dict:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       return cls()
  
  
  
  
  
  
  
  @dataclass
  class 		__magic_name__    (							lowerCamelCase__							):
  
     '''simple docstring'''
  
  
  
  
     lowerCamelCase__				:			ScoreSdeVeSchedulerState
     lowerCamelCase__				:			jnp.ndarray
     lowerCamelCase__				:			Optional[jnp.ndarray]												=     None
  
  
  
  
  
  
  
  class 		__magic_name__    (							lowerCamelCase__	, lowerCamelCase__							):
  
     '''simple docstring'''
  
  
  
  
     @property
     def 				_UpperCAmelCase				(       self )							->					int:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       return True
  
  
  
     @register_to_config
     def __init__(       self,					lowercase_ = 2000,					lowercase_ = 0.15,					lowercase_ = 0.01,					lowercase_ = 1348.0,					lowercase_ = 1E-5,					lowercase_ = 1,					)							->					Optional[int]:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       pass
  
  
  
     def 				_UpperCAmelCase				(       self )							->					Dict:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       a__						=ScoreSdeVeSchedulerState.create()
       return self.set_sigmas(
           lowercase_,					self.config.num_train_timesteps,					self.config.sigma_min,					self.config.sigma_max,					self.config.sampling_eps,					)
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_ = (),					lowercase_ = None )							->					ScoreSdeVeSchedulerState:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       a__						=sampling_eps if sampling_eps is not None else self.config.sampling_eps
  
       a__						=jnp.linspace(1,					lowercase_,					lowercase_ )
       return state.replace(timesteps=lowercase_ )
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_ = None,					lowercase_ = None,					lowercase_ = None,					)							->					ScoreSdeVeSchedulerState:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       a__						=sigma_min if sigma_min is not None else self.config.sigma_min
       a__						=sigma_max if sigma_max is not None else self.config.sigma_max
       a__						=sampling_eps if sampling_eps is not None else self.config.sampling_eps
       if state.timesteps is None:
         a__						=self.set_timesteps(lowercase_,					lowercase_,					lowercase_ )
  
       a__						=jnp.exp(jnp.linspace(jnp.log(lowercase_ ),					jnp.log(lowercase_ ),					lowercase_ ) )
       a__						=jnp.array([sigma_min * (sigma_max / sigma_min) ** t for t in state.timesteps] )
  
       return state.replace(discrete_sigmas=lowercase_,					sigmas=lowercase_ )
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_ )							->					str:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       return jnp.where(timesteps == 0,					jnp.zeros_like(lowercase_ ),					state.discrete_sigmas[timesteps - 1] )
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ = True,					)							->					Union[FlaxSdeVeOutput, Tuple]:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       if state.timesteps is None:
         raise ValueError(
             '''`state.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler''' )
  
       a__						=timestep * jnp.ones(
           sample.shape[0],					)
       a__						=(timestep * (len(state.timesteps ) - 1)).long()
  
       a__						=state.discrete_sigmas[timesteps]
       a__						=self.get_adjacent_sigma(lowercase_,					lowercase_,					lowercase_ )
       a__						=jnp.zeros_like(lowercase_ )
       a__						=(sigma**2 - adjacent_sigma**2) ** 0.5
  
       # equation 6 in the paper: the model_output modeled by the network is grad_x log pt(x)
       # also equation 47 shows the analog from SDE models to ancestral sampling methods
       a__						=diffusion.flatten()
       a__						=broadcast_to_shape_from_left(lowercase_,					sample.shape )
       a__						=drift - diffusion**2 * model_output
  
       #  equation 6: sample noise for the diffusion term of
       a__						=random.split(lowercase_,					num=1 )
       a__						=random.normal(key=lowercase_,					shape=sample.shape )
       a__						=sample - drift  # subtract because `dt` is a small negative timestep
       # TODO is the variable diffusion the correct scaling term for the noise?
       a__						=prev_sample_mean + diffusion * noise  # add impact of diffusion field g
  
       if not return_dict:
         return (prev_sample, prev_sample_mean, state)
  
       return FlaxSdeVeOutput(prev_sample=lowercase_,					prev_sample_mean=lowercase_,					state=lowercase_ )
  
  
  
     def 				_UpperCAmelCase				(       self,					lowercase_,					lowercase_,					lowercase_,					lowercase_,					lowercase_ = True,					)							->					Union[FlaxSdeVeOutput, Tuple]:
  
  
  
  
  
  
       """simple docstring"""
  
  
  
  
  
  
       if state.timesteps is None:
         raise ValueError(
             '''`state.timesteps` is not set, you need to run \'set_timesteps\' after creating the scheduler''' )
  
       # For small batch sizes, the paper "suggest replacing norm(z) with sqrt(d), where d is the dim. of z"
       # sample noise for correction
       a__						=random.split(lowercase_,					num=1 )
       a__						=random.normal(key=lowercase_,					shape=sample.shape )
  
       # compute step size from the model_output, the noise, and the snr
       a__						=jnp.linalg.norm(lowercase_ )
       a__						=jnp.linalg.norm(lowercase_ )
       a__						=(self.config.snr * noise_norm / grad_norm) ** 2 * 2
       a__						=step_size * jnp.ones(sample.shape[0] )
  
       # compute corrected sample: model_output term and noise term
       a__						=step_size.flatten()
       a__						=broadcast_to_shape_from_left(lowercase_,					sample.shape )
       a__						=sample + step_size * model_output
       a__						=prev_sample_mean + ((step_size * 2) ** 0.5) * noise
  
       if not return_dict:
         return (prev_sample, state)
  
       return FlaxSdeVeOutput(prev_sample=lowercase_,					state=lowercase_ )
  
  
  
     def __len__(       self )							->					int:
  
  
  
  
  
  
       """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
  