supertonic2-qualcomm-quantized / generate_calibration_data.py

Claude Opus 4.6

Initial commit: Supertonic2 QNN quantized TTS pipeline

fed42c7 9 months ago

17.8 kB

#!/usr/bin/env python3
"""
Generate Calibration Data for Supertonic-2 ONNX Models
Creates representative input/output pairs for:
1. duration_predictor.onnx
2. text_encoder.onnx
3. vector_estimator.onnx (multiple denoising steps)
4. vocoder.onnx
Used for:
- Quantization calibration (QNN/SNPE conversion)
- Accuracy validation after porting to QCS6490
- Performance profiling
Output structure:
calibration_data/
├── sample_001/
│   ├── metadata.json
│   ├── duration_predictor/
│   │   ├── input_text_ids.raw
│   │   ├── input_style_dp.raw
│   │   ├── input_text_mask.raw
│   │   └── output_duration.raw
│   ├── text_encoder/
│   │   ├── input_text_ids.raw
│   │   ├── input_style_ttl.raw
│   │   ├── input_text_mask.raw
│   │   └── output_text_emb.raw
│   ├── vector_estimator/
│   │   ├── step_000/
│   │   │   ├── input_noisy_latent.raw
│   │   │   ├── input_text_emb.raw
│   │   │   ├── input_style_ttl.raw
│   │   │   ├── input_text_mask.raw
│   │   │   ├── input_latent_mask.raw
│   │   │   ├── input_current_step.raw
│   │   │   ├── input_total_step.raw
│   │   │   └── output_latent.raw
│   │   ├── step_001/...
│   │   └── step_N/...
│   └── vocoder/
│       ├── input_latent.raw
│       └── output_waveform.raw
└── sample_002/...
"""

import argparse
import json
import shutil
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np

from supertonic_inference import SupertonicTTS


# Calibration test cases with varying complexity
# All English, F1 voice, 10 diffusion steps
CALIBRATION_SAMPLES = [
    # Short utterances (5-10 words)
    {
        "text": "Hello world, this is a test.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 42,
        "category": "short"
    },
    {
        "text": "Good morning everyone.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 100,
        "category": "short"
    },
    {
        "text": "The weather is nice today.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 150,
        "category": "short"
    },

    # Medium utterances (15-25 words)
    {
        "text": "The quick brown fox jumps over the lazy dog. This sentence contains every letter of the alphabet.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 200,
        "category": "medium"
    },
    {
        "text": "Machine learning models require careful calibration to ensure accuracy on edge devices.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 300,
        "category": "medium"
    },
    {
        "text": "Text to speech synthesis has improved dramatically with the advent of neural networks and deep learning.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 350,
        "category": "medium"
    },

    # Longer utterances (35-50 words)
    {
        "text": "Artificial intelligence has revolutionized many aspects of our daily lives, from smartphone assistants to recommendation systems. The deployment of neural networks on edge devices requires careful optimization, including quantization and calibration, to maintain accuracy while meeting strict latency and power constraints.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 400,
        "category": "long"
    },
    {
        "text": "Qualcomm's Hexagon processor provides dedicated hardware acceleration for neural network inference. By leveraging the tensor processing units and optimized kernels, developers can achieve real-time performance for complex models like text-to-speech systems on mobile and edge devices.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 450,
        "category": "long"
    },

    # Edge cases and special content
    {
        "text": "Numbers: 1, 2, 3, 4, 5. Dates: January 1st, 2024. Time: 3:45 PM.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 500,
        "category": "numbers"
    },
    {
        "text": "Dr. Smith's presentation at 9:00 AM covered topics including: AI, ML, NLP, and TTS. The conference runs from Mon. to Fri.",
        "voice": "F1",
        "lang": "en",
        "steps": 10,
        "seed": 600,
        "category": "abbreviations"
    },
]


class CalibrationDataGenerator:
    """Generate and save calibration data for all 4 ONNX models"""

    def __init__(self, model_dir: str = "model/onnx", output_dir: str = "calibration_data"):
        self.tts = SupertonicTTS(model_dir=model_dir)
        self.output_dir = Path(output_dir)

    def save_tensor(self, filepath: Path, tensor: np.ndarray):
        """Save tensor as raw binary file with shape metadata"""
        # Save raw data - preserve original dtype
        tensor.tofile(filepath)

        # Save shape metadata as JSON
        shape_file = filepath.with_suffix('.shape.json')
        metadata = {
            'shape': list(tensor.shape),
            'dtype': str(tensor.dtype),
            'size_bytes': tensor.nbytes
        }
        with open(shape_file, 'w') as f:
            json.dump(metadata, f, indent=2)

    def generate_sample(self, sample_config: Dict, sample_idx: int):
        """Generate calibration data for one sample"""

        sample_dir = self.output_dir / f"sample_{sample_idx:03d}"
        sample_dir.mkdir(parents=True, exist_ok=True)

        # Save sample metadata
        metadata_file = sample_dir / "metadata.json"
        with open(metadata_file, 'w') as f:
            json.dump(sample_config, f, indent=2)

        print(f"\n{'='*80}")
        print(f"Sample {sample_idx:03d}: {sample_config['category']}")
        print(f"Text: '{sample_config['text'][:60]}...'")
        print(f"Voice: {sample_config['voice']} | Lang: {sample_config['lang']} | Steps: {sample_config['steps']}")
        print(f"{'='*80}")

        # Load voice style
        style_ttl, style_dp = self.tts.load_voice_style(sample_config['voice'])

        # Step 1: Text to IDs (preprocessing)
        print("[1/4] Text preprocessing...")
        text_ids, text_mask = self.tts.text_to_ids(sample_config['text'], sample_config['lang'])
        original_len = text_ids.shape[1]
        print(f"  Tokens: {original_len}")

        # Pad/truncate to fixed length 128 for QNN conversion
        max_len = 128
        if original_len > max_len:
            print(f"  WARNING: Truncating from {original_len} to {max_len}")
            text_ids = text_ids[:, :max_len]
            text_mask = text_mask[:, :, :max_len]
        elif original_len < max_len:
            # Pad text_ids with zeros
            padding = np.zeros((1, max_len - original_len), dtype=text_ids.dtype)
            text_ids = np.concatenate([text_ids, padding], axis=1)
            # Pad text_mask with zeros
            padding_mask = np.zeros((1, 1, max_len - original_len), dtype=text_mask.dtype)
            text_mask = np.concatenate([text_mask, padding_mask], axis=2)
            print(f"  Padded from {original_len} to {max_len}")

        # Step 2: Duration Predictor
        print("[2/4] Duration predictor...")
        dp_dir = sample_dir / "duration_predictor"
        dp_dir.mkdir(exist_ok=True)

        # Save inputs
        self.save_tensor(dp_dir / "input_text_ids.raw", text_ids)
        self.save_tensor(dp_dir / "input_style_dp.raw", style_dp)
        self.save_tensor(dp_dir / "input_text_mask.raw", text_mask)

        # Run model and save output
        duration_raw = self.tts.duration_predictor.run(None, {
            "text_ids": text_ids,
            "style_dp": style_dp,
            "text_mask": text_mask
        })[0]

        self.save_tensor(dp_dir / "output_duration.raw", duration_raw)
        print(f"  Duration: {duration_raw[0]:.2f}s")

        # Step 3: Text Encoder
        print("[3/4] Text encoder...")
        te_dir = sample_dir / "text_encoder"
        te_dir.mkdir(exist_ok=True)

        # Save inputs
        self.save_tensor(te_dir / "input_text_ids.raw", text_ids)
        self.save_tensor(te_dir / "input_style_ttl.raw", style_ttl)
        self.save_tensor(te_dir / "input_text_mask.raw", text_mask)

        # Run model and save output
        text_emb = self.tts.text_encoder.run(None, {
            "text_ids": text_ids,
            "style_ttl": style_ttl,
            "text_mask": text_mask
        })[0]

        self.save_tensor(te_dir / "output_text_emb.raw", text_emb)
        print(f"  Text embedding: {text_emb.shape}")

        # Step 4: Vector Estimator (multiple steps)
        print("[4/4] Vector estimator (flow matching)...")
        ve_dir = sample_dir / "vector_estimator"
        ve_dir.mkdir(exist_ok=True)

        # Initialize latent noise
        duration = duration_raw / sample_config.get('speed', 1.0)
        wav_length = int(duration[0] * self.tts.sample_rate)
        chunk_size = self.tts.base_chunk_size * self.tts.chunk_compress_factor
        latent_len = (wav_length + chunk_size - 1) // chunk_size
        latent_dim = self.tts.latent_dim * self.tts.chunk_compress_factor

        # Set seed for reproducibility
        np.random.seed(sample_config['seed'])
        noisy_latent = np.random.randn(1, latent_dim, latent_len).astype(np.float32)

        # Create latent mask
        latent_length = np.array([latent_len], dtype=np.int64)
        latent_mask_ids = np.arange(latent_len) < latent_length[:, None]
        latent_mask = latent_mask_ids.astype(np.float32).reshape(1, 1, -1)
        noisy_latent = noisy_latent * latent_mask

        # Pad/truncate latent to fixed length 192 for QNN conversion
        max_latent_len = 192
        original_latent_len = latent_len
        if latent_len > max_latent_len:
            print(f"  WARNING: Truncating latent from {latent_len} to {max_latent_len}")
            noisy_latent = noisy_latent[:, :, :max_latent_len]
            latent_mask = latent_mask[:, :, :max_latent_len]
            latent_len = max_latent_len
        elif latent_len < max_latent_len:
            # Pad noisy_latent with zeros
            padding = np.zeros((1, latent_dim, max_latent_len - latent_len), dtype=noisy_latent.dtype)
            noisy_latent = np.concatenate([noisy_latent, padding], axis=2)
            # Pad latent_mask with zeros
            padding_mask = np.zeros((1, 1, max_latent_len - latent_len), dtype=latent_mask.dtype)
            latent_mask = np.concatenate([latent_mask, padding_mask], axis=2)
            print(f"  Padded latent from {original_latent_len} to {max_latent_len}")
            latent_len = max_latent_len

        # Diffusion loop - save each step
        diffusion_steps = sample_config['steps']
        total_step = np.array([diffusion_steps], dtype=np.float32)

        for step in range(diffusion_steps):
            step_dir = ve_dir / f"step_{step:03d}"
            step_dir.mkdir(exist_ok=True)

            current_step = np.array([step], dtype=np.float32)

            # Save inputs for this step
            self.save_tensor(step_dir / "input_noisy_latent.raw", noisy_latent)
            self.save_tensor(step_dir / "input_text_emb.raw", text_emb)
            self.save_tensor(step_dir / "input_style_ttl.raw", style_ttl)
            self.save_tensor(step_dir / "input_text_mask.raw", text_mask)
            self.save_tensor(step_dir / "input_latent_mask.raw", latent_mask)
            self.save_tensor(step_dir / "input_current_step.raw", current_step)
            self.save_tensor(step_dir / "input_total_step.raw", total_step)

            # Run model
            noisy_latent = self.tts.vector_estimator.run(None, {
                "noisy_latent": noisy_latent,
                "text_emb": text_emb,
                "style_ttl": style_ttl,
                "text_mask": text_mask,
                "latent_mask": latent_mask,
                "current_step": current_step,
                "total_step": total_step
            })[0]

            # Save output
            self.save_tensor(step_dir / "output_latent.raw", noisy_latent)

            if (step + 1) % 5 == 0 or step == diffusion_steps - 1:
                print(f"  Step {step + 1}/{diffusion_steps}")

        # Step 5: Vocoder
        print("[5/5] Vocoder...")
        voc_dir = sample_dir / "vocoder"
        voc_dir.mkdir(exist_ok=True)

        # Save input (final denoised latent)
        self.save_tensor(voc_dir / "input_latent.raw", noisy_latent)

        # Run model and save output
        wav = self.tts.vocoder.run(None, {"latent": noisy_latent})[0]
        wav_trimmed = wav[0, :wav_length]

        self.save_tensor(voc_dir / "output_waveform.raw", wav_trimmed)
        print(f"  Waveform: {len(wav_trimmed)} samples @ {self.tts.sample_rate} Hz")
        print(f"  Duration: {len(wav_trimmed)/self.tts.sample_rate:.2f}s")

        print(f"\n✓ Sample {sample_idx:03d} complete: {sample_dir}")

    def generate_all(self, samples: List[Dict] = None):
        """Generate calibration data for all samples"""

        if samples is None:
            samples = CALIBRATION_SAMPLES

        # Clear existing calibration data
        if self.output_dir.exists():
            print(f"Clearing existing calibration data at {self.output_dir}...")
            shutil.rmtree(self.output_dir)

        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Generate summary
        summary = {
            "num_samples": len(samples),
            "samples": [],
            "model_info": {
                "sample_rate": self.tts.sample_rate,
                "base_chunk_size": self.tts.base_chunk_size,
                "chunk_compress_factor": self.tts.chunk_compress_factor,
                "latent_dim": self.tts.latent_dim,
                "effective_chunk_size": self.tts.base_chunk_size * self.tts.chunk_compress_factor
            }
        }

        # Generate each sample
        for idx, sample_config in enumerate(samples, start=1):
            try:
                self.generate_sample(sample_config, idx)
                summary["samples"].append({
                    "sample_id": f"sample_{idx:03d}",
                    "text": sample_config["text"],
                    "voice": sample_config["voice"],
                    "lang": sample_config["lang"],
                    "category": sample_config["category"],
                    "diffusion_steps": sample_config["steps"],
                    "seed": sample_config["seed"]
                })
            except Exception as e:
                print(f"\n✗ Error generating sample {idx:03d}: {e}")
                import traceback
                traceback.print_exc()
                continue

        # Save summary
        summary_file = self.output_dir / "calibration_summary.json"
        with open(summary_file, 'w') as f:
            json.dump(summary, f, indent=2)

        print(f"\n{'='*80}")
        print(f"✓ Calibration data generation complete!")
        print(f"  Total samples: {len(summary['samples'])}")
        print(f"  Output directory: {self.output_dir}")
        print(f"  Summary: {summary_file}")
        print(f"{'='*80}\n")

        # Print statistics
        print("Sample Statistics:")
        print(f"  Short utterances:  {sum(1 for s in summary['samples'] if 'short' in s['category'])}")
        print(f"  Medium utterances: {sum(1 for s in summary['samples'] if 'medium' in s['category'])}")
        print(f"  Long utterances:   {sum(1 for s in summary['samples'] if 'long' in s['category'])}")
        print(f"  English samples:   {sum(1 for s in summary['samples'] if s['lang'] == 'en')}")
        print(f"  Spanish samples:   {sum(1 for s in summary['samples'] if s['lang'] == 'es')}")
        print(f"  Korean samples:    {sum(1 for s in summary['samples'] if s['lang'] == 'ko')}")
        print(f"  Unique voices:     {len(set(s['voice'] for s in summary['samples']))}")

        return summary


def main():
    parser = argparse.ArgumentParser(
        description="Generate calibration data for Supertonic-2 ONNX models",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Usage Examples:
  # Generate all calibration samples
  python generate_calibration_data.py
  # Custom output directory
  python generate_calibration_data.py --output-dir /path/to/calibration
  # Use custom model directory
  python generate_calibration_data.py --model-dir /path/to/onnx/models
Output Structure:
  calibration_data/
  ├── sample_001/
  │   ├── metadata.json
  │   ├── duration_predictor/ (inputs + output)
  │   ├── text_encoder/ (inputs + output)
  │   ├── vector_estimator/
  │   │   ├── step_000/ (inputs + output)
  │   │   ├── step_001/
  │   │   └── ...
  │   └── vocoder/ (input + output)
  └── calibration_summary.json
Files are saved as:
  - .raw: Binary float32 data
  - .shape.json: Shape and dtype metadata
"""
    )

    parser.add_argument(
        "--model-dir",
        type=str,
        default="model/onnx",
        help="Path to ONNX models directory (default: model/onnx)"
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="calibration_data",
        help="Output directory for calibration data (default: calibration_data)"
    )

    args = parser.parse_args()

    # Generate calibration data
    generator = CalibrationDataGenerator(
        model_dir=args.model_dir,
        output_dir=args.output_dir
    )

    generator.generate_all()


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	Generate Calibration Data for Supertonic-2 ONNX Models

	Creates representative input/output pairs for:
	1. duration_predictor.onnx
	2. text_encoder.onnx
	3. vector_estimator.onnx (multiple denoising steps)
	4. vocoder.onnx

	Used for:
	- Quantization calibration (QNN/SNPE conversion)
	- Accuracy validation after porting to QCS6490
	- Performance profiling

	Output structure:
	calibration_data/
	├── sample_001/
	│ ├── metadata.json
	│ ├── duration_predictor/
	│ │ ├── input_text_ids.raw
	│ │ ├── input_style_dp.raw
	│ │ ├── input_text_mask.raw
	│ │ └── output_duration.raw
	│ ├── text_encoder/
	│ │ ├── input_text_ids.raw
	│ │ ├── input_style_ttl.raw
	│ │ ├── input_text_mask.raw
	│ │ └── output_text_emb.raw
	│ ├── vector_estimator/
	│ │ ├── step_000/
	│ │ │ ├── input_noisy_latent.raw
	│ │ │ ├── input_text_emb.raw
	│ │ │ ├── input_style_ttl.raw
	│ │ │ ├── input_text_mask.raw
	│ │ │ ├── input_latent_mask.raw
	│ │ │ ├── input_current_step.raw
	│ │ │ ├── input_total_step.raw
	│ │ │ └── output_latent.raw
	│ │ ├── step_001/...
	│ │ └── step_N/...
	│ └── vocoder/
	│ ├── input_latent.raw
	│ └── output_waveform.raw
	└── sample_002/...
	"""

	import argparse
	import json
	import shutil
	from pathlib import Path
	from typing import Dict, List, Tuple

	import numpy as np

	from supertonic_inference import SupertonicTTS


	# Calibration test cases with varying complexity
	# All English, F1 voice, 10 diffusion steps
	CALIBRATION_SAMPLES = [
	# Short utterances (5-10 words)
	{
	"text": "Hello world, this is a test.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 42,
	"category": "short"
	},
	{
	"text": "Good morning everyone.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 100,
	"category": "short"
	},
	{
	"text": "The weather is nice today.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 150,
	"category": "short"
	},

	# Medium utterances (15-25 words)
	{
	"text": "The quick brown fox jumps over the lazy dog. This sentence contains every letter of the alphabet.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 200,
	"category": "medium"
	},
	{
	"text": "Machine learning models require careful calibration to ensure accuracy on edge devices.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 300,
	"category": "medium"
	},
	{
	"text": "Text to speech synthesis has improved dramatically with the advent of neural networks and deep learning.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 350,
	"category": "medium"
	},

	# Longer utterances (35-50 words)
	{
	"text": "Artificial intelligence has revolutionized many aspects of our daily lives, from smartphone assistants to recommendation systems. The deployment of neural networks on edge devices requires careful optimization, including quantization and calibration, to maintain accuracy while meeting strict latency and power constraints.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 400,
	"category": "long"
	},
	{
	"text": "Qualcomm's Hexagon processor provides dedicated hardware acceleration for neural network inference. By leveraging the tensor processing units and optimized kernels, developers can achieve real-time performance for complex models like text-to-speech systems on mobile and edge devices.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 450,
	"category": "long"
	},

	# Edge cases and special content
	{
	"text": "Numbers: 1, 2, 3, 4, 5. Dates: January 1st, 2024. Time: 3:45 PM.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 500,
	"category": "numbers"
	},
	{
	"text": "Dr. Smith's presentation at 9:00 AM covered topics including: AI, ML, NLP, and TTS. The conference runs from Mon. to Fri.",
	"voice": "F1",
	"lang": "en",
	"steps": 10,
	"seed": 600,
	"category": "abbreviations"
	},
	]


	class CalibrationDataGenerator:
	"""Generate and save calibration data for all 4 ONNX models"""

	def __init__(self, model_dir: str = "model/onnx", output_dir: str = "calibration_data"):
	self.tts = SupertonicTTS(model_dir=model_dir)
	self.output_dir = Path(output_dir)

	def save_tensor(self, filepath: Path, tensor: np.ndarray):
	"""Save tensor as raw binary file with shape metadata"""
	# Save raw data - preserve original dtype
	tensor.tofile(filepath)

	# Save shape metadata as JSON
	shape_file = filepath.with_suffix('.shape.json')
	metadata = {
	'shape': list(tensor.shape),
	'dtype': str(tensor.dtype),
	'size_bytes': tensor.nbytes
	}
	with open(shape_file, 'w') as f:
	json.dump(metadata, f, indent=2)

	def generate_sample(self, sample_config: Dict, sample_idx: int):
	"""Generate calibration data for one sample"""

	sample_dir = self.output_dir / f"sample_{sample_idx:03d}"
	sample_dir.mkdir(parents=True, exist_ok=True)

	# Save sample metadata
	metadata_file = sample_dir / "metadata.json"
	with open(metadata_file, 'w') as f:
	json.dump(sample_config, f, indent=2)

	print(f"\n{'='*80}")
	print(f"Sample {sample_idx:03d}: {sample_config['category']}")
	print(f"Text: '{sample_config['text'][:60]}...'")
	print(f"Voice: {sample_config['voice']} \| Lang: {sample_config['lang']} \| Steps: {sample_config['steps']}")
	print(f"{'='*80}")

	# Load voice style
	style_ttl, style_dp = self.tts.load_voice_style(sample_config['voice'])

	# Step 1: Text to IDs (preprocessing)
	print("[1/4] Text preprocessing...")
	text_ids, text_mask = self.tts.text_to_ids(sample_config['text'], sample_config['lang'])
	original_len = text_ids.shape[1]
	print(f" Tokens: {original_len}")

	# Pad/truncate to fixed length 128 for QNN conversion
	max_len = 128
	if original_len > max_len:
	print(f" WARNING: Truncating from {original_len} to {max_len}")
	text_ids = text_ids[:, :max_len]
	text_mask = text_mask[:, :, :max_len]
	elif original_len < max_len:
	# Pad text_ids with zeros
	padding = np.zeros((1, max_len - original_len), dtype=text_ids.dtype)
	text_ids = np.concatenate([text_ids, padding], axis=1)
	# Pad text_mask with zeros
	padding_mask = np.zeros((1, 1, max_len - original_len), dtype=text_mask.dtype)
	text_mask = np.concatenate([text_mask, padding_mask], axis=2)
	print(f" Padded from {original_len} to {max_len}")

	# Step 2: Duration Predictor
	print("[2/4] Duration predictor...")
	dp_dir = sample_dir / "duration_predictor"
	dp_dir.mkdir(exist_ok=True)

	# Save inputs
	self.save_tensor(dp_dir / "input_text_ids.raw", text_ids)
	self.save_tensor(dp_dir / "input_style_dp.raw", style_dp)
	self.save_tensor(dp_dir / "input_text_mask.raw", text_mask)

	# Run model and save output
	duration_raw = self.tts.duration_predictor.run(None, {
	"text_ids": text_ids,
	"style_dp": style_dp,
	"text_mask": text_mask
	})[0]

	self.save_tensor(dp_dir / "output_duration.raw", duration_raw)
	print(f" Duration: {duration_raw[0]:.2f}s")

	# Step 3: Text Encoder
	print("[3/4] Text encoder...")
	te_dir = sample_dir / "text_encoder"
	te_dir.mkdir(exist_ok=True)

	# Save inputs
	self.save_tensor(te_dir / "input_text_ids.raw", text_ids)
	self.save_tensor(te_dir / "input_style_ttl.raw", style_ttl)
	self.save_tensor(te_dir / "input_text_mask.raw", text_mask)

	# Run model and save output
	text_emb = self.tts.text_encoder.run(None, {
	"text_ids": text_ids,
	"style_ttl": style_ttl,
	"text_mask": text_mask
	})[0]

	self.save_tensor(te_dir / "output_text_emb.raw", text_emb)
	print(f" Text embedding: {text_emb.shape}")

	# Step 4: Vector Estimator (multiple steps)
	print("[4/4] Vector estimator (flow matching)...")
	ve_dir = sample_dir / "vector_estimator"
	ve_dir.mkdir(exist_ok=True)

	# Initialize latent noise
	duration = duration_raw / sample_config.get('speed', 1.0)
	wav_length = int(duration[0] * self.tts.sample_rate)
	chunk_size = self.tts.base_chunk_size * self.tts.chunk_compress_factor
	latent_len = (wav_length + chunk_size - 1) // chunk_size
	latent_dim = self.tts.latent_dim * self.tts.chunk_compress_factor

	# Set seed for reproducibility
	np.random.seed(sample_config['seed'])
	noisy_latent = np.random.randn(1, latent_dim, latent_len).astype(np.float32)

	# Create latent mask
	latent_length = np.array([latent_len], dtype=np.int64)
	latent_mask_ids = np.arange(latent_len) < latent_length[:, None]
	latent_mask = latent_mask_ids.astype(np.float32).reshape(1, 1, -1)
	noisy_latent = noisy_latent * latent_mask

	# Pad/truncate latent to fixed length 192 for QNN conversion
	max_latent_len = 192
	original_latent_len = latent_len
	if latent_len > max_latent_len:
	print(f" WARNING: Truncating latent from {latent_len} to {max_latent_len}")
	noisy_latent = noisy_latent[:, :, :max_latent_len]
	latent_mask = latent_mask[:, :, :max_latent_len]
	latent_len = max_latent_len
	elif latent_len < max_latent_len:
	# Pad noisy_latent with zeros
	padding = np.zeros((1, latent_dim, max_latent_len - latent_len), dtype=noisy_latent.dtype)
	noisy_latent = np.concatenate([noisy_latent, padding], axis=2)
	# Pad latent_mask with zeros
	padding_mask = np.zeros((1, 1, max_latent_len - latent_len), dtype=latent_mask.dtype)
	latent_mask = np.concatenate([latent_mask, padding_mask], axis=2)
	print(f" Padded latent from {original_latent_len} to {max_latent_len}")
	latent_len = max_latent_len

	# Diffusion loop - save each step
	diffusion_steps = sample_config['steps']
	total_step = np.array([diffusion_steps], dtype=np.float32)

	for step in range(diffusion_steps):
	step_dir = ve_dir / f"step_{step:03d}"
	step_dir.mkdir(exist_ok=True)

	current_step = np.array([step], dtype=np.float32)

	# Save inputs for this step
	self.save_tensor(step_dir / "input_noisy_latent.raw", noisy_latent)
	self.save_tensor(step_dir / "input_text_emb.raw", text_emb)
	self.save_tensor(step_dir / "input_style_ttl.raw", style_ttl)
	self.save_tensor(step_dir / "input_text_mask.raw", text_mask)
	self.save_tensor(step_dir / "input_latent_mask.raw", latent_mask)
	self.save_tensor(step_dir / "input_current_step.raw", current_step)
	self.save_tensor(step_dir / "input_total_step.raw", total_step)

	# Run model
	noisy_latent = self.tts.vector_estimator.run(None, {
	"noisy_latent": noisy_latent,
	"text_emb": text_emb,
	"style_ttl": style_ttl,
	"text_mask": text_mask,
	"latent_mask": latent_mask,
	"current_step": current_step,
	"total_step": total_step
	})[0]

	# Save output
	self.save_tensor(step_dir / "output_latent.raw", noisy_latent)

	if (step + 1) % 5 == 0 or step == diffusion_steps - 1:
	print(f" Step {step + 1}/{diffusion_steps}")

	# Step 5: Vocoder
	print("[5/5] Vocoder...")
	voc_dir = sample_dir / "vocoder"
	voc_dir.mkdir(exist_ok=True)

	# Save input (final denoised latent)
	self.save_tensor(voc_dir / "input_latent.raw", noisy_latent)

	# Run model and save output
	wav = self.tts.vocoder.run(None, {"latent": noisy_latent})[0]
	wav_trimmed = wav[0, :wav_length]

	self.save_tensor(voc_dir / "output_waveform.raw", wav_trimmed)
	print(f" Waveform: {len(wav_trimmed)} samples @ {self.tts.sample_rate} Hz")
	print(f" Duration: {len(wav_trimmed)/self.tts.sample_rate:.2f}s")

	print(f"\n✓ Sample {sample_idx:03d} complete: {sample_dir}")

	def generate_all(self, samples: List[Dict] = None):
	"""Generate calibration data for all samples"""

	if samples is None:
	samples = CALIBRATION_SAMPLES

	# Clear existing calibration data
	if self.output_dir.exists():
	print(f"Clearing existing calibration data at {self.output_dir}...")
	shutil.rmtree(self.output_dir)

	self.output_dir.mkdir(parents=True, exist_ok=True)

	# Generate summary
	summary = {
	"num_samples": len(samples),
	"samples": [],
	"model_info": {
	"sample_rate": self.tts.sample_rate,
	"base_chunk_size": self.tts.base_chunk_size,
	"chunk_compress_factor": self.tts.chunk_compress_factor,
	"latent_dim": self.tts.latent_dim,
	"effective_chunk_size": self.tts.base_chunk_size * self.tts.chunk_compress_factor
	}
	}

	# Generate each sample
	for idx, sample_config in enumerate(samples, start=1):
	try:
	self.generate_sample(sample_config, idx)
	summary["samples"].append({
	"sample_id": f"sample_{idx:03d}",
	"text": sample_config["text"],
	"voice": sample_config["voice"],
	"lang": sample_config["lang"],
	"category": sample_config["category"],
	"diffusion_steps": sample_config["steps"],
	"seed": sample_config["seed"]
	})
	except Exception as e:
	print(f"\n✗ Error generating sample {idx:03d}: {e}")
	import traceback
	traceback.print_exc()
	continue

	# Save summary
	summary_file = self.output_dir / "calibration_summary.json"
	with open(summary_file, 'w') as f:
	json.dump(summary, f, indent=2)

	print(f"\n{'='*80}")
	print(f"✓ Calibration data generation complete!")
	print(f" Total samples: {len(summary['samples'])}")
	print(f" Output directory: {self.output_dir}")
	print(f" Summary: {summary_file}")
	print(f"{'='*80}\n")

	# Print statistics
	print("Sample Statistics:")
	print(f" Short utterances: {sum(1 for s in summary['samples'] if 'short' in s['category'])}")
	print(f" Medium utterances: {sum(1 for s in summary['samples'] if 'medium' in s['category'])}")
	print(f" Long utterances: {sum(1 for s in summary['samples'] if 'long' in s['category'])}")
	print(f" English samples: {sum(1 for s in summary['samples'] if s['lang'] == 'en')}")
	print(f" Spanish samples: {sum(1 for s in summary['samples'] if s['lang'] == 'es')}")
	print(f" Korean samples: {sum(1 for s in summary['samples'] if s['lang'] == 'ko')}")
	print(f" Unique voices: {len(set(s['voice'] for s in summary['samples']))}")

	return summary


	def main():
	parser = argparse.ArgumentParser(
	description="Generate calibration data for Supertonic-2 ONNX models",
	formatter_class=argparse.RawDescriptionHelpFormatter,
	epilog="""
	Usage Examples:
	# Generate all calibration samples
	python generate_calibration_data.py

	# Custom output directory
	python generate_calibration_data.py --output-dir /path/to/calibration

	# Use custom model directory
	python generate_calibration_data.py --model-dir /path/to/onnx/models

	Output Structure:
	calibration_data/
	├── sample_001/
	│ ├── metadata.json
	│ ├── duration_predictor/ (inputs + output)
	│ ├── text_encoder/ (inputs + output)
	│ ├── vector_estimator/
	│ │ ├── step_000/ (inputs + output)
	│ │ ├── step_001/
	│ │ └── ...
	│ └── vocoder/ (input + output)
	└── calibration_summary.json

	Files are saved as:
	- .raw: Binary float32 data
	- .shape.json: Shape and dtype metadata
	"""
	)

	parser.add_argument(
	"--model-dir",
	type=str,
	default="model/onnx",
	help="Path to ONNX models directory (default: model/onnx)"
	)
	parser.add_argument(
	"--output-dir",
	type=str,
	default="calibration_data",
	help="Output directory for calibration data (default: calibration_data)"
	)

	args = parser.parse_args()

	# Generate calibration data
	generator = CalibrationDataGenerator(
	model_dir=args.model_dir,
	output_dir=args.output_dir
	)

	generator.generate_all()


	if __name__ == "__main__":
	main()