




'''simple docstring'''







# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
# and https://github.com/hojonathanho/diffusion

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

import flax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)





@flax.struct.dataclass
class 	lowerCamelCase_				:




					'''simple docstring'''






					lowercase__       =					4_2
					lowercase__       =					4_2

					# setable values
					lowercase__       =					4_2
					lowercase__       =					4_2
					lowercase__       =					None






					@classmethod
					def 							__SCREAMING_SNAKE_CASE      (       cls,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    ):
							return cls(
							    common=_lowerCAmelCase,    final_alpha_cumprod=_lowerCAmelCase,    init_noise_sigma=_lowerCAmelCase,    timesteps=_lowerCAmelCase,    )





@dataclass
class 	lowerCamelCase_				(			snake_case_				):




					'''simple docstring'''






					lowercase__       =					4_2





class 	lowerCamelCase_				(			snake_case_,			snake_case_				):




					'''simple docstring'''






					lowercase__       =					[e.name for e in FlaxKarrasDiffusionSchedulers]

					lowercase__       =					4_2






					@property
					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return True






					@register_to_config
					def __init__(       self,    __UpperCamelCase = 1_000,    __UpperCamelCase = 0.00_01,    __UpperCamelCase = 0.02,    __UpperCamelCase = "linear",    __UpperCamelCase = None,    __UpperCamelCase = True,    __UpperCamelCase = 0,    __UpperCamelCase = "epsilon",    __UpperCamelCase = jnp.floataa,    ):
							_snake_case      			=			dtype






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase = None					):
							if common is None:
									_snake_case      			=			CommonSchedulerState.create(self					)

							# At every step in ddim, we are looking into the previous alphas_cumprod
							# For the final step, there is no previous alphas_cumprod because we are already at 0
							# `set_alpha_to_one` decides whether we set this parameter simply to one or
							# whether we use the final alpha of the "non-previous" one.
							_snake_case      			=			(
							    jnp.array(1.0,    dtype=self.dtype					) if self.config.set_alpha_to_one else common.alphas_cumprod[0]
							)

							# standard deviation of the initial noise distribution
							_snake_case      			=			jnp.array(1.0,    dtype=self.dtype					)

							_snake_case      			=			jnp.arange(0,    self.config.num_train_timesteps					).round()[::-1]

							return DDIMSchedulerState.create(
							    common=_lowerCAmelCase,    final_alpha_cumprod=_lowerCAmelCase,    init_noise_sigma=_lowerCAmelCase,    timesteps=_lowerCAmelCase,    )






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase = None					):
							return sample






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase = ()					):
							_snake_case      			=			self.config.num_train_timesteps // num_inference_steps
							# creates integer timesteps by multiplying by ratio
							# rounding to avoid issues when num_inference_step is power of 3
							_snake_case      			=			(jnp.arange(0,    _lowerCAmelCase					) * step_ratio).round()[::-1] + self.config.steps_offset

							return state.replace(
							    num_inference_steps=_lowerCAmelCase,    timesteps=_lowerCAmelCase,    )






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase					):
							_snake_case      			=			state.common.alphas_cumprod[timestep]
							_snake_case      			=			jnp.where(
							    prev_timestep >= 0,    state.common.alphas_cumprod[prev_timestep],    state.final_alpha_cumprod					)
							_snake_case      			=			1 - alpha_prod_t
							_snake_case      			=			1 - alpha_prod_t_prev

							_snake_case      			=			(beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

							return variance






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase = 0.0,    __UpperCamelCase = True,    ):
							if state.num_inference_steps is None:
									raise ValueError(
									    "Number of inference steps is \'None\', you need to run \'set_timesteps\' after creating the scheduler"					)

							# See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
							# Ideally, read DDIM paper in-detail understanding

							# Notation (<variable name> -> <name in paper>
							# - pred_noise_t -> e_theta(x_t, t)
							# - pred_original_sample -> f_theta(x_t, t) or x_0
							# - std_dev_t -> sigma_t
							# - eta -> η
							# - pred_sample_direction -> "direction pointing to x_t"
							# - pred_prev_sample -> "x_t-1"

							# 1. get previous step value (=t-1)
							_snake_case      			=			timestep - self.config.num_train_timesteps // state.num_inference_steps

							_snake_case      			=			state.common.alphas_cumprod
							_snake_case      			=			state.final_alpha_cumprod

							# 2. compute alphas, betas
							_snake_case      			=			alphas_cumprod[timestep]
							_snake_case      			=			jnp.where(prev_timestep >= 0,    alphas_cumprod[prev_timestep],    _lowerCAmelCase					)

							_snake_case      			=			1 - alpha_prod_t

							# 3. compute predicted original sample from predicted noise also called
							# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
							if self.config.prediction_type == "epsilon":
									_snake_case      			=			(sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
									_snake_case      			=			model_output
							elif self.config.prediction_type == "sample":
									_snake_case      			=			model_output
									_snake_case      			=			(sample - alpha_prod_t ** 0.5 * pred_original_sample) / beta_prod_t ** 0.5
							elif self.config.prediction_type == "v_prediction":
									_snake_case      			=			(alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
									_snake_case      			=			(alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
							else:
									raise ValueError(
									    f'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or'''
									    " `v_prediction`"					)

							# 4. compute variance: "sigma_t(η)" -> see formula (16)
							# σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
							_snake_case      			=			self._get_variance(_lowerCAmelCase,    _lowerCAmelCase,    _lowerCAmelCase					)
							_snake_case      			=			eta * variance ** 0.5

							# 5. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
							_snake_case      			=			(1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * pred_epsilon

							# 6. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
							_snake_case      			=			alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction

							if not return_dict:
									return (prev_sample, state)

							return FlaxDDIMSchedulerOutput(prev_sample=_lowerCAmelCase,    state=_lowerCAmelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    ):
							return add_noise_common(state.common,    _lowerCAmelCase,    _lowerCAmelCase,    _lowerCAmelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase,    ):
							return get_velocity_common(state.common,    _lowerCAmelCase,    _lowerCAmelCase,    _lowerCAmelCase					)






					def __len__(       self					):
							return self.config.num_train_timesteps










import copy

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


UpperCAmelCase_          =				logging.get_logger(__name__)

UpperCAmelCase_          =				{
    "facebook/sam-vit-huge": "https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json",
    "facebook/sam-vit-large": "https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json",
    "facebook/sam-vit-base": "https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json",
}





class 		UpperCAmelCase    (    snake_case_					):





   def __init__(						self							,  _lowerCAmelCase=256							,  _lowerCAmelCase=1_024							,  _lowerCAmelCase=16							,  _lowerCAmelCase=16							,  _lowerCAmelCase=4							,  _lowerCAmelCase="gelu"							,  _lowerCAmelCase=1E-6							,  **_lowerCAmelCase							,  ):
         super().__init__(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	hidden_size
         _lowerCAmelCase			      =	image_size
         _lowerCAmelCase			      =	patch_size
         _lowerCAmelCase			      =	image_size // patch_size
         _lowerCAmelCase			      =	mask_input_channels
         _lowerCAmelCase			      =	num_point_embeddings
         _lowerCAmelCase			      =	hidden_act
         _lowerCAmelCase			      =	layer_norm_eps






class 		UpperCAmelCase    (    snake_case_					):





   def __init__(						self							,  _lowerCAmelCase=256							,  _lowerCAmelCase="relu"							,  _lowerCAmelCase=2_048							,  _lowerCAmelCase=2							,  _lowerCAmelCase=8							,  _lowerCAmelCase=2							,  _lowerCAmelCase=3							,  _lowerCAmelCase=3							,  _lowerCAmelCase=256							,  _lowerCAmelCase=1E-6							,  **_lowerCAmelCase							,  ):
         super().__init__(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	hidden_size
         _lowerCAmelCase			      =	hidden_act
         _lowerCAmelCase			      =	mlp_dim
         _lowerCAmelCase			      =	num_hidden_layers
         _lowerCAmelCase			      =	num_attention_heads
         _lowerCAmelCase			      =	attention_downsample_rate
         _lowerCAmelCase			      =	num_multimask_outputs
         _lowerCAmelCase			      =	iou_head_depth
         _lowerCAmelCase			      =	iou_head_hidden_dim
         _lowerCAmelCase			      =	layer_norm_eps






class 		UpperCAmelCase    (    snake_case_					):





   def __init__(						self							,  _lowerCAmelCase=768							,  _lowerCAmelCase=256							,  _lowerCAmelCase=12							,  _lowerCAmelCase=12							,  _lowerCAmelCase=3							,  _lowerCAmelCase=1_024							,  _lowerCAmelCase=16							,  _lowerCAmelCase="gelu"							,  _lowerCAmelCase=1E-06							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=1E-10							,  _lowerCAmelCase=True							,  _lowerCAmelCase=4.0							,  _lowerCAmelCase=True							,  _lowerCAmelCase=True							,  _lowerCAmelCase=14							,  _lowerCAmelCase=[2, 5, 8, 11]							,  _lowerCAmelCase=128							,  _lowerCAmelCase=None							,  **_lowerCAmelCase							,  ):
         super().__init__(**_lowerCAmelCase      )

         _lowerCAmelCase			      =	hidden_size
         _lowerCAmelCase			      =	output_channels
         _lowerCAmelCase			      =	num_hidden_layers
         _lowerCAmelCase			      =	num_attention_heads
         _lowerCAmelCase			      =	num_channels
         _lowerCAmelCase			      =	image_size
         _lowerCAmelCase			      =	patch_size
         _lowerCAmelCase			      =	hidden_act
         _lowerCAmelCase			      =	layer_norm_eps
         _lowerCAmelCase			      =	attention_dropout
         _lowerCAmelCase			      =	initializer_range
         _lowerCAmelCase			      =	qkv_bias
         _lowerCAmelCase			      =	mlp_ratio
         _lowerCAmelCase			      =	use_abs_pos
         _lowerCAmelCase			      =	use_rel_pos
         _lowerCAmelCase			      =	window_size
         _lowerCAmelCase			      =	global_attn_indexes
         _lowerCAmelCase			      =	num_pos_feats
         _lowerCAmelCase			      =	int(hidden_size * mlp_ratio      ) if mlp_dim is None else mlp_dim






class 		UpperCAmelCase    (    snake_case_					):
   SCREAMING_SNAKE_CASE__					      =  '''sam'''
   SCREAMING_SNAKE_CASE__					      =  True





   def __init__(						self							,  _lowerCAmelCase=None							,  _lowerCAmelCase=None							,  _lowerCAmelCase=None							,  _lowerCAmelCase=0.02							,  **_lowerCAmelCase							,  ):
         super().__init__(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	vision_config if vision_config is not None else {}
         _lowerCAmelCase			      =	prompt_encoder_config if prompt_encoder_config is not None else {}
         _lowerCAmelCase			      =	mask_decoder_config if mask_decoder_config is not None else {}

         if isinstance(_lowerCAmelCase							,  _lowerCAmelCase      ):
               _lowerCAmelCase			      =	vision_config.to_dict()
         if isinstance(_lowerCAmelCase							,  _lowerCAmelCase      ):
               _lowerCAmelCase			      =	prompt_encoder_config.to_dict()
         if isinstance(_lowerCAmelCase							,  _lowerCAmelCase      ):
               _lowerCAmelCase			      =	mask_decoder_config.to_dict()

         _lowerCAmelCase			      =	SamVisionConfig(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	SamPromptEncoderConfig(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	SamMaskDecoderConfig(**_lowerCAmelCase      )
         _lowerCAmelCase			      =	initializer_range





   def 						__lowerCAmelCase				(						self      ):
         _lowerCAmelCase			      =	copy.deepcopy(self.__dict__      )
         _lowerCAmelCase			      =	self.vision_config.to_dict()
         _lowerCAmelCase			      =	self.prompt_encoder_config.to_dict()
         _lowerCAmelCase			      =	self.mask_decoder_config.to_dict()
         _lowerCAmelCase			      =	self.__class__.model_type
         return output