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","."]
code: def setup(self):
      block_out_channels = self.block_out_channels
      time_embed_dim = block_out_channels[0] * 4
  
      if self.num_attention_heads is not None:
          raise ValueError(
              "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."
          )
  
      # If `num_attention_heads` is not defined (which is the case for most models)
      # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
      # The reason for this behavior is to correct for incorrectly named variables that were introduced
      # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
      # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
      # which is why we correct for the naming here.
      num_attention_heads = (
          self.num_attention_heads or self.attention_head_dim
      )
  
      # input
      self.conv_in = nn.Conv(
          block_out_channels[0],
          kernel_size=(3, 3),
          strides=(1, 1),
          padding=((1, 1), (1, 1)),
          dtype=self.dtype,
      )
  
      # time
      self.time_proj = FlaxTimesteps(
          block_out_channels[0],
          flip_sin_to_cos=self.flip_sin_to_cos,
          freq_shift=self.config.freq_shift,
      )
      self.time_embedding = FlaxTimestepEmbedding(
          time_embed_dim, dtype=self.dtype
      )
  
      only_cross_attention = self.only_cross_attention
      if isinstance(only_cross_attention, bool):
          only_cross_attention = (only_cross_attention,) * len(
              self.down_block_types
          )
  
      if isinstance(num_attention_heads, int):
          num_attention_heads = (num_attention_heads,) * len(
              self.down_block_types
          )
  
      # transformer layers per block
      transformer_layers_per_block = self.transformer_layers_per_block
      if isinstance(transformer_layers_per_block, int):
          transformer_layers_per_block = [
              transformer_layers_per_block
          ] * len(self.down_block_types)
  
      # addition embed types
      if self.addition_embed_type is None:
          self.add_embedding = None
      elif self.addition_embed_type == "text_time":
          if self.addition_time_embed_dim is None:
              raise ValueError(
                  f"addition_embed_type {self.addition_embed_type} requires `addition_time_embed_dim` to not be None"
              )
          self.add_time_proj = FlaxTimesteps(
              self.addition_time_embed_dim,
              self.flip_sin_to_cos,
              self.freq_shift,
          )
          self.add_embedding = FlaxTimestepEmbedding(
              time_embed_dim, dtype=self.dtype
          )
      else:
          raise ValueError(
              f"addition_embed_type: {self.addition_embed_type} must be None or `text_time`."
          )
  
      # down
      down_blocks = []
      output_channel = block_out_channels[0]
      for i, down_block_type in enumerate(self.down_block_types):
          input_channel = output_channel
          output_channel = block_out_channels[i]
          is_final_block = i == len(block_out_channels) - 1
  
          if down_block_type == "CrossAttnDownBlock2D":
              down_block = FlaxCrossAttnDownBlock2D(
                  in_channels=input_channel,
                  out_channels=output_channel,
                  dropout=self.dropout,
                  num_layers=self.layers_per_block,
                  transformer_layers_per_block=transformer_layers_per_block[
                      i
                  ],
                  num_attention_heads=num_attention_heads[i],
                  add_downsample=not is_final_block,
                  use_linear_projection=self.use_linear_projection,
                  only_cross_attention=only_cross_attention[i],
                  use_memory_efficient_attention=self.use_memory_efficient_attention,
                  split_head_dim=self.split_head_dim,
                  dtype=self.dtype,
              )
          else:
              down_block = FlaxDownBlock2D(
                  in_channels=input_channel,
                  out_channels=output_channel,
                  dropout=self.dropout,
                  num_layers=self.layers_per_block,
                  add_downsample=not is_final_block,
                  dtype=self.dtype,
              )
  
          down_blocks.append(down_block)
      self.down_blocks = down_blocks
  
      # mid
      self.mid_block = FlaxUNetMidBlock2DCrossAttn(
          in_channels=block_out_channels[-1],
          dropout=self.dropout,
          num_attention_heads=num_attention_heads[-1],
          transformer_layers_per_block=transformer_layers_per_block[-1],
          use_linear_projection=self.use_linear_projection,
          use_memory_efficient_attention=self.use_memory_efficient_attention,
          split_head_dim=self.split_head_dim,
          dtype=self.dtype,
      )
  
      # up
      up_blocks = []
      reversed_block_out_channels = list(reversed(block_out_channels))
      reversed_num_attention_heads = list(reversed(num_attention_heads))
      only_cross_attention = list(reversed(only_cross_attention))
      output_channel = reversed_block_out_channels[0]
      reversed_transformer_layers_per_block = list(
          reversed(transformer_layers_per_block)
      )
      for i, up_block_type in enumerate(self.up_block_types):
          prev_output_channel = output_channel
          output_channel = reversed_block_out_channels[i]
          input_channel = reversed_block_out_channels[
              min(i + 1, len(block_out_channels) - 1)
          ]
  
          is_final_block = i == len(block_out_channels) - 1
  
          if up_block_type == "CrossAttnUpBlock2D":
              up_block = FlaxCrossAttnUpBlock2D(
                  in_channels=input_channel,
                  out_channels=output_channel,
                  prev_output_channel=prev_output_channel,
                  num_layers=self.layers_per_block + 1,
                  transformer_layers_per_block=reversed_transformer_layers_per_block[
                      i
                  ],
                  num_attention_heads=reversed_num_attention_heads[i],
                  add_upsample=not is_final_block,
                  dropout=self.dropout,
                  use_linear_projection=self.use_linear_projection,
                  only_cross_attention=only_cross_attention[i],
                  use_memory_efficient_attention=self.use_memory_efficient_attention,
                  split_head_dim=self.split_head_dim,
                  dtype=self.dtype,
              )
          else:
              up_block = FlaxUpBlock2D(
                  in_channels=input_channel,
                  out_channels=output_channel,
                  prev_output_channel=prev_output_channel,
                  num_layers=self.layers_per_block + 1,
                  add_upsample=not is_final_block,
                  dropout=self.dropout,
                  dtype=self.dtype,
              )
  
          up_blocks.append(up_block)
          prev_output_channel = output_channel
      self.up_blocks = up_blocks
  
      # out
      self.conv_norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-5)
      self.conv_out = nn.Conv(
          self.out_channels,
          kernel_size=(3, 3),
          strides=(1, 1),
          padding=((1, 1), (1, 1)),
          dtype=self.dtype,
      )
  
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 2 for calling the FlaxUNet2DConditionModel obj's underlying member method code with example usage: obj.setup() without return types