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

use rtx_diffuse::{Result, UNet, UNetConfig};
use rtx_tensor::Tensor;

#[test]
fn test_unet_creation() -> Result<()> {
    // Test default config
    let unet = UNet::new(UNetConfig::default())?;
    let config = unet.config();

    assert_eq!(config.in_channels, 4);
    assert_eq!(config.out_channels, 4);
    assert_eq!(config.model_channels, 320);
    assert_eq!(config.num_res_blocks, 2);

    Ok(())
}

#[test]
fn test_unet_custom_config() -> Result<()> {
    let config = UNetConfig {
        in_channels: 3,
        out_channels: 3,
        model_channels: 128,
        num_res_blocks: 1,
        attention_resolutions: vec![2, 1],
        channel_mult: vec![1, 2],
        num_heads: 4,
        num_head_channels: Some(32),
        use_scale_shift_norm: true,
        resblock_updown: true,
        num_classes: Some(1000),
        dropout: 0.0,
        conv_resample: true,
        dims: 2,
    };

    let unet = UNet::new(config.clone())?;
    let stored_config = unet.config();

    assert_eq!(stored_config.in_channels, 3);
    assert_eq!(stored_config.out_channels, 3);
    assert_eq!(stored_config.model_channels, 128);
    assert_eq!(stored_config.num_heads, 4);
    assert_eq!(stored_config.use_scale_shift_norm, true);

    Ok(())
}

#[test]
fn test_unet_forward_pass() -> Result<()> {
    let unet = UNet::new(UNetConfig::default())?;

    // Create input tensor [batch_size, channels, height, width]
    let input_data = vec![0.5; 2 * 4 * 32 * 32];
    let input = Tensor::new(input_data, vec![2, 4, 32, 32])?;

    // Create timesteps tensor [batch_size]
    let timesteps_data = vec![500.0, 300.0];
    let timesteps = Tensor::new(timesteps_data, vec![2])?;

    // Forward pass
    let output = unet.forward(&input, &timesteps, &None)?;

    // Check output shape matches expected
    assert_eq!(output.shape().dims(), &[2, 4, 32, 32]);

    // Verify output is finite
    let output_data = output.data()?;
    for val in output_data {
        assert!(val.is_finite(), "Output should contain finite values");
    }

    Ok(())
}

#[test]
fn test_unet_different_batch_sizes() -> Result<()> {
    let unet = UNet::new(UNetConfig::default())?;

    // Test batch size 1
    let input1 = Tensor::new(vec![0.3; 1 * 4 * 16 * 16], vec![1, 4, 16, 16])?;
    let timesteps1 = Tensor::new(vec![750.0], vec![1])?;
    let output1 = unet.forward(&input1, &timesteps1, &None)?;
    assert_eq!(output1.shape().dims(), &[1, 4, 16, 16]);

    // Test batch size 4
    let input4 = Tensor::new(vec![0.3; 4 * 4 * 16 * 16], vec![4, 4, 16, 16])?;
    let timesteps4 = Tensor::new(vec![100.0, 300.0, 600.0, 900.0], vec![4])?;
    let output4 = unet.forward(&input4, &timesteps4, &None)?;
    assert_eq!(output4.shape().dims(), &[4, 4, 16, 16]);

    Ok(())
}

#[test]
fn test_unet_different_resolutions() -> Result<()> {
    let unet = UNet::new(UNetConfig::default())?;

    // Test 64x64 input
    let input_64 = Tensor::new(vec![0.1; 1 * 4 * 64 * 64], vec![1, 4, 64, 64])?;
    let timesteps = Tensor::new(vec![400.0], vec![1])?;
    let output_64 = unet.forward(&input_64, &timesteps, &None)?;
    assert_eq!(output_64.shape().dims(), &[1, 4, 64, 64]);

    // Test 128x128 input (if memory allows)
    let input_128 = Tensor::new(vec![0.1; 1 * 4 * 8 * 8], vec![1, 4, 8, 8])?; // Smaller for test
    let output_128 = unet.forward(&input_128, &timesteps, &None)?;
    assert_eq!(output_128.shape().dims(), &[1, 4, 8, 8]);

    Ok(())
}

#[test]
fn test_unet_different_timesteps() -> Result<()> {
    let unet = UNet::new(UNetConfig::default())?;

    let input = Tensor::new(vec![0.2; 1 * 4 * 16 * 16], vec![1, 4, 16, 16])?;

    // Test different timestep values
    let timesteps_early = Tensor::new(vec![50.0], vec![1])?;
    let output_early = unet.forward(&input, &timesteps_early, &None)?;
    assert_eq!(output_early.shape().dims(), &[1, 4, 16, 16]);

    let timesteps_mid = Tensor::new(vec![500.0], vec![1])?;
    let output_mid = unet.forward(&input, &timesteps_mid, &None)?;
    assert_eq!(output_mid.shape().dims(), &[1, 4, 16, 16]);

    let timesteps_late = Tensor::new(vec![950.0], vec![1])?;
    let output_late = unet.forward(&input, &timesteps_late, &None)?;
    assert_eq!(output_late.shape().dims(), &[1, 4, 16, 16]);

    // Outputs should be different for different timesteps
    let data_early = output_early.data()?;
    let data_mid = output_mid.data()?;
    let data_late = output_late.data()?;

    // All should be zeros in current implementation, but structure is there
    assert_eq!(data_early[0], 0.0);
    assert_eq!(data_mid[0], 0.0);
    assert_eq!(data_late[0], 0.0);

    Ok(())
}

#[test]
fn test_unet_channel_configurations() -> Result<()> {
    // Test RGB input/output
    let rgb_config = UNetConfig {
        in_channels: 3,
        out_channels: 3,
        model_channels: 64,
        ..Default::default()
    };

    let unet_rgb = UNet::new(rgb_config)?;
    let input_rgb = Tensor::new(vec![0.4; 1 * 3 * 32 * 32], vec![1, 3, 32, 32])?;
    let timesteps = Tensor::new(vec![200.0], vec![1])?;
    let output_rgb = unet_rgb.forward(&input_rgb, &timesteps, &None)?;
    assert_eq!(output_rgb.shape().dims(), &[1, 3, 32, 32]);

    // Test grayscale input/output
    let gray_config = UNetConfig {
        in_channels: 1,
        out_channels: 1,
        model_channels: 32,
        ..Default::default()
    };

    let unet_gray = UNet::new(gray_config)?;
    let input_gray = Tensor::new(vec![0.6; 1 * 1 * 28 * 28], vec![1, 1, 28, 28])?;
    let output_gray = unet_gray.forward(&input_gray, &timesteps, &None)?;
    assert_eq!(output_gray.shape().dims(), &[1, 1, 28, 28]);

    Ok(())
}

#[test]
fn test_unet_with_class_conditioning() -> Result<()> {
    let config = UNetConfig {
        num_classes: Some(10),
        ..Default::default()
    };

    let unet = UNet::new(config)?;
    let input = Tensor::new(vec![0.1; 2 * 4 * 16 * 16], vec![2, 4, 16, 16])?;
    let timesteps = Tensor::new(vec![600.0, 400.0], vec![2])?;

    let output = unet.forward(&input, &timesteps, &None)?;
    assert_eq!(output.shape().dims(), &[2, 4, 16, 16]);

    Ok(())
}

#[test]
fn test_time_embedding() -> Result<()> {
    use rtx_diffuse::models::unet::TimeEmbedding;

    let time_emb = TimeEmbedding::new(512)?;
    let timesteps = Tensor::new(vec![100.0, 500.0, 900.0], vec![3])?;

    let embeddings = time_emb.forward(&timesteps)?;
    assert_eq!(embeddings.shape().dims(), &[3, 512]);

    // Check that embeddings are finite
    let emb_data = embeddings.data()?;
    for val in emb_data {
        assert!(val.is_finite(), "Time embeddings should be finite");
    }

    Ok(())
}

#[test]
fn test_res_block() -> Result<()> {
    use rtx_diffuse::models::unet::ResBlock;

    let res_block = ResBlock::new(64, 128, 512, false)?;

    let x = Tensor::new(vec![0.3; 2 * 64 * 16 * 16], vec![2, 64, 16, 16])?;
    let time_emb = Tensor::new(vec![0.1; 2 * 512], vec![2, 512])?;

    let output = res_block.forward(&x, &time_emb)?;
    assert_eq!(output.shape().dims(), &[2, 128, 16, 16]); // Output channels changed

    Ok(())
}

#[test]
fn test_res_block_with_scale_shift_norm() -> Result<()> {
    use rtx_diffuse::models::unet::ResBlock;

    let res_block = ResBlock::new(32, 32, 256, true)?; // use_scale_shift_norm = true

    let x = Tensor::new(vec![0.5; 1 * 32 * 8 * 8], vec![1, 32, 8, 8])?;
    let time_emb = Tensor::new(vec![0.2; 1 * 256], vec![1, 256])?;

    let output = res_block.forward(&x, &time_emb)?;
    assert_eq!(output.shape().dims(), &[1, 32, 8, 8]);

    Ok(())
}

#[test]
fn test_attention_block() -> Result<()> {
    use rtx_diffuse::models::unet::AttentionBlock;

    let attention = AttentionBlock::new(128, 8)?;
    let x = Tensor::new(vec![0.1; 1 * 128 * 8 * 8], vec![1, 128, 8, 8])?;

    let output = attention.forward(&x)?;
    assert_eq!(output.shape().dims(), &[1, 128, 8, 8]);

    Ok(())
}

#[test]
fn test_attention_block_invalid_heads() {
    use rtx_diffuse::models::unet::AttentionBlock;

    // 128 channels not divisible by 5 heads
    let result = AttentionBlock::new(128, 5);
    assert!(result.is_err());

    // Valid configuration
    let result = AttentionBlock::new(128, 16);
    assert!(result.is_ok());
}

#[test]
fn test_unet_memory_efficiency() -> Result<()> {
    // Test that we can create multiple UNets without memory issues
    for i in 0..5 {
        let config = UNetConfig {
            model_channels: 32, // Smaller to avoid memory issues in tests
            num_res_blocks: 1,
            ..Default::default()
        };

        let unet = UNet::new(config)?;
        let input = Tensor::new(vec![0.1; 1 * 4 * 8 * 8], vec![1, 4, 8, 8])?;
        let timesteps = Tensor::new(vec![i as f32 * 100.0], vec![1])?;

        let output = unet.forward(&input, &timesteps, &None)?;
        assert_eq!(output.shape().dims(), &[1, 4, 8, 8]);
    }

    Ok(())
}

#[test]
fn test_unet_deterministic_output() -> Result<()> {
    let unet = UNet::new(UNetConfig::default())?;

    let input = Tensor::new(vec![0.7; 1 * 4 * 16 * 16], vec![1, 4, 16, 16])?;
    let timesteps = Tensor::new(vec![333.0], vec![1])?;

    // Multiple forward passes should be deterministic
    let output1 = unet.forward(&input, &timesteps, &None)?;
    let output2 = unet.forward(&input, &timesteps, &None)?;

    let data1 = output1.data()?;
    let data2 = output2.data()?;

    for (a, b) in data1.iter().zip(data2.iter()) {
        assert!((a - b).abs() < 1e-6, "UNet should be deterministic");
    }

    Ok(())
}