app.py

import streamlit as st
from PIL import Image
import numpy as np
import torch
import gc
import os
import tempfile
import math
import imageio
import traceback

# --- Attempt to import moviepy for video processing ---
try:
    import moviepy.editor as mpy
    MOVIEPY_AVAILABLE = True
except (ImportError, OSError) as e:
    MOVIEPY_AVAILABLE = False
    st.warning(
        "MoviePy library is not available or ffmpeg is missing. "
        "Video syncing features will be disabled. "
        "If running locally, install with: pip install moviepy. Ensure ffmpeg is installed."
    )
    print(f"MoviePy load error: {e}")


# --- Model Configuration ---
IMAGE_CAPTION_MODEL = "Salesforce/blip-image-captioning-base"
AUDIO_GEN_MODEL = "facebook/musicgen-small"

# --- Constants ---
DEFAULT_NUM_FRAMES = 2        # Fewer frames for faster processing on free tier
DEFAULT_AUDIO_DURATION_S = 5  # Shorter audio for faster generation
MAX_FRAMES_TO_SHOW_UI = 3
DEVICE = torch.device("cpu") # Explicitly use CPU for Hugging Face free tier

# --- Page Setup ---
st.set_page_config(page_title="AI Video Sound Designer (HF Space)", layout="wide", page_icon="🎬")

st.title("🎬 AI Video Sound Designer (for Hugging Face Spaces)")
st.markdown("""
    Upload a short MP4 video. The tool will:
    1. Extract frames from the video.
    2. Analyze frames using an image captioning model to generate sound ideas.
    3. Synthesize audio using MusicGen based on these ideas.
    4. Optionally, combine the new audio with your video.
    ---
    **Note:** Processing on CPU (especially audio generation) can be slow. Please be patient!
""")

# --- Utility Functions ---
def clear_memory(model_obj=None, processor_obj=None):
    """Clears model objects from memory and runs garbage collection."""
    if model_obj:
        del model_obj
    if processor_obj:
        del processor_obj
    gc.collect()
    if torch.cuda.is_available(): # Though we target CPU, good practice
        torch.cuda.empty_cache()
    print("Memory cleared.")

@st.cache_resource
def load_image_caption_model_and_processor():
    """Loads the image captioning model and processor."""
    try:
        from transformers import BlipProcessor, BlipForConditionalGeneration
        st.write(f"Loading Image Captioning Model: {IMAGE_CAPTION_MODEL} (this might take a moment)...")
        processor = BlipProcessor.from_pretrained(IMAGE_CAPTION_MODEL)
        model = BlipForConditionalGeneration.from_pretrained(IMAGE_CAPTION_MODEL).to(DEVICE)
        st.toast("Image Captioning model loaded!", icon="🖼️")
        return processor, model
    except Exception as e:
        st.error(f"Error loading image captioning model: {e}")
        st.error(traceback.format_exc())
        return None, None

@st.cache_resource
def load_audio_gen_model_and_processor():
    """Loads the audio generation model and processor."""
    try:
        from transformers import AutoProcessor, MusicgenForConditionalGeneration
        st.write(f"Loading Audio Generation Model: {AUDIO_GEN_MODEL} (this might take a while on CPU)...")
        processor = AutoProcessor.from_pretrained(AUDIO_GEN_MODEL)
        model = MusicgenForConditionalGeneration.from_pretrained(AUDIO_GEN_MODEL).to(DEVICE)
        st.toast("Audio Generation model loaded! (CPU generation will be slow)", icon="🎶")
        return processor, model
    except Exception as e:
        st.error(f"Error loading audio generation model: {e}")
        st.error(traceback.format_exc())
        return None, None

def extract_frames_from_video(video_path, num_frames):
    """Extracts a specified number of frames evenly from a video."""
    frames = []
    reader = None
    try:
        reader = imageio.get_reader(video_path, "ffmpeg")
        total_frames = reader.count_frames()
        if total_frames == 0: # If count_frames fails, try metadata
            meta = reader.get_meta_data()
            duration = meta.get('duration')
            fps = meta.get('fps', 25)
            if duration:
                total_frames = int(duration * fps)
            else: # Fallback if duration isn't available
                 st.warning("Could not determine video length. Will attempt to read initial frames.")
                 # Try to read a few frames anyway if count fails
                 for i, frame_data in enumerate(reader):
                     if i < num_frames * 5: # Read a bit more than needed to find distinct frames
                         frames.append(Image.fromarray(frame_data).convert("RGB"))
                     if len(frames) >= num_frames:
                         break
                 if reader: reader.close()
                 return frames[::len(frames)//num_frames] if frames else []


        if total_frames < num_frames:
            indices = np.arange(total_frames)
        else:
            indices = np.linspace(0, total_frames - 1, num_frames, dtype=int, endpoint=True)
        
        actual_frames_extracted = 0
        for i in indices:
            if actual_frames_extracted >= num_frames:
                break
            try:
                frame_data = reader.get_data(i)
                frames.append(Image.fromarray(frame_data).convert("RGB"))
                actual_frames_extracted +=1
            except Exception as e:
                st.warning(f"Skipping problematic frame {i}: {e}")
                continue
        return frames
    except (imageio.core.fetching.NeedDownloadError, OSError) as e_ffmpeg:
        st.error(f"FFmpeg not found or failed: {e_ffmpeg}. Please ensure ffmpeg is installed and in PATH if running locally.")
        return []
    except Exception as e:
        st.error(f"Error extracting frames: {e}")
        st.error(traceback.format_exc())
        return []
    finally:
        if reader:
            reader.close()

def generate_sound_prompt_from_frames(frames, caption_processor, caption_model):
    """Generates sound descriptions from frames using BLIP."""
    if not frames:
        return "ambient background noise"
    
    descriptions = []
    instruction = "A short description of this image, focusing on elements that might produce sound:"
    
    with st.spinner(f"Generating sound ideas from {len(frames)} frames..."):
        for i, frame in enumerate(frames):
            try:
                inputs = caption_processor(images=frame, text=instruction, return_tensors="pt").to(DEVICE)
                # For BLIP, generate is typically used like this.
                # You might need to adjust max_length based on desired description length.
                generated_ids = caption_model.generate(**inputs, max_length=50) # Keep descriptions short
                description = caption_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
                if description:
                    descriptions.append(description)
                st.progress((i + 1) / len(frames), text=f"Frame {i+1}/{len(frames)} analyzed.")
            except Exception as e:
                st.warning(f"Could not get description for a frame: {e}")
                continue
    
    if not descriptions:
        return "general ambiance, subtle environmental sounds" # Fallback
    
    # Simple combination: join unique descriptions
    unique_descriptions = list(dict.fromkeys(descriptions))
    combined_prompt = ". ".join(unique_descriptions)
    # Further processing to make it more like a sound design brief
    final_prompt = f"Sounds for a scene featuring: {combined_prompt}. Focus on atmosphere, key sound events, and textures."
    return final_prompt

def generate_audio_from_prompt(prompt, duration_s, audio_processor, audio_model, guidance, temp):
    """Generates audio using MusicGen."""
    try:
        inputs = audio_processor(text=[prompt], return_tensors="pt", padding=True).to(DEVICE)
        
        # MusicGen has a max sequence length for the prompt, often around 2048 tokens.
        # Forcing it to 512 to be safe on CPU and for typical descriptions.
        # The processor handles truncation.
        if inputs.input_ids.shape[1] > 512:
             st.warning(f"Prompt is long ({inputs.input_ids.shape[1]} tokens), might be truncated by MusicGen.")
        # inputs['input_ids'] = inputs['input_ids'][:, :512]
        # inputs['attention_mask'] = inputs['attention_mask'][:, :512]


        # Calculate max_new_tokens based on duration and model's token/sec rate
        # musicgen-small typically 50 tokens/second. Max output length ~2048 tokens.
        tokens_per_second = audio_model.config.audio_encoder.token_per_second # typically 50 for musicgen
        max_new_tokens = min(int(duration_s * tokens_per_second), 1500) # Cap at 1500 (30s) as a practical limit

        with st.spinner(f"Synthesizing {duration_s}s audio... (CPU: This will take several minutes!)"):
            # For CPU, do_sample=False might be faster but less diverse. Try True first.
            audio_values = audio_model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                guidance_scale=guidance,
                temperature=temp,
                # No pad_token_id for MusicGen's generate function, it uses eos_token_id for padding by default if needed
            )
        
        audio_array = audio_values[0, 0].cpu().numpy()
        sampling_rate = audio_model.config.audio_encoder.sampling_rate
        
        # Normalize
        if np.abs(audio_array).max() > 0:
            audio_array = audio_array / np.abs(audio_array).max() * 0.9
        return audio_array, sampling_rate
    except Exception as e:
        st.error(f"Error generating audio: {e}")
        st.error(traceback.format_exc())
        return None, None

def combine_audio_video(video_path, audio_array, sampling_rate, mix_original):
    """Combines generated audio with the video using MoviePy."""
    if not MOVIEPY_AVAILABLE:
        st.error("MoviePy is not available. Cannot combine audio and video.")
        return None

    output_video_path = None
    temp_audio_path = None
    video_clip = None
    generated_audio_clip = None
    final_clip = None

    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_audio:
            # Scipy.io.wavfile can also be used here, or soundfile
            import scipy.io.wavfile
            scipy.io.wavfile.write(tmp_audio.name, sampling_rate, audio_array)
            temp_audio_path = tmp_audio.name
        
        video_clip = mpy.VideoFileClip(video_path)
        generated_audio_clip = mpy.AudioFileClip(temp_audio_path)

        # Loop or trim generated audio to match video duration
        if generated_audio_clip.duration < video_clip.duration:
            generated_audio_clip = generated_audio_clip.fx(mpy.afx.audio_loop, duration=video_clip.duration)
        elif generated_audio_clip.duration > video_clip.duration:
            generated_audio_clip = generated_audio_clip.subclip(0, video_clip.duration)

        final_audio = generated_audio_clip
        if mix_original and video_clip.audio:
            # Adjust volumes for mixing
            original_audio = video_clip.audio.volumex(0.5) # Lower original audio
            generated_audio = generated_audio_clip.volumex(0.8) # Keep generated slightly louder
            final_audio = mpy.CompositeAudioClip([original_audio, generated_audio])
            final_audio = final_audio.set_duration(video_clip.duration) # Ensure composite duration matches

        final_clip = video_clip.set_audio(final_audio)
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video_out:
            output_video_path = tmp_video_out.name

        final_clip.write_videofile(
            output_video_path,
            codec="libx264",
            audio_codec="aac",
            temp_audiofile_path=os.path.dirname(temp_audio_path), # Ensure moviepy can write temp audio here
            threads=2, # Limit threads on free tier
            logger=None # or 'bar' for progress
        )
        return output_video_path

    except Exception as e:
        st.error(f"Error combining audio and video: {e}")
        st.error(traceback.format_exc())
        return None
    finally:
        # Close clips to release resources
        if video_clip: video_clip.close()
        if generated_audio_clip: generated_audio_clip.close()
        # if final_clip: final_clip.close() # final_clip is usually the same as video_clip with modified audio

        if temp_audio_path and os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)
        # The output_video_path is handled by the caller (downloaded, then potentially cleaned up)

# --- Sidebar for Settings ---
with st.sidebar:
    st.header("⚙️ Settings")
    num_frames_analysis = st.slider("Number of Frames to Analyze", 1, 5, DEFAULT_NUM_FRAMES, 1,
                                   help="More frames provide more context but increase analysis time.")
    audio_duration = st.slider("Target Audio Duration (seconds)", 3, 15, DEFAULT_AUDIO_DURATION_S, 1,
                               help="Shorter durations generate much faster on CPU.")
    
    st.subheader("MusicGen Parameters")
    guidance = st.slider("Guidance Scale (MusicGen)", 1.0, 7.0, 3.0, 0.5,
                         help="Higher values make audio follow prompt more closely. Default is 3.0.")
    temperature = st.slider("Temperature (MusicGen)", 0.5, 1.5, 1.0, 0.1,
                            help="Controls randomness. Higher is more diverse. Default is 1.0.")
    
    if MOVIEPY_AVAILABLE:
        st.subheader("Video Output")
        mix_audio = st.checkbox("Mix with original video audio", value=False)
    else:
        mix_audio = False # Disable if moviepy not available


# --- Main Application Logic ---
uploaded_file = st.file_uploader("📤 Upload your MP4 video file (short clips recommended):", type=["mp4", "mov", "avi"])

# Initialize session state for generated file paths
if 'generated_audio_file' not in st.session_state:
    st.session_state.generated_audio_file = None
if 'output_video_file' not in st.session_state:
    st.session_state.output_video_file = None


if uploaded_file is not None:
    st.video(uploaded_file)
    
    # Use a button to trigger processing
    if st.button("✨ Generate Sound Design!", type="primary", use_container_width=True):
        # --- Clear previous results ---
        if st.session_state.generated_audio_file and os.path.exists(st.session_state.generated_audio_file):
            os.remove(st.session_state.generated_audio_file)
        st.session_state.generated_audio_file = None
        if st.session_state.output_video_file and os.path.exists(st.session_state.output_video_file):
            os.remove(st.session_state.output_video_file)
        st.session_state.output_video_file = None
        clear_memory()

        video_bytes = uploaded_file.read()
        temp_video_path = None

        try:
            with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_vid:
                tmp_vid.write(video_bytes)
                temp_video_path = tmp_vid.name
            
            # === Stage 1: Frame Extraction ===
            st.subheader("1. Extracting Frames")
            with st.spinner("Extracting frames from video..."):
                frames = extract_frames_from_video(temp_video_path, num_frames_analysis)
            
            if not frames:
                st.error("No frames extracted. Cannot proceed.")
                st.stop()
            
            st.success(f"Extracted {len(frames)} frames.")
            if frames:
                cols_to_show = min(len(frames), MAX_FRAMES_TO_SHOW_UI)
                if cols_to_show > 0:
                    st.write("Sampled Frames:")
                    cols = st.columns(cols_to_show)
                    for i, frame_img in enumerate(frames[:cols_to_show]):
                        cols[i].image(frame_img, caption=f"Frame {i+1}", use_column_width=True)
            
            # === Stage 2: Image Captioning (Sound Prompt Generation) ===
            st.subheader("2. Generating Sound Ideas (Image Analysis)")
            caption_processor, caption_model = load_image_caption_model_and_processor()
            if caption_processor and caption_model:
                sound_prompt = generate_sound_prompt_from_frames(frames, caption_processor, caption_model)
                st.info(f"✍️ **Generated Sound Prompt:** {sound_prompt}")
                
                # Unload captioning model immediately
                clear_memory(caption_model, caption_processor)
            else:
                st.error("Failed to load image captioning model. Using a default prompt.")
                sound_prompt = "ambient nature sounds with a gentle breeze" # Fallback
            
            # === Stage 3: Audio Generation ===
            st.subheader("3. Synthesizing Audio (MusicGen)")
            st.warning("🎧 Audio generation on CPU can take several minutes. Please be patient!")
            audio_processor, audio_model = load_audio_gen_model_and_processor()
            generated_audio_array, sr = None, None # Initialize

            if audio_processor and audio_model:
                generated_audio_array, sr = generate_audio_from_prompt(sound_prompt, audio_duration, audio_processor, audio_model, guidance, temperature)
                # Unload audio model immediately
                clear_memory(audio_model, audio_processor)
            else:
                st.error("Failed to load audio generation model. Cannot generate audio.")

            if generated_audio_array is not None and sr is not None:
                st.success("Audio generated!")
                st.audio(generated_audio_array, sample_rate=sr)
                
                # Save audio for download
                with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_audio_out:
                    import scipy.io.wavfile # or soundfile
                    scipy.io.wavfile.write(tmp_audio_out.name, sr, generated_audio_array)
                    st.session_state.generated_audio_file = tmp_audio_out.name
                
                with open(st.session_state.generated_audio_file, "rb") as f:
                    st.download_button(
                        "📥 Download Generated Audio (.wav)", 
                        f, 
                        file_name="generated_sound.wav",
                        mime="audio/wav"
                    )

                # === Stage 4: (Optional) Video and Audio Syncing ===
                if MOVIEPY_AVAILABLE:
                    st.subheader("4. Combining Audio with Video")
                    with st.spinner("Processing video with new audio... (can be slow)"):
                        output_video_file_path = combine_audio_video(temp_video_path, generated_audio_array, sr, mix_audio)
                    
                    if output_video_file_path and os.path.exists(output_video_file_path):
                        st.success("Video processing complete!")
                        st.video(output_video_file_path)
                        st.session_state.output_video_file = output_video_file_path
                        
                        with open(output_video_file_path, "rb") as f_vid:
                            st.download_button(
                                "🎬 Download Video with New Sound (.mp4)",
                                f_vid,
                                file_name="video_with_new_sound.mp4",
                                mime="video/mp4"
                            )
                    elif MOVIEPY_AVAILABLE: # Only show error if moviepy was expected to work
                         st.error("Failed to combine audio and video.")
            else:
                st.error("Audio generation failed. Cannot proceed to video syncing.")

        except Exception as e:
            st.error(f"An unexpected error occurred in the main processing pipeline: {e}")
            st.error(traceback.format_exc())
        finally:
            if temp_video_path and os.path.exists(temp_video_path):
                os.remove(temp_video_path)
            # Models are cleared within their stages using clear_memory()
            # Generated download files (audio/video) are kept in session_state until next run or session ends
            print("Main processing finished or errored. Temp video (if any) cleaned up.")
            clear_memory() # Final catch-all clear

    # Display download buttons if files were generated in a previous run within the session
    elif st.session_state.generated_audio_file and os.path.exists(st.session_state.generated_audio_file):
        st.markdown("---")
        st.write("Previously generated audio:")
        st.audio(st.session_state.generated_audio_file)
        with open(st.session_state.generated_audio_file, "rb") as f:
            st.download_button(
                "📥 Download Previously Generated Audio (.wav)",
                f,
                file_name="generated_sound_previous.wav",
                mime="audio/wav",
                key="prev_audio_dl"
            )
    if st.session_state.output_video_file and os.path.exists(st.session_state.output_video_file) and MOVIEPY_AVAILABLE:
        st.markdown("---")
        st.write("Previously generated video with new sound:")
        st.video(st.session_state.output_video_file)
        with open(st.session_state.output_video_file, "rb") as f_vid:
            st.download_button(
                "🎬 Download Previously Generated Video (.mp4)",
                f_vid,
                file_name="video_with_new_sound_previous.mp4",
                mime="video/mp4",
                key="prev_video_dl"
            )

else:
    st.info("☝️ Upload a video to get started.")

st.markdown("---")
st.markdown("Made for Hugging Face Spaces. Model loading & generation can be slow on CPU.")