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

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"
@@ -30,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
 
@@ -39,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
 
@@ -60,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)
@@ -77,17 +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"):
@@ -118,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,7 +1,9 @@
 import streamlit as st
 from engine import DescribeVideo, GenerateAudio
+import os
+from moviepy.editor import VideoFileClip
 
-
+# Define model maps
 video_model_map = {
     "Fast": "flash",
     "Quality": "pro",
@@ -13,79 +15,175 @@ 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",
+    "Jazz": "Jazz",
+    "Classical": "Classical",
+    "Blues": "Blues",
+    "Country": "Country",
+    "EDM": "Electronic/Dance",
+    "Metal": "Metal",
+    "Disco": "Disco",
+    "Lo-Fi": "Lo-Fi",
+}
 
-st.set_page_config(page_title="VidTune: Where Videos Find Their Melody", layout="centered")
+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.")
+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", "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)
+
+# 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),
+)
+st.session_state.num_samples = st.sidebar.slider(
+    "Number of samples", 1, 5, st.session_state.num_samples
+)
+
+# 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 = 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"])
-    
-    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'
+
+# 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:
+    st.session_state.uploaded_video = uploaded_video
+    with open("temp.mp4", mode="wb") as w:
+        w.write(uploaded_video.getvalue())
+
+# Video Player
+if os.path.exists("temp.mp4") and uploaded_video is not None:
+    st.video(uploaded_video)
+
+# Submit button if video is not uploaded
+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",
+            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"]
+
+        st.success("Video description generated successfully.")
+
+        # Display Video Description and Music Prompt
+        st.text_area(
+            "Video Description",
+            video_description["Content Description"],
+            disabled=True,
+            height=120,
+        )
+        music_prompt = st.text_area(
+            "Music Prompt",
+            music_prompt,
+            disabled=False,
+            height=120,
+        )
+
+    # Generate Music
+    with st.spinner("Generating music..."):
+        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.balloons()