diff --git a/backend/main.py b/backend/main.py
index 8fabb39..fea6c86 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -96,8 +96,9 @@ async def stream_code(websocket: WebSocket):
     async def process_chunk(content):
         await websocket.send_json({"type": "chunk", "value": content})
 
-    if params.get("resultImg") and params["resultImg"]:
-        prompt_messages = assemble_prompt(params["image"], params["resultImg"])
+    if params.get("resultImage") and params["resultImage"]:
+        prompt_messages = assemble_prompt(params["image"], params["resultImage"])
+
     else:
         prompt_messages = assemble_prompt(params["image"])
 
diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index 01ca77a..4c843d6 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -49,7 +49,8 @@ function App() {
     },
     "setting"
   );
-  const [isImgCompare, setIsImgCompare] = useState<boolean>(false);
+  const [shouldIncludeResultImage, setShouldIncludeResultImage] =
+    useState<boolean>(false);
   const wsRef = useRef<WebSocket>(null);
 
   const takeScreenshot = async (): Promise<string> => {
@@ -127,7 +128,7 @@ function App() {
   // Subsequent updates
   async function doUpdate() {
     const updatedHistory = [...history, generatedCode, updateInstruction];
-    if (isImgCompare) {
+    if (shouldIncludeResultImage) {
       const resultImage = await takeScreenshot();
       doGenerateCode({
         generationType: "update",
@@ -206,18 +207,20 @@ function App() {
               {appState === AppState.CODE_READY && (
                 <div>
                   <div className="grid w-full gap-2">
-                    <div className="flex justify-between items-center gap-x-2">
-                      <div className="font-500">Auto Image Comparison</div>
-                      <Switch
-                        checked={isImgCompare}
-                        onCheckedChange={setIsImgCompare}
-                      />
-                    </div>
                     <Textarea
                       placeholder="Tell the AI what to change..."
                       onChange={(e) => setUpdateInstruction(e.target.value)}
                       value={updateInstruction}
                     />
+                    <div className="flex justify-between items-center gap-x-2">
+                      <div className="font-500 text-xs text-slate-700">
+                        Include screenshot of current version?
+                      </div>
+                      <Switch
+                        checked={shouldIncludeResultImage}
+                        onCheckedChange={setShouldIncludeResultImage}
+                      />
+                    </div>
                     <Button onClick={doUpdate}>Update</Button>
                   </div>
                   <div className="flex items-center gap-x-2 mt-2">