diff --git a/frontend/package.json b/frontend/package.json
index 7a34ab2..78ce51e 100644
--- a/frontend/package.json
+++ b/frontend/package.json
@@ -22,6 +22,7 @@
     "classnames": "^2.3.2",
     "clsx": "^2.0.0",
     "codemirror": "^6.0.1",
+    "html2canvas": "^1.4.1",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "react-dropzone": "^14.2.3",
diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index 6bb0c80..d533d89 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -12,6 +12,7 @@ import {
   FaMobile,
   FaUndo,
 } from "react-icons/fa";
+import { Switch } from "./components/ui/switch";
 import { Button } from "@/components/ui/button";
 import { Textarea } from "@/components/ui/textarea";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "./components/ui/tabs";
@@ -24,6 +25,8 @@ import { OnboardingNote } from "./components/OnboardingNote";
 import { usePersistedState } from "./hooks/usePersistedState";
 import { UrlInputSection } from "./components/UrlInputSection";
 import TermsOfServiceDialog from "./components/TermsOfServiceDialog";
+import html2canvas from 'html2canvas';
+
 
 function App() {
   const [appState, setAppState] = useState<AppStatus>(
@@ -43,8 +46,20 @@ function App() {
     },
     "setting"
   );
+  const [isImgCompare, setIsImgCompare] = useState<boolean>(false);
   const wsRef = useRef<WebSocket>(null);
 
+  const takeScreenshot = async (): Promise<string> => {
+		const iframeElement = document.querySelector('#preview-desktop') as HTMLIFrameElement;
+    if (!iframeElement?.contentWindow?.document.body) {
+      return '';
+    }
+
+    const canvas = await html2canvas(iframeElement.contentWindow.document.body);
+		const png = canvas.toDataURL('image/png');
+    return png;
+  }
+
   const downloadCode = () => {
     // Create a blob from the generated code
     const blob = new Blob([generatedCode], { type: "text/html" });
@@ -72,6 +87,8 @@ function App() {
 
   const stop = () => {
     wsRef.current?.close?.(USER_CLOSE_WEB_SOCKET_CODE);
+    // make sure stop can correct the state even if the websocket is already closed
+    setAppState(AppStatus.CODE_READY);
   }
 
   function doGenerateCode(params: CodeGenerationParams) {
@@ -103,14 +120,24 @@ function App() {
   }
 
   // Subsequent updates
-  function doUpdate() {
+  async function doUpdate() {
     const updatedHistory = [...history, generatedCode, updateInstruction];
-
-    doGenerateCode({
-      generationType: "update",
-      image: referenceImages[0],
-      history: updatedHistory,
-    });
+    let resultImg = referenceImages[0];
+    if (isImgCompare) {
+      resultImg = await takeScreenshot();
+      doGenerateCode({
+        generationType: "update",
+        image: referenceImages[0],
+        resultImg: resultImg,
+        history: updatedHistory,
+      });
+    } else {
+      doGenerateCode({
+        generationType: "update",
+        image: referenceImages[0],
+        history: updatedHistory,
+      });
+    }
 
     setHistory(updatedHistory);
     setGeneratedCode("");
@@ -160,6 +187,15 @@ function App() {
               {appState === AppStatus.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)}
diff --git a/frontend/src/components/Preview.tsx b/frontend/src/components/Preview.tsx
index de0134f..497bfa9 100644
--- a/frontend/src/components/Preview.tsx
+++ b/frontend/src/components/Preview.tsx
@@ -12,6 +12,7 @@ function Preview({ code, device }: Props) {
   return (
     <div className="flex justify-center mx-2">
       <iframe
+        id={`preview-${device}`}
         title="Preview"
         srcDoc={throttledCode}
         className={classNames(
diff --git a/frontend/yarn.lock b/frontend/yarn.lock
index 2109c60..4779a96 100644
--- a/frontend/yarn.lock
+++ b/frontend/yarn.lock
@@ -1107,6 +1107,11 @@ balanced-match@^1.0.0:
   resolved "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz"
   integrity sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw==
 
+base64-arraybuffer@^1.0.2:
+  version "1.0.2"
+  resolved "https://registry.npmjs.org/base64-arraybuffer/-/base64-arraybuffer-1.0.2.tgz#1c37589a7c4b0746e34bd1feb951da2df01c1bdc"
+  integrity sha512-I3yl4r9QB5ZRY3XuJVEPfc2XhZO6YweFPI+UovAzn+8/hb3oJ6lnysaFcjVpkCPfVWFUDvoZ8kmVDP7WyRtYtQ==
+
 binary-extensions@^2.0.0:
   version "2.2.0"
   resolved "https://registry.npmjs.org/binary-extensions/-/binary-extensions-2.2.0.tgz"
@@ -1272,6 +1277,13 @@ cross-spawn@^7.0.2:
     shebang-command "^2.0.0"
     which "^2.0.1"
 
+css-line-break@^2.1.0:
+  version "2.1.0"
+  resolved "https://registry.npmjs.org/css-line-break/-/css-line-break-2.1.0.tgz#bfef660dfa6f5397ea54116bb3cb4873edbc4fa0"
+  integrity sha512-FHcKFCZcAha3LwfVBhCQbW2nCNbkZXn7KVUJcsT5/P8YmfsVja0FMPJr0B903j/E69HUphKiV9iQArX8SDYA4w==
+  dependencies:
+    utrie "^1.0.2"
+
 cssesc@^3.0.0:
   version "3.0.0"
   resolved "https://registry.npmjs.org/cssesc/-/cssesc-3.0.0.tgz"
@@ -1680,6 +1692,14 @@ hasown@^2.0.0:
   dependencies:
     function-bind "^1.1.2"
 
+html2canvas@^1.4.1:
+  version "1.4.1"
+  resolved "https://registry.npmjs.org/html2canvas/-/html2canvas-1.4.1.tgz#7cef1888311b5011d507794a066041b14669a543"
+  integrity sha512-fPU6BHNpsyIhr8yyMpTLLxAbkaK8ArIBcmZIRiBLiDhjeqvXolaEmDGmELFuX9I4xDcaKKcJl+TKZLqruBbmWA==
+  dependencies:
+    css-line-break "^2.1.0"
+    text-segmentation "^1.0.3"
+
 ignore@^5.2.0, ignore@^5.2.4:
   version "5.2.4"
   resolved "https://registry.npmjs.org/ignore/-/ignore-5.2.4.tgz"
@@ -2379,6 +2399,13 @@ tailwindcss@^3.3.5:
     resolve "^1.22.2"
     sucrase "^3.32.0"
 
+text-segmentation@^1.0.3:
+  version "1.0.3"
+  resolved "https://registry.npmjs.org/text-segmentation/-/text-segmentation-1.0.3.tgz#52a388159efffe746b24a63ba311b6ac9f2d7943"
+  integrity sha512-iOiPUo/BGnZ6+54OsWxZidGCsdU8YbE4PSpdPinp7DeMtUJNJBoJ/ouUSTJjHkh1KntHaltHl/gDs2FC4i5+Nw==
+  dependencies:
+    utrie "^1.0.2"
+
 text-table@^0.2.0:
   version "0.2.0"
   resolved "https://registry.npmjs.org/text-table/-/text-table-0.2.0.tgz"
@@ -2502,6 +2529,13 @@ util-deprecate@^1.0.2:
   resolved "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz"
   integrity sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==
 
+utrie@^1.0.2:
+  version "1.0.2"
+  resolved "https://registry.npmjs.org/utrie/-/utrie-1.0.2.tgz#d42fe44de9bc0119c25de7f564a6ed1b2c87a645"
+  integrity sha512-1MLa5ouZiOmQzUbjbu9VmjLzn1QLXBhwpUa7kdLUQK+KQ5KA9I1vk5U4YHe/X2Ch7PYnJfWuWT+VbuxbGwljhw==
+  dependencies:
+    base64-arraybuffer "^1.0.2"
+
 vite-plugin-checker@^0.6.2:
   version "0.6.2"
   resolved "https://registry.npmjs.org/vite-plugin-checker/-/vite-plugin-checker-0.6.2.tgz"