ai-edge-torch-nightly 0.2.0.dev20240611__py3-none-any.whl → 0.2.0.dev20240619__py3-none-any.whl

ai_edge_torch/generative/layers/unet/model_config.py

@@ -22,16 +22,28 @@ from typing import List, Optional
 import ai_edge_torch.generative.layers.model_config as layers_cfg
 
 
-@dataclass
+@enum.unique
 class SamplingType(enum.Enum):
   NEAREST = enum.auto()
   BILINEAR = enum.auto()
+  AVERAGE = enum.auto()
+  CONVOLUTION = enum.auto()
 
 
 @dataclass
-class SamplingConfig:
+class UpSamplingConfig:
+  mode: SamplingType
   scale_factor: float
+
+
+@dataclass
+class DownSamplingConfig:
   mode: SamplingType
+  in_channels: int
+  kernel_size: int
+  stride: int
+  padding: int
+  out_channels: Optional[int] = None
 
 
 @dataclass
@@ -46,9 +58,38 @@ class ResidualBlock2DConfig:
 
 @dataclass
 class AttentionBlock2DConfig:
-  dims: int
+  dim: int
+  normalization_config: layers_cfg.NormalizationConfig
+  attention_config: layers_cfg.AttentionConfig
+  enable_hlfb: bool = True
+  attention_batch_size: int = 1
+
+
+@dataclass
+class CrossAttentionBlock2DConfig:
+  query_dim: int
+  cross_dim: int
   normalization_config: layers_cfg.NormalizationConfig
   attention_config: layers_cfg.AttentionConfig
+  enable_hlfb: bool = True
+  attention_batch_size: int = 1
+
+
+@dataclass
+class FeedForwardBlock2DConfig:
+  dim: int
+  hidden_dim: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  use_bias: bool
+
+
+@dataclass
+class TransformerBlock2Dconfig:
+  pre_conv_normalization_config: layers_cfg.NormalizationConfig
+  attention_block_config: AttentionBlock2DConfig
+  cross_attention_block_config: CrossAttentionBlock2DConfig
+  feed_forward_block_config: FeedForwardBlock2DConfig
 
 
 @dataclass
@@ -58,14 +99,62 @@ class UpDecoderBlock2DConfig:
   normalization_config: layers_cfg.NormalizationConfig
   activation_config: layers_cfg.ActivationConfig
   num_layers: int
-  # Optional time embedding channels if the residual blocks take a time embedding context as input
+  # Optional time embedding channels if the residual blocks take a time embedding as input
   time_embedding_channels: Optional[int] = None
   # Whether to add upsample operation after residual blocks
   add_upsample: bool = True
   # Whether to add a conv2d layer after upsample
   upsample_conv: bool = True
   # Optional sampling config if add_upsample is True.
-  sampling_config: Optional[SamplingConfig] = None
+  sampling_config: Optional[UpSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2Dconfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
+
+
+@dataclass
+class SkipUpDecoderBlock2DConfig:
+  in_channels: int
+  out_channels: int
+  # The dimension of output channels of previous connected block
+  prev_out_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  num_layers: int
+  # Optional time embedding channels if the residual blocks take a time embedding as input
+  time_embedding_channels: Optional[int] = None
+  # Whether to add upsample operation after residual blocks
+  add_upsample: bool = True
+  # Whether to add a conv2d layer after upsample
+  upsample_conv: bool = True
+  # Optional sampling config if add_upsample is True.
+  sampling_config: Optional[UpSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2Dconfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
+
+
+@dataclass
+class DownEncoderBlock2DConfig:
+  in_channels: int
+  out_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  num_layers: int
+  # Padding for the downsampling convolution.
+  padding: int = 1
+  # Optional time embedding channels if the residual blocks take a time embedding as input
+  time_embedding_channels: Optional[int] = None
+  # Whether to add downsample operation after residual blocks
+  add_downsample: bool = True
+  # Optional sampling config if add_upsample is True.
+  sampling_config: Optional[DownSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2Dconfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
 
 
 @dataclass
@@ -78,6 +167,10 @@ class MidBlock2DConfig:
   time_embedding_channels: Optional[int] = None
   # Optional config of attention blocks interleaved with residual blocks
   attention_block_config: Optional[AttentionBlock2DConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2Dconfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
 
 
 @dataclass
@@ -115,3 +208,62 @@ class AutoEncoderConfig:
 
   # The configuration of middle blocks, that is, after the last block of encoder and before the first block of decoder.
   mid_block_config: MidBlock2DConfig
+
+
+@dataclass
+class DiffusionModelConfig:
+  """Configurations of Diffusion model."""
+
+  # Number of channels in the input tensor.
+  in_channels: int
+
+  # Number of channels in the output tensor.
+  out_channels: int
+
+  # The output channels of each block.
+  block_out_channels: List[int]
+
+  # The layesr number of each block.
+  layers_per_block: int
+
+  # The padding to use for the downsampling.
+  downsample_padding: int
+
+  # Normalization config used in residual blocks.
+  residual_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation config used in residual blocks
+  residual_activation_type: layers_cfg.ActivationType
+
+  # The batch size used in transformer blocks, for attention layers.
+  transformer_batch_size: int
+
+  # The number of attention heads used in transformer blocks.
+  transformer_num_attention_heads: int
+
+  # The dimension of cross attention used in transformer blocks.
+  transformer_cross_attention_dim: int
+
+  # Normalization config used in prev conv layer of transformer blocks.
+  transformer_pre_conv_norm_config: layers_cfg.NormalizationConfig
+
+  # Normalization config used in transformer blocks.
+  transformer_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation type of feed forward used in transformer blocks.
+  transformer_ff_activation_type: layers_cfg.ActivationType
+
+  # Number of layers in mid block.
+  mid_block_layers: int
+
+  # Dimension of time embedding.
+  time_embedding_dim: int
+
+  # Time embedding dimensions for blocks.
+  time_embedding_blocks_dim: int
+
+  # Normalization config used for final layer
+  final_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation type used in final layer
+  final_activation_type: layers_cfg.ActivationType

ai_edge_torch/generative/test/test_model_conversion.py

@@ -55,6 +55,30 @@ class TestModelConversion(unittest.TestCase):
           )
       )
 
+  def test_toy_model_with_multi_batches(self):
+    config = toy_model_with_kv_cache.get_model_config()
+    config.batch_size = 2
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+    idx, input_pos = torch.tensor([[1], [2]], dtype=torch.long), torch.tensor(
+        [10], dtype=torch.int64
+    )
+
+    edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              pytorch_model,
+              (idx, input_pos),
+              num_valid_inputs=1,
+              atol=1e-5,
+              rtol=1e-5,
+          )
+      )
+
   def test_toy_model_with_kv_cache_with_hlfb(self):
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = True

ai_edge_torch/generative/test/test_quantize.py

@@ -116,6 +116,7 @@ class TestQuantizeConvert(unittest.TestCase):
       ]
   )
   def test_quantize_convert_toy_sizes(self, quant_config, expected_compression):
+    self.skipTest("b/346896669")
     config = toy_model_with_kv_cache.get_model_config()
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
     idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(