Create app.py

app.py

@@ -0,0 +1,189 @@
+import streamlit as st
+import imageio
+import numpy as np
+from PIL import Image
+from transformers import pipeline
+import soundfile as sf
+import torch
+import os
+import tempfile
+import math
+import gc  # Garbage collector
+
+# Try importing moviepy, with fallback
+try:
+    import moviepy.editor as mpy
+except ModuleNotFoundError:
+    st.error("The 'moviepy' library is not installed. Please install it (`pip install moviepy`) 🚨")
+    st.stop()
+except OSError as e:
+    st.error(f"Error initializing moviepy: {e} 🚨")
+    st.warning("Ensure ffmpeg is installed and accessible in your PATH.")
+
+# --- Constants & Defaults ---
+MODEL_MUSICGEN = "facebook/musicgen-small"
+MODEL_MOONDREAM = "vikhyatk/moondream2"
+DEFAULT_AUDIO_DURATION_S = 15
+DEFAULT_FRAMES = 5
+DEFAULT_GUIDANCE = 5.0
+DEFAULT_TEMPERATURE = 0.8
+MAX_FRAMES_TO_SHOW = 5
+
+# --- Page Config ---
+st.set_page_config(
+    page_title="AI Video Sound Designer (Moondream2)",
+    page_icon="🎬",
+    layout="wide"
+)
+
+# --- Cached Loaders ---
+@st.cache_resource
+def load_moondream2():
+    return pipeline(
+        "image-text-to-text", 
+        model=MODEL_MOONDREAM, 
+        trust_remote_code=True,
+        device=0 if torch.cuda.is_available() else -1
+    )
+
+@st.cache_resource
+def load_musicgen_model():
+    from transformers import AutoProcessor, MusicgenForConditionalGeneration
+    processor = AutoProcessor.from_pretrained(MODEL_MUSICGEN)
+    model = MusicgenForConditionalGeneration.from_pretrained(MODEL_MUSICGEN)
+    if torch.cuda.is_available():
+        model = model.half().to("cuda")
+    return processor, model
+
+# --- Utilities ---
+def clear_gpu_memory():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+
+# --- Frame Extraction ---
+def extract_frames(video_path, num_frames):
+    frames = []
+    reader = imageio.get_reader(video_path, "ffmpeg")
+    meta = reader.get_meta_data()
+    fps = meta.get('fps', 24)
+    duration = meta.get('duration', 5)
+    total = int(fps * duration)
+    indices = np.linspace(0, total - 1, num_frames, dtype=int)
+    for i in indices:
+        frames.append(Image.fromarray(reader.get_data(i)).convert("RGB"))
+    reader.close()
+    return frames
+
+# --- Sound Prompt Generation ---
+def generate_sound_prompt(frames, pipe):
+    instruction = (
+        "Describe only the sounds implied by this image: ambient noise, textures, "
+        "actions producing sound, atmosphere. Be concise but evocative."
+    )
+    descriptions = []
+    for frame in frames:
+        out = pipe(image=frame, text=instruction)
+        text = out[0].get('generated_text', '').strip()
+        if text:
+            descriptions.append(text)
+    combined = "; ".join(dict.fromkeys(descriptions))
+    return combined or "ambient background noise"
+
+# --- Audio Generation ---
+def generate_audio(prompt, duration, processor, model, guidance, temp):
+    device = next(model.parameters()).device
+    inputs = processor(text=[prompt], return_tensors="pt", padding=True).to(device)
+    inputs = {k: (v.to(model.dtype) if v.dtype.is_floating_point else v) for k, v in inputs.items()}
+    tokens_per_sec = 50
+    max_tokens = min(int(duration * tokens_per_sec), 1500)
+    with torch.inference_mode():
+        audio_tensor = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            guidance_scale=guidance,
+            temperature=temp,
+            pad_token_id=model.config.eos_token_id
+        )
+    arr = audio_tensor[0].cpu().float().numpy()
+    peak = np.max(np.abs(arr)) or 1e-6
+    arr = np.clip(arr / peak * 0.9, -1.0, 1.0)
+    clear_gpu_memory()
+    return arr, model.config.audio_encoder.sampling_rate
+
+# --- Sync Audio/Video ---
+def sync_audio_video(video_path, audio_arr, sr, mix):
+    tmp_wav = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
+    sf.write(tmp_wav.name, audio_arr, sr)
+    video = mpy.VideoFileClip(video_path)
+    sound = mpy.AudioFileClip(tmp_wav.name)
+    # Loop or trim to match
+    if sound.duration < video.duration:
+        loops = math.ceil(video.duration / sound.duration)
+        sound = mpy.concatenate_audioclips([sound] * loops).subclip(0, video.duration)
+    else:
+        sound = sound.subclip(0, video.duration)
+    if mix and video.audio:
+        sound = mpy.CompositeAudioClip([video.audio.volumex(0.5), sound.volumex(0.5)])
+    final = video.set_audio(sound)
+    out = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
+    final.write_videofile(out, codec='libx264', audio_codec='aac', threads=os.cpu_count())
+    video.close(); sound.close(); final.close()
+    os.remove(tmp_wav.name)
+    clear_gpu_memory()
+    return out
+
+# --- UI ---
+st.title("🎬 AI Video Sound Designer (Moondream2)")
+st.markdown("Upload an MP4 video to generate immersive sound effects or download standalone audio.")
+
+uploaded = st.file_uploader("Upload MP4 Video", type=['mp4'])
+# Sidebar Settings
+with st.sidebar:
+    st.header("Settings")
+    n_frames = st.slider("Frames to analyze", 1, 10, DEFAULT_FRAMES)
+    duration = st.slider("Audio Duration (s)", 5, 30, DEFAULT_AUDIO_DURATION_S)
+    mix = st.checkbox("Mix with original audio", False)
+    gs = st.slider("Guidance Scale", 1.0, 10.0, DEFAULT_GUIDANCE, 0.5)
+    temp = st.slider("Temperature", 0.1, 2.0, DEFAULT_TEMPERATURE, 0.1)
+
+if uploaded:
+    with st.spinner("Extracting frames and analyzing visuals..."):
+        tmp_vid = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+        tmp_vid.write(uploaded.getbuffer()); tmp_vid.flush()
+        frames = extract_frames(tmp_vid.name, n_frames)
+        st.subheader("Sample Frames")
+        for col, img in zip(st.columns(min(len(frames), MAX_FRAMES_TO_SHOW)), frames):
+            col.image(img, use_column_width=True)
+
+    # Sequential: first Moondream
+    moondream = load_moondream2()
+    prompt = generate_sound_prompt(frames, moondream)
+    del moondream; clear_gpu_memory()
+    st.info(f"🧠 Sound Prompt: {prompt}")
+
+    # Then MusicGen
+    with st.spinner("Synthesizing audio..."):
+        proc, model = load_musicgen_model()
+        audio_arr, sr = generate_audio(prompt, duration, proc, model, gs, temp)
+
+    # Playback & Downloads
+    st.subheader("Generated Sound Effect")
+    st.audio(audio_arr, sr)
+    # Save to temp file
+    wav_tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
+    sf.write(wav_tmp.name, audio_arr, sr)
+    with open(wav_tmp.name, 'rb') as f:
+        st.download_button("Download Audio Only", f, file_name='sound_effect.wav')
+
+    # Video sync
+    with st.spinner("Syncing audio with video..."):
+        out_video = sync_audio_video(tmp_vid.name, audio_arr, sr, mix)
+        st.video(out_video)
+        with open(out_video, 'rb') as vf:
+            st.download_button("Download Video with Sound", vf, file_name='sound_designed.mp4')
+    # Cleanup
+    os.remove(tmp_vid.name)
+else:
+    st.info("Upload a video above to get started.")