import spaces
import gradio as gr
import json
import torch
import wavio
from tqdm import tqdm
from huggingface_hub import snapshot_download
from pydub import AudioSegment
from gradio import Markdown

import torch
from diffusers import DiffusionPipeline,AudioPipelineOutput
from transformers import CLIPTextModel, T5EncoderModel, AutoModel, T5Tokenizer, T5TokenizerFast
from typing import Union
from diffusers.utils.torch_utils import randn_tensor
from tqdm import tqdm
from TangoFlux import TangoFluxInference
import torchaudio



tangoflux = TangoFluxInference(name="declare-lab/TangoFlux")



@spaces.GPU(duration=15)
def gradio_generate(prompt, steps, guidance,duration=10):

    output = tangoflux.generate(prompt,steps=steps,guidance_scale=guidance,duration=duration)
    output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
    
    
    #wavio.write(output_filename, output_wave, rate=44100, sampwidth=2)
    unique_filename = f"output_{uuid.uuid4().hex}.wav"
    print(f"Saving audio to file: {unique_filename}")

    # Save to file
    torchaudio.save(unique_filename, output, sample_rate)
    print(f"Audio saved: {unique_filename}")

    # Return the path to the generated audio file
    return unique_filename

    #if (output_format == "mp3"):
     #   AudioSegment.from_wav("temp.wav").export("temp.mp3", format = "mp3")
      #  output_filename = "temp.mp3"

    #return output_filename

description_text = """
<p><a href="https://huggingface.co/spaces/declare-lab/tango2/blob/main/app.py?duplicate=true"> <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a> For faster inference without waiting in queue, you may duplicate the space and upgrade to a GPU in the settings. <br/><br/>
Generate audio using Tango2 by providing a text prompt. Tango2 was built from Tango and was trained on <a href="https://huggingface.co/datasets/declare-lab/audio-alpaca">Audio-alpaca</a>
<br/><br/> This is the demo for Tango2 for text to audio generation: <a href="https://arxiv.org/abs/2404.09956">Read our paper.</a>
<p/>
"""
# Gradio input and output components
input_text = gr.Textbox(lines=2, label="Prompt")
#output_format = gr.Radio(label = "Output format", info = "The file you can dowload", choices =  "wav"], value = "wav")
output_audio = gr.Audio(label="Generated Audio", type="filepath")
denoising_steps = gr.Slider(minimum=10, maximum=100, value=25, step=5, label="Steps", interactive=True)
guidance_scale = gr.Slider(minimum=1, maximum=10, value=4.5, step=0.5, label="Guidance Scale", interactive=True)
duration_scale = gr.Slider(minimum=1, maximum=30, value=10, step=1, label="Duration", interactive=True)

interface = gr.Interface(
    fn=gradio_generate,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your text prompt here"),
        gr.Slider(0, 30, value=10, label="Duration in Seconds"),
        gr.Slider(10, 150, value=50, step=5, label="Number of Diffusion Steps"),
        gr.Slider(1, 10, value=4.5, step=0.5, label="CFG Scale")
    ],
    outputs=gr.Audio(type="filepath", label="Generated Audio"),
    title="TangoFlux Generator",
    description="Generate variable-length stereo audio at 44.1kHz from text prompts using TangoFlux.",
    examples=[
    [
        "Create a serene soundscape of a quiet beach at sunset.",  # Text prompt
 
        15,  # Duration in Seconds
        100,  # Number of Diffusion Steps
        4.5,  # CFG Scale
    ],
    
    ["Rock beat played in a treated studio, session drumming on an acoustic kit.",
        30,  # Duration in Seconds
        100,  # Number of Diffusion Steps
        7,  # CFG Scale
     
    ]
])

interface.launch()