diff --git a/Source/SelectionsGallery/YPSelectionsGalleryCell.swift b/Source/SelectionsGallery/YPSelectionsGalleryCell.swift
index d51abc695..38a3c424d 100644
--- a/Source/SelectionsGallery/YPSelectionsGalleryCell.swift
+++ b/Source/SelectionsGallery/YPSelectionsGalleryCell.swift
@@ -9,11 +9,17 @@
 import UIKit
 import Stevia
 
+public protocol YPSelectionsGalleryCellDelegate: class {
+    func selectionsGalleryCellDidTapRemove(cell: YPSelectionsGalleryCell)
+}
+
 public class YPSelectionsGalleryCell: UICollectionViewCell {
     
+    weak var delegate: YPSelectionsGalleryCellDelegate?
     let imageView = UIImageView()
     let editIcon = UIView()
     let editSquare = UIView()
+    let removeButton = UIButton()
     
     override init(frame: CGRect) {
         super.init(frame: frame)
@@ -21,7 +27,8 @@ public class YPSelectionsGalleryCell: UICollectionViewCell {
         sv(
             imageView,
             editIcon,
-            editSquare
+            editSquare,
+            removeButton
         )
         
         imageView.fillContainer()
@@ -30,6 +37,8 @@ public class YPSelectionsGalleryCell: UICollectionViewCell {
         editSquare.CenterY == editIcon.CenterY
         editSquare.CenterX == editIcon.CenterX
         
+        removeButton.top(12).trailing(12)
+        
         layer.shadowColor = UIColor.black.cgColor
         layer.shadowOpacity = 0.2
         layer.shadowOffset = CGSize(width: 4, height: 7)
@@ -47,6 +56,14 @@ public class YPSelectionsGalleryCell: UICollectionViewCell {
             v.layer.borderWidth = 1
             v.layer.borderColor = UIColor.ypLabel.cgColor
         }
+        removeButton.setImage(YPConfig.icons.removeImage, for: .normal)
+    
+        removeButton.addTarget(self, action: #selector(removeButtonTapped), for: .touchUpInside)
+    }
+    
+    @objc
+    func removeButtonTapped() {
+        delegate?.selectionsGalleryCellDidTapRemove(cell: self)
     }
     
     func setEditable(_ editable: Bool) {
@@ -54,20 +71,10 @@ public class YPSelectionsGalleryCell: UICollectionViewCell {
         self.editSquare.isHidden = !editable
     }
     
-    func addRemoveButton(target: Any?, action: Selector) {
-        let image = YPConfig.icons.removeImage
-        let removeButton = UIButton(frame: CGRect(origin: CGPoint.zero, size: image.size))
-        removeButton.setBackgroundImage(image, for: UIControl.State())
-        sv(removeButton)
-        removeButton.top(12).trailing(12)
-        removeButton.addTarget(target, action: action, for: .touchUpInside)
-    }
-    
     required public init?(coder aDecoder: NSCoder) {
         fatalError("init(coder:) has not been implemented")
     }
     
-    
     public override var isHighlighted: Bool {
         didSet {
             UIView.animate(withDuration: 0.5,
diff --git a/Source/SelectionsGallery/YPSelectionsGalleryVC.swift b/Source/SelectionsGallery/YPSelectionsGalleryVC.swift
index d5683258e..6a94c1c59 100644
--- a/Source/SelectionsGallery/YPSelectionsGalleryVC.swift
+++ b/Source/SelectionsGallery/YPSelectionsGalleryVC.swift
@@ -8,7 +8,7 @@
 
 import UIKit
 
-public class YPSelectionsGalleryVC: UIViewController {
+public class YPSelectionsGalleryVC: UIViewController, YPSelectionsGalleryCellDelegate {
     
     override public var prefersStatusBarHidden: Bool { return YPConfig.hidesStatusBar }
     
@@ -63,15 +63,13 @@ public class YPSelectionsGalleryVC: UIViewController {
         didFinishHandler?(self, items)
     }
     
-    @objc
-    private func removeButtonDidClick(sender: UIButton) {
-        guard let cell = sender.superview as? UICollectionViewCell, let indexPath = v.collectionView.indexPath(for: cell) else {
-            return
+    public func selectionsGalleryCellDidTapRemove(cell: YPSelectionsGalleryCell) {
+        if let indexPath = v.collectionView.indexPath(for: cell) {
+            items.remove(at: indexPath.row)
+            v.collectionView.performBatchUpdates({
+                v.collectionView.deleteItems(at: [indexPath])
+            }, completion: { _ in })
         }
-        items.remove(at: indexPath.row)
-        v.collectionView.performBatchUpdates({
-            v.collectionView.deleteItems(at: [indexPath])
-        }, completion: { _ in })
     }
 }
 
@@ -87,6 +85,7 @@ extension YPSelectionsGalleryVC: UICollectionViewDataSource {
                                                             for: indexPath) as? YPSelectionsGalleryCell else {
             return UICollectionViewCell()
         }
+        cell.delegate = self
         let item = items[indexPath.row]
         switch item {
         case .photo(let photo):
@@ -96,11 +95,7 @@ extension YPSelectionsGalleryVC: UICollectionViewDataSource {
             cell.imageView.image = video.thumbnail
             cell.setEditable(YPConfig.showsVideoTrimmer)
         }
-        
-        if !YPConfig.gallery.hidesRemoveButton {
-            cell.addRemoveButton(target: self, action: #selector(removeButtonDidClick(sender:)))
-        }
-        
+        cell.removeButton.isHidden = YPConfig.gallery.hidesRemoveButton
         return cell
     }
 }