Update app.py

app.py

@@ -10,8 +10,8 @@ import logging
 from scipy.io.wavfile import write as write_wav
 from scipy import signal
 from moviepy.editor import VideoFileClip, AudioFileClip
-from transformers import AutoProcessor, AutoModelForAudioGeneration
-import requests  # Add this line
+from transformers import AutoProcessor, AutoModelForCausalLM
+import requests
 
 # Set up logging for better debug tracking
 logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -48,9 +48,9 @@ with open("categories_places365.txt", "r") as f:
 try:
     logging.info("Loading AudioGen Medium and MusicGen Medium models...")
     audiogen_processor = AutoProcessor.from_pretrained("facebook/audiogen-medium")
-    audiogen_model = AutoModelForAudioGeneration.from_pretrained("facebook/audiogen-medium")
+    audiogen_model = AutoModelForCausalLM.from_pretrained("facebook/audiogen-medium")
     musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-medium")
-    musicgen_model = AutoModelForAudioGeneration.from_pretrained("facebook/musicgen-medium")
+    musicgen_model = AutoModelForCausalLM.from_pretrained("facebook/musicgen-medium")
     
     # Move models to GPU if available
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -61,4 +61,159 @@ except Exception as e:
     logging.error(f"Error loading AudioGen/MusicGen models: {e}")
     raise
 
-# Rest of the code remains the same...
+# 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...")
+        inputs = audiogen_processor(
+            text=[f"Ambient sounds of {scene}"],
+            padding=True,
+            return_tensors="pt",
+        ).to(audiogen_model.device)  # Move inputs to the same device as the model
+        with torch.no_grad():
+            audio = audiogen_model.generate(**inputs, max_new_tokens=duration * 50)  # Adjust tokens for duration
+        audio = audio.cpu().numpy().squeeze()
+        audio_path = "generated_audio_audiogen.wav"
+        write_wav(audio_path, 16000, audio)  # 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...")
+        inputs = musicgen_processor(
+            text=[f"Calm music for {scene}"],
+            padding=True,
+            return_tensors="pt",
+        ).to(musicgen_model.device)  # Move inputs to the same device as the model
+        with torch.no_grad():
+            music = musicgen_model.generate(**inputs, max_new_tokens=duration * 50)  # Adjust tokens for duration
+        music = music.cpu().numpy().squeeze()
+        music_path = "generated_music_musicgen.wav"
+        write_wav(music_path, 16000, music)  # 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