Merge pull request #1 from animikhaich/web-app-dev

.gitignore

@@ -167,4 +167,7 @@ cython_debug/
 *.mp3
 *.mp4
 
-creds.json
+creds.json
+
+# Ignore the test file
+test.py

engine/__init__.py

@@ -1 +1,2 @@
 from .video_descriptor import DescribeVideo
+from .audio_generator import GenerateAudio

engine/audio_generator.py

@@ -1 +1,128 @@
-# TODO: Add from model server
+import os
+import warnings
+
+warnings.simplefilter("ignore")
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+import io
+import torch
+import numpy as np
+from audiocraft.models import musicgen
+from scipy.io.wavfile import write as wav_write
+
+try:
+    from logger import logging
+except:
+    import logging
+
+
+class GenerateAudio:
+    def __init__(self, model="musicgen-stereo-small"):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model_name = self.get_model_name(model)
+        self.model = self.get_model(self.model_name, self.device)
+        self.generated_audio = None
+        self.sampling_rate = None
+
+    @staticmethod
+    def get_model(model, device):
+        try:
+            model = musicgen.MusicGen.get_pretrained(model, device=device)
+            logging.info(f"Loaded model: {model}")
+            return model
+        except Exception as e:
+            logging.error(f"Failed to load model: {e}")
+            raise ValueError(f"Failed to load model: {e}")
+            return
+
+    @staticmethod
+    def get_model_name(model_name):
+        if model_name.startswith("facebook/"):
+            return model_name
+        return f"facebook/{model_name}"
+    
+    @staticmethod
+    def duration_sanity_check(duration):
+        if duration < 1:
+            logging.warning("Duration is less than 1 second. Setting duration to 1 second.")
+            return 1
+        elif duration > 30:
+            logging.warning("Duration is greater than 30 seconds. Setting duration to 30 seconds.")
+            return 30
+        return duration
+
+    @staticmethod
+    def prompts_sanity_check(prompts):
+        if isinstance(prompts, str):
+            prompts = [prompts]
+        elif not isinstance(prompts, list):
+            raise ValueError("Prompts should be a string or a list of strings.")
+        else:
+            for prompt in prompts:
+                if not isinstance(prompt, str):
+                    raise ValueError("Prompts should be a string or a list of strings.")
+            if len(prompts) > 8: # Too many prompts will cause OOM error
+                raise ValueError("Maximum number of prompts allowed is 8.")
+        return prompts
+    
+
+    def generate_audio(self, prompts, duration=10):
+        duration = self.duration_sanity_check(duration)
+        prompts = self.prompts_sanity_check(prompts)
+
+        try:
+            self.model.set_generation_params(duration=duration)
+            result = self.model.generate(prompts, progress=False)
+            self.result = result.cpu().numpy().T
+            self.result = self.result.transpose((2, 0, 1))
+            self.sampling_rate = self.model.sample_rate
+            logging.info(
+                f"Generated audio with shape: {self.result.shape}, sample rate: {self.sampling_rate} Hz"
+            )
+            print(f"Generated audio with shape: {self.result.shape}, sample rate: {self.sampling_rate} Hz")
+            return self.sampling_rate, self.result
+        except Exception as e:
+            logging.error(f"Failed to generate audio: {e}")
+            raise ValueError(f"Failed to generate audio: {e}")
+
+    def save_audio(self, audio_dir="generated_audio"):
+        if self.result is None:
+            raise ValueError("Audio is not generated yet.")
+        if self.sampling_rate is None:
+            raise ValueError("Sampling rate is not available.")
+
+        paths = []
+        os.makedirs(audio_dir, exist_ok=True)
+        for i, audio in enumerate(self.result):
+            path = os.path.join(audio_dir, f"audio_{i}.wav")
+            wav_write(path, self.sampling_rate, audio)
+            paths.append(path)
+        return paths
+
+    def get_audio_buffer(self):
+        if self.result is None:
+            raise ValueError("Audio is not generated yet.")
+        if self.sampling_rate is None:
+            raise ValueError("Sampling rate is not available.")
+
+        buffers = []
+        for audio in self.result:
+            buffer = io.BytesIO()
+            wav_write(buffer, self.sampling_rate, audio)
+            buffer.seek(0)
+            buffers.append(buffer)
+        return buffers
+
+if __name__ == "__main__":
+    audio_gen = GenerateAudio()
+    sample_rate, result = audio_gen.generate_audio(
+        [
+            "A piano playing a jazz melody", 
+            "A guitar playing a rock riff", 
+            "A LoFi music for coding"
+        ], 
+        duration=10
+    )
+    paths = audio_gen.save_audio()
+    print(f"Saved audio to: {paths}")
+    buffers = audio_gen.get_audio_buffer()
+    print(f"Audio buffers: {buffers}")

engine/video_descriptor.py

@@ -1,8 +1,7 @@
+import os
 from warnings import simplefilter
 
 simplefilter("ignore")
-import os
-
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 import json
 import time
@@ -78,6 +77,9 @@ class DescribeVideo:
 
         return json.loads(cleaned_response.text.strip("```json\n"))
 
+    def __call__(self, video_path):
+        return self.describe_video(video_path)
+
     def reset_safety_settings(self):
         logging.info("Resetting safety settings")
         self.is_safety_set = False

main.py

@@ -1,67 +1,91 @@
 import streamlit as st
+from engine import DescribeVideo, GenerateAudio
 
-def main():
-    st.set_page_config(page_title="VidTune: Where Videos Find Their Melody", layout="centered")
 
-    # Title and Description
-    st.title("VidTune: Where Videos Find Their Melody")
-    st.write("VidTune is a web application that allows users to upload videos and generate melodies matching the mood of the video.")
+video_model_map = {
+    "Fast": "flash",
+    "Quality": "pro",
+}
+
+music_model_map = {
+    "Fast": "musicgen-stereo-small",
+    "Balanced": "musicgen-stereo-medium",
+    "Quality": "musicgen-stereo-large",
+}
+
+
+st.set_page_config(page_title="VidTune: Where Videos Find Their Melody", layout="centered")
+
+# Title and Description
+st.title("VidTune: Where Videos Find Their Melody")
+st.write("VidTune is a web application that allows users to upload videos and generate melodies matching the mood of the video.")
+
+
+# Sidebar
+st.sidebar.title("Settings")
+video_model = st.sidebar.selectbox("Select Video Descriptor", ["Fast", "Balanced", "Quality"], index=0)
+music_model = st.sidebar.selectbox("Select Music Generator", ["Fast", "Balanced", "Quality"], index=0)
+num_samples = st.sidebar.slider("Number of samples", 1, 8, 3)
+generate_button = st.sidebar.button("Generate Music")
+
+video_descriptor = DescribeVideo(model=video_model_map[video_model])
+audio_generator = GenerateAudio(model=music_model_map[music_model])
+
+video_description = None
+
+# Main Page (Page 1)
+if 'page' not in st.session_state:
+    st.session_state.page = 'main'
+
+if st.session_state.page == 'main':
+    st.header("Video to Music")
+    uploaded_video = st.file_uploader("Upload Video", type=["mp4"])
     
-    # Main Page (Page 1)
-    if 'page' not in st.session_state:
-        st.session_state.page = 'main'
+    if uploaded_video is not None:
+        st.session_state.uploaded_video = uploaded_video
+        with open("temp.mp4", mode='wb') as w:
+            w.write(uploaded_video.getvalue())
+        video_description = video_descriptor.describe_video("temp.mp4")
+        
+        st.session_state.page = 'video_to_music'
     
     if st.session_state.page == 'main':
-        st.header("Video to Music")
-        uploaded_video = st.file_uploader("Upload Video", type=["mp4"])
-        if uploaded_video is not None:
-            st.session_state.uploaded_video = uploaded_video
-            st.session_state.page = 'video_to_music'
-        
-        if st.session_state.page == 'main':
-            st.header("Prompt to Music")
-            prompt = st.text_area("Prompt")
-            if st.button("Generate"):
-                st.session_state.prompt = prompt
-                st.session_state.page = 'prompt_to_music'
+        st.header("Prompt to Music")
+        prompt = st.text_area("Prompt")
+        if generate_button:
+            st.session_state.prompt = prompt
+            st.session_state.page = 'prompt_to_music'
+
+# Page 2a (If the user uploads a video)
+if st.session_state.page == 'video_to_music':    
+    st.video(st.session_state.uploaded_video)
     
-    # Page 2a (If the user uploads a video)
-    if st.session_state.page == 'video_to_music':
-        st.sidebar.title("Settings")
-        device = st.sidebar.selectbox("Select Device", ["GPU", "CPU"], index=0)
-        num_samples = st.sidebar.slider("Number of samples", 1, 10, 3)
-        
-        st.video(st.session_state.uploaded_video)
-        
-        st.text_area("Video Description", "This is a fixed video description", disabled=True)
-        st.text_area("Music Description")
-        
-        if st.button("Generate Music"):
-            st.session_state.page = 'result'
-            st.session_state.device = device
-            st.session_state.num_samples = num_samples
+    st.text_area("Video Description", "This is a fixed video description", disabled=True)
+    st.text_area("Music Description")
     
-    # Page 2b (If user selects "Prompt to Music" in Page 1)
-    if st.session_state.page == 'prompt_to_music':
-        st.sidebar.title("Settings")
-        device = st.sidebar.selectbox("Select Device", ["GPU", "CPU"], index=0)
-        num_samples = st.sidebar.slider("Number of samples", 1, 10, 3)
-        
-        if st.button("Generate Music"):
-            st.session_state.page = 'result'
-            st.session_state.device = device
-            st.session_state.num_samples = num_samples
+    if generate_button:
+        st.session_state.page = 'result'
+        st.session_state.device = device
+        st.session_state.num_samples = num_samples
+
+# Page 2b (If user selects "Prompt to Music" in Page 1)
+if st.session_state.page == 'prompt_to_music':
+    st.sidebar.title("Settings")
+    device = st.sidebar.selectbox("Select Device", ["GPU", "CPU"], index=0)
+    num_samples = st.sidebar.slider("Number of samples", 1, 10, 3)
     
-    # Page 3 (Results Page)
-    if st.session_state.page == 'result':
-        st.header("Generated Music")
-        for i in range(st.session_state.num_samples):
-            st.write(f"Music Sample {i+1}")
-            st.audio(f"Generated Music {i+1}.mp3", format='audio/mp3')
-            st.download_button(f"Download Music {i+1}", f"Generated Music {i+1}.mp3")
-        
-        if st.button("Start Over"):
-            st.session_state.page = 'main'
+    if generate_button:
+        st.session_state.page = 'result'
+        st.session_state.device = device
+        st.session_state.num_samples = num_samples
 
-if __name__ == "__main__":
-    main()
+# Page 3 (Results Page)
+if st.session_state.page == 'result':
+    st.header("Generated Music")
+    for i in range(st.session_state.num_samples):
+        st.write(f"Music Sample {i+1}")
+        st.audio(f"Generated Music {i+1}.mp3", format='audio/mp3')
+        st.download_button(f"Download Music {i+1}", f"Generated Music {i+1}.mp3")
+    
+    if st.button("Start Over"):
+        st.session_state.page = 'main'