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

ai_edge_torch/generative/utilities/t5_loader.py

@@ -318,7 +318,7 @@ class ModelLoader:
     q_name = names.attn_query_proj.format(idx)
     k_name = names.attn_key_proj.format(idx)
     v_name = names.attn_value_proj.format(idx)
-    # model.encoder.transformer_blocks[0].atten_func.q.weight
+    # model.encoder.transformer_blocks[0].atten_func.q_projection.weight
     if fuse_attention:
       converted_state[f"{prefix}.atten_func.attn.weight"] = self._fuse_qkv(
           config,
@@ -334,18 +334,34 @@ class ModelLoader:
             state.pop(f"{v_name}.bias"),
         )
     else:
-      converted_state[f"{prefix}.atten_func.q.weight"] = state.pop(f"{q_name}.weight")
-      converted_state[f"{prefix}.atten_func.k.weight"] = state.pop(f"{k_name}.weight")
-      converted_state[f"{prefix}.atten_func.v.weight"] = state.pop(f"{v_name}.weight")
+      converted_state[f"{prefix}.atten_func.q_projection.weight"] = state.pop(
+          f"{q_name}.weight"
+      )
+      converted_state[f"{prefix}.atten_func.k_projection.weight"] = state.pop(
+          f"{k_name}.weight"
+      )
+      converted_state[f"{prefix}.atten_func.v_projection.weight"] = state.pop(
+          f"{v_name}.weight"
+      )
       if config.attn_config.qkv_use_bias:
-        converted_state[f"{prefix}.atten_func.q.bias"] = state.pop(f"{q_name}.bias")
-        converted_state[f"{prefix}.atten_func.k.bias"] = state.pop(f"{k_name}.bias")
-        converted_state[f"{prefix}.atten_func.v.bias"] = state.pop(f"{v_name}.bias")
+        converted_state[f"{prefix}.atten_func.q_projection.bias"] = state.pop(
+            f"{q_name}.bias"
+        )
+        converted_state[f"{prefix}.atten_func.k_projection.bias"] = state.pop(
+            f"{k_name}.bias"
+        )
+        converted_state[f"{prefix}.atten_func.v_projection.bias"] = state.pop(
+            f"{v_name}.bias"
+        )
 
     o_name = names.attn_output_proj.format(idx)
-    converted_state[f"{prefix}.atten_func.proj.weight"] = state.pop(f"{o_name}.weight")
+    converted_state[f"{prefix}.atten_func.output_projection.weight"] = state.pop(
+        f"{o_name}.weight"
+    )
     if config.attn_config.output_proj_use_bias:
-      converted_state[f"{prefix}.atten_func.proj.bias"] = state.pop(f"{o_name}.bias")
+      converted_state[f"{prefix}.atten_func.output_projection.bias"] = state.pop(
+          f"{o_name}.bias"
+      )
 
   def _map_cross_attention(
       self,
@@ -383,32 +399,32 @@ class ModelLoader:
             state.pop(f"{v_name}.bias"),
         )
     else:
-      converted_state[f"{prefix}.cross_atten_func.q.weight"] = state.pop(
+      converted_state[f"{prefix}.cross_atten_func.q_projection.weight"] = state.pop(
           f"{q_name}.weight"
       )
-      converted_state[f"{prefix}.cross_atten_func.k.weight"] = state.pop(
+      converted_state[f"{prefix}.cross_atten_func.k_projection.weight"] = state.pop(
           f"{k_name}.weight"
       )
-      converted_state[f"{prefix}.cross_atten_func.v.weight"] = state.pop(
+      converted_state[f"{prefix}.cross_atten_func.v_projection.weight"] = state.pop(
           f"{v_name}.weight"
       )
       if config.attn_config.qkv_use_bias:
-        converted_state[f"{prefix}.cross_atten_func.q.bias"] = state.pop(
+        converted_state[f"{prefix}.cross_atten_func.q_projection.bias"] = state.pop(
             f"{q_name}.bias"
         )
-        converted_state[f"{prefix}.cross_atten_func.k.bias"] = state.pop(
+        converted_state[f"{prefix}.cross_atten_func.k_projection.bias"] = state.pop(
             f"{k_name}.bias"
         )
-        converted_state[f"{prefix}.cross_atten_func.v.bias"] = state.pop(
+        converted_state[f"{prefix}.cross_atten_func.v_projection.bias"] = state.pop(
             f"{v_name}.bias"
         )
 
     o_name = names.cross_attn_output_proj.format(idx)
-    converted_state[f"{prefix}.cross_atten_func.proj.weight"] = state.pop(
+    converted_state[f"{prefix}.cross_atten_func.output_projection.weight"] = state.pop(
         f"{o_name}.weight"
     )
     if config.attn_config.output_proj_use_bias:
-      converted_state[f"{prefix}.cross_atten_func.proj.bias"] = state.pop(
+      converted_state[f"{prefix}.cross_atten_func.output_projection.bias"] = state.pop(
           f"{o_name}.bias"
       )
 

{ai_edge_torch_nightly-0.2.0.dev20240611.dist-info → ai_edge_torch_nightly-0.2.0.dev20240618.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240611
+Version: 0.2.0.dev20240618
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.2.0.dev20240611.dist-info → ai_edge_torch_nightly-0.2.0.dev20240618.dist-info}/RECORD

@@ -15,20 +15,21 @@ ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_mark.py,sha256=uXCHC23pWN-3JmDtAErWbSUnL8jjlQgUAy4gqtfDsQU,1560
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=RAgU31B98PQmXEIM3GOjgS0q9aRe2whJhGXpW2EjoqY,12438
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py,sha256=tCx7J-WIFnxFCeRBtqJ159jWLgK9_9DCJrR4mkeBuYE,982
-ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=wr59GFss8fP8Vy--BaBj34Bto0N16gXxQj6OuTXH8cE,10030
+ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=cfY6RTWQTGXNoQxKHaDcBYR9QdkVQXOWjKhuxvglocw,10383
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/utils.py,sha256=bItkXVaPA9THcFypAmqldpkLuD8WpOFmKlhVbBJJkPk,2076
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/__init__.py,sha256=B-zisphkH7aRCUOJNdwHnTA0fQXuDpN08q3Qjy5bL6E,715
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=8uHJbIwPMTgeSfYVba163pkXSQkHLxFwar_8A1AhgAM,2279
-ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=FlNKt2EhIKnlVEeUWTiv5sz446YKU6Yy1H0Gd6VRgkU,6432
+ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=lklGxE1R32vsjFbhLLBDEFL4pfLi_iTgI9Ftb6Grezk,7156
 ai_edge_torch/convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/convert/test/test_convert.py,sha256=2qPmmGqnfV_o1gfsSdjGq3-JR1b323ligiy5MdAv9NA,8021
 ai_edge_torch/convert/test/test_convert_composites.py,sha256=_Ojc-H6GOS5s8ek3_8eRBL_AiCs-k3srziPJ2R4Ulrg,7255
 ai_edge_torch/convert/test/test_convert_multisig.py,sha256=kMaGnHe9ylfyU68qCifYcaGwJqyejKz--QQt9jS2oUA,4537
-ai_edge_torch/debug/__init__.py,sha256=TKvmnjVk3asvYcVh6C-LPr6srgAF_nppSAupWEXqwPY,707
-ai_edge_torch/debug/culprit.py,sha256=vklaxBUfINdo44OsH7csILK70N41gEThCGchGEfbTZw,12789
+ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
+ai_edge_torch/debug/culprit.py,sha256=urtCKPXORPvn6oyDxDSCSjgvngUnjjcsUMwAOeIl15E,14236
 ai_edge_torch/debug/utils.py,sha256=hjVmQVVl1dKxEF0D6KB4a3ouQ3wBkTsebOX2YsUObZM,1430
 ai_edge_torch/debug/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/debug/test/test_culprit.py,sha256=9An_n9p_RWTAYdHYTCO-__EJlbnjclCDo8tDhOzMlwk,3731
+ai_edge_torch/debug/test/test_search_model.py,sha256=0guAEon5cvwBpPXk6J0wVOKj7TXMDaiuomEEQmHgO5o,1590
 ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -41,9 +42,9 @@ ai_edge_torch/generative/examples/phi2/phi2.py,sha256=PMhKC6JCAMYSj2F3UmWHWK4rTc
 ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=Lo4Dq7a3Kg-lyH56iqGtqCo5UaClQHRCTDdNagXGTo8,3535
 ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=qU1wVEcn_biwCuDguZljhlLGzpLIqgqC31Dh_lXquQc,3720
-ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=NmgDo5uAefrhMUbYku0TKHlqzO0NVWI_M1ue8tddQR4,4024
-ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=meW8t-3BDdjFs5vCAf76cn6lGx49a_GcEvnVa9R5if4,11106
-ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=_gEeUxa9Xyd3iLb_fyeUefHKuELVDorDlQs8e7wdXKg,7878
+ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=wVEjsKd5JCIiYf5GF19rOXs2NHscZh0D69mxaS4f0Sk,4182
+ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=RgxedILk7iNMb0mhE4VkCs6d7BnFzYhR3vspUkC0-1o,11425
+ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=sRevfsmCun7zbceJbOstLKNUsLwzQDsGm7Mi2JmlREg,26021
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=mgbxkeFDMkNIGmnbcFTIFPu8EWKokghiviYIOB2lE3Q,3437
 ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=FCbnwlkpYYb-tF7KscbSYjNEdg7XnuLju1cDuIRoQv8,8277
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=r9RqbyNvuvXOGu3ojtl7ZmbC7o4Pt8aUKAhN1yCdtEc,3397
@@ -56,7 +57,7 @@ ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py,sha256=5i
 ai_edge_torch/generative/examples/t5/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/t5/convert_to_tflite.py,sha256=bWtwtUacvJOEDUpuYvLTgkP7oTkXKJA-Tf4FPxlD1Cw,4536
 ai_edge_torch/generative/examples/t5/t5.py,sha256=L6YrVzUEzP-Imb8W28LdukFGrx1aWSzz1kyYK_9RFZM,21087
-ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=rRgwCEdVtzcJEaGbbBjw8HxCxrCX3pXA5nelawdYiME,9036
+ai_edge_torch/generative/examples/t5/t5_attention.py,sha256=rkMwi-NJGBXHm5S57Rsj1LbcoVdyRkS7GmIBuU6F_2E,8274
 ai_edge_torch/generative/examples/test_models/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/test_models/toy_model.py,sha256=CUXsQ_IU96NaCg9jyfeKI0Zz2iWDkJUsPJyPR1Pgz7I,3813
 ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py,sha256=zwCmCnhr-vhBwHqv9i7xMasdBGVNqAGxZvWsncsJn58,5543
@@ -65,19 +66,19 @@ ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=47DEQpj8HBSa-_TI
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=E4I5OlC4zyl5cxiiu7uTED-zcwYRu210lP1zuT3xLBE,2566
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=IFRLPG9wz_aLl_zV_6CETCjSM03ukA6bZqqyDLVACuw,5651
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/attention.py,sha256=Z8gXHYs6h8gaRiYAdvYUbHzg_2EmqfxiChsf_SYraAc,7902
+ai_edge_torch/generative/layers/attention.py,sha256=AW0Qo3uOIe6p1rJNJ6zR_r4fqL2y-6QJHh0yUd-5Yb0,11966
 ai_edge_torch/generative/layers/attention_utils.py,sha256=hXhuyKblPPxKIRzlAf1YNlwHgpbj-6nReRLhRHELx5k,6350
 ai_edge_torch/generative/layers/builder.py,sha256=jAyrR5hsSI0aimKZumyvxdJ1GovERIfsK0g-dezX2gs,4163
 ai_edge_torch/generative/layers/feed_forward.py,sha256=4j2QaSCw59Jkk_ixKDpKEj7FLRauzuExTiSNRzAjAhE,2820
 ai_edge_torch/generative/layers/kv_cache.py,sha256=4uiZLO3om5G3--kT04Jt0esEYznbkJ7QLzSHfb8mjc4,3090
-ai_edge_torch/generative/layers/model_config.py,sha256=toWECENDWgay9hsZcy4C89qph0KI3CpaeFqFc8Fr-Xk,4584
+ai_edge_torch/generative/layers/model_config.py,sha256=aQLtOPdGpehfnb4aGO-iILLAsRU5t7j6opyezPEUY_w,4673
 ai_edge_torch/generative/layers/normalization.py,sha256=M27eW3TcNK20oaXClXtfnu0lLWrAGrSKSsbegRWnj3c,1867
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=12SsCuoRuLNCwnFGe_pHDOZEBwBcqXs87Aj0PaWWw4E,1383
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=dYafGC205QE5CLIbBTCI-7eVvEGZEHzs1toPEhemeDs,3391
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=7mHyJYq9lq5zVYp4mEz-R8Az3FFngi711YC20KP6ED8,10066
-ai_edge_torch/generative/layers/unet/builder.py,sha256=iH0_nuY9TF2ap5h1JbGNCOonPTfrXQHcF8U0slrIREM,1210
-ai_edge_torch/generative/layers/unet/model_config.py,sha256=sbtbDEHmMV9GLKngwjsNvqm8wovLxnlidkQbXdXkXKs,4060
+ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=H45wsXA6iJi_Mjd66NiQrh7i1fx05r9o_FI-fSnhVts,26538
+ai_edge_torch/generative/layers/unet/builder.py,sha256=NmJiZ2-e1wbv9jnvI3VCyUJlONV5ZAOz-RTc7ipAZ5U,1872
+ai_edge_torch/generative/layers/unet/model_config.py,sha256=FrIO-CR8aRIV2i8aFqom_4S7WCEDLMyYwo6U0oFyn7A,9097
 ai_edge_torch/generative/quantize/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/quantize/example.py,sha256=t-YwyKSPAG-OZC1DfH-0vfie2RHHpTSQjxUY-tmhu5g,1543
 ai_edge_torch/generative/quantize/quant_attrs.py,sha256=n1Fm8BFC8gJa_oiwwAOOghJyHtOXYZ4q-5ZRy4pHrIw,1957
@@ -89,12 +90,12 @@ ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/__init__.py,sha256=47DE
 ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/translate_recipe.py,sha256=qUB4f2DoB14dLkNPWf6TZodpT81mfAJeWM-lCAmkuHY,5735
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/test/loader_test.py,sha256=N88CbrLW7Q2x1EyurwdXQ6YjsA-ySQcPxpZH3QOGp-M,3317
-ai_edge_torch/generative/test/test_model_conversion.py,sha256=i_SAW-hD8SaHuopMZI9IuXXDFn5uSTJa1nKZhaC3dAQ,6811
-ai_edge_torch/generative/test/test_quantize.py,sha256=NVlMixAxVpDUabEvp6zTHHgIDgHFsMRwlf5MuyDwrPg,5355
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=LsPTrLC1I4JW2GowTS3V9Eu257vLHr2Yj5f_qaFUX84,7589
+ai_edge_torch/generative/test/test_quantize.py,sha256=IjCbCPWzIgXk3s7y7SJsg2usIxhOqs3PuhFvEYR4Sdw,5388
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/autoencoder_loader.py,sha256=G2Nosy33JzkjGALPR4JjvffdFX1JWOj2zjbbuaDJEgg,10065
 ai_edge_torch/generative/utilities/loader.py,sha256=Hs92478j1g4jQGvbdP1aWvOy907HjwqQZE-NFy6HELo,11326
-ai_edge_torch/generative/utilities/t5_loader.py,sha256=guDTv-12UUvJGl4eDvvZX3t4rRKewfXO8SpcYXM6gbc,16156
+ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=7ChqrnthD7I-Be6vkRvYTRhbGQ3tqMbikLpjY5HpSzE,30890
+ai_edge_torch/generative/utilities/t5_loader.py,sha256=h1FQzt4x8wiQMX4NzYNVIaJGLr_YKH0sojBvy0amexM,16503
 ai_edge_torch/hlfb/__init__.py,sha256=rrje8a2iuKboBoV96bVq7nlS9HsnuEMbHE5JiWmCxFA,752
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=2VXnHcGf23VOuP-1GriGIpuL98leBB8twp_qaScMnmc,4799
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=YV2YKBkh7y7j7sd7EA81vf_1hUKUvTRiy1pfqZustXc,1539
@@ -109,8 +110,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDd
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.2.0.dev20240611.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240611.dist-info/METADATA,sha256=WPGu2pq6N57fBtpunyFhunPe73UK_SVbqlZQsZwjWGo,1748
-ai_edge_torch_nightly-0.2.0.dev20240611.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240611.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240611.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/METADATA,sha256=aMhby_ftyg_8pWf8klYbHTCP7rMDcmuSTeryoRKt4U0,1748
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240618.dist-info/RECORD,,

ai_edge_torch/generative/utilities/autoencoder_loader.py

@@ -1,298 +0,0 @@
-# Copyright 2024 The AI Edge Torch Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-# Common utility functions for data loading etc.
-from dataclasses import dataclass
-from typing import Dict, List, Tuple
-
-import torch
-
-import ai_edge_torch.generative.layers.model_config as layers_config
-import ai_edge_torch.generative.layers.unet.model_config as unet_config
-import ai_edge_torch.generative.utilities.loader as loader
-
-
-@dataclass
-class ResidualBlockTensorNames:
-  norm_1: str = None
-  conv_1: str = None
-  norm_2: str = None
-  conv_2: str = None
-  residual_layer: str = None
-
-
-@dataclass
-class AttnetionBlockTensorNames:
-  norm: str = None
-  fused_qkv_proj: str = None
-  output_proj: str = None
-
-
-@dataclass
-class MidBlockTensorNames:
-  residual_block_tensor_names: List[ResidualBlockTensorNames]
-  attention_block_tensor_names: List[AttnetionBlockTensorNames]
-
-
-@dataclass
-class UpDecoderBlockTensorNames:
-  residual_block_tensor_names: List[ResidualBlockTensorNames]
-  upsample_conv: str = None
-
-
-def _map_to_converted_state(
-    state: Dict[str, torch.Tensor],
-    state_param: str,
-    converted_state: Dict[str, torch.Tensor],
-    converted_state_param: str,
-):
-  converted_state[f"{converted_state_param}.weight"] = state.pop(
-      f"{state_param}.weight"
-  )
-  if f"{state_param}.bias" in state:
-    converted_state[f"{converted_state_param}.bias"] = state.pop(f"{state_param}.bias")
-
-
-class AutoEncoderModelLoader(loader.ModelLoader):
-
-  @dataclass
-  class TensorNames:
-    quant_conv: str = None
-    post_quant_conv: str = None
-    conv_in: str = None
-    conv_out: str = None
-    final_norm: str = None
-    mid_block_tensor_names: MidBlockTensorNames = None
-    up_decoder_blocks_tensor_names: List[UpDecoderBlockTensorNames] = None
-
-  def __init__(self, file_name: str, names: TensorNames):
-    """AutoEncoderModelLoader constructor. Can be used to load encoder and decoder models.
-
-    Args:
-        file_name (str): Path to the checkpoint. Can be a directory or an
-          exact file.
-        names (TensorNames): An instance of `TensorNames` to determine mappings.
-    """
-    self._file_name = file_name
-    self._names = names
-    self._loader = self._get_loader()
-
-  def load(
-      self, model: torch.nn.Module, strict: bool = True
-  ) -> Tuple[List[str], List[str]]:
-    """Load the model from the checkpoint.
-
-    Args:
-        model (torch.nn.Module): The pytorch model that needs to be loaded.
-        strict (bool, optional): Whether the converted keys are strictly
-          matched. Defaults to True.
-
-    Returns:
-        missing_keys (List[str]): a list of str containing the missing keys.
-        unexpected_keys (List[str]): a list of str containing the unexpected keys.
-
-    Raises:
-        ValueError: If conversion results in unmapped tensors and strict mode is
-          enabled.
-    """
-    state = self._loader(self._file_name)
-    converted_state = dict()
-    if self._names.quant_conv is not None:
-      _map_to_converted_state(
-          state, self._names.quant_conv, converted_state, "quant_conv"
-      )
-    if self._names.post_quant_conv is not None:
-      _map_to_converted_state(
-          state, self._names.post_quant_conv, converted_state, "post_quant_conv"
-      )
-    if self._names.conv_in is not None:
-      _map_to_converted_state(state, self._names.conv_in, converted_state, "conv_in")
-    if self._names.conv_out is not None:
-      _map_to_converted_state(state, self._names.conv_out, converted_state, "conv_out")
-    if self._names.final_norm is not None:
-      _map_to_converted_state(
-          state, self._names.final_norm, converted_state, "final_norm"
-      )
-    self._map_mid_block(
-        state,
-        converted_state,
-        model.config.mid_block_config,
-        self._names.mid_block_tensor_names,
-    )
-
-    reversed_block_out_channels = list(reversed(model.config.block_out_channels))
-    block_out_channels = reversed_block_out_channels[0]
-    for i, out_channels in enumerate(reversed_block_out_channels):
-      prev_output_channel = block_out_channels
-      block_out_channels = out_channels
-      not_final_block = i < len(reversed_block_out_channels) - 1
-      self._map_up_decoder_block(
-          state,
-          converted_state,
-          f"up_decoder_blocks.{i}",
-          unet_config.UpDecoderBlock2DConfig(
-              in_channels=prev_output_channel,
-              out_channels=block_out_channels,
-              normalization_config=model.config.normalization_config,
-              activation_type=model.config.activation_type,
-              num_layers=model.config.layers_per_block,
-              add_upsample=not_final_block,
-              upsample_conv=True,
-          ),
-          self._names.up_decoder_blocks_tensor_names[i],
-      )
-    if strict and state:
-      raise ValueError(
-          f"Failed to map all tensor. Remaing tensor are: {list(state.keys())}"
-      )
-    return model.load_state_dict(converted_state, strict=strict)
-
-  def _map_residual_block(
-      self,
-      state: Dict[str, torch.Tensor],
-      converted_state: Dict[str, torch.Tensor],
-      tensor_names: ResidualBlockTensorNames,
-      converted_state_param_prefix: str,
-      config: unet_config.ResidualBlock2DConfig,
-  ):
-    _map_to_converted_state(
-        state,
-        tensor_names.norm_1,
-        converted_state,
-        f"{converted_state_param_prefix}.norm_1",
-    )
-    _map_to_converted_state(
-        state,
-        tensor_names.conv_1,
-        converted_state,
-        f"{converted_state_param_prefix}.conv_1",
-    )
-    _map_to_converted_state(
-        state,
-        tensor_names.norm_2,
-        converted_state,
-        f"{converted_state_param_prefix}.norm_2",
-    )
-    _map_to_converted_state(
-        state,
-        tensor_names.conv_2,
-        converted_state,
-        f"{converted_state_param_prefix}.conv_2",
-    )
-    if config.in_channels != config.out_channels:
-      _map_to_converted_state(
-          state,
-          tensor_names.residual_layer,
-          converted_state,
-          f"{converted_state_param_prefix}.residual_layer",
-      )
-
-  def _map_attention_block(
-      self,
-      state: Dict[str, torch.Tensor],
-      converted_state: Dict[str, torch.Tensor],
-      tensor_names: AttnetionBlockTensorNames,
-      converted_state_param_prefix: str,
-      config: unet_config.AttentionBlock2DConfig,
-  ):
-    if config.normalization_config.type != layers_config.NormalizationType.NONE:
-      _map_to_converted_state(
-          state,
-          tensor_names.norm,
-          converted_state,
-          f"{converted_state_param_prefix}.norm",
-      )
-    attention_layer_prefix = f"{converted_state_param_prefix}.attention"
-    _map_to_converted_state(
-        state,
-        tensor_names.fused_qkv_proj,
-        converted_state,
-        f"{attention_layer_prefix}.qkv_projection",
-    )
-    _map_to_converted_state(
-        state,
-        tensor_names.output_proj,
-        converted_state,
-        f"{attention_layer_prefix}.output_projection",
-    )
-
-  def _map_mid_block(
-      self,
-      state: Dict[str, torch.Tensor],
-      converted_state: Dict[str, torch.Tensor],
-      config: unet_config.MidBlock2DConfig,
-      tensor_names: MidBlockTensorNames,
-  ):
-    converted_state_param_prefix = "mid_block"
-    residual_block_config = unet_config.ResidualBlock2DConfig(
-        in_channels=config.in_channels,
-        out_channels=config.in_channels,
-        time_embedding_channels=config.time_embedding_channels,
-        normalization_config=config.normalization_config,
-        activation_type=config.activation_type,
-    )
-    self._map_residual_block(
-        state,
-        converted_state,
-        tensor_names.residual_block_tensor_names[0],
-        f"{converted_state_param_prefix}.resnets.0",
-        residual_block_config,
-    )
-    for i in range(config.num_layers):
-      if config.attention_block_config:
-        self._map_attention_block(
-            state,
-            converted_state,
-            tensor_names.attention_block_tensor_names[i],
-            f"{converted_state_param_prefix}.attentions.{i}",
-            config.attention_block_config,
-        )
-      self._map_residual_block(
-          state,
-          converted_state,
-          tensor_names.residual_block_tensor_names[i + 1],
-          f"{converted_state_param_prefix}.resnets.{i+1}",
-          residual_block_config,
-      )
-
-  def _map_up_decoder_block(
-      self,
-      state: Dict[str, torch.Tensor],
-      converted_state: Dict[str, torch.Tensor],
-      converted_state_param_prefix: str,
-      config: unet_config.UpDecoderBlock2DConfig,
-      tensor_names: UpDecoderBlockTensorNames,
-  ):
-    for i in range(config.num_layers):
-      input_channels = config.in_channels if i == 0 else config.out_channels
-      self._map_residual_block(
-          state,
-          converted_state,
-          tensor_names.residual_block_tensor_names[i],
-          f"{converted_state_param_prefix}.resnets.{i}",
-          unet_config.ResidualBlock2DConfig(
-              in_channels=input_channels,
-              out_channels=config.out_channels,
-              time_embedding_channels=config.time_embedding_channels,
-              normalization_config=config.normalization_config,
-              activation_type=config.activation_type,
-          ),
-      )
-    if config.add_upsample and config.upsample_conv:
-      _map_to_converted_state(
-          state,
-          tensor_names.upsample_conv,
-          converted_state,
-          f"{converted_state_param_prefix}.upsample_conv",
-      )