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 uuid
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)
    filename = 'temp.wav'
    #print(f"Saving audio to file: {unique_filename}")

    # Save to file
    torchaudio.save(filename, output, 44100)
   

    # Return the path to the generated audio file
    return filename

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

    #return output_filename

description_text = """
Generate high quality and faithful audio in just a few seconds using <b>TangoFlux</b> by providing a text prompt. <b>TangoFlux</b> was trained from scratch and underwent alignment to follow human instructions using a new method called <b>Claped-Ranked Preference Optimization (CRPO)</b>.
<div style="display: flex; gap: 10px; align-items: center;">
  <a href="https://arxiv.org/abs/2412.21037">
    <img src="https://img.shields.io/badge/Read_the_Paper-blue?link=https%3A%2F%2Fopenreview.net%2Fattachment%3Fid%3DtpJPlFTyxd%26name%3Dpdf" alt="arXiv">
  </a>
  <a href="https://huggingface.co/declare-lab/TangoFlux">
    <img src="https://img.shields.io/badge/TangoFlux-Huggingface-violet?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fdeclare-lab%2FTangoFlux" alt="Static Badge">
  </a>
  <a href="https://tangoflux.github.io/">
    <img src="https://img.shields.io/badge/Demos-declare--lab-brightred?style=flat" alt="Static Badge">
  </a>
  <a href="https://huggingface.co/spaces/declare-lab/TangoFlux">
    <img src="https://img.shields.io/badge/TangoFlux-Huggingface_Space-8A2BE2?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fdeclare-lab%2FTangoFlux" alt="Static Badge">
  </a>
  <a href="https://huggingface.co/datasets/declare-lab/CRPO">
    <img src="https://img.shields.io/badge/TangoFlux_Dataset-Huggingface-red?logo=huggingface&link=https%3A%2F%2Fhuggingface.co%2Fdatasets%2Fdeclare-lab%2FTangoFlux" alt="Static Badge">
  </a>
  <a href="https://github.com/declare-lab/TangoFlux">
    <img src="https://img.shields.io/badge/Github-brown?logo=github&link=https%3A%2F%2Fgithub.com%2Fdeclare-lab%2FTangoFlux" alt="Static Badge">
  </a>
</div>
"""
# 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)


# Gradio interface
gr_interface = gr.Interface(
    fn=gradio_generate,
    inputs=[input_text, denoising_steps, guidance_scale,duration_scale],
    outputs=output_audio,
    title="TangoFlux: Super Fast and Faithful Text to Audio Generation with Flow Matching and Clap-Ranked Preference Optimization",
    description=description_text,
    allow_flagging=False,
    examples=[
["A parade marches through a town square, with drumbeats pounding, children clapping, and a horse neighing amidst the commotion"],
["A soccer ball hits a goalpost with a metallic clang, followed by cheers, clapping, and the distant hum of a commentator’s voice"],
["The deep growl of an alligator ripples through the swamp as reeds sway with a soft rustle and a turtle splashes into the murky water"],        
["A basketball bounces rhythmically on a court, shoes squeak against the floor, and a referee’s whistle cuts through the air"],
["A train conductor blows a sharp whistle, metal wheels screech on the rails, and passengers murmur while settling into their seats"],
["A fork scrapes a plate, water drips slowly into a sink, and the faint hum of a refrigerator lingers in the background"],
["Alarms blare with rising urgency as fragments clatter against a metallic hull, interrupted by a faint hiss of escaping air"],
["Tiny pops and hisses of chemical reactions intermingle with the rhythmic pumping of a centrifuge and the soft whirr of air filtration"],
["A train conductor blows a sharp whistle, metal wheels screech on the rails, and passengers murmur while settling into their seats"],
["Simulate a forest ambiance with birds chirping and wind rustling through the leaves"],
["Quiet whispered conversation gradually fading into distant jet engine roar diminishing into silence"],
["Clear sound of bicycle tires crunching on loose gravel and dirt, followed by deep male laughter echoing"],
["Multiple ducks quacking loudly with splashing water and piercing wild animal shriek in background"],
["Create a serene soundscape of a quiet beach at sunset"],
["A pile of coins spills onto a wooden table with a metallic clatter, followed by the hushed murmur of a tavern crowd and the creak of a swinging door"],
["Generate an energetic and bustling city street scene with distant traffic and close conversations"],
["Powerful ocean waves crashing and receding on sandy beach with distant seagulls"],
["Gentle female voice cooing and baby responding with happy gurgles and giggles"],
["Clear male voice speaking, sharp popping sound, followed by genuine group laughter"],
["Stream of water hitting empty ceramic cup, pitch rising as cup fills up"],
["Massive crowd erupting in thunderous applause and excited cheering"],
["Deep rolling thunder with bright lightning strikes crackling through sky"],
["Aggressive dog barking and distressed cat meowing as racing car roars past at high speed"],
["Peaceful stream bubbling and birds singing, interrupted by sudden explosive gunshot"],
["Man speaking outdoors, goat bleating loudly, metal gate scraping closed, ducks quacking frantically, wind howling into microphone"],
["Series of loud aggressive dog barks echoing"],
["Multiple distinct cat meows at different pitches"],
["Rhythmic wooden table tapping overlaid with steady water pouring sound"],
["Sustained crowd applause with camera clicks and amplified male announcer voice"],
["Two sharp gunshots followed by panicked birds taking flight with rapid wing flaps"],
["Melodic human whistling harmonizing with natural birdsong"],
["Deep rhythmic snoring with clear breathing patterns"],
["Multiple racing engines revving and accelerating with sharp whistle piercing through"],
["Massive stadium crowd cheering as thunder crashes and lightning strikes"],
["Heavy helicopter blades chopping through air with engine and wind noise"],
["Dog barking excitedly and man shouting as race car engine roars past"],
["Quiet speech and then and airplane flying away"],
["A bicycle peddling on dirt and gravel followed by a man speaking then laughing"],
["Ducks quack and water splashes with some animal screeching in the background"],
["Describe the sound of the ocean"],
["A woman and a baby are having a conversation"],
["A man speaks followed by a popping noise and laughter"],
["A cup is filled from a faucet"],
["An audience cheering and clapping"],
["Rolling thunder with lightning strikes"],
["A dog barking and a cat mewing and a racing car passes by"],
["Gentle water stream, birds chirping and sudden gun shot"],
["A dog barking"],
["A cat meowing"],
["Wooden table tapping sound while water pouring"],
["Applause from a crowd with distant clicking and a man speaking over a loudspeaker"],
["two gunshots followed by birds flying away while chirping"],
["Whistling with birds chirping"],
["A person snoring"],
["Motor vehicles are driving with loud engines and a person whistles"],
["People cheering in a stadium while thunder and lightning strikes"],
["A helicopter is in flight"],
["A dog barking and a man talking and a racing car passes by"],
    ],
    cache_examples="lazy", # Turn on to cache.
)



gr_interface.queue(15).launch()