initial test

app.py

@@ -0,0 +1,192 @@
+import os
+import json
+import gradio as gr
+import tempfile
+from PIL import Image, ImageDraw, ImageFont
+import cv2
+from typing import Tuple, Optional
+import torch
+import spaces
+from pathlib import Path
+import time
+
+# Import your highlight detection code
+from video_highlight_detector import (
+    load_model,
+    BatchedVideoHighlightDetector,
+    get_video_duration_seconds
+)
+
+def load_examples(json_path: str) -> dict:
+    """Load pre-computed examples from JSON file"""
+    with open(json_path, 'r') as f:
+        return json.load(f)
+
+def format_duration(seconds: int) -> str:
+    """Convert seconds to MM:SS or HH:MM:SS format"""
+    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 add_watermark(video_path: str, output_path: str):
+    """Add watermark to video using ffmpeg"""
+    watermark_text = "🤗 SmolVLM2 Highlight"
+    command = f"""ffmpeg -i {video_path} -vf \
+        "drawtext=text='{watermark_text}':fontcolor=white:fontsize=24:box=1:[email protected]:\
+        boxborderw=5:x=w-tw-10:y=h-th-10" \
+        -codec:a copy {output_path}"""
+    os.system(command)
+
+def process_video(
+    video_path: str,
+    progress = gr.Progress()
+) -> Tuple[str, str, str, str]:
+    """
+    Process video and return paths to:
+    - Processed video with watermark
+    - Video description
+    - Highlight types
+    - Error message (if any)
+    """
+    try:
+        # Check video duration
+        duration = get_video_duration_seconds(video_path)
+        if duration > 1200:  # 20 minutes
+            return None, None, None, "Video must be shorter than 20 minutes"
+
+        # Load model (could be cached)
+        progress(0.1, desc="Loading model...")
+        model, processor = load_model()
+        detector = BatchedVideoHighlightDetector(model, processor)
+
+        # Analyze video content
+        progress(0.2, desc="Analyzing video content...")
+        video_description = detector.analyze_video_content(video_path)
+        
+        # Determine highlights
+        progress(0.3, desc="Determining highlight types...")
+        highlight_types = detector.determine_highlights(video_description)
+
+        # Create highlight video
+        progress(0.4, desc="Detecting and extracting highlights...")
+        with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
+            temp_output = tmp_file.name
+        
+        detector.create_highlight_video(video_path, temp_output)
+        
+        # Add watermark
+        progress(0.9, desc="Adding watermark...")
+        output_path = temp_output.replace('.mp4', '_watermark.mp4')
+        add_watermark(temp_output, output_path)
+        
+        # Cleanup
+        os.unlink(temp_output)
+        
+        # Truncate description and highlights if too long
+        video_description = video_description[:500] + "..." if len(video_description) > 500 else video_description
+        highlight_types = highlight_types[:500] + "..." if len(highlight_types) > 500 else highlight_types
+        
+        return output_path, video_description, highlight_types, None
+
+    except Exception as e:
+        return None, None, None, f"Error processing video: {str(e)}"
+
+
+def create_ui(examples_path: str):
+    """Create the Gradio interface with optional thumbnails"""
+    examples_data = load_examples(examples_path)
+
+    with gr.Blocks() as app:
+        gr.Markdown("# Video Highlight Generator")
+        gr.Markdown("Upload a video (max 20 minutes) and get an automated highlight reel!")
+        
+        # Pre-computed examples section
+        with gr.Row():
+            gr.Markdown("## Example Results")
+        
+        for example in examples_data["examples"]:
+            with gr.Row():
+                with gr.Column():
+                    # Use thumbnail if available, otherwise default to video
+                    video_component = gr.Video(
+                        example["original"]["url"],
+                        label=f"Original ({format_duration(example['original']['duration_seconds'])})",
+                        thumbnail=example["original"].get("thumbnail_url", None)
+                    )
+                    gr.Markdown(example["title"])
+                
+                with gr.Column():
+                    gr.Video(
+                        example["highlights"]["url"],
+                        label=f"Highlights ({format_duration(example['highlights']['duration_seconds'])})",
+                        thumbnail=example["highlights"].get("thumbnail_url", None)
+                    )
+                    with gr.Accordion("Analysis", open=False):
+                        gr.Markdown(example["analysis"]["video_description"])
+                        gr.Markdown(example["analysis"]["highlight_types"])
+
+        # Upload section
+        gr.Markdown("## Try It Yourself!")
+        with gr.Row():
+            input_video = gr.Video(
+                label="Upload your video (max 20 minutes)",
+                source="upload"
+            )
+
+        # Results section (initially hidden)
+        with gr.Row(visible=False) as results_row:
+            with gr.Column():
+                video_description = gr.Markdown(label="Video Analysis")
+            with gr.Column():
+                highlight_types = gr.Markdown(label="Detected Highlights")
+        
+        with gr.Row(visible=False) as output_row:
+            output_video = gr.Video(label="Highlight Video")
+            download_btn = gr.Button("Download Highlights")
+
+        # Error message
+        error_msg = gr.Markdown(visible=False)
+
+        # Process video when uploaded
+        def on_upload(video):
+            results_row.visible = False
+            output_row.visible = False
+            error_msg.visible = False
+            
+            if not video:
+                error_msg.visible = True
+                error_msg.value = "Please upload a video"
+                return None, None, None, error_msg
+            
+            output_path, desc, highlights, err = process_video(video)
+            
+            if err:
+                error_msg.visible = True
+                error_msg.value = err
+                return None, None, None, error_msg
+            
+            results_row.visible = True
+            output_row.visible = True
+            return output_path, desc, highlights, ""
+
+        input_video.change(
+            on_upload,
+            inputs=[input_video],
+            outputs=[output_video, video_description, highlight_types, error_msg]
+        )
+
+        # Download button
+        download_btn.click(
+            lambda x: x,
+            inputs=[output_video],
+            outputs=[output_video]
+        )
+
+    return app
+
+if __name__ == "__main__":
+    app = create_ui("video_spec.json")
+    app.launch()

modeling_smolvlm.py

@@ -0,0 +1,297 @@
+import torch
+from torch import nn
+from transformers import Idefics3Model, Idefics3ForConditionalGeneration
+from typing import Dict, Any, List, Optional, Union, Tuple
+from transformers.cache_utils import Cache, DynamicCache
+
+from transformers.utils import add_start_docstrings_to_model_forward, logging
+from transformers.models.idefics3.modeling_idefics3 import IDEFICS3_INPUTS_DOCSTRING, Idefics3BaseModelOutputWithPast
+
+logger = logging.get_logger(__name__)
+
+class SmolVLMModel(Idefics3Model):
+    """
+    A subclass of Idefics3Model. We do *not* remove or block the call to inputs_merger
+    in forward. Instead, we override inputs_merger here with custom logic.
+    """
+    def inputs_merger(
+        self,
+        input_ids: torch.LongTensor,
+        inputs_embeds: torch.Tensor,
+        image_hidden_states: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Merge text embeddings with image embeddings out-of-place (no in-place indexing).
+    
+        The shapes are something like:
+          - input_ids:          (B, T)
+          - inputs_embeds:      (B, T, D)
+          - image_hidden_states:(N, S, D) where N is total images across the batch,
+            S is #patches (or #slots) per image, D is embedding dim.
+    
+        Logic:
+          1) For each sample in the batch, find <image> tokens in the text.
+          2) If zero <image> tokens => text-only. Concatenate a zero-length slice
+             from image_hidden_states but do NOT advance the offset. This ensures
+             the model's image encoder is still in the computation graph, but we
+             skip "consuming" any image block for a text-only sample.
+          3) If there are <image> tokens, they appear in multiples of S for each image
+             (because each image is S embeddings). We chunk those positions into groups
+             of S. For each chunk => we consume one block from image_hidden_states[offset]
+             (which is shape (S, D)), and place each row into the text in place of a token.
+    
+        Returns:
+          A tensor of (B, T, D).
+        """
+    
+        ##############################################
+        # 1) Basic shape checks
+        ##############################################
+        #old_merger_outputs = self.inputs_merger_old(input_ids, inputs_embeds, image_hidden_states)
+        B, T, D_text = inputs_embeds.shape
+        N, S, D_img  = image_hidden_states.shape
+        if D_text != D_img:
+            raise ValueError(
+                f"Text embedding dim {D_text} != image embedding dim {D_img}"
+            )
+    
+        ##############################################
+        # 2) We'll track how many images we've used so far across the entire batch
+        ##############################################
+        image_offset = 0
+    
+        # We'll store one merged tensor per batch sample
+        merged_outputs: List[torch.Tensor] = []
+    
+        ##############################################
+        # 3) Iterate through each sample
+        ##############################################
+        for b_idx, (cur_ids, cur_embeds) in enumerate(zip(input_ids, inputs_embeds)):
+            # Find positions of <image> tokens in the text
+            image_positions = (cur_ids == self.image_token_id).nonzero(as_tuple=True)[0]
+            num_image_tokens = len(image_positions)
+    
+            # If no <image> => text-only
+            if num_image_tokens == 0:
+                # We do not consume any row from image_hidden_states; 
+                # but we do a zero-length slice so the image encoder is in the graph.
+                empty_slice = image_hidden_states[0][:0, :]  # shape (0, D)
+                # Concatenate text plus that empty slice.
+                # NOTE: this is important for DeepSpeed.
+                merged_text_only = torch.cat([cur_embeds, empty_slice], dim=0)
+                merged_outputs.append(merged_text_only)
+                continue
+    
+            # Otherwise, we have at least one <image> token.
+            # Typically, if each image is S embeddings, we expect the total # of <image> tokens
+            # in this sample to be multiple of S => each group of S tokens = 1 image
+            if num_image_tokens % S != 0:
+                raise ValueError(
+                    f"Sample {b_idx} has {num_image_tokens} <image> tokens, not a multiple of S={S}. "
+                    "Cannot map them to blocks of shape (S, D)."
+                )
+    
+            # We'll chunk image_positions into groups of size S
+            positions_list = image_positions.tolist()
+            # Example: if num_image_tokens=162 and S=81 => we have 2 images => 2 chunks each of length 81
+            chunks = [
+                positions_list[i : i + S]
+                for i in range(0, num_image_tokens, S)
+            ]
+    
+            # We'll build a list of segments: text, then image row(s), text, etc.
+            segments = []
+            text_start = 0
+    
+            # For each chunk (each chunk => 1 image)
+            for chunk in chunks:
+                # image_hidden_states[image_offset] => shape (S, D)
+                cur_block = image_hidden_states[image_offset]
+                image_offset += 1
+    
+                # We'll iterate over the S positions in ascending order
+                for i_s, pos in enumerate(chunk):
+                    # Add text from [text_start..pos)
+                    if pos > text_start:
+                        segments.append(cur_embeds[text_start:pos])
+                    # Then add one row from cur_block => shape (1, D)
+                    row_of_block = cur_block[i_s : i_s + 1, :]
+                    segments.append(row_of_block)
+                    # skip the <image> token
+                    text_start = pos + 1
+    
+            # leftover text after the final <image> token
+            if text_start < T:
+                segments.append(cur_embeds[text_start:])
+    
+            # cat them into a single (T_b, D) tensor
+            merged_sample = torch.cat(segments, dim=0)
+            merged_outputs.append(merged_sample)
+            
+        merged_outputs = torch.stack(merged_outputs)
+        #assert (old_merger_outputs==merged_outputs).all()
+        return merged_outputs
+
+
+    @add_start_docstrings_to_model_forward(
+        """
+        Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to
+        the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where
+        max_num_images is the maximum number of images among the batch_size samples in the batch.
+        Padding images are not needed beyond padding the pixel_values at the entrance of the model.
+        For efficiency, we only pass through the vision_model's forward the real images by
+        discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where
+        image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3.
+        """,
+        IDEFICS3_INPUTS_DOCSTRING,
+    )
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        pixel_attention_mask: Optional[torch.BoolTensor] = None,
+        image_hidden_states: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, Idefics3BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.training and self.text_model.gradient_checkpointing and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+            )
+            use_cache = False
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        past_seen_tokens = 0
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = DynamicCache()
+            past_seen_tokens = past_key_values.get_seq_length()
+
+        if inputs_embeds is not None and input_ids is None and past_seen_tokens == 0:
+            raise ValueError("When first calling the model, if input_embeds are passed, input_ids should not be None.")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.text_model.get_input_embeddings()(input_ids).to(self.device)
+
+        # START VISUAL INPUTS INTEGRATION
+        if pixel_values is not None and image_hidden_states is not None:
+            raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time")
+        elif pixel_values is not None:
+            batch_size, num_images, num_channels, height, width = pixel_values.shape
+            pixel_values = pixel_values.to(dtype=self.dtype)  # fp16 compatibility
+            pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:])
+
+            # Remove padding images - padding images are full 0.
+            nb_values_per_image = pixel_values.shape[1:].numel()
+            real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image
+            
+            if not any(real_images_inds):
+                # no images, leave one empty image.
+                real_images_inds[0] = True
+                
+            pixel_values = pixel_values[real_images_inds].contiguous()
+            
+            # Handle the vision attention mask
+            if pixel_attention_mask is None:
+                pixel_attention_mask = torch.ones(
+                    size=(pixel_values.size(0), pixel_values.size(2), pixel_values.size(3)),
+                    dtype=torch.bool,
+                    device=pixel_values.device,
+                )
+            else:
+                # Remove padding images from the mask
+                pixel_attention_mask = pixel_attention_mask.view(
+                    batch_size * num_images, *pixel_attention_mask.shape[2:]
+                )
+                pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous()
+
+            patch_size = self.config.vision_config.patch_size
+            patches_subgrid = pixel_attention_mask.unfold(dimension=1, size=patch_size, step=patch_size)
+            patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size)
+            patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool()
+
+            # Get sequence from the vision encoder
+            image_hidden_states = self.vision_model(
+                pixel_values=pixel_values,
+                patch_attention_mask=patch_attention_mask,
+            ).last_hidden_state
+            
+            # Modality projection & resampling
+            image_hidden_states = self.connector(image_hidden_states)
+
+        elif image_hidden_states is not None:
+            image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device)
+
+        if past_seen_tokens == 0 and inputs_embeds is not None and image_hidden_states is not None:
+            # When we generate, we don't want to replace the potential image_token_id that we generated by images
+            # that simply don't exist
+            inputs_embeds = self.inputs_merger(
+                input_ids=input_ids,
+                inputs_embeds=inputs_embeds,
+                image_hidden_states=image_hidden_states,
+            )
+
+        outputs = self.text_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        if not return_dict:
+            return tuple(v for v in [*outputs, image_hidden_states] if v is not None)
+
+        return Idefics3BaseModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=image_hidden_states,
+        )
+
+
+
+
+class SmolVLMForConditionalGeneration(Idefics3ForConditionalGeneration):
+    """
+    A subclass of Idefics3ForConditionalGeneration that uses MyIdefics3Model
+    instead of the default Idefics3Model.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        # Instead of the original self.model = Idefics3Model(config),
+        # we point to our custom class.
+        self.model = SmolVLMModel(config)
+
+        # We *keep* the same lm_head from the parent, or re-init if you prefer:
+        self.lm_head = nn.Linear(
+            config.text_config.hidden_size, config.text_config.vocab_size, bias=False
+        )
+
+        # If parent sets up any post_init() logic:
+        self.post_init()

requirements.txt

@@ -0,0 +1,4 @@
+Pillow
+opencv-python
+num2words
+ffmpeg-python

video_highlight_detector.py

@@ -0,0 +1,785 @@
+import torch
+from typing import List, Tuple, Dict, Optional
+from tqdm import tqdm
+import logging
+from PIL import Image
+import datetime
+from num2words import num2words
+import subprocess
+import sys
+from modeling_smolvlm import SmolVLMForConditionalGeneration
+from transformers import AutoProcessor, AutoTokenizer
+import json
+import math
+logger = logging.getLogger(__name__)
+logging.basicConfig(
+    level=logging.INFO
+)
+
+
+
+SYSTEM_MESSAGE = (
+    "Carefully watch the video and pay attention to the cause and sequence of events, "
+    "the detail and movement of objects, and the action and pose of persons. "
+    "Based on your observations, answer the question with yes or no."
+    " <end_of_utterance>"
+)
+
+FRAME_TIMESTAMP_MESSAGE = "Frame from"
+DEFAULT_VIDEO_INTRO = (
+    "You are provided the following series of {frame_count} frames "
+    "from a {video_duration} [H:MM:SS] video.\n"
+)
+
+# ----------------------------------------------------------------------
+# Helper functions for resizing, etc.
+# ----------------------------------------------------------------------
+
+def round_by_factor(number: float, factor: int) -> int:
+    return round(number / factor) * factor
+
+def ceil_by_factor(number: float, factor: int) -> int:
+    return math.ceil(number / factor) * factor
+
+def floor_by_factor(number: float, factor: int) -> int:
+    return math.floor(number / factor) * factor
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int,
+    min_pixels: int,
+    max_pixels: int,
+    max_ratio: float,
+) -> Tuple[int, int]:
+    """
+    Rescale (height, width) so that:
+      - aspect ratio <= max_ratio
+      - total area in [min_pixels, max_pixels]
+      - each dimension is multiple of factor
+    """
+    ratio = max(height, width) / min(height, width)
+    if ratio > max_ratio:
+        raise ValueError(f"Aspect ratio {ratio:.2f} > {max_ratio}")
+
+    h_ = max(factor, round_by_factor(height, factor))
+    w_ = max(factor, round_by_factor(width, factor))
+    area = h_ * w_
+
+    if area > max_pixels:
+        scale = math.sqrt((height * width) / max_pixels)
+        h_ = floor_by_factor(height / scale, factor)
+        w_ = floor_by_factor(width / scale, factor)
+    elif area < min_pixels:
+        scale = math.sqrt(min_pixels / (height * width))
+        h_ = ceil_by_factor(height * scale, factor)
+        w_ = ceil_by_factor(width * scale, factor)
+    return h_, w_
+
+def _smart_nframes(
+    total_frames: int,
+    video_fps: float,
+    frame_factor: int = 1,
+    target_fps: float = 2.0,
+    min_frames: int = 4,
+    max_frames: int = 32
+) -> int:
+    """
+    Decide how many frames to pick from a range based on target FPS.
+    Result is clamped to [min_frames, max_frames] and must be multiple of frame_factor.
+    """
+    minf = ceil_by_factor(min_frames, frame_factor)
+    maxf = floor_by_factor(min(max_frames, total_frames), frame_factor)
+    val = total_frames / video_fps * target_fps
+    val = min(max(val, minf), maxf)
+    nframes = round_by_factor(val, frame_factor)
+
+    if not (frame_factor <= nframes <= total_frames):
+        raise ValueError(f"Invalid nframes={nframes}, out of range.")
+    return int(nframes)
+
+
+def get_video_duration_seconds(video_path: str) -> float:
+    """
+    Use ffprobe to retrieve the total duration of a video (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"])
+
+def get_fixed_30s_segments(video_path: str) -> list:
+    """
+    Produce a list of (start_sec, end_sec) tuples in 30-second blocks
+    for the entire video.
+    """
+    duration = get_video_duration_seconds(video_path)
+    segments = []
+    start = 0.0
+    block_size = 10.0
+
+    while start < duration:
+        end = min(start + block_size, duration)
+        segments.append((start, end))
+        start = end
+    
+    return segments
+
+
+
+class SmartVideoFrameExtractor:
+    """
+    This class extracts frames from a specific portion of a video
+    (defined by start_frame and end_frame or start_sec and end_sec).
+    """
+    def __init__(
+        self,
+        frame_factor: int = 1,
+        min_pixels: int = 384 * 384,
+        max_pixels: int = 384 * 384 * 4,
+        max_ratio: float = 2.0
+    ):
+        self.frame_factor = frame_factor
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.max_ratio = max_ratio
+
+        try:
+            import decord
+            self.reader = "decord"
+            decord.bridge.set_bridge("torch")
+        except ImportError:
+            self.reader = "torchvision"
+            logger.info("Decord not found, falling back to torchvision")
+
+    def extract_frames(
+        self,
+        video_path: str,
+        start_sec: float,
+        end_sec: float,
+        target_fps: float = 1.0,
+        min_frames: int = 4,
+        max_frames: int = 32
+    ) -> Tuple[List[Image.Image], List[str]]:
+        """Extract frames from [start_sec, end_sec] using decord or torchvision."""
+        if self.reader == "decord":
+            return self._extract_frames_decord(
+                video_path, start_sec, end_sec, target_fps, min_frames, max_frames
+            )
+        else:
+            return self._extract_frames_torchvision(
+                video_path, start_sec, end_sec, target_fps, min_frames, max_frames
+            )
+
+    def _extract_frames_decord(
+        self,
+        video_path: str,
+        start_sec: float,
+        end_sec: float,
+        target_fps: float,
+        min_frames: int,
+        max_frames: int
+    ) -> Tuple[List[Image.Image], List[str]]:
+        """Extract frames with decord from a certain segment."""
+        import decord
+        from decord import VideoReader
+
+        vr = VideoReader(video_path)
+        total_frames = len(vr)
+        video_fps = vr.get_avg_fps()
+
+        # Convert start/end times to frame indices
+        start_frame = int(start_sec * video_fps)
+        end_frame = min(int(end_sec * video_fps), total_frames - 1)
+        if start_frame >= end_frame:
+            return [], []
+
+        working_frames = end_frame - start_frame + 1
+        nframes = _smart_nframes(
+            working_frames,
+            video_fps,
+            self.frame_factor,
+            target_fps,
+            min_frames,
+            max_frames
+        )
+        indices = torch.linspace(start_frame, end_frame, nframes).round().long()
+
+        frames_tensor = vr.get_batch(indices).cpu()  # NHWC
+        frames = []
+        timestamps = []
+
+        for i, frame_idx in enumerate(indices):
+            frame = frames_tensor[i].numpy()
+            pil_image = Image.fromarray(frame).convert("RGB")
+
+            # Compute timestamp
+            sec = frame_idx.item() / video_fps
+            mm = int(sec // 60)
+            ss = int(sec % 60)
+            timestamps.append(f"{mm:02d}:{ss:02d}")
+
+            # Resize
+            w, h = pil_image.size
+            rh, rw = smart_resize(
+                h, w,
+                factor=8,
+                min_pixels=self.min_pixels,
+                max_pixels=self.max_pixels,
+                max_ratio=self.max_ratio
+            )
+            pil_image = pil_image.resize((rw, rh), Image.Resampling.LANCZOS)
+            frames.append(pil_image)
+
+        return frames, timestamps, end_sec - start_sec
+
+    def _extract_frames_torchvision(
+        self,
+        video_path: str,
+        start_sec: float,
+        end_sec: float,
+        target_fps: float,
+        min_frames: int,
+        max_frames: int
+    ) -> Tuple[List[Image.Image], List[str]]:
+        """Extract frames with torchvision from a certain segment."""
+        from torchvision import io
+
+        # Read entire video (beware of memory usage on large videos!)
+        vid, _, info = io.read_video(
+            video_path,
+            start_pts=0,
+            end_pts=None,
+            pts_unit="sec",
+            output_format="TCHW"
+        )
+
+        total_frames = vid.size(0)
+        video_fps = info["video_fps"]
+
+        # Convert start/end times to frame indices
+        start_frame = int(start_sec * video_fps)
+        end_frame = min(int(end_sec * video_fps), total_frames - 1)
+        if start_frame >= end_frame:
+            return [], []
+
+        working_frames = end_frame - start_frame + 1
+        nframes = _smart_nframes(
+            working_frames,
+            video_fps,
+            self.frame_factor,
+            target_fps,
+            min_frames,
+            max_frames
+        )
+        indices = torch.linspace(start_frame, end_frame, nframes).round().long()
+
+        frames = []
+        timestamps = []
+        for idx in indices:
+            frame = vid[idx].permute(1, 2, 0).numpy()
+            pil_image = Image.fromarray(frame).convert("RGB")
+
+            sec = idx.item() / video_fps
+            mm = int(sec // 60)
+            ss = int(sec % 60)
+            timestamps.append(f"{mm:02d}:{ss:02d}")
+
+            w, h = pil_image.size
+            rh, rw = smart_resize(
+                h, w,
+                factor=8,
+                min_pixels=self.min_pixels,
+                max_pixels=self.max_pixels,
+                max_ratio=self.max_ratio
+            )
+            pil_image = pil_image.resize((rw, rh), Image.Resampling.LANCZOS)
+            frames.append(pil_image)
+
+        return frames, timestamps, end_sec - start_sec
+
+
+class BatchedVideoHighlightDetector:
+    """
+    Optimized version of video highlight detection that processes multiple segments
+    in parallel using batched inference.
+    """
+    def __init__(
+        self,
+        model,
+        processor,
+        device="cuda",
+        batch_size=8,
+        max_frames_per_segment=32,
+        target_fps=1.0
+    ):
+        self.model = model
+        self.processor = processor
+        self.device = device
+        self.batch_size = batch_size
+        self.max_frames_per_segment = max_frames_per_segment
+        self.target_fps = target_fps
+
+    def _extract_frames_batch(
+        self,
+        video_path: str,
+        segments: List[Tuple[float, float]]
+    ) -> List[Tuple[List[Image.Image], List[str], float]]:
+        """
+        Extract frames from multiple segments in parallel using decord's batch capabilities.
+        """
+        import decord
+        from decord import VideoReader
+        decord.bridge.set_bridge("torch")
+
+        # Open video once for all segments
+        vr = VideoReader(video_path)
+        video_fps = vr.get_avg_fps()
+        results = []
+
+        for start_sec, end_sec in segments:
+            # Convert time to frame indices
+            start_frame = int(start_sec * video_fps)
+            end_frame = min(int(end_sec * video_fps), len(vr) - 1)
+            
+            # Calculate number of frames to sample
+            segment_duration = end_sec - start_sec
+            desired_frames = min(
+                int(segment_duration * self.target_fps),
+                self.max_frames_per_segment
+            )
+            
+            # Generate frame indices
+            indices = torch.linspace(start_frame, end_frame, desired_frames).round().long()
+            
+            # Extract frames
+            frames_tensor = vr.get_batch(indices).cpu()  # NHWC format
+            
+            # Convert to PIL and generate timestamps
+            frames = []
+            timestamps = []
+            for i, frame_idx in enumerate(indices):
+                frame = frames_tensor[i].numpy()
+                pil_image = Image.fromarray(frame).convert("RGB")
+                
+                # Resize maintaining aspect ratio
+                w, h = pil_image.size
+                scale = min(384 / w, 384 / h)
+                new_w = int(w * scale)
+                new_h = int(h * scale)
+                pil_image = pil_image.resize((new_w, new_h), Image.Resampling.LANCZOS)
+                
+                frames.append(pil_image)
+                
+                # Generate timestamp
+                sec = frame_idx.item() / video_fps
+                mm = int(sec // 60)
+                ss = int(sec % 60)
+                timestamps.append(f"{mm:02d}:{ss:02d}")
+            
+            results.append((frames, timestamps, segment_duration))
+            
+        return results
+
+    def _prepare_batch_inputs(
+        self,
+        frame_batches: List[Tuple[List[Image.Image], List[str], float]],
+        highlight_types: str
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Convert a batch of frame sequences into model inputs.
+        """
+        conversations = []
+        all_frames = []
+        
+        for frames, timestamps, duration in frame_batches:
+            # Build conversation for each segment
+            conversation = [
+                {
+                    "role": "system",
+                    "content": [{
+                        "type": "text",
+                        "text": "You are a helpful assistant that analyzes videos for specific moments of interest."
+                    }]
+                },
+                {
+                    "role": "user",
+                    "content": []
+                }
+            ]
+            
+            # Add video intro
+            conversation[1]["content"].append({
+                "type": "text",
+                "text": f"You are provided the following series of {num2words(len(frames))} frames from a {str(datetime.timedelta(seconds=duration))} [H:MM:SS] video.\n"
+            })
+            
+            # Add frames with timestamps
+            for ts, frame in zip(timestamps, frames):
+                conversation[1]["content"].extend([
+                    {
+                        "type": "text",
+                        "text": f"Frame from {ts}:"
+                    },
+                    {
+                        "type": "image"
+                    }
+                ])
+            
+            # Add highlight check question
+            conversation[1]["content"].append({
+                "type": "text",
+                "text": f"""Do you see any of the following types of highlight moments in these frames? 
+                
+                Potential highlights to look for:
+                {highlight_types}
+                
+                Only answer yes if you see any of those moments and answer no if you don't."""
+            })
+            
+            conversations.append(conversation)
+            all_frames.extend(frames)
+
+        # Convert to model inputs using processor
+        prompts = [
+            self.processor.apply_chat_template(conv, add_generation_prompt=True)
+            for conv in conversations
+        ]
+        
+        # Create batched inputs
+        model_inputs = self.processor(
+            text=prompts,
+            images=all_frames,
+            return_tensors="pt",
+            padding=True
+        ).to(self.device)
+        
+        return model_inputs
+
+    def _process_segment_batch(
+        self,
+        video_path: str,
+        segments: List[Tuple[float, float]],
+        highlight_types: str
+    ) -> List[bool]:
+        """
+        Process a batch of segments and return which ones contain highlights.
+        """
+        # Extract frames for all segments in batch
+        frame_batches = self._extract_frames_batch(video_path, segments)
+        
+        # Prepare model inputs
+        model_inputs = self._prepare_batch_inputs(frame_batches, highlight_types)
+        
+        # Generate responses for entire batch
+        outputs = self.model.generate(
+            **model_inputs,
+            max_new_tokens=256,
+            num_beams=5,
+            temperature=0.7,
+            do_sample=True,
+            use_cache=True
+        )
+        
+        # Process responses
+        responses = [
+            self.processor.decode(output, skip_special_tokens=True).lower().split("assistant:")[1]
+            for output in outputs
+        ]
+        
+        # Check for "yes" in responses
+        return ["yes" in response for response in responses]
+
+    def create_highlight_video(self, video_path: str, output_path: str) -> List[Tuple[float, float]]:
+        """
+        Main function that executes the batched highlight detection pipeline.
+        """
+        # Step 1: Analyze video content
+        logger.info("Step 1: Analyzing video content...")
+        video_description = self.analyze_video_content(video_path)
+        logger.info(f"Video description: {video_description}")
+        
+        # Step 2: Determine highlight types
+        logger.info("Step 2: Determining highlight types...")
+        highlight_types = self.determine_highlights(video_description)
+        logger.info(f"Looking for highlights: {highlight_types}")
+        
+        # Step 3: Get all segments
+        segments = self._get_fixed_30s_segments(video_path)
+        
+        # Step 4: Process segments in batches
+        logger.info("Step 3: Detecting highlight segments in batches...")
+        kept_segments = []
+        
+        for i in tqdm(range(0, len(segments), self.batch_size)):
+            batch_segments = segments[i:i + self.batch_size]
+            keep_flags = self._process_segment_batch(video_path, batch_segments, highlight_types)
+            
+            for segment, keep in zip(batch_segments, keep_flags):
+                if keep:
+                    kept_segments.append(segment)
+                    logger.info(f"\tKeeping segment {segment}")
+        
+        # Step 5: Create final video
+        if kept_segments:
+            logger.info(f"Creating highlight video with {len(kept_segments)} segments...")
+            self._concatenate_scenes(video_path, kept_segments, output_path)
+        else:
+            logger.info("No highlights detected")
+            
+        return kept_segments
+
+       
+    def analyze_video_content(self, video_path: str, sample_rate: float = 0.2) -> str:
+        """
+        Step 1: Sample frames from the full video and get a general description
+        """
+        extractor = SmartVideoFrameExtractor()
+        duration = get_video_duration_seconds(video_path)
+        
+        # Sample frames from entire video
+        frames, timestamps, duration_seconds = extractor.extract_frames(
+            video_path, 
+            start_sec=0, 
+            end_sec=duration,
+            target_fps=sample_rate,
+            max_frames=32  # Limit total frames to not overwhelm model
+        )
+        
+        # Build conversation asking for video 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."
+        conversation = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": system_message}]
+            },
+            {
+                "role": "user",
+                "content": []
+            }
+        ]
+        
+        # Add video intro using DEFAULT_VIDEO_INTRO
+        conversation[1]["content"].append({
+            "type": "text",
+            "text": DEFAULT_VIDEO_INTRO.format(
+                frame_count=num2words(len(frames)), 
+                video_duration=str(datetime.timedelta(seconds=duration_seconds))
+            )
+        })
+        
+        # Add frames with timestamps
+        for ts, frame in zip(timestamps, frames):
+            conversation[1]["content"].extend([
+                {
+                    "type": "text",
+                    "text": f"{FRAME_TIMESTAMP_MESSAGE} {ts}:"
+                },
+                {
+                    "type": "image"
+                }
+            ])
+        
+        # Add question
+        conversation[1]["content"].append({
+            "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."
+        })
+        
+        # Get model response
+        prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True)
+        model_inputs = self.processor(
+            text=prompt,
+            images=frames,
+            return_tensors="pt"
+        ).to(self.model.device)
+        
+        outputs = self.model.generate(
+            **model_inputs,
+            max_new_tokens=512,
+            num_beams=5,
+            temperature=0.7,
+            do_sample=True,
+            use_cache=True
+        )
+        return self.processor.decode(outputs[0], skip_special_tokens=True).split("Assistant:")[1]
+
+    def determine_highlights(self, video_description: str) -> str:
+        """
+        Step 2: Based on video description, determine what would constitute highlights
+        """
+        conversation = [{
+            "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. For sports content, you typically select only 3-5 of the most remarkable moments that would make viewers say 'I can't believe that happened!'"}]
+        }, {
+            "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 higlight."""
+            }]
+        }]
+#                 Based on this description, what unique segments should be included in a highlight video? list moments that cannot be missed and their description, nothing else."""
+
+        #                Based on this description, what unique segments should be included in a highlight video? list moments that cannot be missed."""
+
+        prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True)
+        model_inputs = self.processor(text=prompt, return_tensors="pt").to(self.model.device)
+        
+        outputs = self.model.generate(
+            **model_inputs, 
+            max_new_tokens=256,
+            num_beams=5,
+            temperature=0.7,
+            do_sample=True
+        )
+        return self.processor.decode(outputs[0], skip_special_tokens=True).split("Assistant:")[1]
+
+
+    def _get_fixed_30s_segments(self, video_path: str) -> List[Tuple[float, float]]:
+        """Helper to get video segments"""
+        duration = self._get_video_duration_seconds(video_path)
+        segments = []
+        start = 0.0
+        block_size = 10.0
+        
+        while start < duration:
+            end = min(start + block_size, duration)
+            segments.append((start, end))
+            start = end
+        
+        return segments
+
+    def _get_video_duration_seconds(self, video_path: str) -> float:
+        """Helper to get video duration"""
+        import json
+        import subprocess
+        
+        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"])
+
+    def _concatenate_scenes(
+        self,
+        video_path: str,
+        scene_times: List[Tuple[float, float]],
+        output_path: str
+    ):
+        """
+        Concatenate selected (start_sec, end_sec) scenes from 'video_path' into 'output_path'
+        using a complex ffmpeg filter instead of multiple intermediate files.
+        """
+
+        if not scene_times:
+            logger.warning("No scenes to concatenate, skipping.")
+            return
+
+        # Build the filter_complex string
+        # For each scene i, we create two filter chains: one for video [vN] and one for audio [aN].
+        # Then we feed them into the concat filter.
+        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}]")
+
+        # Now build the actual concat invocation.
+        # n = number of segments to concat, v=1 video stream, a=1 audio stream
+        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
+
+        # Build the ffmpeg command
+        cmd = [
+            "ffmpeg",
+            "-y",  # overwrite
+            "-i", video_path,
+            "-filter_complex", filter_complex,
+            "-map", "[outv]",
+            "-map", "[outa]",
+            "-c:v", "libx264",   # or any codec of your choice
+            "-c:a", "aac",       # or any audio codec of your choice
+            output_path
+        ]
+
+        logger.info(f"Running ffmpeg command: {' '.join(cmd)}")
+        subprocess.run(cmd, check=True)
+        logger.info(f"Final video saved to: {output_path}")
+
+
+
+def load_model(
+    checkpoint_path: Optional[str] = None,
+    base_model_id: str = "HuggingFaceTB/SmolVLM-2.2B-Instruct",
+    device: str = "cuda"
+):
+    """Load the model and processor."""
+    # For demonstration, we set the target size
+    video_target_size = 384
+
+    processor = AutoProcessor.from_pretrained(base_model_id)
+    # Configure the image processor
+    processor.image_processor.size = {"longest_edge": video_target_size}
+    processor.image_processor.do_resize = True
+    processor.image_processor.do_image_splitting = False
+
+    if checkpoint_path:
+        model = SmolVLMForConditionalGeneration.from_pretrained(
+            checkpoint_path,
+            torch_dtype=torch.bfloat16,
+            device_map=device
+        )
+    else:
+        model =  SmolVLMForConditionalGeneration.from_pretrained(
+            base_model_id,
+            torch_dtype=torch.bfloat16,
+            device_map=device
+        )
+    
+    return model, processor
+
+
+def main():
+    checkpoint_path = "/fsx/miquel/smolvlmvideo/checkpoints/final-visionUnfrozen-balanced/checkpoint-6550"
+    base_model_id = "HuggingFaceTB/SmolVLM-2.2B-Instruct"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    model, processor = load_model(checkpoint_path, base_model_id, device)
+    detector = BatchedVideoHighlightDetector(model, processor, device=device)
+
+    if len(sys.argv) < 3:
+        print("Usage: python video_highlight_detector.py <input_video> <output_video>")
+        sys.exit(1)
+
+    video_path = sys.argv[1]
+    output_path = sys.argv[2]
+
+    # Create highlight video
+    highlight_segments = detector.create_highlight_video(video_path, output_path)
+    print(f"Created highlight video with {len(highlight_segments)} segments")
+
+
+if __name__ == "__main__":
+    main()

video_spec.json

@@ -0,0 +1,62 @@
+{
+    "version": "1.0",
+    "examples": [
+      {
+        "id": "Example 1",
+        "title": "Football Match Highlights",
+        "description": "Champions League semifinal match",
+        "original": {
+          "url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/barcamadrid.mp4",
+          "duration_seconds": 6114
+        },
+        "highlights": {
+          "url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/barcamadridhighlights.mp4",
+          "duration_seconds": 130,
+          "thumbnail_url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/thumbnail_barcamadrid.png" 
+
+        },
+        "analysis": {
+          "video_description": "This is a high-stakes football match between Barcelona and Madrid.",
+          "highlight_types": "- Goals scored\n- Player interactions\n- Vibe at the stadium"
+        }
+      },
+      {
+        "id": "Example 2",
+        "title": "Football Match Highlights",
+        "description": "Champions League semifinal match",
+        "original": {
+          "url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/barcamadrid.mp4",
+          "duration_seconds": 6114
+        },
+        "highlights": {
+          "url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/barcamadridhighlights.mp4",
+          "duration_seconds": 130,
+          "thumbnail_url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/thumbnail_barcamadrid.png" 
+
+        },
+        "analysis": {
+          "video_description": "This is a high-stakes football match between Barcelona and Madrid.",
+          "highlight_types": "- Goals scored\n- Player interactions\n- Vibe at the stadium"
+        }
+      },
+      {
+        "id": "Example 3",
+        "title": "Football Match Highlights",
+        "description": "Champions League semifinal match",
+        "original": {
+          "url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/barcamadrid.mp4",
+          "duration_seconds": 6114
+        },
+        "highlights": {
+          "url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/barcamadridhighlights.mp4",
+          "duration_seconds": 130,
+          "thumbnail_url": "https://huggingface.co/datasets/mfarre/servedfiles/resolve/main/thumbnail_barcamadrid.png" 
+
+        },
+        "analysis": {
+          "video_description": "This is a high-stakes football match between Barcelona and Madrid.",
+          "highlight_types": "- Goals scored\n- Player interactions\n- Vibe at the stadium"
+        }
+      }
+    ]
+  }