rustytorch/crates/models/rtx-diffuse/src/lcm_demo.rs

//! LCM Sampler Demo
//!
//! Demonstrates the LCM sampler implementation capabilities.

#[cfg(test)]
mod demo {
    use crate::{LCMConfig, LCMPredictionType, LCMSampler, NoiseGenerator, NoiseSchedule};
    use rtx_tensor::{Device, Tensor};

    #[test]
    fn test_lcm_demo_basic_usage() {
        // Create noise generator
        let noise_gen =
            NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();

        // Configure LCM for ultra-fast 2-step generation
        let config = LCMConfig {
            num_steps: 2,
            guidance_scale: 7.5,
            consistency_weight: 1.0,
            prediction_type: LCMPredictionType::Epsilon,
            solver_order: 2,
            training_mode: false,
            use_karras_sigmas: true,
            distillation_weight: 0.5,
        };

        // Create LCM sampler
        let mut sampler = LCMSampler::new(config, noise_gen).unwrap();

        // Mock diffusion model
        let mock_model = |x: &Tensor, t: u32| -> crate::Result<Tensor> {
            // Simple mock: gradually reduce noise based on timestep
            let noise_scale = (t as f32 / 1000.0) * 0.1;
            x.scalar_mul(noise_scale)
                .map_err(|e| crate::DiffusionError::TensorError(e.to_string()))
        };

        // Generate sample
        let input_shape = vec![1, 4, 64, 64];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let noisy_latent = Tensor::randn(&input_shape, &device).unwrap();

        println!("LCM Demo: Starting 2-step generation...");
        let result = sampler.sample(&noisy_latent, 999, mock_model);

        assert!(result.is_ok());
        let generated = result.unwrap();
        assert_eq!(generated.dims(), input_shape.as_slice());

        println!("LCM Demo: Generated tensor shape: {:?}", generated.dims());
        println!(
            "LCM Demo: NFE (Number of Function Evaluations): {}",
            sampler.stats().nfe
        );
        println!(
            "LCM Demo: Average sampling time: {:.2}ms",
            sampler.stats().avg_sampling_time_ms
        );

        // Verify ultra-fast sampling (should be exactly 2 model calls)
        assert_eq!(sampler.stats().nfe, 2);
    }

    #[test]
    fn test_lcm_demo_guided_generation() {
        let noise_gen = NoiseGenerator::new(
            NoiseSchedule::Linear {
                beta_start: 0.0001,
                beta_end: 0.02,
            },
            1000,
            Some(42),
        )
        .unwrap();

        let config = LCMConfig {
            num_steps: 4,
            guidance_scale: 10.0, // High guidance for strong conditioning
            ..Default::default()
        };

        let mut sampler = LCMSampler::new(config, noise_gen).unwrap();

        // Mock guided model (unconditional + conditional)
        let guided_model = |x: &Tensor, t: u32| -> crate::Result<(Tensor, Tensor)> {
            let noise_scale = (t as f32 / 1000.0) * 0.1;
            let uncond = x
                .scalar_mul(noise_scale)
                .map_err(|e| crate::DiffusionError::TensorError(e.to_string()))?;
            let cond = x
                .scalar_mul(noise_scale * 1.3)
                .map_err(|e| crate::DiffusionError::TensorError(e.to_string()))?; // Stronger conditioning
            Ok((uncond, cond))
        };

        let input_shape = vec![2, 4, 32, 32];
        let device = Device::cuda(0).unwrap_or(Device::default());
        let noisy_latent = Tensor::randn(&input_shape, &device).unwrap();

        println!("LCM Demo: Starting guided 4-step generation...");
        let result = sampler.sample_guided(&noisy_latent, 900, guided_model);

        assert!(result.is_ok());
        let generated = result.unwrap();
        assert_eq!(generated.dims(), input_shape.as_slice());

        println!("LCM Demo: Guided generation complete");
        println!("LCM Demo: NFE with CFG: {}", sampler.stats().nfe);

        // Should be 8 model calls (2 per step for CFG)
        assert_eq!(sampler.stats().nfe, 8);
    }

    #[test]
    #[ignore = "Pre-existing Metal shader compilation issue"]
    fn test_lcm_demo_consistency_training() {
        let noise_gen =
            NoiseGenerator::new(NoiseSchedule::Cosine { s: 0.008 }, 1000, Some(42)).unwrap();

        let config = LCMConfig {
            training_mode: true,
            consistency_weight: 2.0,
            distillation_weight: 1.0,
            ..Default::default()
        };

        let sampler = LCMSampler::new(config, noise_gen).unwrap();

        // Mock training data
        let batch_size = 4;
        let channels = 3;
        let height = 64;
        let width = 64;

        let device = Device::cuda(0).unwrap_or(Device::default());
        let x0 = Tensor::randn(&vec![batch_size, channels, height, width], &device).unwrap();
        let noise = Tensor::randn(&vec![batch_size, channels, height, width], &device).unwrap();
        let teacher_output =
            Tensor::randn(&vec![batch_size, channels, height, width], &device).unwrap();
        let student_output =
            Tensor::randn(&vec![batch_size, channels, height, width], &device).unwrap();

        // Test consistency loss
        let consistency_loss = sampler.compute_consistency_loss(&x0, &noise, 100, 300);
        assert!(consistency_loss.is_ok());

        let loss_val = consistency_loss.unwrap().to_scalar::<f32>().unwrap();
        println!("LCM Demo: Consistency loss: {:.6}", loss_val);
        assert!(loss_val >= 0.0);

        // Test distillation loss
        let distillation_loss =
            sampler.compute_distillation_loss(&x0, &teacher_output, &student_output, 500);
        assert!(distillation_loss.is_ok());

        let distill_val = distillation_loss.unwrap().to_scalar::<f32>().unwrap();
        println!("LCM Demo: Distillation loss: {:.6}", distill_val);
        assert!(distill_val >= 0.0);

        // Test advanced consistency loss
        let timesteps = vec![50, 150, 250, 350, 450];
        let advanced_loss = sampler.compute_advanced_consistency_loss(&x0, &noise, &timesteps);
        assert!(advanced_loss.is_ok());

        let advanced_val = advanced_loss.unwrap().to_scalar::<f32>().unwrap();
        println!("LCM Demo: Advanced consistency loss: {:.6}", advanced_val);
        assert!(advanced_val >= 0.0);

        println!("LCM Demo: All training losses computed successfully!");
    }
}