(model, output_path)
| 52 | logging.info("Finished Conditioners Model conversion.\n") |
| 53 | |
| 54 | def export_dit(model, output_path) -> None: |
| 55 | dit_model = get_dit_module(model=model) |
| 56 | dit_example_mapping = get_dit_example_input_mapping() |
| 57 | |
| 58 | # Quantize the models' linear layers to int8 per-channel |
| 59 | logging.info("Starting Dit Model conversion...\n") |
| 60 | |
| 61 | from torchao.quantization.granularity import PerAxis, PerGroup |
| 62 | from torchao.quantization.quant_api import ( |
| 63 | Int8DynamicActivationIntxWeightConfig, |
| 64 | quantize_, |
| 65 | ) |
| 66 | from torchao.utils import unwrap_tensor_subclass |
| 67 | |
| 68 | with torch.no_grad(): |
| 69 | quantize_( |
| 70 | dit_model, |
| 71 | Int8DynamicActivationIntxWeightConfig( |
| 72 | weight_dtype=torch.int8, |
| 73 | weight_granularity=PerAxis(0), |
| 74 | ), |
| 75 | ) |
| 76 | dit_model = unwrap_tensor_subclass(dit_model) |
| 77 | |
| 78 | logging.info("quantized model: %s", dit_model) |
| 79 | |
| 80 | # Export the model to ExecuTorch format |
| 81 | exported_program: ExportedProgram = torch.export.export(dit_model, args=(), kwargs=dit_example_mapping, dynamic_shapes=None) |
| 82 | edge: EdgeProgramManager = to_edge_transform_and_lower( |
| 83 | exported_program, |
| 84 | partitioner=[ |
| 85 | XnnpackDynamicallyQuantizedPartitioner(), |
| 86 | XnnpackPartitioner()], |
| 87 | ) |
| 88 | exec_prog = edge.to_executorch() |
| 89 | |
| 90 | with open(os.path.join(output_path, "dit_model.pte"), "wb") as file: |
| 91 | exec_prog.write_to_file(file) |
| 92 | |
| 93 | logging.info("Finished Dit Model conversion.\n") |
| 94 | |
| 95 | def export_autoencoder(model, output_path) -> None: |
| 96 | # Load the AutoEncoder part of the model |
no test coverage detected