app.py

import json
import os
import numpy as np
from threading import Thread
import torch
import time
from MusicgenStreamer import MusicgenStreamer
from AudioLogger import AudioLogger
from transformers import set_seed, AutoProcessor

import torchaudio

import datetime
from diffusers import AutoPipelineForText2Image
from PIL import Image
import random
import numpy as np

import torch
from transformers import MusicgenForConditionalGeneration, AutoProcessor
import gradio as gr
import vertexai
import pathlib
import json

# import vertexai.generative_models as
from vertexai.generative_models import GenerativeModel, Part,Image
import vertexai.preview.generative_models as generative_models
import time
import typing_extensions as typing
import os
import json
import tempfile

# Retrieve the JSON string from the environment variable
google_creds_json = os.getenv('GOOGLE_APPLICATION_CREDENTIALS_JSON')

if google_creds_json:
    # Create a temporary file to store the credentials
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
        temp_file.write(google_creds_json.encode('utf-8'))
        temp_file_path = temp_file.name

    # Set the environment variable to point to this temporary file
    os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = temp_file_path
else:
    raise ValueError("GOOGLE_APPLICATION_CREDENTIALS_JSON environment variable is not set.")
vertexai.init(project="wubble2024", location="asia-southeast1")
torch.set_default_device("cpu")
# Initialize the text-to-image pipeline
pipe = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/sdxl-turbo",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
)
pipe.to("cuda")
torch.set_default_device("cuda")

generation_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small").to("cuda")
generation_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
generation_model = torch.compile(generation_model)


system_prompt = """You are a music generation assistant developed by Wubble AI. Your job is to understand the inputs that the user is giving and generate an appropriate data which may be used to generate audio, thus giving you the capabilities to creaste audio and music.
The user can also request you to extend an audio. In that case, simply set the extend_flag to True.

Always return output in the given JSON format where the keys are:
song_title: str: A title for the song
model_response: str: response sent back to the user
music_caption: str: will be fed into the musicgen model to generate music. should be a natural language prompt for the required music. can contain instruments required, tempo, genre, mood, etc.
cover_description: str: description for the generation of a cover image
generation_flag: bool: set True if user needs a music sample
extend_flag: bool: set True if the user is requesting you to extend an audio
extended_duration: float: duration of the extended audio (in seconds). the full duration of the original + the extension amount


Generate samples (by setting the generation_flag to true and setting a music_caption, song_title, and cover_description) as much as possible.
You will be provided with the generated audio in the next user input.
If prompted to modify an audio or music, try to modify the previous music description accordingly.
if asked to add, remove, or extract a stem from the audio, modify the music description accordingly and set generation_flag to true.
All fields are required.
Important: Only return a single piece of valid JSON text."""

class GeminiOutput(typing.TypedDict):
  song_title: str
  model_response:str
  music_caption: str
  cover_description: str
  generation_flag: bool
  extend_flag: bool
  extended_duration: float

model = GenerativeModel(
  "gemini-1.5-flash-001",system_instruction=[system_prompt
    ],generation_config={"response_mime_type": "application/json","response_schema":GeminiOutput})



safety_settings = {
    generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
}




max_retries = 1
# multiple audios and multiple images
def predict(chat, prompt_input, image_path_list = None, audio_path_list = None,  video_path_list = None,  top_p = 0.8, temperature = 0.6, max_gen_tokens = 512):
    generation_config = {
      "max_output_tokens": max_gen_tokens,
      "temperature": temperature,
      "top_p": top_p,
    }

    model_input = []
    if audio_path_list is None:
      audio_path_list = []
    if image_path_list is None:
      image_path_list = []
    if video_path_list is None:
      video_path_list = []

    print(f"audio_paths: {audio_path_list}")

    for image_path in image_path_list:
      if image_path is not None:
        _,file_extension = os.path.splitext(image_path)
        image_data = Part.from_data(pathlib.Path(image_path).read_bytes(), mime_type=f"image/{file_extension.replace('.','')}")
        model_input.append(image_data)
    for audio_path in audio_path_list:
      if audio_path is not None:
        print(audio_path)
        _,file_extension = os.path.splitext(audio_path)
        audio_data = Part.from_data(pathlib.Path(audio_path).read_bytes(), mime_type=f"audio/{file_extension}")
        model_input.append(audio_data)
    for video_path in video_path_list:
      if video_path is not None:
        _,file_extension = os.path.splitext(video_path)
        video_data = Part.from_data(pathlib.Path(video_path).read_bytes(), mime_type=f"video/{file_extension.replace('.','')}")
        model_input.append(video_data)
    if prompt_input is not None:
      model_input.append(prompt_input)
    global retries
    retries = 0
    
    def valiate_gemini_response(response):
      global retries
      global max_retries
      json_response = None
      try:
        json_response = json.loads(response.split('\n',1)[1].rsplit('\n',1)[0])

      except:
        pass
      if json_response is None:
        try:
          json_response = json.loads(response.split('\n',1)[1].rsplit('\n',2)[0])
        except:
          if retries < max_retries:
            print('Give the previous output in the correct JSON format.')
            response = chat.send_message(
                ['Give the previous output in the correct JSON format.'],
                generation_config=generation_config,
                safety_settings=safety_settings)
            retries += 1
            return valiate_gemini_response(response)
          else:
            print("Gemini output not JSON compatible")

      if json_response is not None:
        if json_response['generation_flag'] is True and ((json_response['music_caption'] is None or json_response['music_caption']  == "") or (json_response['cover_description'] is None or json_response['cover_description']  == "") or  (json_response['song_title'] is None or json_response['song_title']  == "")): #XOR gate here?
          if retries < max_retries:
            print('The generation flag is set to true but there is no music caption and/or cover_description and/or song title associated . correct the error and respond in the correct JSON format')
            response = chat.send_message(
                ['The generation flag is set to true but there is no music caption and/or cover_description and/or song title associated. correct the error and respond in the correct JSON format'],
                generation_config=generation_config,
                safety_settings=safety_settings)
            retries += 1
            return valiate_gemini_response(response)
          else:
            print("Gemini generation flag is true but no music caption/cover description/song_title given")
        elif json_response['generation_flag'] is False and ((json_response['music_caption'] is not None and json_response['music_caption']  != "") and (json_response['cover_description'] is not None and json_response['cover_description']  != "") and (json_response['song_title'] is not None and json_response['song_title']  != "")):

          json_response['generation_flag'] = True
          return json_response
        elif json_response['extend_flag'] is True and (json_response['extended_duration'] is None or json_response['extended_duration']  == ""):
          json_response['extended_duration'] = 60.0
          return json_response
        elif json_response['extend_flag'] is False and (json_response['extended_duration'] is not None and json_response['extended_duration']  != "" and json_response['extended_duration']  > 15.0):
          json_response['extend_flag'] = True
          return json_response
        else:
          return json_response

    # validation
    start  = time.time()
    response = chat.send_message(
        model_input,
        generation_config=generation_config,
        safety_settings=safety_settings
    )
    print(f"Input Prompt: {prompt_input}, Response: {response.candidates[0].content.parts[0].text}")
    json_response = valiate_gemini_response(response.candidates[0].content.parts[0].text)
    # print(json_response)
    # # print(response['candidates']['content']['parts'])
    print(f"inference time = {time.time() - start}")
    # print(message, history)

    return json_response

@torch.inference_mode()
def extend_music(log, audio_path, audio_length_in_s, prompt, play_steps_in_s, stride=None, extend_stride=18.0, seed=None, top_k = 50, top_p = 1.0, temperature = 1.0,do_sample=True, guidance_scale=3):
    if audio_path is not None:
        sample_rate = 32000
        waveform, sr = torchaudio.load(audio_path)
        if sample_rate != sr:
            waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=sample_rate)
        audio = torch.mean(waveform, 0)
    if seed is not None:
      print(f"Seed: {seed}")
      set_seed(seed)
    else:
      seed = np.random.randint(0, 2**32 - 1)
      print(f"Seed: {seed}")
      set_seed(seed)
    if top_p is not None:
      print(f"top_p: {top_p}")
    if top_k is not None:
      print(f"top_k: {top_k}")
    if temperature is not None:
      print(f"temperature: {temperature}")
    device = "cuda"
    print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} Generating Music...")
    sampling_rate = generation_model.audio_encoder.config.sampling_rate
    frame_rate = generation_model.audio_encoder.config.frame_rate
    max_duration = 30.0
    max_tokens = 1503

    play_steps = int(frame_rate * play_steps_in_s) # for streamer

    audio_duration = round(len(audio)/sampling_rate)
    generation_duration = audio_length_in_s - audio_duration
    assert generation_duration > 0, f"extended duration ({generation_duration}s) should be greater than the duration of the audio input ({audio_duration}s)"

    total_gen_tokens = int(sampling_rate * generation_duration) # number of tokens to be generated
    audio_condition = None # initialize audio condition as none
    current_gen_length = 0 # number of tokens currently generated
    stride_tokens = int(sampling_rate * extend_stride) # number of tokens to generate each time
    context_duration = max_duration - extend_stride
    context_tokens = int(sampling_rate * context_duration) # number of tokens required for context
    initial_streamer = True
    initial_audio_flag = False
    if len(audio) > context_tokens:
        audio_condition = audio[-context_tokens:]
    else:
        audio_condition = audio
        extend_stride = max_duration - audio_duration
        stride_tokens = int(sampling_rate * extend_stride) # number of tokens to generate each time
        context_duration = audio_duration
        context_tokens = int(sampling_rate * context_duration) # number of tokens required for context
    print(len(audio_condition)/sampling_rate)
    global_audio = list(audio)
    time_offset = 0
    while current_gen_length < total_gen_tokens:
        chunk_duration = min(generation_duration - time_offset, max_duration - context_duration) # each new chunk can be of duration min(whats left to generate, max possible generation)
        max_gen_len = int(chunk_duration * frame_rate)
        print(time_offset,chunk_duration,max_gen_len)

        streamer = MusicgenStreamer(generation_model, device=device, stride = stride, play_steps=play_steps,
                                    duration =audio_length_in_s, initial_streamer=initial_streamer, audio_context = list(audio_condition))
        initial_streamer = False
        assert audio_condition is not None, "Audio condition is none"
        print(f"Tokens of Audio condition: {int((len(audio_condition)/32000)*50)}")
        gen_inputs = generation_processor(audio = audio_condition,text=prompt, sampling_rate = sampling_rate,padding="max_length",
                                    max_length=128,return_tensors="pt").to(device)

        generation_kwargs = dict(**gen_inputs,streamer=streamer,
                                max_new_tokens=max_gen_len,do_sample=do_sample, guidance_scale=guidance_scale, top_k = top_k, top_p = top_p, temperature=temperature)
        thread = Thread(target=generation_model.generate, kwargs=generation_kwargs)
        thread.start()

        total_audio = []
        offset = 0
        start_time = time.time()
        offset_flag = True
        for new_audio in streamer:
            if len(new_audio) > 0:
                if initial_audio_flag is False:
                    print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}. {len(audio)}")
                    yield audio
                    initial_audio_flag = True
                total_audio += list(new_audio)
                print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}. {time.time() - start_time} Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds, {round(new_audio.shape[0])}, {len(total_audio)}, max: {np.max(new_audio)}")
                yield new_audio
            start_time = time.time()
        current_gen_length += len(total_audio)+(0.06*32000)
        global_audio += list(total_audio)
        total_audio = list(audio_condition) + total_audio
        audio_condition = total_audio[-context_tokens:]

        time_offset = current_gen_length / sampling_rate
        # max_duration = 30 - len(audio_condition)/sampling_rate
        print(current_gen_length, round(current_gen_length/32000), len(global_audio)/32000, (len(total_audio))/32000)


    log.log_audio(task = 'extension', audio = global_audio, duration = audio_length_in_s, prompt = prompt, generation_flag = True,
                  music_caption= None, play_steps_in_s = play_steps_in_s, extend_stride = extend_stride, seed = seed,
                  top_p = top_p, top_k = top_k, temperature = temperature,do_sample=do_sample, guidance_scale=guidance_scale, context_audio_path=audio_path)
    random.seed()
    np.random.seed()
    torch.manual_seed(torch.initial_seed())
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(torch.initial_seed())

@torch.inference_mode()
def generate_music(log, audio_length_in_s, user_prompt, music_caption, play_steps_in_s,stride=None, extend_stride=18.0, seed=None, top_k = 50, top_p = 1.0, temperature = 1.0,do_sample=True, guidance_scale=3, generation_flag=False):
    if music_caption is not None:
        prompt = music_caption
    else:
        prompt = user_prompt
    if seed is not None:
      print(f"Seed: {seed}")
      set_seed(seed)
    else:
      seed = np.random.randint(0, 2**32 - 1)
      print(f"Seed: {seed}")
      set_seed(seed)
    if top_p is not None:
      print(f"top_p: {top_p}")
    if top_k is not None:
      print(f"top_k: {top_k}")
    if temperature is not None:
      print(f"temperature: {temperature}")
    device = "cuda"
    print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} Generating Music...")
    sampling_rate = generation_model.audio_encoder.config.sampling_rate
    frame_rate = generation_model.audio_encoder.config.frame_rate
    max_duration = 30.0
    max_tokens = 1503


    play_steps = int(frame_rate * play_steps_in_s) # for streamer

    generation_duration = audio_length_in_s
    total_gen_tokens = int(sampling_rate * generation_duration) # number of tokens to be generated
    audio_condition = None # initialize audio condition as none
    current_gen_length = 0 # number of tokens currently generated
    stride_tokens = int(sampling_rate * extend_stride) # number of tokens to generate each time
    context_duration = max_duration - extend_stride
    context_tokens = int(sampling_rate * context_duration) # number of tokens required for context
    initial_streamer = True
    initial_audio_flag = False
    global_audio = []
    time_offset = 0

    while current_gen_length < total_gen_tokens:
        chunk_duration = min(generation_duration - time_offset, max_duration) # each new chunk can be of duration min(whats left to generate, max possible generation)
        max_gen_len = int(chunk_duration * frame_rate)
        print(time_offset,chunk_duration,max_gen_len)

        streamer = MusicgenStreamer(generation_model, device=device, stride = stride, play_steps=play_steps,
                                    duration =audio_length_in_s, initial_streamer=initial_streamer, audio_context = audio_condition)
        initial_streamer = False

        if audio_condition is not None:
          print(f"Tokens of Audio condition: {int((len(audio_condition)/32000)*50)}")

          gen_inputs = generation_processor(audio = audio_condition,text=prompt, sampling_rate = sampling_rate,padding="max_length",
                                    max_length=128,return_tensors="pt").to(device)
        else:
          gen_inputs = generation_processor(text=prompt, padding='max_length',
                                            max_length=128, truncation=True, return_tensors="pt").to(device)

        generation_kwargs = dict(**gen_inputs,streamer=streamer,
                                max_new_tokens=max_gen_len,do_sample=do_sample, guidance_scale=guidance_scale, top_k = top_k, top_p = top_p, temperature=temperature)
        thread = Thread(target=generation_model.generate, kwargs=generation_kwargs)
        thread.start()

        total_audio = []
        offset = 0
        start_time = time.time()
        offset_flag = True
        for new_audio in streamer:
            if len(new_audio) > 0:
                total_audio += list(new_audio)
                print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}. {time.time() - start_time} Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds, {round(new_audio.shape[0])}, {len(total_audio)}, max: {np.max(new_audio)}")
                yield new_audio
            start_time = time.time()
        current_gen_length += len(total_audio)+(0.06*32000)
        global_audio += list(total_audio)
        if audio_condition is not None:
          total_audio = list(audio_condition) + total_audio
        audio_condition = total_audio[-context_tokens:]
        max_duration = extend_stride
        time_offset = current_gen_length / sampling_rate
        # max_duration = 30 - len(audio_condition)/sampling_rate
        print(current_gen_length, round(current_gen_length/32000), len(global_audio)/32000, (len(total_audio))/32000)

    log.log_audio(task = 'generation', audio = global_audio, duration = audio_length_in_s,prompt= user_prompt, generation_flag=generation_flag,
                  music_caption = music_caption, play_steps_in_s = play_steps_in_s, extend_stride = extend_stride, seed = seed,
                  top_p = top_p, top_k = top_k, temperature = temperature,do_sample=do_sample, guidance_scale=guidance_scale, context_audio_path=None)
    random.seed()
    np.random.seed()
    torch.manual_seed(torch.initial_seed())
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(torch.initial_seed())

def generate_image(prompt,num_inference_steps=1, guidance_scale=3, seed=None):
    if seed is None:
      seed = np.random.randint(0, 2**32 - 1)
      print(f"Random Seed: {seed}")

    start_time = time.time()
    image = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
    end_time = time.time()
    inference_time = end_time - start_time
    print(f" Inference Time for Cover Generation: {inference_time}")
    return image



def main(history, prompt, image, audio, video, duration=15, play_steps_in_s=2, top_p=0.9, temperature=1,extend_stride=18.0, top_k = 65,do_sample=True, guidance_scale=3, use_original_parameters_for_extension= True):
  print(f"history: {history}")
  if history is None:
    chat = model.start_chat()
    log = AudioLogger()
    history = []
  else:
    chat = history[-1]['chat']
    log = history[-1]['log']

  duration += 1
  cover = None
  cover_description = None
  music_caption  =None
  sampling_rate = 32000
  generated_audio = []
  audio_stream = None
  audio_paths = []
  if len(log.history) > 0:
    last_log = log.get_last_log()
    if last_log is not None:
      prev_audio = last_log['audio_path']
      audio_paths = [prev_audio]
  audio_paths.append(audio)
  gemini_output = predict(chat, prompt, [image], audio_paths, [video])
  print(f"Gemini output to main function: {gemini_output}")
  if gemini_output['extend_flag'] == True and gemini_output['extended_duration'] > 0: # if extension mode, then we dont need to run
    if len(log.history) == 0:
      raise gr.Error("Please generate music first")
    prev_log = log.get_last_log()
    if use_original_parameters_for_extension:
      # extend_music(audio_path, audio_length_in_s, prompt, play_steps_in_s, stride=None, extend_stride=18.0, seed=None, top_k = 50, top_p = 1.0, temperature = 1.0,do_sample=True, guidance_scale=3):
        audio_stream = extend_music(log=log, audio_path = prev_log['audio_path'], audio_length_in_s = gemini_output['extended_duration'], prompt = prev_log['prompt'], **prev_log['audio_gen_args']) # keep previous seed for extension? keep previous generation args for extension?
    else:
        audio_stream = extend_music(log=log, audio_path = prev_log['audio_path'], audio_length_in_s = gemini_output['extended_duration'], prompt = prev_log['prompt'], play_steps_in_s=play_steps_in_s, extend_stride=extend_stride, seed=prev_log['seed'],
                    top_k = top_k, top_p = top_p, temperature = temperature,do_sample=do_sample, guidance_scale=guidance_scale) # keep previous seed for extension? keep previous generation args for extension?

    for audio_chunk in audio_stream:
      generated_audio += list(audio_chunk)
      yield gemini_output['model_response'], gemini_output['song_title'], history[-1]['music_caption'], (sampling_rate, np.asarray(audio_chunk)), history[-1]['cover'], cover_description, history
      
    history.append({"chat": chat, "log": log, "prompt": prompt, "model_response": gemini_output['model_response'], "music_caption": music_caption, "audio": generated_audio, "cover": cover, "cover_description": cover_description})

  elif gemini_output['generation_flag']: # should we give the actual audio here as well
    cover_description = gemini_output['cover_description']
    cover = generate_image(cover_description, guidance_scale=1)
    music_caption = gemini_output['music_caption']
    audio_stream = generate_music(log=log, audio_length_in_s = duration, user_prompt = prompt, music_caption=music_caption, play_steps_in_s=play_steps_in_s,stride=None,
            extend_stride=extend_stride, seed=None, top_k = top_k, top_p = top_p, temperature = temperature,do_sample=do_sample, guidance_scale=guidance_scale, generation_flag=gemini_output['generation_flag'])
    for audio_chunk in audio_stream:
      generated_audio += list(audio_chunk)
      yield gemini_output['model_response'], gemini_output['song_title'], music_caption, (sampling_rate, audio_chunk), cover, cover_description, history
      
    history.append({"chat": chat, "log": log, "prompt": prompt, "model_response": gemini_output['model_response'], "music_caption": music_caption, "audio": generated_audio, "cover": cover, "cover_description": cover_description})

  else:
    if len(log.history) > 0:
      last_log = log.get_last_log()
      if last_log is not None:
        audio = last_log['audio_path']
      yield gemini_output['model_response'], gemini_output['song_title'], music_caption, audio, cover, cover_description, history
      
      history.append({"chat": chat, "log": log, "prompt": prompt, "model_response": gemini_output['model_response'], "music_caption": music_caption, "audio": audio, "cover": cover, "cover_description": cover_description})
    else:
      yield gemini_output['model_response'], gemini_output['song_title'], music_caption, None, cover, cover_description, history

      history.append({"chat": chat, "log": log, "prompt": prompt, "model_response": gemini_output['model_response'], "music_caption": music_caption, "audio": None, "cover": cover, "cover_description": cover_description})




demo = gr.Interface(
    fn=main,
    inputs=[gr.State(),gr.Text(label="Prompt", placeholder="Chat with Wubble AI here!")],
    additional_inputs= [
            gr.Image(type="filepath"),
            gr.Audio(type="filepath",),
            gr.Video(),
            gr.Slider(10, 200, value=15, step=5, label="Duration"),
            gr.Slider(0.5, 5, value=2, step=0.5, label="Stream Steps"),
            gr.Slider(0, 1, value=1, step=0.1, label="Top P"),
            gr.Slider(0, 2, value=1.4, step=0.1, label="Temperature"),
            gr.Slider(10, 20, value=18, step=0.1, label="Extend Stride"),
            gr.Slider(45, 120, value=65, step=10, label="top_k"),
            gr.Checkbox(value=True),
            gr.Slider(0, 6, value=3, step=1, label="Guidance Scale"),
            gr.Checkbox(value=True)
            ],
    outputs=[gr.Textbox(label="Generated Text Output"),
             gr.Textbox(label="Song Title"),
             gr.Textbox(label="Music generation caption"),
              gr.Audio(label="Generated Music", streaming=True, autoplay=True),
             gr.Image(label="Cover Image"),
             gr.Textbox(label="Cover description"),
             gr.State(),
    ],
    flagging_options=['Good', 'Bad'],
    title= "Wubble AI Demo",
    allow_flagging="manual",
    concurrency_limit=4)


demo.queue().launch(debug=True)