rustytorch/crates/models/rtx-diffuse/tests/unipc_integration_test.rs

//! Integration tests for UniPC sampler
//!
//! Tests the complete workflow including scheduler integration
//! and compatibility with existing diffusion infrastructure.

use rtx_diffuse::{
    DiffusionScheduler, NoiseGenerator, NoiseSchedule, PredictionType, Result, SchedulerType,
    UniPCConfig, UniPCSampler,
};
use rtx_tensor::{Device, Tensor};

fn create_test_noise_generator() -> NoiseGenerator {
    NoiseGenerator::new(
        NoiseSchedule::Linear {
            beta_start: 0.0001,
            beta_end: 0.02,
        },
        1000,
        Some(42),
    )
    .unwrap()
}

fn create_test_tensor(shape: Vec<usize>) -> Tensor {
    let device = Device::cuda(0).unwrap_or(Device::default());
    Tensor::randn(&shape, &device).unwrap()
}

#[test]
fn test_unipc_scheduler_integration() {
    // Create UniPC scheduler
    let noise_gen = create_test_noise_generator();
    let scheduler = DiffusionScheduler::new(
        SchedulerType::UniPC {
            predictor_order: 2,
            corrector_order: 1,
            use_corrector: true,
        },
        noise_gen,
        10, // 10 inference steps for fast sampling
    )
    .unwrap();

    let sample = create_test_tensor(vec![1, 3, 32, 32]);
    let model_output = create_test_tensor(vec![1, 3, 32, 32]);

    // Test scheduler step
    let result = scheduler.step(&model_output, 500, &sample, None);
    assert!(result.is_ok());
    assert_eq!(result.unwrap().dims(), sample.dims());

    // Test timesteps configuration
    let timesteps = scheduler.timesteps();
    assert_eq!(timesteps.len(), 10);
}

#[test]
fn test_unipc_standalone_sampler() {
    let mut sampler =
        UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();

    let sample = create_test_tensor(vec![1, 3, 32, 32]);
    let model_output = create_test_tensor(vec![1, 3, 32, 32]);

    // Test full sampling workflow
    let result = sampler.step(&model_output, 500, &sample).unwrap();
    assert_eq!(result.dims(), sample.dims());
    assert_eq!(sampler.stats().nfe, 1);
}

#[test]
fn test_unipc_fast_sampling() {
    let mut sampler = UniPCSampler::new(
        UniPCConfig {
            predictor_order: 2,
            corrector_order: 1,
            use_corrector: true,
            adaptive_order: false,
            prediction_type: PredictionType::Epsilon,
            variance_reduction: true,
            corrector_iterations: 1,
            max_order: 3,
        },
        create_test_noise_generator(),
    )
    .unwrap();

    // Configure fast timesteps (7 steps)
    let timesteps = sampler.configure_fast_timesteps(7).unwrap();
    assert_eq!(timesteps.len(), 7);

    let mut current_sample = create_test_tensor(vec![1, 3, 32, 32]);
    let model_output = create_test_tensor(vec![1, 3, 32, 32]);

    // Simulate 7-step sampling
    for timestep in timesteps {
        let step_result = sampler
            .step(&model_output, timestep, &current_sample)
            .unwrap();
        current_sample = step_result;
    }

    // Verify sampling completed successfully
    assert_eq!(sampler.stats().nfe, 7);
    assert!(sampler.stats().predictor_steps >= 7);
    assert_eq!(current_sample.dims(), &[1, 3, 32, 32]);
}

#[test]
fn test_unipc_adaptive_order() {
    let mut sampler = UniPCSampler::new(
        UniPCConfig {
            adaptive_order: true,
            max_order: 3,
            ..Default::default()
        },
        create_test_noise_generator(),
    )
    .unwrap();

    // Test adaptive order selection
    let (pred_order, corr_order) = sampler.select_adaptive_order(0.01, 10).unwrap();
    assert!(pred_order >= 1 && pred_order <= 3);
    assert!(corr_order >= 1 && corr_order <= 3);

    // High error should prefer lower order
    let (low_order, _) = sampler.select_adaptive_order(0.5, 5).unwrap();
    let (high_order, _) = sampler.select_adaptive_order(0.01, 15).unwrap();
    assert!(low_order <= high_order);
}

#[test]
fn test_prediction_type_conversions() {
    let sampler = UniPCSampler::new(UniPCConfig::default(), create_test_noise_generator()).unwrap();

    let sample = create_test_tensor(vec![1, 3, 32, 32]);
    let model_output = create_test_tensor(vec![1, 3, 32, 32]);

    // Test all prediction type conversions
    let conversions = [
        (PredictionType::Epsilon, PredictionType::Data),
        (PredictionType::Epsilon, PredictionType::VPrediction),
        (PredictionType::Data, PredictionType::Epsilon),
        (PredictionType::VPrediction, PredictionType::Epsilon),
    ];

    for (from_type, to_type) in conversions {
        let result =
            sampler.convert_prediction_type(&model_output, 500, &sample, from_type, to_type);
        assert!(result.is_ok());
        assert_eq!(result.unwrap().dims(), model_output.dims());
    }
}

#[test]
fn test_unipc_error_conditions() {
    // Test invalid configuration
    let noise_gen = create_test_noise_generator();
    assert!(
        UniPCSampler::new(
            UniPCConfig {
                predictor_order: 0,
                ..Default::default()
            },
            noise_gen.clone()
        )
        .is_err()
    );
    assert!(
        UniPCSampler::new(
            UniPCConfig {
                corrector_order: 5,
                ..Default::default()
            },
            noise_gen.clone()
        )
        .is_err()
    );

    // Test invalid timestep configuration
    let sampler = UniPCSampler::new(UniPCConfig::default(), noise_gen).unwrap();
    assert!(sampler.configure_fast_timesteps(0).is_err());
    assert!(sampler.configure_fast_timesteps(2000).is_err());
}