Update app.py

app.py

@@ -1,145 +1,106 @@
 import streamlit as st
-import cv2
+import imageio
 import numpy as np
-from transformers import (
-    pipeline,
-    AutoProcessor,
-    MusicgenForCausalLM,
-    set_seed,
-)
+from PIL import Image
+from transformers import AutoProcessor, BlipForConditionalGeneration, MusicgenForConditionalGeneration
+import soundfile as sf
 import torch
-import tempfile
 import os
-import ffmpeg
-from PIL import Image
-import warnings
-import io
-
-warnings.filterwarnings("ignore")
-
-def analyze_video(video_path):
-    """Analyze video with optimized BLIP processing"""
-    cap = cv2.VideoCapture(video_path)
-    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    sample_rate = max(frame_count // 30, 1)  # Increased frame sampling
-
-    frames = []
-    try:
-        for i in range(0, frame_count, sample_rate):
-            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
-            ret, frame = cap.read()
-            if ret:
-                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-                small_frame = cv2.resize(frame_rgb, (128, 128)) # Reduced resolution
-                frames.append(Image.fromarray(small_frame))
-    finally:
-        cap.release()
-
-    if not frames:
-        return ""
 
-    outputs = st.session_state.model(frames)
-    captions = [out[0]['generated_text'] for out in outputs]
+# Set page title
+st.title("Video Sound Effect Generator")
 
-    return " ".join(list(set(captions)))
-
-def generate_audio(description, duration):
-    """Optimized audio generation with corrected parameters"""
-    inputs = st.session_state.processor(
-        text=[description],
-        padding=True,
-        return_tensors="pt",
-    ).to(st.session_state.device)
-
-    set_seed(42)  # For reproducibility
-
-    max_new_tokens = int(duration * 16)  # 16 tokens per second
+# File uploader for video
+uploaded_file = st.file_uploader(
+    "Upload a short video (MP4, max 10 seconds, high resolution)", 
+    type=["mp4"]
+)
 
+if uploaded_file is not None:
     try:
-        audio = st.session_state.music_gen.generate(
-            **inputs,
-            do_sample=True,
-            top_k=250,
-            max_new_tokens=max_new_tokens,
-        )
-        if audio.shape[0] == 1 and len(audio.shape) == 3:
-            return audio[0][0].cpu().numpy()
-        else:
-            st.error(f"Unexpected audio shape: {audio.shape}")
-            return None
-
-    except torch.cuda.OutOfMemoryError:
-        st.error("Out of GPU memory. Try a shorter video or smaller model.")
-        return None
-    except Exception as e:
-        st.error(f"Error during audio generation: {e}")
-        return None
-
-def process_video(uploaded_file):
-    """Generates audio from video description"""
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        video_path = os.path.join(tmp_dir, "input.mp4")
-        with open(video_path, "wb") as f:
+        # Save the uploaded video temporarily
+        with open("temp_video.mp4", "wb") as f:
             f.write(uploaded_file.getbuffer())
-
-        probe = ffmpeg.probe(video_path)
-        duration = min(float(probe['format']['duration']), 10)
-
-        description = analyze_video(video_path)
-        audio_array = generate_audio(description, duration)
-
-        if audio_array is None:
-            return None
-
-        return audio_array
-
-# Streamlit UI
-st.set_page_config(page_title="Video Sound FX", layout="wide")
-st.title("🎬 Video Sound Generator (Audio Only)")
-
-if 'initialized' not in st.session_state:
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    if device.type == 'cpu':
-        st.warning("Running on CPU. Audio generation will be slow.")
-
-    st.session_state.model = pipeline(
-        "image-to-text",
-        model="Salesforce/blip-image-captioning-base",
-        device=device
-    )
-
-    st.session_state.music_gen = MusicgenForCausalLM.from_pretrained(
-        "facebook/musicgen-small"
-    ).to(device)
-
-    st.session_state.processor = AutoProcessor.from_pretrained(
-        "facebook/musicgen-small"
-    )
-
-    st.session_state.device = device
-    st.session_state.initialized = True
-
-uploaded_file = st.file_uploader("Upload video (MP4, max 10s and low resolution)", type=["mp4"])
-
-if uploaded_file and st.button("Generate Audio"):
-    with st.spinner("Processing..."):
-        try:
-            audio_array = process_video(uploaded_file)
-            if audio_array is not None:
-                st.success("Audio Generation Complete!")
-
-                # Convert numpy array to bytes for download
-                audio_bytes = (audio_array * 32767).astype(np.int16).tobytes()
-
+        
+        # Check video duration
+        video = imageio.get_reader("temp_video.mp4")
+        fps = video.get_meta_data()['fps']
+        num_frames = len(list(video.iter_data()))
+        duration = num_frames / fps
+        
+        if duration > 10:
+            st.error("Video is too long. Please upload a video of maximum 10 seconds.")
+        else:
+            st.success("Video uploaded successfully!")
+            
+            # Extract 10 evenly spaced frames
+            num_frames_to_extract = 10
+            step = max(1, num_frames // num_frames_to_extract)
+            frames = [
+                Image.fromarray(video.get_data(i)) 
+                for i in range(0, num_frames, step)
+            ][:num_frames_to_extract]
+            
+            # Load BLIP model with caching
+            @st.cache_resource
+            def load_blip_model():
+                processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+                model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+                return processor, model
+
+            processor, model = load_blip_model()
+            
+            # Generate text descriptions for each frame
+            descriptions = []
+            for frame in frames:
+                inputs = processor(images=frame, return_tensors="pt")
+                out = model.generate(**inputs)
+                description = processor.decode(out[0], skip_special_tokens=True)
+                descriptions.append(description)
+            
+            # Combine descriptions into a single prompt
+            text_prompt = ". ".join(descriptions)
+            st.write("Generated text prompt:", text_prompt)
+            
+            # Load MusicGen model with caching
+            @st.cache_resource
+            def load_musicgen_model():
+                processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+                model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+                return processor, model
+
+            musicgen_processor, musicgen_model = load_musicgen_model()
+            
+            # Generate sound effect
+            inputs = musicgen_processor(
+                text=[text_prompt],
+                padding=True,
+                return_tensors="pt",
+            )
+            audio_values = musicgen_model.generate(**inputs, max_new_tokens=512)
+            audio_array = audio_values[0].numpy()
+            sample_rate = musicgen_model.config.audio_encoder.sampling_rate
+            
+            # Save audio to a WAV file
+            sf.write("output.wav", audio_array, sample_rate)
+            
+            # Provide audio playback and download options
+            st.audio("output.wav", format="audio/wav")
+            with open("output.wav", "rb") as audio_file:
                 st.download_button(
-                    "Download Audio (WAV)",
-                    data=audio_bytes,
-                    file_name="generated_audio.wav",
+                    label="Download Sound Effect",
+                    data=audio_file,
+                    file_name="sound_effect.wav",
                     mime="audio/wav"
                 )
-                st.audio(audio_bytes, format='audio/wav')
-            else:
-                st.error("Error during audio generation. Please check the video and try again. Ensure video is max 10 seconds and low resolution.")
-
-        except Exception as e:
-            st.error(f"Error: {str(e)}")
+    
+    except Exception as e:
+        st.error(f"An error occurred: {str(e)}")
+        st.write("Please try uploading a different video or check your connection.")
+    
+    finally:
+        # Clean up temporary files
+        if os.path.exists("temp_video.mp4"):
+            os.remove("temp_video.mp4")
+        if os.path.exists("output.wav"):
+            os.remove("output.wav")