ai-edge-torch-nightly 0.7.0.dev20250929__py3-none-any.whl → 0.8.0.dev20251206__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -133,10 +133,11 @@ def convert_signatures(
     exported_program = fx_infra.safe_run_decompositions(
         exported_program,
         fx_infra.decomp.pre_convert_decomp(),
+        can_skip=False,
     )
     return exported_program
 
-  exported_programs: torch.export.ExportedProgram = [
+  exported_programs = [
       export(
           mod=sig.module,
           args=sig.args,

ai_edge_torch/fx_infra/_safe_run_decompositions.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """ExportedProgram.run_decompositions wrapper to handle unexpected export behavior."""
+import operator
 import torch
 
 
@@ -26,8 +27,39 @@ _DUMMY_DECOMP_TABLE = {
     torch._ops.OperatorBase(): lambda: None,
 }
 
+_BUILTIN_OPERATORS = {
+    getattr(operator, name)
+    for name in dir(operator)
+    if not name.startswith("_")
+}
+
+
+def _require_decomp(
+    exported_program: torch.export.ExportedProgram, decomp_table
+):
+  """Checks if the exported program requires decompositions."""
+  for node in exported_program.graph.nodes:
+    if "call_" not in str(node.op):
+      continue
+
+    op = node.target
+    if isinstance(op, torch._ops.OpOverloadPacket):
+      op = op.default
 
-def safe_run_decompositions(exported_program, decomp_table=None):
+    if op in decomp_table:
+      return True
+
+    if (
+        not isinstance(op, (torch._ops.OpOverload, torch._ops.OperatorBase))
+        and op not in _BUILTIN_OPERATORS
+    ):
+      # Python function that requires to be retraced via run_decompositions.
+      return True
+
+  return False
+
+
+def safe_run_decompositions(exported_program, decomp_table=None, can_skip=True):
   """Wrapper for ExportedProgram.run_decompositions to handle unexpected export behavior."""
 
   if decomp_table is not None and not decomp_table:
@@ -35,6 +67,9 @@ def safe_run_decompositions(exported_program, decomp_table=None):
     # instead for backward compatibility.
     decomp_table = _DUMMY_DECOMP_TABLE
 
+  if can_skip and not _require_decomp(exported_program, decomp_table):
+    return exported_program
+
   for node in exported_program.graph.nodes:
     if node.target == torch.ops.aten.view.default:
       # Passes or torch.export may generate aten.view nodes not respecting the

ai_edge_torch/generative/examples/amd_llama_135m/convert_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.amd_llama_135m import amd_llama_135m
 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("amd-llama-135m")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = amd_llama_135m.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(amd_llama_135m.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.deepseek import deepseek
 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(
     'deepseek', default_mask_as_input=True, default_transpose_kv_cache=True
@@ -27,24 +25,7 @@ flags = converter.define_conversion_flags(
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = deepseek.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(deepseek.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.gemma import gemma1
 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("gemma-2b")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = gemma1.build_2b_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(gemma1.build_2b_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.gemma import gemma2
 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(
     "gemma2-2b", default_mask_as_input=True, default_transpose_kv_cache=True
@@ -27,24 +25,7 @@ flags = converter.define_conversion_flags(
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = gemma2.build_2b_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(gemma2.build_2b_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.gemma3 import gemma3
 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(
     'gemma3-1b', default_mask_as_input=True, default_transpose_kv_cache=True
@@ -33,36 +31,14 @@ _MODEL_SIZE = flags.DEFINE_string(
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
   if _MODEL_SIZE.value == '1b':
-    pytorch_model = gemma3.build_model_1b(
-        checkpoint_path,
-        custom_loader=loader.maybe_get_custom_loader(
-            checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-        ),
-        mask_cache_size=converter.get_mask_cache_size_from_flags(),
-    )
+    model_builder = gemma3.build_model_1b
   elif _MODEL_SIZE.value == '270m':
-    pytorch_model = gemma3.build_model_270m(
-        checkpoint_path,
-        custom_loader=loader.maybe_get_custom_loader(
-            checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-        ),
-        mask_cache_size=converter.get_mask_cache_size_from_flags(),
-    )
+    model_builder = gemma3.build_model_270m
   else:
     raise ValueError(f'Unsupported model size: {_MODEL_SIZE.value}')
 
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(model_builder)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/hammer/convert_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.hammer import hammer
 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('hammer')
 
@@ -37,24 +35,7 @@ _BUILDER = {
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(_BUILDER[_MODEL_SIZE.value])
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/llama/convert_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.llama import llama
 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('llama')
 
@@ -37,24 +35,7 @@ _BUILDER = {
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(_BUILDER[_MODEL_SIZE.value])
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/openelm/convert_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.openelm import openelm
 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("openelm")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = openelm.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(openelm.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/phi/convert_phi3_to_tflite.py

@@ -18,31 +18,12 @@
 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(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(phi3.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/phi/convert_phi4_to_tflite.py

@@ -18,31 +18,12 @@
 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(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(phi4.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/phi/convert_to_tflite.py

@@ -19,31 +19,12 @@
 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(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(phi2.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/qwen/convert_to_tflite.py

@@ -18,8 +18,6 @@
 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')
 
@@ -38,24 +36,7 @@ _BUILDER = {
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(_BUILDER[_MODEL_SIZE.value])
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/qwen/convert_v3_to_tflite.py

@@ -18,8 +18,6 @@
 from absl import app
 from ai_edge_torch.generative.examples.qwen import qwen3
 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')
 
@@ -38,24 +36,7 @@ _BUILDER = {
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = _BUILDER[_MODEL_SIZE.value](
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(_BUILDER[_MODEL_SIZE.value])
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/smollm/convert_to_tflite.py

@@ -18,41 +18,12 @@
 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('smollm')
 
-_DECODE_BATCH_SIZE = flags.DEFINE_integer(
-    'decode_batch_size',
-    1,
-    'The batch size for the decode signature.',
-)
-
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = smollm.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-
-  export_config = export_cfg.get_from_flags()
-  export_config.decode_batch_size = _DECODE_BATCH_SIZE.value
-
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config,
-  )
+  converter.build_and_convert_to_tflite_from_flags(smollm.build_model)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/smollm/convert_v2_to_tflite.py

@@ -18,41 +18,12 @@
 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')
 
-_DECODE_BATCH_SIZE = flags.DEFINE_integer(
-    'decode_batch_size',
-    1,
-    'The batch size for the decode signature.',
-)
-
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = smollm.build_model_v2(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-
-  export_config = export_cfg.get_from_flags()
-  export_config.decode_batch_size = _DECODE_BATCH_SIZE.value
-
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config,
-  )
+  converter.build_and_convert_to_tflite_from_flags(smollm.build_model_v2)
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py

@@ -138,9 +138,7 @@ def convert_stable_diffusion_to_tflite(
   if not os.path.exists(output_dir):
     pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)
 
-  quant_config = (
-      quant_recipes.full_int8_weight_only_recipe() if quantize else None
-  )
+  quant_config = quant_recipes.full_weight_only_recipe() if quantize else None
 
   # TODO(yichunk): convert to multi signature tflite model.
   # CLIP text encoder

ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py

@@ -69,6 +69,6 @@ class KLMSSampler(SamplerInterface):
           continue
         y *= x - self.sigmas[t - j]
         y /= self.sigmas[t - i] - self.sigmas[t - j]
-      lms_coeff = np.trapz(y=y, x=x)
+      lms_coeff = np.trapezoid(y=y, x=x)
       latents += lms_coeff * output
     return latents

ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py

@@ -18,31 +18,12 @@
 from absl import app
 from ai_edge_torch.generative.examples.tiny_llama import tiny_llama
 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("tiny_llama")
 
 
 def main(_):
-  checkpoint_path = flags.FLAGS.checkpoint_path
-  pytorch_model = tiny_llama.build_model(
-      checkpoint_path,
-      custom_loader=loader.maybe_get_custom_loader(
-          checkpoint_path, flags.FLAGS.custom_checkpoint_loader
-      ),
-      mask_cache_size=converter.get_mask_cache_size_from_flags(),
-  )
-  converter.convert_to_tflite(
-      pytorch_model,
-      output_path=flags.FLAGS.output_path,
-      output_name_prefix=flags.FLAGS.output_name_prefix,
-      prefill_seq_len=flags.FLAGS.prefill_seq_lens,
-      kv_cache_max_len=flags.FLAGS.kv_cache_max_len,
-      quantize=flags.FLAGS.quantize,
-      lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=export_config.get_from_flags(),
-  )
+  converter.build_and_convert_to_tflite_from_flags(tiny_llama.build_model)
 
 
 if __name__ == '__main__':