Update rust_highlight/src/lib.rs

rust_highlight/src/lib.rs

@@ -1,10 +1,7 @@
 // rust_highlight/src/lib.rs
-
-
-
 use pyo3::prelude::*;
 use pyo3::types::PyModule;
-use opencv::core::{Mat, Point, Scalar, CV_8UC3};
+use opencv::core::{Mat, Point, Scalar, CV_8UC3, Vector};
 use opencv::imgproc::{circle, get_text_size, line, put_text, HersheyFonts, LineTypes};
 use opencv::prelude::*;
 use std::process::{Command, Stdio};
@@ -12,6 +9,7 @@ use std::io::Write;
 use std::time::Instant;
 use std::path::Path;
 use std::f64::consts::PI;
+use rayon::prelude::*;
 
 #[pyfunction]
 fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f64, clips_dir: String) -> PyResult<Option<String>> {
@@ -21,7 +19,7 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
 
     let skip_spaces = false;
     let fps: f64 = 30.0;
-    let animation_frames_per_char: usize = 2;
+    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;
@@ -32,21 +30,22 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
     let header_font_scale: f64 = 2.0;
     let default_thickness: i32 = 2;
     let header_thickness: i32 = 3;
-    let default_text_color = Scalar::new(0.0, 0.0, 0.0, 0.0);  // BGR Black
-    let header_text_color = Scalar::new(255.0, 0.0, 0.0, 0.0);  // BGR Blue
-    let bg_color = Scalar::new(255.0, 255.0, 255.0, 0.0);  // BGR White
+    let default_text_color = Scalar::new(0.0, 0.0, 0.0, 0.0);
+    let header_text_color = Scalar::new(255.0, 0.0, 0.0, 0.0);
+    let bg_color = Scalar::new(255.0, 255.0, 255.0, 0.0);
     let ffmpeg_preset = "ultrafast";
     let crf = "28";
-    // Pen settings
-    let pen_color = Scalar::new(0.0, 0.0, 255.0, 0.0);  // BGR Red
+    let pen_color = Scalar::new(0.0, 0.0, 255.0, 0.0);
     let pen_tip_radius: i32 = 5;
     let pen_length: i32 = 20;
     let pen_thickness: i32 = 2;
     let pen_base_angle: i32 = 45;
     let pen_movement_amplitude: i32 = 10;
 
-    let silent_video_name = format!("silent_video{}.mp4", id);
-    let silent_video_path = format!("{}/{}", clips_dir, silent_video_name);
+    let animation_video_name = format!("anim_video{}.mp4", id);
+    let animation_video_path = format!("{}/{}", clips_dir, animation_video_name);
+    let static_frame_name = format!("static_{}.png", id);
+    let static_frame_path = format!("{}/{}", clips_dir, static_frame_name);
     let final_video_name = format!("clip{}.mp4", id);
     let final_video_path = format!("{}/{}", clips_dir, final_video_name);
 
@@ -67,26 +66,19 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
     };
 
     let total_glyphs = visible_indices.len();
-    println!("Wrapped lines: {} lines, total glyphs (counted): {}", wrapped_lines.len(), total_glyphs);
+    println!("Wrapped lines: {} lines, total glyphs: {}", wrapped_lines.len(), total_glyphs);
     if total_glyphs == 0 {
         println!("No text to animate.");
         return Ok(None);
     }
 
-    // Calculate animation duration (1/4 of audio duration)
+    // Calculate durations
     let animation_duration = duration / 4.0;
     let static_duration = duration - animation_duration;
     
-    // Calculate total frames needed
-    let animation_frames = total_glyphs * animation_frames_per_char;
-    let animation_total_frames = (animation_duration * fps) as usize;
-    let static_frames = (static_duration * fps) as usize;
-    
-    println!("Animation duration: {:.3}s ({} frames), Static duration: {:.3}s ({} frames)", 
-             animation_duration, animation_total_frames, static_duration, static_frames);
+    println!("Animation duration: {:.3}s, Static duration: {:.3}s", animation_duration, static_duration);
 
     // Pre-calc line heights and y_positions
-    let mut line_heights: Vec<i32> = Vec::new();
     let mut y_positions: Vec<i32> = Vec::new();
     let mut y = margin_y;
     for (i, line) in wrapped_lines.iter().enumerate() {
@@ -98,43 +90,21 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
         let size = get_text_size(&line_for_size, font, font_scale, thickness, &mut base_line).unwrap();
         let h = size.height;
         let lh = h + base_line + line_spacing;
-        line_heights.push(lh);
         y_positions.push(y);
         y += lh;
     }
 
-    // FFmpeg for silent video
-    let mut child = Command::new("ffmpeg")
-        .arg("-y")
-        .arg("-f").arg("rawvideo")
-        .arg("-pix_fmt").arg("bgr24")
-        .arg("-s").arg(format!("{}x{}", width, height))
-        .arg("-r").arg(fps.to_string())
-        .arg("-i").arg("-")
-        .arg("-an")
-        .arg("-c:v").arg("libx264")
-        .arg("-preset").arg(ffmpeg_preset)
-        .arg("-crf").arg(crf)
-        .arg("-pix_fmt").arg("yuv420p")
-        .arg(&silent_video_path)
-        .stdin(Stdio::piped())
-        .spawn()
-        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to spawn FFmpeg: {}", e)))?;
-
-    let mut stdin = child.stdin.take().unwrap();
-
     let t0 = Instant::now();
-    let mut frames_sent: usize = 0;
-    let mut prev_visible_sub = String::new();
 
-    // Pre-render base frame with full text (optimization)
-    let base_frame = render_frame(
+    // STEP 1: Pre-render the final static frame (used for Phase 2)
+    println!("Rendering static frame...");
+    let static_frame = render_frame(
         &full_text,
         -1,
         0,
         0.0,
-        width as i32,
-        height as i32,
+        width,
+        height,
         &line_styles,
         &y_positions,
         margin_x,
@@ -154,13 +124,20 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
         pen_movement_amplitude,
     )?;
 
-    // Buffer for batch writing
-    let frame_size = (width * height * 3) as usize;
-    let mut buffer: Vec<u8> = Vec::with_capacity(frame_size * 100);
+    // Save static frame as PNG
+    let mut params = Vector::new();
+    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
+    let mut frame_specs: Vec<(String, i32, i32, usize)> = Vec::new();
+    let mut prev_visible_sub = String::new();
 
-    // PHASE 1: Animation phase (first 1/4 of duration)
     for &idx_in_full in visible_indices.iter() {
-        let visible_sub = &full_text[0..=idx_in_full];
+        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 last_line = lines.last().unwrap();
@@ -176,121 +153,123 @@ fn generate_video_clip(id: usize, text: String, audio_path: String, duration: f6
             let pen_y = y_positions[line_idx] + h / 2;
 
             for anim_step in 0..animation_frames_per_char {
-                let anim_offset = (anim_step as f64) / (animation_frames_per_char as f64);
-                let frame_img = render_frame(
-                    visible_sub,
-                    pen_x,
-                    pen_y,
-                    anim_offset,
-                    width as i32,
-                    height as i32,
-                    &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,
-                )?;
-                
-                buffer.extend_from_slice(frame_img.data_bytes().unwrap());
-                frames_sent += 1;
-                
-                // Flush buffer every 100 frames
-                if buffer.len() >= frame_size * 100 {
-                    stdin.write_all(&buffer)
-                        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to write frames: {}", e)))?;
-                    buffer.clear();
-                }
+                frame_specs.push((visible_sub.clone(), pen_x, pen_y, anim_step));
             }
-            prev_visible_sub = visible_sub.to_string();
+            prev_visible_sub = visible_sub;
         }
     }
 
-    // Flush remaining animation frames
-    if !buffer.is_empty() {
-        stdin.write_all(&buffer)
-            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to write frames: {}", e)))?;
-        buffer.clear();
-    }
+    println!("Total animation frames to render: {}", frame_specs.len());
+
+    // Parallel rendering using rayon
+    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 3: Write animation frames to FFmpeg
+    let mut child = Command::new("ffmpeg")
+        .arg("-y")
+        .arg("-f").arg("rawvideo")
+        .arg("-pix_fmt").arg("bgr24")
+        .arg("-s").arg(format!("{}x{}", width, height))
+        .arg("-r").arg(fps.to_string())
+        .arg("-i").arg("-")
+        .arg("-an")
+        .arg("-c:v").arg("libx264")
+        .arg("-preset").arg(ffmpeg_preset)
+        .arg("-crf").arg(crf)
+        .arg("-pix_fmt").arg("yuv420p")
+        .arg(&animation_video_path)
+        .stdin(Stdio::piped())
+        .spawn()
+        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to spawn FFmpeg: {}", e)))?;
 
-    // PHASE 2: Static phase (remaining 3/4 of duration)
-    // Reuse the pre-rendered base frame
-    println!("Animation complete. Rendering {} static frames...", static_frames);
-    
-    let base_frame_bytes = base_frame.data_bytes().unwrap();
-    
-    // Write static frames in batches
-    let batch_size = 300;
-    let full_batches = static_frames / batch_size;
-    let remainder = static_frames % batch_size;
-    
-    for _ in 0..full_batches {
-        for _ in 0..batch_size {
-            buffer.extend_from_slice(base_frame_bytes);
-        }
-        stdin.write_all(&buffer)
-            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to write static frames: {}", e)))?;
-        buffer.clear();
-        frames_sent += batch_size;
+    let mut stdin = child.stdin.take().unwrap();
+
+    // 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 remaining frames
-    if remainder > 0 {
-        for _ in 0..remainder {
-            buffer.extend_from_slice(base_frame_bytes);
-        }
-        stdin.write_all(&buffer)
-            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to write static frames: {}", e)))?;
-        frames_sent += remainder;
-    }
+    stdin.write_all(&buffer)
+        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to write animation frames: {}", e)))?;
 
     drop(stdin);
-    child.wait().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("FFmpeg failed: {}", e)))?;
+    child.wait().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("FFmpeg animation failed: {}", e)))?;
 
-    let elapsed = t0.elapsed().as_secs_f64();
-    println!("Total frames sent: {}, elapsed time: {:.3} seconds", frames_sent, elapsed);
+    println!("Animation video created in {:.3}s", t0.elapsed().as_secs_f64());
 
-    if !Path::new(&silent_video_path).exists() {
-        println!("Silent video generation failed.");
-        return Ok(None);
-    }
+    // 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;
 
-    let rendered_duration = frames_sent as f64 / fps;
-    println!("Rendered video duration: {:.3}s, Target audio duration: {:.3}s", rendered_duration, 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]",
+        speed_multiplier,
+        static_duration
+    );
 
-    // Combine with audio (no speed adjustment needed as we matched the duration)
     let mut combine_child = Command::new("ffmpeg")
         .arg("-y")
-        .arg("-i").arg(&silent_video_path)
+        .arg("-i").arg(&animation_video_path)
+        .arg("-loop").arg("1")
+        .arg("-i").arg(&static_frame_path)
         .arg("-i").arg(&audio_path)
+        .arg("-filter_complex").arg(&filter_complex)
+        .arg("-map").arg("[outv]")
+        .arg("-map").arg("2:a:0")
         .arg("-c:v").arg("libx264")
         .arg("-preset").arg("ultrafast")
         .arg("-crf").arg("28")
         .arg("-pix_fmt").arg("yuv420p")
         .arg("-c:a").arg("aac")
         .arg("-shortest")
-        .arg("-map").arg("0:v:0")
-        .arg("-map").arg("1:a:0")
         .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)))?;
 
-    // Clean up
-    std::fs::remove_file(&silent_video_path)
-        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("Failed to remove silent video: {}", e)))?;
+    let elapsed = t0.elapsed().as_secs_f64();
+    println!("Total processing time: {:.3}s", elapsed);
+
+    // Clean up temporary files
+    let _ = std::fs::remove_file(&animation_video_path);
+    let _ = std::fs::remove_file(&static_frame_path);
 
     Ok(Some(final_video_path))
 }
@@ -402,7 +381,6 @@ fn render_frame(
         }
     }
 
-    // Only draw pen if pen_x > 0
     if pen_x > 0 {
         let offset_y = (pen_movement_amplitude as f64 * (anim_offset * PI).sin()) as i32;
         let pen_tip_y = pen_y + offset_y;
@@ -422,9 +400,4 @@ fn render_frame(
 fn rust_highlight(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(generate_video_clip, m)?)?;
     Ok(())
-
-
-
-
-    
 }