Update rust_highlight/src/lib.rs

rust_highlight/src/lib.rs

@@ -10,16 +10,18 @@ use std::time::Instant;
 use std::path::Path;
 use std::f64::consts::PI;
 use rayon::prelude::*;
+use rayon::ThreadPoolBuilder;
 
 #[pyfunction]
 fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f64, clips_dir: String) -> PyResult<Option<String>> {
+    // Validate audio file exists
     if !Path::new(&audio_path).exists() {
         return Err(pyo3::exceptions::PyFileNotFoundError::new_err(format!("Audio not found: {}", audio_path)));
     }
 
     let skip_spaces = false;
     let fps: f64 = 30.0;
-    let animation_frames_per_char: usize = 1;  // Reduced from 2 for speed
+    let animation_frames_per_char: usize = 1;
     let width: i32 = 1280;
     let height: i32 = 720;
     let margin_x: i32 = 40;
@@ -53,14 +55,16 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
     let text_area_width = width - 2 * margin_x;
     let (wrapped_lines, line_styles) = wrap_text_cv(&text, font, default_font_scale, default_thickness, text_area_width, header_font_scale, header_thickness);
 
-    let full_text = wrapped_lines.join("\n");
+    let full_text = wrapped_lines.join("
+");
     if full_text.is_empty() {
         println!("No text to animate.");
         return Ok(None);
     }
 
     let visible_indices: Vec<usize> = if skip_spaces {
-        full_text.char_indices().filter(|&(_, ch)| ch != ' ' && ch != '\n' && ch != '\t').map(|(i, _)| i).collect()
+        full_text.char_indices().filter(|&(_, ch)| ch != ' ' && ch != '
+' && ch != '\t').map(|(i, _)| i).collect()
     } else {
         (0..full_text.len()).collect()
     };
@@ -72,10 +76,16 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
         return Ok(None);
     }
 
-    // Calculate durations
+    // Calculate durations with validation
     let animation_duration = duration / 4.0;
     let static_duration = duration - animation_duration;
     
+    if static_duration <= 0.0 {
+        return Err(pyo3::exceptions::PyValueError::new_err(
+            format!("Invalid static_duration: {:.3}s (audio duration {:.3}s too short)", static_duration, duration)
+        ));
+    }
+    
     println!("Animation duration: {:.3}s, Static duration: {:.3}s", animation_duration, static_duration);
 
     // Pre-calc line heights and y_positions
@@ -96,7 +106,7 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
 
     let t0 = Instant::now();
 
-    // STEP 1: Pre-render the final static frame (used for Phase 2)
+    // STEP 1: Pre-render the final static frame
     println!("Rendering static frame...");
     let static_frame = render_frame(
         &full_text,
@@ -129,17 +139,23 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
     opencv::imgcodecs::imwrite(&static_frame_path, &static_frame, &params)
         .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to save static frame: {}", e)))?;
 
-    // STEP 2: Render animation frames in parallel (HUGE speedup)
-    println!("Rendering animation frames in parallel...");
-    
-    // Collect frame data for parallel processing
+    // Verify static frame was created
+    if !Path::new(&static_frame_path).exists() {
+        return Err(pyo3::exceptions::PyFileNotFoundError::new_err(
+            format!("Static frame not created at: {}", static_frame_path)
+        ));
+    }
+
+    // STEP 2: Build frame specifications
+    println!("Building frame specifications...");
     let mut frame_specs: Vec<(String, i32, i32, usize)> = Vec::new();
     let mut prev_visible_sub = String::new();
 
     for &idx_in_full in visible_indices.iter() {
         let visible_sub = full_text[0..=idx_in_full].to_string();
         if visible_sub != prev_visible_sub {
-            let lines: Vec<&str> = visible_sub.split('\n').collect();
+            let lines: Vec<&str> = visible_sub.split('
+').collect();
             let last_line = lines.last().unwrap();
             let line_idx = lines.len() - 1;
             let is_header = line_styles[line_idx];
@@ -161,7 +177,8 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
 
     println!("Total animation frames to render: {}", frame_specs.len());
 
-    // STEP 3: Start FFmpeg process first
+    // STEP 3: Start FFmpeg process for animation
+    println!("Starting FFmpeg process for animation...");
     let mut child = Command::new("ffmpeg")
         .arg("-y")
         .arg("-f").arg("rawvideo")
@@ -176,74 +193,111 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
         .arg("-pix_fmt").arg("yuv420p")
         .arg(&animation_video_path)
         .stdin(Stdio::piped())
+        .stderr(Stdio::piped())
         .spawn()
         .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to spawn FFmpeg: {}", e)))?;
 
     let mut stdin = child.stdin.take().unwrap();
 
-    // Parallel rendering using rayon and streaming to FFmpeg
-    let animation_frames: Vec<Vec<u8>> = frame_specs
-        .par_iter()
-        .map(|(visible_sub, pen_x, pen_y, anim_step)| {
-            let anim_offset = (*anim_step as f64) / (animation_frames_per_char as f64);
-            let frame = render_frame(
-                visible_sub,
-                *pen_x,
-                *pen_y,
-                anim_offset,
-                width,
-                height,
-                &line_styles,
-                &y_positions,
-                margin_x,
-                font,
-                default_font_scale,
-                header_font_scale,
-                default_thickness,
-                header_thickness,
-                default_text_color,
-                header_text_color,
-                bg_color,
-                pen_color,
-                pen_tip_radius,
-                pen_length,
-                pen_thickness,
-                pen_base_angle,
-                pen_movement_amplitude,
-            ).unwrap();
-            frame.data_bytes().unwrap().to_vec()
-        })
-        .collect();
+    // STEP 4: Parallel rendering with limited thread pool
+    println!("Rendering animation frames in parallel...");
+    
+    // Create limited thread pool to avoid OpenCV contention
+    let pool = ThreadPoolBuilder::new()
+        .num_threads(4)
+        .build()
+        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to build thread pool: {}", e)))?;
+
+    let animation_frames: Vec<Vec<u8>> = pool.install(|| {
+        frame_specs
+            .par_iter()
+            .map(|(visible_sub, pen_x, pen_y, anim_step)| {
+                let anim_offset = (*anim_step as f64) / (animation_frames_per_char as f64);
+                let frame = render_frame(
+                    visible_sub,
+                    *pen_x,
+                    *pen_y,
+                    anim_offset,
+                    width,
+                    height,
+                    &line_styles,
+                    &y_positions,
+                    margin_x,
+                    font,
+                    default_font_scale,
+                    header_font_scale,
+                    default_thickness,
+                    header_thickness,
+                    default_text_color,
+                    header_text_color,
+                    bg_color,
+                    pen_color,
+                    pen_tip_radius,
+                    pen_length,
+                    pen_thickness,
+                    pen_base_angle,
+                    pen_movement_amplitude,
+                ).unwrap();
+                frame.data_bytes().unwrap().to_vec()
+            })
+            .collect()
+    });
 
     println!("Animation frames rendered in {:.3}s", t0.elapsed().as_secs_f64());
 
-    // Write all animation frames in one large batch
+    // STEP 5: Write frames to FFmpeg in chunks to avoid broken pipe
+    println!("Writing frames to FFmpeg...");
+    const CHUNK_SIZE: usize = 1024 * 1024; // 1MB chunks
     let mut buffer: Vec<u8> = Vec::with_capacity(animation_frames.len() * width as usize * height as usize * 3);
+    
     for frame_data in &animation_frames {
         buffer.extend_from_slice(frame_data);
     }
     
-    stdin.write_all(&buffer)
-        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to write animation frames: {}", e)))?;
+    // Write in chunks to prevent pipe overflow
+    for chunk in buffer.chunks(CHUNK_SIZE) {
+        stdin.write_all(chunk)
+            .map_err(|e| pyo3::exceptions::PyIOError::new_err(
+                format!("Failed to write frame chunk to FFmpeg (broken pipe?): {}", e)
+            ))?;
+    }
 
     drop(stdin);
-    child.wait().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("FFmpeg animation failed: {}", e)))?;
+    
+    // Wait and check FFmpeg exit status
+    let output = child.wait_with_output()
+        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("FFmpeg animation failed: {}", e)))?;
+    
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(pyo3::exceptions::PyRuntimeError::new_err(
+            format!("FFmpeg animation exited with status {:?}:
+{}", output.status.code(), stderr)
+        ));
+    }
+
+    // Verify animation video was created
+    if !Path::new(&animation_video_path).exists() {
+        return Err(pyo3::exceptions::PyFileNotFoundError::new_err(
+            format!("Animation video not created at: {}", animation_video_path)
+        ));
+    }
 
     println!("Animation video created in {:.3}s", t0.elapsed().as_secs_f64());
 
-    // STEP 4: Combine animation + static frame + audio using FFmpeg filters
+    // STEP 6: Combine animation + static frame + audio
+    println!("Combining videos with FFmpeg filters...");
+    
     let animation_actual_duration = frame_specs.len() as f64 / fps;
     let speed_multiplier = animation_duration / animation_actual_duration;
 
-    println!("Combining videos with FFmpeg filters...");
-    
     let filter_complex = format!(
-        "[0:v]setpts={}*PTS[v0];[1:v]loop=loop=-1:size=1:start=0,trim=duration={}[v1];[v0][v1]concat=n=2:v=1:a=0[outv]",
+        "[0:v]setpts={}*PTS[v0];[1:v]loop=loop=-1:size=1:start=0,trim=duration={:.3}[v1];[v0][v1]concat=n=2:v=1:a=0[outv]",
         speed_multiplier,
         static_duration
     );
 
-    let mut combine_child = Command::new("ffmpeg")
+    let combine_output = Command::new("ffmpeg")
         .arg("-y")
         .arg("-i").arg(&animation_video_path)
         .arg("-loop").arg("1")
@@ -252,6 +306,7 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
         .arg("-filter_complex").arg(&filter_complex)
         .arg("-map").arg("[outv]")
         .arg("-map").arg("2:a:0")
+        .arg("-vsync").arg("cfr")
         .arg("-c:v").arg("libx264")
         .arg("-preset").arg("ultrafast")
         .arg("-crf").arg("28")
@@ -259,10 +314,16 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
         .arg("-c:a").arg("aac")
         .arg("-shortest")
         .arg(&final_video_path)
-        .spawn()
-        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to spawn FFmpeg for combine: {}", e)))?;
-
-    combine_child.wait().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("FFmpeg combine failed: {}", e)))?;
+        .output()
+        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to run FFmpeg combine: {}", e)))?;
+
+    if !combine_output.status.success() {
+        let stderr = String::from_utf8_lossy(&combine_output.stderr);
+        return Err(pyo3::exceptions::PyRuntimeError::new_err(
+            format!("FFmpeg combine exited with status {:?}:
+{}", combine_output.status.code(), stderr)
+        ));
+    }
 
     let elapsed = t0.elapsed().as_secs_f64();
     println!("Total processing time: {:.3}s", elapsed);
@@ -271,6 +332,14 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
     let _ = std::fs::remove_file(&animation_video_path);
     let _ = std::fs::remove_file(&static_frame_path);
 
+    // Final verification before returning
+    if !Path::new(&final_video_path).exists() {
+        return Err(pyo3::exceptions::PyFileNotFoundError::new_err(
+            format!("Final video was not created at: {}", final_video_path)
+        ));
+    }
+
+    println!("✓ Video successfully generated: {}", final_video_path);
     Ok(Some(final_video_path))
 }
 
@@ -363,7 +432,8 @@ fn render_frame(
     let mut img = Mat::new_rows_cols_with_default(height, width, CV_8UC3, bg_color)
         .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create Mat: {}", e)))?;
 
-    let lines: Vec<&str> = visible_text.split('\n').collect();
+    let lines: Vec<&str> = visible_text.split('
+').collect();
     for (idx, &line) in lines.iter().enumerate() {
         let is_header = line_styles[idx];
         let font_scale = if is_header { header_font_scale } else { default_font_scale };
@@ -400,5 +470,4 @@ fn render_frame(
 fn rust_highlight(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(generate_video_clip, m)?)?;
     Ok(())
-}
-
+}