OmniVoice / ChangeToOnnx.py

Upload 3 files

10e4df2 verified 6 months ago

2.06 kB

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "onnx>=1.21.0",
#     "onnxruntime>=1.26.0",
#     "onnxscript>=0.7.0",
#     "torch>=2.12.0",
#     "torchvision>=0.27.0",
# ]
# ///
from pathlib import Path
import torch

# Load the uploaded PyTorch module
model_path = "embedding.pt"
import onnxruntime as ort


model = torch.load(model_path, map_location="cpu",weights_only=False)
model.eval()

# Create a dummy input for export
dummy_input = torch.randint(
    0,
    #config.num_audio_codebook * config.audio_vocab_size
     8 * 1025 ,
    (1, 8, 128),
    dtype=torch.long
)

# Export to ONNX
onnx_path = "embedding.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,export_params=True,
    #input_names=["input_ids"],
    # output_names=["embeddings"],
    # dynamic_axes={
    #     "input_ids": {0: "batch_size", 1: "sequence_length"},
    #     "embeddings": {0: "batch_size", 1: "sequence_length"},
    # },
    external_data=False,
    opset_version=19,
)

print(f"ONNX model saved to: {onnx_path}")
print(f"Embedding shape: ({model.num_embeddings}, {model.embedding_dim})")
session = ort.InferenceSession(onnx_path)

print(session.get_inputs())
print(session.get_outputs())


###
model_path = "linear_head.pt"
 


model = torch.load(model_path, map_location="cpu",weights_only=False)
model.eval()

# Create a dummy input for export
dummy_input = torch.randn(1, 1024, dtype=torch.float32)

# Export to ONNX
onnx_path = "liner_head.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    export_params=True,
    input_names=["in_features"],
    output_names=["out_features"],
    external_data=False,
    # dynamic_axes={
    #     "input_ids": {0: "batch_size", 1: "sequence_length"},
    #     "embeddings": {0: "batch_size", 1: "sequence_length"},
    # },
    opset_version=19,
)

print(f"ONNX model saved to: {onnx_path}")
print(f"Embedding shape: ({model.in_features}, {model.out_features})")
session = ort.InferenceSession(onnx_path)

print(session.get_inputs())
print(session.get_outputs())

	# /// script
	# requires-python = ">=3.12"
	# dependencies = [
	# "onnx>=1.21.0",
	# "onnxruntime>=1.26.0",
	# "onnxscript>=0.7.0",
	# "torch>=2.12.0",
	# "torchvision>=0.27.0",
	# ]
	# ///
	from pathlib import Path
	import torch

	# Load the uploaded PyTorch module
	model_path = "embedding.pt"
	import onnxruntime as ort


	model = torch.load(model_path, map_location="cpu",weights_only=False)
	model.eval()

	# Create a dummy input for export
	dummy_input = torch.randint(
	0,
	#config.num_audio_codebook * config.audio_vocab_size
	8 * 1025 ,
	(1, 8, 128),
	dtype=torch.long
	)

	# Export to ONNX
	onnx_path = "embedding.onnx"
	torch.onnx.export(
	model,
	dummy_input,
	onnx_path,export_params=True,
	#input_names=["input_ids"],
	# output_names=["embeddings"],
	# dynamic_axes={
	# "input_ids": {0: "batch_size", 1: "sequence_length"},
	# "embeddings": {0: "batch_size", 1: "sequence_length"},
	# },
	external_data=False,
	opset_version=19,
	)

	print(f"ONNX model saved to: {onnx_path}")
	print(f"Embedding shape: ({model.num_embeddings}, {model.embedding_dim})")
	session = ort.InferenceSession(onnx_path)

	print(session.get_inputs())
	print(session.get_outputs())


	###
	model_path = "linear_head.pt"



	model = torch.load(model_path, map_location="cpu",weights_only=False)
	model.eval()

	# Create a dummy input for export
	dummy_input = torch.randn(1, 1024, dtype=torch.float32)

	# Export to ONNX
	onnx_path = "liner_head.onnx"
	torch.onnx.export(
	model,
	dummy_input,
	onnx_path,
	export_params=True,
	input_names=["in_features"],
	output_names=["out_features"],
	external_data=False,
	# dynamic_axes={
	# "input_ids": {0: "batch_size", 1: "sequence_length"},
	# "embeddings": {0: "batch_size", 1: "sequence_length"},
	# },
	opset_version=19,
	)

	print(f"ONNX model saved to: {onnx_path}")
	print(f"Embedding shape: ({model.in_features}, {model.out_features})")
	session = ort.InferenceSession(onnx_path)

	print(session.get_inputs())
	print(session.get_outputs())