diff --git a/Signal/src/ViewControllers/ConversationView/ConversationViewController.m b/Signal/src/ViewControllers/ConversationView/ConversationViewController.m
index 88deca352..1652d4b34 100644
--- a/Signal/src/ViewControllers/ConversationView/ConversationViewController.m
+++ b/Signal/src/ViewControllers/ConversationView/ConversationViewController.m
@@ -2048,18 +2048,12 @@ typedef enum : NSUInteger {
         [self becomeFirstResponder];
     }
 
-    if (![viewItem.interaction isKindOfClass:[TSMessage class]]) {
-        OWSFailDebug(@"Unexpected viewItem.interaction");
-        return;
-    }
-    TSMessage *mediaMessage = (TSMessage *)viewItem.interaction;
-
     MediaGallery *mediaGallery =
         [[MediaGallery alloc] initWithThread:self.thread
                         uiDatabaseConnection:self.uiDatabaseConnection
                                      options:MediaGalleryOptionSliderEnabled | MediaGalleryOptionShowAllMediaButton];
 
-    [mediaGallery presentDetailViewFromViewController:self mediaMessage:mediaMessage replacingView:imageView];
+    [mediaGallery presentDetailViewFromViewController:self mediaAttachment:attachmentStream replacingView:imageView];
 }
 
 - (void)didTapVideoViewItem:(id<ConversationViewItem>)viewItem
@@ -2077,18 +2071,12 @@ typedef enum : NSUInteger {
         [self becomeFirstResponder];
     }
 
-    if (![viewItem.interaction isKindOfClass:[TSMessage class]]) {
-        OWSFailDebug(@"Unexpected viewItem.interaction");
-        return;
-    }
-    TSMessage *mediaMessage = (TSMessage *)viewItem.interaction;
-
     MediaGallery *mediaGallery =
         [[MediaGallery alloc] initWithThread:self.thread
                         uiDatabaseConnection:self.uiDatabaseConnection
                                      options:MediaGalleryOptionSliderEnabled | MediaGalleryOptionShowAllMediaButton];
 
-    [mediaGallery presentDetailViewFromViewController:self mediaMessage:mediaMessage replacingView:imageView];
+    [mediaGallery presentDetailViewFromViewController:self mediaAttachment:attachmentStream replacingView:imageView];
 }
 
 - (void)didTapAudioViewItem:(id<ConversationViewItem>)viewItem attachmentStream:(TSAttachmentStream *)attachmentStream
diff --git a/Signal/src/ViewControllers/MediaDetailViewController.m b/Signal/src/ViewControllers/MediaDetailViewController.m
index 4502c0012..9d04aa9fb 100644
--- a/Signal/src/ViewControllers/MediaDetailViewController.m
+++ b/Signal/src/ViewControllers/MediaDetailViewController.m
@@ -350,6 +350,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)didPressShare:(id)sender
 {
+    OWSFailDebug(@"TODO: support sharing individual attachment, not viewItem");
+
     OWSLogInfo(@"didPressShare");
     if (!self.viewItem) {
         OWSFailDebug(@"share should only be available when a viewItem is present");
@@ -361,6 +363,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)didPressDelete:(id)sender
 {
+    OWSFailDebug(@"TODO: support sharing individual attachment, not viewItem");
+
     OWSLogInfo(@"didPressDelete");
     if (!self.viewItem) {
         OWSFailDebug(@"delete should only be available when a viewItem is present");
diff --git a/Signal/src/ViewControllers/MediaGalleryViewController.swift b/Signal/src/ViewControllers/MediaGalleryViewController.swift
index 2319f74a2..a831392ff 100644
--- a/Signal/src/ViewControllers/MediaGalleryViewController.swift
+++ b/Signal/src/ViewControllers/MediaGalleryViewController.swift
@@ -39,13 +39,13 @@ public class MediaGalleryItem: Equatable, Hashable {
     // MARK: Equatable
 
     public static func == (lhs: MediaGalleryItem, rhs: MediaGalleryItem) -> Bool {
-        return lhs.message.uniqueId == rhs.message.uniqueId
+        return lhs.attachmentStream.uniqueId == rhs.attachmentStream.uniqueId
     }
 
     // MARK: Hashable
 
     public var hashValue: Int {
-        return message.hashValue
+        return attachmentStream.hashValue
     }
 }
 
@@ -281,10 +281,10 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
     private var originRect: CGRect?
 
     @objc
-    public func presentDetailView(fromViewController: UIViewController, mediaMessage: TSMessage, replacingView: UIView) {
+    public func presentDetailView(fromViewController: UIViewController, mediaAttachment: TSAttachment, replacingView: UIView) {
         var galleryItem: MediaGalleryItem?
         uiDatabaseConnection.read { transaction in
-            galleryItem = self.buildGalleryItem(message: mediaMessage, transaction: transaction)!
+            galleryItem = self.buildGalleryItem(attachment: mediaAttachment, transaction: transaction)!
         }
 
         guard let initialDetailItem = galleryItem else {
@@ -414,8 +414,8 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
     func pushTileView(fromNavController: OWSNavigationController) {
         var mostRecentItem: MediaGalleryItem?
         self.uiDatabaseConnection.read { transaction in
-            if let message = self.mediaGalleryFinder.mostRecentMediaMessage(transaction: transaction) {
-                mostRecentItem = self.buildGalleryItem(message: message, transaction: transaction)
+            if let attachment = self.mediaGalleryFinder.mostRecentMediaAttachment(transaction: transaction) {
+                mostRecentItem = self.buildGalleryItem(attachment: attachment, transaction: transaction)
             }
         }
 
@@ -633,9 +633,13 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
     var hasFetchedOldest = false
     var hasFetchedMostRecent = false
 
-    func buildGalleryItem(message: TSMessage, transaction: YapDatabaseReadTransaction) -> MediaGalleryItem? {
-        // TODO: Support multi-image messages.
-        guard let attachmentStream = message.attachments(with: transaction).first as? TSAttachmentStream else {
+    func buildGalleryItem(attachment: TSAttachment, transaction: YapDatabaseReadTransaction) -> MediaGalleryItem? {
+        guard let attachmentStream = attachment as? TSAttachmentStream else {
+            owsFailDebug("gallery doesn't yet support showing undownloaded attachments")
+            return nil
+        }
+
+        guard let message = attachmentStream.fetchAlbumMessage(with: transaction) else {
             owsFailDebug("attachment was unexpectedly empty")
             return nil
         }
@@ -662,7 +666,7 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
         Bench(title: "fetching gallery items") {
             self.uiDatabaseConnection.read { transaction in
 
-                let initialIndex: Int = Int(self.mediaGalleryFinder.mediaIndex(message: item.message, transaction: transaction))
+                let initialIndex: Int = Int(self.mediaGalleryFinder.mediaIndex(attachment: item.attachmentStream, transaction: transaction))
                 let mediaCount: Int = Int(self.mediaGalleryFinder.mediaCount(transaction: transaction))
 
                 let requestRange: Range<Int> = { () -> Range<Int> in
@@ -711,14 +715,14 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
 
                 Logger.debug("fetching set: \(unfetchedSet)")
                 let nsRange: NSRange = NSRange(location: unfetchedSet.min()!, length: unfetchedSet.count)
-                self.mediaGalleryFinder.enumerateMediaMessages(range: nsRange, transaction: transaction) { (message: TSMessage) in
+                self.mediaGalleryFinder.enumerateMediaAttachments(range: nsRange, transaction: transaction) { (attachment: TSAttachment) in
 
-                    guard !self.deletedMessages.contains(message) else {
-                        Logger.debug("skipping \(message) which has been deleted.")
+                    guard !self.deletedAttachments.contains(attachment) else {
+                        Logger.debug("skipping \(attachment) which has been deleted.")
                         return
                     }
 
-                    guard let item: MediaGalleryItem = self.buildGalleryItem(message: message, transaction: transaction) else {
+                    guard let item: MediaGalleryItem = self.buildGalleryItem(attachment: attachment, transaction: transaction) else {
                         owsFailDebug("unexpectedly failed to buildGalleryItem")
                         return
                     }
@@ -792,7 +796,7 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
         dataSourceDelegates.append(Weak(value: dataSourceDelegate))
     }
 
-    var deletedMessages: Set<TSMessage> = Set()
+    var deletedAttachments: Set<TSAttachment> = Set()
     var deletedGalleryItems: Set<MediaGalleryItem> = Set()
 
     func delete(items: [MediaGalleryItem], initiatedBy: MediaGalleryDataSourceDelegate) {
@@ -806,8 +810,9 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
         self.editingDatabaseConnection.asyncReadWrite { transaction in
             for item in items {
                 let message = item.message
-                message.remove(with: transaction)
-                self.deletedMessages.insert(message)
+                let attachment = item.attachmentStream
+                message.removeAttachment(attachment, transaction: transaction)
+                self.deletedAttachments.insert(attachment)
             }
         }
 
@@ -928,6 +933,6 @@ class MediaGallery: NSObject, MediaGalleryDataSource, MediaTileViewControllerDel
         self.uiDatabaseConnection.read { (transaction: YapDatabaseReadTransaction) in
             count = self.mediaGalleryFinder.mediaCount(transaction: transaction)
         }
-        return Int(count) - deletedMessages.count
+        return Int(count) - deletedAttachments.count
     }
 }
diff --git a/Signal/src/ViewControllers/MessageDetailViewController.swift b/Signal/src/ViewControllers/MessageDetailViewController.swift
index d026a050f..bd15cf48a 100644
--- a/Signal/src/ViewControllers/MessageDetailViewController.swift
+++ b/Signal/src/ViewControllers/MessageDetailViewController.swift
@@ -619,14 +619,14 @@ class MessageDetailViewController: OWSViewController, MediaGalleryDataSourceDele
         let mediaGallery = MediaGallery(thread: self.thread, uiDatabaseConnection: self.uiDatabaseConnection)
 
         mediaGallery.addDataSourceDelegate(self)
-        mediaGallery.presentDetailView(fromViewController: self, mediaMessage: self.message, replacingView: imageView)
+        mediaGallery.presentDetailView(fromViewController: self, mediaAttachment: attachmentStream, replacingView: imageView)
     }
 
     func didTapVideoViewItem(_ viewItem: ConversationViewItem, attachmentStream: TSAttachmentStream, imageView: UIView) {
         let mediaGallery = MediaGallery(thread: self.thread, uiDatabaseConnection: self.uiDatabaseConnection)
 
         mediaGallery.addDataSourceDelegate(self)
-        mediaGallery.presentDetailView(fromViewController: self, mediaMessage: self.message, replacingView: imageView)
+        mediaGallery.presentDetailView(fromViewController: self, mediaAttachment: attachmentStream, replacingView: imageView)
     }
 
     func didTapContactShare(_ viewItem: ConversationViewItem) {
diff --git a/SignalServiceKit/src/Messages/Attachments/TSAttachment.h b/SignalServiceKit/src/Messages/Attachments/TSAttachment.h
index 8bd3df862..dbbbeee70 100644
--- a/SignalServiceKit/src/Messages/Attachments/TSAttachment.h
+++ b/SignalServiceKit/src/Messages/Attachments/TSAttachment.h
@@ -6,6 +6,8 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+@class TSMessage;
+
 typedef NS_ENUM(NSUInteger, TSAttachmentType) {
     TSAttachmentTypeDefault = 0,
     TSAttachmentTypeVoiceMessage = 1,
@@ -41,6 +43,7 @@ typedef NS_ENUM(NSUInteger, TSAttachmentType) {
 
 @property (nonatomic, readonly, nullable) NSString *caption;
 @property (nonatomic, readonly, nullable) NSString *albumMessageId;
+- (nullable TSMessage *)fetchAlbumMessageWithTransaction:(YapDatabaseReadTransaction *)transaction;
 
 #pragma mark -
 
diff --git a/SignalServiceKit/src/Messages/Attachments/TSAttachment.m b/SignalServiceKit/src/Messages/Attachments/TSAttachment.m
index 3f9e66824..7a05bfa7d 100644
--- a/SignalServiceKit/src/Messages/Attachments/TSAttachment.m
+++ b/SignalServiceKit/src/Messages/Attachments/TSAttachment.m
@@ -4,6 +4,7 @@
 
 #import "TSAttachment.h"
 #import "MIMETypeUtil.h"
+#import "TSMessage.h"
 #import <SignalCoreKit/NSString+SSK.h>
 #import <SignalCoreKit/iOSVersions.h>
 
@@ -243,6 +244,16 @@ NSUInteger const TSAttachmentSchemaVersion = 4;
     return _contentType.filterFilename;
 }
 
+#pragma mark - Relationships
+
+- (nullable TSMessage *)fetchAlbumMessageWithTransaction:(YapDatabaseReadTransaction *)transaction
+{
+    if (self.albumMessageId == nil) {
+        return nil;
+    }
+    return [TSMessage fetchObjectWithUniqueID:self.albumMessageId transaction:transaction];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/SignalServiceKit/src/Messages/Interactions/TSInteraction.m b/SignalServiceKit/src/Messages/Interactions/TSInteraction.m
index ff234c97e..ff05ac78a 100644
--- a/SignalServiceKit/src/Messages/Interactions/TSInteraction.m
+++ b/SignalServiceKit/src/Messages/Interactions/TSInteraction.m
@@ -98,7 +98,7 @@ NSString *NSStringFromOWSInteractionType(OWSInteractionType value)
 {
     OWSAssertDebug(timestamp > 0);
 
-    self = [super initWithUniqueId:nil];
+    self = [super initWithUniqueId:[[NSUUID UUID] UUIDString]];
 
     if (!self) {
         return self;
@@ -172,6 +172,7 @@ NSString *NSStringFromOWSInteractionType(OWSInteractionType value)
 }
 
 - (void)saveWithTransaction:(YapDatabaseReadWriteTransaction *)transaction {
+    OWSAssertDebug(self.uniqueId);
     if (!self.uniqueId) {
         self.uniqueId = [OWSPrimaryStorage getAndIncrementMessageIdWithTransaction:transaction];
     }
diff --git a/SignalServiceKit/src/Messages/Interactions/TSMessage.h b/SignalServiceKit/src/Messages/Interactions/TSMessage.h
index 263775a90..fcd90bba8 100644
--- a/SignalServiceKit/src/Messages/Interactions/TSMessage.h
+++ b/SignalServiceKit/src/Messages/Interactions/TSMessage.h
@@ -42,6 +42,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (BOOL)hasAttachments;
 - (NSArray<TSAttachment *> *)attachmentsWithTransaction:(YapDatabaseReadTransaction *)transaction;
+- (void)removeAttachment:(TSAttachment *)attachment
+             transaction:(YapDatabaseReadWriteTransaction *)transaction NS_SWIFT_NAME(removeAttachment(_:transaction:));
 
 - (BOOL)isMediaAlbumWithTransaction:(YapDatabaseReadTransaction *)transaction;
 
diff --git a/SignalServiceKit/src/Messages/Interactions/TSMessage.m b/SignalServiceKit/src/Messages/Interactions/TSMessage.m
index 82755d552..3289008de 100644
--- a/SignalServiceKit/src/Messages/Interactions/TSMessage.m
+++ b/SignalServiceKit/src/Messages/Interactions/TSMessage.m
@@ -214,6 +214,19 @@ static const NSUInteger OWSMessageSchemaVersion = 4;
     return [attachments copy];
 }
 
+- (void)removeAttachment:(TSAttachment *)attachment transaction:(YapDatabaseReadWriteTransaction *)transaction;
+{
+    OWSAssertDebug([self.attachmentIds containsObject:attachment.uniqueId]);
+    [attachment removeWithTransaction:transaction];
+
+    [self.attachmentIds removeObject:attachment.uniqueId];
+
+    // TODO - Should we delete self if we delete the last attachment?
+    // Or should that depend on whether message.body == nil
+
+    [self saveWithTransaction:transaction];
+}
+
 - (BOOL)isMediaAlbumWithTransaction:(YapDatabaseReadTransaction *)transaction
 {
     NSArray<TSAttachment *> *attachments = [self attachmentsWithTransaction:transaction];
diff --git a/SignalServiceKit/src/Messages/OWSMessageSender.m b/SignalServiceKit/src/Messages/OWSMessageSender.m
index 6f955e33a..d2ff700a5 100644
--- a/SignalServiceKit/src/Messages/OWSMessageSender.m
+++ b/SignalServiceKit/src/Messages/OWSMessageSender.m
@@ -87,6 +87,7 @@ void AssertIsOnSendingQueue()
                        contentType:(NSString *)contentType
                     sourceFilename:(nullable NSString *)sourceFilename
                            caption:(nullable NSString *)caption
+                    albumMessageId:(nullable NSString *)albumMessageId
 {
     self = [super init];
     if (!self) {
@@ -97,6 +98,7 @@ void AssertIsOnSendingQueue()
     _contentType = contentType;
     _sourceFilename = sourceFilename;
     _caption = caption;
+    _albumMessageId = albumMessageId;
 
     return self;
 }
@@ -455,10 +457,13 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
                failure:(void (^)(NSError *error))failureHandler
 {
     OWSAssertDebug(dataSource);
+
+    NSString *albumMessageId = message.uniqueId;
     OWSOutgoingAttachmentInfo *attachmentInfo = [[OWSOutgoingAttachmentInfo alloc] initWithDataSource:dataSource
                                                                                           contentType:contentType
                                                                                        sourceFilename:sourceFilename
-                                                                                              caption:nil];
+                                                                                              caption:nil
+                                                                                       albumMessageId:albumMessageId];
     [OutgoingMessagePreparer prepareAttachments:@[ attachmentInfo ]
                                       inMessage:message
                               completionHandler:^(NSError *_Nullable error) {
@@ -1826,14 +1831,15 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
 
         [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
             for (TSAttachmentStream *attachmentStream in attachmentStreams) {
-                [attachmentStream saveWithTransaction:transaction];
-
                 [outgoingMessage.attachmentIds addObject:attachmentStream.uniqueId];
                 if (attachmentStream.sourceFilename) {
                     outgoingMessage.attachmentFilenameMap[attachmentStream.uniqueId] = attachmentStream.sourceFilename;
                 }
             }
             [outgoingMessage saveWithTransaction:transaction];
+            for (TSAttachmentStream *attachmentStream in attachmentStreams) {
+                [attachmentStream saveWithTransaction:transaction];
+            }
         }];
 
         completionHandler(nil);
diff --git a/SignalServiceKit/src/Storage/OWSMediaGalleryFinder.h b/SignalServiceKit/src/Storage/OWSMediaGalleryFinder.h
index 3e15bfd78..5fcbb6108 100644
--- a/SignalServiceKit/src/Storage/OWSMediaGalleryFinder.h
+++ b/SignalServiceKit/src/Storage/OWSMediaGalleryFinder.h
@@ -5,7 +5,7 @@
 NS_ASSUME_NONNULL_BEGIN
 
 @class OWSStorage;
-@class TSMessage;
+@class TSAttachment;
 @class TSThread;
 @class YapDatabaseReadTransaction;
 
@@ -18,15 +18,20 @@ NS_ASSUME_NONNULL_BEGIN
 // How many media items a thread has
 - (NSUInteger)mediaCountWithTransaction:(YapDatabaseReadTransaction *)transaction NS_SWIFT_NAME(mediaCount(transaction:));
 
-// The ordinal position of a message within a thread's media gallery
-- (NSUInteger)mediaIndexForMessage:(TSMessage *)message transaction:(YapDatabaseReadTransaction *)transaction NS_SWIFT_NAME(mediaIndex(message:transaction:));
+// The ordinal position of an attachment within a thread's media gallery
+- (NSUInteger)mediaIndexForAttachment:(TSAttachment *)attachment
+                          transaction:(YapDatabaseReadTransaction *)transaction
+    NS_SWIFT_NAME(mediaIndex(attachment:transaction:));
 
-- (nullable TSMessage *)oldestMediaMessageWithTransaction:(YapDatabaseReadTransaction *)transaction NS_SWIFT_NAME(oldestMediaMessage(transaction:));
-- (nullable TSMessage *)mostRecentMediaMessageWithTransaction:(YapDatabaseReadTransaction *)transaction NS_SWIFT_NAME(mostRecentMediaMessage(transaction:));
+- (nullable TSAttachment *)oldestMediaAttachmentWithTransaction:(YapDatabaseReadTransaction *)transaction
+    NS_SWIFT_NAME(oldestMediaAttachment(transaction:));
+- (nullable TSAttachment *)mostRecentMediaAttachmentWithTransaction:(YapDatabaseReadTransaction *)transaction
+    NS_SWIFT_NAME(mostRecentMediaAttachment(transaction:));
 
-- (void)enumerateMediaMessagesWithRange:(NSRange)range
-                            transaction:(YapDatabaseReadTransaction *)transaction
-                                  block:(void (^)(TSMessage *))messageBlock NS_SWIFT_NAME(enumerateMediaMessages(range:transaction:block:));
+- (void)enumerateMediaAttachmentsWithRange:(NSRange)range
+                               transaction:(YapDatabaseReadTransaction *)transaction
+                                     block:(void (^)(TSAttachment *))attachmentBlock
+    NS_SWIFT_NAME(enumerateMediaAttachments(range:transaction:block:));
 
 #pragma mark - Extension registration
 
diff --git a/SignalServiceKit/src/Storage/OWSMediaGalleryFinder.m b/SignalServiceKit/src/Storage/OWSMediaGalleryFinder.m
index 03a3cf1e3..9a8ea69b0 100644
--- a/SignalServiceKit/src/Storage/OWSMediaGalleryFinder.m
+++ b/SignalServiceKit/src/Storage/OWSMediaGalleryFinder.m
@@ -43,15 +43,15 @@ static NSString *const OWSMediaGalleryFinderExtensionName = @"OWSMediaGalleryFin
     return [[self galleryExtensionWithTransaction:transaction] numberOfItemsInGroup:self.mediaGroup];
 }
 
-- (NSUInteger)mediaIndexForMessage:(TSMessage *)message transaction:(YapDatabaseReadTransaction *)transaction
+- (NSUInteger)mediaIndexForAttachment:(TSAttachment *)attachment transaction:(YapDatabaseReadTransaction *)transaction
 {
     NSString *groupId;
     NSUInteger index;
 
     BOOL wasFound = [[self galleryExtensionWithTransaction:transaction] getGroup:&groupId
                                                                            index:&index
-                                                                          forKey:message.uniqueId
-                                                                    inCollection:[TSMessage collection]];
+                                                                          forKey:attachment.uniqueId
+                                                                    inCollection:[TSAttachment collection]];
 
     OWSAssertDebug(wasFound);
     OWSAssertDebug([self.mediaGroup isEqual:groupId]);
@@ -59,19 +59,19 @@ static NSString *const OWSMediaGalleryFinderExtensionName = @"OWSMediaGalleryFin
     return index;
 }
 
-- (nullable TSMessage *)oldestMediaMessageWithTransaction:(YapDatabaseReadTransaction *)transaction
+- (nullable TSAttachment *)oldestMediaAttachmentWithTransaction:(YapDatabaseReadTransaction *)transaction
 {
     return [[self galleryExtensionWithTransaction:transaction] firstObjectInGroup:self.mediaGroup];
 }
 
-- (nullable TSMessage *)mostRecentMediaMessageWithTransaction:(YapDatabaseReadTransaction *)transaction
+- (nullable TSAttachment *)mostRecentMediaAttachmentWithTransaction:(YapDatabaseReadTransaction *)transaction
 {
     return [[self galleryExtensionWithTransaction:transaction] lastObjectInGroup:self.mediaGroup];
 }
 
-- (void)enumerateMediaMessagesWithRange:(NSRange)range
-                            transaction:(YapDatabaseReadTransaction *)transaction
-                                  block:(void (^)(TSMessage *))messageBlock
+- (void)enumerateMediaAttachmentsWithRange:(NSRange)range
+                               transaction:(YapDatabaseReadTransaction *)transaction
+                                     block:(void (^)(TSAttachment *))attachmentBlock
 {
 
     [[self galleryExtensionWithTransaction:transaction]
@@ -83,9 +83,8 @@ static NSString *const OWSMediaGalleryFinderExtensionName = @"OWSMediaGalleryFin
                                 id _Nonnull object,
                                 NSUInteger index,
                                 BOOL *_Nonnull stop) {
-
-                                OWSAssertDebug([object isKindOfClass:[TSMessage class]]);
-                                messageBlock((TSMessage *)object);
+                                OWSAssertDebug([object isKindOfClass:[TSAttachment class]]);
+                                attachmentBlock((TSAttachment *)object);
                             }];
 }
 
@@ -124,63 +123,81 @@ static NSString *const OWSMediaGalleryFinderExtensionName = @"OWSMediaGalleryFin
 
 + (YapDatabaseAutoView *)mediaGalleryDatabaseExtension
 {
-    YapDatabaseViewSorting *sorting = [YapDatabaseViewSorting withObjectBlock:^NSComparisonResult(YapDatabaseReadTransaction * _Nonnull transaction, NSString * _Nonnull group, NSString * _Nonnull collection1, NSString * _Nonnull key1, id  _Nonnull object1, NSString * _Nonnull collection2, NSString * _Nonnull key2, id  _Nonnull object2) {
-        
-        if (![object1 isKindOfClass:[TSMessage class]]) {
-            OWSFailDebug(@"Unexpected object while sorting: %@", [object1 class]);
-            return NSOrderedSame;
-        }
-        TSMessage *message1 = (TSMessage *)object1;
-        
-        if (![object2 isKindOfClass:[TSMessage class]]) {
-            OWSFailDebug(@"Unexpected object while sorting: %@", [object2 class]);
-            return NSOrderedSame;
-        }
-        TSMessage *message2 = (TSMessage *)object2;
-        
-        return [@(message1.timestampForSorting) compare:@(message2.timestampForSorting)];
-    }];
-    
-    YapDatabaseViewGrouping *grouping = [YapDatabaseViewGrouping withObjectBlock:^NSString * _Nullable(YapDatabaseReadTransaction * _Nonnull transaction, NSString * _Nonnull collection, NSString * _Nonnull key, id  _Nonnull object) {
-        
-        if (![object isKindOfClass:[TSMessage class]]) {
-            return nil;
-        }
-        TSMessage *message = (TSMessage *)object;
-
-        BOOL allAttachmentsAreMedia = message.attachmentIds.count > 0;
-        for (NSString *attachmentId in message.attachmentIds) {
-            OWSAssertDebug(attachmentId.length > 0);
-            if (![self attachmentIdShouldAppearInMediaGallery:attachmentId transaction:transaction]) {
-                allAttachmentsAreMedia = NO;
-                break;
+    YapDatabaseViewSorting *sorting =
+        [YapDatabaseViewSorting withObjectBlock:^NSComparisonResult(YapDatabaseReadTransaction *_Nonnull transaction,
+            NSString *_Nonnull group,
+            NSString *_Nonnull collection1,
+            NSString *_Nonnull key1,
+            id _Nonnull object1,
+            NSString *_Nonnull collection2,
+            NSString *_Nonnull key2,
+            id _Nonnull object2) {
+            if (![object1 isKindOfClass:[TSAttachment class]]) {
+                OWSFailDebug(@"Unexpected object while sorting: %@", [object1 class]);
+                return NSOrderedSame;
+            }
+            TSAttachment *attachment1 = (TSAttachment *)object1;
+
+            if (![object2 isKindOfClass:[TSAttachment class]]) {
+                OWSFailDebug(@"Unexpected object while sorting: %@", [object2 class]);
+                return NSOrderedSame;
+            }
+            TSAttachment *attachment2 = (TSAttachment *)object2;
+
+            TSMessage *_Nullable message1 = [attachment1 fetchAlbumMessageWithTransaction:transaction];
+            TSMessage *_Nullable message2 = [attachment2 fetchAlbumMessageWithTransaction:transaction];
+            if (message1 == nil || message2 == nil) {
+                OWSFailDebug(@"couldn't find albumMessage");
+                return NSOrderedSame;
+            }
+
+            if ([message1.uniqueId isEqualToString:message2.uniqueId]) {
+                NSUInteger index1 = [message1.attachmentIds indexOfObject:attachment1.uniqueId];
+                NSUInteger index2 = [message1.attachmentIds indexOfObject:attachment2.uniqueId];
+
+                if (index1 == NSNotFound || index2 == NSNotFound) {
+                    OWSFailDebug(@"couldn't find attachmentId in it's albumMessage");
+                    return NSOrderedSame;
+                }
+                return [@(index1) compare:@(index2)];
+            } else {
+                return [@(message1.timestampForSorting) compare:@(message2.timestampForSorting)];
+            }
+        }];
+
+    YapDatabaseViewGrouping *grouping =
+        [YapDatabaseViewGrouping withObjectBlock:^NSString *_Nullable(YapDatabaseReadTransaction *_Nonnull transaction,
+            NSString *_Nonnull collection,
+            NSString *_Nonnull key,
+            id _Nonnull object) {
+            // Don't include nil or not yet downloaded attachments.
+            if (![object isKindOfClass:[TSAttachmentStream class]]) {
+                return nil;
+            }
+
+            TSAttachmentStream *attachment = (TSAttachmentStream *)object;
+            if (attachment.albumMessageId == nil) {
+                return nil;
+            }
+
+            if (!attachment.isValidVisualMedia) {
+                return nil;
+            }
+
+            TSMessage *message = [attachment fetchAlbumMessageWithTransaction:transaction];
+            if (message == nil) {
+                OWSFailDebug(@"message was unexpectedly nil");
+                return nil;
             }
-        }
 
-        if (allAttachmentsAreMedia) {
             return [self mediaGroupWithThreadId:message.uniqueThreadId];
-        }
+        }];
 
-        return nil;
-    }];
-    
     YapDatabaseViewOptions *options = [YapDatabaseViewOptions new];
-    options.allowedCollections = [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:TSMessage.collection]];
+    options.allowedCollections =
+        [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:TSAttachment.collection]];
 
-    return [[YapDatabaseAutoView alloc] initWithGrouping:grouping sorting:sorting versionTag:@"3" options:options];
-}
-
-+ (BOOL)attachmentIdShouldAppearInMediaGallery:(NSString *)attachmentId transaction:(YapDatabaseReadTransaction *)transaction
-{
-    TSAttachmentStream *attachment = [TSAttachmentStream fetchObjectWithUniqueID:attachmentId
-                                                                     transaction:transaction];
-
-    // Don't include nil or not yet downloaded attachments.
-    if (![attachment isKindOfClass:[TSAttachmentStream class]]) {
-        return NO;
-    }
-
-    return attachment.isValidVisualMedia;
+    return [[YapDatabaseAutoView alloc] initWithGrouping:grouping sorting:sorting versionTag:@"4" options:options];
 }
 
 @end