Update handler.py

handler.py

@@ -4,6 +4,8 @@ import numpy as np
 import os
 import requests
 import tempfile
+import gc
+import time
 from transformers import VideoLlavaProcessor, VideoLlavaForConditionalGeneration
 
 class EndpointHandler:
@@ -28,7 +30,11 @@ class EndpointHandler:
         temp_path = temp_file.name
         temp_file.close()
         try:
-            response = requests.get(video_url, stream=True)
+            # Added timeout (30s) to prevent hanging on bad URLs
+            response = requests.get(video_url, stream=True, timeout=30)
+            if response.status_code != 200:
+                raise ValueError(f"Failed to download: {response.status_code}")
+                
             with open(temp_path, 'wb') as f:
                 for chunk in response.iter_content(chunk_size=8192):
                     if chunk:
@@ -48,20 +54,39 @@ class EndpointHandler:
                 break
             if i >= start_index and i in indices:
                 frames.append(frame)
-        # Return as list of numpy arrays, which acts like a "list of images" for the processor
+        
+        # Guard clause: If video is corrupted or empty
+        if not frames:
+            raise ValueError("Video decoding failed: No frames found.")
+
+        # Return as list of numpy arrays
         return [x.to_ndarray(format="rgb24") for x in frames]
 
     def __call__(self, data):
+        start_time = time.time()
         print("\n--- NEW REQUEST ---")
+        
+        container = None
+        video_path = None
+        
         try:
             # 1. EXTRACT DATA
-            inputs = data.pop("inputs", "What is happening in this video?")
+            inputs = data.pop("inputs", "Describe this video.")
             video_url = data.pop("video", None)
             parameters = data.pop("parameters", {})
             
-            num_frames = parameters.get("num_frames", 8)
-            max_new_tokens = parameters.get("max_new_tokens", 250)
-            temperature = parameters.get("temperature", 0.1)
+            # Default to 8 frames because LanguageBind is trained on 8 frames.
+            # Only change this if you are sure the model handles interpolation.
+            num_frames = parameters.pop("num_frames", 8)
+            
+            # Clean parameters for generation (pass everything else to the model)
+            gen_kwargs = {
+                "max_new_tokens": parameters.pop("max_new_tokens", 500),
+                "temperature": parameters.pop("temperature", 0.7),
+                "top_p": parameters.pop("top_p", 0.9),
+                "do_sample": True
+            }
+            gen_kwargs.update(parameters) # Merge any other params user sent
 
             if not video_url:
                 return {"error": "Missing 'video' URL."}
@@ -71,23 +96,27 @@ class EndpointHandler:
             video_path = self.download_video(video_url)
             container = av.open(video_path)
             
-            # 3. SAMPLE FRAMES
+            # 3. SMART FRAME SAMPLING
             total_frames = container.streams.video[0].frames
             if total_frames == 0:
+                # Fallback for videos with missing metadata
                 total_frames = sum(1 for _ in container.decode(video=0))
                 container.seek(0)
             
-            # Ensure we don't request more frames than exist
+            # Clamp frames to available count
             frames_to_use = min(total_frames, num_frames)
             if frames_to_use < 1: frames_to_use = 1
-                
+            
             indices = np.linspace(0, total_frames - 1, frames_to_use, dtype=int)
             clip = self.read_video_pyav(container, indices)
             print(f"Processed {len(clip)} frames.")
 
-            # 4. PREPARE INPUTS
-            # Note: VideoLlava expects specific prompt formatting
-            full_prompt = f"USER: <video>{inputs} ASSISTANT:"
+            # 4. PREPARE PROMPT
+            # Check if user already added the template to avoid double-templating
+            if "USER:" in inputs:
+                full_prompt = inputs
+            else:
+                full_prompt = f"USER: <video>\n{inputs}\nASSISTANT:"
             
             model_inputs = self.processor(
                 text=full_prompt,
@@ -96,26 +125,39 @@ class EndpointHandler:
             ).to(self.model.device)
 
             # 5. GENERATE
-            print("Generating...")
+            print(f"Generating with params: {gen_kwargs}")
             with torch.inference_mode():
                 generate_ids = self.model.generate(
                     **model_inputs,
-                    max_new_tokens=max_new_tokens,
-                    temperature=temperature,
-                    do_sample=True if temperature > 0 else False
+                    **gen_kwargs
                 )
 
             result = self.processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-            final_output = result.split("ASSISTANT:")[-1].strip()
             
-            print(f"Result: {final_output[:50]}...")
+            # Clean output based on prompt structure
+            if "ASSISTANT:" in result:
+                final_output = result.split("ASSISTANT:")[-1].strip()
+            else:
+                final_output = result
+
+            # LOG TIME
+            duration = time.time() - start_time
+            print(f"✅ Success! Total time: {duration:.2f} seconds.")
+            print(f"Result preview: {final_output[:50]}...")
+            
             return [{"generated_text": final_output}]
 
         except Exception as e:
             import traceback
             traceback.print_exc()
             return {"error": str(e)}
+            
         finally:
-            if 'container' in locals() and container: container.close()
-            if 'video_path' in locals() and video_path and os.path.exists(video_path):
-                os.unlink(video_path)
+            # 6. CLEANUP (Crucial for long-running endpoints)
+            if container: container.close()
+            if video_path and os.path.exists(video_path):
+                os.unlink(video_path)
+            
+            # Clear GPU memory
+            torch.cuda.empty_cache()
+            gc.collect()