diff --git a/src/components/LazyMessage.tsx b/src/components/LazyMessage.tsx
index 7f11126..2a6364f 100644
--- a/src/components/LazyMessage.tsx
+++ b/src/components/LazyMessage.tsx
@@ -201,7 +201,7 @@ const LazyMessage = observer(
               <button
                 key={imageUrl}
                 className="h-56 place-content-center overflow-hidden rounded-md border border-base-content/30 bg-base-content/30"
-                onClick={() => lightboxStore.setLightboxMessageById(message.uniqId, imageUrl)}
+                onClick={() => lightboxStore.setLightboxMessageById(baseMessage.uniqId, imageUrl)}
               >
                 <CachedImage
                   className="max-h-56 min-w-20 max-w-56 rounded-md object-contain object-center"
diff --git a/src/features/lightbox/LightboxStore.ts b/src/features/lightbox/LightboxStore.ts
index 15c6ed1..1495b19 100644
--- a/src/features/lightbox/LightboxStore.ts
+++ b/src/features/lightbox/LightboxStore.ts
@@ -25,11 +25,13 @@ export const LightboxStore = types
     get lightboxSlides() {
       if (!this.lightboxMessage) return []
 
-      const lightBoxSources: Array<Slide & { uniqId: string }> = []
+      const lightBoxSources: Array<Slide & { baseUniqId: string  }> = []
 
       let userPrompt: string | undefined
 
-      this.chat!.messages.forEach(message => {
+      this.chat!.messages.forEach(baseMessage => {
+        const message = baseMessage.selectedVariation
+
         if (message.fromBot === false) {
           userPrompt = message.content
         }
@@ -38,7 +40,7 @@ export const LightboxStore = types
           lightBoxSources.push({
             description: userPrompt + ` (${index + 1}/${message.imageUrls.length})`,
             src: imageUrl,
-            uniqId: message.uniqId,
+            baseUniqId: baseMessage.uniqId,
           })
         })
       })
@@ -53,8 +55,8 @@ export const LightboxStore = types
     },
   }))
   .actions(self => ({
-    setLightboxMessageById(uniqId: string, imageUrl: string) {
-      self.messageId = uniqId
+    setLightboxMessageById(baseUniqId: string, imageUrl: string) {
+      self.messageId = baseUniqId
       self.imageUrl = imageUrl
     },
 
diff --git a/src/features/lightbox/components/LazyLightbox.tsx b/src/features/lightbox/components/LazyLightbox.tsx
index 5cda077..e5e5d96 100644
--- a/src/features/lightbox/components/LazyLightbox.tsx
+++ b/src/features/lightbox/components/LazyLightbox.tsx
@@ -15,13 +15,13 @@ import 'yet-another-react-lightbox/plugins/thumbnails.css'
 import 'yet-another-react-lightbox/plugins/captions.css'
 
 const LazyLightbox = observer(() => {
-  const [slides, setSlides] = useState<Array<Slide & { uniqId: string }>>([])
+  const [slides, setSlides] = useState<Array<Slide & { baseUniqId: string }>>([])
   const { chat, lightboxSlides, imageUrlIndex } = lightboxStore
 
   const getAllSlideImages = async () => {
     if (!chat) return
 
-    const slides: Array<Slide & { uniqId: string }> = []
+    const slides: Array<Slide & { baseUniqId: string }> = []
 
     for (const slide of lightboxSlides) {
       const src = await CachedStorage.get(slide.src)
@@ -62,9 +62,9 @@ const LazyLightbox = observer(() => {
       carousel={{ finite: true }}
       on={{
         view: ({ index }) => {
-          const { src, uniqId } = lightboxSlides[index]
+          const { src, baseUniqId } = lightboxSlides[index]
 
-          return lightboxStore.setLightboxMessageById(uniqId, src)
+          return lightboxStore.setLightboxMessageById(baseUniqId, src)
         },
       }}
       slides={slides}