tests/test_denoiser.cs
using HumanoidMocapX.Inference;
using HumanoidMocapX;

static class DenoiserTests
{
    internal const int Width = 8, Heads = 2;

    internal static void Run()
    {
        Rotary();
        var weights = Fixture();
        MotionTests.Equal(304, weights.Count);
        const int frames = 3;
        var model = new FlowDenoiser(new FlowWeights(weights), frames, Width, Heads);
        var state = Wave(frames * FlowSampler.Channels, .019f);
        var body = Wave(frames * FlowDenoiser.FeatureChannels, .037f);
        var camera = Wave(frames * FlowDenoiser.CameraChannels, .17f);
        var output = new float[state.Length];
        MotionTests.Throws<InvalidOperationException>(() => model.Predict(state, .3f, output));
        model.SetVideo(body, camera);
        foreach (var time in new[] { 0f, .37f, .95f })
        {
            model.Predict(state, time, output);
            var expected = DenoiserReference.Predict(weights, state.Select(v => (double)v).ToArray(), body, camera, time);
            for (var i = 0; i < output.Length; i++) MotionTests.Near(expected[i], output[i], .00003);
        }
        var first = output.ToArray();
        body.AsSpan().Clear(); camera.AsSpan().Clear();
        model.Predict(state, .95f, output);
        for (var i = 0; i < output.Length; i++) MotionTests.Near(first[i], output[i], 0);
        model.Predict(state, .5f, output);
        var before = GC.GetAllocatedBytesForCurrentThread();
        for (var i = 0; i < 3; i++) model.Predict(state, .5f, output);
        MotionTests.Equal(0L, GC.GetAllocatedBytesForCurrentThread() - before);
        MotionTests.Throws<ArgumentException>(() => model.Predict(state, .5f, state));
        MotionTests.Throws<ArgumentOutOfRangeException>(() => model.Predict(state, float.NaN, output));
        using var cancelled = new CancellationTokenSource(); cancelled.Cancel();
        MotionTests.Throws<OperationCanceledException>(() => model.Predict(state, .5f, output, cancelled.Token));
        MotionTests.Throws<OperationCanceledException>(() => model.SetVideo(body, camera, cancelled.Token));
        MotionTests.Throws<InvalidOperationException>(() => model.Predict(state, .5f, output));
        model.SetVideo(body, camera);
        var expectedState = state.Select(v => (double)v).ToArray();
        for (var step = 0; step < FlowSampler.Steps; step++)
        {
            var time = step / (float)FlowSampler.Steps;
            var prediction = DenoiserReference.Predict(weights, expectedState, body, camera, time);
            var dt = (step + 1) / (float)FlowSampler.Steps - time;
            for (var i = 0; i < expectedState.Length; i++) expectedState[i] += (prediction[i] - expectedState[i]) / Math.Max(1f - time, .05f) * dt;
        }
        FlowSampler.Sample(state, output, model.Predict);
        for (var i = 0; i < state.Length; i++) MotionTests.Near(expectedState[i], state[i], .00005);
        Console.WriteLine("PASS: all eighteen denoiser blocks, conditioning/final projection, sampler integration and zero-allocation prediction");
        GpuFixtures();
    }

    static void GpuFixtures()
    {
        foreach (var (frames, heads) in new[] { (1, 1), (3, 2), (67, 1) })
        {
            var width = 64 * heads;
            var model = new FlowDenoiser(new FlowWeights(Fixture(width, heads)), frames, width, heads);
            var body = Wave(frames * FlowDenoiser.FeatureChannels, .037f);
            var camera = Wave(frames * FlowDenoiser.CameraChannels, .17f);
            var initial = Wave(frames * FlowSampler.Channels, .019f);
            var state = (float[])initial.Clone(); var prediction = new float[state.Length];
            model.SetVideo(body, camera); FlowSampler.Sample(state, prediction, model.Predict);
            MotionTests.Check(!state.AsSpan().SequenceEqual(initial), "GPU-check reference fixture must generate nontrivial motion");
            var first = (float[])state.Clone();
            for (var i = 0; i < body.Length; i++) body[i] *= -.4f;
            for (var i = 0; i < camera.Length; i++) camera[i] *= .7f;
            initial.CopyTo(state, 0);
            model.SetVideo(body, camera); FlowSampler.Sample(state, prediction, model.Predict);
            MotionTests.Check(!state.AsSpan().SequenceEqual(first), "GPU-check reference fixture must respond to changed video conditioning");
        }
        Console.WriteLine("PASS: CPU reference fixtures for GPU comparisons, 64-channel heads and 67-frame attention windows");
    }

    static void Rotary()
    {
        const int frames = 5, width = 8, heads = 2, headWidth = width / heads;
        var q = Wave(frames * width, .3f); var k = Wave(q.Length, .2f); var original = q.ToArray();
        var rotary = new FlowRotary(frames, width, heads);
        rotary.Apply(q, k);
        var factors = new float[frames * headWidth]; rotary.CopyFactors(factors);
        MotionTests.Throws<ArgumentException>(() => rotary.CopyFactors(new float[1]));
        for (var frame = 0; frame < frames; frame++)
            for (var head = 0; head < heads; head++)
                for (var pair = 0; pair < headWidth / 2; pair++)
                {
                    var at = frame * width + head * headWidth + pair * 2;
                    var angle = frame * Math.Pow(10000, -2d * pair / headWidth);
                    MotionTests.Near(Math.Cos(angle), factors[frame * headWidth + pair * 2]);
                    MotionTests.Near(Math.Sin(angle), factors[frame * headWidth + pair * 2 + 1]);
                    var expected = new System.Numerics.Complex(original[at], original[at + 1]) * System.Numerics.Complex.FromPolarCoordinates(1, angle);
                    MotionTests.Near(expected.Real, q[at]); MotionTests.Near(expected.Imaginary, q[at + 1]);
                }
        MotionTests.Throws<ArgumentException>(() => new FlowRotary(5001, width, heads));
        MotionTests.Throws<ArgumentException>(() => new FlowRotary(1, 6, 2));
    }

    internal static float[] Wave(int count, float frequency) => Enumerable.Range(0, count).Select(i => MathF.Sin(i * frequency)).ToArray();

    internal static Dictionary<string, Tensor> Fixture(int width = Width, int heads = Heads)
    {
        var weights = RefinerTests.Fixture(width, heads); int serial = 37;
        var amplitude = .07f * MathF.Sqrt(Width / (float)width);
        void Add(string name, int[] shape, bool norm = false)
        {
            var count = shape.Aggregate(1, (a, b) => a * b); var data = new float[count];
            for (var i = 0; i < count; i++) data[i] = (norm ? 1 : 0) + amplitude * MathF.Sin((serial++ % 503) * .73f);
            weights.Add(name, new Tensor(name, data, shape));
        }
        void Linear(string name, int inputs, int outputs)
        { Add(name + ".weight", new[] { inputs, outputs }); Add(name + ".bias", new[] { outputs }); }
        Linear("input_encoder", FlowSampler.Channels, width);
        Linear("ctxt_encoder.feature", FlowDenoiser.FeatureChannels, width);
        Linear("ctxt_encoder.camera_R", FlowDenoiser.CameraChannels, width);
        Linear("timestep_encoder.blocks.0", width, width); Linear("timestep_encoder.blocks.2", width, width);
        Linear("vtxt_encoder.linears.0", FlowDenoiser.TextChannels, width); Linear("vtxt_encoder.linears.2", width, width);
        Linear("final_layer.adaLN_modulation.linear", width, 2 * width); Linear("final_layer.linear", width, FlowSampler.Channels);
        for (var index = 0; index < FlowDenoiser.DoubleBlocks; index++)
            foreach (var branch in new[] { "motion", "text" })
            {
                var name = "double_blocks." + index + "." + branch;
                Linear(name + "_mod.linear", width, 6 * width); Linear(name + "_qkv", width, 3 * width);
                Linear(name + "_out_proj", width, width); Linear(name + "_mlp.fc1", width, 4 * width); Linear(name + "_mlp.fc2", 4 * width, width);
                Add(name + "_q_norm.weight", new[] { width / heads }, true); Add(name + "_k_norm.weight", new[] { width / heads }, true);
            }
        for (var index = 0; index < FlowDenoiser.SingleBlocks; index++)
        {
            var name = "single_blocks." + index;
            Linear(name + ".modulation.linear", width, 3 * width); Linear(name + ".linear1", width, 7 * width); Linear(name + ".linear2", 5 * width, width);
            Add(name + ".q_norm.weight", new[] { width / heads }, true); Add(name + ".k_norm.weight", new[] { width / heads }, true);
        }
        return weights;
    }
}