Update app.py

app.py

@@ -1,42 +1,51 @@
 import streamlit as st
 from PIL import Image
 import torch
-import skvideo.io
-from diffusers import I2VGenXLPipeline
-from diffusers.utils import export_to_video, load_image
 import numpy as np
-import imageio
 from moviepy.editor import ImageSequenceClip
 from transformers import MusicgenForConditionalGeneration, AutoProcessor
 from scipy.io import wavfile
 import ffmpeg
 
+# Function to generate video frames
 def generate_video(image, prompt, negative_prompt, video_length):
     generator = torch.manual_seed(8888)
     device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
     print(f"Using device: {device}")
     pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float32)
     pipeline.to(device)
-    frames = pipeline(
-        prompt=prompt,
-        image=image,
-        num_inference_steps=2,
-        negative_prompt=negative_prompt,
-        guidance_scale=9.0,
-        generator=generator,
-        num_frames=video_length*20
-    ).frames[0]
+
+    frames = []
+    total_frames = video_length * 20  # Assuming 20 frames per second
+
+    # Generate frames with progress tracking
+    for i in range(total_frames):
+        frame = pipeline(
+            prompt=prompt,
+            image=image,
+            num_inference_steps=2,
+            negative_prompt=negative_prompt,
+            guidance_scale=9.0,
+            generator=generator,
+            num_frames=1
+        ).frames[0]
+        frames.append(frame)
+        st.progress((i + 1) / total_frames)  # Update progress bar
+
     return frames
 
+# Function to export frames to video
 def export_frames_to_video(frames, output_file):
     frames_np = [np.array(frame) for frame in frames]
     clip = ImageSequenceClip(frames_np, fps=30)
     clip.write_videofile(output_file, codec='libx264', audio=False)
 
+# Function to generate music
 def generate_music(prompt, unconditional=False):
     model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
     model.to(device)
+
     if unconditional:
         unconditional_inputs = model.get_unconditional_inputs(num_samples=1)
         audio_values = model.generate(**unconditional_inputs, do_sample=True, max_new_tokens=256)
@@ -48,15 +57,18 @@ def generate_music(prompt, unconditional=False):
             return_tensors="pt",
         )
         audio_values = model.generate(**inputs.to(device), do_sample=True, guidance_scale=3, max_new_tokens=256)
+
     sampling_rate = model.config.audio_encoder.sampling_rate
     return audio_values[0].cpu().numpy(), sampling_rate
 
+# Function to combine audio and video
 def combine_audio_video(audio_file, video_file, output_file):
     audio = ffmpeg.input(audio_file)
     video = ffmpeg.input(video_file)
     output = ffmpeg.output(video, audio, output_file, vcodec='copy', acodec='aac')
     ffmpeg.run(output)
 
+# Streamlit UI
 st.title("AI-Powered Video and Music Generation")
 
 st.sidebar.title("Options")
@@ -74,13 +86,20 @@ unconditional = st.sidebar.checkbox("Generate unconditional music")
 if st.sidebar.button("Generate Video and Music"):
     if image is not None:
         image = Image.open(image)
+        
+        # Video generation with progress bar
+        st.write("Generating video...")
         frames = generate_video(image, prompt, negative_prompt, video_length)
         export_frames_to_video(frames, "output_video.mp4")
         st.video("output_video.mp4")
 
+    # Music generation with progress bar
+    st.write("Generating music...")
     audio_values, sampling_rate = generate_music(music_prompt, unconditional)
     wavfile.write("musicgen_out.wav", sampling_rate, audio_values)
     st.audio("musicgen_out.wav")
 
+    # Combine audio and video
+    st.write("Combining audio and video...")
     combine_audio_video("musicgen_out.wav", "output_video.mp4", "combined_output.mp4")
     st.video("combined_output.mp4")