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, AutoModelForCausalLM
 import requests
+from audiocraft.models import AudioGen, MusicGen  # Add this line
 
 # Set up logging for better debug tracking
 logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -47,15 +47,8 @@ with open("categories_places365.txt", "r") as f:
 # Load AudioGen Medium and MusicGen Medium models
 try:
     logging.info("Loading AudioGen Medium and MusicGen Medium models...")
-    audiogen_processor = AutoProcessor.from_pretrained("facebook/audiogen-medium")
-    audiogen_model = AutoModelForCausalLM.from_pretrained("facebook/audiogen-medium")
-    musicgen_processor = AutoProcessor.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"
-    audiogen_model.to(device)
-    musicgen_model.to(device)
+    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}")
@@ -111,16 +104,11 @@ def analyze_video(video_path):
 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()
+        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"
-        write_wav(audio_path, 16000, audio)  # Save as WAV file
+        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:
@@ -131,16 +119,11 @@ def generate_audio_audiogen(scene, duration=10):
 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()
+        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"
-        write_wav(music_path, 16000, music)  # Save as WAV file
+        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: