Function main

examples/models/llava/export_llava.py:292–361

()

def main():
    parser = ArgumentParser()
    parser.add_argument(
        "--use-sdpa-with-kv-cache",
        default=True,
        action=BooleanOptionalAction,
        help="Use sdpa_with_kv_cache custom op in LLava text model.",
    )
    parser.add_argument(
        "--max-context-len",
        required=True,
        type=int,
        help="Maximum context length for the text model.",
    )
    parser.add_argument(
        "--max-seq-len",
        default=768,
        type=int,
        help="Maximum sequence length for the text model.",
    )
    parser.add_argument(
        "--pte-name",
        default="llava_combined_xnnpack.pte",
        help="Name of the exported ExecuTorch program.",
    )
    parser.add_argument(
        "--with-artifacts",
        default=False,
        action=BooleanOptionalAction,
        help="Generate artifacts for llava runner.",
    )
    parser.add_argument(
        "--profile_memory",
        required=False,
        action="store_true",
        help="Generate chrome trace of activation memory for intermediate tensors.",
    )
    args = parser.parse_args()

    # Create LlmConfig from args
    llm_config = create_llava_config_from_args(args)

    logging.info(
        f"Exporting Llava model to ExecuTorch with sdpa_with_kv_cache: {llm_config.model.use_sdpa_with_kv_cache}, max_seq_len: {llm_config.export.max_seq_length}, max_context_len: {llm_config.export.max_context_length}"
    )

    llava_model = LlavaModel(
        use_sdpa_with_kv_cache_op=llm_config.model.use_sdpa_with_kv_cache,
        max_seq_len=llm_config.export.max_seq_length,
        max_context_len=llm_config.export.max_context_length,
    )

    executorch_program = export_all(llava_model)

    # memory profiling
    if llm_config.debug.profile_memory:
        for method_name in executorch_program.methods:
            generate_memory_trace(