Update processing/two_stage/two_stage_processor.py

processing/two_stage/two_stage_processor.py

@@ -1,18 +1,28 @@
 #!/usr/bin/env python3
 """
-ENHANCED Two-Stage Alpha Channel Processing System ✅ 2025-09-09
+ENHANCED Two-Stage Alpha Channel Processing System ✅
+(Full file, ready to drop in)
+Date: 2025-09-09
+
 VIDEO-TO-VIDEO PIPELINE with proper audio support via AudioProcessor
 
-Key fixes:
-- Uses AudioProcessor for proper ffmpeg-based audio handling
-- Ensures full video duration processing (no 5-second truncation)
-- Direct alpha compositing without green screen
+Key fixes in this build:
+- Uses AudioProcessor for ffmpeg-based audio extraction & muxing
+- Ensures full video duration processing (no 5-second truncation unless you pass trim_seconds explicitly)
+- Direct alpha compositing (no green screen intermediate), supports image or video backgrounds
+- Chunked processing for long videos, with seamless reassembly (video first, audio added after)
+- Robust MatAnyone integration (supports both known signature and positional fallback)
+- Defensive handling for SAM2 mask creation; combines first N frames for a stable trimap
+- Careful memory cleanup between stages (GPU/system) to avoid leaks in long runs
+- Extra logging and debug artifacts to help diagnose issues quickly
 """
 
 # ==============================================================================
 # IMPORTS
 # ==============================================================================
 
+from __future__ import annotations
+
 import os
 import sys
 import cv2
@@ -23,35 +33,39 @@
 import gc
 import time
 import traceback
-import subprocess
 from pathlib import Path
 from typing import Optional, Tuple, Dict, Any, List
-from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip, ImageClip, concatenate_videoclips, VideoClip
-import moviepy.video.fx.all as vfx
+
+# MoviePy (only used for probing durations, optional trimming, concat *video only*, and fallback audio mux)
+from moviepy.editor import (
+    VideoFileClip, AudioFileClip, concatenate_videoclips
+)
+import moviepy.video.fx.all as vfx  # noqa: F401 (kept for future effects)
 
 # Import the AudioProcessor for proper audio handling
+# Search both "processing.audio" and local "audio"
 try:
     from processing.audio import AudioProcessor
-except ImportError:
+except Exception:
     try:
         from audio import AudioProcessor  # Fallback if in different location
-    except ImportError:
+    except Exception:
         AudioProcessor = None
-        logging.warning("AudioProcessor not found - audio features will be limited")
 
-# Setup logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    datefmt='%Y-%m-%d %H:%M:%S'
-)
+# Setup logging (do not override global logging level outside this module)
 logger = logging.getLogger(__name__)
+if not logger.handlers:
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
 
-# PyTorch memory management
+# PyTorch memory management (optional, if available on the host)
 try:
     import torch
     TORCH_AVAILABLE = True
-except ImportError:
+except Exception:
     TORCH_AVAILABLE = False
     logger.warning("PyTorch not available, GPU memory management disabled")
 
@@ -66,11 +80,11 @@ class ProcessingConfig:
     REFERENCE_FRAMES = 3  # Number of frames to use for reference mask
     
     # Chunked processing settings
-    MAX_CHUNK_DURATION = 300  # 5 minutes in seconds
+    MAX_CHUNK_DURATION = 300  # seconds (5 minutes). Videos longer than this will be chunked.
     CHUNK_OVERLAP_FRAMES = 5
-    MAX_PROCESSING_RESOLUTION = None
+    MAX_PROCESSING_RESOLUTION = None  # e.g., 1920 for downscaling if needed
     
-    # Quality settings
+    # Quality settings (for writes we control)
     VIDEO_CODEC = 'libx264'
     VIDEO_BITRATE = '8000k'
     AUDIO_CODEC = 'aac'
@@ -91,7 +105,7 @@ class ProcessingConfig:
     CLEAR_CACHE_AFTER_STAGE = True
     FORCE_GARBAGE_COLLECTION = True
     MEMORY_USAGE_THRESHOLD = 70
-    MEMORY_CHECK_INTERVAL = 30
+    MEMORY_CHECK_INTERVAL = 30  # every N frames during compositing
     
     # Debug settings
     SAVE_DEBUG_FILES = True
@@ -102,8 +116,9 @@ class ProcessingConfig:
     AUDIO_QUALITY = 'high'
     USE_FFMPEG_FOR_AUDIO = True  # Prefer ffmpeg over MoviePy
 
+
 # ==============================================================================
-# MEMORY MANAGER (keeping your original)
+# MEMORY MANAGER
 # ==============================================================================
 
 class MemoryManager:
@@ -113,54 +128,64 @@ def __init__(self, config: ProcessingConfig):
         self.config = config
         self.initial_memory = self.get_current_memory()
         self.peak_memory = self.initial_memory
-        self.stage_memories = {}
+        self.stage_memories: Dict[str, Dict[str, float]] = {}
     
     @staticmethod
-    def clear_gpu_cache():
+    def clear_gpu_cache() -> float:
         if TORCH_AVAILABLE and torch.cuda.is_available():
             before = torch.cuda.memory_allocated()
             torch.cuda.empty_cache()
-            torch.cuda.synchronize()
+            try:
+                torch.cuda.synchronize()
+            except Exception:
+                pass
             after = torch.cuda.memory_allocated()
             freed = (before - after) / 1024**3
             logger.info(f"GPU cache cleared: {freed:.2f}GB freed")
             return freed
-        return 0
+        return 0.0
     
     @staticmethod
-    def force_garbage_collection():
+    def force_garbage_collection() -> Tuple[int, float]:
         before = 0
         if TORCH_AVAILABLE and torch.cuda.is_available():
-            before = torch.cuda.memory_allocated()
+            try:
+                before = torch.cuda.memory_allocated()
+            except Exception:
+                before = 0
         
         collected = gc.collect()
         
         after = 0
         if TORCH_AVAILABLE and torch.cuda.is_available():
-            after = torch.cuda.memory_allocated()
+            try:
+                after = torch.cuda.memory_allocated()
+            except Exception:
+                after = 0
         
-        freed = (before - after) / 1024**3 if before > 0 else 0
+        freed = (before - after) / 1024**3 if before > 0 else 0.0
         logger.info(f"Garbage collection: {collected} objects freed, {freed:.2f}GB GPU memory recovered")
         return collected, freed
     
     def get_current_memory(self) -> Dict[str, float]:
         memory = {
-            'gpu_allocated': 0,
-            'gpu_reserved': 0,
-            'gpu_free': 0,
-            'gpu_total': 0,
-            'gpu_usage_percent': 0
+            'gpu_allocated': 0.0,
+            'gpu_reserved': 0.0,
+            'gpu_free': 0.0,
+            'gpu_total': 0.0,
+            'gpu_usage_percent': 0.0
         }
         
         if TORCH_AVAILABLE and torch.cuda.is_available():
-            memory['gpu_allocated'] = torch.cuda.memory_allocated() / 1024**3
-            memory['gpu_reserved'] = torch.cuda.memory_reserved() / 1024**3
-            
-            props = torch.cuda.get_device_properties(0)
-            memory['gpu_total'] = props.total_memory / 1024**3
-            memory['gpu_free'] = memory['gpu_total'] - memory['gpu_allocated']
-            memory['gpu_usage_percent'] = (memory['gpu_allocated'] / memory['gpu_total']) * 100
-        
+            try:
+                memory['gpu_allocated'] = torch.cuda.memory_allocated() / 1024**3
+                memory['gpu_reserved'] = torch.cuda.memory_reserved() / 1024**3
+                props = torch.cuda.get_device_properties(0)
+                memory['gpu_total'] = props.total_memory / 1024**3
+                memory['gpu_free'] = memory['gpu_total'] - memory['gpu_allocated']
+                memory['gpu_usage_percent'] = (memory['gpu_allocated'] / memory['gpu_total']) * 100 if memory['gpu_total'] > 0 else 0
+            except Exception:
+                pass
         return memory
     
     def get_gpu_memory_info(self) -> str:
@@ -176,13 +201,13 @@ def should_clear_memory(self) -> bool:
         return memory['gpu_usage_percent'] > self.config.MEMORY_USAGE_THRESHOLD
     
     def cleanup_stage(self, stage_name: str, force: bool = False):
-        logger.info(f"Cleaning up after {stage_name}...")
+        logger.info(f"Cleaning up after {stage_name}…")
         
         before = self.get_current_memory()
         
-        gpu_freed = 0
+        gpu_freed = 0.0
         gc_objects = 0
-        gc_memory = 0
+        gc_memory = 0.0
         
         if force or self.should_clear_memory():
             gpu_freed = self.clear_gpu_cache()
@@ -191,16 +216,15 @@ def cleanup_stage(self, stage_name: str, force: bool = False):
         after = self.get_current_memory()
         
         self.stage_memories[stage_name] = {
-            'before': before,
-            'after': after,
-            'freed': gpu_freed + gc_memory
+            'before_alloc': before['gpu_allocated'],
+            'after_alloc': after['gpu_allocated'],
+            'freed_total': gpu_freed + gc_memory
         }
         
         if after['gpu_allocated'] > self.peak_memory.get('gpu_allocated', 0):
             self.peak_memory = after
         
-        logger.info(f"Memory cleanup completed for {stage_name}: "
-                   f"{gpu_freed + gc_memory:.2f}GB total freed")
+        logger.info(f"Memory cleanup completed for {stage_name}: {gpu_freed + gc_memory:.2f}GB total freed")
     
     def cleanup_model(self, model_handler: Any, handler_name: str, deep_cleanup: bool = False):
         try:
@@ -214,40 +238,34 @@ def cleanup_model(self, model_handler: Any, handler_name: str, deep_cleanup: boo
             
             if deep_cleanup:
                 if hasattr(model_handler, 'model'):
-                    if model_handler.model is not None:
-                        del model_handler.model
-                        model_handler.model = None
+                    try:
+                        delattr(model_handler, 'model')
                         logger.info(f"{handler_name}: model deleted")
-                
+                    except Exception:
+                        pass
                 if hasattr(model_handler, 'predictor'):
-                    if model_handler.predictor is not None:
-                        del model_handler.predictor
-                        model_handler.predictor = None
+                    try:
+                        delattr(model_handler, 'predictor')
                         logger.info(f"{handler_name}: predictor deleted")
-                    
+                    except Exception:
+                        pass
         except Exception as e:
             logger.warning(f"Failed to cleanup {handler_name}: {e}")
     
     def get_memory_report(self) -> str:
-        report = ["="*60]
-        report.append("MEMORY USAGE REPORT")
-        report.append("="*60)
-        
+        report = ["="*60, "MEMORY USAGE REPORT", "="*60]
         current = self.get_current_memory()
         report.append(f"Current GPU Memory: {current['gpu_allocated']:.2f}GB / {current['gpu_total']:.2f}GB ({current['gpu_usage_percent']:.0f}%)")
-        report.append(f"Peak GPU Memory: {self.peak_memory.get('gpu_allocated', 0):.2f}GB")
-        
+        peak = self.peak_memory.get('gpu_allocated', 0.0)
+        report.append(f"Peak GPU Memory: {peak:.2f}GB")
         if self.stage_memories:
             report.append("\nStage Memory Usage:")
             for stage, mem in self.stage_memories.items():
-                report.append(f"  {stage}:")
-                report.append(f"    Before: {mem['before']['gpu_allocated']:.2f}GB")
-                report.append(f"    After:  {mem['after']['gpu_allocated']:.2f}GB")
-                report.append(f"    Freed:  {mem['freed']:.2f}GB")
-        
+                report.append(f"  {stage}: before={mem['before_alloc']:.2f}GB after={mem['after_alloc']:.2f}GB freed={mem['freed_total']:.2f}GB")
         report.append("="*60)
         return "\n".join(report)
 
+
 # ==============================================================================
 # QUALITY MANAGER
 # ==============================================================================
@@ -280,6 +298,7 @@ class QualityManager:
     def get_profile(cls, quality: str = 'medium') -> Dict[str, Any]:
         return cls.PROFILES.get(quality, cls.PROFILES['medium'])
 
+
 # ==============================================================================
 # ALPHA REFINEMENT
 # ==============================================================================
@@ -299,18 +318,18 @@ def refine_alpha_matte(alpha: np.ndarray, config: ProcessingConfig) -> np.ndarra
         
         # Temporal smoothing
         if config.ALPHA_SMOOTHING > 0:
-            kernel_size = config.ALPHA_SMOOTHING
-            if kernel_size % 2 == 0:
-                kernel_size += 1
-            alpha = cv2.medianBlur((alpha * 255).astype(np.uint8), kernel_size).astype(np.float32) / 255.0
+            k = config.ALPHA_SMOOTHING
+            if k % 2 == 0:
+                k += 1
+            alpha = cv2.medianBlur((alpha * 255).astype(np.uint8), k).astype(np.float32) / 255.0
         
         # Edge feathering
         if config.ALPHA_EDGE_BLUR > 0:
-            kernel_size = config.ALPHA_EDGE_BLUR * 2 + 1
-            alpha = cv2.GaussianBlur(alpha, (kernel_size, kernel_size), 0)
+            ks = config.ALPHA_EDGE_BLUR * 2 + 1
+            alpha = cv2.GaussianBlur(alpha, (ks, ks), 0)
         
         # Contrast adjustment
-        if config.ALPHA_CONTRAST != 1.0:
+        if abs(config.ALPHA_CONTRAST - 1.0) > 1e-6:
             alpha = np.power(alpha, config.ALPHA_CONTRAST)
         
         return np.clip(alpha, 0, 1)
@@ -326,21 +345,22 @@ def defringe_edges(composite: np.ndarray, alpha: np.ndarray, radius: int = 1) ->
         kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (radius*2+1, radius*2+1))
         dilated = cv2.dilate(alpha, kernel)
         eroded = cv2.erode(alpha, kernel)
-        edge_mask = dilated - eroded
+        edge_mask = np.clip(dilated - eroded, 0, 1)
         
         blurred = cv2.GaussianBlur(composite, (3, 3), 0)
         
         edge_mask_3d = np.stack([edge_mask] * 3, axis=2)
         result = composite * (1 - edge_mask_3d) + blurred * edge_mask_3d
         
-        return result.astype(np.uint8)
+        return np.clip(result, 0, 255).astype(np.uint8)
+
 
 # ==============================================================================
 # CHUNKED VIDEO PROCESSOR
 # ==============================================================================
 
 class ChunkedVideoProcessor:
-    """Handles splitting and reassembling videos."""
+    """Handles splitting and reassembling videos (video-only; audio added later)."""
     
     def __init__(self, temp_dir: str, config: ProcessingConfig):
         self.temp_dir = temp_dir
@@ -351,7 +371,7 @@ def __init__(self, temp_dir: str, config: ProcessingConfig):
     def should_chunk_video(self, video_path: str) -> bool:
         try:
             with VideoFileClip(video_path) as clip:
-                duration = clip.duration
+                duration = clip.duration or 0
                 should_chunk = duration > self.config.MAX_CHUNK_DURATION
                 logger.info(f"Video duration: {duration:.1f}s, chunking: {should_chunk}")
                 return should_chunk
@@ -362,14 +382,14 @@ def should_chunk_video(self, video_path: str) -> bool:
     def split_video_into_chunks(self, video_path: str) -> List[str]:
         try:
             with VideoFileClip(video_path) as clip:
-                duration = clip.duration
-                chunk_paths = []
+                duration = clip.duration or 0.0
+                chunk_paths: List[str] = []
                 
                 chunk_duration = self.config.MAX_CHUNK_DURATION
-                current_time = 0
+                current_time = 0.0
                 chunk_index = 0
                 
-                while current_time < duration:
+                while current_time < duration - 1e-6:
                     end_time = min(current_time + chunk_duration, duration)
                     
                     remaining_after = duration - end_time
@@ -391,11 +411,11 @@ def split_video_into_chunks(self, video_path: str) -> List[str]:
                     chunk_clip.close()
                     
                     chunk_paths.append(chunk_path)
-                    logger.info(f"Created chunk {chunk_index}: {current_time:.1f}s - {end_time:.1f}s")
+                    logger.info(f"Created chunk {chunk_index}: {current_time:.2f}s - {end_time:.2f}s")
                     
                     if end_time < duration:
-                        overlap_seconds = self.config.CHUNK_OVERLAP_FRAMES / clip.fps
-                        current_time = end_time - overlap_seconds
+                        overlap_seconds = self.config.CHUNK_OVERLAP_FRAMES / (clip.fps or 30.0)
+                        current_time = max(0.0, end_time - overlap_seconds)
                     else:
                         current_time = duration
                     
@@ -416,15 +436,10 @@ def reassemble_chunks(self, processed_chunk_paths: List[str], output_path: str)
                 logger.info("Single chunk, copied directly")
                 return output_path
             
-            chunk_clips = []
-            for chunk_path in processed_chunk_paths:
-                clip = VideoFileClip(chunk_path)
-                chunk_clips.append(clip)
-            
-            logger.info(f"Concatenating {len(chunk_clips)} chunks...")
+            # Compose video-only
+            chunk_clips = [VideoFileClip(p) for p in processed_chunk_paths]
+            logger.info(f"Concatenating {len(chunk_clips)} chunks…")
             final_clip = concatenate_videoclips(chunk_clips, method="compose")
-            
-            # Write WITHOUT audio (we'll add it separately)
             final_clip.write_videofile(
                 output_path,
                 codec=self.config.VIDEO_CODEC,
@@ -432,10 +447,9 @@ def reassemble_chunks(self, processed_chunk_paths: List[str], output_path: str)
                 verbose=False,
                 logger=None
             )
-            
             final_clip.close()
-            for clip in chunk_clips:
-                clip.close()
+            for c in chunk_clips:
+                c.close()
             
             logger.info(f"Successfully reassembled video: {output_path}")
             return output_path
@@ -444,6 +458,7 @@ def reassemble_chunks(self, processed_chunk_paths: List[str], output_path: str)
             logger.error(f"Chunk reassembly failed: {e}")
             raise
 
+
 # ==============================================================================
 # MAIN TWOSTAGEPROCESSOR CLASS
 # ==============================================================================
@@ -465,8 +480,12 @@ def __init__(self, sam2_handler, matanyone_handler, temp_dir: Optional[str] = No
         
         # Initialize AudioProcessor for proper audio handling
         if AudioProcessor and self.config.USE_FFMPEG_FOR_AUDIO:
-            self.audio_processor = AudioProcessor(temp_dir=self.temp_dir)
-            logger.info("AudioProcessor initialized for ffmpeg-based audio handling")
+            try:
+                self.audio_processor = AudioProcessor(temp_dir=self.temp_dir)
+                logger.info("AudioProcessor initialized for ffmpeg-based audio handling")
+            except Exception as e:
+                self.audio_processor = None
+                logger.warning(f"AudioProcessor failed to initialize ({e}). Using MoviePy fallback for audio.")
         else:
             self.audio_processor = None
             logger.warning("AudioProcessor not available - using MoviePy fallback for audio")
@@ -481,19 +500,26 @@ def process_video(self,
                      background_path: str,
                      output_path: str,
                      quality: str = 'medium',
-                     trim_seconds: Optional[int] = None,
+                     trim_seconds: Optional[float] = None,
                      callback: Optional[callable] = None,
-                     **kwargs) -> Tuple[str, str]:
+                     **kwargs) -> Tuple[Optional[str], str]:
         """
         Main processing pipeline with proper audio handling.
+
+        Returns:
+            (final_output_path, status_message)
         """
         try:
             logger.info(f"🎬 Enhanced Two-Stage Alpha Pipeline: {video_path}")
             logger.info(f"🎯 Background: {background_path}")
             logger.info(f"📁 Temp: {self.temp_dir}")
             logger.info(f"🎤 Audio: {'Enabled (ffmpeg)' if self.audio_processor else 'Enabled (MoviePy)'}")
+            if trim_seconds is not None:
+                logger.info(f"✂️  Trim requested: first {trim_seconds}s will be processed")
+            else:
+                logger.info("⏱️  No trimming requested; full duration will be processed")
             
-            # Extract original audio using AudioProcessor if available
+            # Extract original audio (if any)
             original_audio_path = None
             if self.config.PRESERVE_ORIGINAL_AUDIO:
                 original_audio_path = self._extract_audio(video_path)
@@ -502,22 +528,17 @@ def process_video(self,
                 else:
                     logger.info("ℹ️ No audio found in source video")
             
-            # Check video duration
-            with VideoFileClip(video_path) as clip:
-                duration = clip.duration
-                logger.info(f"📏 Input video duration: {duration:.1f} seconds")
-            
             # Determine if chunking is needed
             needs_chunking = self.chunked_processor.should_chunk_video(video_path)
             
             if needs_chunking:
-                logger.info("Using chunked processing for long video...")
+                logger.info("Using chunked processing for long video…")
                 result = self._process_chunked_video(
                     video_path, background_path, output_path, 
                     quality, original_audio_path, callback
                 )
             else:
-                logger.info("Processing full video in single pass...")
+                logger.info("Processing full video in single pass…")
                 result = self._process_single_video(
                     video_path, background_path, output_path, 
                     quality, trim_seconds, original_audio_path, callback
@@ -525,10 +546,13 @@ def process_video(self,
             
             # Verify final output
             if result[0] and os.path.exists(result[0]):
-                with VideoFileClip(result[0]) as clip:
-                    final_duration = clip.duration
-                    has_audio = clip.audio is not None
-                    logger.info(f"✅ Final output: {final_duration:.1f}s, Audio: {has_audio}")
+                try:
+                    with VideoFileClip(result[0]) as clip:
+                        final_duration = clip.duration or 0.0
+                        has_audio = clip.audio is not None
+                        logger.info(f"✅ Final output: {final_duration:.1f}s, Audio: {has_audio}")
+                except Exception:
+                    pass
             
             logger.info(self.memory_manager.get_memory_report())
             return result
@@ -539,12 +563,15 @@ def process_video(self,
             logger.error(f"Traceback: {traceback.format_exc()}")
             return None, error_msg
     
+    # ------------------------------------------------------------------
+    # Audio helpers
+    # ------------------------------------------------------------------
     def _extract_audio(self, video_path: str) -> Optional[str]:
         """Extract audio using AudioProcessor (ffmpeg) or MoviePy fallback."""
         try:
             if self.audio_processor:
                 # Use AudioProcessor with ffmpeg
-                logger.info("Extracting audio using ffmpeg...")
+                logger.info("Extracting audio using ffmpeg…")
                 audio_path = self.audio_processor.extract_audio(
                     video_path=video_path,
                     output_path=os.path.join(self.temp_dir, "original_audio.aac"),
@@ -554,7 +581,7 @@ def _extract_audio(self, video_path: str) -> Optional[str]:
                 return audio_path
             else:
                 # Fallback to MoviePy
-                logger.info("Extracting audio using MoviePy...")
+                logger.info("Extracting audio using MoviePy…")
                 audio_path = os.path.join(self.temp_dir, "original_audio.m4a")
                 with VideoFileClip(video_path) as clip:
                     if clip.audio is not None:
@@ -572,21 +599,23 @@ def _extract_audio(self, video_path: str) -> Optional[str]:
             logger.warning(f"Could not extract audio: {e}")
             return None
     
-    def _add_audio_to_video(self, video_path: str, audio_path: str, output_path: str) -> str:
+    def _add_audio_to_video(self, video_path: str, audio_path: Optional[str], output_path: str) -> str:
         """Add audio to video using AudioProcessor (ffmpeg) or MoviePy fallback."""
         try:
             if self.audio_processor and audio_path:
-                # Use AudioProcessor with ffmpeg for lossless muxing
-                logger.info("Adding audio using ffmpeg (lossless if possible)...")
+                # Using ffmpeg via AudioProcessor.
+                # Note: AudioProcessor.add_audio_to_video accepts an input for "original_video".
+                # It's OK to pass an audio-only file here; the function maps [1:a:0] as audio.
+                logger.info("Adding audio using ffmpeg (lossless where possible)…")
                 return self.audio_processor.add_audio_to_video(
-                    original_video=audio_path,  # Source of audio
-                    processed_video=video_path,  # Video without audio
+                    original_video=audio_path,       # Source of audio (audio-only or original video)
+                    processed_video=video_path,      # Video without audio
                     output_path=output_path,
                     audio_quality=self.config.AUDIO_QUALITY
                 )
             elif audio_path:
                 # Fallback to MoviePy
-                logger.info("Adding audio using MoviePy...")
+                logger.info("Adding audio using MoviePy…")
                 with VideoFileClip(video_path) as video:
                     with AudioFileClip(audio_path) as audio:
                         final = video.set_audio(audio)
@@ -602,34 +631,42 @@ def _add_audio_to_video(self, video_path: str, audio_path: str, output_path: str
                         )
                 return output_path
             else:
-                # No audio to add
+                # No audio to add; just copy
                 shutil.copy2(video_path, output_path)
                 return output_path
                 
         except Exception as e:
             logger.error(f"Failed to add audio: {e}")
             # Return video without audio rather than failing completely
-            shutil.copy2(video_path, output_path)
+            try:
+                shutil.copy2(video_path, output_path)
+            except Exception:
+                pass
             return output_path
     
+    # ------------------------------------------------------------------
+    # Chunked processing
+    # ------------------------------------------------------------------
     def _process_chunked_video(self, video_path: str, background_path: str,
                               output_path: str, quality: str,
                               original_audio_path: Optional[str],
-                              callback: Optional[callable]) -> Tuple[str, str]:
+                              callback: Optional[callable]) -> Tuple[Optional[str], str]:
         """Process long video in chunks."""
         try:
             if callback:
-                callback("Splitting video into chunks...", 5)
+                callback("Splitting video into chunks…", 5)
             
             chunk_paths = self.chunked_processor.split_video_into_chunks(video_path)
+            if not chunk_paths:
+                raise RuntimeError("No chunks were created")
             
-            processed_chunks = []
+            processed_chunks: List[str] = []
             total_chunks = len(chunk_paths)
             
             for i, chunk_path in enumerate(chunk_paths):
                 if callback:
-                    progress = 10 + (i * 70 // total_chunks)
-                    callback(f"Processing chunk {i+1}/{total_chunks}...", progress)
+                    progress = 10 + (i * 70 // max(1, total_chunks))
+                    callback(f"Processing chunk {i+1}/{total_chunks}…", progress)
                 
                 logger.info(f"Processing chunk {i+1}/{total_chunks}")
                 
@@ -649,19 +686,19 @@ def _process_chunked_video(self, video_path: str, background_path: str,
                 self.memory_manager.cleanup_stage(f"Chunk_{i+1}", force=True)
             
             if callback:
-                callback("Reassembling video...", 85)
+                callback("Reassembling video…", 85)
             
-            # Reassemble chunks
+            # Reassemble chunks (video only)
             temp_output = os.path.join(self.temp_dir, "reassembled_no_audio.mp4")
-            final_path = self.chunked_processor.reassemble_chunks(processed_chunks, temp_output)
+            final_video_no_audio = self.chunked_processor.reassemble_chunks(processed_chunks, temp_output)
             
             # Add original audio back
             if original_audio_path:
                 if callback:
-                    callback("Adding audio track...", 95)
-                final_path = self._add_audio_to_video(temp_output, original_audio_path, output_path)
+                    callback("Adding audio track…", 95)
+                final_path = self._add_audio_to_video(final_video_no_audio, original_audio_path, output_path)
             else:
-                shutil.move(temp_output, output_path)
+                shutil.move(final_video_no_audio, output_path)
                 final_path = output_path
             
             if callback:
@@ -671,43 +708,59 @@ def _process_chunked_video(self, video_path: str, background_path: str,
             
         except Exception as e:
             logger.error(f"Chunked processing failed: {e}")
-            raise
+            return None, f"Chunked processing failed: {e}"
     
+    # ------------------------------------------------------------------
+    # Single-pass processing
+    # ------------------------------------------------------------------
     def _process_single_video(self, video_path: str, background_path: str,
                              output_path: str, quality: str,
-                             trim_seconds: Optional[int],
+                             trim_seconds: Optional[float],
                              original_audio_path: Optional[str],
-                             callback: Optional[callable]) -> Tuple[str, str]:
-        """Process single video/chunk."""
+                             callback: Optional[callable]) -> Tuple[Optional[str], str]:
+        """Process single video/chunk (video-only here; audio added at the end)."""
         try:
+            # Optional trim (explicit only)
+            input_for_masks = video_path
+            if trim_seconds is not None and trim_seconds > 0:
+                # If you pass trim_seconds, we trim for *this* run.
+                trimmed_path = os.path.join(self.temp_dir, "trimmed_input.mp4")
+                self._trim_video(video_path, trimmed_path, trim_seconds)
+                input_for_masks = trimmed_path
+                logger.info(f"Trim applied: {trim_seconds}s (only because trim_seconds was provided)")
+            else:
+                logger.info("No trimming applied in this run")
+            
             # Stage 1: Reference mask
             if callback:
-                callback("Creating reference mask...", 10)
-            reference_mask_path = self._stage1_create_reference_mask(video_path)
+                callback("Creating reference mask…", 10)
+            reference_mask_path = self._stage1_create_reference_mask(input_for_masks)
             self.memory_manager.cleanup_stage("SAM2", force=True)
             
             # Stage 2: MatAnyone
             if callback:
-                callback("Processing with MatAnyone...", 30)
+                callback("Processing with MatAnyone…", 30)
             alpha_video_path = self._stage2_matanyone_processing(
-                video_path, reference_mask_path, trim_seconds
+                video_path=input_for_masks, reference_mask_path=reference_mask_path
             )
             self.memory_manager.cleanup_stage("MatAnyone", force=True)
             
-            # Stage 3: Alpha compositing
+            # Stage 3: Alpha compositing (video-only)
             if callback:
-                callback("Compositing with background...", 70)
-            
-            # Composite WITHOUT audio first
+                callback("Compositing with background…", 70)
             temp_output = os.path.join(self.temp_dir, "composite_no_audio.mp4")
             composite_path = self._stage3_alpha_composite(
-                video_path, alpha_video_path, background_path, temp_output, quality
+                original_video_path=input_for_masks,
+                alpha_video_path=alpha_video_path,
+                background_path=background_path,
+                output_path=temp_output,
+                quality=quality
             )
             
-            # Add audio if provided
+            # Add audio if we have it
             if original_audio_path:
                 if callback:
-                    callback("Adding audio track...", 90)
+                    callback("Adding audio track…", 90)
                 final_path = self._add_audio_to_video(composite_path, original_audio_path, output_path)
             else:
                 shutil.move(composite_path, output_path)
@@ -720,48 +773,47 @@ def _process_single_video(self, video_path: str, background_path: str,
                 
         except Exception as e:
             logger.error(f"Single video processing failed: {e}")
-            raise
+            return None, f"Single video processing failed: {e}"
     
-    # [Keeping all your original stage methods unchanged]
+    # ------------------------------------------------------------------
+    # Stage 1: Reference mask creation
+    # ------------------------------------------------------------------
     def _stage1_create_reference_mask(self, video_path: str) -> str:
-        """Stage 1: Create reference mask from first frames using SAM2."""
+        """Create reference mask from first frames using SAM2."""
+        cap = None
         try:
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
                 raise RuntimeError(f"Cannot open video: {video_path}")
             
-            frames = []
+            frames: List[np.ndarray] = []
             frame_count = 0
             
             while frame_count < self.config.REFERENCE_FRAMES:
                 ret, frame = cap.read()
                 if not ret:
                     break
-                
                 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 frames.append(frame_rgb)
                 frame_count += 1
             
-            cap.release()
-            
             if not frames:
                 raise RuntimeError("No frames extracted for reference mask")
             
             logger.info(f"Extracted {len(frames)} reference frames")
             
-            masks = []
+            masks: List[np.ndarray] = []
             for i, frame in enumerate(frames):
                 try:
-                    logger.info(f"Creating mask for frame {i+1}/{len(frames)}...")
+                    logger.info(f"Creating mask for frame {i+1}/{len(frames)}…")
                     mask = self._create_sam2_mask(frame)
-                    
                     if mask is not None:
+                        if mask.dtype == np.uint8:
+                            mask = mask.astype(np.float32) / 255.0
                         masks.append(mask)
-                        
                         if self.config.SAVE_DEBUG_FILES:
                             debug_path = os.path.join(self.temp_dir, f"debug_mask_{i}.png")
-                            cv2.imwrite(debug_path, (mask * 255).astype(np.uint8))
-                            
+                            cv2.imwrite(debug_path, np.clip(mask * 255, 0, 255).astype(np.uint8))
                 except Exception as e:
                     logger.warning(f"Failed to create mask for frame {i}: {e}")
                     continue
@@ -770,107 +822,125 @@ def _stage1_create_reference_mask(self, video_path: str) -> str:
                 raise RuntimeError("No valid reference masks created")
             
             combined_mask = self._combine_reference_masks(masks)
-            
             reference_mask_path = os.path.join(self.temp_dir, "reference_mask.png")
-            cv2.imwrite(reference_mask_path, (combined_mask * 255).astype(np.uint8))
-            
+            cv2.imwrite(reference_mask_path, np.clip(combined_mask * 255, 0, 255).astype(np.uint8))
             logger.info(f"✅ Reference mask created: {reference_mask_path}")
             return reference_mask_path
             
         except Exception as e:
             logger.error(f"Stage 1 failed: {e}")
             raise
+        finally:
+            try:
+                if cap is not None:
+                    cap.release()
+            except Exception:
+                pass
     
     def _create_sam2_mask(self, frame_rgb: np.ndarray) -> Optional[np.ndarray]:
-        """Create mask using SAM2Handler."""
+        """Create mask using SAM2Handler; try a few signatures."""
         try:
             methods_to_try = [
                 lambda: self.sam2_handler.create_mask(frame_rgb),
                 lambda: self.sam2_handler.create_mask(frame_rgb, None, None, None),
                 lambda: self.sam2_handler.predict(frame_rgb) if hasattr(self.sam2_handler, 'predict') else None
             ]
-            
             for method in methods_to_try:
                 try:
                     mask = method()
                     if mask is not None:
+                        if isinstance(mask, dict) and 'mask' in mask:
+                            mask = mask['mask']
                         if mask.dtype == np.uint8:
                             mask = mask.astype(np.float32) / 255.0
-                        return mask
-                except (TypeError, AttributeError) as e:
-                    logger.debug(f"Method failed: {e}")
+                        return np.clip(mask, 0.0, 1.0)
+                except (TypeError, AttributeError):
                     continue
-            
             logger.warning("All SAM2 methods failed")
             return None
-            
         except Exception as e:
             logger.error(f"SAM2 mask creation failed: {e}")
             return None
     
-    def _combine_reference_masks(self, masks: list) -> np.ndarray:
-        """Combine multiple masks with voting."""
+    def _combine_reference_masks(self, masks: List[np.ndarray]) -> np.ndarray:
+        """Combine multiple masks with simple voting/averaging."""
         if len(masks) == 1:
             return masks[0]
-        
         stacked = np.stack(masks, axis=0)
         avg_mask = np.mean(stacked, axis=0)
-        
         threshold = 0.3
-        result = np.where(avg_mask > threshold, avg_mask, 0.0)
-        
+        result = np.where(avg_mask > threshold, avg_mask, 0.0).astype(np.float32)
         logger.info(f"Combined {len(masks)} masks with threshold {threshold}")
         return result
     
-    def _stage2_matanyone_processing(self, video_path: str, 
-                                     reference_mask_path: str, 
-                                     trim_seconds: Optional[int]) -> str:
-        """Stage 2: Process with MatAnyone."""
+    # ------------------------------------------------------------------
+    # Stage 2: MatAnyone
+    # ------------------------------------------------------------------
+    def _stage2_matanyone_processing(self, video_path: str, reference_mask_path: str) -> str:
+        """Run MatAnyone to produce an alpha video (no trimming here)."""
         try:
-            input_video = video_path
-            if trim_seconds:
-                trimmed_path = os.path.join(self.temp_dir, "trimmed_input.mp4")
-                self._trim_video(video_path, trimmed_path, trim_seconds)
-                input_video = trimmed_path
-            
             matanyone_dir = os.path.join(self.temp_dir, "matanyone_out")
             os.makedirs(matanyone_dir, exist_ok=True)
             
             profile = QualityManager.get_profile('high')
             max_size = profile.get('max_size', 1920)
             
+            # Known kwargs signature (from your logs):
+            # ['input_path','mask_path','output_path','n_warmup','r_erode','r_dilate','suffix','save_image','max_size']
+            tried = False
             try:
                 self.matanyone_handler.process_video(
-                    input_path=input_video,
+                    input_path=video_path,
                     mask_path=reference_mask_path,
                     output_path=matanyone_dir,
-                    n_warmup=0,  # CRITICAL: Don't limit to 5 seconds!
+                    n_warmup=0,     # ✅ do NOT limit to 5 seconds
                     r_erode=0,
                     r_dilate=15,
                     suffix='pha',
                     save_image=False,
                     max_size=max_size
                 )
+                tried = True
             except Exception as e:
-                logger.warning(f"Explicit kwargs failed: {e}, trying positional")
-                self.matanyone_handler.process_video(
-                    input_video, reference_mask_path, matanyone_dir
-                )
+                logger.warning(f"Explicit kwargs failed: {e}")
+            
+            if not tried:
+                # Positional fallback (older builds sometimes expose a simpler signature)
+                try:
+                    self.matanyone_handler.process_video(video_path, reference_mask_path, matanyone_dir)
+                    tried = True
+                except Exception as e:
+                    logger.error(f"Positional call to MatAnyone failed: {e}")
+            
+            if not tried:
+                raise RuntimeError("MatAnyone.process_video could not be invoked with known signatures")
             
-            # Find output
+            # Find alpha video in output
             alpha_output_path = None
             for filename in os.listdir(matanyone_dir):
-                if filename.endswith('_pha.mp4') or filename == 'pha.mp4':
+                low = filename.lower()
+                if low.endswith('_pha.mp4') or low == 'pha.mp4' or ('alpha' in low and low.endswith('.mp4')):
                     alpha_output_path = os.path.join(matanyone_dir, filename)
                     break
             
             if not alpha_output_path or not os.path.exists(alpha_output_path):
-                raise RuntimeError("MatAnyone did not create alpha video")
+                # Sometimes MatAnyone may place files inside a nested folder; check subdirs
+                for root, _, files in os.walk(matanyone_dir):
+                    for f in files:
+                        low = f.lower()
+                        if low.endswith('_pha.mp4') or low == 'pha.mp4' or ('alpha' in low and low.endswith('.mp4')):
+                            alpha_output_path = os.path.join(root, f)
+                            break
+                    if alpha_output_path:
+                        break
             
-            # Verify
+            if not alpha_output_path or not os.path.exists(alpha_output_path):
+                raise RuntimeError("MatAnyone did not create an alpha video")
+            
+            # Verify basic properties
             cap = cv2.VideoCapture(alpha_output_path)
-            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-            fps = cap.get(cv2.CAP_PROP_FPS)
+            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+            fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
             duration = frame_count / fps if fps > 0 else 0
             cap.release()
             
@@ -884,12 +954,18 @@ def _stage2_matanyone_processing(self, video_path: str,
             logger.error(f"Stage 2 failed: {e}")
             raise
     
+    # ------------------------------------------------------------------
+    # Stage 3: Alpha compositing (video-only)
+    # ------------------------------------------------------------------
     def _stage3_alpha_composite(self, original_video_path: str, 
                                 alpha_video_path: str,
                                 background_path: str, 
                                 output_path: str, 
                                 quality: str) -> str:
-        """Stage 3: Alpha compositing WITHOUT audio (audio added separately)."""
+        """Composite original RGB over background using alpha video (no audio here)."""
+        original_cap = None
+        alpha_cap = None
+        bg_cap = None
         try:
             original_cap = cv2.VideoCapture(original_video_path)
             alpha_cap = cv2.VideoCapture(alpha_video_path)
@@ -897,25 +973,30 @@ def _stage3_alpha_composite(self, original_video_path: str,
             if not original_cap.isOpened() or not alpha_cap.isOpened():
                 raise RuntimeError("Cannot open videos for compositing")
             
-            fps = original_cap.get(cv2.CAP_PROP_FPS)
-            width = int(original_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-            height = int(original_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            total_frames = int(original_cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            fps = float(original_cap.get(cv2.CAP_PROP_FPS) or 30.0)
+            width = int(original_cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
+            height = int(original_cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
+            total_frames = int(original_cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
             
-            # Load background
-            if background_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
+            # Load background (image or looping video)
+            bg_is_video = False
+            bg_image = None
+            if background_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.webm')):
                 bg_cap = cv2.VideoCapture(background_path)
+                if not bg_cap.isOpened():
+                    raise RuntimeError(f"Cannot open background video: {background_path}")
                 bg_is_video = True
             else:
                 bg_image = cv2.imread(background_path)
                 if bg_image is None:
-                    raise RuntimeError(f"Cannot load background: {background_path}")
+                    raise RuntimeError(f"Cannot load background image: {background_path}")
                 bg_image = cv2.resize(bg_image, (width, height))
-                bg_is_video = False
             
             # Setup video writer (NO AUDIO)
             fourcc = cv2.VideoWriter_fourcc(*'mp4v')
             out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+            if not out.isOpened():
+                raise RuntimeError("Failed to open VideoWriter for composite")
             
             frame_count = 0
             while True:
@@ -925,65 +1006,79 @@ def _stage3_alpha_composite(self, original_video_path: str,
                 if not ret_orig or not ret_alpha:
                     break
                 
-                # Get background frame
+                # Background frame
                 if bg_is_video:
                     ret_bg, bg_frame = bg_cap.read()
                     if not ret_bg:
-                        bg_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+                        bg_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)  # loop
                         ret_bg, bg_frame = bg_cap.read()
                     bg_frame = cv2.resize(bg_frame, (width, height))
                 else:
                     bg_frame = bg_image.copy()
                 
-                # Process alpha
-                if len(alpha_frame.shape) == 3:
+                # Alpha extraction (single channel 0..1)
+                if alpha_frame.ndim == 3:
                     alpha_mask = cv2.cvtColor(alpha_frame, cv2.COLOR_BGR2GRAY)
                 else:
                     alpha_mask = alpha_frame
-                
                 if alpha_mask.shape[:2] != orig_frame.shape[:2]:
-                    alpha_mask = cv2.resize(alpha_mask, (width, height))
+                    alpha_mask = cv2.resize(alpha_mask, (width, height), interpolation=cv2.INTER_LINEAR)
                 
-                # Refine alpha
+                # Refine & normalize
                 alpha_refined = self.alpha_refiner.refine_alpha_matte(alpha_mask, self.config)
                 
-                # Composite
-                alpha_3d = np.stack([alpha_refined] * 3, axis=2)
-                composite = alpha_3d * orig_frame + (1 - alpha_3d) * bg_frame
+                # Composite (use float math, then cast)
+                a3 = np.dstack([alpha_refined] * 3).astype(np.float32)
+                fg = orig_frame.astype(np.float32)
+                bg = bg_frame.astype(np.float32)
+                mixed = a3 * fg + (1.0 - a3) * bg
                 
-                # Optional defringing
+                # Optional defringing (light blur at edges)
                 if self.config.ENABLE_DEFRINGING:
-                    composite = self.alpha_refiner.defringe_edges(
-                        composite.astype(np.uint8),
+                    mixed = self.alpha_refiner.defringe_edges(
+                        mixed.astype(np.uint8),
                         alpha_refined,
                         self.config.DEFRINGE_RADIUS
                     )
                 
-                out.write(composite.astype(np.uint8))
+                out.write(mixed.astype(np.uint8))
                 frame_count += 1
                 
-                if frame_count % self.config.MEMORY_CHECK_INTERVAL == 0:
+                if frame_count % max(1, self.config.MEMORY_CHECK_INTERVAL) == 0:
                     logger.info(f"Compositing: {frame_count}/{total_frames}")
             
-            # Cleanup
-            original_cap.release()
-            alpha_cap.release()
             out.release()
-            if bg_is_video:
-                bg_cap.release()
-            
-            logger.info(f"✅ Compositing completed: {frame_count} frames")
+            logger.info(f"✅ Compositing completed: {frame_count} frames → {output_path}")
             return output_path
             
         except Exception as e:
             logger.error(f"Stage 3 compositing failed: {e}")
             raise
+        finally:
+            try:
+                if original_cap is not None:
+                    original_cap.release()
+            except Exception:
+                pass
+            try:
+                if alpha_cap is not None:
+                    alpha_cap.release()
+            except Exception:
+                pass
+            try:
+                if bg_cap is not None:
+                    bg_cap.release()
+            except Exception:
+                pass
     
-    def _trim_video(self, input_path: str, output_path: str, seconds: int):
-        """Trim video to specified duration."""
+    # ------------------------------------------------------------------
+    # Small helpers
+    # ------------------------------------------------------------------
+    def _trim_video(self, input_path: str, output_path: str, seconds: float):
+        """Trim video to specified duration (video-only)."""
         try:
             with VideoFileClip(input_path) as clip:
-                trimmed = clip.subclip(0, min(seconds, clip.duration))
+                trimmed = clip.subclip(0, min(seconds, clip.duration or seconds))
                 trimmed.write_videofile(
                     output_path,
                     codec=self.config.VIDEO_CODEC,
@@ -991,7 +1086,7 @@ def _trim_video(self, input_path: str, output_path: str, seconds: int):
                     verbose=False,
                     logger=None
                 )
-            logger.info(f"Video trimmed to {seconds}s")
+            logger.info(f"Video trimmed to {seconds}s → {output_path}")
         except Exception as e:
             logger.error(f"Video trimming failed: {e}")
             raise
@@ -1012,7 +1107,7 @@ def cleanup(self):
     
     def get_processing_stats(self) -> Dict[str, Any]:
         """Get processing statistics."""
-        stats = {
+        stats: Dict[str, Any] = {
             'temp_dir': self.temp_dir,
             'memory_info': self.memory_manager.get_gpu_memory_info(),
             'audio_processor': 'ffmpeg' if self.audio_processor else 'MoviePy',
@@ -1022,26 +1117,30 @@ def get_processing_stats(self) -> Dict[str, Any]:
                 'chunk_duration': self.config.MAX_CHUNK_DURATION
             }
         }
-        
-        if self.audio_processor:
-            stats['audio_stats'] = self.audio_processor.get_stats()
-        
+        try:
+            if self.audio_processor:
+                stats['audio_stats'] = self.audio_processor.get_stats()
+        except Exception:
+            pass
         return stats
 
+
 # ==============================================================================
 # EXPORT INTERFACE
 # ==============================================================================
 
-def create_two_stage_processor(sam2_handler, matanyone_handler, **kwargs):
+def create_two_stage_processor(sam2_handler, matanyone_handler, **kwargs) -> TwoStageProcessor:
     """Factory function to create TwoStageProcessor."""
     return TwoStageProcessor(sam2_handler, matanyone_handler, **kwargs)
 
+
 if __name__ == "__main__":
     logger.info("Enhanced TwoStageProcessor with Audio Support")
     logger.info("Features:")
     logger.info("  ✅ ffmpeg-based audio via AudioProcessor")
-    logger.info("  ✅ Full video duration processing (no 5-second limit)")
+    logger.info("  ✅ Full video duration processing (no 5-second limit unless you pass trim_seconds)")
     logger.info("  ✅ Direct alpha compositing (no green screen)")
     logger.info("  ✅ Memory-efficient chunked processing")
+    logger.info("  ✅ Robust MatAnyone integration with signature fallbacks")
     logger.info("")
     logger.info("Use create_two_stage_processor(sam2_handler, matanyone_handler) to initialize")