diff --git a/notebooks/AGIPD/Characterize_AGIPD_Gain_FlatFields_Summary.ipynb b/notebooks/AGIPD/Characterize_AGIPD_Gain_FlatFields_Summary.ipynb
index 31bbb92df0ac33675e97b2b98e667ab09228a78c..fca571903e01ac972a3a4a4c1f03c44ead950819 100644
--- a/notebooks/AGIPD/Characterize_AGIPD_Gain_FlatFields_Summary.ipynb
+++ b/notebooks/AGIPD/Characterize_AGIPD_Gain_FlatFields_Summary.ipynb
@@ -29,7 +29,7 @@
     "use_dir_creation_date = True # use the creation data of the input dir for database queries\n",
     "cal_db_interface = \"tcp://max-exfl-cal001:8015#8045\" # the database interface to use\n",
     "cal_db_timeout = 30000 # in milli seconds\n",
-    "local_output = False # output constants locally\n",
+    "local_output = True # output constants locally\n",
     "db_output = False # output constants to database\n",
     "\n",
     "# Fit parameters\n",
@@ -561,17 +561,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def get_trains_data(run_folder, source, include, tid=None):\n",
+    "def get_trains_data(run_data, source, include, tid=None):\n",
     "    \"\"\"\n",
     "    Load single train for all module\n",
     "    \n",
-    "    :param run_folder: Path to folder with data\n",
+    "    :param run_data: Corrcted data to be loaded from\n",
     "    :param source: Data source to be loaded\n",
     "    :param include: Inset of file name to be considered \n",
     "    :param tid: Train Id to be loaded. First train is considered if None is given\n",
     "    \n",
     "    \"\"\"\n",
-    "    run_data = RunDirectory(run_folder, include)\n",
     "    if tid:\n",
     "        tid, data = run_data.select('*/DET/*', source).train_from_id(tid)\n",
     "        return tid, stack_detector_data(data, source, modules=nmods)\n",
@@ -580,25 +579,12 @@
     "            return tid, stack_detector_data(data, source, modules=nmods)\n",
     "    return None, None\n",
     "\n",
-    "\n",
     "include = '*S00000*'\n",
-    "tid, orig = get_trains_data(f'{proc_folder}/r{run:04d}/', 'image.data', include)\n",
+    "run_data = RunDirectory(f'{proc_folder}/r{run:04d}/', include)\n",
+    "tid, orig = get_trains_data(run_data, 'image.data', include)\n",
     "orig = orig[cell_range[0]:cell_range[1], ...]"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Get cell Ids from 1st train\n",
-    "run_data = RunDirectory(run_folder, include)\n",
-    "sel = run_data.select(f'{karabo_id}/DET/0CH0:xtdf', 'image.cellId')\n",
-    "tid0, cell = sel.train_from_index(0)\n",
-    "stacked_cells = stack_detector_data(cell, 'image.cellId')[0, :, 0]"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -610,9 +596,14 @@
     "corrections = const_data['slopesFF'] # (16,shape[0],512,128) shape[0]= cell_range[1]-cell_range[0] /\n",
     "corrections = np.moveaxis(corrections, 1, 0) # (shape[0],16,512,128)\n",
     "rel_corr = corrections/np.nanmedian(corrections)\n",
+    "\n",
     "# this is needed if LitFrame is enabled in DAQ to avoid shape mismatch \n",
     "# and correction of the right cells\n",
     "if np.diff(cell_range)[0] == mem_cells:\n",
+    "    sel = run_data.select(f'{karabo_id}/DET/0CH0:xtdf', 'image.cellId')\n",
+    "    _, cell = sel.train_from_index(0)\n",
+    "    stacked_cells = stack_detector_data(cell, 'image.cellId')[:, 0]\n",
+    "    \n",
     "    rel_corr = rel_corr[stacked_cells[0]:stacked_cells[-1]+1]\n",
     "\n",
     "corrected = orig / rel_corr"