(model, output_path, dtype = torch.float)
| 48 | logging.info("Conditioners model has been saved to %s", os.path.abspath(os.path.join(output_path, "conditioners_float32.tflite"))) |
| 49 | |
| 50 | def export_dit(model, output_path, dtype = torch.float) -> None: |
| 51 | |
| 52 | logging.info("Starting DiT Model conversion to LiteRT format...\n") |
| 53 | |
| 54 | with torch.no_grad(): |
| 55 | dit_model = get_dit_module(model=model) |
| 56 | dit_model = dit_model.to(dtype).eval().requires_grad_(False) |
| 57 | dit_model_example_input = get_dit_example_input_mapping(dtype) |
| 58 | |
| 59 | # Create the dynamic weights int8 quantization config |
| 60 | quant_config_audiogen_int8 = quant_config.QuantConfig( |
| 61 | generative_recipe=quant_recipe.GenerativeQuantRecipe( |
| 62 | default=quant_recipe_utils.create_layer_quant_dynamic(), |
| 63 | ) |
| 64 | ) |
| 65 | |
| 66 | # Workaround for some issue in LiteRT that occurs at runtime |
| 67 | rotary_pos_emb_res = ( |
| 68 | dit_model.model.transformer.rotary_pos_emb.forward_from_seq_len(257) |
| 69 | ) |
| 70 | def rotary_emb_const(_): |
| 71 | return rotary_pos_emb_res |
| 72 | dit_model.model.transformer.rotary_pos_emb.forward_from_seq_len = rotary_emb_const |
| 73 | |
| 74 | # Export the DiT to LiteRT format |
| 75 | edge_model = litert_torch.convert( |
| 76 | dit_model, sample_args=None, sample_kwargs=dit_model_example_input, quant_config=quant_config_audiogen_int8 |
| 77 | ) |
| 78 | |
| 79 | edge_model.export(os.path.join(output_path, "dit_model.tflite")) |
| 80 | logging.info("DiT model has been saved to %s", os.path.abspath(os.path.join(output_path, "dit_model.tflite"))) |
| 81 | |
| 82 | def export_autoencoder(model, output_path, dtype = torch.float) -> None: |
| 83 |
no test coverage detected