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
    output = output[:,:int(duration*44100)]
    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>CLAP-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
    
    * Powered by **Stability AI**''',
    description=description_text,
    allow_flagging=False,
    examples=[
["Generate the festive sounds of a fireworks show: explosions lighting up the sky, crowd cheering, and the faint music playing in the background!! Celebration of the new year!"],
["Melodic human whistling harmonizing with natural birdsong"],
["A parade marches through a town square, with drumbeats pounding, children clapping, and a horse neighing amidst the commotion"],
["Quiet speech and then and airplane flying away"],
["A soccer ball hits a goalpost with a metallic clang, followed by cheers, clapping, and the distant hum of a commentator’s voice"],
["A basketball bounces rhythmically on a court, shoes squeak against the floor, and a referee’s whistle cuts through the air"],
["Dripping water echoes sharply, a distant growl reverberates through the cavern, and soft scraping metal suggests something lurking unseen"],
["A cow is mooing whilst a lion is roaring in the background as a hunter shoots. A flock of birds subsequently fly away from the trees."],
["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"],
["Gentle female voice cooing and baby responding with happy gurgles and giggles"],
["A fork scrapes a plate, water drips slowly into a sink, and the faint hum of a refrigerator lingers in the background"],
["Powerful ocean waves crashing and receding on sandy beach with distant seagulls"],
["Emulate the lively sounds of a retro arcade: 8-bit game music, coins clinking. People cheering occasionally when players winning"],
["Simulate a forest ambiance with birds chirping and wind rustling through the leaves"],
["A train conductor blows a sharp whistle, metal wheels screech on the rails, and passengers murmur while settling into their seats"],
["Generate an energetic and bustling city street scene with distant traffic and close conversations"],
["Alarms blare with rising urgency as fragments clatter against a metallic hull, interrupted by a faint hiss of escaping air"],
["Create a serene soundscape of a quiet beach at sunset"],
["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"],
["Emulate the lively sounds of a retro arcade: 8-bit game music, coins clinking. People cheering occasionally when players winning"],
["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 the underwater soundscape: gentle waves, faint whale calls, and the occasional clink of scuba gear"],
["Recreate the sounds of an active volcano: rumbling earth, lava bubbling, and the occasional loud explosive roar of an eruption"],
["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"],
["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"],
["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"],
["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()