ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

ai_edge_torch/generative/layers/unet/blocks_2d.py

@@ -13,10 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import List, Optional, Tuple
-
-import torch
-from torch import nn
+from typing import List, Optional, Tuple, Union
 
 from ai_edge_torch.generative.layers.attention import CrossAttention
 from ai_edge_torch.generative.layers.attention import SelfAttention
@@ -24,6 +21,8 @@ import ai_edge_torch.generative.layers.builder as layers_builder
 import ai_edge_torch.generative.layers.model_config as layers_cfg
 import ai_edge_torch.generative.layers.unet.builder as unet_builder
 import ai_edge_torch.generative.layers.unet.model_config as unet_cfg
+import torch
+from torch import nn
 
 
 class ResidualBlock2D(nn.Module):
@@ -41,26 +40,38 @@ class ResidualBlock2D(nn.Module):
         config.in_channels, config.normalization_config
     )
     self.conv_1 = nn.Conv2d(
-        config.in_channels, config.out_channels, kernel_size=3, stride=1, padding=1
+        config.in_channels,
+        config.hidden_channels,
+        kernel_size=3,
+        stride=1,
+        padding=1,
     )
     if config.time_embedding_channels is not None:
       self.time_emb_proj = nn.Linear(
-          config.time_embedding_channels, config.out_channels
+          config.time_embedding_channels, config.hidden_channels
       )
     else:
       self.time_emb_proj = None
     self.norm_2 = layers_builder.build_norm(
-        config.out_channels, config.normalization_config
+        config.hidden_channels, config.normalization_config
     )
     self.conv_2 = nn.Conv2d(
-        config.out_channels, config.out_channels, kernel_size=3, stride=1, padding=1
+        config.hidden_channels,
+        config.out_channels,
+        kernel_size=3,
+        stride=1,
+        padding=1,
     )
     self.act_fn = layers_builder.get_activation(config.activation_config)
     if config.in_channels == config.out_channels:
       self.residual_layer = nn.Identity()
     else:
       self.residual_layer = nn.Conv2d(
-          config.in_channels, config.out_channels, kernel_size=1, stride=1, padding=0
+          config.in_channels,
+          config.out_channels,
+          kernel_size=1,
+          stride=1,
+          padding=0,
       )
 
   def forward(
@@ -94,7 +105,6 @@ class AttentionBlock2D(nn.Module):
   """2D self attention block
 
   x = SelfAttention(Norm(input_tensor)) + x
-
   """
 
   def __init__(self, config: unet_cfg.AttentionBlock2DConfig):
@@ -105,12 +115,13 @@ class AttentionBlock2D(nn.Module):
     """
     super().__init__()
     self.config = config
-    self.norm = layers_builder.build_norm(config.dim, config.normalization_config)
+    self.norm = layers_builder.build_norm(
+        config.dim, config.normalization_config
+    )
     self.attention = SelfAttention(
         config.attention_batch_size,
         config.dim,
         config.attention_config,
-        0,
         enable_hlfb=config.enable_hlfb,
     )
 
@@ -125,20 +136,23 @@ class AttentionBlock2D(nn.Module):
     """
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
-    if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+    if (
+        self.config.normalization_config.type
+        == layers_cfg.NormalizationType.GROUP_NORM
+    ):
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
-      x = input_tensor.view(B, C, H * W)
-      x = x.transpose(-1, -2)
+      x = torch.permute(input_tensor, (0, 2, 3, 1))
       x = self.norm(x)
+      x = x.view(B, H * W, C)
     x = x.contiguous()  # Prevent BATCH_MATMUL op in converted tflite.
     x = self.attention(x)
-    x = x.transpose(-1, -2)
-    x = x.view(B, C, H, W)
+    x = x.view(B, H, W, C)
+    residual = torch.permute(residual, (0, 2, 3, 1))
     x = x + residual
+    x = torch.permute(x, (0, 3, 1, 2))
     return x
 
 
@@ -146,24 +160,27 @@ class CrossAttentionBlock2D(nn.Module):
   """2D cross attention block
 
   x = CrossAttention(Norm(input_tensor), context) + x
-
   """
 
   def __init__(self, config: unet_cfg.CrossAttentionBlock2DConfig):
     """Initialize an instance of the AttentionBlock2D.
 
     Args:
-      config (unet_cfg.CrossAttentionBlock2DConfig): the configuration of this block.
+      config (unet_cfg.CrossAttentionBlock2DConfig): the configuration of this
+        block.
     """
     super().__init__()
     self.config = config
-    self.norm = layers_builder.build_norm(config.query_dim, config.normalization_config)
+    self.norm = layers_builder.build_norm(
+        config.query_dim, config.normalization_config
+    )
     self.attention = CrossAttention(
         config.attention_batch_size,
         config.query_dim,
         config.cross_dim,
+        config.hidden_dim,
+        config.output_dim,
         config.attention_config,
-        0,
         enable_hlfb=config.enable_hlfb,
     )
 
@@ -174,26 +191,30 @@ class CrossAttentionBlock2D(nn.Module):
 
     Args:
       input_tensor (torch.Tensor): the input tensor.
-      context_tensor (torch.Tensor): the context tensor to apply cross attention on.
+      context_tensor (torch.Tensor): the context tensor to apply cross attention
+        on.
 
     Returns:
       output activation tensor after cross attention.
     """
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
-    if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+    if (
+        self.config.normalization_config.type
+        == layers_cfg.NormalizationType.GROUP_NORM
+    ):
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
-      x = input_tensor.view(B, C, H * W)
-      x = x.transpose(-1, -2)
+      x = torch.permute(input_tensor, (0, 2, 3, 1))
       x = self.norm(x)
+      x = x.view(B, H * W, C)
     x = self.attention(x, context_tensor)
-    x = x.transpose(-1, -2)
-    x = x.view(B, C, H, W)
+    x = x.view(B, H, W, C)
+    residual = torch.permute(residual, (0, 2, 3, 1))
     x = x + residual
+    x = torch.permute(x, (0, 3, 1, 2))
     return x
 
 
@@ -201,7 +222,6 @@ class FeedForwardBlock2D(nn.Module):
   """2D feed forward block
 
   x = w2(Activation(w1(Norm(x)))) + x
-
   """
 
   def __init__(
@@ -211,7 +231,9 @@ class FeedForwardBlock2D(nn.Module):
     super().__init__()
     self.config = config
     self.act = layers_builder.get_activation(config.activation_config)
-    self.norm = layers_builder.build_norm(config.dim, config.normalization_config)
+    self.norm = layers_builder.build_norm(
+        config.dim, config.normalization_config
+    )
     if config.activation_config.type == layers_cfg.ActivationType.GE_GLU:
       self.w1 = nn.Identity()
       self.w2 = nn.Linear(config.hidden_dim, config.dim)
@@ -222,23 +244,25 @@ class FeedForwardBlock2D(nn.Module):
   def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
     residual = input_tensor
     B, C, H, W = input_tensor.shape
-    x = input_tensor
-    if self.config.normalization_config.type == layers_cfg.NormalizationType.GROUP_NORM:
-      x = self.norm(x)
-      x = input_tensor.view(B, C, H * W)
+    if (
+        self.config.normalization_config.type
+        == layers_cfg.NormalizationType.GROUP_NORM
+    ):
+      x = self.norm(input_tensor)
+      x = x.view(B, C, H * W)
       x = x.transpose(-1, -2)
     else:
-      x = input_tensor.view(B, C, H * W)
-      x = x.transpose(-1, -2)
+      x = torch.permute(input_tensor, (0, 2, 3, 1))
       x = self.norm(x)
+      x = x.view(B, H * W, C)
     x = self.w1(x)
     x = self.act(x)
     x = self.w2(x)
-
-    x = x.transpose(-1, -2)  # (B, C, HW)
-    x = x.view((B, C, H, W))
-
-    return x + residual
+    x = x.view(B, H, W, C)
+    residual = torch.permute(residual, (0, 2, 3, 1))
+    x = x + residual
+    x = torch.permute(x, (0, 3, 1, 2))
+    return x
 
 
 class TransformerBlock2D(nn.Module):
@@ -268,15 +292,14 @@ class TransformerBlock2D(nn.Module):
   └─────────┬─────────┘
             ▼
       hidden_states
-
-
   """
 
   def __init__(self, config: unet_cfg.TransformerBlock2DConfig):
     """Initialize an instance of the TransformerBlock2D.
 
     Args:
-      config (unet_cfg.TransformerBlock2Dconfig): the configuration of this block.
+      config (unet_cfg.TransformerBlock2Dconfig): the configuration of this
+        block.
     """
     super().__init__()
     self.config = config
@@ -290,7 +313,9 @@ class TransformerBlock2D(nn.Module):
         padding=0,
     )
     self.self_attention = AttentionBlock2D(config.attention_block_config)
-    self.cross_attention = CrossAttentionBlock2D(config.cross_attention_block_config)
+    self.cross_attention = CrossAttentionBlock2D(
+        config.cross_attention_block_config
+    )
     self.feed_forward = FeedForwardBlock2D(config.feed_forward_block_config)
     self.conv_out = nn.Conv2d(
         config.attention_block_config.dim,
@@ -304,7 +329,8 @@ class TransformerBlock2D(nn.Module):
 
     Args:
       input_tensor (torch.Tensor): the input tensor.
-      context_tensor (torch.Tensor): the context tensor to apply cross attention on.
+      context_tensor (torch.Tensor): the context tensor to apply cross attention
+        on.
 
     Returns:
       output activation tensor after transformer block.
@@ -352,7 +378,8 @@ class DownEncoderBlock2D(nn.Module):
     """Initialize an instance of the DownEncoderBlock2D.
 
     Args:
-      config (unet_cfg.DownEncoderBlock2DConfig): the configuration of this block.
+      config (unet_cfg.DownEncoderBlock2DConfig): the configuration of this
+        block.
     """
     super().__init__()
     self.config = config
@@ -364,6 +391,7 @@ class DownEncoderBlock2D(nn.Module):
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=input_channels,
+                  hidden_channels=config.out_channels,
                   out_channels=config.out_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
@@ -374,7 +402,9 @@ class DownEncoderBlock2D(nn.Module):
       if config.transformer_block_config:
         transformers.append(TransformerBlock2D(config.transformer_block_config))
     self.resnets = nn.ModuleList(resnets)
-    self.transformers = nn.ModuleList(transformers) if len(transformers) > 0 else None
+    self.transformers = (
+        nn.ModuleList(transformers) if len(transformers) > 0 else None
+    )
     if config.add_downsample:
       self.downsampler = unet_builder.build_downsampling(config.sampling_config)
     else:
@@ -386,15 +416,18 @@ class DownEncoderBlock2D(nn.Module):
       time_emb: Optional[torch.Tensor] = None,
       context_tensor: Optional[torch.Tensor] = None,
       output_hidden_states: bool = False,
-  ) -> torch.Tensor | Tuple[torch.Tensor, List[torch.Tensor]]:
+  ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
     """Forward function of the DownEncoderBlock2D.
 
     Args:
       input_tensor (torch.Tensor): the input tensor.
-      time_emb (torch.Tensor): optional time embedding tensor, if the block is configured to accept
-        time embedding.
-      context_tensor (torch.Tensor): optional context tensor, if the block if configured to use transofrmer block.
-      output_hidden_states (bool): whether to output hidden states, usually for skip connections.
+      time_emb (torch.Tensor): optional time embedding tensor, if the block is
+        configured to accept time embedding.
+      context_tensor (torch.Tensor): optional context tensor, if the block if
+        configured to use transofrmer block.
+      output_hidden_states (bool): whether to output hidden states, usually for
+        skip connections.
+
     Returns:
       output hidden_states tensor after DownEncoderBlock2D.
     """
@@ -460,6 +493,7 @@ class UpDecoderBlock2D(nn.Module):
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=input_channels,
+                  hidden_channels=config.out_channels,
                   out_channels=config.out_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
@@ -470,12 +504,18 @@ class UpDecoderBlock2D(nn.Module):
       if config.transformer_block_config:
         transformers.append(TransformerBlock2D(config.transformer_block_config))
     self.resnets = nn.ModuleList(resnets)
-    self.transformers = nn.ModuleList(transformers) if len(transformers) > 0 else None
+    self.transformers = (
+        nn.ModuleList(transformers) if len(transformers) > 0 else None
+    )
     if config.add_upsample:
       self.upsampler = unet_builder.build_upsampling(config.sampling_config)
       if config.upsample_conv:
         self.upsample_conv = nn.Conv2d(
-            config.out_channels, config.out_channels, kernel_size=3, stride=1, padding=1
+            config.out_channels,
+            config.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
         )
     else:
       self.upsampler = None
@@ -490,9 +530,10 @@ class UpDecoderBlock2D(nn.Module):
 
     Args:
       input_tensor (torch.Tensor): the input tensor.
-      time_emb (torch.Tensor): optional time embedding tensor, if the block is configured to accept
-        time embedding.
-      context_tensor (torch.Tensor): optional context tensor, if the block if configured to use transofrmer block.
+      time_emb (torch.Tensor): optional time embedding tensor, if the block is
+        configured to accept time embedding.
+      context_tensor (torch.Tensor): optional context tensor, if the block if
+        configured to use transofrmer block.
 
     Returns:
       output hidden_states tensor after UpDecoderBlock2D.
@@ -543,7 +584,8 @@ class SkipUpDecoderBlock2D(nn.Module):
     """Initialize an instance of the SkipUpDecoderBlock2D.
 
     Args:
-      config (unet_cfg.SkipUpDecoderBlock2DConfig): the configuration of this block.
+      config (unet_cfg.SkipUpDecoderBlock2DConfig): the configuration of this
+        block.
     """
     super().__init__()
     self.config = config
@@ -551,13 +593,18 @@ class SkipUpDecoderBlock2D(nn.Module):
     transformers = []
     for i in range(config.num_layers):
       res_skip_channels = (
-          config.in_channels if (i == config.num_layers - 1) else config.out_channels
+          config.in_channels
+          if (i == config.num_layers - 1)
+          else config.out_channels
+      )
+      resnet_in_channels = (
+          config.prev_out_channels if i == 0 else config.out_channels
       )
-      resnet_in_channels = config.prev_out_channels if i == 0 else config.out_channels
       resnets.append(
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=resnet_in_channels + res_skip_channels,
+                  hidden_channels=config.out_channels,
                   out_channels=config.out_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
@@ -568,12 +615,18 @@ class SkipUpDecoderBlock2D(nn.Module):
       if config.transformer_block_config:
         transformers.append(TransformerBlock2D(config.transformer_block_config))
     self.resnets = nn.ModuleList(resnets)
-    self.transformers = nn.ModuleList(transformers) if len(transformers) > 0 else None
+    self.transformers = (
+        nn.ModuleList(transformers) if len(transformers) > 0 else None
+    )
     if config.add_upsample:
       self.upsampler = unet_builder.build_upsampling(config.sampling_config)
       if config.upsample_conv:
         self.upsample_conv = nn.Conv2d(
-            config.out_channels, config.out_channels, kernel_size=3, stride=1, padding=1
+            config.out_channels,
+            config.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
         )
     else:
       self.upsampler = None
@@ -589,10 +642,12 @@ class SkipUpDecoderBlock2D(nn.Module):
 
     Args:
       input_tensor (torch.Tensor): the input tensor.
-      skip_connection_tensors (List[torch.Tensor]): the skip connection tensors from encoder blocks.
-      time_emb (torch.Tensor): optional time embedding tensor, if the block is configured to accept
-        time embedding.
-      context_tensor (torch.Tensor): optional context tensor, if the block if configured to use transofrmer block.
+      skip_connection_tensors (List[torch.Tensor]): the skip connection tensors
+        from encoder blocks.
+      time_emb (torch.Tensor): optional time embedding tensor, if the block is
+        configured to accept time embedding.
+      context_tensor (torch.Tensor): optional context tensor, if the block if
+        configured to use transofrmer block.
 
     Returns:
       output hidden_states tensor after SkipUpDecoderBlock2D.
@@ -654,6 +709,7 @@ class MidBlock2D(nn.Module):
         ResidualBlock2D(
             unet_cfg.ResidualBlock2DConfig(
                 in_channels=config.in_channels,
+                hidden_channels=config.in_channels,
                 out_channels=config.in_channels,
                 time_embedding_channels=config.time_embedding_channels,
                 normalization_config=config.normalization_config,
@@ -672,6 +728,7 @@ class MidBlock2D(nn.Module):
           ResidualBlock2D(
               unet_cfg.ResidualBlock2DConfig(
                   in_channels=config.in_channels,
+                  hidden_channels=config.in_channels,
                   out_channels=config.in_channels,
                   time_embedding_channels=config.time_embedding_channels,
                   normalization_config=config.normalization_config,
@@ -681,7 +738,9 @@ class MidBlock2D(nn.Module):
       )
     self.resnets = nn.ModuleList(resnets)
     self.attentions = nn.ModuleList(attentions) if len(attentions) > 0 else None
-    self.transformers = nn.ModuleList(transformers) if len(transformers) > 0 else None
+    self.transformers = (
+        nn.ModuleList(transformers) if len(transformers) > 0 else None
+    )
 
   def forward(
       self,
@@ -693,10 +752,10 @@ class MidBlock2D(nn.Module):
 
     Args:
       input_tensor (torch.Tensor): the input tensor.
-      time_emb (torch.Tensor): optional time embedding tensor, if the block is configured to accept
-        time embedding.
-      context_tensor (torch.Tensor): optional context tensor, if the block if configured to use
-        transofrmer block.
+      time_emb (torch.Tensor): optional time embedding tensor, if the block is
+        configured to accept time embedding.
+      context_tensor (torch.Tensor): optional context tensor, if the block if
+        configured to use transofrmer block.
 
     Returns:
       output hidden_states tensor after MidBlock2D.

ai_edge_torch/generative/layers/unet/builder.py

@@ -14,9 +14,8 @@
 # ==============================================================================
 # Builder utils for individual components.
 
-from torch import nn
-
 import ai_edge_torch.generative.layers.unet.model_config as unet_config
+from torch import nn
 
 
 def build_upsampling(config: unet_config.UpSamplingConfig):
@@ -30,10 +29,14 @@ def build_upsampling(config: unet_config.UpSamplingConfig):
 
 def build_downsampling(config: unet_config.DownSamplingConfig):
   if config.mode == unet_config.SamplingType.AVERAGE:
-    return nn.AvgPool2d(config.kernel_size, config.stride, padding=config.padding)
+    return nn.AvgPool2d(
+        config.kernel_size, config.stride, padding=config.padding
+    )
   elif config.mode == unet_config.SamplingType.CONVOLUTION:
     out_channels = (
-        config.in_channels if config.out_channels is None else config.out_channels
+        config.in_channels
+        if config.out_channels is None
+        else config.out_channels
     )
     padding = (0, 1, 0, 1) if config.padding == 0 else config.padding
     return nn.Conv2d(

ai_edge_torch/generative/layers/unet/model_config.py

@@ -14,8 +14,7 @@
 # ==============================================================================
 
 # UNet configuration class.
-from dataclasses import dataclass
-from dataclasses import field
+import dataclasses
 import enum
 from typing import List, Optional
 
@@ -30,13 +29,13 @@ class SamplingType(enum.Enum):
   CONVOLUTION = enum.auto()
 
 
-@dataclass
+@dataclasses.dataclass
 class UpSamplingConfig:
   mode: SamplingType
   scale_factor: float
 
 
-@dataclass
+@dataclasses.dataclass
 class DownSamplingConfig:
   mode: SamplingType
   in_channels: int
@@ -46,9 +45,10 @@ class DownSamplingConfig:
   out_channels: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class ResidualBlock2DConfig:
   in_channels: int
+  hidden_channels: int
   out_channels: int
   normalization_config: layers_cfg.NormalizationConfig
   activation_config: layers_cfg.ActivationConfig
@@ -56,7 +56,7 @@ class ResidualBlock2DConfig:
   time_embedding_channels: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class AttentionBlock2DConfig:
   dim: int
   normalization_config: layers_cfg.NormalizationConfig
@@ -65,17 +65,19 @@ class AttentionBlock2DConfig:
   attention_batch_size: int = 1
 
 
-@dataclass
+@dataclasses.dataclass
 class CrossAttentionBlock2DConfig:
   query_dim: int
   cross_dim: int
+  hidden_dim: int
+  output_dim: int
   normalization_config: layers_cfg.NormalizationConfig
   attention_config: layers_cfg.AttentionConfig
   enable_hlfb: bool = True
   attention_batch_size: int = 1
 
 
-@dataclass
+@dataclasses.dataclass
 class FeedForwardBlock2DConfig:
   dim: int
   hidden_dim: int
@@ -84,7 +86,7 @@ class FeedForwardBlock2DConfig:
   use_bias: bool
 
 
-@dataclass
+@dataclasses.dataclass
 class TransformerBlock2DConfig:
   pre_conv_normalization_config: layers_cfg.NormalizationConfig
   attention_block_config: AttentionBlock2DConfig
@@ -92,7 +94,7 @@ class TransformerBlock2DConfig:
   feed_forward_block_config: FeedForwardBlock2DConfig
 
 
-@dataclass
+@dataclasses.dataclass
 class UpDecoderBlock2DConfig:
   in_channels: int
   out_channels: int
@@ -113,7 +115,7 @@ class UpDecoderBlock2DConfig:
   context_dim: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class SkipUpDecoderBlock2DConfig:
   in_channels: int
   out_channels: int
@@ -136,7 +138,7 @@ class SkipUpDecoderBlock2DConfig:
   context_dim: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class DownEncoderBlock2DConfig:
   in_channels: int
   out_channels: int
@@ -157,7 +159,7 @@ class DownEncoderBlock2DConfig:
   context_dim: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class MidBlock2DConfig:
   in_channels: int
   normalization_config: layers_cfg.NormalizationConfig
@@ -173,7 +175,7 @@ class MidBlock2DConfig:
   context_dim: Optional[int] = None
 
 
-@dataclass
+@dataclasses.dataclass
 class AutoEncoderConfig:
   """Configurations of encoder/decoder in the autoencoder model."""
 
@@ -210,7 +212,7 @@ class AutoEncoderConfig:
   mid_block_config: MidBlock2DConfig
 
 
-@dataclass
+@dataclasses.dataclass
 class DiffusionModelConfig:
   """Configurations of Diffusion model."""
 

ai_edge_torch/generative/quantize/example.py

@@ -13,22 +13,21 @@
 # limitations under the License.
 # ==============================================================================
 
-import numpy as np
-import torch
-
 import ai_edge_torch
-from ai_edge_torch.generative.examples.gemma import gemma
+from ai_edge_torch.generative.examples.gemma import gemma1
 from ai_edge_torch.generative.quantize import quant_recipes
+import numpy as np
+import torch
 
 
 def main():
   # Build a PyTorch model as usual
-  config = gemma.get_fake_model_config_2b_for_test()
-  model = gemma.Gemma(config)
+  config = gemma1.get_fake_model_config()
+  model = gemma1.Gemma(config)
   idx = torch.from_numpy(np.array([[1, 2, 3, 4]]))
-  tokens = torch.full((1, 10), 0, dtype=torch.long, device="cpu")
+  tokens = torch.full((1, 10), 0, dtype=torch.int, device="cpu")
   tokens[0, :4] = idx
-  input_pos = torch.arange(0, 10)
+  input_pos = torch.arange(0, 10, dtype=torch.int)
 
   # Create a quantization recipe to be applied to the model
   quant_config = quant_recipes.full_int8_dynamic_recipe()