Spaces:
Runtime error
Runtime error
| # 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 |