declare-lab/tango2
92
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from typing import Tuple, Union15 16import flax17import flax.linen as nn18import jax19import jax.numpy as jnp20from flax.core.frozen_dict import FrozenDict21 22from ..configuration_utils import ConfigMixin, flax_register_to_config23from ..utils import BaseOutput24from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps25from .modeling_flax_utils import FlaxModelMixin26from .unet_2d_blocks_flax import (27 FlaxCrossAttnDownBlock2D,28 FlaxCrossAttnUpBlock2D,29 FlaxDownBlock2D,30 FlaxUNetMidBlock2DCrossAttn,31 FlaxUpBlock2D,32)33 34 35@flax.struct.dataclass36class FlaxUNet2DConditionOutput(BaseOutput):37 """38 Args:39 sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`):40 Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.41 """42 43 sample: jnp.ndarray44 45 46@flax_register_to_config47class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):48 r"""49 FlaxUNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a50 timestep and returns sample shaped output.51 52 This model inherits from [`FlaxModelMixin`]. Check the superclass documentation for the generic methods the library53 implements for all the models (such as downloading or saving, etc.)54 55 Also, this model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module)56 subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to57 general usage and behavior.58 59 Finally, this model supports inherent JAX features such as:60 - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)61 - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)62 - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)63 - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)64 65 Parameters:66 sample_size (`int`, *optional*):67 The size of the input sample.68 in_channels (`int`, *optional*, defaults to 4):69 The number of channels in the input sample.70 out_channels (`int`, *optional*, defaults to 4):71 The number of channels in the output.72 down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):73 The tuple of downsample blocks to use. The corresponding class names will be: "FlaxCrossAttnDownBlock2D",74 "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D"75 up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`):76 The tuple of upsample blocks to use. The corresponding class names will be: "FlaxUpBlock2D",77 "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D"78 block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):79 The tuple of output channels for each block.80 layers_per_block (`int`, *optional*, defaults to 2):81 The number of layers per block.82 attention_head_dim (`int` or `Tuple[int]`, *optional*, defaults to 8):83 The dimension of the attention heads.84 cross_attention_dim (`int`, *optional*, defaults to 768):85 The dimension of the cross attention features.86 dropout (`float`, *optional*, defaults to 0):87 Dropout probability for down, up and bottleneck blocks.88 flip_sin_to_cos (`bool`, *optional*, defaults to `True`):89 Whether to flip the sin to cos in the time embedding.90 freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.91 92 """93 94 sample_size: int = 3295 in_channels: int = 496 out_channels: int = 497 down_block_types: Tuple[str] = (98 "CrossAttnDownBlock2D",99 "CrossAttnDownBlock2D",100 "CrossAttnDownBlock2D",101 "DownBlock2D",102 )103 up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")104 only_cross_attention: Union[bool, Tuple[bool]] = False105 block_out_channels: Tuple[int] = (320, 640, 1280, 1280)106 layers_per_block: int = 2107 attention_head_dim: Union[int, Tuple[int]] = 8108 cross_attention_dim: int = 1280109 dropout: float = 0.0110 use_linear_projection: bool = False111 dtype: jnp.dtype = jnp.float32112 flip_sin_to_cos: bool = True113 freq_shift: int = 0114 115 def init_weights(self, rng: jax.random.KeyArray) -> FrozenDict:116 # init input tensors117 sample_shape = (1, self.in_channels, self.sample_size, self.sample_size)118 sample = jnp.zeros(sample_shape, dtype=jnp.float32)119 timesteps = jnp.ones((1,), dtype=jnp.int32)120 encoder_hidden_states = jnp.zeros((1, 1, self.cross_attention_dim), dtype=jnp.float32)121 122 params_rng, dropout_rng = jax.random.split(rng)123 rngs = {"params": params_rng, "dropout": dropout_rng}124 125 return self.init(rngs, sample, timesteps, encoder_hidden_states)["params"]126 127 def setup(self):128 block_out_channels = self.block_out_channels129 time_embed_dim = block_out_channels[0] * 4130 131 # input132 self.conv_in = nn.Conv(133 block_out_channels[0],134 kernel_size=(3, 3),135 strides=(1, 1),136 padding=((1, 1), (1, 1)),137 dtype=self.dtype,138 )139 140 # time141 self.time_proj = FlaxTimesteps(142 block_out_channels[0], flip_sin_to_cos=self.flip_sin_to_cos, freq_shift=self.config.freq_shift143 )144 self.time_embedding = FlaxTimestepEmbedding(time_embed_dim, dtype=self.dtype)145 146 only_cross_attention = self.only_cross_attention147 if isinstance(only_cross_attention, bool):148 only_cross_attention = (only_cross_attention,) * len(self.down_block_types)149 150 attention_head_dim = self.attention_head_dim151 if isinstance(attention_head_dim, int):152 attention_head_dim = (attention_head_dim,) * len(self.down_block_types)153 154 # down155 down_blocks = []156 output_channel = block_out_channels[0]157 for i, down_block_type in enumerate(self.down_block_types):158 input_channel = output_channel159 output_channel = block_out_channels[i]160 is_final_block = i == len(block_out_channels) - 1161 162 if down_block_type == "CrossAttnDownBlock2D":163 down_block = FlaxCrossAttnDownBlock2D(164 in_channels=input_channel,165 out_channels=output_channel,166 dropout=self.dropout,167 num_layers=self.layers_per_block,168 attn_num_head_channels=attention_head_dim[i],169 add_downsample=not is_final_block,170 use_linear_projection=self.use_linear_projection,171 only_cross_attention=only_cross_attention[i],172 dtype=self.dtype,173 )174 else:175 down_block = FlaxDownBlock2D(176 in_channels=input_channel,177 out_channels=output_channel,178 dropout=self.dropout,179 num_layers=self.layers_per_block,180 add_downsample=not is_final_block,181 dtype=self.dtype,182 )183 184 down_blocks.append(down_block)185 self.down_blocks = down_blocks186 187 # mid188 self.mid_block = FlaxUNetMidBlock2DCrossAttn(189 in_channels=block_out_channels[-1],190 dropout=self.dropout,191 attn_num_head_channels=attention_head_dim[-1],192 use_linear_projection=self.use_linear_projection,193 dtype=self.dtype,194 )195 196 # up197 up_blocks = []198 reversed_block_out_channels = list(reversed(block_out_channels))199 reversed_attention_head_dim = list(reversed(attention_head_dim))200 only_cross_attention = list(reversed(only_cross_attention))201 output_channel = reversed_block_out_channels[0]202 for i, up_block_type in enumerate(self.up_block_types):203 prev_output_channel = output_channel204 output_channel = reversed_block_out_channels[i]205 input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]206 207 is_final_block = i == len(block_out_channels) - 1208 209 if up_block_type == "CrossAttnUpBlock2D":210 up_block = FlaxCrossAttnUpBlock2D(211 in_channels=input_channel,212 out_channels=output_channel,213 prev_output_channel=prev_output_channel,214 num_layers=self.layers_per_block + 1,215 attn_num_head_channels=reversed_attention_head_dim[i],216 add_upsample=not is_final_block,217 dropout=self.dropout,218 use_linear_projection=self.use_linear_projection,219 only_cross_attention=only_cross_attention[i],220 dtype=self.dtype,221 )222 else:223 up_block = FlaxUpBlock2D(224 in_channels=input_channel,225 out_channels=output_channel,226 prev_output_channel=prev_output_channel,227 num_layers=self.layers_per_block + 1,228 add_upsample=not is_final_block,229 dropout=self.dropout,230 dtype=self.dtype,231 )232 233 up_blocks.append(up_block)234 prev_output_channel = output_channel235 self.up_blocks = up_blocks236 237 # out238 self.conv_norm_out = nn.GroupNorm(num_groups=32, epsilon=1e-5)239 self.conv_out = nn.Conv(240 self.out_channels,241 kernel_size=(3, 3),242 strides=(1, 1),243 padding=((1, 1), (1, 1)),244 dtype=self.dtype,245 )246 247 def __call__(248 self,249 sample,250 timesteps,251 encoder_hidden_states,252 down_block_additional_residuals=None,253 mid_block_additional_residual=None,254 return_dict: bool = True,255 train: bool = False,256 ) -> Union[FlaxUNet2DConditionOutput, Tuple]:257 r"""258 Args:259 sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor260 timestep (`jnp.ndarray` or `float` or `int`): timesteps261 encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states262 return_dict (`bool`, *optional*, defaults to `True`):263 Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a264 plain tuple.265 train (`bool`, *optional*, defaults to `False`):266 Use deterministic functions and disable dropout when not training.267 268 Returns:269 [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] or `tuple`:270 [`~models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`.271 When returning a tuple, the first element is the sample tensor.272 """273 # 1. time274 if not isinstance(timesteps, jnp.ndarray):275 timesteps = jnp.array([timesteps], dtype=jnp.int32)276 elif isinstance(timesteps, jnp.ndarray) and len(timesteps.shape) == 0:277 timesteps = timesteps.astype(dtype=jnp.float32)278 timesteps = jnp.expand_dims(timesteps, 0)279 280 t_emb = self.time_proj(timesteps)281 t_emb = self.time_embedding(t_emb)282 283 # 2. pre-process284 sample = jnp.transpose(sample, (0, 2, 3, 1))285 sample = self.conv_in(sample)286 287 # 3. down288 down_block_res_samples = (sample,)289 for down_block in self.down_blocks:290 if isinstance(down_block, FlaxCrossAttnDownBlock2D):291 sample, res_samples = down_block(sample, t_emb, encoder_hidden_states, deterministic=not train)292 else:293 sample, res_samples = down_block(sample, t_emb, deterministic=not train)294 down_block_res_samples += res_samples295 296 if down_block_additional_residuals is not None:297 new_down_block_res_samples = ()298 299 for down_block_res_sample, down_block_additional_residual in zip(300 down_block_res_samples, down_block_additional_residuals301 ):302 down_block_res_sample += down_block_additional_residual303 new_down_block_res_samples += (down_block_res_sample,)304 305 down_block_res_samples = new_down_block_res_samples306 307 # 4. mid308 sample = self.mid_block(sample, t_emb, encoder_hidden_states, deterministic=not train)309 310 if mid_block_additional_residual is not None:311 sample += mid_block_additional_residual312 313 # 5. up314 for up_block in self.up_blocks:315 res_samples = down_block_res_samples[-(self.layers_per_block + 1) :]316 down_block_res_samples = down_block_res_samples[: -(self.layers_per_block + 1)]317 if isinstance(up_block, FlaxCrossAttnUpBlock2D):318 sample = up_block(319 sample,320 temb=t_emb,321 encoder_hidden_states=encoder_hidden_states,322 res_hidden_states_tuple=res_samples,323 deterministic=not train,324 )325 else:326 sample = up_block(sample, temb=t_emb, res_hidden_states_tuple=res_samples, deterministic=not train)327 328 # 6. post-process329 sample = self.conv_norm_out(sample)330 sample = nn.silu(sample)331 sample = self.conv_out(sample)332 sample = jnp.transpose(sample, (0, 3, 1, 2))333 334 if not return_dict:335 return (sample,)336 337 return FlaxUNet2DConditionOutput(sample=sample)338