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

//! Real Local Image Upload Test
//!
//! This example demonstrates how to upload your actual local image file
//! and use it with the AI Model Face & Hair Detail Fix template.

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;

/// Test with real local image upload
async fn test_with_real_local_image() -> Result<(), Box<dyn std::error::Error>> {
    println!("🚀 Starting Real Local Image Upload Test");
    
    // Initialize client
    let mut client = ComfyUIClient::new(ComfyUIClientConfig {
        base_url: "http://192.168.0.193:8188".to_string(),
        ..Default::default()
    })?;

    // Connect to ComfyUI server
    println!("📡 Connecting to ComfyUI server...");
    client.connect().await?;
    println!("✅ Connected successfully!");

    // Register the template
    println!("📝 Registering AI Model Face & Hair Fix 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_TEMPLATE.metadata.id)
        .ok_or("Failed to register template")?;
    println!("✅ Template registered!");

    // Your actual local image path
    let local_image_path = "20250808-111737.png";
    
    println!("\n📤 Uploading your image: {}", local_image_path);
    println!("⏳ This may take a moment depending on file size and network speed...");

    match upload_and_execute(&client, template, local_image_path).await {
        Ok(_) => println!("✅ Upload and execution completed successfully!"),
        Err(e) => {
            println!("\n💥 Upload/Execution failed: {}", e);
            show_troubleshooting_tips(&e);
        }
    }

    // Disconnect
    println!("\n🔌 Disconnecting from ComfyUI server...");
    client.disconnect().await?;
    println!("👋 Disconnected successfully!");

    Ok(())
}

/// Upload image and execute template
async fn upload_and_execute(
    client: &ComfyUIClient,
    template: &comfyui_sdk::WorkflowTemplate,
    image_path: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    // Step 1: Upload image to ComfyUI server
    println!("\n🔄 Step 1: Uploading image to ComfyUI server...");
    
    if !Path::new(image_path).exists() {
        return Err(format!("File not found: {}", image_path).into());
    }

    let upload_response = client.upload_image(image_path, false).await?;
    let uploaded_filename = upload_response.name;
    println!("✅ Upload successful! Server filename: {}", uploaded_filename);

    // Step 2: Execute AI Model Face & Hair Enhancement
    println!("\n🔄 Step 2: Executing AI Model Face & Hair Enhancement...");
    
    // Create execution 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: {}% ({}/{}) - Node: {}", 
                    percentage, progress.progress, progress.max, progress.node_id);
            })
            .with_executing(|node_id| {
                println!("🔄 Executing node: {}", node_id);
            })
            .with_executed(|result| {
                println!("✅ Node completed - Prompt ID: {}", result.prompt_id);
            })
            .with_error(|error| {
                println!("❌ Execution error: {}", error.message);
            })
    );

    // Prepare parameters
    let mut parameters = HashMap::new();
    parameters.insert("input_image".to_string(), json!(uploaded_filename));
    parameters.insert("face_prompt".to_string(), 
        json!("beautiful woman, perfect skin, detailed facial features, professional photography"));
    parameters.insert("face_denoise".to_string(), json!("0.25"));

    let execution_options = ExecutionOptions {
        timeout: Some(std::time::Duration::from_secs(180)), // 3 minutes timeout
        priority: None,
    };

    let result = client.execute_template_with_callbacks(
        template,
        parameters,
        execution_options,
        callbacks,
    ).await?;

    if result.success {
        println!("\n🎉 AI Model Enhancement completed successfully!");
        println!("📊 Total execution time: {}ms ({:.1}s)", 
            result.execution_time, result.execution_time as f64 / 1000.0);
        println!("🆔 Prompt ID: {}", result.prompt_id);
        
        if let Some(outputs) = &result.outputs {
            println!("\n📁 Generated outputs:");
            println!("{}", serde_json::to_string_pretty(outputs)?);
            
            // Convert outputs to HTTP URLs
            let image_urls = client.outputs_to_urls(outputs);
            println!("\n🖼️  Enhanced Image URLs:");
            for (index, url) in image_urls.iter().enumerate() {
                println!("   {}. {}", index + 1, url);
            }
            
            println!("\n💡 How to use these URLs:");
            println!("   - Copy the URL and paste in your browser to view");
            println!("   - Right-click and \"Save As\" to download");
            println!("   - Use in your application to display the enhanced image");
        }
    } else {
        println!("\n❌ AI Model Enhancement failed!");
        if let Some(error) = &result.error {
            println!("Error details: {}", error.message);
        }
    }

    Ok(())
}

/// Show troubleshooting tips based on error
fn show_troubleshooting_tips(error: &Box<dyn std::error::Error>) {
    let error_msg = error.to_string();
    
    if error_msg.contains("File not found") || error_msg.contains("No such file") {
        println!("\n💡 Troubleshooting Tips:");
        println!("1. 🔍 Check if the file path is correct");
        println!("2. 📁 Make sure the file exists at the specified location");
        println!("3. 🔐 Ensure you have read permissions for the file");
        println!("4. 📝 Try using absolute path: /home/user/images/20250808-111737.png");
        println!("5. 🖼️  Verify the file is a valid image format (PNG, JPG, etc.)");
    } else if error_msg.contains("Failed to upload") || error_msg.contains("HTTP") {
        println!("\n💡 Network/Server Issues:");
        println!("1. 🌐 Check if ComfyUI server is running and accessible");
        println!("2. 📡 Verify network connection to the server");
        println!("3. 💾 Check if server has enough disk space");
        println!("4. 🔒 Ensure server allows file uploads");
    }
}

/// Alternative method using one-step upload and execute (conceptual)
async fn test_one_step_method() -> Result<(), Box<dyn std::error::Error>> {
    println!("\n\n🚀 Testing One-Step Method (Upload + Execute)");
    
    let mut client = ComfyUIClient::new(ComfyUIClientConfig {
        base_url: "http://192.168.0.193:8188".to_string(),
        ..Default::default()
    })?;

    client.connect().await?;
    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_TEMPLATE.metadata.id)
        .ok_or("Failed to register template")?;

    let local_image_path = "20250808-111737.png";
    
    println!("🔄 One-step upload and execute...");
    
    // For now, we'll implement this as upload + execute since we don't have
    // a dedicated one-step method yet
    match upload_and_execute_one_step(&client, template, local_image_path).await {
        Ok(_) => println!("🎉 One-step method completed successfully!"),
        Err(e) => println!("⚠️  One-step method failed: {}", e),
    }

    client.disconnect().await?;
    Ok(())
}

/// One-step upload and execute (simplified version)
async fn upload_and_execute_one_step(
    client: &ComfyUIClient,
    template: &comfyui_sdk::WorkflowTemplate,
    image_path: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    // Upload image
    let upload_response = client.upload_image(image_path, false).await?;
    
    // Create simple 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: {}% - Node: {}", percentage, progress.node_id);
            })
    );

    // Execute template
    let mut parameters = HashMap::new();
    parameters.insert("input_image".to_string(), json!(upload_response.name));
    parameters.insert("face_prompt".to_string(), 
        json!("stunning model, flawless skin, professional photography, high quality"));
    parameters.insert("face_denoise".to_string(), json!("0.3"));

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

    if result.success {
        if let Some(outputs) = &result.outputs {
            let image_urls = client.outputs_to_urls(outputs);
            println!("🖼️  Result URLs: {:?}", image_urls);
        }
    }

    Ok(())
}

/// Show usage examples
fn show_usage_examples() {
    println!("\n📚 Complete Usage Guide for Local Images:");
    println!("");
    println!("🔧 Basic Setup:");
    println!("```rust");
    println!("use comfyui_sdk::{{ComfyUIClient, ComfyUIClientConfig, AI_MODEL_FACE_HAIR_FIX_TEMPLATE}};");
    println!("");
    println!("let mut client = ComfyUIClient::new(ComfyUIClientConfig {{");
    println!("    base_url: \"http://your-comfyui-server:8188\".to_string(),");
    println!("    ..Default::default()");
    println!("}});");
    println!("```");
    println!("");
    println!("📤 Method 1 - Manual Upload:");
    println!("```rust");
    println!("// Upload image first");
    println!("let upload_response = client.upload_image(");
    println!("    \"/path/to/your/image.png\", false");
    println!(").await?;");
    println!("");
    println!("// Then execute template");
    println!("let result = client.execute_template(template, parameters, options).await?;");
    println!("```");
    println!("");
    println!("🚀 Method 2 - Combined Upload and Execute:");
    println!("```rust");
    println!("// Upload and execute in sequence");
    println!("let upload_response = client.upload_image(image_path, false).await?;");
    println!("let mut parameters = HashMap::new();");
    println!("parameters.insert(\"input_image\".to_string(), json!(upload_response.name));");
    println!("let result = client.execute_template_with_callbacks(");
    println!("    template, parameters, options, callbacks");
    println!(").await?;");
    println!("```");
}

/// Run all tests
async fn run_all_tests() -> Result<(), Box<dyn std::error::Error>> {
    show_usage_examples();
    test_with_real_local_image().await?;
    // Uncomment to test one-step method as well
    // test_one_step_method().await?;
    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    run_all_tests().await
}

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

    #[tokio::test]
    async fn test_show_usage_examples() {
        show_usage_examples();
        // This test just ensures the function runs without panicking
    }
}