Update app.py

app.py

@@ -1,26 +1,20 @@
 import gradio as gr
 import torch
-import numpy as np
 from diffusers import I2VGenXLPipeline
 from transformers import MusicgenForConditionalGeneration, AutoProcessor
 from PIL import Image
 from moviepy.editor import ImageSequenceClip
+import numpy as np
 import io
-import ffmpeg
 import scipy.io.wavfile
+import ffmpeg
 
 def generate_video(image, prompt, negative_prompt, video_length):
     generator = torch.manual_seed(8888)
-
-    # Set the device to CPU or a non-NVIDIA GPU
     device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
-    print(f"Using device: {device}")
-
-    # Load the pipeline
     pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float32)
-    pipeline.to(device)  # Move the model to the selected device
+    pipeline.to(device)
 
-    # Generate frames with progress tracking
     frames = []
     total_frames = video_length * 30  # Assuming 30 frames per second
 
@@ -35,15 +29,11 @@ def generate_video(image, prompt, negative_prompt, video_length):
             num_frames=1
         ).frames[0]
         frames.append(np.array(frame))
+        yield (i + 1) / total_frames  # Update progress
 
-        # Update progress
-        yield (i + 1) / total_frames  # Yield progress
-
-    # Create a video clip from the frames
     output_file = "output_video.mp4"
-    clip = ImageSequenceClip(frames, fps=30)  # Set the frames per second
+    clip = ImageSequenceClip(frames, fps=30)
     clip.write_videofile(output_file, codec='libx264', audio=False)
-
     return output_file
 
 def generate_music(prompt, unconditional=False):
@@ -51,7 +41,6 @@ def generate_music(prompt, unconditional=False):
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
     model.to(device)
 
-    # Generate music
     if unconditional:
         unconditional_inputs = model.get_unconditional_inputs(num_samples=1)
         audio_values = model.generate(**unconditional_inputs, do_sample=True, max_new_tokens=256)
@@ -66,9 +55,10 @@ def generate_music(prompt, unconditional=False):
 
     sampling_rate = model.config.audio_encoder.sampling_rate
     audio_file = "musicgen_out.wav"
-    # Save the generated audio
-    scipy.io.wavfile.write(audio_file, sampling_rate, audio_values[0].cpu().numpy())
-    
+    audio_data = audio_values[0].cpu().numpy()
+    audio_data = np.clip(audio_data, -1.0, 1.0)
+    audio_data = (audio_data * 32767).astype(np.int16)
+    scipy.io.wavfile.write(audio_file, sampling_rate, audio_data)
     return audio_file
 
 def combine_audio_video(audio_file, video_file):
@@ -79,26 +69,16 @@ def combine_audio_video(audio_file, video_file):
     ffmpeg.run(output)
     return output_file
 
-# Gradio interface
 def interface(image_path, prompt, negative_prompt, video_length, music_prompt, unconditional):
-    # Convert the uploaded image path to a PIL Image
     image = Image.open(image_path)
-    
-    # Generate video and track progress
     video_file = generate_video(image, prompt, negative_prompt, video_length)
-    
-    # Generate music
     audio_file = generate_music(music_prompt, unconditional)
-    
-    # Combine audio and video
     combined_file = combine_audio_video(audio_file, video_file)
-    
     return combined_file
 
-# Create Gradio Blocks
 with gr.Blocks() as demo:
     gr.Markdown("# AI-Powered Video and Music Generation")
-    
+
     with gr.Row():
         image_input = gr.Image(type="filepath", label="Upload Image")
         prompt_input = gr.Textbox(label="Enter the Video Prompt")
@@ -110,13 +90,11 @@ with gr.Blocks() as demo:
     generate_button = gr.Button("Generate Video and Music")
     output_video = gr.Video(label="Output Video with Sound")
 
-    # Define the button action
     generate_button.click(
         interface,
         inputs=[image_input, prompt_input, negative_prompt_input, video_length_input, music_prompt_input, unconditional_checkbox],
         outputs=output_video,
-        show_progress=True  # Show progress bar
+        show_progress=True
     )
 
-# Launch the Gradio app
 demo.launch()