class_docstr: A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
  shaped output.
  
  This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for it's generic methods
  implemented for all models (such as downloading or saving).
  
  This model is also a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module)
  subclass. Use it as a regular Flax Linen module and refer to the Flax documentation for all matters related to its
  general usage and behavior.
  
  Inherent JAX features such as the following are supported:
  - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
  - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
  - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
  - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)
  
  Parameters:
      sample_size (`int`, *optional*):
          The size of the input sample.
      in_channels (`int`, *optional*, defaults to 4):
          The number of channels in the input sample.
      out_channels (`int`, *optional*, defaults to 4):
          The number of channels in the output.
      down_block_types (`Tuple[str]`, *optional*, defaults to `("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")`):
          The tuple of downsample blocks to use.
      up_block_types (`Tuple[str]`, *optional*, defaults to `("FlaxUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D")`):
          The tuple of upsample blocks to use.
      block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
          The tuple of output channels for each block.
      layers_per_block (`int`, *optional*, defaults to 2):
          The number of layers per block.
      attention_head_dim (`int` or `Tuple[int]`, *optional*, defaults to 8):
          The dimension of the attention heads.
      num_attention_heads (`int` or `Tuple[int]`, *optional*):
          The number of attention heads.
      cross_attention_dim (`int`, *optional*, defaults to 768):
          The dimension of the cross attention features.
      dropout (`float`, *optional*, defaults to 0):
          Dropout probability for down, up and bottleneck blocks.
      flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
          Whether to flip the sin to cos in the time embedding.
      freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
      use_memory_efficient_attention (`bool`, *optional*, defaults to `False`):
          Enable memory efficient attention as described [here](https://arxiv.org/abs/2112.05682).
      split_head_dim (`bool`, *optional*, defaults to `False`):
          Whether to split the head dimension into a new axis for the self-attention computation. In most cases,
          enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL.
class_docstr_tok: ["A","conditional","2D","UNet","model","that","takes","a","noisy","sample",",","conditional","state",",","and","a","timestep","and","returns","a","sample","shaped","output",".","This","model","inherits","from","[","`","FlaxModelMixin","`","]",".","Check","the","superclass","documentation","for","it","'s","generic","methods","implemented","for","all","models","(","such","as","downloading","or","saving",")",".","This","model","is","also","a","Flax","Linen","[","flax.linen.Module","]","(","https",":","\/\/flax.readthedocs.io\/en\/latest\/flax.linen.html","#","module",")","subclass",".","Use","it","as","a","regular","Flax","Linen","module","and","refer","to","the","Flax","documentation","for","all","matters","related","to","its","general","usage","and","behavior",".","Inherent","JAX","features","such","as","the","following","are","supported",":","-","[","Just-In-Time","(","JIT",")","compilation","]","(","https",":","\/\/jax.readthedocs.io\/en\/latest\/jax.html","#","just-in-time-compilation-jit",")","-","[","Automatic","Differentiation","]","(","https",":","\/\/jax.readthedocs.io\/en\/latest\/jax.html","#","automatic-differentiation",")","-","[","Vectorization","]","(","https",":","\/\/jax.readthedocs.io\/en\/latest\/jax.html","#","vectorization-vmap",")","-","[","Parallelization","]","(","https",":","\/\/jax.readthedocs.io\/en\/latest\/jax.html","#","parallelization-pmap",")","Parameters",":","sample_size","(","`","int","`",",","*","optional","*",")",":","The","size","of","the","input","sample",".","in_channels","(","`","int","`",",","*","optional","*",",","defaults","to","4",")",":","The","number","of","channels","in","the","input","sample",".","out_channels","(","`","int","`",",","*","optional","*",",","defaults","to","4",")",":","The","number","of","channels","in","the","output",".","down_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","``","FlaxCrossAttnDownBlock2D","''",",","``","FlaxCrossAttnDownBlock2D","''",",","``","FlaxCrossAttnDownBlock2D","''",",","``","FlaxDownBlock2D","''",")","`",")",":","The","tuple","of","downsample","blocks","to","use",".","up_block_types","(","`","Tuple","[","str","]","`",",","*","optional","*",",","defaults","to","`","(","``","FlaxUpBlock2D","''",",","``","FlaxCrossAttnUpBlock2D","''",",","``","FlaxCrossAttnUpBlock2D","''",",","``","FlaxCrossAttnUpBlock2D","''",")","`",")",":","The","tuple","of","upsample","blocks","to","use",".","block_out_channels","(","`","Tuple","[","int","]","`",",","*","optional","*",",","defaults","to","`","(","320",",","640",",","1280",",","1280",")","`",")",":","The","tuple","of","output","channels","for","each","block",".","layers_per_block","(","`","int","`",",","*","optional","*",",","defaults","to","2",")",":","The","number","of","layers","per","block",".","attention_head_dim","(","`","int","`","or","`","Tuple","[","int","]","`",",","*","optional","*",",","defaults","to","8",")",":","The","dimension","of","the","attention","heads",".","num_attention_heads","(","`","int","`","or","`","Tuple","[","int","]","`",",","*","optional","*",")",":","The","number","of","attention","heads",".","cross_attention_dim","(","`","int","`",",","*","optional","*",",","defaults","to","768",")",":","The","dimension","of","the","cross","attention","features",".","dropout","(","`","float","`",",","*","optional","*",",","defaults","to","0",")",":","Dropout","probability","for","down",",","up","and","bottleneck","blocks",".","flip_sin_to_cos","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","to","flip","the","sin","to","cos","in","the","time","embedding",".","freq_shift","(","`","int","`",",","*","optional","*",",","defaults","to","0",")",":","The","frequency","shift","to","apply","to","the","time","embedding",".","use_memory_efficient_attention","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Enable","memory","efficient","attention","as","described","[","here","]","(","https",":","\/\/arxiv.org\/abs\/2112.05682",")",".","split_head_dim","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","to","split","the","head","dimension","into","a","new","axis","for","the","self-attention","computation",".","In","most","cases",",","enabling","this","flag","should","speed","up","the","computation","for","Stable","Diffusion","2.x","and","Stable","Diffusion","XL","."]
docstr: Args:
      sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor
      timestep (`jnp.ndarray` or `float` or `int`): timesteps
      encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states
      added_cond_kwargs: (`dict`, *optional*):
          A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
          are passed along to the UNet blocks.
      down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
          A tuple of tensors that if specified are added to the residuals of down unet blocks.
      mid_block_additional_residual: (`torch.Tensor`, *optional*):
          A tensor that if specified is added to the residual of the middle unet block.
      return_dict (`bool`, *optional*, defaults to `True`):
          Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a
          plain tuple.
      train (`bool`, *optional*, defaults to `False`):
          Use deterministic functions and disable dropout when not training.
  
  Returns:
      [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`:
      [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`.
      When returning a tuple, the first element is the sample tensor.
docstr_tok: ["Args",":","sample","(","`","jnp.ndarray","`",")",":","(","batch",",","channel",",","height",",","width",")","noisy","inputs","tensor","timestep","(","`","jnp.ndarray","`","or","`","float","`","or","`","int","`",")",":","timesteps","encoder_hidden_states","(","`","jnp.ndarray","`",")",":","(","batch_size",",","sequence_length",",","hidden_size",")","encoder","hidden","states","added_cond_kwargs",":","(","`","dict","`",",","*","optional","*",")",":","A","kwargs","dictionary","containing","additional","embeddings","that","if","specified","are","added","to","the","embeddings","that","are","passed","along","to","the","UNet","blocks",".","down_block_additional_residuals",":","(","`","tuple","`","of","`","torch.Tensor","`",",","*","optional","*",")",":","A","tuple","of","tensors","that","if","specified","are","added","to","the","residuals","of","down","unet","blocks",".","mid_block_additional_residual",":","(","`","torch.Tensor","`",",","*","optional","*",")",":","A","tensor","that","if","specified","is","added","to","the","residual","of","the","middle","unet","block",".","return_dict","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","models.unet_2d_condition_flax.FlaxUNet2DConditionOutput","`","]","instead","of","a","plain","tuple",".","train","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Use","deterministic","functions","and","disable","dropout","when","not","training",".","Returns",":","[","`","~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput","`","]","or","`","tuple","`",":","[","`","~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput","`","]","if","`","return_dict","`","is","True",",","otherwise","a","`","tuple","`",".","When","returning","a","tuple",",","the","first","element","is","the","sample","tensor","."]
code: def __call__(
      self,
      sample,
      timesteps,
      encoder_hidden_states,
      added_cond_kwargs: Optional[Union[Dict, FrozenDict]] = None,
      down_block_additional_residuals=None,
      mid_block_additional_residual=None,
      return_dict: bool = True,
      train: bool = False,
  ) -> Union[FlaxUNet2DConditionOutput, Tuple]:
      r"""
      Args:
          sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor
          timestep (`jnp.ndarray` or `float` or `int`): timesteps
          encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states
          added_cond_kwargs: (`dict`, *optional*):
              A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
              are passed along to the UNet blocks.
          down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
              A tuple of tensors that if specified are added to the residuals of down unet blocks.
          mid_block_additional_residual: (`torch.Tensor`, *optional*):
              A tensor that if specified is added to the residual of the middle unet block.
          return_dict (`bool`, *optional*, defaults to `True`):
              Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a
              plain tuple.
          train (`bool`, *optional*, defaults to `False`):
              Use deterministic functions and disable dropout when not training.
  
      Returns:
          [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`:
          [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`.
          When returning a tuple, the first element is the sample tensor.
      """
      # 1. time
      if not isinstance(timesteps, jnp.ndarray):
          timesteps = jnp.array([timesteps], dtype=jnp.int32)
      elif (
          isinstance(timesteps, jnp.ndarray)
          and len(timesteps.shape) == 0
      ):
          timesteps = timesteps.astype(dtype=jnp.float32)
          timesteps = jnp.expand_dims(timesteps, 0)
  
      t_emb = self.time_proj(timesteps)
      t_emb = self.time_embedding(t_emb)
  
      # additional embeddings
      aug_emb = None
      if self.addition_embed_type == "text_time":
          if added_cond_kwargs is None:
              raise ValueError(
                  f"Need to provide argument `added_cond_kwargs` for {self.__class__} when using `addition_embed_type={self.addition_embed_type}`"
              )
          text_embeds = added_cond_kwargs.get("text_embeds")
          if text_embeds is None:
              raise ValueError(
                  f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
              )
          time_ids = added_cond_kwargs.get("time_ids")
          if time_ids is None:
              raise ValueError(
                  f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
              )
          # compute time embeds
          time_embeds = self.add_time_proj(
              jnp.ravel(time_ids)
          )  # (1, 6) => (6,) => (6, 256)
          time_embeds = jnp.reshape(
              time_embeds, (text_embeds.shape[0], -1)
          )
          add_embeds = jnp.concatenate(
              [text_embeds, time_embeds], axis=-1
          )
          aug_emb = self.add_embedding(add_embeds)
  
      t_emb = t_emb + aug_emb if aug_emb is not None else t_emb
  
      # 2. pre-process
      sample = jnp.transpose(sample, (0, 2, 3, 1))
      sample = self.conv_in(sample)
  
      # 3. down
      down_block_res_samples = (sample,)
      for down_block in self.down_blocks:
          if isinstance(down_block, FlaxCrossAttnDownBlock2D):
              sample, res_samples = down_block(
                  sample,
                  t_emb,
                  encoder_hidden_states,
                  deterministic=not train,
              )
          else:
              sample, res_samples = down_block(
                  sample, t_emb, deterministic=not train
              )
          down_block_res_samples += res_samples
  
      if down_block_additional_residuals is not None:
          new_down_block_res_samples = ()
  
          for (
              down_block_res_sample,
              down_block_additional_residual,
          ) in zip(
              down_block_res_samples, down_block_additional_residuals
          ):
              down_block_res_sample += down_block_additional_residual
              new_down_block_res_samples += (down_block_res_sample,)
  
          down_block_res_samples = new_down_block_res_samples
  
      # 4. mid
      sample = self.mid_block(
          sample, t_emb, encoder_hidden_states, deterministic=not train
      )
  
      if mid_block_additional_residual is not None:
          sample += mid_block_additional_residual
  
      # 5. up
      for up_block in self.up_blocks:
          res_samples = down_block_res_samples[
              -(self.layers_per_block + 1) :
          ]
          down_block_res_samples = down_block_res_samples[
              : -(self.layers_per_block + 1)
          ]
          if isinstance(up_block, FlaxCrossAttnUpBlock2D):
              sample = up_block(
                  sample,
                  temb=t_emb,
                  encoder_hidden_states=encoder_hidden_states,
                  res_hidden_states_tuple=res_samples,
                  deterministic=not train,
              )
          else:
              sample = up_block(
                  sample,
                  temb=t_emb,
                  res_hidden_states_tuple=res_samples,
                  deterministic=not train,
              )
  
      # 6. post-process
      sample = self.conv_norm_out(sample)
      sample = nn.silu(sample)
      sample = self.conv_out(sample)
      sample = jnp.transpose(sample, (0, 3, 1, 2))
  
      if not return_dict:
          return (sample,)
  
      return FlaxUNet2DConditionOutput(sample=sample)
  
file_path: llmga/llmga/diffusers/src/diffusers/models/unet_2d_condition_flax.py
imports: from typing import Dict, Optional, Tuple, Union
  import flax
  import flax.linen
  import jax
  import jax.numpy
  from flax.core.frozen_dict import FrozenDict
  from ..configuration_utils import ConfigMixin, flax_register_to_config
  from ..utils import BaseOutput
  from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
  from .modeling_flax_utils import FlaxModelMixin
  from .unet_2d_blocks_flax import FlaxCrossAttnDownBlock2D, FlaxCrossAttnUpBlock2D, FlaxDownBlock2D, FlaxUNetMidBlock2DCrossAttn, FlaxUpBlock2D
label_desc: Use image node_id 3 for calling the FlaxUNet2DConditionModel obj's underlying member method code with example usage: obj.__call__(sample, timesteps, encoder_hidden_states, added_cond_kwargs, down_block_additional_residuals, mid_block_additional_residual, return_dict, train) and returns: FlaxUNet2DConditionOutput, sample