ai-edge-torch-nightly 0.5.0.dev20250515__py3-none-any.whl → 0.5.0.dev20250516__py3-none-any.whl

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -96,6 +96,7 @@ def get_decoder2_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       type=cfg.NormalizationType.RMS_NORM,
       epsilon=1e-6,
       zero_centered=True,
+      enable_hlfb=True,
   )
   ff_config = cfg.FeedForwardConfig(
       type=cfg.FeedForwardType.GATED,

ai_edge_torch/generative/examples/paligemma/image_encoder.py

@@ -66,7 +66,8 @@ class SiglipVisionEncoder(nn.Module):
         config.image_embedding.image_size // config.image_embedding.patch_size
     ) ** 2
     self.tok_embedding_position = nn.Parameter(
-        torch.zeros((num_patches, config.embedding_dim))
+        torch.zeros((num_patches, config.embedding_dim)),
+        requires_grad=False,
     )
 
     self.transformer_blocks = nn.ModuleList(

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -16,7 +16,7 @@
 """Example of building a full-stack of PaliGemma model."""
 
 import dataclasses
-from typing import Optional
+from typing import Callable, Dict, Optional
 
 from ai_edge_torch.generative.examples.paligemma import decoder
 from ai_edge_torch.generative.examples.paligemma import decoder2
@@ -139,7 +139,12 @@ def get_fake_model_config(get_decoder_config, **kwargs) -> PaliGemmaConfig:
   )
 
 
-def build_model(checkpoint_path: str, version: int = 2, **kwargs) -> PaliGemma:
+def build_model(
+    checkpoint_path: str,
+    version: int = 2,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs,
+) -> PaliGemma:
   if version == 1:
     decoder_class = decoder.Decoder
     decoder_tensor_names = decoder.TENSOR_NAMES
@@ -153,15 +158,17 @@ def build_model(checkpoint_path: str, version: int = 2, **kwargs) -> PaliGemma:
   model = PaliGemma(config, decoder_class)
   # Load the parameters of image encoder.
   loader = loading_utils.ModelLoader(
-      checkpoint_path, image_encoder.TENSOR_NAMES
+      checkpoint_path, image_encoder.TENSOR_NAMES, custom_loader
   )
   loader.load(model.image_encoder, strict=False)
   # Load the parameters of decoder.
-  loader = loading_utils.ModelLoader(checkpoint_path, decoder_tensor_names)
+  loader = loading_utils.ModelLoader(
+      checkpoint_path, decoder_tensor_names, custom_loader
+  )
   loader.load(model.decoder, strict=False)
 
   # Load the parameters of image projection.
-  loader = loading_utils.ModelLoader(checkpoint_path, None)
+  loader = loading_utils.ModelLoader(checkpoint_path, None, custom_loader)
   state = loader.get_state()
   converted_state = dict()
   converted_state["weight"] = state.pop(f"{PROJECTION_TENSOR_NAME}.weight")

ai_edge_torch/generative/examples/paligemma/verify.py

@@ -21,6 +21,7 @@ from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.paligemma import paligemma
 from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.utilities import transformers_verifier
 from ai_edge_torch.generative.utilities import verifier
 import kagglehub
 from PIL import Image
@@ -39,10 +40,15 @@ _IMAGE_URL = flags.DEFINE_string(
     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true",
     "The image URI to encode.",
 )
-_PROMPTS = flags.DEFINE_string(
-    "prompts",
+_PROMPTS_WITH_IMAGE = flags.DEFINE_string(
+    "prompts_with_image",
     "<image><bos>describe en",
-    "The input prompts to generate answers.",
+    "The input prompts to generate answers with an image.",
+)
+_PROMPTS_TEXT_ONLY = flags.DEFINE_multi_string(
+    "prompts_text_only",
+    "What is the meaning of life?",
+    "The input prompts to generate answers only with text.",
 )
 _MAX_NEW_TOKENS = flags.DEFINE_integer(
     "max_new_tokens",
@@ -84,6 +90,7 @@ def main(_):
   reauthored_model = paligemma.build_model(
       reauthored_checkpoint, version=int(_VERSION.value)
   )
+  wrapped_reauthored_model = ReauthoredPaliGemmaWrapper(reauthored_model)
 
   logging.info("Loading the processor from: %s", checkpoint)
   # It works only when GemmaTokenizerFast is available. In some environments,
@@ -91,9 +98,25 @@ def main(_):
   # sentencepiece model file properly.
   processor = transformers.AutoProcessor.from_pretrained(checkpoint)
 
+  logging.info("Verifying with text-only prompts...")
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=wrapped_reauthored_model,
+      tokenizer=verifier.TokenizerWrapper(processor.tokenizer),
+      generate_prompts=_PROMPTS_TEXT_ONLY.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      verify_inputs=False,  # Numeric check not working. Disable it for now.
+      atol=1e-04,
+  )
+
+  logging.info("Verifying with image input...")
   logging.info("Loading the image from: %s", _IMAGE_URL.value)
   image = Image.open(requests.get(_IMAGE_URL.value, stream=True).raw)
-  inputs = processor(text=_PROMPTS.value, images=image, return_tensors="pt")
+  inputs = processor(
+      text=_PROMPTS_WITH_IMAGE.value, images=image, return_tensors="pt"
+  )
 
   logging.info("Verifying the reauthored model with model.forward()...")
   logging.info("Forwarding the original model...")
@@ -104,7 +127,6 @@ def main(_):
   logging.info("outputs_original: %s", outputs_original)
 
   logging.info("Forwarding the reauthored model...")
-  wrapped_reauthored_model = ReauthoredPaliGemmaWrapper(reauthored_model)
   outputs_reauthored = wrapped_reauthored_model.forward(
       tokens=inputs["input_ids"],
       pixel_values=inputs["pixel_values"],

ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py

@@ -19,13 +19,19 @@ from absl import app
 from ai_edge_torch.generative.examples.phi import phi3
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags("phi3")
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = phi3.build_model(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py

@@ -19,13 +19,19 @@ from absl import app
 from ai_edge_torch.generative.examples.phi import phi4
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags("phi4")
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = phi4.build_model(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/phi/convert_to_tflite.py

@@ -20,13 +20,19 @@ from absl import app
 from ai_edge_torch.generative.examples.phi import phi2
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags("phi2")
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = phi2.build_model(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/phi/phi2.py

@@ -15,9 +15,11 @@
 
 """Example of building a Phi-2 model."""
 
+from typing import Callable, Dict
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
 from torch import nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
@@ -98,10 +100,15 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Phi2,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/phi/phi3.py

@@ -17,7 +17,7 @@
 
 from functools import partial
 import math
-from typing import Tuple
+from typing import Callable, Dict, Tuple
 
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
@@ -208,11 +208,16 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> torch.nn.Module:
   """Instantiates the model instance and load checkpoint if provided."""
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Phi3_5Mini,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/phi/phi4.py

@@ -17,7 +17,7 @@
 
 from functools import partial
 import math
-from typing import Tuple
+from typing import Callable, Dict, Tuple
 
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
@@ -157,11 +157,16 @@ def get_fake_model_config(kv_cache_max_len: int = 128) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> torch.nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> torch.nn.Module:
   """Instantiates the model instance and load checkpoint if provided."""
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Phi4Mini,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/qwen/convert_to_tflite.py

@@ -19,6 +19,7 @@ from absl import app
 from ai_edge_torch.generative.examples.qwen import qwen
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags('qwen')
 
@@ -37,8 +38,13 @@ _BUILDER = {
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
   converter.convert_to_tflite(
       pytorch_model,

ai_edge_torch/generative/examples/qwen/qwen.py

@@ -15,8 +15,10 @@
 
 """Example of building Qwen 2.5 models."""
 
+from typing import Callable, Dict
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
+import torch
 from torch import nn
 
 TENSOR_NAMES = model_builder.TENSOR_NAMES
@@ -108,28 +110,43 @@ def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
   return config
 
 
-def build_3b_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_3b_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_3b_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Qwen,
+      custom_loader=custom_loader,
   )
 
 
-def build_1_5b_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_1_5b_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_1_5b_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Qwen,
+      custom_loader=custom_loader,
   )
 
 
-def build_0_5b_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_0_5b_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_0_5b_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=Qwen,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py

@@ -19,6 +19,7 @@ from absl import app
 from ai_edge_torch.generative.examples.qwen_vl import qwen_vl
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags('qwen_vl')
 
@@ -35,8 +36,12 @@ _IMAGE_WIDTH = flags.DEFINE_integer(
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = qwen_vl.build_model(
-      flags.FLAGS.checkpoint_path,
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
       kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
       image_size=(_IMAGE_HEIGHT.value, _IMAGE_WIDTH.value),
   )

ai_edge_torch/generative/examples/qwen_vl/decoder.py

@@ -97,8 +97,7 @@ def get_decoder_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
       intermediate_size=11008,
   )
   norm_config = cfg.NormalizationConfig(
-      type=cfg.NormalizationType.RMS_NORM,
-      epsilon=1e-06,
+      type=cfg.NormalizationType.RMS_NORM, epsilon=1e-06, enable_hlfb=True
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,

ai_edge_torch/generative/examples/qwen_vl/image_encoder.py

@@ -16,7 +16,7 @@
 """Example of building an image encoder of Qwen 2.5 VL model."""
 
 import dataclasses
-from typing import List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import attention_utils
@@ -385,13 +385,21 @@ def build_image_encoder(
   return encoder
 
 
-def load_image_encoder(checkpoint_path: str, encoder: QwenVLImageEncoder):
-  loader = loading_utils.ModelLoader(checkpoint_path, TENSOR_NAMES)
+def load_image_encoder(
+    checkpoint_path: str,
+    encoder: QwenVLImageEncoder,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+):
+  loader = loading_utils.ModelLoader(
+      checkpoint_path, TENSOR_NAMES, custom_loader
+  )
   # Loose the strictness because only image encoder is being loaded.
   loader.load(encoder, strict=False)
 
   # Load merger weights.
-  merger_loader = loading_utils.ModelLoader(checkpoint_path, None)
+  merger_loader = loading_utils.ModelLoader(
+      checkpoint_path, None, custom_loader
+  )
   state = merger_loader.get_state()
   w1_state = dict()
   w1_state["weight"] = state.pop(f"{MERGER_TENSOR_NAMES.ff_up_proj}.weight")

ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py

@@ -16,7 +16,7 @@
 """Example of building a full-stack of Qwen 2.5 VL model."""
 
 import dataclasses
-from typing import List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 from ai_edge_torch.generative.examples.qwen_vl import decoder
 from ai_edge_torch.generative.examples.qwen_vl import image_encoder
@@ -204,12 +204,20 @@ def get_fake_model_config(**kwargs) -> QwenVLConfig:
   )
 
 
-def build_model(checkpoint_path: str, **kwargs) -> QwenVL:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> QwenVL:
   config = get_model_config(**kwargs)
   model = QwenVL(config)
-  image_encoder.load_image_encoder(checkpoint_path, model.image_encoder)
+  image_encoder.load_image_encoder(
+      checkpoint_path, model.image_encoder, custom_loader
+  )
   # Load the parameters of decoder.
-  loader = loading_utils.ModelLoader(checkpoint_path, decoder.TENSOR_NAMES)
+  loader = loading_utils.ModelLoader(
+      checkpoint_path, decoder.TENSOR_NAMES, custom_loader
+  )
   loader.load(model.decoder, strict=False)
   model.eval()
   return model

ai_edge_torch/generative/examples/qwen_vl/verify.py

@@ -22,6 +22,7 @@ from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.qwen_vl import qwen_vl
 from ai_edge_torch.generative.layers import kv_cache
+from ai_edge_torch.generative.utilities import transformers_verifier
 from ai_edge_torch.generative.utilities import verifier
 from PIL import Image
 import requests
@@ -33,10 +34,15 @@ _IMAGE_URL = flags.DEFINE_string(
     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg?download=true",
     "The image URI to encode.",
 )
-_PROMPTS = flags.DEFINE_string(
-    "prompts",
+_PROMPTS_WITH_IMAGE = flags.DEFINE_string(
+    "prompts_with_image",
     "<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>",
-    "The input prompts to generate answers.",
+    "The input prompts to generate answers with an image.",
+)
+_PROMPTS_TEXT_ONLY = flags.DEFINE_multi_string(
+    "prompts_text_only",
+    "What is the meaning of life?",
+    "The input prompts to generate answers only with text.",
 )
 _MAX_NEW_TOKENS = flags.DEFINE_integer(
     "max_new_tokens",
@@ -68,13 +74,29 @@ def main(_):
   reauthored_checkpoint = pathlib.Path(cached_config_file).parent
   logging.info("Building the reauthored model from: %s", reauthored_checkpoint)
   reauthored_model = qwen_vl.build_model(str(reauthored_checkpoint))
+  wrapped_reauthored_model = ReauthoredQwenVLWrapper(reauthored_model)
 
   logging.info("Loading the processor from: %s", checkpoint)
   processor = transformers.AutoProcessor.from_pretrained(checkpoint)
 
+  logging.info("Verifying with text-only prompts...")
+  verifier.verify_reauthored_model(
+      original_model=transformers_verifier.TransformersModelWrapper(
+          original_model
+      ),
+      reauthored_model=wrapped_reauthored_model,
+      tokenizer=verifier.TokenizerWrapper(processor.tokenizer),
+      generate_prompts=_PROMPTS_TEXT_ONLY.value,
+      max_new_tokens=_MAX_NEW_TOKENS.value,
+      atol=1e-04,
+  )
+
+  logging.info("Verifying with image input...")
   logging.info("Loading the image from: %s", _IMAGE_URL.value)
   image = Image.open(requests.get(_IMAGE_URL.value, stream=True).raw)
-  inputs = processor(text=_PROMPTS.value, images=image, return_tensors="pt")
+  inputs = processor(
+      text=_PROMPTS_WITH_IMAGE.value, images=image, return_tensors="pt"
+  )
 
   logging.info("Verifying the reauthored model with model.forward()...")
   logging.info("Forwarding the original model...")
@@ -87,7 +109,6 @@ def main(_):
   logging.info("outputs_original: %s", outputs_original)
 
   logging.info("Forwarding the reauthored model...")
-  wrapped_reauthored_model = ReauthoredQwenVLWrapper(reauthored_model)
   grid_thw = inputs["image_grid_thw"].tolist()
   config = reauthored_model.config.image_encoder_config.image_embedding
   reauthored_model.image_encoder.set_image_size(

ai_edge_torch/generative/examples/smollm/convert_to_tflite.py

@@ -15,12 +15,12 @@
 
 """Example of converting SmolLM model to multi-signature tflite model."""
 
-import os
 from absl import app
 from absl import flags
 from ai_edge_torch.generative.examples.smollm import smollm
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config as export_cfg
+from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags('smollm')
 
@@ -32,8 +32,13 @@ _DECODE_BATCH_SIZE = flags.DEFINE_integer(
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = smollm.build_model(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
 
   export_config = export_cfg.get_from_flags()

ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py

@@ -19,6 +19,7 @@ from absl import app
 from ai_edge_torch.generative.examples.smollm import smollm
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config as export_cfg
+from ai_edge_torch.generative.utilities import loader
 
 flags = converter.define_conversion_flags('smollm2')
 
@@ -30,8 +31,13 @@ _DECODE_BATCH_SIZE = flags.DEFINE_integer(
 
 
 def main(_):
+  checkpoint_path = flags.FLAGS.checkpoint_path
   pytorch_model = smollm.build_model_v2(
-      flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
+      checkpoint_path,
+      custom_loader=loader.maybe_get_custom_loader(
+          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
+      ),
+      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
   )
 
   export_config = export_cfg.get_from_flags()

ai_edge_torch/generative/examples/smollm/smollm.py

@@ -15,8 +15,10 @@
 
 """Example of building a SmolLM model."""
 
+from typing import Callable, Dict
 import ai_edge_torch.generative.layers.model_config as cfg
 from ai_edge_torch.generative.utilities import model_builder
+import torch
 from torch import nn
 
 TENSOR_NAMES = model_builder.TENSOR_NAMES
@@ -80,12 +82,17 @@ def get_fake_model_config(**kwargs) -> cfg.ModelConfig:
   return config
 
 
-def build_model(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_model(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=SmolLM,
+      custom_loader=custom_loader,
   )
 
 
@@ -118,10 +125,15 @@ def get_fake_model_config_v2(**kwargs) -> cfg.ModelConfig:
   return config
 
 
-def build_model_v2(checkpoint_path: str, **kwargs) -> nn.Module:
+def build_model_v2(
+    checkpoint_path: str,
+    custom_loader: Callable[[str], Dict[str, torch.Tensor]] = None,
+    **kwargs
+) -> nn.Module:
   return model_builder.build_decoder_only_model(
       checkpoint_path=checkpoint_path,
       config=get_model_config_v2(**kwargs),
       tensor_names=TENSOR_NAMES,
       model_class=SmolLM2,
+      custom_loader=custom_loader,
   )

ai_edge_torch/generative/examples/smollm/verify.py

@@ -43,8 +43,8 @@ _MODEL_VERSION = flags.DEFINE_enum(
     "The version of SmolLm to verify.",
 )
 _CHECKPOINT = {
-    "v1": "HuggingFaceTB/SmolLM-135M",
-    "v2": "HuggingFaceTB/SmolLM2-135M",
+    "v1": "HuggingFaceTB/SmolLM-135M-Instruct",
+    "v2": "HuggingFaceTB/SmolLM2-135M-Instruct",
 }
 
 _BUILDER = {

ai_edge_torch/generative/examples/stable_diffusion/clip.py

@@ -57,7 +57,8 @@ class CLIP(nn.Module):
     super().__init__()
     self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
     self.tok_embedding_position = nn.Parameter(
-        torch.zeros((config.max_seq_len, config.embedding_dim))
+        torch.zeros((config.max_seq_len, config.embedding_dim)),
+        requires_grad=False,
     )
 
     self.config = config