update alg

app.py

@@ -108,18 +108,30 @@ class VideoHighlightDetector:
         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]
 
-    def determine_highlights(self, video_description: str) -> str:
-        """Determine what constitutes highlights based on video description."""
+    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 any video 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. Generally, highlights should be relatively 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": "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 any video of this type."}]
+                "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\nList potential highlight moments to look for in this video:"""}]
+                "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,
@@ -265,7 +277,7 @@ def create_ui(examples_path: str, model_path: str):
                     gr.update(visible=False)
                 ]
 
-                detector = VideoHighlightDetector(model_path=model_path)
+                detector = VideoHighlightDetector(model_path=model_path, batch_size=16)
 
                 yield [
                     None,
@@ -287,20 +299,21 @@ def create_ui(examples_path: str, model_path: str):
                     gr.update(visible=True)
                 ]
 
-                # Determine highlight types
-                highlights = detector.determine_highlights(video_desc)
-                formatted_highlights = f"### Highlight Criteria:\n{highlights}"
+                highlights1 = detector.determine_highlights(video_desc, prompt_num=1)
+                highlights2 = detector.determine_highlights(video_desc, prompt_num=2)
+                formatted_highlights = f"### Highlight Criteria:\nSet 1:\n{highlights1}\n\nSet 2:\n{highlights2}"
                 
                 # Process video in segments
                 segment_length = 10.0
-                kept_segments = []
-                segment_descriptions = []
+                kept_segments1 = []
+                kept_segments2 = []
+                segment_descriptions1 = []
+                segment_descriptions2 = []
                 segments_processed = 0
                 total_segments = int(duration / segment_length)
 
                 for start_time in range(0, int(duration), int(segment_length)):
                     end_time = min(start_time + segment_length, duration)
-                    segments_processed +=1
                     progress = int((segments_processed / total_segments) * 100)
 
                     yield [
@@ -325,24 +338,56 @@ def create_ui(examples_path: str, model_path: str):
                         ]
                         subprocess.run(cmd, check=True)
                         
-                        if detector.process_segment(temp_segment.name, highlights):
-                            # Get segment description
+                        # Process with both highlight sets
+                        if detector.process_segment(temp_segment.name, highlights1):
+                            description = detector.analyze_segment(temp_segment.name)
+                            kept_segments1.append((start_time, end_time))
+                            segment_descriptions1.append(description)
+                            
+                        if detector.process_segment(temp_segment.name, highlights2):
                             description = detector.analyze_segment(temp_segment.name)
-                            kept_segments.append((start_time, end_time))
-                            segment_descriptions.append(description)
+                            kept_segments2.append((start_time, end_time))
+                            segment_descriptions2.append(description)
+
+                    segments_processed += 1
+
+                # 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)
                 
-                if kept_segments:
+                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
+                if (0 < percent2 <= percent1 or percent1 == 0):
+                    final_segments = kept_segments2
+                    segment_descriptions = segment_descriptions2
+                    selected_set = "2"
+                    percent_used = percent2
+                else:
+                    final_segments = kept_segments1
+                    segment_descriptions = segment_descriptions1
+                    selected_set = "1"
+                    percent_used = percent1
+
+                if final_segments:
                     # Create XSPF playlist
-                    playlist_content = create_xspf_playlist(video, kept_segments, segment_descriptions)
+                    playlist_content = create_xspf_playlist(video, final_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
                     
+                    completion_message = f"Processing complete! Using highlight set {selected_set} ({percent_used:.1f}% of video). You can download the playlist."
+                    
                     yield [
                         gr.update(value=playlist_path, visible=True),
-                        "Processing complete! You can download the playlist.",
+                        completion_message,
                         formatted_desc,
                         formatted_highlights,
                         gr.update(visible=True)