test

example.py

@@ -0,0 +1,458 @@
+import os
+import json
+import gradio as gr
+import tempfile
+import torch
+import spaces
+from pathlib import Path
+from transformers import AutoProcessor, AutoModelForImageTextToText
+import subprocess
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+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 get_video_duration_seconds(video_path: str) -> float:
+    """Use ffprobe to get video duration in seconds."""
+    cmd = [
+        "ffprobe",
+        "-v", "quiet",
+        "-print_format", "json",
+        "-show_format",
+        video_path
+    ]
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    info = json.loads(result.stdout)
+    return float(info["format"]["duration"])
+
+class VideoHighlightDetector:
+    def __init__(
+        self,
+        model_path: str,
+        device: str = "cuda",
+        batch_size: int = 8
+    ):
+        self.device = device
+        self.batch_size = batch_size
+        
+        # Initialize model and processor
+        self.processor = AutoProcessor.from_pretrained(model_path)
+        self.model = AutoModelForImageTextToText.from_pretrained(
+            model_path,
+            torch_dtype=torch.bfloat16,
+            # _attn_implementation="flash_attention_2"
+        ).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 = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": system_message}]
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "video", "path": video_path},
+                    {"type": "text", "text": "What type of video is this and what's happening in it? Be specific about the content type and general activities you observe."}
+                ]
+            }
+        ]
+        
+        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=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, prompt_num: int = 1) -> str:
+        """Determine what constitutes highlights based on video description with different prompts."""
+        system_prompts = {
+            1: "You are a highlight editor. List archetypal dramatic moments that would make compelling highlights if they appear in the video. Each moment should be specific enough to be recognizable but generic enough to potentially exist in other videos of this type.",
+            2: "You are a helpful visual-language assistant that can understand videos and edit. You are tasked helping the user to create highlight reels for videos. Highlights should be rare and important events in the video in question."
+        }
+        user_prompts = {
+            1: "List potential highlight moments to look for in this video:",
+            2: "List dramatic moments that would make compelling highlights if they appear in the video. Each moment should be specific enough to be recognizable but generic enough to potentially exist in any video of this type:"
+        }
+
+        
+        messages = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": system_prompts[prompt_num]}]
+            },
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": f"""Here is a description of a video:\n\n{video_description}\n\n{user_prompts[prompt_num]}"""}]
+            }
+        ]
+
+        print(f"Using prompt {prompt_num} for highlight detection")
+        print(messages)
+        
+        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=256, do_sample=True, temperature=0.7)
+        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": "system",
+                "content": [{"type": "text", "text": "You are a video highlight analyzer. Your role is to identify moments that have high dramatic value, focusing on displays of skill, emotion, personality, or tension. Compare video segments against provided example highlights to find moments with similar emotional impact and visual interest, even if the specific actions differ."}]
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "video", "path": video_path},
+                    {"type": "text", "text": f"""Given these highlight examples:\n{highlight_types}\n\nDoes this video contain a moment that matches the core action of one of the highlights? Answer with:\n'yes' or 'no'\nIf yes, justify it"""}]
+            }
+        ]
+
+
+        print(messages)
+
+        
+        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=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_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!")
+        
+        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(
+                    label="Upload your video (max 30 minutes)",
+                    interactive=True
+                )
+                process_btn = gr.Button("Process Video", variant="primary")
+            
+            with gr.Column(scale=1):
+                output_video = gr.Video(
+                    label="Highlight Video",
+                    visible=False,
+                    interactive=False,
+                )
+
+                status = gr.Markdown()
+        
+                analysis_accordion = gr.Accordion(
+                    "Chain of thought details", 
+                    open=True, 
+                    visible=False
+                )
+                
+                with analysis_accordion:
+                    video_description = gr.Markdown("", elem_id="video_desc")
+                    highlight_types = gr.Markdown("", elem_id="highlight_types")
+
+        @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 [
+                    "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 [
+                        "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=16
+                )
+
+                yield [
+                    "Analyzing video content...",
+                    "",
+                    "",
+                    gr.update(visible=False),
+                    gr.update(visible=True)
+                ]
+                
+                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 (2 variations)...",
+                    formatted_desc,
+                    "",
+                    gr.update(visible=False),
+                    gr.update(visible=True)
+                ]
+                
+                # Get two different sets of highlights
+                highlights1 = detector.determine_highlights(video_desc, prompt_num=1)
+                highlights2 = detector.determine_highlights(video_desc, prompt_num=2)
+                formatted_highlights = f"### Highlights to search for:\nSet 1:\n{highlights1[:500] + '...' if len(highlights1) > 500 else highlights1}\n\nSet 2:\n{highlights2[:500] + '...' if len(highlights2) > 500 else highlights2}"
+
+                # Split video into segments
+                temp_dir = "temp_segments"
+                os.makedirs(temp_dir, exist_ok=True)
+                
+                segment_length = 10.0
+                duration = get_video_duration_seconds(video)
+                kept_segments1 = []
+                kept_segments2 = []
+                segments_processed = 0
+                total_segments = int(duration / segment_length)
+
+                for start_time in range(0, int(duration), int(segment_length)):
+                    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 with both highlight sets
+                    if detector.process_segment(segment_path, highlights1):
+                        print("KEEPING SEGMENT FOR SET 1")
+                        kept_segments1.append((start_time, end_time))
+                    
+                    if detector.process_segment(segment_path, highlights2):
+                        print("KEEPING SEGMENT FOR SET 2")
+                        kept_segments2.append((start_time, end_time))
+                    
+                    # Clean up segment file
+                    os.remove(segment_path)
+                    segments_processed += 1
+
+
+                # Remove temp directory
+                os.rmdir(temp_dir)
+
+                # Calculate percentages of video kept for each highlight set
+                total_duration = duration
+                duration1 = sum(end - start for start, end in kept_segments1)
+                duration2 = sum(end - start for start, end in kept_segments2)
+                
+                percent1 = (duration1 / total_duration) * 100
+                percent2 = (duration2 / total_duration) * 100
+                
+                print(f"Highlight set 1: {percent1:.1f}% of video")
+                print(f"Highlight set 2: {percent2:.1f}% of video")
+
+                # Choose the set with lower percentage unless it's zero
+                final_segments = kept_segments2 if (0 < percent2 <= percent1 or percent1 == 0) else kept_segments1
+
+                # Create final video
+                if final_segments:
+                    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
+                        temp_output = tmp_file.name
+                        detector._concatenate_scenes(video, final_segments, temp_output)
+                    
+                    selected_set = "2" if final_segments == kept_segments2 else "1"
+                    percent_used = percent2 if final_segments == kept_segments2 else percent1
+                    
+                    completion_message = f"Processing complete! Used highlight set {selected_set} ({percent_used:.1f}% of video)"
+                    
+                    yield [
+                        completion_message,
+                        formatted_desc,
+                        formatted_highlights,
+                        gr.update(value=temp_output, visible=True),
+                        gr.update(visible=True)
+                    ]
+                else:
+                    yield [
+                        "No highlights detected in the video with either set of criteria.",
+                        formatted_desc,
+                        formatted_highlights,
+                        gr.update(visible=False),
+                        gr.update(visible=True)
+                    ]
+
+            except Exception as e:
+                logger.exception("Error processing video")
+                yield [
+                    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=[
+                status,
+                video_description,
+                highlight_types,
+                output_video,
+                analysis_accordion
+            ],
+            queue=True,
+        )
+
+    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()

requirements.txt

@@ -1,8 +1,9 @@
-torch==2.1.2
-torchvision==0.16.2
+torch
+torchvision
 git+https://github.com/huggingface/transformers.git
-num2words==0.5.13
-gradio==4.19.2
-av==10.0.0
-numpy==1.24.3
-Pillow==10.0.0
+gradio
+huggingface_hub
+spaces
+av
+numpy
+Pillow

src/app.py

@@ -1,31 +1,108 @@
 import gradio as gr
 from video_processor.processor import VideoAnalyzer
 import logging
+import torch
+import spaces
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-analyzer = VideoAnalyzer()
-
-def process_video(video_path):
-    """Process video and return description"""
+@spaces.GPU
+def on_process(video):
+    # Clear all components when starting new processing
+    yield [
+        "",  # Clear status
+        "",  # Clear description
+        gr.update(visible=False)  # Hide accordion
+    ]
+    
+    if not video:
+        yield [
+            "Please upload a video",
+            "",
+            gr.update(visible=False)
+        ]
+        return
+    
     try:
-        logger.info(f"Processing video: {video_path}")
-        results = analyzer.process_video(video_path)
-        return results[0]["description"]
+        # Initialize analyzer
+        yield [
+            "Initializing video analyzer...",
+            "",
+            gr.update(visible=False)
+        ]
+        
+        analyzer = VideoAnalyzer()
+        
+        # Process video
+        yield [
+            "Analyzing video content...",
+            "",
+            gr.update(visible=True)
+        ]
+        
+        logger.info(f"Processing video: {video}")
+        result = analyzer.process_video(video)
+        description = result[0]["description"]
+        
+        # Format output
+        formatted_desc = f"### Analysis:\n{description}"
+        
+        yield [
+            "Processing complete!",
+            formatted_desc,
+            gr.update(visible=True)
+        ]
+        
     except Exception as e:
-        logger.error(f"Error processing video: {e}")
-        return str(e)
+        logger.exception("Error processing video")
+        yield [
+            f"Error processing video: {str(e)}",
+            "",
+            gr.update(visible=False)
+        ]
+    finally:
+        # Clean up
+        torch.cuda.empty_cache()
 
 # Create Gradio interface
-demo = gr.Interface(
-    fn=process_video,
-    inputs=gr.Video(label="Upload Video"),
-    outputs=gr.Textbox(label="Video Description"),
-    title="SmolVLM Video Analyzer",
-    description="Upload a video to get a detailed description of its contents."
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# SmolVLM Video Analyzer")
+    gr.Markdown("Upload a video to get a detailed analysis of its content.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_video = gr.Video(
+                label="Upload your video",
+                interactive=True
+            )
+            process_btn = gr.Button("Process Video", variant="primary")
+            
+        with gr.Column(scale=1):
+            status = gr.Markdown()
+            analysis_accordion = gr.Accordion(
+                "Analysis Details",
+                open=True,
+                visible=False
+            )
+            with analysis_accordion:
+                video_description = gr.Markdown("")
+
+    process_btn.click(
+        on_process,
+        inputs=[input_video],
+        outputs=[
+            status,
+            video_description,
+            analysis_accordion
+        ],
+        queue=True,
+    )
 
 if __name__ == "__main__":
-    demo.launch() 
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    ) 

src/video_processor/processor.py

@@ -31,7 +31,6 @@ class VideoAnalyzer:
             torch_dtype=torch.bfloat16
         )
 
-        
         self.model = AutoModelForImageTextToText.from_pretrained(
             self.model_path,
             torch_dtype=torch.bfloat16,