diff --git a/autofit/aggregator/summary/aggregate_fits.py b/autofit/aggregator/summary/aggregate_fits.py
index e7d4ea209..14fce2b44 100644
--- a/autofit/aggregator/summary/aggregate_fits.py
+++ b/autofit/aggregator/summary/aggregate_fits.py
@@ -35,7 +35,7 @@ class FITSFit(Enum):
 
 
 class AggregateFITS:
-    def __init__(self, aggregator: Aggregator):
+    def __init__(self, aggregator: Union[Aggregator, List[SearchOutput]]):
         """
         A class for extracting fits files from the aggregator.
 
@@ -80,7 +80,7 @@ def _hdus(
             )
         return row
 
-    def extract_fits(self, hdus: List[Enum]) -> List[fits.HDUList]:
+    def extract_fits(self, hdus: List[Enum]) -> fits.HDUList:
         """
         Extract the HDUs from the fits files for every search in the aggregator.
 
@@ -101,7 +101,7 @@ def extract_fits(self, hdus: List[Enum]) -> List[fits.HDUList]:
 
         return fits.HDUList(output)
 
-    def extract_csv(self, filename : str) -> List[Dict]:
+    def extract_csv(self, filename: str) -> List[Dict]:
         """
         Extract .csv files which store imaging results that are typically on irregular grids and thus don't suit
         a .fits file.
diff --git a/autofit/aggregator/summary/aggregate_images.py b/autofit/aggregator/summary/aggregate_images.py
index ee233699f..402ee8cb6 100644
--- a/autofit/aggregator/summary/aggregate_images.py
+++ b/autofit/aggregator/summary/aggregate_images.py
@@ -103,7 +103,7 @@ def image_at_coordinates(
 class AggregateImages:
     def __init__(
         self,
-        aggregator: Aggregator,
+        aggregator: Union[Aggregator, List[SearchOutput]],
     ):
         """
         Extracts images from the aggregator and combines them into one