| 200 | |
| 201 | |
| 202 | def export_all(llava_model: LlavaModel): |
| 203 | llava = llava_model.get_eager_model() |
| 204 | |
| 205 | ( |
| 206 | prompt_before_image, |
| 207 | resized, |
| 208 | prompt_after_image, |
| 209 | ) = llava_model.get_inputs_for_prefill() |
| 210 | |
| 211 | image_encoder_ep = export_image_encoder( |
| 212 | llava, resized, llava_model._get_image_dynamic_shapes() |
| 213 | ) |
| 214 | |
| 215 | embeddings = llava.prefill_embedding( |
| 216 | prompt_before_image, resized, prompt_after_image |
| 217 | ) |
| 218 | |
| 219 | text_model_ep = export_text_model( |
| 220 | llava, embeddings, llava_model._get_prompt_dynamic_shapes() |
| 221 | ) |
| 222 | |
| 223 | token_embedding_ep = export_token_embedding(llava, prompt_before_image) |
| 224 | |
| 225 | lowered_and_edge = to_edge_transform_and_lower( |
| 226 | { |
| 227 | "vision_encoder": image_encoder_ep, |
| 228 | "token_embedding": token_embedding_ep, |
| 229 | "text_decoder": text_model_ep, |
| 230 | }, |
| 231 | partitioner={ |
| 232 | "vision_encoder": [XnnpackPartitioner()], |
| 233 | "text_decoder": [ |
| 234 | # First partition the DQLinear nodes, then partition the rest of the nodes, |
| 235 | # to avoid multiple DQLinear nodes in the same partition, |
| 236 | # to avoid holding multiple unpacked and packed weight buffers in memory, |
| 237 | # to reduce peak memory footprint. |
| 238 | XnnpackPartitioner( |
| 239 | config_precisions=ConfigPrecisionType.DYNAMIC_QUANT, |
| 240 | per_op_mode=True, |
| 241 | ), |
| 242 | XnnpackPartitioner(), |
| 243 | ], |
| 244 | }, |
| 245 | constant_methods={"get_max_seq_len": llava_model.max_seq_len}, |
| 246 | compile_config=EdgeCompileConfig(_check_ir_validity=False), |
| 247 | ) |
| 248 | |
| 249 | executorch_program = lowered_and_edge.to_executorch( |
| 250 | ExecutorchBackendConfig( |
| 251 | extract_delegate_segments=True, |
| 252 | passes=[ |
| 253 | QuantFusionPass(), |
| 254 | ], |
| 255 | memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), |
| 256 | sym_shape_eval_pass={ |
| 257 | "vision_encoder": ConstraintBasedSymShapeEvalPass(), |
| 258 | "text_decoder": ConstraintBasedSymShapeEvalPass(), |
| 259 | "token_embedding": HintBasedSymShapeEvalPass(), |