File size: 4,048 Bytes
4edc9aa 0254260 4edc9aa 0254260 4edc9aa 0254260 4edc9aa 0254260 4edc9aa 0254260 4edc9aa 0254260 4edc9aa 0254260 4edc9aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import sys
import argparse
from pathlib import Path
import torch
import torch.nn as nn
from omegaconf import OmegaConf
ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(ROOT))
sys.path.append(str(ROOT / 'flow_matching'))
sys.path.append(str(ROOT / 'flow_matching/Matcha-TTS'))
from flow_matching.src.stage1.medarc_architecture import MultiSubjectConvLinearEncoder
from flow_matching.src.stage2.CFM import CFM
def export_stage1(cfg):
print('Exporting Stage 1...')
subjects_list = cfg.get('subjects', [1, 2, 3, 5])
feat_dims = (32,)
model = MultiSubjectConvLinearEncoder(
num_subjects=len(subjects_list),
feat_dims=feat_dims,
**cfg.stage1.model
)
weights_path = ROOT / 'output' / 'debug_run' / 'stage1_best.pt'
if weights_path.exists():
state_dict = torch.load(weights_path, map_location='cpu', weights_only=True)
model.load_state_dict(state_dict, strict=False)
print('Loaded Stage 1 weights.')
model.eval()
dummy_input = [torch.randn(2, 10, 32)]
save_path = 'stage1_model.onnx'
try:
torch.onnx.export(
model,
dummy_input,
save_path,
opset_version=14,
input_names=['features'],
output_names=['embeddings']
)
print(f'Saved {save_path}')
except Exception as e:
print(f'Failed exporting Stage 1: {e}')
def export_stage2(cfg, target_dim=1000):
print('\nExporting Stage 2...')
cfm_params = cfg.stage2.cfm
velocity_net_params = cfg.stage2.velocity_net
source_ve_params = cfg.stage2.source_ve
transport_params = cfg.stage2.transport
model = CFM(
feat_dim=target_dim,
cfm_params=cfm_params,
velocity_net_params=velocity_net_params,
source_ve_params=source_ve_params,
transport_params=transport_params,
)
weights_path = ROOT / 'output' / 'debug_run' / 'stage2_epoch_0.pt'
if weights_path.exists():
state_dict = torch.load(weights_path, map_location='cpu', weights_only=True)
model.load_state_dict(state_dict, strict=False)
print('Loaded Stage 2 weights.')
model.eval()
class Stage2ExportWrapper(nn.Module):
def __init__(self, cfm_model: nn.Module, steps: int):
super().__init__()
self.cfm_model = cfm_model
self.steps = steps
def forward(self, mu):
return self.cfm_model(mu, n_timesteps=self.steps)
export_wrapper = Stage2ExportWrapper(model, steps=10)
export_wrapper.eval()
# Dummy input: (B, C, T)
B, C, T = 1, target_dim, 10
mu = torch.randn(B, C, T)
save_path = 'stage2_model.onnx'
try:
torch.onnx.export(
export_wrapper,
mu,
save_path,
opset_version=14,
input_names=['mu'],
output_names=['output']
)
print(f'Saved {save_path}')
except Exception as e:
print(f'Failed exporting Stage 2: {e}')
if __name__ == '__main__':
config_path = ROOT / 'output' / 'debug_run' / 'config.yaml'
if config_path.exists():
cfg = OmegaConf.load(config_path)
else:
# Dummy config if config not found
cfg = OmegaConf.create(
{
'stage1': {'model': {}},
'stage2': {
'cfm': {'solver': 'euler', 'kld_weight': 3.0, 'kld_target_std': 1.0, 'detach_ut': False},
'velocity_net': {'hidden_dim': 256, 'n_blocks': 2, 'n_heads': 4, 'dropout': 0.1, 'max_seq_len': 2048, 'temporal_attn_layers': 1},
'source_ve': {'depth': 2, 'num_heads': 4, 'num_queries': 8, 'dropout': 0.1, 'use_variational': True},
'transport': {'path_type': 'Linear', 'prediction': 'velocity', 'time_dist_type': 'uniform', 'time_dist_shift': 1.0},
},
}
)
export_stage1(cfg)
export_stage2(cfg, target_dim=1000)
|