import spaces
import torch
import numpy as np
import gradio as gr
from diffusers import StableAudioPipeline

# 1. Load pipeline on CPU initialized with FP16
pipe = StableAudioPipeline.from_pretrained(
    "stabilityai/stable-audio-open-1.0", 
    torch_dtype=torch.float16,
    token=True
)

@spaces.GPU(duration=120)
def generate_audio(
    prompt: str,
    negative_prompt: str,
    seconds: float,
    num_steps: int,
    guidance_scale: float,
    seed: int,
    num_variations: int
):
    # Move model to CUDA inside ZeroGPU function
    pipe.to("cuda")

    # Handle seed generation for reproducibility or dynamic variation
    if seed == -1:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    
    generator = torch.Generator(device="cuda").manual_seed(int(seed))

    # Run inference with parameters
    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        audio_start_in_s=0.0,
        audio_end_in_s=seconds,
        num_inference_steps=num_steps,
        guidance_scale=guidance_scale,
        num_waveforms_per_prompt=num_variations,
        generator=generator
    )

    # Convert audio tensor outputs to tuple format for Gradio (Sample Rate, Audio Array)
    audio_outputs = []
    for waveform in output.audios:
        # Transpose tensor from (channels, samples) to (samples, channels)
        audio_data = waveform.T.cpu().numpy()
        audio_outputs.append((44100, audio_data))

    # Pad list to match max 4 variation outputs in Gradio UI
    while len(audio_outputs) < 4:
        audio_outputs.append(None)

    return (*audio_outputs, f"Used Seed: {int(seed)}")

# Build Gradio UI
with gr.Blocks(title="Stable Audio Studio Pro") as demo:
    gr.Markdown("# 🎵 Stable Audio Generator (With Inference & Variation Controls)")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt", 
                value="Heavy metal door slam in an echoing dungeon, crisp transient impact",
                lines=2
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt", 
                value="low quality, distortion, noise, background speech",
                lines=2
            )
            
            # --- Inference Parameters ---
            with gr.Accordion("⚙️ Inference Parameters", open=True):
                seconds = gr.Slider(minimum=1, maximum=47, value=10, step=0.5, label="Duration (Seconds)")
                num_steps = gr.Slider(minimum=20, maximum=200, value=100, step=5, label="Inference Steps")
                guidance_scale = gr.Slider(minimum=1.0, maximum=15.0, value=7.0, step=0.5, label="Guidance Scale (CFG)")

            # --- Variations Configuration ---
            with gr.Accordion("🎲 Variations & Seed Settings", open=True):
                num_variations = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Number of Variations")
                seed = gr.Number(label="Seed (-1 for Random)", value=-1, precision=0)
            
            generate_btn = gr.Button("Generate Sound Effects", variant="primary")
            seed_info = gr.Textbox(label="Active Seed Used", interactive=False)

        # --- Output Audio Variations ---
        with gr.Column(scale=1):
            gr.Markdown("### Generated Audio Variations")
            audio_out_1 = gr.Audio(label="Variation 1")
            audio_out_2 = gr.Audio(label="Variation 2")
            audio_out_3 = gr.Audio(label="Variation 3")
            audio_out_4 = gr.Audio(label="Variation 4")

    generate_btn.click(
        fn=generate_audio,
        inputs=[
            prompt, 
            negative_prompt, 
            seconds, 
            num_steps, 
            guidance_scale, 
            seed, 
            num_variations
        ],
        outputs=[
            audio_out_1, 
            audio_out_2, 
            audio_out_3, 
            audio_out_4, 
            seed_info
        ]
    )

demo.launch()