generate_calibration_data.py

#!/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()