app.py

import spaces
import gradio as gr
import json
import torch
import wavio
from tqdm import tqdm
from huggingface_hub import snapshot_download
from models import AudioDiffusion, DDPMScheduler
from audioldm.audio.stft import TacotronSTFT
from audioldm.variational_autoencoder import AutoencoderKL
from pydub import AudioSegment
from gradio import Markdown

from diffusers.models.unet_2d_condition import UNet2DConditionModel
from diffusers import DiffusionPipeline, AudioPipelineOutput
from transformers import T5EncoderModel, T5Tokenizer, T5TokenizerFast, pipeline
from typing import Union
from diffusers.utils.torch_utils import randn_tensor
from tqdm import tqdm
from langdetect import detect, DetectorFactory

# Ensure consistent results from langdetect
DetectorFactory.seed = 0

class Tango2Pipeline(DiffusionPipeline):
    
    def __init__(
        self,
        vae: AutoencoderKL,
        text_encoder: T5EncoderModel,
        tokenizer: Union[T5Tokenizer, T5TokenizerFast],
        unet: UNet2DConditionModel,
        scheduler: DDPMScheduler
    ):
        super().__init__()
        self.register_modules(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            unet=unet,
            scheduler=scheduler
        )
        
    def _encode_prompt(self, prompt):
        device = self.text_encoder.device
        
        batch = self.tokenizer(
            prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
        )
        input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)

        encoder_hidden_states = self.text_encoder(
                input_ids=input_ids, attention_mask=attention_mask
            )[0]

        boolean_encoder_mask = (attention_mask == 1).to(device)
        
        return encoder_hidden_states, boolean_encoder_mask
        
    def _encode_text_classifier_free(self, prompt, num_samples_per_prompt):
        device = self.text_encoder.device
        batch = self.tokenizer(
            prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
        )
        input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)

        with torch.no_grad():
            prompt_embeds = self.text_encoder(
                input_ids=input_ids, attention_mask=attention_mask
            )[0]
                
        prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
        attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)

        # get unconditional embeddings for classifier free guidance
        uncond_tokens = [""] * len(prompt)

        max_length = prompt_embeds.shape[1]
        uncond_batch = self.tokenizer(
            uncond_tokens, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt",
        )
        uncond_input_ids = uncond_batch.input_ids.to(device)
        uncond_attention_mask = uncond_batch.attention_mask.to(device)

        with torch.no_grad():
            negative_prompt_embeds = self.text_encoder(
                input_ids=uncond_input_ids, attention_mask=uncond_attention_mask
            )[0]
                
        negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
        uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)

        # For classifier free guidance, we need to do two forward passes.
        # We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes
        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
        prompt_mask = torch.cat([uncond_attention_mask, attention_mask])
        boolean_prompt_mask = (prompt_mask == 1).to(device)

        return prompt_embeds, boolean_prompt_mask
        
    def prepare_latents(self, batch_size, inference_scheduler, num_channels_latents, dtype, device):
        shape = (batch_size, num_channels_latents, 256, 16)
        latents = randn_tensor(shape, generator=None, device=device, dtype=dtype)
        # scale the initial noise by the standard deviation required by the scheduler
        latents = latents * inference_scheduler.init_noise_sigma
        return latents
    
    @torch.no_grad()
    def inference(self, prompt, inference_scheduler, num_steps=20, guidance_scale=3, num_samples_per_prompt=1, 
                  disable_progress=True):
        device = self.text_encoder.device
        classifier_free_guidance = guidance_scale > 1.0
        batch_size = len(prompt) * num_samples_per_prompt

        if classifier_free_guidance:
            prompt_embeds, boolean_prompt_mask = self._encode_text_classifier_free(prompt, num_samples_per_prompt)
        else:
            prompt_embeds, boolean_prompt_mask = self._encode_prompt(prompt)
            prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
            boolean_prompt_mask = boolean_prompt_mask.repeat_interleave(num_samples_per_prompt, 0)

        inference_scheduler.set_timesteps(num_steps, device=device)
        timesteps = inference_scheduler.timesteps

        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(batch_size, inference_scheduler, num_channels_latents, prompt_embeds.dtype, device)

        num_warmup_steps = len(timesteps) - num_steps * inference_scheduler.order
        progress_bar = tqdm(range(num_steps), disable=disable_progress)

        for i, t in enumerate(timesteps):
            # expand the latents if we are doing classifier free guidance
            latent_model_input = torch.cat([latents] * 2) if classifier_free_guidance else latents
            latent_model_input = inference_scheduler.scale_model_input(latent_model_input, t)

            noise_pred = self.unet(
                latent_model_input, t, encoder_hidden_states=prompt_embeds,
                encoder_attention_mask=boolean_prompt_mask
            ).sample

            # perform guidance
            if classifier_free_guidance:
                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

            # compute the previous noisy sample x_t -> x_t-1
            latents = inference_scheduler.step(noise_pred, t, latents).prev_sample

            # call the callback, if provided
            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % inference_scheduler.order == 0):
                progress_bar.update(1)

        return latents
        
    @torch.no_grad()
    def __call__(self, prompt, steps=100, guidance=3, samples=1, disable_progress=True):
        """ Generate audio for a single prompt string. """
        with torch.no_grad():
            latents = self.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress=disable_progress)
            mel = self.vae.decode_first_stage(latents)
            wave = self.vae.decode_to_waveform(mel)

        return AudioPipelineOutput(audios=wave)

# Automatic device detection
if torch.cuda.is_available():
    device_type = "cuda"
    device_selection = "cuda:0"
else:
    device_type = "cpu"
    device_selection = "cpu"

class Tango:
    def __init__(self, name="declare-lab/tango2", device=device_selection):
        
        path = snapshot_download(repo_id=name)
        
        vae_config = json.load(open("{}/vae_config.json".format(path)))
        stft_config = json.load(open("{}/stft_config.json".format(path)))
        main_config = json.load(open("{}/main_config.json".format(path)))
        
        self.vae = AutoencoderKL(**vae_config).to(device)
        self.stft = TacotronSTFT(**stft_config).to(device)
        self.model = AudioDiffusion(**main_config).to(device)
        
        vae_weights = torch.load("{}/pytorch_model_vae.bin".format(path), map_location=device)
        stft_weights = torch.load("{}/pytorch_model_stft.bin".format(path), map_location=device)
        main_weights = torch.load("{}/pytorch_model_main.bin".format(path), map_location=device)
        
        self.vae.load_state_dict(vae_weights)
        self.stft.load_state_dict(stft_weights)
        self.model.load_state_dict(main_weights)

        print ("Successfully loaded checkpoint from:", name)
        
        self.vae.eval()
        self.stft.eval()
        self.model.eval()
        
        self.scheduler = DDPMScheduler.from_pretrained(main_config["scheduler_name"], subfolder="scheduler")
        
    def chunks(self, lst, n):
        """ Yield successive n-sized chunks from a list. """
        for i in range(0, len(lst), n):
            yield lst[i:i + n]
        
    def generate(self, prompt, steps=200, guidance=8, samples=1, disable_progress=True):
        """ Generate audio for a single prompt string. """
        with torch.no_grad():
            latents = self.model.inference([prompt], self.scheduler, steps, guidance, samples, disable_progress=disable_progress)
            mel = self.vae.decode_first_stage(latents)
            wave = self.vae.decode_to_waveform(mel)
        return wave[0]
    
    def generate_for_batch(self, prompts, steps=200, guidance=8, samples=1, batch_size=8, disable_progress=True):
        """ Generate audio for a list of prompt strings. """
        outputs = []
        for k in tqdm(range(0, len(prompts), batch_size)):
            batch = prompts[k: k+batch_size]
            with torch.no_grad():
                latents = self.model.inference(batch, self.scheduler, steps, guidance, samples, disable_progress=disable_progress)
                mel = self.vae.decode_first_stage(latents)
                wave = self.vae.decode_to_waveform(mel)
                outputs += [item for item in wave]
        if samples == 1:
            return outputs
        else:
            return list(self.chunks(outputs, samples))

# Initialize TANGO
tango = Tango(device=device_selection)
tango.vae.to(device_type)
tango.stft.to(device_type)
tango.model.to(device_type)

pipe = Tango2Pipeline(
    vae=tango.vae,
    text_encoder=tango.model.text_encoder,
    tokenizer=tango.model.tokenizer,
    unet=tango.model.unet,
    scheduler=tango.scheduler
)

# Initialize Translation Pipeline
translation_pipeline = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

def adjust_audio_length(audio_path, desired_length_sec, output_format):
    """
    Adjust the audio to the desired length.
    If the audio is shorter, pad with silence.
    If longer, trim the audio.
    """
    audio = AudioSegment.from_file(audio_path)
    desired_length_ms = desired_length_sec * 1000  # Convert to milliseconds

    if len(audio) < desired_length_ms:
        # Pad with silence
        padding = AudioSegment.silent(duration=desired_length_ms - len(audio))
        audio += padding
    elif len(audio) > desired_length_ms:
        # Trim the audio
        audio = audio[:desired_length_ms]

    # Export the adjusted audio
    adjusted_path = f"adjusted.{output_format}"
    audio.export(adjusted_path, format=output_format)
    return adjusted_path

@spaces.GPU(duration=60)
def gradio_generate(prompt, output_format, steps, guidance, audio_length):
    """
    Generate audio based on the prompt, translate if necessary, and adjust its length.
    """
    # Detect language
    try:
        lang = detect(prompt)
    except:
        lang = "unknown"

    # If the prompt is in Korean, translate to English
    if lang == "ko":
        translated = translation_pipeline(prompt)[0]['translation_text']
        print(f"Translated Prompt: {translated}")
        prompt_to_use = translated
    else:
        prompt_to_use = prompt

    # Generate audio using the pipeline
    output_wave = pipe(prompt_to_use, steps, guidance)
    output_wave = output_wave.audios[0]
    temp_wav = "temp.wav"
    wavio.write(temp_wav, output_wave, rate=16000, sampwidth=2)

    # Adjust audio length
    adjusted_path = adjust_audio_length(temp_wav, audio_length, output_format)

    return adjusted_path

# 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 download",
    choices=["mp3", "wav"],
    value="wav"
)
audio_length = gr.Slider(
    minimum=4,
    maximum=10,
    step=1,
    label="Audio Length (seconds)",
    value=6,
    interactive=True
)
output_audio = gr.Audio(label="Generated Audio", type="filepath")
denoising_steps = gr.Slider(
    minimum=100,
    maximum=200,
    step=1,
    label="Steps",
    value=200,  # Changed from 100 to 200
    interactive=True
)
guidance_scale = gr.Slider(
    minimum=1,
    maximum=10,
    step=0.1,
    label="Guidance Scale",
    value=8,  # Changed from 3 to 8
    interactive=True
)

# Gradio interface
gr_interface = gr.Interface(
    theme="Nymbo/Nymbo_Theme",
    fn=gradio_generate,
    inputs=[input_text, output_format, denoising_steps, guidance_scale, audio_length],
    outputs=[output_audio],
    title="Tango2: Text to SoundFX",
    allow_flagging=False,
    examples=[
        ["조용한 말소리 후 비행기가 멀어지는 소리"],
        ["사람들이 환호하고 박수치는 소리"],
        ["강한 바람 소리와 빗소리"],        
        ["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 man talking followed by a goat baaing then a metal gate sliding shut as ducks quack and wind blows into a microphone."],
        ["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.
)

# Launch Gradio app
gr_interface.queue(10).launch()