Function export

kleidiai-examples/audiogen-et/scripts/export_sao.py:126–150

(args)

    logging.info("Finished AutoEncoder Model conversion.\n")

def export(args) -> None:

    torch.manual_seed(0)
    device = torch.device("cpu")

    # Load the model configuration
    logging.info("Loading the AudioGen Checkpoint...")
    with open(args.model_config, encoding="utf-8") as f:
        model_config = json.load(f)
    model, model_config = load_model(
        model_config,
        args.ckpt_path,
        pretrained_name=None,
        device=device,
    )
    logging.info("Model is loaded...")

    # --------- Conditioners Model ---------
    export_conditioners(model, args.output_path)

    # --------- Dit Model ----------------
    export_dit(model, args.output_path)

    # --------- AutoEncoder Model ---------
    export_autoencoder(model, args.output_path)

def main():
    parser = argparse.ArgumentParser()