1 files changed, 20 insertions, 21 deletions

diff --git a/apps/plugins/pictureflow/pictureflow.c b/apps/plugins/pictureflow/pictureflow.c
index 546afad217..8ad62b0238 100644
--- a/apps/plugins/pictureflow/pictureflow.c
+++ b/apps/plugins/pictureflow/pictureflow.c
@@ -27,6 +27,7 @@
 #include <albumart.h>
 #include "lib/read_image.h"
 #include "lib/pluginlib_actions.h"
+#include "lib/pluginlib_exit.h"
 #include "lib/helper.h"
 #include "lib/configfile.h"
 #include "lib/grey.h"
@@ -34,7 +35,7 @@
 #include "lib/feature_wrappers.h"
 #include "lib/buflib.h"
 
-PLUGIN_HEADER
+
 
 /******************************* Globals ***********************************/
 
@@ -2068,9 +2069,8 @@ void update_scroll_animation(void)
 /**
   Cleanup the plugin
 */
-void cleanup(void *parameter)
+void cleanup(void)
 {
-    (void) parameter;
     int i;
 #ifdef HAVE_ADJUSTABLE_CPU_FREQ
     rb->cpu_boost(false);
@@ -2610,19 +2610,6 @@ int main(void)
     cache_version = CACHE_VERSION;
     configfile_save(CONFIG_FILE, config, CONFIG_NUM_ITEMS, CONFIG_VERSION);
 
-
-#ifdef USEGSLIB
-    long grey_buf_used;
-    if (!grey_init(buf, buf_size, GREY_BUFFERED|GREY_ON_COP,
-                   LCD_WIDTH, LCD_HEIGHT, &grey_buf_used))
-    {
-        error_wait("Greylib init failed!");
-        return PLUGIN_ERROR;
-    }
-    grey_setfont(FONT_UI);
-    buf_size -= grey_buf_used;
-    buf = (void*)(grey_buf_used + (char*)buf);
-#endif
     buflib_init(&buf_ctx, (void *)buf, buf_size);
 
     if (!(empty_slide_hid = read_pfraw(EMPTY_SLIDE, 0)))
@@ -2841,9 +2828,7 @@ int main(void)
             }
             break;
         default:
-            if (rb->default_event_handler_ex(button, cleanup, NULL)
-                == SYS_USB_CONNECTED)
-                return PLUGIN_USB_CONNECTED;
+            exit_on_usb(button);
             break;
         }
     }
@@ -2855,6 +2840,7 @@ enum plugin_status plugin_start(const void *parameter)
 {
     int ret, i;
     (void) parameter;
+    atexit(cleanup);
 
     FOR_NB_SCREENS(i)
         rb->viewportmanager_theme_enable(i, false, NULL);
@@ -2873,6 +2859,21 @@ enum plugin_status plugin_start(const void *parameter)
     }
 #endif
 #endif
+
+#ifdef USEGSLIB
+    long grey_buf_used;
+    if (!grey_init(buf, buf_size, GREY_BUFFERED|GREY_ON_COP,
+                   LCD_WIDTH, LCD_HEIGHT, &grey_buf_used))
+    {
+        error_wait("Greylib init failed!");
+        return PLUGIN_ERROR;
+    }
+    grey_setfont(FONT_UI);
+    buf_size -= grey_buf_used;
+    buf = (void*)(grey_buf_used + (char*)buf);
+#endif
+
+    atexit(cleanup);
     ret = main();
     if ( ret == PLUGIN_OK || ret == PLUGIN_GOTO_WPS) {
         if (configfile_save(CONFIG_FILE, config, CONFIG_NUM_ITEMS,
@@ -2882,7 +2883,5 @@ enum plugin_status plugin_start(const void *parameter)
             ret = PLUGIN_ERROR;
         }
     }
-
-    cleanup(NULL);
     return ret;
 }