diff --git a/src/encmain.c b/src/encmain.c
index 97d31469..1f4b910b 100644
--- a/src/encmain.c
+++ b/src/encmain.c
@@ -224,7 +224,10 @@ int main(int argc, char *argv[])
       }
 
       image_t *img_out = NULL;
-      api->encoder_encode(enc, img_in, &img_out, &output_stream);
+      if (1 != api->encoder_encode(enc, img_in, &img_out, &output_stream)) {
+        fprintf(stderr, "Failed to encode image.\n");
+        goto exit_failure;
+      }
 
       if (img_out != NULL) {
         state = &enc->states[enc->cur_state_num];
@@ -245,14 +248,11 @@ int main(int argc, char *argv[])
     
     //Compute stats for the remaining encoders
     {
-      int first_enc = current_encoder_state;
-      do {
-        double frame_psnr[3] = { 0.0, 0.0, 0.0 };
-        image_t *img_out = NULL;
-        api->encoder_encode(enc, NULL, &img_out, &output_stream);
-
+      image_t *img_out = NULL;
+      while (1 == api->encoder_encode(enc, NULL, &img_out, &output_stream)) {
         if (img_out != NULL) {
-          encoder_state_t *state = &enc->states[current_encoder_state];
+          double frame_psnr[3] = { 0.0, 0.0, 0.0 };
+          encoder_state_t *state = &enc->states[enc->cur_state_num];
 
           encoder_compute_stats(state, recout, frame_psnr, &bitstream_length);
           print_frame_info(state, frame_psnr);
@@ -262,10 +262,9 @@ int main(int argc, char *argv[])
           psnr_sum[2] += frame_psnr[2];
 
           image_free(img_out);
+          img_out = NULL;
         }
-
-        current_encoder_state = (current_encoder_state + 1) % (encoder->owf + 1);
-      } while (current_encoder_state != first_enc);
+      }
     }
     
     GET_TIME(&encoding_end_real_time);
diff --git a/src/kvazaar.c b/src/kvazaar.c
index ac374e2b..c4a73cbe 100644
--- a/src/kvazaar.c
+++ b/src/kvazaar.c
@@ -73,6 +73,7 @@ static kvz_encoder * kvazaar_open(config_t *cfg)
   encoder->num_encoder_states = cfg->owf + 1;
   encoder->cur_state_num = 0;
   encoder->frames_started = 0;
+  encoder->frames_done = 0;
   encoder->states = MALLOC(encoder_state_t, encoder->num_encoder_states);
   if (!encoder->states) {
     goto kvazaar_open_failure;
@@ -122,6 +123,11 @@ static int kvazaar_encode(kvz_encoder *enc, kvz_picture *img_in, kvz_picture **i
     encode_one_frame(state);
   }
 
+  // If we have finished encoding as many frames as we have started, we are done.
+  if (enc->frames_done == enc->frames_started) {
+    return 0;
+  }
+
   enc->cur_state_num = (enc->cur_state_num + 1) % (enc->num_encoder_states);
   encoder_state_t *state = &enc->states[enc->cur_state_num];
 
@@ -136,9 +142,11 @@ static int kvazaar_encode(kvz_encoder *enc, kvz_picture *img_in, kvz_picture **i
     }
 
     *img_out = image_make_subimage(state->tile->frame->rec, 0, 0, state->tile->frame->width, state->tile->frame->height);
+
+    enc->frames_done += 1;
   }
 
-  return 0;
+  return 1;
 }
 
 kvz_api kvz_8bit_api = {
diff --git a/src/kvazaar.h b/src/kvazaar.h
index 7f0ea1f5..8de20ad7 100644
--- a/src/kvazaar.h
+++ b/src/kvazaar.h
@@ -82,6 +82,7 @@ typedef struct kvz_encoder {
   unsigned num_encoder_states;
   unsigned cur_state_num;
   unsigned frames_started;
+  unsigned frames_done;
 
   size_t bitstream_length;
 } kvz_encoder;