



# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim

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

import flax
import jax
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   _A   :


						'''simple docstring'''





						_snake_case			:     CommonSchedulerState

						# setable values
						_snake_case			:     jnp.ndarray
						_snake_case			:     jnp.ndarray
						_snake_case			:     Optional[int]      							=      None



						@classmethod
						def     _snake_case	(	cls	:							Tuple							,      lowerCamelCase	:							CommonSchedulerState							,      lowerCamelCase	:							jnp.ndarray							,      lowerCamelCase	:							jnp.ndarray	):






							'''simple docstring'''



							return cls(common=lowerCamelCase							,      init_noise_sigma=lowerCamelCase							,      timesteps=lowerCamelCase	)

@dataclass
class   _A   ( _lowercase	):


						'''simple docstring'''





						_snake_case			:     DDPMSchedulerState

class   _A   ( _lowercase  ,						_lowercase	):


						'''simple docstring'''





						_snake_case			:     List[str]      							=      [e.name for e in FlaxKarrasDiffusionSchedulers]

						_snake_case			:     jnp.dtype



						@property
						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							return True



						@register_to_config
						def __init__(	self	:							Tuple							,      lowerCamelCase	:							int = 1_000							,      lowerCamelCase	:							float = 0.0001							,      lowerCamelCase	:							float = 0.02							,      lowerCamelCase	:							str = "linear"							,      lowerCamelCase	:							Optional[jnp.ndarray] = None							,      lowerCamelCase	:							str = "fixed_small"							,      lowerCamelCase	:							bool = True							,      lowerCamelCase	:							str = "epsilon"							,      lowerCamelCase	:							jnp.dtype = jnp.floataa							,      ):






							'''simple docstring'''



							__lowercase		       =			dtype



						def     _snake_case	(	self	:							Union[str, Any]							,      lowerCamelCase	:							Optional[CommonSchedulerState] = None	):






							'''simple docstring'''



							if common is None:
								__lowercase		       =			CommonSchedulerState.create(self	)

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

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

							return DDPMSchedulerState.create(
							    common=lowerCamelCase							,      init_noise_sigma=lowerCamelCase							,      timesteps=lowerCamelCase							,      )



						def     _snake_case	(	self	:							List[Any]							,      lowerCamelCase	:							DDPMSchedulerState							,      lowerCamelCase	:							jnp.ndarray							,      lowerCamelCase	:							Optional[int] = None	):






							'''simple docstring'''



							return sample



						def     _snake_case	(	self	:							Union[str, Any]							,      lowerCamelCase	:							DDPMSchedulerState							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Tuple = ()	):






							'''simple docstring'''



							__lowercase		       =			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
							__lowercase		       =			(jnp.arange(0							,      lowerCamelCase	) * step_ratio).round()[::-1]

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



						def     _snake_case	(	self	:							Tuple							,      lowerCamelCase	:							DDPMSchedulerState							,      lowerCamelCase	:							int							,      lowerCamelCase	:							int=None							,      lowerCamelCase	:							Tuple=None	):






							'''simple docstring'''



							__lowercase		       =			state.common.alphas_cumprod[t]
							__lowercase		       =			jnp.where(t > 0							,      state.common.alphas_cumprod[t - 1]							,      jnp.array(1.0							,      dtype=self.dtype	)	)

							# For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
							# and sample from it to get previous sample
							# x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
							__lowercase		       =			(1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t]

							if variance_type is None:
								__lowercase		       =			self.config.variance_type

							# hacks - were probably added for training stability
							if variance_type == "fixed_small":
								__lowercase		       =			jnp.clip(lowerCamelCase							,      a_min=1e-20	)
							# for rl-diffuser https://arxiv.org/abs/2205.09991
							elif variance_type == "fixed_small_log":
								__lowercase		       =			jnp.log(jnp.clip(lowerCamelCase							,      a_min=1e-20	)	)
							elif variance_type == "fixed_large":
								__lowercase		       =			state.common.betas[t]
							elif variance_type == "fixed_large_log":
								# Glide max_log
								__lowercase		       =			jnp.log(state.common.betas[t]	)
							elif variance_type == "learned":
								return predicted_variance
							elif variance_type == "learned_range":
								__lowercase		       =			variance
								__lowercase		       =			state.common.betas[t]
								__lowercase		       =			(predicted_variance + 1) / 2
								__lowercase		       =			frac * max_log + (1 - frac) * min_log

							return variance



						def     _snake_case	(	self	:							Dict							,      lowerCamelCase	:							DDPMSchedulerState							,      lowerCamelCase	:							jnp.ndarray							,      lowerCamelCase	:							int							,      lowerCamelCase	:							jnp.ndarray							,      lowerCamelCase	:							Optional[jax.random.KeyArray] = None							,      lowerCamelCase	:							bool = True							,      ):






							'''simple docstring'''



							__lowercase		       =			timestep

							if key is None:
								__lowercase		       =			jax.random.PRNGKey(0	)

							if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
								__lowercase       ,    __lowercase		       =			jnp.split(lowerCamelCase							,      sample.shape[1]							,      axis=1	)
							else:
								__lowercase		       =			None

							# 1. compute alphas, betas
							__lowercase		       =			state.common.alphas_cumprod[t]
							__lowercase		       =			jnp.where(t > 0							,      state.common.alphas_cumprod[t - 1]							,      jnp.array(1.0							,      dtype=self.dtype	)	)
							__lowercase		       =			1 - alpha_prod_t
							__lowercase		       =			1 - alpha_prod_t_prev

							# 2. compute predicted original sample from predicted noise also called
							# "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
							if self.config.prediction_type == "epsilon":
								__lowercase		       =			(sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
							elif self.config.prediction_type == "sample":
								__lowercase		       =			model_output
							elif self.config.prediction_type == "v_prediction":
								__lowercase		       =			(alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
							else:
								raise ValueError(
								    f"""prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` """
								    " for the FlaxDDPMScheduler."	)

							# 3. Clip "predicted x_0"
							if self.config.clip_sample:
								__lowercase		       =			jnp.clip(lowerCamelCase							,      -1							,      1	)

							# 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
							# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
							__lowercase		       =			(alpha_prod_t_prev ** 0.5 * state.common.betas[t]) / beta_prod_t
							__lowercase		       =			state.common.alphas[t] ** 0.5 * beta_prod_t_prev / beta_prod_t

							# 5. Compute predicted previous sample µ_t
							# See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
							__lowercase		       =			pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample

							# 6. Add noise
							def random_variance():
								__lowercase		       =			jax.random.split(lowerCamelCase							,      num=1	)
								__lowercase		       =			jax.random.normal(lowerCamelCase							,      shape=model_output.shape							,      dtype=self.dtype	)
								return (self._get_variance(lowerCamelCase							,      lowerCamelCase							,      predicted_variance=lowerCamelCase	) ** 0.5) * noise

							__lowercase		       =			jnp.where(t > 0							,      random_variance()							,      jnp.zeros(model_output.shape							,      dtype=self.dtype	)	)

							__lowercase		       =			pred_prev_sample + variance

							if not return_dict:
								return (pred_prev_sample, state)

							return FlaxDDPMSchedulerOutput(prev_sample=lowerCamelCase							,      state=lowerCamelCase	)



						def     _snake_case	(	self	:							Optional[int]							,      lowerCamelCase	:							DDPMSchedulerState							,      lowerCamelCase	:							jnp.ndarray							,      lowerCamelCase	:							jnp.ndarray							,      lowerCamelCase	:							jnp.ndarray							,      ):






							'''simple docstring'''



							return add_noise_common(state.common							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)



						def     _snake_case	(	self	:							Tuple							,      lowerCamelCase	:							DDPMSchedulerState							,      lowerCamelCase	:							jnp.ndarray							,      lowerCamelCase	:							jnp.ndarray							,      lowerCamelCase	:							jnp.ndarray							,      ):






							'''simple docstring'''



							return get_velocity_common(state.common							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)



						def __len__(	self	:							Tuple	):






							'''simple docstring'''



							return self.config.num_train_timesteps








import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin

class   _A   ( _lowercase  ,						_lowercase	):


						'''simple docstring'''





						@register_to_config
						def __init__(	self	:							Optional[Any]							,      *,
						    lowerCamelCase	:							int = 4							,      lowerCamelCase	:							int = 768							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Optional[int]							,      ):






							'''simple docstring'''



							super().__init__()

							__lowercase		       =			nn.Parameter(torch.zeros(lowerCamelCase	)	)

							# parameters for additional clip time embeddings
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)

							# parameters for encoder hidden states
							__lowercase		       =			clip_extra_context_tokens
							__lowercase		       =			nn.Linear(
							    lowerCamelCase							,      self.clip_extra_context_tokens * cross_attention_dim	)
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			nn.LayerNorm(lowerCamelCase	)



						def     _snake_case	(	self	:							Union[str, Any]							,      *, lowerCamelCase	:							Any							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							Tuple	):






							'''simple docstring'''



							if do_classifier_free_guidance:
								# Add the classifier free guidance embeddings to the image embeddings
								__lowercase		       =			image_embeddings.shape[0]
								__lowercase		       =			self.learned_classifier_free_guidance_embeddings.unsqueeze(0	)
								__lowercase		       =			classifier_free_guidance_embeddings.expand(
								    lowerCamelCase							,      -1	)
								__lowercase		       =			torch.cat([classifier_free_guidance_embeddings, image_embeddings]							,      dim=0	)

							# The image embeddings batch size and the text embeddings batch size are equal
							assert image_embeddings.shape[0] == prompt_embeds.shape[0]

							__lowercase		       =			prompt_embeds.shape[0]

							# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
							# adding CLIP embeddings to the existing timestep embedding, ...
							__lowercase		       =			self.embedding_proj(lowerCamelCase	)
							__lowercase		       =			self.clip_image_embeddings_project_to_time_embeddings(lowerCamelCase	)
							__lowercase		       =			time_projected_image_embeddings + time_projected_prompt_embeds

							# ... and by projecting CLIP embeddings into four
							# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
							__lowercase		       =			self.clip_extra_context_tokens_proj(lowerCamelCase	)
							__lowercase		       =			clip_extra_context_tokens.reshape(lowerCamelCase							,      -1							,      self.clip_extra_context_tokens	)
							__lowercase		       =			clip_extra_context_tokens.permute(0							,      2							,      1	)

							__lowercase		       =			self.encoder_hidden_states_proj(lowerCamelCase	)
							__lowercase		       =			self.text_encoder_hidden_states_norm(lowerCamelCase	)
							__lowercase		       =			torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]							,      dim=1	)

							return text_encoder_hidden_states, additive_clip_time_embeddings



