1 files changed, 19 insertions, 17 deletions

diff --git a/third_party/aom/aom_dsp/flow_estimation/disflow.c b/third_party/aom/aom_dsp/flow_estimation/disflow.c
index 82b531c729..f511a6eb49 100644
--- a/third_party/aom/aom_dsp/flow_estimation/disflow.c
+++ b/third_party/aom/aom_dsp/flow_estimation/disflow.c
@@ -603,9 +603,9 @@ static void upscale_flow_component(double *flow, int cur_width, int cur_height,
 
 // make sure flow_u and flow_v start at 0
 static bool compute_flow_field(const ImagePyramid *src_pyr,
-                               const ImagePyramid *ref_pyr, FlowField *flow) {
+                               const ImagePyramid *ref_pyr, int n_levels,
+                               FlowField *flow) {
   bool mem_status = true;
-  assert(src_pyr->n_levels == ref_pyr->n_levels);
 
   double *flow_u = flow->u;
   double *flow_v = flow->v;
@@ -613,7 +613,7 @@ static bool compute_flow_field(const ImagePyramid *src_pyr,
   double *tmpbuf0;
   double *tmpbuf;
 
-  if (src_pyr->n_levels < 2) {
+  if (n_levels < 2) {
     // tmpbuf not needed
     tmpbuf0 = NULL;
     tmpbuf = NULL;
@@ -639,7 +639,7 @@ static bool compute_flow_field(const ImagePyramid *src_pyr,
   // correspondences by interpolating this flow field, and then refine the
   // correspondences themselves. This is both faster and gives better output
   // compared to refining the flow field at level 0 and then interpolating.
-  for (int level = src_pyr->n_levels - 1; level >= 1; --level) {
+  for (int level = n_levels - 1; level >= 1; --level) {
     const PyramidLayer *cur_layer = &src_pyr->layers[level];
     const int cur_width = cur_layer->width;
     const int cur_height = cur_layer->height;
@@ -762,29 +762,31 @@ static void free_flow_field(FlowField *flow) {
 // Following the convention in flow_estimation.h, the flow vectors are computed
 // at fixed points in `src` and point to the corresponding locations in `ref`,
 // regardless of the temporal ordering of the frames.
-bool av1_compute_global_motion_disflow(TransformationType type,
-                                       YV12_BUFFER_CONFIG *src,
-                                       YV12_BUFFER_CONFIG *ref, int bit_depth,
-                                       MotionModel *motion_models,
-                                       int num_motion_models,
-                                       bool *mem_alloc_failed) {
+bool av1_compute_global_motion_disflow(
+    TransformationType type, YV12_BUFFER_CONFIG *src, YV12_BUFFER_CONFIG *ref,
+    int bit_depth, int downsample_level, MotionModel *motion_models,
+    int num_motion_models, bool *mem_alloc_failed) {
   // Precompute information we will need about each frame
   ImagePyramid *src_pyramid = src->y_pyramid;
   CornerList *src_corners = src->corners;
   ImagePyramid *ref_pyramid = ref->y_pyramid;
-  if (!aom_compute_pyramid(src, bit_depth, src_pyramid)) {
-    *mem_alloc_failed = true;
-    return false;
-  }
-  if (!av1_compute_corner_list(src_pyramid, src_corners)) {
+
+  const int src_layers =
+      aom_compute_pyramid(src, bit_depth, DISFLOW_PYRAMID_LEVELS, src_pyramid);
+  const int ref_layers =
+      aom_compute_pyramid(ref, bit_depth, DISFLOW_PYRAMID_LEVELS, ref_pyramid);
+
+  if (src_layers < 0 || ref_layers < 0) {
     *mem_alloc_failed = true;
     return false;
   }
-  if (!aom_compute_pyramid(ref, bit_depth, ref_pyramid)) {
+  if (!av1_compute_corner_list(src, bit_depth, downsample_level, src_corners)) {
     *mem_alloc_failed = true;
     return false;
   }
 
+  assert(src_layers == ref_layers);
+
   const int src_width = src_pyramid->layers[0].width;
   const int src_height = src_pyramid->layers[0].height;
   assert(ref_pyramid->layers[0].width == src_width);
@@ -796,7 +798,7 @@ bool av1_compute_global_motion_disflow(TransformationType type,
     return false;
   }
 
-  if (!compute_flow_field(src_pyramid, ref_pyramid, flow)) {
+  if (!compute_flow_field(src_pyramid, ref_pyramid, src_layers, flow)) {
     *mem_alloc_failed = true;
     free_flow_field(flow);
     return false;