Update app.py

app.py

@@ -1,526 +1,464 @@
 import streamlit as st
-import imageio
-import numpy as np
 from PIL import Image
-from transformers import pipeline, AutoProcessor, MusicgenForConditionalGeneration
-import soundfile as sf
+import numpy as np
 import torch
+import gc
 import os
 import tempfile
 import math
-import gc  # Garbage collector
-import traceback # For detailed error logging
+import imageio
+import traceback
 
-# Try importing moviepy, with fallback
+# --- Attempt to import moviepy for video processing ---
 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}. This might be due to a missing ffmpeg. 🚨")
-    st.warning("Ensure ffmpeg is installed and accessible in your system's PATH (e.g., `sudo apt-get install ffmpeg` or `conda install ffmpeg`).")
-    st.stop()
-
-# --- Constants & Defaults ---
-MODEL_MUSICGEN = "facebook/musicgen-small"
-MODEL_MOONDREAM = "vikhyatk/moondream2" # Using official Moondream2
-DEFAULT_AUDIO_DURATION_S = 10 # Reduced default for faster testing
-DEFAULT_FRAMES_TO_ANALYZE = 3 # Reduced for faster default processing
-DEFAULT_GUIDANCE = 3.0 # MusicGen default
-DEFAULT_TEMPERATURE = 1.0 # MusicGen default
-MAX_FRAMES_TO_SHOW_UI = 3
-MAX_MOONDREAM_DESCRIPTION_TOKENS = 70 # For concise sound descriptions
-MUSICGEN_PROMPT_MAX_TOKENS = 1024 # Safety limit for MusicGen prompt
+    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}")
+
 
-IS_CUDA_AVAILABLE = torch.cuda.is_available()
-DEVICE = torch.device("cuda:0" if IS_CUDA_AVAILABLE else "cpu")
+# --- Model Configuration ---
+IMAGE_CAPTION_MODEL = "Salesforce/blip-image-captioning-base"
+AUDIO_GEN_MODEL = "facebook/musicgen-small"
 
-# --- Page Config ---
-st.set_page_config(
-    page_title="AI Video Sound Designer",
-    page_icon="🎬",
-    layout="wide"
-)
+# --- 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.")
 
-# --- Cached Model Loaders ---
 @st.cache_resource
-def load_moondream_model():
+def load_image_caption_model_and_processor():
+    """Loads the image captioning model and processor."""
     try:
-        model_kwargs = {"torch_dtype": torch.float16} if IS_CUDA_AVAILABLE else {}
-        pipe = pipeline(
-            "image-to-text", # Correct pipeline type for Moondream-like models
-            model=MODEL_MOONDREAM,
-            trust_remote_code=True, # Moondream requires this
-            model_kwargs=model_kwargs,
-            device=DEVICE
-        )
-        st.toast(f"Moondream2 loaded on: {str(DEVICE).upper()} ({model_kwargs.get('torch_dtype', 'float32')})", icon="🤖")
-        return pipe
+        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 Moondream2 model ({MODEL_MOONDREAM}): {e}")
+        st.error(f"Error loading image captioning model: {e}")
         st.error(traceback.format_exc())
-        return None
+        return None, None
 
 @st.cache_resource
-def load_musicgen_model_and_processor():
+def load_audio_gen_model_and_processor():
+    """Loads the audio generation model and processor."""
     try:
-        processor = AutoProcessor.from_pretrained(MODEL_MUSICGEN)
-        model = MusicgenForConditionalGeneration.from_pretrained(MODEL_MUSICGEN)
-        
-        dtype_str = "float32 (full)"
-        if IS_CUDA_AVAILABLE:
-            model = model.half().to(DEVICE)
-            dtype_str = "float16 (half)"
-        else:
-            model = model.to(DEVICE) # Ensure it's on CPU if not CUDA
-            
-        st.toast(f"MusicGen ({MODEL_MUSICGEN}) loaded on: {str(DEVICE).upper()} with {dtype_str} precision.", icon="🎶")
+        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 MusicGen model ({MODEL_MUSICGEN}): {e}")
+        st.error(f"Error loading audio generation model: {e}")
         st.error(traceback.format_exc())
         return None, None
 
-# --- Utilities ---
-def clear_gpu_memory():
-    if IS_CUDA_AVAILABLE:
-        torch.cuda.empty_cache()
-    gc.collect()
-
-# --- Frame Extraction ---
-def extract_frames(video_path: str, num_frames_to_extract: int) -> list[Image.Image]:
+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)
         
-        total_video_frames = 0
-        # Try to get frame count using different methods for robustness
-        try:
-            total_video_frames = reader.count_frames()
-        except Exception: # count_frames might not be implemented or fail
-            pass
-
-        if not isinstance(total_video_frames, (int, float)) or total_video_frames <= 0:
-            meta_data = reader.get_meta_data()
-            total_video_frames = meta_data.get('nframes') # Check 'nframes' in metadata
-            if not isinstance(total_video_frames, (int, float)) or total_video_frames <= 0:
-                fps = meta_data.get('fps', 25)
-                duration = meta_data.get('duration')
-                if duration and fps:
-                    total_video_frames = int(fps * duration)
-                else: # Cannot determine length
-                    st.warning("Could not reliably determine video length. Frame selection may be suboptimal.")
-                    # Fallback: try to read up to a certain number of frames to estimate
-                    # This part can be complex, for now, assume if above fails, it's problematic.
-                    total_video_frames = 0 # Indicate failure to determine
-
-        if total_video_frames < 1:
-            st.error("Video appears to have 0 frames or its length could not be determined accurately.")
-            if reader: reader.close()
-            return []
-
-        num_to_sample = max(1, min(num_frames_to_extract, int(total_video_frames)))
-        
-        indices = np.linspace(0, int(total_video_frames) - 1, num_to_sample, dtype=int, endpoint=True)
-        indices = sorted(list(set(indices))) # Ensure unique and sorted indices
-
+        actual_frames_extracted = 0
         for i in indices:
-            frame_data = reader.get_data(i)
-            frames.append(Image.fromarray(frame_data).convert("RGB"))
-        
+            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} 🚨.")
-        st.warning("Please install ffmpeg and ensure it's in your system's PATH (e.g., `sudo apt-get install ffmpeg` or `conda install 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"Could not extract frames: {e}")
+        st.error(f"Error extracting frames: {e}")
         st.error(traceback.format_exc())
         return []
     finally:
         if reader:
             reader.close()
-            
-    if not frames and num_frames_to_extract > 0:
-         st.warning("No frames were extracted. The video might be empty, corrupted, or in an unsupported format.")
-    return frames
-
-# --- Sound Prompt Generation ---
-def generate_sound_prompt(frames: list[Image.Image], moondream_pipe) -> str:
-    instruction = (
-        "Describe only the sounds implied by this image. Focus on: ambient noise, distinct sound events, "
-        "sound textures, actions producing sound, and the overall atmosphere. Be concise and evocative."
-    )
+
+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 = []
-    with st.spinner(f"Analyzing {len(frames)} frames with Moondream2..."):
+    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):
-            # st.progress((i + 1) / len(frames), text=f"Analyzing frame {i+1}/{len(frames)}") # More granular progress
             try:
-                # Moondream2 pipeline expects 'images' and 'prompt'
-                # generate_kwargs can control output length for the captioner
-                out = moondream_pipe(
-                    images=frame, 
-                    prompt=instruction, 
-                    generate_kwargs={"max_new_tokens": MAX_MOONDREAM_DESCRIPTION_TOKENS}
-                )
-                text = out[0].get('generated_text', '').strip()
-                if text:
-                    descriptions.append(text)
+                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 analyze frame {i+1} with Moondream2: {e}")
+                st.warning(f"Could not get description for a frame: {e}")
                 continue
     
     if not descriptions:
-        return "ambient background noise, general atmosphere" # Fallback if all frames fail
-
-    # Combine unique descriptions, trying to maintain some order
-    combined = "; ".join(list(dict.fromkeys(descriptions))) 
-    return combined
-
-# --- Audio Generation ---
-def generate_audio(prompt: str, duration_s: int, musicgen_processor, musicgen_model, guidance_scale: float, temperature: float):
+        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:
-        # Check prompt length against a practical limit for MusicGen (e.g., T5 encoder limit)
-        # MusicGen processor will truncate, but a warning is good.
-        # This tokenization is for estimation; actual truncation is by the processor.
-        prompt_tokens = musicgen_processor.tokenizer.tokenize(prompt)
-        if len(prompt_tokens) > MUSICGEN_PROMPT_MAX_TOKENS:
-            st.warning(
-                f"Generated sound prompt is very long ({len(prompt_tokens)} tokens) and will be truncated by MusicGen. "
-                f"This might affect audio quality. Consider using fewer analysis frames or if descriptions are too verbose."
-            )
-            # Truncate prompt manually if desired, or let processor handle it.
-            # For simplicity, we let the processor handle it.
-
-        inputs = musicgen_processor(
-            text=[prompt], 
-            return_tensors="pt", 
-            padding=True
-        )
-        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-        
-        # Critical for .half() models: ensure non-float tensors are not cast to half
-        inputs = {
-            k: (v.to(musicgen_model.dtype) if v.dtype.is_floating_point else v) 
-            for k, v in inputs.items()
-        }
-
-        # Determine max_new_tokens based on duration
-        # MusicGen's default is 50 tokens/second.
-        # Max sequence length for musicgen-small's decoder is typically 2048.
-        # We need to leave space for the prompt tokens.
-        # Max generated tokens (1500 for 30s) is a safe upper bound.
-        tokens_per_second = musicgen_model.config.audio_encoder.token_per_second
-        max_new_tokens = min(int(duration_s * tokens_per_second), 1500) 
+        inputs = audio_processor(text=[prompt], return_tensors="pt", padding=True).to(DEVICE)
         
-        with torch.inference_mode():
-            audio_tensor = musicgen_model.generate(
+        # 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_scale,
-                temperature=temperature,
-                pad_token_id=musicgen_model.config.eos_token_id 
+                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
             )
         
-        # Post-processing
-        arr = audio_tensor[0].cpu().float().numpy()
-        # Normalize audio
-        peak = np.max(np.abs(arr))
-        if peak == 0:  # Avoid division by zero for silent audio
-            arr = np.zeros_like(arr) # return silence
-        else:
-            arr = arr / peak * 0.9 # Normalize with some headroom
-        arr = np.clip(arr, -1.0, 1.0)
+        audio_array = audio_values[0, 0].cpu().numpy()
+        sampling_rate = audio_model.config.audio_encoder.sampling_rate
         
-        sampling_rate = musicgen_model.config.audio_encoder.sampling_rate
-        return arr, 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 during audio generation: {e}")
+        st.error(f"Error generating audio: {e}")
         st.error(traceback.format_exc())
         return None, None
-    finally:
-        clear_gpu_memory()
 
+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
 
-# --- Sync Audio/Video ---
-def sync_audio_video(video_path: str, audio_arr: np.ndarray, sampling_rate: int, mix_original_audio: bool) -> str | None:
-    tmp_wav_path = None
     output_video_path = None
+    temp_audio_path = None
     video_clip = None
-    audio_clip_generated = None
+    generated_audio_clip = None
     final_clip = None
 
     try:
-        # Write generated audio to a temporary WAV file
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp_wav_file:
-            sf.write(tmp_wav_file.name, audio_arr, sampling_rate)
-            tmp_wav_path = tmp_wav_file.name
-
+        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)
-        audio_clip_generated = mpy.AudioFileClip(tmp_wav_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
 
-        # Match audio duration to video duration (loop or trim)
-        if audio_clip_generated.duration < video_clip.duration:
-            num_loops = math.ceil(video_clip.duration / audio_clip_generated.duration)
-            audio_clip_generated = mpy.concatenate_audioclips([audio_clip_generated] * num_loops)
-        
-        audio_clip_generated = audio_clip_generated.subclip(0, video_clip.duration) # Trim to exact video duration
-
-        final_audio = audio_clip_generated
-        if mix_original_audio and video_clip.audio:
-            # Simple mix: lower volume of both and combine
-            # Adjust volumes as needed for better mixing
-            original_audio_processed = video_clip.audio.volumex(0.6) 
-            generated_audio_processed = audio_clip_generated.volumex(0.8) # Give generated slightly more presence
-            final_audio = mpy.CompositeAudioClip([original_audio_processed, generated_audio_processed])
-        
         final_clip = video_clip.set_audio(final_audio)
         
-        # Create a temporary file for the output video
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_out_vid_file:
-            output_video_path = tmp_out_vid_file.name
-        
+        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', 
-            threads=max(1, os.cpu_count() // 2), # Use half CPU cores, or at least 1
-            logger=None # Suppress verbose moviepy logs
+            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 syncing audio with video: {e}")
+        st.error(f"Error combining audio and video: {e}")
         st.error(traceback.format_exc())
         return None
     finally:
-        # Close moviepy clips to release file locks
+        # Close clips to release resources
         if video_clip: video_clip.close()
-        if audio_clip_generated: audio_clip_generated.close()
-        if final_clip: final_clip.close()
-        
-        if tmp_wav_path and os.path.exists(tmp_wav_path):
-            os.remove(tmp_wav_path)
-        # The output_video_path is returned and should be handled (e.g., deleted after download) by the caller
-        clear_gpu_memory()
-
-# --- Main UI ---
-st.title("🎬 AI Video Sound Designer")
-st.markdown(
-    "Upload an MP4 video, and this tool will analyze its visuals using **Moondream2** "
-    "to generate relevant sound prompts, then synthesize sound effects or music using **MusicGen**. "
-    "You can download the audio or the video with new synchronized sound."
-)
-st.markdown("---")
+        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 Settings
+# --- Sidebar for Settings ---
 with st.sidebar:
-    st.header("⚙️ Generation Settings")
-    
-    num_frames_to_analyze = st.slider(
-        "Frames to Analyze (for sound ideas)", 
-        min_value=1, max_value=10, value=DEFAULT_FRAMES_TO_ANALYZE,
-        help="Number of frames sampled from the video to generate sound descriptions. More frames can give diverse ideas but take longer to analyze."
-    )
+    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.")
     
-    audio_duration_s = st.slider(
-        "Target Audio Duration (seconds)", 
-        min_value=5, max_value=30, value=DEFAULT_AUDIO_DURATION_S,
-        help="Duration of the generated audio. If shorter than video, it will be looped."
-    )
-
     st.subheader("MusicGen Parameters")
-    guidance_scale = st.slider(
-        "Guidance Scale", 
-        min_value=1.0, max_value=10.0, value=DEFAULT_GUIDANCE, step=0.5,
-        help="Higher values make the audio follow the prompt more closely, but can reduce diversity. (MusicGen default: 3.0)"
-    )
-    temperature = st.slider(
-        "Temperature", 
-        min_value=0.1, max_value=2.0, value=DEFAULT_TEMPERATURE, step=0.1,
-        help="Controls randomness. Higher values mean more diversity/creativity, lower values more deterministic. (MusicGen default: 1.0)"
-    )
+    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.")
     
-    st.subheader("Video Output")
-    mix_with_original_audio = st.checkbox(
-        "Mix with original video audio", value=False,
-        help="If checked, the generated sound will be mixed with the video's existing audio. Otherwise, existing audio is replaced."
-    )
+    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
 
-    st.markdown("---")
-    st.caption(f"Using Moondream2: `{MODEL_MOONDREAM}`")
-    st.caption(f"Using MusicGen: `{MODEL_MUSICGEN}`")
-    st.caption(f"Processing on: `{str(DEVICE).upper()}`")
-    # if st.button("Clear Model Cache & Reload"):
-    #     st.cache_resource.clear()
-    #     st.experimental_rerun()
 
+# --- Main Application Logic ---
+uploaded_file = st.file_uploader("📤 Upload your MP4 video file (short clips recommended):", type=["mp4", "mov", "avi"])
 
-# File Uploader
-uploaded_video_file = st.file_uploader("📤 Upload MP4 Video", type=['mp4', 'mov', 'avi']) # Added more types
+# 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 'processed_video_path' not in st.session_state:
-    st.session_state.processed_video_path = None
-if 'generated_audio_path' not in st.session_state:
-    st.session_state.generated_audio_path = None
 
-
-if uploaded_video_file:
-    st.info(f"Uploaded: `{uploaded_video_file.name}` ({uploaded_video_file.size / (1024*1024):.2f} MB)")
+if uploaded_file is not None:
+    st.video(uploaded_file)
     
-    # Use a button to trigger processing to avoid re-processing on every widget change
+    # Use a button to trigger processing
     if st.button("✨ Generate Sound Design!", type="primary", use_container_width=True):
-        # Clear previous results from session state if any
-        if st.session_state.processed_video_path and os.path.exists(st.session_state.processed_video_path):
-            os.remove(st.session_state.processed_video_path)
-        st.session_state.processed_video_path = None
-        if st.session_state.generated_audio_path and os.path.exists(st.session_state.generated_audio_path):
-            os.remove(st.session_state.generated_audio_path)
-        st.session_state.generated_audio_path = None
-
-        # Path for the uploaded video, will be cleaned up in finally
-        temp_uploaded_video_path = None
-        
+        # --- 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:
-            # --- 1. Save uploaded video to a temporary file ---
-            with st.spinner("Preparing video..."):
-                with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_video_file.name)[1]) as tfile:
-                    tfile.write(uploaded_video_file.read())
-                    temp_uploaded_video_path = tfile.name
+            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
             
-            # --- 2. Extract Frames ---
-            with st.spinner(f"Extracting {num_frames_to_analyze} frames from video..."):
-                extracted_frames = extract_frames(temp_uploaded_video_path, num_frames_to_analyze)
-
-            if not extracted_frames:
-                st.error("Failed to extract frames. Cannot proceed.")
-                st.stop()
-
-            st.subheader("🖼️ Sampled Frames for Sound Analysis")
-            cols_to_show = min(len(extracted_frames), MAX_FRAMES_TO_SHOW_UI)
-            if cols_to_show > 0:
-                cols = st.columns(cols_to_show)
-                for i in range(cols_to_show):
-                    cols[i].image(extracted_frames[i], caption=f"Frame {i+1}", use_column_width=True)
-            elif extracted_frames:
-                 st.write(f"{len(extracted_frames)} frames extracted (not shown due to display limit).")
+            # === 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)
             
-            # --- 3. Load Moondream2 and Generate Sound Prompt ---
-            moondream_pipe = load_moondream_model()
-            if not moondream_pipe:
-                st.error("Moondream2 model could not be loaded. Cannot generate sound prompt.")
+            if not frames:
+                st.error("No frames extracted. Cannot proceed.")
                 st.stop()
             
-            sound_prompt = generate_sound_prompt(extracted_frames, moondream_pipe)
-            del moondream_pipe # Release Moondream model from memory
-            clear_gpu_memory()
-
-            if not sound_prompt:
-                st.error("Failed to generate a sound prompt from video frames.")
-                st.stop()
-            st.info(f"🧠 **Generated Sound Prompt:** {sound_prompt}")
-
-            # --- 4. Load MusicGen and Generate Audio ---
-            musicgen_processor, musicgen_model = load_musicgen_model_and_processor()
-            if not musicgen_processor or not musicgen_model:
-                st.error("MusicGen model/processor could not be loaded. Cannot generate audio.")
-                st.stop()
-
-            with st.spinner(f"Synthesizing {audio_duration_s}s audio with MusicGen... (This can take a few minutes)"):
-                generated_audio_arr, generated_sr = generate_audio(
-                    sound_prompt, audio_duration_s, musicgen_processor, musicgen_model, guidance_scale, temperature
-                )
-            
-            del musicgen_processor, musicgen_model # Release MusicGen model from memory
-            clear_gpu_memory()
-
-            if generated_audio_arr is None or generated_sr is None:
-                st.error("Audio generation failed.")
-                st.stop()
-
-            st.subheader("🔊 Generated Sound Effect")
-            st.audio(generated_audio_arr, sample_rate=generated_sr)
+            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)
             
-            # Save generated audio to a temporary file for download
-            with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp_audio_file:
-                sf.write(tmp_audio_file.name, generated_audio_arr, generated_sr)
-                st.session_state.generated_audio_path = tmp_audio_file.name # Store path in session_state
-            
-            with open(st.session_state.generated_audio_path, 'rb') as f_audio:
-                st.download_button(
-                    label="📥 Download Generated Audio Only (.wav)",
-                    data=f_audio,
-                    file_name=f"{os.path.splitext(uploaded_video_file.name)[0]}_sound_effect.wav",
-                    mime='audio/wav'
-                )
+            # === 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
             
-            # --- 5. Sync Audio with Video ---
-            with st.spinner("Synchronizing generated audio with video..."):
-                output_video_path = sync_audio_video(
-                    temp_uploaded_video_path, generated_audio_arr, generated_sr, mix_with_original_audio
-                )
-
-            if output_video_path and os.path.exists(output_video_path):
-                st.subheader("🎥 Video with New Sound Design")
-                st.video(output_video_path)
-                st.session_state.processed_video_path = output_video_path # Store path
-
-                with open(output_video_path, 'rb') as f_video:
+            # === 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(
-                        label="🎬 Download Video with New Sound (.mp4)",
-                        data=f_video,
-                        file_name=f"{os.path.splitext(uploaded_video_file.name)[0]}_sound_designed.mp4",
-                        mime='video/mp4'
+                        "📥 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("Failed to create video with new sound.")
+                st.error("Audio generation failed. Cannot proceed to video syncing.")
 
         except Exception as e:
-            st.error(f"An unexpected error occurred during processing: {e}")
+            st.error(f"An unexpected error occurred in the main processing pipeline: {e}")
             st.error(traceback.format_exc())
         finally:
-            # Clean up the initially uploaded temporary video file
-            if temp_uploaded_video_path and os.path.exists(temp_uploaded_video_path):
-                os.remove(temp_uploaded_video_path)
-            clear_gpu_memory()
-            # Note: Temp files for download (st.session_state.processed_video_path, st.session_state.generated_audio_path)
-            # are not deleted here. They would ideally be cleaned up when the session ends or a new file is processed.
-            # Streamlit doesn't have a direct session end hook for cleanup easily.
-            # For robust cleanup of download files, more complex session management or a background task would be needed.
-            # Current approach: they are overwritten or deleted on next "Generate" click.
-
-    # Display results from session state if they exist (e.g., after a settings change without re-upload)
-    # This part is tricky because changing settings should ideally re-trigger processing with new settings.
-    # The button "Generate Sound Design!" helps control this. If there are previous results and no button click,
-    # they might still be shown. This part is simplified for now.
-    elif st.session_state.generated_audio_path and os.path.exists(st.session_state.generated_audio_path):
-        st.subheader("🔊 Previously Generated Sound Effect")
-        st.audio(st.session_state.generated_audio_path)
-        with open(st.session_state.generated_audio_path, 'rb') as f_audio:
+            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(
-                label="📥 Download Generated Audio Only (.wav)",
-                data=f_audio,
-                file_name=f"{os.path.splitext(uploaded_video_file.name)[0]}_sound_effect.wav",
-                mime='audio/wav',
-                key="prev_audio_download" # Unique key
+                "📥 Download Previously Generated Audio (.wav)",
+                f,
+                file_name="generated_sound_previous.wav",
+                mime="audio/wav",
+                key="prev_audio_dl"
             )
-    if st.session_state.processed_video_path and os.path.exists(st.session_state.processed_video_path):
-        st.subheader("🎥 Previously Generated Video with New Sound Design")
-        st.video(st.session_state.processed_video_path)
-        with open(st.session_state.processed_video_path, 'rb') as f_video:
+    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(
-                label="🎬 Download Video with New Sound (.mp4)",
-                data=f_video,
-                file_name=f"{os.path.splitext(uploaded_video_file.name)[0]}_sound_designed.mp4",
-                mime='video/mp4',
-                key="prev_video_download" # Unique key
+                "🎬 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 file to begin!")
+    st.info("☝️ Upload a video to get started.")
 
 st.markdown("---")
-st.caption("Developed with ❤️ using Streamlit, Hugging Face Transformers, MoviePy, and ImageIO.")
+st.markdown("Made for Hugging Face Spaces. Model loading & generation can be slow on CPU.")