Pre-Final Changes, Release Ready

engine/audio_generator.py

@@ -1,5 +1,6 @@
 import os
 import warnings
+import traceback
 
 warnings.simplefilter("ignore")
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -9,13 +10,11 @@ import numpy as np
 from audiocraft.models import musicgen
 from scipy.io.wavfile import write as wav_write
 
-import logging
+try:
+    from logger import logging
+except:
+    import logging
 
-FORMAT = "%(asctime)s: %(levelname)s: %(message)s"
-logging.basicConfig(filename='logs.log', level=logging.INFO, format=FORMAT)
-STDERRLOGGER = logging.StreamHandler()
-STDERRLOGGER.setFormatter(logging.Formatter(FORMAT))
-logging.getLogger().addHandler(STDERRLOGGER)
 
 class GenerateAudio:
     def __init__(self, model="musicgen-stereo-small"):
@@ -32,7 +31,9 @@ class GenerateAudio:
             logging.info(f"Loaded model: {model}")
             return model
         except Exception as e:
-            logging.error(f"Failed to load model: {e}")
+            logging.error(
+                f"Failed to load model: {e}, Traceback: {traceback.format_exc()}"
+            )
             raise ValueError(f"Failed to load model: {e}")
             return
 
@@ -41,14 +42,18 @@ class GenerateAudio:
         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.")
+            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.")
+            logging.warning(
+                "Duration is greater than 30 seconds. Setting duration to 30 seconds."
+            )
             return 30
         return duration
 
@@ -62,16 +67,16 @@ class GenerateAudio:
             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
+            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.sampling_rate = self.model.sample_rate
             if duration <= 30:
                 self.model.set_generation_params(duration=duration)
                 result = self.model.generate(prompts, progress=False)
@@ -79,18 +84,23 @@ class GenerateAudio:
                 self.model.set_generation_params(duration=30)
                 result = self.model.generate(prompts, progress=False)
                 self.model.set_generation_params(duration=duration)
-                result = self.model.generate_with_chroma(prompts, result, melody_sample_rate=self.sampling_rate, progress=False)
+                result = self.model.generate_with_chroma(
+                    prompts,
+                    result,
+                    melody_sample_rate=self.sampling_rate,
+                    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}")
+            logging.error(
+                f"Failed to generate audio: {e}, Traceback: {traceback.format_exc()}"
+            )
             raise ValueError(f"Failed to generate audio: {e}")
 
     def save_audio(self, audio_dir="generated_audio"):
@@ -121,17 +131,18 @@ class GenerateAudio:
             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
+            "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}")
+    print(f"Audio buffers: {buffers}")

engine/video_descriptor.py

@@ -1,5 +1,6 @@
 import os
 from warnings import simplefilter
+import traceback
 
 simplefilter("ignore")
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -43,26 +44,38 @@ class DescribeVideo:
         self.safety_settings = self.get_safety_settings()
 
         genai.configure(api_key=__api_key)
-        self.mllm_model = genai.GenerativeModel(self.model, system_instruction=gemini_instructions)
+        self.mllm_model = genai.GenerativeModel(
+            self.model, system_instruction=gemini_instructions
+        )
 
         logging.info(f"Initialized DescribeVideo with model: {self.model}")
 
     def describe_video(self, video_path, genre, bpm, user_keywords):
         video_file = genai.upload_file(video_path)
-        logging.info(f"Uploaded video: {video_path}")
 
         while video_file.state.name == "PROCESSING":
             time.sleep(0.25)
             video_file = genai.get_file(video_file.name)
 
         if video_file.state.name == "FAILED":
-            logging.error(f"Failed to upload video: {video_file.state.name}")
+            logging.error(
+                f"Failed to upload video: {video_file.state.name}, Traceback: {traceback.format_exc()}"
+            )
             raise ValueError(f"Failed to upload video: {video_file.state.name}")
-        
-        additional_keywords = ", ".join([genre, user_keywords, bpm]) + "bpm"
+
+        additional_keywords = ", ".join(filter(None, [genre, user_keywords])) + (
+            f", {bpm} bpm" if bpm else ""
+        )
+
+        logging.info(f"Uploaded video: {video_path} and config: {additional_keywords}")
+
+        user_prompt = "Explain what is happening in this video."
+
+        if additional_keywords:
+            user_prompt += f" The following keywords are provided by the user for generating the music prompt: {additional_keywords}"
 
         response = self.mllm_model.generate_content(
-            [video_file, f"Explain what is happening in this video. The following keywords are provided by the user for generating the music prompt: {additional_keywords}"],
+            [video_file, user_prompt],
             request_options={"timeout": 600},
             safety_settings=self.safety_settings,
         )
@@ -116,7 +129,9 @@ class DescribeVideo:
 
         api_key = creds.get("google_api_key", None)
         if api_key is None or not isinstance(api_key, str):
-            logging.error(f"Google API key not found in {path}")
+            logging.error(
+                f"Google API key not found in {path}, Traceback: {traceback.format_exc()}"
+            )
             raise ValueError(f"Gemini API key not found in {path}")
         return api_key
 
@@ -129,7 +144,7 @@ class DescribeVideo:
 
         if model not in models:
             logging.error(
-                f"Invalid model name '{model}'. Valid options are: {', '.join(models.keys())}"
+                f"Invalid model name '{model}'. Valid options are: {', '.join(models.keys())}, Traceback: {traceback.format_exc()}"
             )
             raise ValueError(
                 f"Invalid model name '{model}'. Valid options are: {', '.join(models.keys())}"

main.py

@@ -1,10 +1,9 @@
 import streamlit as st
 from engine import DescribeVideo, GenerateAudio
 import os
-import time
 from moviepy.editor import VideoFileClip
 
-
+# Define model maps
 video_model_map = {
     "Fast": "flash",
     "Quality": "pro",
@@ -16,7 +15,13 @@ music_model_map = {
     "Quality": "musicgen-stereo-large",
 }
 
+# music_model_map = {
+#     "Fast": "facebook/musicgen-melody",
+#     "Quality": "facebook/musicgen-melody-large",
+# }
+
 genre_map = {
+    "None": None,
     "Pop": "Pop",
     "Rock": "Rock",
     "Hip Hop": "Hip-Hop/Rap",
@@ -30,7 +35,6 @@ genre_map = {
     "Lo-Fi": "Lo-Fi",
 }
 
-
 st.set_page_config(
     page_title="VidTune: Where Videos Find Their Melody", layout="centered"
 )
@@ -41,29 +45,90 @@ st.write(
     "VidTune is a web application that allows users to upload videos and generate melodies matching the mood of the video."
 )
 
+# Initialize session state for advanced settings and other inputs
+if "show_advanced" not in st.session_state:
+    st.session_state.show_advanced = False
+if "video_model" not in st.session_state:
+    st.session_state.video_model = "Fast"
+if "music_model" not in st.session_state:
+    st.session_state.music_model = "Fast"
+if "num_samples" not in st.session_state:
+    st.session_state.num_samples = 3
+if "music_genre" not in st.session_state:
+    st.session_state.music_genre = None
+if "music_bpm" not in st.session_state:
+    st.session_state.music_bpm = 100
+if "user_keywords" not in st.session_state:
+    st.session_state.user_keywords = None
 
 # Sidebar
 st.sidebar.title("Settings")
-video_model = st.sidebar.selectbox(
-    "Select Video Descriptor", ["Fast", "Quality"], index=0
+
+# Basic Settings
+st.session_state.video_model = st.sidebar.selectbox(
+    "Select Video Descriptor",
+    ["Fast", "Quality"],
+    index=["Fast", "Quality"].index(st.session_state.video_model),
+)
+st.session_state.music_model = st.sidebar.selectbox(
+    "Select Music Generator",
+    ["Fast", "Balanced", "Quality"],
+    index=["Fast", "Balanced", "Quality"].index(st.session_state.music_model),
 )
-music_model = st.sidebar.selectbox(
-    "Select Music Generator", ["Fast", "Balanced", "Quality"], index=0
+st.session_state.num_samples = st.sidebar.slider(
+    "Number of samples", 1, 5, st.session_state.num_samples
 )
-music_genre = st.sidebar.selectbox("Select Music Genre", list(genre_map.keys()))
-num_samples = st.sidebar.slider("Number of samples", 1, 5, 3)
+
+# Sidebar for advanced settings
+with st.sidebar:
+    # Create a placeholder for the advanced settings button
+    placeholder = st.empty()
+
+    # Button to toggle advanced settings
+    if placeholder.button("Advanced"):
+        st.session_state.show_advanced = not st.session_state.show_advanced
+        st.rerun()  # Refresh the layout after button click
+
+# Display advanced settings if enabled
+if st.session_state.show_advanced:
+    # Advanced settings
+    st.session_state.music_bpm = st.sidebar.slider("Beats Per Minute", 35, 180, 100)
+    st.session_state.music_genre = st.sidebar.selectbox(
+        "Select Music Genre",
+        list(genre_map.keys()),
+        index=(
+            list(genre_map.keys()).index(st.session_state.music_genre)
+            if st.session_state.music_genre in genre_map.keys()
+            else 0
+        ),
+    )
+    st.session_state.user_keywords = st.sidebar.text_input(
+        "User Keywords",
+        value=st.session_state.user_keywords,
+        help="Enter keywords separated by commas.",
+    )
+else:
+    st.session_state.music_genre = None
+    st.session_state.music_bpm = None
+    st.session_state.user_keywords = None
+
+# Generate Button
 generate_button = st.sidebar.button("Generate Music")
 
-video_descriptor = None
-audio_descriptor = None
-video_description = None
 
-# Initialize Video Descriptor and Audio Generator
-if video_descriptor is None or audio_descriptor is None:
-    video_descriptor = DescribeVideo(model=video_model_map[video_model])
-    audio_generator = GenerateAudio(model=music_model_map[music_model])
+# Cache the model loading
+@st.cache_resource
+def load_models(video_model_key, music_model_key):
+    video_descriptor = DescribeVideo(model=video_model_map[video_model_key])
+    audio_generator = GenerateAudio(model=music_model_map[music_model_key])
+    return video_descriptor, audio_generator
 
 
+# Load models
+video_descriptor, audio_generator = load_models(
+    st.session_state.video_model, st.session_state.music_model
+)
+
 # Video Uploader
 uploaded_video = st.file_uploader("Upload Video", type=["mp4"])
 if uploaded_video is not None:
@@ -80,11 +145,15 @@ if generate_button and uploaded_video is None:
     st.error("Please upload a video before generating music.")
     st.stop()
 
-
 # Submit Button and music generation if video is uploaded
 if generate_button and uploaded_video is not None:
     with st.spinner("Analyzing video..."):
-        video_description = video_descriptor.describe_video("temp.mp4")
+        video_description = video_descriptor.describe_video(
+            "temp.mp4",
+            genre=st.session_state.music_genre,
+            bpm=st.session_state.music_bpm,
+            user_keywords=st.session_state.user_keywords,
+        )
         video_duration = VideoFileClip("temp.mp4").duration
         music_prompt = video_description["Music Prompt"]
 
@@ -106,75 +175,15 @@ if generate_button and uploaded_video is not None:
 
     # Generate Music
     with st.spinner("Generating music..."):
-        music_prompt = [music_prompt] * num_samples
+        if video_duration > 30:
+            st.warning(
+                "Due to hardware limitations, the maximum music length is capped at 30 seconds."
+            )
+        music_prompt = [music_prompt] * st.session_state.num_samples
         audio_generator.generate_audio(music_prompt, duration=video_duration)
         audio_paths = audio_generator.save_audio()
         st.success("Music generated successfully.")
         for i, audio_path in enumerate(audio_paths):
             st.audio(audio_path, format="audio/wav")
-            # st.download_button(
-            #     label=f"Download Music {i+1}",
-            #     data=open(audio_path, "rb"),
-            #     file_name=f"Generated Music {i+1}.wav",
-            #     mime="audio/wav",
-            # )
-
-
-# # 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"])
-
-#     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("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)
-
-#     st.text_area(
-#         "Video Description", "This is a fixed video description", disabled=True
-#     )
-#     st.text_area("Music Description")
-
-#     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)
-
-#     if generate_button:
-#         st.session_state.page = "result"
-#         st.session_state.device = device
-#         st.session_state.num_samples = num_samples
-
-# # 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"
+
+        st.balloons()