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)