# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""
Example script for text-to-audio generation using Stable Audio Open.

This script demonstrates how to generate audio from text prompts using
the Stable Audio Open model with vLLM-Omni.

Usage:
    python text_to_audio.py --prompt "The sound of a dog barking"
    python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0
    python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality"
    python text_to_audio.py --prompt "A soft synth pad" --cache-backend tea_cache
"""

import argparse
import time
from pathlib import Path

import numpy as np
import torch

from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate audio with Stable Audio Open.")
    parser.add_argument(
        "--model",
        default="stabilityai/stable-audio-open-1.0",
        help="Stable Audio model name or local path.",
    )
    parser.add_argument(
        "--prompt",
        default="The sound of a hammer hitting a wooden surface.",
        help="Text prompt for audio generation.",
    )
    parser.add_argument(
        "--negative-prompt",
        default="Low quality.",
        help="Negative prompt for classifier-free guidance.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for deterministic results.",
    )
    parser.add_argument(
        "--guidance-scale",
        type=float,
        default=7.0,
        help="Classifier-free guidance scale.",
    )
    parser.add_argument(
        "--audio-start",
        type=float,
        default=0.0,
        help="Audio start time in seconds.",
    )
    parser.add_argument(
        "--audio-length",
        type=float,
        default=10.0,
        help="Audio length in seconds (max ~47s for stable-audio-open-1.0).",
    )
    parser.add_argument(
        "--num-inference-steps",
        type=int,
        default=100,
        help="Number of denoising steps for the diffusion sampler.",
    )
    parser.add_argument(
        "--num-waveforms",
        type=int,
        default=1,
        help="Number of audio waveforms to generate for the given prompt.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="stable_audio_output.wav",
        help="Path to save the generated audio (WAV format).",
    )
    parser.add_argument(
        "--sample-rate",
        type=int,
        default=44100,
        help="Sample rate for output audio (Stable Audio uses 44100 Hz).",
    )
    parser.add_argument(
        "--cache-backend",
        type=str,
        default=None,
        choices=["tea_cache"],
        help=(
            "Cache backend to use for acceleration. "
            "Stable Audio currently supports 'tea_cache'. "
            "Default: None (no cache acceleration)."
        ),
    )
    parser.add_argument(
        "--tea-cache-rel-l1-thresh",
        type=float,
        default=0.2,
        help="[tea_cache] Threshold for accumulated relative L1 distance.",
    )
    parser.add_argument(
        "--enable-diffusion-pipeline-profiler",
        action="store_true",
        help="Enable diffusion pipeline profiler to display stage durations.",
    )
    parser.add_argument(
        "--enable-cpu-offload", action="store_true", help="Enable ModelWise cpu offloading to save gpu memory"
    )
    parser.add_argument(
        "--enable-layerwise-offload", action="store_true", help="Enable Layerwise cpu offloading to save gpu memory"
    )
    parser.add_argument(
        "--use-hsdp",
        action="store_true",
        help="Enable HSDP for Stable Audio DiT weight sharding.",
    )
    parser.add_argument(
        "--hsdp-shard-size",
        type=int,
        default=1,
        help="Number of GPUs to shard Stable Audio DiT weights across when HSDP is enabled.",
    )
    parser.add_argument(
        "--hsdp-replicate-size",
        type=int,
        default=1,
        help="Number of HSDP replica groups. Default 1 means pure sharding.",
    )
    parser.add_argument(
        "--tensor-parallel-size",
        type=int,
        default=1,
        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
    )
    parser.add_argument(
        "--ulysses-degree",
        type=int,
        default=1,
        help="Number of GPUs used for ulysses sequence parallelism.",
    )
    parser.add_argument(
        "--ulysses-mode",
        type=str,
        default="strict",
        choices=["strict", "advanced_uaa"],
        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
    )
    parser.add_argument(
        "--ring-degree",
        type=int,
        default=1,
        help="Number of GPUs used for ring sequence parallelism.",
    )
    parser.add_argument(
        "--cfg-parallel-size",
        type=int,
        default=1,
        choices=[1, 2],
        help="Number of GPUs used for classifier free guidance parallel size.",
    )
    parser.add_argument(
        "--vae-patch-parallel-size",
        type=int,
        default=1,
        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
    )
    return parser.parse_args()


def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 44100):
    """Save audio data to a WAV file."""
    try:
        import soundfile as sf

        sf.write(output_path, audio_data, sample_rate)
    except ImportError:
        try:
            import scipy.io.wavfile as wav

            # Ensure audio is in the correct format for scipy
            if audio_data.dtype == np.float32 or audio_data.dtype == np.float64:
                # Normalize to int16 range
                audio_data = np.clip(audio_data, -1.0, 1.0)
                audio_data = (audio_data * 32767).astype(np.int16)
            wav.write(output_path, sample_rate, audio_data)
        except ImportError:
            raise ImportError(
                "Either 'soundfile' or 'scipy' is required to save audio files. "
                "Install with: pip install soundfile or pip install scipy"
            )


def main():
    args = parse_args()
    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
    cache_config = None
    if args.cache_backend == "tea_cache":
        cache_config = {
            "rel_l1_thresh": args.tea_cache_rel_l1_thresh,
        }

    print(f"\n{'=' * 60}")
    print("Stable Audio Open - Text-to-Audio Generation")
    print(f"{'=' * 60}")
    print(f"  Model: {args.model}")
    print(f"  Prompt: {args.prompt}")
    print(f"  Negative prompt: {args.negative_prompt}")
    print(f"  Audio length: {args.audio_length}s")
    print(f"  Inference steps: {args.num_inference_steps}")
    print(f"  Guidance scale: {args.guidance_scale}")
    print(f"  Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
    print(f"  ModelWise Offload: {'Enabled' if args.enable_cpu_offload else 'None'}")
    print(f"  LayerWise Offload: {'Enabled' if args.enable_layerwise_offload else 'None'}")
    if args.use_hsdp:
        print(f"  HSDP: enabled (shard_size={args.hsdp_shard_size}, replicate_size={args.hsdp_replicate_size})")
    else:
        print("  HSDP: disabled")
    print(f"  Seed: {args.seed}")
    print(f"{'=' * 60}\n")

    parallel_config = DiffusionParallelConfig(
        use_hsdp=args.use_hsdp,
        hsdp_shard_size=args.hsdp_shard_size,
        hsdp_replicate_size=args.hsdp_replicate_size,
    )

    # Initialize Omni with Stable Audio model
    omni = Omni(
        model=args.model,
        parallel_config=parallel_config,
        cache_backend=args.cache_backend,
        cache_config=cache_config,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
        enable_cpu_offload=args.enable_cpu_offload,
        enable_layerwise_offload=args.enable_layerwise_offload,
    )

    # Calculate audio end time
    audio_end_in_s = args.audio_start + args.audio_length

    # Time profiling for generation
    generation_start = time.perf_counter()

    # Generate audio
    outputs = omni.generate(
        {
            "prompt": args.prompt,
            "negative_prompt": args.negative_prompt,
        },
        OmniDiffusionSamplingParams(
            generator=generator,
            guidance_scale=args.guidance_scale,
            num_inference_steps=args.num_inference_steps,
            num_outputs_per_prompt=args.num_waveforms,
            extra_args={
                "audio_start_in_s": args.audio_start,
                "audio_end_in_s": audio_end_in_s,
            },
        ),
    )

    generation_end = time.perf_counter()
    generation_time = generation_end - generation_start

    print(f"Total generation time: {generation_time:.2f} seconds")

    # Process and save audio
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    suffix = output_path.suffix or ".wav"
    stem = output_path.stem or "stable_audio_output"

    # Extract audio from omni.generate() outputs
    if not outputs:
        raise ValueError("No output generated from omni.generate()")

    output = outputs[0]
    if not hasattr(output, "request_output") or not output.request_output:
        raise ValueError("No request_output found in OmniRequestOutput")
    request_output = output.request_output
    if not hasattr(request_output, "multimodal_output"):
        raise ValueError("No multimodal_output found in request_output")

    audio = request_output.multimodal_output.get("audio")
    if audio is None:
        raise ValueError("No audio output found in request_output")

    # Handle different output formats
    if isinstance(audio, torch.Tensor):
        audio = audio.cpu().float().numpy()

    # Audio shape is typically [batch, channels, samples] or [channels, samples]
    if audio.ndim == 3:
        # [batch, channels, samples]
        if args.num_waveforms <= 1:
            audio_data = audio[0].T  # [samples, channels]
            save_audio(audio_data, str(output_path), args.sample_rate)
            print(f"Saved generated audio to {output_path}")
        else:
            for idx in range(audio.shape[0]):
                audio_data = audio[idx].T  # [samples, channels]
                save_path = output_path.parent / f"{stem}_{idx}{suffix}"
                save_audio(audio_data, str(save_path), args.sample_rate)
                print(f"Saved generated audio to {save_path}")
    elif audio.ndim == 2:
        # [channels, samples]
        audio_data = audio.T  # [samples, channels]
        save_audio(audio_data, str(output_path), args.sample_rate)
        print(f"Saved generated audio to {output_path}")
    else:
        # [samples] - mono audio
        save_audio(audio, str(output_path), args.sample_rate)
        print(f"Saved generated audio to {output_path}")

    print(f"\nGenerated {args.audio_length}s of audio at {args.sample_rate} Hz")


if __name__ == "__main__":
    main()