mixvideo-v2/cargos/comfyui-sdk/examples/simple_local_image.rs

//! Simple Local Image Upload and Enhancement
//!
//! A simplified version that focuses on the core functionality:
//! 1. Upload a local image to ComfyUI server
//! 2. Execute AI Model Face & Hair enhancement
//! 3. Get the enhanced image URLs

use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use comfyui_sdk::{
    ComfyUIClient, ComfyUIClientConfig, ExecutionOptions, 
    AI_MODEL_FACE_HAIR_FIX_TEMPLATE
};
use comfyui_sdk::utils::SimpleCallbacks;
use serde_json::json;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("🚀 Simple Local Image Enhancement Test");
    
    // Configuration
    let server_url = "http://192.168.0.193:8188";
    let image_path = "20250808-111737.png";
    
    // Step 1: Initialize client and connect
    println!("📡 Connecting to ComfyUI server at {}...", server_url);
    let mut client = ComfyUIClient::new(ComfyUIClientConfig {
        base_url: server_url.to_string(),
        ..Default::default()
    })?;
    
    client.connect().await?;
    println!("✅ Connected!");

    // Step 2: Register template
    println!("📝 Registering AI enhancement template...");
    client.templates().register_from_data(AI_MODEL_FACE_HAIR_FIX_TEMPLATE.clone())?;
    let template = client.templates_ref()
        .get_by_id("ai-model-face-hair-fix")
        .ok_or("Template not found")?;
    println!("✅ Template ready!");

    // Step 3: Check if image exists
    if !Path::new(image_path).exists() {
        println!("❌ Image file not found: {}", image_path);
        println!("💡 Please make sure the file exists in the current directory");
        return Ok(());
    }

    // Step 4: Upload image
    println!("📤 Uploading image: {}...", image_path);
    let upload_response = client.upload_image(image_path, false).await?;
    println!("✅ Upload successful! Server filename: {}", upload_response.name);

    // Step 5: Setup progress callbacks
    let callbacks = Arc::new(
        SimpleCallbacks::new()
            .with_progress(|progress| {
                let percentage = (progress.progress as f64 / progress.max as f64 * 100.0) as u32;
                println!("⏳ Progress: {}%", percentage);
            })
            .with_executing(|node_id| {
                println!("🔄 Processing node: {}", node_id);
            })
            .with_error(|error| {
                println!("❌ Error: {}", error.message);
            })
    );

    // Step 6: Execute enhancement
    println!("🎨 Starting AI enhancement...");
    let mut parameters = HashMap::new();
    parameters.insert("input_image".to_string(), json!(upload_response.name));
    parameters.insert("face_prompt".to_string(), 
        json!("beautiful woman, perfect skin, detailed facial features"));
    parameters.insert("face_denoise".to_string(), json!("0.25"));

    let result = client.execute_template_with_callbacks(
        template,
        parameters,
        ExecutionOptions {
            timeout: Some(std::time::Duration::from_secs(180)),
            priority: None,
        },
        callbacks,
    ).await?;

    // Step 7: Show results
    if result.success {
        println!("\n🎉 Enhancement completed successfully!");
        println!("⏱️  Execution time: {:.1}s", result.execution_time as f64 / 1000.0);
        
        if let Some(outputs) = &result.outputs {
            let image_urls = client.outputs_to_urls(outputs);
            
            if !image_urls.is_empty() {
                println!("\n🖼️  Enhanced image URLs:");
                for (i, url) in image_urls.iter().enumerate() {
                    println!("   {}. {}", i + 1, url);
                }
                
                println!("\n💡 To view your enhanced image:");
                println!("   - Copy the URL above and paste it in your browser");
                println!("   - Right-click and 'Save As' to download the image");
            } else {
                println!("⚠️  No image URLs found in outputs");
            }
        }
    } else {
        println!("❌ Enhancement failed!");
        if let Some(error) = &result.error {
            println!("Error: {}", error.message);
        }
    }

    // Step 8: Cleanup
    client.disconnect().await?;
    println!("\n👋 Disconnected from server");
    
    Ok(())
}

/// Helper function to validate image file
fn validate_image_file(path: &str) -> Result<(), String> {
    let path = Path::new(path);
    
    if !path.exists() {
        return Err(format!("File does not exist: {}", path.display()));
    }
    
    if !path.is_file() {
        return Err(format!("Path is not a file: {}", path.display()));
    }
    
    // Check file extension
    if let Some(extension) = path.extension() {
        let ext = extension.to_string_lossy().to_lowercase();
        match ext.as_str() {
            "png" | "jpg" | "jpeg" | "bmp" | "tiff" | "webp" => Ok(()),
            _ => Err(format!("Unsupported image format: {}", ext)),
        }
    } else {
        Err("File has no extension".to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_validate_image_file() {
        // Test with non-existent file
        assert!(validate_image_file("non_existent.png").is_err());
        
        // Test with valid extensions
        // Note: These tests would pass if the files existed
        // assert!(validate_image_file("test.png").is_ok());
        // assert!(validate_image_file("test.jpg").is_ok());
    }
}