ai-edge-torch-nightly 0.2.0.dev20240710__py3-none-any.whl → 0.2.0.dev20240712__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/decoder.py

@@ -22,29 +22,31 @@ import ai_edge_torch.generative.layers.unet.blocks_2d as blocks_2d
 import ai_edge_torch.generative.layers.unet.model_config as unet_cfg
 import ai_edge_torch.generative.utilities.stable_diffusion_loader as stable_diffusion_loader
 
-TENSORS_NAMES = stable_diffusion_loader.AutoEncoderModelLoader.TensorNames(
-    post_quant_conv="0",
-    conv_in="1",
+TENSOR_NAMES = stable_diffusion_loader.AutoEncoderModelLoader.TensorNames(
+    post_quant_conv="first_stage_model.post_quant_conv",
+    conv_in="first_stage_model.decoder.conv_in",
     mid_block_tensor_names=stable_diffusion_loader.MidBlockTensorNames(
         residual_block_tensor_names=[
             stable_diffusion_loader.ResidualBlockTensorNames(
-                norm_1="2.groupnorm_1",
-                norm_2="2.groupnorm_2",
-                conv_1="2.conv_1",
-                conv_2="2.conv_2",
+                norm_1="first_stage_model.decoder.mid.block_1.norm1",
+                norm_2="first_stage_model.decoder.mid.block_1.norm2",
+                conv_1="first_stage_model.decoder.mid.block_1.conv1",
+                conv_2="first_stage_model.decoder.mid.block_1.conv2",
             ),
             stable_diffusion_loader.ResidualBlockTensorNames(
-                norm_1="4.groupnorm_1",
-                norm_2="4.groupnorm_2",
-                conv_1="4.conv_1",
-                conv_2="4.conv_2",
+                norm_1="first_stage_model.decoder.mid.block_2.norm1",
+                norm_2="first_stage_model.decoder.mid.block_2.norm2",
+                conv_1="first_stage_model.decoder.mid.block_2.conv1",
+                conv_2="first_stage_model.decoder.mid.block_2.conv2",
             ),
         ],
         attention_block_tensor_names=[
             stable_diffusion_loader.AttentionBlockTensorNames(
-                norm="3.groupnorm",
-                fused_qkv_proj="3.attention.in_proj",
-                output_proj="3.attention.out_proj",
+                norm="first_stage_model.decoder.mid.attn_1.norm",
+                q_proj="first_stage_model.decoder.mid.attn_1.q",
+                k_proj="first_stage_model.decoder.mid.attn_1.k",
+                v_proj="first_stage_model.decoder.mid.attn_1.v",
+                output_proj="first_stage_model.decoder.mid.attn_1.proj_out",
             )
         ],
     ),
@@ -52,99 +54,99 @@ TENSORS_NAMES = stable_diffusion_loader.AutoEncoderModelLoader.TensorNames(
         stable_diffusion_loader.UpDecoderBlockTensorNames(
             residual_block_tensor_names=[
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="5.groupnorm_1",
-                    norm_2="5.groupnorm_2",
-                    conv_1="5.conv_1",
-                    conv_2="5.conv_2",
+                    norm_1="first_stage_model.decoder.up.3.block.0.norm1",
+                    norm_2="first_stage_model.decoder.up.3.block.0.norm2",
+                    conv_1="first_stage_model.decoder.up.3.block.0.conv1",
+                    conv_2="first_stage_model.decoder.up.3.block.0.conv2",
                 ),
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="6.groupnorm_1",
-                    norm_2="6.groupnorm_2",
-                    conv_1="6.conv_1",
-                    conv_2="6.conv_2",
+                    norm_1="first_stage_model.decoder.up.3.block.1.norm1",
+                    norm_2="first_stage_model.decoder.up.3.block.1.norm2",
+                    conv_1="first_stage_model.decoder.up.3.block.1.conv1",
+                    conv_2="first_stage_model.decoder.up.3.block.1.conv2",
                 ),
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="7.groupnorm_1",
-                    norm_2="7.groupnorm_2",
-                    conv_1="7.conv_1",
-                    conv_2="7.conv_2",
+                    norm_1="first_stage_model.decoder.up.3.block.2.norm1",
+                    norm_2="first_stage_model.decoder.up.3.block.2.norm2",
+                    conv_1="first_stage_model.decoder.up.3.block.2.conv1",
+                    conv_2="first_stage_model.decoder.up.3.block.2.conv2",
                 ),
             ],
-            upsample_conv="9",
+            upsample_conv="first_stage_model.decoder.up.3.upsample.conv",
         ),
         stable_diffusion_loader.UpDecoderBlockTensorNames(
             residual_block_tensor_names=[
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="10.groupnorm_1",
-                    norm_2="10.groupnorm_2",
-                    conv_1="10.conv_1",
-                    conv_2="10.conv_2",
+                    norm_1="first_stage_model.decoder.up.2.block.0.norm1",
+                    norm_2="first_stage_model.decoder.up.2.block.0.norm2",
+                    conv_1="first_stage_model.decoder.up.2.block.0.conv1",
+                    conv_2="first_stage_model.decoder.up.2.block.0.conv2",
                 ),
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="11.groupnorm_1",
-                    norm_2="11.groupnorm_2",
-                    conv_1="11.conv_1",
-                    conv_2="11.conv_2",
+                    norm_1="first_stage_model.decoder.up.2.block.1.norm1",
+                    norm_2="first_stage_model.decoder.up.2.block.1.norm2",
+                    conv_1="first_stage_model.decoder.up.2.block.1.conv1",
+                    conv_2="first_stage_model.decoder.up.2.block.1.conv2",
                 ),
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="12.groupnorm_1",
-                    norm_2="12.groupnorm_2",
-                    conv_1="12.conv_1",
-                    conv_2="12.conv_2",
+                    norm_1="first_stage_model.decoder.up.2.block.2.norm1",
+                    norm_2="first_stage_model.decoder.up.2.block.2.norm2",
+                    conv_1="first_stage_model.decoder.up.2.block.2.conv1",
+                    conv_2="first_stage_model.decoder.up.2.block.2.conv2",
                 ),
             ],
-            upsample_conv="14",
+            upsample_conv="first_stage_model.decoder.up.2.upsample.conv",
         ),
         stable_diffusion_loader.UpDecoderBlockTensorNames(
             residual_block_tensor_names=[
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="15.groupnorm_1",
-                    norm_2="15.groupnorm_2",
-                    conv_1="15.conv_1",
-                    conv_2="15.conv_2",
-                    residual_layer="15.residual_layer",
+                    norm_1="first_stage_model.decoder.up.1.block.0.norm1",
+                    norm_2="first_stage_model.decoder.up.1.block.0.norm2",
+                    conv_1="first_stage_model.decoder.up.1.block.0.conv1",
+                    conv_2="first_stage_model.decoder.up.1.block.0.conv2",
+                    residual_layer="first_stage_model.decoder.up.1.block.0.nin_shortcut",
                 ),
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="16.groupnorm_1",
-                    norm_2="16.groupnorm_2",
-                    conv_1="16.conv_1",
-                    conv_2="16.conv_2",
+                    norm_1="first_stage_model.decoder.up.1.block.1.norm1",
+                    norm_2="first_stage_model.decoder.up.1.block.1.norm2",
+                    conv_1="first_stage_model.decoder.up.1.block.1.conv1",
+                    conv_2="first_stage_model.decoder.up.1.block.1.conv2",
                 ),
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="17.groupnorm_1",
-                    norm_2="17.groupnorm_2",
-                    conv_1="17.conv_1",
-                    conv_2="17.conv_2",
+                    norm_1="first_stage_model.decoder.up.1.block.2.norm1",
+                    norm_2="first_stage_model.decoder.up.1.block.2.norm2",
+                    conv_1="first_stage_model.decoder.up.1.block.2.conv1",
+                    conv_2="first_stage_model.decoder.up.1.block.2.conv2",
                 ),
             ],
-            upsample_conv="19",
+            upsample_conv="first_stage_model.decoder.up.1.upsample.conv",
         ),
         stable_diffusion_loader.UpDecoderBlockTensorNames(
             residual_block_tensor_names=[
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="20.groupnorm_1",
-                    norm_2="20.groupnorm_2",
-                    conv_1="20.conv_1",
-                    conv_2="20.conv_2",
-                    residual_layer="20.residual_layer",
+                    norm_1="first_stage_model.decoder.up.0.block.0.norm1",
+                    norm_2="first_stage_model.decoder.up.0.block.0.norm2",
+                    conv_1="first_stage_model.decoder.up.0.block.0.conv1",
+                    conv_2="first_stage_model.decoder.up.0.block.0.conv2",
+                    residual_layer="first_stage_model.decoder.up.0.block.0.nin_shortcut",
                 ),
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="21.groupnorm_1",
-                    norm_2="21.groupnorm_2",
-                    conv_1="21.conv_1",
-                    conv_2="21.conv_2",
+                    norm_1="first_stage_model.decoder.up.0.block.1.norm1",
+                    norm_2="first_stage_model.decoder.up.0.block.1.norm2",
+                    conv_1="first_stage_model.decoder.up.0.block.1.conv1",
+                    conv_2="first_stage_model.decoder.up.0.block.1.conv2",
                 ),
                 stable_diffusion_loader.ResidualBlockTensorNames(
-                    norm_1="22.groupnorm_1",
-                    norm_2="22.groupnorm_2",
-                    conv_1="22.conv_1",
-                    conv_2="22.conv_2",
+                    norm_1="first_stage_model.decoder.up.0.block.2.norm1",
+                    norm_2="first_stage_model.decoder.up.0.block.2.norm2",
+                    conv_1="first_stage_model.decoder.up.0.block.2.conv1",
+                    conv_2="first_stage_model.decoder.up.0.block.2.conv2",
                 ),
             ],
         ),
     ],
-    final_norm="23",
-    conv_out="25",
+    final_norm="first_stage_model.decoder.norm_out",
+    conv_out="first_stage_model.decoder.conv_out",
 )
 
 
@@ -288,6 +290,7 @@ def get_model_config() -> unet_cfg.AutoEncoderConfig:
           output_proj_use_bias=True,
           enable_kv_cache=False,
           qkv_transpose_before_split=True,
+          qkv_fused_interleaved=False,
           rotary_percentage=0.0,
       ),
   )

ai_edge_torch/generative/examples/stable_diffusion/diffusion.py

@@ -26,12 +26,12 @@ _down_encoder_blocks_tensor_names = [
     stable_diffusion_loader.DownEncoderBlockTensorNames(
         residual_block_tensor_names=[
             stable_diffusion_loader.ResidualBlockTensorNames(
-                norm_1=f"unet.encoders.{i*3+j+1}.0.groupnorm_feature",
-                conv_1=f"unet.encoders.{i*3+j+1}.0.conv_feature",
-                norm_2=f"unet.encoders.{i*3+j+1}.0.groupnorm_merged",
-                conv_2=f"unet.encoders.{i*3+j+1}.0.conv_merged",
-                time_embedding=f"unet.encoders.{i*3+j+1}.0.linear_time",
-                residual_layer=f"unet.encoders.{i*3+j+1}.0.residual_layer"
+                norm_1=f"model.diffusion_model.input_blocks.{i*3+j+1}.0.in_layers.0",
+                conv_1=f"model.diffusion_model.input_blocks.{i*3+j+1}.0.in_layers.2",
+                norm_2=f"model.diffusion_model.input_blocks.{i*3+j+1}.0.out_layers.0",
+                conv_2=f"model.diffusion_model.input_blocks.{i*3+j+1}.0.out_layers.3",
+                time_embedding=f"model.diffusion_model.input_blocks.{i*3+j+1}.0.emb_layers.1",
+                residual_layer=f"model.diffusion_model.input_blocks.{i*3+j+1}.0.skip_connection"
                 if (i * 3 + j + 1) in [4, 7]
                 else None,
             )
@@ -39,32 +39,36 @@ _down_encoder_blocks_tensor_names = [
         ],
         transformer_block_tensor_names=[
             stable_diffusion_loader.TransformerBlockTensorNames(
-                pre_conv_norm=f"unet.encoders.{i*3+j+1}.1.groupnorm",
-                conv_in=f"unet.encoders.{i*3+j+1}.1.conv_input",
-                conv_out=f"unet.encoders.{i*3+j+1}.1.conv_output",
+                pre_conv_norm=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.norm",
+                conv_in=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.proj_in",
+                conv_out=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.proj_out",
                 self_attention=stable_diffusion_loader.AttentionBlockTensorNames(
-                    norm=f"unet.encoders.{i*3+j+1}.1.layernorm_1",
-                    fused_qkv_proj=f"unet.encoders.{i*3+j+1}.1.attention_1.in_proj",
-                    output_proj=f"unet.encoders.{i*3+j+1}.1.attention_1.out_proj",
+                    norm=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.norm1",
+                    q_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn1.to_q",
+                    k_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn1.to_k",
+                    v_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn1.to_v",
+                    output_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn1.to_out.0",
                 ),
                 cross_attention=stable_diffusion_loader.CrossAttentionBlockTensorNames(
-                    norm=f"unet.encoders.{i*3+j+1}.1.layernorm_2",
-                    q_proj=f"unet.encoders.{i*3+j+1}.1.attention_2.q_proj",
-                    k_proj=f"unet.encoders.{i*3+j+1}.1.attention_2.k_proj",
-                    v_proj=f"unet.encoders.{i*3+j+1}.1.attention_2.v_proj",
-                    output_proj=f"unet.encoders.{i*3+j+1}.1.attention_2.out_proj",
+                    norm=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.norm2",
+                    q_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn2.to_q",
+                    k_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn2.to_k",
+                    v_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn2.to_v",
+                    output_proj=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.attn2.to_out.0",
                 ),
                 feed_forward=stable_diffusion_loader.FeedForwardBlockTensorNames(
-                    norm=f"unet.encoders.{i*3+j+1}.1.layernorm_3",
-                    ge_glu=f"unet.encoders.{i*3+j+1}.1.linear_geglu_1",
-                    w2=f"unet.encoders.{i*3+j+1}.1.linear_geglu_2",
+                    norm=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.norm3",
+                    ge_glu=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.ff.net.0.proj",
+                    w2=f"model.diffusion_model.input_blocks.{i*3+j+1}.1.transformer_blocks.0.ff.net.2",
                 ),
             )
             for j in range(2)
         ]
         if i < 3
         else None,
-        downsample_conv=f"unet.encoders.{i*3+3}.0" if i < 3 else None,
+        downsample_conv=f"model.diffusion_model.input_blocks.{i*3+3}.0.op"
+        if i < 3
+        else None,
     )
     for i in range(4)
 ]
@@ -72,35 +76,37 @@ _down_encoder_blocks_tensor_names = [
 _mid_block_tensor_names = stable_diffusion_loader.MidBlockTensorNames(
     residual_block_tensor_names=[
         stable_diffusion_loader.ResidualBlockTensorNames(
-            norm_1=f"unet.bottleneck.{i}.groupnorm_feature",
-            conv_1=f"unet.bottleneck.{i}.conv_feature",
-            norm_2=f"unet.bottleneck.{i}.groupnorm_merged",
-            conv_2=f"unet.bottleneck.{i}.conv_merged",
-            time_embedding=f"unet.bottleneck.{i}.linear_time",
+            norm_1=f"model.diffusion_model.middle_block.{i}.in_layers.0",
+            conv_1=f"model.diffusion_model.middle_block.{i}.in_layers.2",
+            norm_2=f"model.diffusion_model.middle_block.{i}.out_layers.0",
+            conv_2=f"model.diffusion_model.middle_block.{i}.out_layers.3",
+            time_embedding=f"model.diffusion_model.middle_block.{i}.emb_layers.1",
         )
         for i in [0, 2]
     ],
     transformer_block_tensor_names=[
         stable_diffusion_loader.TransformerBlockTensorNames(
-            pre_conv_norm=f"unet.bottleneck.{i}.groupnorm",
-            conv_in=f"unet.bottleneck.{i}.conv_input",
-            conv_out=f"unet.bottleneck.{i}.conv_output",
+            pre_conv_norm=f"model.diffusion_model.middle_block.{i}.norm",
+            conv_in=f"model.diffusion_model.middle_block.{i}.proj_in",
+            conv_out=f"model.diffusion_model.middle_block.{i}.proj_out",
             self_attention=stable_diffusion_loader.AttentionBlockTensorNames(
-                norm=f"unet.bottleneck.{i}.layernorm_1",
-                fused_qkv_proj=f"unet.bottleneck.{i}.attention_1.in_proj",
-                output_proj=f"unet.bottleneck.{i}.attention_1.out_proj",
+                norm=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.norm1",
+                q_proj=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.attn1.to_q",
+                k_proj=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.attn1.to_k",
+                v_proj=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.attn1.to_v",
+                output_proj=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.attn1.to_out.0",
             ),
             cross_attention=stable_diffusion_loader.CrossAttentionBlockTensorNames(
-                norm=f"unet.bottleneck.{i}.layernorm_2",
-                q_proj=f"unet.bottleneck.{i}.attention_2.q_proj",
-                k_proj=f"unet.bottleneck.{i}.attention_2.k_proj",
-                v_proj=f"unet.bottleneck.{i}.attention_2.v_proj",
-                output_proj=f"unet.bottleneck.{i}.attention_2.out_proj",
+                norm=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.norm2",
+                q_proj=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.attn2.to_q",
+                k_proj=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.attn2.to_k",
+                v_proj=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.attn2.to_v",
+                output_proj=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.attn2.to_out.0",
             ),
             feed_forward=stable_diffusion_loader.FeedForwardBlockTensorNames(
-                norm=f"unet.bottleneck.{i}.layernorm_3",
-                ge_glu=f"unet.bottleneck.{i}.linear_geglu_1",
-                w2=f"unet.bottleneck.{i}.linear_geglu_2",
+                norm=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.norm3",
+                ge_glu=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.ff.net.0.proj",
+                w2=f"model.diffusion_model.middle_block.{i}.transformer_blocks.0.ff.net.2",
             ),
         )
         for i in [1]
@@ -111,58 +117,59 @@ _up_decoder_blocks_tensor_names = [
     stable_diffusion_loader.SkipUpDecoderBlockTensorNames(
         residual_block_tensor_names=[
             stable_diffusion_loader.ResidualBlockTensorNames(
-                norm_1=f"unet.decoders.{i*3+j}.0.groupnorm_feature",
-                conv_1=f"unet.decoders.{i*3+j}.0.conv_feature",
-                norm_2=f"unet.decoders.{i*3+j}.0.groupnorm_merged",
-                conv_2=f"unet.decoders.{i*3+j}.0.conv_merged",
-                time_embedding=f"unet.decoders.{i*3+j}.0.linear_time",
-                residual_layer=f"unet.decoders.{i*3+j}.0.residual_layer",
+                norm_1=f"model.diffusion_model.output_blocks.{i*3+j}.0.in_layers.0",
+                conv_1=f"model.diffusion_model.output_blocks.{i*3+j}.0.in_layers.2",
+                norm_2=f"model.diffusion_model.output_blocks.{i*3+j}.0.out_layers.0",
+                conv_2=f"model.diffusion_model.output_blocks.{i*3+j}.0.out_layers.3",
+                time_embedding=f"model.diffusion_model.output_blocks.{i*3+j}.0.emb_layers.1",
+                residual_layer=f"model.diffusion_model.output_blocks.{i*3+j}.0.skip_connection",
             )
             for j in range(3)
         ],
         transformer_block_tensor_names=[
             stable_diffusion_loader.TransformerBlockTensorNames(
-                pre_conv_norm=f"unet.decoders.{i*3+j}.1.groupnorm",
-                conv_in=f"unet.decoders.{i*3+j}.1.conv_input",
-                conv_out=f"unet.decoders.{i*3+j}.1.conv_output",
+                pre_conv_norm=f"model.diffusion_model.output_blocks.{i*3+j}.1.norm",
+                conv_in=f"model.diffusion_model.output_blocks.{i*3+j}.1.proj_in",
+                conv_out=f"model.diffusion_model.output_blocks.{i*3+j}.1.proj_out",
                 self_attention=stable_diffusion_loader.AttentionBlockTensorNames(
-                    norm=f"unet.decoders.{i*3+j}.1.layernorm_1",
-                    fused_qkv_proj=f"unet.decoders.{i*3+j}.1.attention_1.in_proj",
-                    output_proj=f"unet.decoders.{i*3+j}.1.attention_1.out_proj",
+                    norm=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.norm1",
+                    q_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn1.to_q",
+                    k_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn1.to_k",
+                    v_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn1.to_v",
+                    output_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn1.to_out.0",
                 ),
                 cross_attention=stable_diffusion_loader.CrossAttentionBlockTensorNames(
-                    norm=f"unet.decoders.{i*3+j}.1.layernorm_2",
-                    q_proj=f"unet.decoders.{i*3+j}.1.attention_2.q_proj",
-                    k_proj=f"unet.decoders.{i*3+j}.1.attention_2.k_proj",
-                    v_proj=f"unet.decoders.{i*3+j}.1.attention_2.v_proj",
-                    output_proj=f"unet.decoders.{i*3+j}.1.attention_2.out_proj",
+                    norm=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.norm2",
+                    q_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn2.to_q",
+                    k_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn2.to_k",
+                    v_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn2.to_v",
+                    output_proj=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.attn2.to_out.0",
                 ),
                 feed_forward=stable_diffusion_loader.FeedForwardBlockTensorNames(
-                    norm=f"unet.decoders.{i*3+j}.1.layernorm_3",
-                    ge_glu=f"unet.decoders.{i*3+j}.1.linear_geglu_1",
-                    w2=f"unet.decoders.{i*3+j}.1.linear_geglu_2",
+                    norm=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.norm3",
+                    ge_glu=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.ff.net.0.proj",
+                    w2=f"model.diffusion_model.output_blocks.{i*3+j}.1.transformer_blocks.0.ff.net.2",
                 ),
             )
             for j in range(3)
         ]
         if i > 0
         else None,
-        upsample_conv=f"unet.decoders.{i*3+2}.2.conv"
+        upsample_conv=f"model.diffusion_model.output_blocks.{i*3+2}.2.conv"
         if 0 < i < 3
-        else (f"unet.decoders.2.1.conv" if i == 0 else None),
+        else (f"model.diffusion_model.output_blocks.2.1.conv" if i == 0 else None),
     )
     for i in range(4)
 ]
 
-
-TENSORS_NAMES = stable_diffusion_loader.DiffusionModelLoader.TensorNames(
+TENSOR_NAMES = stable_diffusion_loader.DiffusionModelLoader.TensorNames(
     time_embedding=stable_diffusion_loader.TimeEmbeddingTensorNames(
-        w1="time_embedding.linear_1",
-        w2="time_embedding.linear_2",
+        w1="model.diffusion_model.time_embed.0",
+        w2="model.diffusion_model.time_embed.2",
     ),
-    conv_in="unet.encoders.0.0",
-    conv_out="final.conv",
-    final_norm="final.groupnorm",
+    conv_in="model.diffusion_model.input_blocks.0.0",
+    conv_out="model.diffusion_model.out.2",
+    final_norm="model.diffusion_model.out.0",
     down_encoder_blocks_tensor_names=_down_encoder_blocks_tensor_names,
     mid_block_tensor_names=_mid_block_tensor_names,
     up_decoder_blocks_tensor_names=_up_decoder_blocks_tensor_names,
@@ -249,6 +256,7 @@ class Diffusion(nn.Module):
         qkv_use_bias=False,
         output_proj_use_bias=True,
         enable_kv_cache=False,
+        qkv_fused_interleaved=False,
     )
 
     # Down encoders.
@@ -280,7 +288,7 @@ class Diffusion(nn.Module):
                         stride=2,
                         padding=config.downsample_padding,
                     ),
-                    transformer_block_config=unet_cfg.TransformerBlock2Dconfig(
+                    transformer_block_config=unet_cfg.TransformerBlock2DConfig(
                         attention_block_config=unet_cfg.AttentionBlock2DConfig(
                             dim=output_channel,
                             attention_batch_size=config.transformer_batch_size,
@@ -340,7 +348,7 @@ class Diffusion(nn.Module):
             ),
             num_layers=config.mid_block_layers,
             time_embedding_channels=config.time_embedding_blocks_dim,
-            transformer_block_config=unet_cfg.TransformerBlock2Dconfig(
+            transformer_block_config=unet_cfg.TransformerBlock2DConfig(
                 attention_block_config=unet_cfg.AttentionBlock2DConfig(
                     dim=mid_block_channels,
                     attention_batch_size=config.transformer_batch_size,
@@ -401,7 +409,7 @@ class Diffusion(nn.Module):
                         mode=unet_cfg.SamplingType.NEAREST,
                         scale_factor=2,
                     ),
-                    transformer_block_config=unet_cfg.TransformerBlock2Dconfig(
+                    transformer_block_config=unet_cfg.TransformerBlock2DConfig(
                         attention_block_config=unet_cfg.AttentionBlock2DConfig(
                             dim=output_channel,
                             attention_batch_size=config.transformer_batch_size,

ai_edge_torch/generative/examples/stable_diffusion/pipeline.py

@@ -167,7 +167,7 @@ def run_tflite_pipeline(
   if input_image:
     if not hasattr(model, 'encoder'):
       raise AttributeError(
-          'Stable Diffusion must be initilaized with encoder to accept input_image.'
+          'Stable Diffusion must be initialized with encoder to accept input_image.'
       )
     input_image = input_image.resize((width, height))
     input_image_np = np.array(input_image).astype(np.float32)

ai_edge_torch/generative/examples/t5/t5_attention.py

@@ -27,6 +27,8 @@ import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention  # NOQA
 from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention_with_hlfb  # NOQA
 
+BATCH_SIZE = 1
+
 
 class EncoderDecoderBlock(nn.Module):
 
@@ -44,6 +46,7 @@ class EncoderDecoderBlock(nn.Module):
 
     super().__init__()
     self.atten_func = T5Attention(
+        BATCH_SIZE,
         config.embedding_dim,
         config.attn_config,
         config.pre_attention_norm_config,
@@ -54,6 +57,7 @@ class EncoderDecoderBlock(nn.Module):
     # For a decoder, we add a cross attention.
     if config.is_decoder:
       self.cross_atten_func = T5Attention(
+          BATCH_SIZE,
           config.embedding_dim,
           config.attn_config,
           config.pre_attention_norm_config,
@@ -127,6 +131,7 @@ class T5Attention(CrossAttention):
 
   def __init__(
       self,
+      batch: int,
       dim: int,
       config: cfg.AttentionConfig,
       norm_config: cfg.NormalizationConfig,
@@ -144,7 +149,7 @@ class T5Attention(CrossAttention):
       enable_hlfb (bool): whether hlfb is enabled or not.
       has_relative_attention_bias (bool): whether we compute relative bias.
     """
-    super().__init__(dim, dim, config, kv_cache_max, enable_hlfb)
+    super().__init__(batch, dim, dim, config, kv_cache_max, enable_hlfb)
     self.pre_atten_norm = builder.build_norm(dim, norm_config)
 
     self.has_relative_attention_bias = has_relative_attention_bias

ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py

@@ -40,7 +40,7 @@ class RemoveSDPACompositeZeroMaskPass(ExportedProgramPassBase):
         if self.is_zero_tensor_node(source):
           # Remove the mark_tensor call on the mask input by
           # replacing the target with an identity function.
-          node.target = lambda *args, **kwargs: args[0]
+          node.target = lambda *args, **kwargs: torch.zeros_like(args[0])
 
     exported_program.graph_module.graph.lint()
     exported_program.graph_module.recompile()

ai_edge_torch/generative/layers/model_config.py

@@ -68,6 +68,10 @@ class AttentionConfig:
   qkv_transpose_before_split: bool = False
   # Whether to use bias with Query, Key, and Value projection.
   qkv_use_bias: bool = False
+  # Whether the fused q, k, v projection weights interleaves q, k, v heads.
+  # If True, the projection weights are in format [q_head_0, k_head_0, v_head_0, q_head_1, k_head_1, v_head_1, ...]
+  # If False, the projection weights are in format [q_head_0, q_head_1, ..., k_head_0, k_head_1, ... v_head_0, v_head_1, ...]
+  qkv_fused_interleaved: bool = True
   # Whether to use bias with attention output projection.
   output_proj_use_bias: bool = False
   enable_kv_cache: bool = True

ai_edge_torch/generative/layers/unet/blocks_2d.py

@@ -272,7 +272,7 @@ class TransformerBlock2D(nn.Module):
 
   """
 
-  def __init__(self, config: unet_cfg.TransformerBlock2Dconfig):
+  def __init__(self, config: unet_cfg.TransformerBlock2DConfig):
     """Initialize an instance of the TransformerBlock2D.
 
     Args:

ai_edge_torch/generative/layers/unet/model_config.py

@@ -85,7 +85,7 @@ class FeedForwardBlock2DConfig:
 
 
 @dataclass
-class TransformerBlock2Dconfig:
+class TransformerBlock2DConfig:
   pre_conv_normalization_config: layers_cfg.NormalizationConfig
   attention_block_config: AttentionBlock2DConfig
   cross_attention_block_config: CrossAttentionBlock2DConfig
@@ -108,7 +108,7 @@ class UpDecoderBlock2DConfig:
   # 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
+  transformer_block_config: Optional[TransformerBlock2DConfig] = None
   # Optional dimension of context tensor if context tensor is given as input.
   context_dim: Optional[int] = None
 
@@ -131,7 +131,7 @@ class SkipUpDecoderBlock2DConfig:
   # 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
+  transformer_block_config: Optional[TransformerBlock2DConfig] = None
   # Optional dimension of context tensor if context tensor is given as input.
   context_dim: Optional[int] = None
 
@@ -152,7 +152,7 @@ class DownEncoderBlock2DConfig:
   # 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
+  transformer_block_config: Optional[TransformerBlock2DConfig] = None
   # Optional dimension of context tensor if context tensor is given as input.
   context_dim: Optional[int] = None
 
@@ -168,7 +168,7 @@ class MidBlock2DConfig:
   # 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
+  transformer_block_config: Optional[TransformerBlock2DConfig] = None
   # Optional dimension of context tensor if context tensor is given as input.
   context_dim: Optional[int] = None
 

ai_edge_torch/generative/utilities/loader.py

@@ -317,9 +317,12 @@ class ModelLoader:
       k: torch.Tensor,
       v: torch.Tensor,
   ) -> torch.Tensor:
-    q_per_kv = config.attn_config.num_heads // config.attn_config.num_query_groups
-    qs = torch.split(q, config.head_dim * q_per_kv)
-    ks = torch.split(k, config.head_dim)
-    vs = torch.split(v, config.head_dim)
-    cycled = [t for group in zip(qs, ks, vs) for t in group]
-    return torch.cat(cycled)
+    if config.attn_config.qkv_fused_interleaved:
+      q_per_kv = config.attn_config.num_heads // config.attn_config.num_query_groups
+      qs = torch.split(q, config.head_dim * q_per_kv)
+      ks = torch.split(k, config.head_dim)
+      vs = torch.split(v, config.head_dim)
+      cycled = [t for group in zip(qs, ks, vs) for t in group]
+      return torch.cat(cycled)
+    else:
+      return torch.cat([q, k, v], dim=0)