Create app.py

app.py

@@ -0,0 +1,126 @@
+import gradio as gr
+import tempfile
+import os
+import numpy as np
+from PIL import Image
+from moviepy.editor import (
+    VideoFileClip, concatenate_videoclips, CompositeVideoClip,
+    TextClip, AudioFileClip, afx
+)
+from transformers import AutoProcessor, BlipForConditionalGeneration, MusicgenForConditionalGeneration
+import torch
+import soundfile as sf
+import imageio
+
+def enhance_prompt(base_description):
+    base = base_description.lower().strip()
+    actions = {
+        "walk": "crisp footsteps on a wooden floor",
+        "run": "rapid footsteps and heavy breathing",
+        "drive": "engine roar and tires screeching",
+        "talk": "soft voices and background murmur",
+        "crash": "loud crash and debris scattering",
+        "fall": "thud of impact and rustling debris"
+    }
+    objects = {
+        "person": "human activity with subtle breathing",
+        "dog": "playful barks and pawsteps",
+        "car": "mechanical hum and tire friction",
+        "tree": "rustling leaves in a breeze",
+        "forest": "gentle wind and distant bird calls"
+    }
+    environments = {
+        "room": "echoing footsteps and muffled sounds",
+        "street": "distant traffic and urban hum",
+        "forest": "wind through trees and twigs snapping",
+        "outside": "open air with faint wind"
+    }
+
+    sound_description = next((sound for k, sound in actions.items() if k in base), "subtle ambient hum")
+    sound_description += next((f" and {sound}" for k, sound in objects.items() if k in base), "")
+    sound_description += next((f" in a {env} with {sound}" for env, sound in environments.items() if env in base), "")
+
+    return f"{base} with {sound_description}"
+
+@torch.inference_mode()
+def generate_caption(frame, processor, model):
+    inputs = processor(images=frame, return_tensors="pt")
+    if torch.cuda.is_available():
+        inputs = {k: v.to("cuda") for k, v in inputs.items()}
+        model = model.to("cuda").half()
+    out = model.generate(**inputs)
+    return processor.decode(out[0], skip_special_tokens=True)
+
+def generate_audio(prompt, musicgen_processor, musicgen_model, duration):
+    inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt")
+    if torch.cuda.is_available():
+        inputs = {k: v.to("cuda") for k, v in inputs.items()}
+        musicgen_model = musicgen_model.to("cuda").half()
+    audio_values = musicgen_model.generate(
+        **inputs, max_new_tokens=256, do_sample=True,
+        guidance_scale=3.0, top_k=50, top_p=0.95
+    )
+    audio_array = audio_values[0].cpu().numpy().flatten()
+    audio_array = audio_array / np.max(np.abs(audio_array)) * 0.9
+    audio_array = np.clip(audio_array, -1.0, 1.0)
+    temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+    sf.write(temp_audio.name, audio_array, 32000)
+    return temp_audio.name
+
+def process_clips(video_files, top_texts, bottom_texts):
+    processor_blip = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+    model_blip = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+    processor_music = AutoProcessor.from_pretrained("facebook/musicgen-small")
+    model_music = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+
+    clips = []
+    for i, file in enumerate(video_files):
+        clip = VideoFileClip(file.name)
+        # Extract frame and caption
+        frame = Image.fromarray(imageio.v2.imread(clip.filename))
+        caption = generate_caption(frame, processor_blip, model_blip)
+        prompt = enhance_prompt(caption)
+
+        # Generate audio
+        audio_path = generate_audio(prompt, processor_music, model_music, clip.duration)
+        audio_clip = AudioFileClip(audio_path).subclip(0, clip.duration)
+        clip = clip.set_audio(audio_clip)
+
+        # Add text overlays
+        overlays = [clip]
+        if top_texts[i]:
+            top_txt = TextClip(top_texts[i], fontsize=30, color='white').set_position("top").set_duration(clip.duration)
+            overlays.append(top_txt)
+        if bottom_texts[i]:
+            btm_txt = TextClip(bottom_texts[i], fontsize=30, color='white').set_position("bottom").set_duration(clip.duration)
+            overlays.append(btm_txt)
+        clip = CompositeVideoClip(overlays).fadein(0.5).fadeout(0.5)
+        clips.append(clip)
+
+    final_video = concatenate_videoclips(clips, method="compose")
+    output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
+    final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")
+    return output_path
+
+def launch_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("# POV Video Generator with AI Sound Effects")
+        video_files = gr.File(file_types=[".mp4"], file_count="multiple", label="Upload your video clips")
+        top_texts = gr.Textbox(label="Top Texts (comma-separated for each clip)")
+        bottom_texts = gr.Textbox(label="Bottom Texts (comma-separated for each clip)")
+        generate_btn = gr.Button("Generate Final POV Video")
+        output_video = gr.Video()
+
+        def run(files, tops, bottoms):
+            top_list = [t.strip() for t in tops.split(",")] if tops else ["" for _ in files]
+            bottom_list = [b.strip() for b in bottoms.split(",")] if bottoms else ["" for _ in files]
+            while len(top_list) < len(files): top_list.append("")
+            while len(bottom_list) < len(files): bottom_list.append("")
+            return process_clips(files, top_list, bottom_list)
+
+        generate_btn.click(fn=run, inputs=[video_files, top_texts, bottom_texts], outputs=output_video)
+
+    demo.launch()
+
+if __name__ == "__main__":
+    launch_interface()