diff --git a/checkmark.svg b/checkmark.svg
deleted file mode 100644
index 5d4c32b..0000000
--- a/checkmark.svg
+++ /dev/null
@@ -1,3 +0,0 @@
-<svg width="21" height="20" viewBox="0 0 21 20" fill="none" xmlns="http://www.w3.org/2000/svg">
-<path fill-rule="evenodd" clip-rule="evenodd" d="M10.3438 19.6875C15.7803 19.6875 20.1875 15.2803 20.1875 9.84375C20.1875 4.4072 15.7803 0 10.3438 0C4.9072 0 0.5 4.4072 0.5 9.84375C0.5 15.2803 4.9072 19.6875 10.3438 19.6875ZM15.3321 6.5547C15.6384 6.09517 15.5142 5.4743 15.0547 5.16795C14.5952 4.8616 13.9743 4.98577 13.6679 5.4453L9.34457 11.9304L7.20711 9.79289C6.81658 9.40237 6.18342 9.40237 5.79289 9.79289C5.40237 10.1834 5.40237 10.8166 5.79289 11.2071L8.79289 14.2071C9.00474 14.419 9.3004 14.5247 9.59854 14.4951C9.89667 14.4656 10.1659 14.304 10.3321 14.0547L15.3321 6.5547Z" fill="#434190"/>
-</svg>
diff --git a/refresh.svg b/refresh.svg
new file mode 100644
index 0000000..c4c702b
--- /dev/null
+++ b/refresh.svg
@@ -0,0 +1,3 @@
+<svg width="14" height="17" viewBox="0 0 14 17" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M13 9.9271C13 13.189 10.3137 15.8333 7 15.8333C3.68629 15.8333 1 13.189 1 9.9271C1 6.66517 3.68629 4.02085 7 4.02085C9 4.02085 10.986 4.99917 12 5.77084M12 5.77084L8.33333 7.08334M12 5.77084L10.6667 1.83334" stroke="#434190" stroke-width="2" stroke-linecap="round"/>
+</svg>
diff --git a/web/src/components/IdeCascadeStudio/IdeCascadeStudio.js b/web/src/components/IdeCascadeStudio/IdeCascadeStudio.js
index a85bc22..e7e37e8 100644
--- a/web/src/components/IdeCascadeStudio/IdeCascadeStudio.js
+++ b/web/src/components/IdeCascadeStudio/IdeCascadeStudio.js
@@ -76,28 +76,27 @@ const IdeCascadeStudio = ({ part, saveCode, loading }) => {
             // Get the canvas image as a Data URL
             const imgBlob = await CascadeController.capture(threejsViewport.environment)
 
-            // // Upload the image to Cloudinary
-            // const cloudinaryImg = await uploadToCloudinary(imgBlob)
+            async function uploadAndUpdateImage(){
+              // Upload the image to Cloudinary
+              const cloudinaryImg = await uploadToCloudinary(imgBlob)
 
-            // // TODO: Save the screenshot as the mainImage if none has been set
-            // // If it has been set, pass along the Blob without uploading
-            // if (!part.mainImage) {
-            //   saveCode({
-            //     input: {
-            //       code,
-            //       title: part?.title,
-            //       userId: currentUser?.sub,
-            //       description: part?.description,
-            //       mainImage: cloudinaryImg.public_id,
-            //     },
-            //     id: part.id,
-            //     isFork: !canEdit,
-            //   })
-            // } else {
-            //   console.log('not saving, passing back into IDE', cloudinaryImg)
-            // }
+              // Save the screenshot as the mainImage
+              saveCode({
+                input: {
+                  code,
+                  title: part?.title,
+                  userId: currentUser?.sub,
+                  description: part?.description,
+                  mainImage: cloudinaryImg.public_id,
+                },
+                id: part.id,
+                isFork: !canEdit,
+              })
 
-            return imgBlob
+              return cloudinaryImg
+            }
+
+            return { image: imgBlob, mainImage: part.mainImage, callback: uploadAndUpdateImage}
           }}
         />
       </div>
diff --git a/web/src/components/IdeToolbar/IdeToolbar.js b/web/src/components/IdeToolbar/IdeToolbar.js
index 4f8ab54..acfb543 100644
--- a/web/src/components/IdeToolbar/IdeToolbar.js
+++ b/web/src/components/IdeToolbar/IdeToolbar.js
@@ -106,11 +106,6 @@ const IdeToolbar = ({
     setIsLoginModalOpen(true)
   }
 
-  const captureScreenshot = async () => {
-    setCaptureState(await onCapture())
-    // console.log({ onCapture: onCapture() })
-  }
-
   const handleDownload = (imgBlob) => {
     const aTag = document.createElement('a')
     document.body.appendChild(aTag)
@@ -260,16 +255,27 @@ const IdeToolbar = ({
               { !captureState
                 ? 'Loading...'
                 : <div className="grid grid-cols-2">
-                  <img src={ window.URL.createObjectURL(captureState) } className="w-32" />
-                  <div className="p-2">
-                    <div className="flex justify-center mb-4">
-                      <Svg name="checkmark" className="mr-2 w-6"/> Part Image Set
-                    </div>
+                  <div className="rounded m-auto" style={{width: 'fit-content', overflow: 'hidden'}}>
+                    <img src={ window.URL.createObjectURL(captureState.image) } className="w-32" />
+                  </div>
+                  <div className="p-2 text-indigo-800">
+                    { (captureState.mainImage && !captureState.updated)
+                      ? <button className="flex m-auto mb-4"
+                          onClick={ async () => {
+                            const cloudinaryImg = await captureState.callback()
+                            setCaptureState({...captureState, mainImage: cloudinaryImg.public_id, updated: true })
+                          }}>
+                          <Svg name="refresh" className="mr-2 w-4 text-indigo-600"/> Set Part Image
+                        </button>
+                      : <div className="flex m-auto mb-4">
+                          <Svg name="checkmark" className="mr-2 w-6 text-indigo-600"/> Part Image Set
+                        </div>
+                    }
                     <Button
                       iconName="save"
-                      className="shadow-md hover:shadow-lg border-indigo-600 border-2 border-opacity-0 hover:border-opacity-100 bg-indigo-800 text-indigo-100 bg-opacity-100"
+                      className="shadow-md hover:shadow-lg border-indigo-600 border-2 border-opacity-0 hover:border-opacity-100 bg-indigo-800 text-indigo-100 text-opacity-100 bg-opacity-100"
                       shouldAnimateHover
-                      onClick={() => handleDownload(captureState)}>
+                      onClick={() => handleDownload(captureState.image)}>
                       Download
                     </Button>
                   </div>
diff --git a/web/src/components/Svg/Svg.js b/web/src/components/Svg/Svg.js
index 4456a9b..20d8769 100644
--- a/web/src/components/Svg/Svg.js
+++ b/web/src/components/Svg/Svg.js
@@ -311,6 +311,18 @@ const Svg = ({ name, className: className2, strokeWidth = 2 }) => {
         />
       </svg>
     ),
+    refresh: (
+      <svg
+        xmlns="http://www.w3.org/2000/svg"
+        viewBox="0 0 14 17"
+        fill="none">
+        <path
+          d="M13 9.9271C13 13.189 10.3137 15.8333 7 15.8333C3.68629 15.8333 1 13.189 1 9.9271C1 6.66517 3.68629 4.02085 7 4.02085C9 4.02085 10.986 4.99917 12 5.77084M12 5.77084L8.33333 7.08334M12 5.77084L10.6667 1.83334"
+          stroke="currentColor"
+          stroke-width="2"
+          stroke-linecap="round"/>
+      </svg>
+    ),
     save: (
       <svg
         xmlns="http://www.w3.org/2000/svg"