diff --git a/reg_with_pic.module b/reg_with_pic.module
index 123b285..98a7df6 100644
--- a/reg_with_pic.module
+++ b/reg_with_pic.module
@@ -61,33 +61,36 @@ function reg_with_pic_form_user_register_form_alter(&$form, &$form_state, $form_
 }
 
 /**
- * Implements hook_user_insert().
+ * Implements hook_entity_presave().
  */
-function reg_with_pic_user_insert(&$edit, $account, $category) {
-  // Set picture if it exists.
-  if (!empty($account->picture)) {
-    $picture = $account->picture;
-    $info = image_get_info($picture->uri);
-    $picture_directory =  file_default_scheme() . '://' . variable_get('user_picture_path', 'pictures');
+function reg_with_pic_entity_presave($account, $type) {
+  // Only handle new users that have an uploaded picture.
+  // user_save() will handle old users.
+  if ($type != 'user' || !$account->is_new || empty($account->picture) || !is_object($account->picture)) {
+    return;
+  }
+  $picture = $account->picture;
+  $info = image_get_info($picture->uri);
+  $picture_directory =  file_default_scheme() . '://' . variable_get('user_picture_path', 'pictures');
+
+  // Set the new uid since it is typically not set at this point. user_save()
+  // will use this uid if set so no problems setting it early here.
+  if (empty($account->uid)) {
+    $account->uid = db_next_id(db_query('SELECT MAX(uid) FROM {users}')->fetchField());
+  }
 
-    // Prepare the pictures directory.
-    file_prepare_directory($picture_directory, FILE_CREATE_DIRECTORY);
-    $destination = file_stream_wrapper_uri_normalize($picture_directory . '/picture-' . $account->uid . '-' . REQUEST_TIME . '.' . $info['extension']);
+  // Prepare the pictures directory.
+  file_prepare_directory($picture_directory, FILE_CREATE_DIRECTORY);
+  $destination = file_stream_wrapper_uri_normalize($picture_directory . '/picture-' . $account->uid . '-' . REQUEST_TIME . '.' . $info['extension']);
 
-    // Move the temporary file into the final location.
-    if ($picture = file_move($picture, $destination, FILE_EXISTS_RENAME)) {
-      $picture->status = FILE_STATUS_PERMANENT;
-      $account->picture = file_save($picture);
-      file_usage_add($picture, 'user', 'user', $account->uid);
-    }
-    // Update user record with picture fid.
-    if (isset($account->picture->fid)) {
-      db_update('users')
-        ->fields(array(
-          'picture' => $account->picture->fid,
-        ))
-        ->condition('uid', $account->uid)
-        ->execute();
-    }
+  // Move the temporary file into the final location.
+  if ($picture = file_move($picture, $destination, FILE_EXISTS_RENAME)) {
+    $picture->status = FILE_STATUS_PERMANENT;
+    $account->picture = file_save($picture);
+    file_usage_add($picture, 'user', 'user', $account->uid);
   }
-}
\ No newline at end of file
+  // Update user record with picture fid.
+  if (isset($account->picture->fid)) {
+    $account->picture = $account->picture->fid;
+  }
+}