Create app.py

app.py

@@ -0,0 +1,202 @@
+# Import libraries
+import os
+import gradio as gr
+import torch
+import soundfile as sf
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+import logging
+from scipy.io.wavfile import write as write_wav
+from scipy import signal
+from moviepy.editor import VideoFileClip, AudioFileClip
+import requests
+from audiocraft.models import AudioGen, MusicGen  # Use audiocraft for AudioGen and MusicGen
+
+# Set up logging for better debug tracking
+logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger()
+
+# Download Places365 class labels
+try:
+    logging.info("Downloading Places365 class labels...")
+    url = "http://places2.csail.mit.edu/models_places365/categories_places365.txt"
+    response = requests.get(url)
+    with open("categories_places365.txt", "wb") as f:
+        f.write(response.content)
+    logging.info("Places365 class labels downloaded successfully.")
+except Exception as e:
+    logging.error(f"Error downloading Places365 class labels: {e}")
+    raise
+
+# Load Places365 model for scene detection (on CPU to save GPU memory)
+try:
+    logging.info("Loading Places365 model for scene detection...")
+    places365 = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
+    places365.eval()
+    places365.to("cpu")  # Move to CPU
+    logging.info("Places365 model loaded successfully.")
+except Exception as e:
+    logging.error(f"Error loading Places365 model: {e}")
+    raise
+
+# Load Places365 class labels
+with open("categories_places365.txt", "r") as f:
+    SCENE_CLASSES = [line.strip().split(" ")[0][3:] for line in f.readlines()]
+
+# Load AudioGen Medium and MusicGen Medium models
+try:
+    logging.info("Loading AudioGen Medium and MusicGen Medium models...")
+    audiogen_model = AudioGen.get_pretrained("facebook/audiogen-medium")
+    musicgen_model = MusicGen.get_pretrained("facebook/musicgen-medium")
+    logging.info("AudioGen Medium and MusicGen Medium models loaded successfully.")
+except Exception as e:
+    logging.error(f"Error loading AudioGen/MusicGen models: {e}")
+    raise
+
+# Function to classify a frame using Places365
+def classify_frame(frame):
+    try:
+        preprocess = transforms.Compose([
+            transforms.Resize(128),  # Smaller resolution
+            transforms.CenterCrop(128),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+        img = Image.fromarray(frame)
+        img = preprocess(img).unsqueeze(0)
+        with torch.no_grad():
+            output = places365(img.to("cpu"))  # Ensure inference on CPU
+        probabilities = F.softmax(output, dim=1)
+        _, predicted = torch.max(probabilities, 1)
+        predicted_index = predicted.item()
+
+        # Ensure the predicted index is within the range of SCENE_CLASSES
+        if predicted_index >= len(SCENE_CLASSES) or predicted_index < 0:
+            logging.warning(f"Predicted class index {predicted_index} is out of range. Defaulting to 'nature'.")
+            return "nature"  # Default scene type
+
+        scene_type = SCENE_CLASSES[predicted_index]
+        logging.info(f"Predicted scene: {scene_type}")
+        return scene_type
+    except Exception as e:
+        logging.error(f"Error classifying frame: {e}")
+        raise
+
+# Function to analyze video content and return the scene type using Places365
+def analyze_video(video_path):
+    try:
+        logging.info(f"Analyzing video: {video_path}")
+        clip = VideoFileClip(video_path)
+        frame = clip.get_frame(0)  # Get the first frame
+        frame = Image.fromarray(frame)  # Convert to PIL image
+        frame = np.array(frame.resize((128, 128)))  # Resize to reduce memory usage
+
+        # Classify the frame using Places365
+        scene_type = classify_frame(frame)
+        logging.info(f"Scene type detected: {scene_type}")
+        return scene_type
+    except Exception as e:
+        logging.error(f"Error analyzing video: {e}")
+        raise
+
+# Function to generate audio using AudioGen Medium
+def generate_audio_audiogen(scene, duration=10):
+    try:
+        logging.info(f"Generating audio for scene: {scene} using AudioGen Medium...")
+        audiogen_model.set_generation_params(duration=duration)
+        descriptions = [f"Ambient sounds of {scene}"]
+        wav = audiogen_model.generate(descriptions)  # Generate audio
+        audio_path = "generated_audio_audiogen.wav"
+        sf.write(audio_path, wav.squeeze().cpu().numpy(), 32000)  # Save as WAV file
+        logging.info(f"Audio generated and saved to: {audio_path}")
+        return audio_path
+    except Exception as e:
+        logging.error(f"Error generating audio with AudioGen Medium: {e}")
+        raise
+
+# Function to generate music using MusicGen Medium
+def generate_music_musicgen(scene, duration=10):
+    try:
+        logging.info(f"Generating music for scene: {scene} using MusicGen Medium...")
+        musicgen_model.set_generation_params(duration=duration)
+        descriptions = [f"Calm music for {scene}"]
+        wav = musicgen_model.generate(descriptions)  # Generate music
+        music_path = "generated_music_musicgen.wav"
+        sf.write(music_path, wav.squeeze().cpu().numpy(), 32000)  # Save as WAV file
+        logging.info(f"Music generated and saved to: {music_path}")
+        return music_path
+    except Exception as e:
+        logging.error(f"Error generating music with MusicGen Medium: {e}")
+        raise
+
+# Function to merge audio and video into a final video file using moviepy
+def merge_audio_video(video_path, audio_path, output_path="output.mp4"):
+    try:
+        logging.info("Merging audio and video using moviepy...")
+        video_clip = VideoFileClip(video_path)
+        audio_clip = AudioFileClip(audio_path)
+        final_clip = video_clip.set_audio(audio_clip)
+        final_clip.write_videofile(output_path, codec="libx264", audio_codec="aac")
+        logging.info(f"Final video saved to: {output_path}")
+        return output_path
+    except Exception as e:
+        logging.error(f"Error merging audio and video: {e}")
+        return None
+
+# Main processing function to handle video upload, scene analysis, and video output
+def process_video(video_path, progress=gr.Progress()):
+    try:
+        progress(0.1, desc="Starting video processing...")
+        logging.info("Starting video processing...")
+
+        # Analyze the video to determine the scene type
+        progress(0.3, desc="Analyzing video...")
+        scene_type = analyze_video(video_path)
+
+        # Generate audio using AudioGen Medium
+        progress(0.5, desc="Generating audio...")
+        audio_path = generate_audio_audiogen(scene_type, duration=10)
+
+        # Generate music using MusicGen Medium
+        progress(0.7, desc="Generating music...")
+        music_path = generate_music_musicgen(scene_type, duration=10)
+
+        # Merge the generated audio with the video and output the final video
+        progress(0.9, desc="Merging audio and video...")
+        output_path = merge_audio_video(video_path, music_path)
+        if not output_path:
+            return "Error: Failed to merge audio and video.", "Logs: Merge failed."
+
+        logging.info("Video processing completed successfully.")
+        return output_path, "Logs: Processing completed."
+    except Exception as e:
+        logging.error(f"Error in process_video: {e}")
+        return f"An error occurred during processing: {e}", f"Logs: {e}"
+
+# Gradio UI for video upload
+def gradio_interface(video_file, progress=gr.Progress()):
+    try:
+        progress(0.1, desc="Starting video processing...")
+        logging.info("Gradio interface triggered.")
+        output_video, logs = process_video(video_file, progress)
+        return output_video, logs
+    except Exception as e:
+        logging.error(f"Error in Gradio interface: {e}")
+        return f"An error occurred: {e}", f"Logs: {e}"
+
+# Launch Gradio app
+try:
+    logging.info("Launching Gradio app...")
+    interface = gr.Interface(
+        fn=gradio_interface,
+        inputs=[gr.Video(label="Upload Video")],
+        outputs=[gr.Video(label="Output Video with Generated Audio"), gr.Textbox(label="Logs", lines=10)],
+        title="Video to Video with Generated Audio and Music",
+        description="Upload a video, and this app will analyze it and generate matching audio and music using AudioGen Medium and MusicGen Medium."
+    )
+    interface.queue()  # Enable queue for long-running tasks
+    interface.launch(share=True)  # Launch the app
+except Exception as e:
+    logging.error(f"Error launching Gradio app: {e}")
+    raise