v1

app.py

@@ -1,13 +1,13 @@
 import os
 import json
 import gradio as gr
-import tempfile
 import torch
 import spaces
 from pathlib import Path
-from transformers import AutoProcessor, AutoModelForVision2Seq
 import subprocess
 import logging
+import xml.etree.ElementTree as ET
+from xml.dom import minidom
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -16,16 +16,13 @@ def load_examples(json_path: str) -> dict:
     with open(json_path, 'r') as f:
         return json.load(f)
 
-def format_duration(seconds: int) -> str:
-    hours = seconds // 3600
-    minutes = (seconds % 3600) // 60
-    secs = seconds % 60
-    if hours > 0:
-        return f"{hours}:{minutes:02d}:{secs:02d}"
-    return f"{minutes}:{secs:02d}"
+def format_duration(seconds: float) -> str:
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    return f"{hours:02d}:{minutes:02d}:{secs:02d}"
 
 def get_video_duration_seconds(video_path: str) -> float:
-    """Use ffprobe to get video duration in seconds."""
     cmd = [
         "ffprobe",
         "-v", "quiet",
@@ -51,12 +48,10 @@ class VideoHighlightDetector:
         self.processor = AutoProcessor.from_pretrained(model_path)
         self.model = AutoModelForVision2Seq.from_pretrained(
             model_path,
-            torch_dtype=torch.bfloat16,
-            # _attn_implementation="flash_attention_2"
+            torch_dtype=torch.bfloat16
         ).to(device)
         
     def analyze_video_content(self, video_path: str) -> str:
-        """Analyze video content to determine its type and description."""
         system_message = "You are a helpful assistant that can understand videos. Describe what type of video this is and what's happening in it."
         messages = [
             {
@@ -83,24 +78,44 @@ class VideoHighlightDetector:
         outputs = self.model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
         return self.processor.decode(outputs[0], skip_special_tokens=True).lower().split("assistant: ")[1]
 
-    def determine_highlights(self, video_description: str) -> str:
-        """Determine what constitutes highlights based on video description."""
+    def analyze_segment(self, video_path: str) -> str:
+        """Analyze a specific video segment and provide a brief description."""
         messages = [
             {
                 "role": "system",
-                "content": [{"type": "text", "text": "You are a professional video editor specializing in creating viral highlight reels.  You understand that the most engaging highlights are brief and focus only on exceptional moments that are statistically rare or particularly dramatic. Moments that would make viewers say 'I can't believe that happened!"}]
+                "content": [{"type": "text", "text": "Describe what is happening in this specific video segment in a brief, concise way."}]
             },
             {
                 "role": "user",
-                "content": [{"type": "text", "text": f"""Here is a description of a video: 
-                
-                {video_description}
-                
-                Based on this description, list which rare segments should be included in a best of the best highlight."""}]
+                "content": [
+                    {"type": "video", "path": video_path},
+                    {"type": "text", "text": "What is happening in this segment? Provide a brief description."}
+                ]
             }
         ]
+        
+        inputs = self.processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt"
+        ).to(self.device)
+        
+        outputs = self.model.generate(**inputs, max_new_tokens=128, do_sample=True, temperature=0.7)
+        return self.processor.decode(outputs[0], skip_special_tokens=True).split("Assistant: ")[1]
 
-        print(messages)
+    def determine_highlights(self, video_description: str) -> str:
+        messages = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": "You are a professional video editor specializing in creating viral highlight reels."}]
+            },
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": f"Based on this description, list which segments should be included in highlights: {video_description}"}]
+            }
+        ]
         
         inputs = self.processor.apply_chat_template(
             messages,
@@ -114,22 +129,15 @@ class VideoHighlightDetector:
         return self.processor.decode(outputs[0], skip_special_tokens=True).split("Assistant: ")[1]
 
     def process_segment(self, video_path: str, highlight_types: str) -> bool:
-        """Process a video segment and determine if it contains highlights."""
         messages = [
             {
                 "role": "user",
                 "content": [
                     {"type": "video", "path": video_path},
-                    {"type": "text", "text": f"""{highlight_types}
-                    
-                     
-                    Do you see any of those elements in the video? answer yes if you do and answer no if you don't."""}
+                    {"type": "text", "text": f"Do you see any of these elements in the video: {highlight_types}? Answer yes or no."}
                 ]
             }
         ]
-
-        print(messages)
-
         
         inputs = self.processor.apply_chat_template(
             messages,
@@ -141,82 +149,53 @@ class VideoHighlightDetector:
         
         outputs = self.model.generate(**inputs, max_new_tokens=64, do_sample=False)
         response = self.processor.decode(outputs[0], skip_special_tokens=True).lower().split("assistant: ")[1]
-        print(f"Segment response {response}")
         return "yes" in response
 
-    def _concatenate_scenes(
-        self,
-        video_path: str,
-        scene_times: list,
-        output_path: str
-    ):
-        """Concatenate selected scenes into final video."""
-        if not scene_times:
-            logger.warning("No scenes to concatenate, skipping.")
-            return
-
-        filter_complex_parts = []
-        concat_inputs = []
-        for i, (start_sec, end_sec) in enumerate(scene_times):
-            filter_complex_parts.append(
-                f"[0:v]trim=start={start_sec}:end={end_sec},"
-                f"setpts=PTS-STARTPTS[v{i}];"
-            )
-            filter_complex_parts.append(
-                f"[0:a]atrim=start={start_sec}:end={end_sec},"
-                f"asetpts=PTS-STARTPTS[a{i}];"
-            )
-            concat_inputs.append(f"[v{i}][a{i}]")
-
-        concat_filter = f"{''.join(concat_inputs)}concat=n={len(scene_times)}:v=1:a=1[outv][outa]"
-        filter_complex = "".join(filter_complex_parts) + concat_filter
-
-        cmd = [
-            "ffmpeg",
-            "-y",
-            "-i", video_path,
-            "-filter_complex", filter_complex,
-            "-map", "[outv]",
-            "-map", "[outa]",
-            "-c:v", "libx264",
-            "-c:a", "aac",
-            output_path
-        ]
-
-        logger.info(f"Running ffmpeg command: {' '.join(cmd)}")
-        subprocess.run(cmd, check=True)
+def create_xspf_playlist(video_path: str, segments: list, descriptions: list) -> str:
+    """Create XSPF playlist from segments with descriptions."""
+    root = ET.Element("playlist", version="1", xmlns="http://xspf.org/ns/0/")
+    
+    # Get video filename for the title
+    video_filename = os.path.basename(video_path)
+    title = ET.SubElement(root, "title")
+    title.text = f"{video_filename} - Highlights"
+    
+    tracklist = ET.SubElement(root, "trackList")
+    
+    for idx, ((start_time, end_time), description) in enumerate(zip(segments, descriptions)):
+        track = ET.SubElement(tracklist, "track")
+        
+        location = ET.SubElement(track, "location")
+        location.text = f"file:///{video_filename}"
+        
+        title = ET.SubElement(track, "title")
+        title.text = f"Highlight {idx + 1}"
+        
+        annotation = ET.SubElement(track, "annotation")
+        annotation.text = description
+        
+        start_meta = ET.SubElement(track, "meta", rel="start")
+        start_meta.text = format_duration(start_time)
+        
+        end_meta = ET.SubElement(track, "meta", rel="end")
+        end_meta.text = format_duration(end_time)
+    
+    # Add VLC extension
+    extension = ET.SubElement(root, "extension", application="http://www.videolan.org/vlc/playlist/0")
+    for i in range(len(segments)):
+        item = ET.SubElement(extension, "vlc:item", tid=str(i))
+    
+    # Convert to string with pretty printing
+    xml_str = minidom.parseString(ET.tostring(root)).toprettyxml(indent="    ")
+    return xml_str
 
 def create_ui(examples_path: str, model_path: str):
     examples_data = load_examples(examples_path)
 
     with gr.Blocks() as app:
-        gr.Markdown("# Video Highlight Generator")
-        gr.Markdown("Upload a video and get an automated highlight reel!")
+        gr.Markdown("# Video Highlight Playlist Generator")
+        gr.Markdown("Upload a video and get an XSPF playlist of highlights!")
         
-        with gr.Row():
-            gr.Markdown("## Example Results")
-
-        with gr.Row():
-            for example in examples_data["examples"]:
-                with gr.Column():
-                    gr.Video(
-                        value=example["original"]["url"],
-                        label=f"Original ({format_duration(example['original']['duration_seconds'])})",
-                        interactive=False
-                    )
-                    gr.Markdown(f"### {example['title']}")
-                
-                with gr.Column():
-                    gr.Video(
-                        value=example["highlights"]["url"],
-                        label=f"Highlights ({format_duration(example['highlights']['duration_seconds'])})",
-                        interactive=False
-                    )
-                    with gr.Accordion("Chain of thought details", open=False):
-                        gr.Markdown(f"### Summary:\n{example['analysis']['video_description']}")
-                        gr.Markdown(f"### Highlights to search for:\n{example['analysis']['highlight_types']}")
-
-        gr.Markdown("## Try It Yourself!")
         with gr.Row():
             with gr.Column(scale=1):
                 input_video = gr.Video(
@@ -226,185 +205,128 @@ def create_ui(examples_path: str, model_path: str):
                 process_btn = gr.Button("Process Video", variant="primary")
             
             with gr.Column(scale=1):
-                output_video = gr.Video(
-                    label="Highlight Video",
+                output_playlist = gr.File(
+                    label="Highlight Playlist (XSPF)",
                     visible=False,
                     interactive=False,
                 )
-
                 status = gr.Markdown()
-        
+                
                 analysis_accordion = gr.Accordion(
-                    "Chain of thought details", 
+                    "Analysis Details", 
                     open=True, 
                     visible=False
                 )
                 
                 with analysis_accordion:
-                    video_description = gr.Markdown("", elem_id="video_desc")
-                    highlight_types = gr.Markdown("", elem_id="highlight_types")
+                    video_description = gr.Markdown("")
+                    highlight_types = gr.Markdown("")
 
         @spaces.GPU
         def on_process(video):
-            # Clear all components when starting new processing
-            yield [
-                "",  # Clear status
-                "",  # Clear video description
-                "",  # Clear highlight types
-                gr.update(value=None, visible=False),  # Clear video
-                gr.update(visible=False)  # Hide accordion
-            ]
-            
             if not video:
-                yield [
+                return [
+                    None,
                     "Please upload a video",
                     "",
                     "",
-                    gr.update(visible=False),
                     gr.update(visible=False)
                 ]
-                return
             
             try:
                 duration = get_video_duration_seconds(video)
-                if duration > 1800:  # 30 minutes
-                    yield [
+                if duration > 18000:  # 300 minutes
+                    return [
+                        None,
                         "Video must be shorter than 30 minutes",
                         "",
                         "",
-                        gr.update(visible=False),
                         gr.update(visible=False)
                     ]
-                    return
 
-                yield [
-                    "Initializing video highlight detector...",
-                    "",
-                    "",
-                    gr.update(visible=False),
-                    gr.update(visible=False)
-                ]
-
-                detector = VideoHighlightDetector(
-                    model_path=model_path,
-                    batch_size=8
-                )
-
-                yield [
-                    "Analyzing video content...",
-                    "",
-                    "",
-                    gr.update(visible=False),
-                    gr.update(visible=True)
-                ]
+                detector = VideoHighlightDetector(model_path=model_path)
                 
+                # Analyze video content
                 video_desc = detector.analyze_video_content(video)
-                formatted_desc = f"### Summary:\n {video_desc[:500] + '...' if len(video_desc) > 500 else video_desc}"
-                
-                yield [
-                    "Determining highlight types...",
-                    formatted_desc,
-                    "",
-                    gr.update(visible=False),
-                    gr.update(visible=True)
-                ]
+                formatted_desc = f"### Video Summary:\n{video_desc}"
                 
+                # Determine highlight types
                 highlights = detector.determine_highlights(video_desc)
-                formatted_highlights = f"### Highlights to search for:\n {highlights[:500] + '...' if len(highlights) > 500 else highlights}"
-
-                # Split video into segments
-                temp_dir = "temp_segments"
-                os.makedirs(temp_dir, exist_ok=True)
+                formatted_highlights = f"### Highlight Criteria:\n{highlights}"
                 
+                # Process video in segments
                 segment_length = 10.0
-                duration = get_video_duration_seconds(video)
                 kept_segments = []
-                segments_processed = 0
-                total_segments = int(duration / segment_length)
-
+                segment_descriptions = []
+                
                 for start_time in range(0, int(duration), int(segment_length)):
-                    segments_processed += 1
-                    progress = int((segments_processed / total_segments) * 100)
-                    
-                    yield [
-                        f"Processing segments... {progress}% complete",
-                        formatted_desc,
-                        formatted_highlights,
-                        gr.update(visible=False),
-                        gr.update(visible=True)
-                    ]
-
-                    # Create segment
-                    segment_path = f"{temp_dir}/segment_{start_time}.mp4"
                     end_time = min(start_time + segment_length, duration)
                     
-                    cmd = [
-                        "ffmpeg",
-                        "-y",
-                        "-i", video,
-                        "-ss", str(start_time),
-                        "-t", str(segment_length),
-                        "-c:v", "libx264",
-                        "-preset", "ultrafast",  # Use ultrafast preset for speed
-                        "-pix_fmt", "yuv420p",   # Ensure compatible pixel format
-                        segment_path
-                    ]
-                    subprocess.run(cmd, check=True)
-
-                    # Process segment
-                    if detector.process_segment(segment_path, highlights):
-                        print("KEEPING SEGMENT")
-                        kept_segments.append((start_time, end_time))
-                    
-                    # Clean up segment file
-                    os.remove(segment_path)
-
-                # Remove temp directory
-                os.rmdir(temp_dir)
-
-                # Create final video
+                    # Create temporary segment
+                    with tempfile.NamedTemporaryFile(suffix='.mp4') as temp_segment:
+                        cmd = [
+                            "ffmpeg",
+                            "-y",
+                            "-i", video,
+                            "-ss", str(start_time),
+                            "-t", str(segment_length),
+                            "-c:v", "libx264",
+                            "-preset", "ultrafast",
+                            temp_segment.name
+                        ]
+                        subprocess.run(cmd, check=True)
+                        
+                        if detector.process_segment(temp_segment.name, highlights):
+                            # Get segment description
+                            description = detector.analyze_segment(temp_segment.name)
+                            kept_segments.append((start_time, end_time))
+                            segment_descriptions.append(description)
+                
                 if kept_segments:
-                    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
-                        temp_output = tmp_file.name
-                        detector._concatenate_scenes(video, kept_segments, temp_output)
+                    # Create XSPF playlist
+                    playlist_content = create_xspf_playlist(video, kept_segments, segment_descriptions)
+                    
+                    # Save playlist to temporary file
+                    with tempfile.NamedTemporaryFile(mode='w', suffix='.xspf', delete=False) as f:
+                        f.write(playlist_content)
+                        playlist_path = f.name
                     
-                    yield [
-                        "Processing complete!",
+                    return [
+                        gr.update(value=playlist_path, visible=True),
+                        "Processing complete! Download the XSPF playlist.",
                         formatted_desc,
                         formatted_highlights,
-                        gr.update(value=temp_output, visible=True),
                         gr.update(visible=True)
                     ]
                 else:
-                    yield [
+                    return [
+                        None,
                         "No highlights detected in the video.",
                         formatted_desc,
                         formatted_highlights,
-                        gr.update(visible=False),
                         gr.update(visible=True)
                     ]
 
             except Exception as e:
                 logger.exception("Error processing video")
-                yield [
+                return [
+                    None,
                     f"Error processing video: {str(e)}",
                     "",
                     "",
-                    gr.update(visible=False),
                     gr.update(visible=False)
                 ]
             finally:
-                # Clean up
                 torch.cuda.empty_cache()
 
         process_btn.click(
             on_process,
             inputs=[input_video],
             outputs=[
+                output_playlist,
                 status,
                 video_description,
                 highlight_types,
-                output_video,
                 analysis_accordion
             ],
             queue=True,
@@ -413,10 +335,5 @@ def create_ui(examples_path: str, model_path: str):
     return app
 
 if __name__ == "__main__":
-    # subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
-    # Initialize CUDA
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    
     app = create_ui("video_spec.json", "HuggingFaceTB/SmolVLM2-2.2B-Instruct")
     app.launch()