()
| 134 | |
| 135 | |
| 136 | def main(): |
| 137 | parser = argparse.ArgumentParser(description='Convert LuxTTS weights to safetensors') |
| 138 | parser.add_argument('--model-dir', type=Path, required=True) |
| 139 | parser.add_argument('--output-dir', type=Path, required=True) |
| 140 | args = parser.parse_args() |
| 141 | |
| 142 | args.output_dir.mkdir(parents=True, exist_ok=True) |
| 143 | |
| 144 | # Convert main model |
| 145 | model_pt = args.model_dir / 'model.pt' |
| 146 | if model_pt.exists(): |
| 147 | convert_model(model_pt, args.output_dir / 'model.safetensors') |
| 148 | |
| 149 | # Convert vocoder |
| 150 | vocos_bin = args.model_dir / 'vocoder' / 'vocos.bin' |
| 151 | if vocos_bin.exists(): |
| 152 | convert_vocoder(vocos_bin, args.output_dir / 'vocos.safetensors') |
| 153 | |
| 154 | # Copy config, tokens, and vocoder config |
| 155 | for src, dst_name in [ |
| 156 | (args.model_dir / 'config.json', 'config.json'), |
| 157 | (args.model_dir / 'tokens.txt', 'tokens.txt'), |
| 158 | (args.model_dir / 'vocoder' / 'config.yaml', 'vocoder_config.yaml'), |
| 159 | ]: |
| 160 | if src.exists(): |
| 161 | shutil.copy2(src, args.output_dir / dst_name) |
| 162 | print(f"Copied {dst_name}") |
| 163 | |
| 164 | # Add architectures to config.json |
| 165 | config_path = args.output_dir / 'config.json' |
| 166 | if config_path.exists(): |
| 167 | with open(config_path) as f: |
| 168 | config = json.load(f) |
| 169 | if 'architectures' not in config: |
| 170 | config['architectures'] = ['LuxTTSForTextToSpeech'] |
| 171 | # Add feature extraction params from vocoder config |
| 172 | if 'n_fft' not in config.get('feature', {}): |
| 173 | config['feature']['n_fft'] = 1024 |
| 174 | config['feature']['hop_length'] = 256 |
| 175 | config['feature']['n_mels'] = 100 |
| 176 | config['feature']['sample_rate'] = config['feature'].pop('sampling_rate', 24000) |
| 177 | with open(config_path, 'w') as f: |
| 178 | json.dump(config, f, indent=2) |
| 179 | print("Updated config.json with architectures and feature params") |
| 180 | |
| 181 | # Generate index file |
| 182 | model_st = args.output_dir / 'model.safetensors' |
| 183 | if model_st.exists(): |
| 184 | from safetensors import safe_open |
| 185 | with safe_open(str(model_st), framework='pt') as f: |
| 186 | keys = list(f.keys()) |
| 187 | index = { |
| 188 | "metadata": {"total_size": model_st.stat().st_size}, |
| 189 | "weight_map": {k: "model.safetensors" for k in keys} |
| 190 | } |
| 191 | with open(args.output_dir / 'model.safetensors.index.json', 'w') as f: |
| 192 | json.dump(index, f, indent=2) |
| 193 | print(f"Generated index with {len(keys)} entries") |
no test coverage detected