diff --git a/karabo_devices/shared_mem_processor/src/onlinemid_karabo/processor.py b/karabo_devices/shared_mem_processor/src/onlinemid_karabo/processor.py
index 5c2cef1a1eb1724f6771bf3d791ab7828e0bf2b5..598be39de2e3add4a064541979f0976eb6483415 100644
--- a/karabo_devices/shared_mem_processor/src/onlinemid_karabo/processor.py
+++ b/karabo_devices/shared_mem_processor/src/onlinemid_karabo/processor.py
@@ -183,8 +183,9 @@ class OnlineMIDProcessor(TrainMatcher.TrainMatcher):
         
         
 
-        schema_update = self.onlineImagePlotter.get_schema_update("plots", 'agipd_data_plot', AgipdModulePlotter)
-        self.updateSchema(schema_update)
+        
+        self.updateSchema(self.onlineImagePlotter.get_schema_update("plots", 'first_bunch', AgipdModulePlotter))
+        self.updateSchema(self.onlineImagePlotter.get_schema_update("plots", 'gain_stage', GainStagePlotter))
 
         self.assembled_output = self.signalSlotable.getOutputChannel("plots")
         
@@ -377,30 +378,10 @@ class OnlineMIDProcessor(TrainMatcher.TrainMatcher):
 
         plot_hash = Hash()
 
-        plot_hash = self.onlineImagePlotter.add_to_hash(plot_hash, 'agipd_data_plot', AgipdModulePlotter, train_id, detector_data_list[0][1][0])
-        
-        
-
+        plot_hash = self.onlineImagePlotter.add_to_hash(plot_hash, 'agipd_data_plot', train_id, detector_data_list[0][1][0])
+        plot_hash = self.onlineImagePlotter.add_to_hash(plot_hash, 'gain_stage', train_id, numPixelPerGainStage)
 
 
-#        if numPixelPerGainStage is not None:
-#            plot_hash = self.onlineImagePlotter.add_to_hash(plot_hash, 'data.accumulating_lit_bunch', GainStagePlotter, train_id, numPixelPerGainStage)
-        
-            
-
-            
-#            plot_hash = self.onlineImagePlotter.add_to_hash(plot_hash, 'data.test_hash', GainStagePlotter, train_id, numPixelPerGainStage)
-
-
-        
-
-        # time_to_handle_list_0 = time.perf_counter()
-        
-        # self.online_processing_memory, online_output_dict = onlinexpcsutils.compute_xpcs_from_detector_data_list(
-        #     detector_data_list, 
-        #     xpcs_processing_memory=self.online_processing_memory,
-        #     xpcs_settings=self.xpcs_settings)       
-        # print("time_to_handle_list 0", time.perf_counter() - time_to_handle_list_0)
 
         '''        
         roi = self.xpcs_settings['qbin']
diff --git a/python_package/src/onlinemid/utils/image_plotting.py b/python_package/src/onlinemid/utils/image_plotting.py
index c50aa7f64a296945d98caeb9a0a62a6b82c91a9d..e9df14a8413da1093cf4ebfd94087cd0a7fdea8c 100644
--- a/python_package/src/onlinemid/utils/image_plotting.py
+++ b/python_package/src/onlinemid/utils/image_plotting.py
@@ -143,6 +143,7 @@ class DataToPlotToHash():
 
         )
 
+        
         schema_update = Schema()
         (
             OUTPUT_CHANNEL(schema_update)
@@ -152,16 +153,12 @@ class DataToPlotToHash():
         )
         return schema_update
         
-    def add_to_hash(self, hash, hash_key, Plotter, *args):
+    def add_to_hash(self, hash, hash_key, *args):
         
         if hash_key not in self.hash_keys_and_plotter:
             print(f"key {hash_key} does not exist, call get_schema_update function in initialization() part of this karabo device!")
             return hash
         
-        if not isinstance(self.hash_keys_and_plotter[hash_key]['plotter'], Plotter):
-            print("wrong type, aborting!")
-            return hash
-                        
         # plotter is initialized, give plotter data
         current_img = self.hash_keys_and_plotter[hash_key]['plotter'].plot(*args)