tir

ui_core_functionality.py

@@ -2,6 +2,7 @@
 """
 BackgroundFX Pro — Core Functionality
 All processing logic, utilities, background generators, and handlers
+Enhanced with ChatGPT's file safety and error handling improvements
 """
 
 import os
@@ -121,6 +122,65 @@ def startup_probe():
         logger.error(f"📊 Disk stats: {_disk_stats(APP_ROOT)}")
         raise RuntimeError(f"Startup probe failed - system not ready: {e}") from e
 
+# ===============================================================================
+# CHATGPT ENHANCED FILE SAFETY UTILITIES
+# ===============================================================================
+
+def new_tmp_path(suffix: str) -> Path:
+    """Generate safe temporary path within TMP_ROOT"""
+    return TMP_ROOT / f"{uuid.uuid4().hex}{suffix}"
+
+def atomic_write_bytes(dst: Path, data: bytes):
+    """Atomic file write to prevent corruption"""
+    tmp = new_tmp_path(dst.suffix + ".part")
+    try:
+        with open(tmp, "wb") as f:
+            f.write(data)
+        tmp.replace(dst)  # atomic on same filesystem
+        logger.debug(f"✅ Atomic write completed: {dst}")
+    except Exception as e:
+        if tmp.exists():
+            tmp.unlink(missing_ok=True)
+        raise e
+
+def safe_name(name: str, default="file") -> str:
+    """Sanitize filename to prevent traversal/unicode issues"""
+    import re
+    base = re.sub(r"[^A-Za-z0-9._-]+", "_", (name or default))
+    return base[:120] or default
+
+def place_uploaded(in_path: str, sub="uploads") -> Path:
+    """Safely handle uploaded files with sanitized names"""
+    target_dir = DATA_ROOT / sub
+    target_dir.mkdir(exist_ok=True, parents=True)
+    out = target_dir / safe_name(Path(in_path).name)
+    shutil.copy2(in_path, out)
+    logger.info(f"📁 Uploaded file placed: {out}")
+    return out
+
+def tmp_video_path(ext=".mp4") -> Path:
+    """Generate temporary video path"""
+    return new_tmp_path(ext)
+
+def tmp_image_path(ext=".png") -> Path:
+    """Generate temporary image path"""
+    return new_tmp_path(ext)
+
+def run_safely(fn: Callable, *args, **kwargs):
+    """Execute function with comprehensive error logging"""
+    try:
+        return fn(*args, **kwargs)
+    except Exception as e:
+        logger.error("PROCESSING FAILED\n%s", "".join(traceback.format_exc()))
+        logger.error("CWD=%s | DATA_ROOT=%s | TMP_ROOT=%s | %s",
+                      os.getcwd(), DATA_ROOT, TMP_ROOT, _disk_stats(APP_ROOT))
+        logger.error("Env: OMP_NUM_THREADS=%s | CUDA=%s | torch=%s | cu=%s",
+                      os.environ.get("OMP_NUM_THREADS"),
+                      os.environ.get("CUDA_VISIBLE_DEVICES", "default"),
+                      torch.__version__,
+                      torch.version.cuda)
+        raise
+
 # ===============================================================================
 # SYSTEM UTILITIES
 # ===============================================================================
@@ -330,6 +390,13 @@ def extract_frame(video_path: str, frame_number: int) -> Optional[np.ndarray]:
         logger.error(f"Failed to extract frame: {e}")
         return None
 
+def ffmpeg_safe_call(inp: Path, out: Path, extra=()):
+    """Safe FFmpeg call with proper path quoting"""
+    cmd = ["ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
+           "-i", str(inp), *extra, str(out)]
+    logger.info("FFMPEG: %s", " ".join(cmd))
+    subprocess.run(cmd, check=True, timeout=300)
+
 # ===============================================================================
 # PROGRESS TRACKING
 # ===============================================================================
@@ -379,7 +446,7 @@ def get_status(self) -> Dict[str, Any]:
 progress_tracker = ProgressTracker()
 
 # ===============================================================================
-# SAFE FILE OPERATIONS
+# SAFE FILE OPERATIONS (ENHANCED)
 # ===============================================================================
 
 def create_job_directory() -> Path:
@@ -437,7 +504,7 @@ def safe_download(url: str, filepath: Path, max_size: int = 500 * 1024 * 1024):
         raise
 
 # ===============================================================================
-# PIPELINE INTEGRATION
+# ENHANCED PIPELINE INTEGRATION
 # ===============================================================================
 
 def process_video_pipeline(
@@ -448,46 +515,77 @@ def process_video_pipeline(
     job_dir: Path,
     progress_callback: Optional[Callable] = None
 ) -> str:
-    """Process video using the two-stage pipeline with detailed monitoring"""
-    try:
+    """Process video using the two-stage pipeline with enhanced safety and monitoring"""
+    
+    # Wrap entire function with ChatGPT's safe execution wrapper
+    def _inner_process():
         logger.info("=" * 60)
-        logger.info("=== STARTING TWO-STAGE PIPELINE DEBUG SESSION ===")
+        logger.info("=== ENHANCED TWO-STAGE PIPELINE (WITH CHATGPT SAFETY) ===")
         logger.info("=" * 60)
         
-        # Pre-flight checks
+        # Pre-flight checks with enhanced validation
         logger.info(f"DEBUG: Video path: {video_path}")
         logger.info(f"DEBUG: Video exists: {Path(video_path).exists()}")
+        logger.info(f"DEBUG: Video file size: {Path(video_path).stat().st_size if Path(video_path).exists() else 'N/A'} bytes")
         logger.info(f"DEBUG: Job directory: {job_dir}")
+        logger.info(f"DEBUG: Job directory writable: {os.access(job_dir, os.W_OK)}")
         logger.info(f"DEBUG: Background image size: {background_image.size if background_image else 'None'}")
         logger.info(f"DEBUG: Background type: {background_type}")
+        logger.info(f"DEBUG: Disk space: {_disk_stats(APP_ROOT)}")
+        
+        # Safely handle uploaded video
+        if not Path(video_path).exists():
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+            
+        # Create safe video path within our controlled environment
+        safe_video_path = place_uploaded(video_path, "videos")
+        logger.info(f"DEBUG: Safe video path: {safe_video_path}")
         
-        # Lazy import to avoid startup issues
+        # Import with error context
         logger.info("DEBUG: Attempting to import two-stage pipeline...")
-        from two_stage_pipeline import process_two_stage as pipeline_process
-        logger.info("✓ Two-stage pipeline imported successfully")
+        try:
+            from two_stage_pipeline import process_two_stage as pipeline_process
+            logger.info("✓ Two-stage pipeline imported successfully")
+        except ImportError as e:
+            logger.error(f"Failed to import two-stage pipeline: {e}")
+            raise
         
-        progress_tracker.update("Initializing two-stage pipeline...")
+        progress_tracker.update("Initializing enhanced two-stage pipeline...")
         
-        # Enhanced progress callback with stage monitoring
-        current_stage = {"stage": "init"}
+        # Enhanced progress callback with stage monitoring and memory tracking
+        current_stage = {"stage": "init", "start_time": time.time()}
         
         def safe_progress_callback(step: str, progress: float = None):
             try:
-                # Stage detection
+                current_time = time.time()
+                elapsed = current_time - current_stage["start_time"]
+                
+                # Stage detection with timing
                 if "Stage 1" in step:
                     if current_stage["stage"] != "stage1":
                         current_stage["stage"] = "stage1"
+                        current_stage["start_time"] = current_time
                         logger.info("🔄 STAGE TRANSITION: Entering Stage 1 (SAM2)")
+                        logger.info(f"Memory before Stage 1: {_disk_stats(APP_ROOT)}")
+                        
                 elif "Stage 2" in step:
                     if current_stage["stage"] != "stage2":
+                        stage1_duration = current_time - current_stage["start_time"]
                         current_stage["stage"] = "stage2"
+                        current_stage["start_time"] = current_time
                         logger.info("🔄 STAGE TRANSITION: Entering Stage 2 (Composition)")
+                        logger.info(f"Stage 1 completed in {stage1_duration:.1f}s")
+                        logger.info(f"Memory after Stage 1: {_disk_stats(APP_ROOT)}")
+                        
                 elif "Done" in step:
                     if current_stage["stage"] != "complete":
+                        stage2_duration = current_time - current_stage["start_time"]
                         current_stage["stage"] = "complete"
                         logger.info("🔄 STAGE TRANSITION: Pipeline Complete")
+                        logger.info(f"Stage 2 completed in {stage2_duration:.1f}s")
+                        logger.info(f"Final memory: {_disk_stats(APP_ROOT)}")
                 
-                logger.info(f"PROGRESS [{current_stage['stage'].upper()}]: {step} ({progress})")
+                logger.info(f"PROGRESS [{current_stage['stage'].upper()}] ({elapsed:.1f}s): {step} ({progress})")
                 progress_tracker.update(step, progress)
                 
                 if progress_callback:
@@ -495,22 +593,25 @@ def safe_progress_callback(step: str, progress: float = None):
                         progress_callback(f"Progress: {progress:.1%} - {step}")
                     else:
                         progress_callback(step)
+                        
+                # Memory warning if Stage 1 takes too long
+                if current_stage["stage"] == "stage1" and elapsed > 15:
+                    logger.warning(f"⚠️ Stage 1 running for {elapsed:.1f}s - monitoring for memory issues")
+                    
             except Exception as e:
                 logger.error(f"Progress callback error: {e}")
         
-        # Validation checks
+        # Validation with enhanced error context
         if background_image is None:
             raise ValueError("Background image is required")
             
         logger.info("DEBUG: Pre-pipeline validation complete")
-        logger.info("DEBUG: About to call two-stage pipeline...")
-        
-        # Monitor job directory before processing
         logger.info(f"DEBUG: Job dir contents before: {list(job_dir.iterdir()) if job_dir.exists() else 'does not exist'}")
         
-        # Call two-stage pipeline with correct parameters
+        # Call two-stage pipeline with safe paths
+        logger.info("DEBUG: Calling two-stage pipeline with enhanced monitoring...")
         result_path = pipeline_process(
-            video_path=video_path,
+            video_path=str(safe_video_path),
             background_image=background_image,
             workdir=job_dir,
             progress=safe_progress_callback,
@@ -520,32 +621,46 @@ def safe_progress_callback(step: str, progress: float = None):
         logger.info(f"DEBUG: Pipeline returned: {result_path}")
         logger.info(f"DEBUG: Result path type: {type(result_path)}")
         
-        # Post-processing validation
+        # Post-processing validation with enhanced checks
         if result_path:
             result_file = Path(result_path)
             logger.info(f"DEBUG: Result file exists: {result_file.exists()}")
             if result_file.exists():
-                logger.info(f"DEBUG: Result file size: {result_file.stat().st_size} bytes")
+                file_size = result_file.stat().st_size
+                logger.info(f"DEBUG: Result file size: {file_size} bytes")
+                if file_size == 0:
+                    raise RuntimeError("Pipeline produced empty output file")
             logger.info(f"DEBUG: Job dir contents after: {list(job_dir.iterdir())}")
+            
+            # Verify the output is a valid video
+            try:
+                cap = cv2.VideoCapture(str(result_file))
+                if cap.isOpened():
+                    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+                    logger.info(f"DEBUG: Output video frame count: {frame_count}")
+                    cap.release()
+                else:
+                    logger.warning("⚠️ Output file may not be a valid video")
+            except Exception as e:
+                logger.warning(f"⚠️ Could not verify output video: {e}")
         
         if not result_path or not Path(result_path).exists():
             raise RuntimeError("Two-stage pipeline processing failed - no output produced")
             
         logger.info("=" * 60)
-        logger.info(f"✅ TWO-STAGE PIPELINE COMPLETED: {result_path}")
+        logger.info(f"✅ ENHANCED TWO-STAGE PIPELINE COMPLETED: {result_path}")
         logger.info("=" * 60)
         return result_path
-        
+    
+    # Execute with ChatGPT's comprehensive error wrapper
+    try:
+        return run_safely(_inner_process)
     except Exception as e:
-        logger.error("=" * 60)
-        logger.error(f"❌ DETAILED ERROR ANALYSIS:")
-        logger.error(f"   Error type: {type(e).__name__}")
-        logger.error(f"   Error message: {e}")
-        logger.error(f"   Current stage: {current_stage.get('stage', 'unknown')}")
-        logger.error(f"   Job directory exists: {job_dir.exists() if 'job_dir' in locals() else 'unknown'}")
-        if 'job_dir' in locals() and job_dir.exists():
-            logger.error(f"   Job directory contents: {list(job_dir.iterdir())}")
-        logger.error("   Full traceback:")
-        logger.error(f"{traceback.format_exc()}")
-        logger.error("=" * 60)
+        # Enhanced error cleanup
+        logger.error("🧹 Performing error cleanup...")
+        clear_gpu_memory()
+        
+        # Additional error context
+        logger.error(f"Job directory state: {list(job_dir.iterdir()) if job_dir.exists() else 'does not exist'}")
+        
         raise