Update rust_highlight/src/lib.rs

rust_highlight/src/lib.rs

@@ -10,18 +10,16 @@ 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;
+    let animation_frames_per_char: usize = 1;  // Reduced from 2 for speed
     let width: i32 = 1280;
     let height: i32 = 720;
     let margin_x: i32 = 40;
@@ -55,8 +53,7 @@ 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("
-");
+    let full_text = wrapped_lines.join("\n");
     if full_text.is_empty() {
         println!("No text to animate.");
         return Ok(None);
@@ -75,16 +72,10 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
         return Ok(None);
     }
 
-    // Calculate durations with validation
+    // Calculate durations
     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
@@ -105,7 +96,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
+    // STEP 1: Pre-render the final static frame (used for Phase 2)
     println!("Rendering static frame...");
     let static_frame = render_frame(
         &full_text,
@@ -138,15 +129,10 @@ 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)))?;
 
-    // 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...");
+    // STEP 2: Render animation frames in parallel (HUGE speedup)
+    println!("Rendering animation frames in parallel...");
+    
+    // Collect frame data for parallel processing
     let mut frame_specs: Vec<(String, i32, i32, usize)> = Vec::new();
     let mut prev_visible_sub = String::new();
 
@@ -175,8 +161,7 @@ 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 for animation
-    println!("Starting FFmpeg process for animation...");
+    // STEP 3: Start FFmpeg process first
     let mut child = Command::new("ffmpeg")
         .arg("-y")
         .arg("-f").arg("rawvideo")
@@ -191,111 +176,74 @@ 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();
 
-    // 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()
-    });
+    // 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();
 
     println!("Animation frames rendered in {:.3}s", t0.elapsed().as_secs_f64());
 
-    // 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
+    // Write all animation frames in one large batch
     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);
     }
     
-    // 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)
-            ))?;
-    }
+    stdin.write_all(&buffer)
+        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to write animation frames: {}", e)))?;
 
     drop(stdin);
-    
-    // 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)
-        ));
-    }
+    child.wait().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("FFmpeg animation failed: {}", e)))?;
 
     println!("Animation video created in {:.3}s", t0.elapsed().as_secs_f64());
 
-    // STEP 6: Combine animation + static frame + audio
-    println!("Combining videos with FFmpeg filters...");
-    
+    // STEP 4: Combine animation + static frame + audio using 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={:.3}[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={}[v1];[v0][v1]concat=n=2:v=1:a=0[outv]",
         speed_multiplier,
         static_duration
     );
 
-    let combine_output = Command::new("ffmpeg")
+    let mut combine_child = Command::new("ffmpeg")
         .arg("-y")
         .arg("-i").arg(&animation_video_path)
         .arg("-loop").arg("1")
@@ -304,7 +252,6 @@ 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")
@@ -312,16 +259,10 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
         .arg("-c:a").arg("aac")
         .arg("-shortest")
         .arg(&final_video_path)
-        .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)
-        ));
-    }
+        .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)))?;
 
     let elapsed = t0.elapsed().as_secs_f64();
     println!("Total processing time: {:.3}s", elapsed);
@@ -330,14 +271,6 @@ 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))
 }