diff --git a/autolens/imaging/model/analysis.py b/autolens/imaging/model/analysis.py
index 4d71aee3f..4ca751448 100644
--- a/autolens/imaging/model/analysis.py
+++ b/autolens/imaging/model/analysis.py
@@ -117,7 +117,9 @@ def fit_from(
 
         dataset_model = self.dataset_model_via_instance_from(instance=instance)
 
-        adapt_images = self.adapt_images_via_instance_from(instance=instance)
+        adapt_images = self.adapt_images_via_instance_from(
+            instance=instance, galaxies=tracer.galaxies
+        )
 
         return FitImaging(
             dataset=self.dataset,
diff --git a/autolens/interferometer/model/analysis.py b/autolens/interferometer/model/analysis.py
index 19500dc9f..fc9b4d733 100644
--- a/autolens/interferometer/model/analysis.py
+++ b/autolens/interferometer/model/analysis.py
@@ -180,7 +180,9 @@ def fit_from(self, instance: af.ModelInstance) -> FitInterferometer:
             instance=instance,
         )
 
-        adapt_images = self.adapt_images_via_instance_from(instance=instance)
+        adapt_images = self.adapt_images_via_instance_from(
+            instance=instance, galaxies=tracer.galaxies
+        )
 
         return FitInterferometer(
             dataset=self.dataset,
diff --git a/autolens/lens/to_inversion.py b/autolens/lens/to_inversion.py
index 5e060901e..f1c74da70 100644
--- a/autolens/lens/to_inversion.py
+++ b/autolens/lens/to_inversion.py
@@ -270,24 +270,12 @@ def adapt_galaxy_image_pg_list(self) -> List[List[np.ndarray]]:
                 )
 
                 for galaxy in galaxies_with_pixelization_list:
-                    try:
-                        image = self.adapt_images.galaxy_image_dict[galaxy]
-                    except (AttributeError, KeyError, TypeError):
+                    if self.adapt_images is None:
                         image = None
-
-                    # Bug fix whereby for certain models the galaxy doesnt pair correctly.
-
-                    if image is None and len(galaxies_with_pixelization_list) == 1:
-                        galaxy_list = self.adapt_images.galaxy_image_dict.keys()
-                        galaxy_with_pixelization = [
-                            galaxy
-                            for galaxy in galaxy_list
-                            if galaxy.has(cls=aa.Pixelization)
-                        ][0]
-
-                        image = self.adapt_images.galaxy_image_dict[
-                            galaxy_with_pixelization
-                        ]
+                    else:
+                        image = self.adapt_images.image_for_galaxy(
+                            galaxy, self.tracer.galaxies
+                        )
 
                     plane_image_list.append(image)
 
@@ -330,6 +318,7 @@ def image_plane_mesh_grid_pg_list(self) -> List[List]:
                 adapt_images=self.adapt_images,
                 settings=self.settings,
                 xp=self._xp,
+                path_galaxies=self.tracer.galaxies,
             )
 
             image_plane_mesh_grid_list = to_inversion.image_plane_mesh_grid_list