Integrate RES audio updates on the audio sigma schedule

h3_dit.c

@@ -2518,24 +2518,14 @@ int h3_dit_denoise(h3_dit *dit, float *video_latent, float *audio_latent,
         ok = h3_dit_forward(dit, step, video_latent, audio_latent,
                             video_velocity, audio_velocity,
                             error, error_size);
-        float sigma = dit->sigmas.video[step];
-        float timestep = 1.0f - sigma;
-        float sigma_from_timestep = 1.0f - timestep;
-        float audio_slope = (float)h3_time_shift_slope(
-            sigma, H3_VIDEO_SIGMA_SHIFT, H3_AUDIO_SIGMA_SHIFT);
         if (ok) {
-            for (size_t index = 0; index < video_count; index++)
-                video_denoised[index] = video_latent[index] +
-                    sigma_from_timestep * video_velocity[index];
-            for (size_t index = 0; index < audio_count; index++)
-                audio_denoised[index] = audio_latent[index] +
-                    sigma_from_timestep * audio_velocity[index] * audio_slope;
-            ok = h3_res_step(video_next, video_latent, video_denoised,
-                             step ? old_video : NULL, video_count,
-                             dit->sigmas.video, step, dit->sigmas.steps) &&
-                 h3_res_step(audio_next, audio_latent, audio_denoised,
-                             step ? old_audio : NULL, audio_count,
-                             dit->sigmas.video, step, dit->sigmas.steps);
+            ok = h3_res_velocity_step(
+                     video_next, video_latent, video_velocity, video_denoised,
+                     step ? old_video : NULL, video_count, dit->sigmas.video,
+                     step, dit->sigmas.steps) &&
+                 h3_audio_res_velocity_step(
+                     audio_next, audio_latent, audio_velocity, audio_denoised,
+                     step ? old_audio : NULL, audio_count, &dit->sigmas, step);
             if (!ok) fail(error, error_size, "RES solver rejected step %d", step);
         }
         if (ok) {

h3_host.c

@@ -636,6 +636,29 @@ int h3_res_step(float *output, const float *sample, const float *denoised,
     return 1;
 }
 
+int h3_res_velocity_step(float *output, const float *sample,
+                         const float *velocity, float *denoised,
+                         const float *old_denoised, size_t count,
+                         const float *sigmas, int step, int total_steps) {
+    if (!output || !sample || !velocity || !denoised || !sigmas || step < 0 ||
+        total_steps < 1 || step >= total_steps) return 0;
+    float sigma = sigmas[step];
+    for (size_t index = 0; index < count; index++)
+        denoised[index] = sample[index] + sigma * velocity[index];
+    return h3_res_step(output, sample, denoised, old_denoised, count,
+                       sigmas, step, total_steps);
+}
+
+int h3_audio_res_velocity_step(float *output, const float *sample,
+                               const float *velocity, float *denoised,
+                               const float *old_denoised, size_t count,
+                               const h3_sigma_schedule *schedule, int step) {
+    if (!schedule) return 0;
+    return h3_res_velocity_step(output, sample, velocity, denoised,
+                                old_denoised, count, schedule->audio,
+                                step, schedule->steps);
+}
+
 int h3_euler_velocity_step(float *sample, const float *velocity, size_t count,
                            float sigma, float sigma_next) {
     if (!sample || !velocity || !isfinite(sigma) || !isfinite(sigma_next) ||

h3_host.h

@@ -136,6 +136,17 @@ int h3_resize_rgb24_high_quality(const uint8_t *input, int frames,
 int h3_res_step(float *output, const float *sample, const float *denoised,
                 const float *old_denoised, size_t count,
                 const float *sigmas, int step, int total_steps);
+/* Convert data-ward velocity with the same sigma grid used by RES. The caller
+ * retains denoised for the following multistep update. */
+int h3_res_velocity_step(float *output, const float *sample,
+                         const float *velocity, float *denoised,
+                         const float *old_denoised, size_t count,
+                         const float *sigmas, int step, int total_steps);
+/* Apply an audio velocity on the audio half of an H3 sigma schedule. */
+int h3_audio_res_velocity_step(float *output, const float *sample,
+                               const float *velocity, float *denoised,
+                               const float *old_denoised, size_t count,
+                               const h3_sigma_schedule *schedule, int step);
 int h3_euler_velocity_step(float *sample, const float *velocity, size_t count,
                            float sigma, float sigma_next);
 

tests/test_h3.c

@@ -318,6 +318,30 @@ static void test_rng_and_solver(void) {
     CHECK(close_enough(output[0], 0.25, 1e-7));
     CHECK(h3_res_step(output, sample, denoised, old, 1, sigmas, 1, 3));
     CHECK(isfinite(output[0]));
+
+    /* Constant velocity has the exact endpoint x(0) = x(1) + velocity.
+     * On the native audio sigma grid, x + sigma * velocity stays constant. */
+    const int res_steps[] = {4, 7, 20};
+    for (size_t case_index = 0;
+         case_index < sizeof(res_steps) / sizeof(*res_steps); case_index++) {
+        h3_sigma_schedule schedule;
+        int steps = res_steps[case_index];
+        CHECK(h3_serving_schedule_build(steps, &schedule));
+        float res_sample = 1.0f;
+        float res_velocity = 0.5f;
+        float res_denoised = 0.0f;
+        float old_res_denoised = 0.0f;
+        float res_output = 0.0f;
+        for (int step = 0; step < steps; step++) {
+            CHECK(h3_audio_res_velocity_step(
+                &res_output, &res_sample, &res_velocity, &res_denoised,
+                step ? &old_res_denoised : NULL, 1, &schedule, step));
+            old_res_denoised = res_denoised;
+            res_sample = res_output;
+        }
+        CHECK(close_enough(res_sample, 1.5, 1e-6));
+    }
+
     float velocity[] = {2.0f, -4.0f};
     float euler[] = {1.0f, 3.0f};
     CHECK(h3_euler_velocity_step(euler, velocity, 2, 0.75f, 0.25f));