From 4290c7054873f275e3a98a611587585a8057fbab Mon Sep 17 00:00:00 2001 From: clean99 Date: Tue, 21 Nov 2023 21:14:25 +0800 Subject: [PATCH 01/10] feat: update prompt to support image comparison --- backend/prompts.py | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/backend/prompts.py b/backend/prompts.py index 3761404..be092ef 100644 --- a/backend/prompts.py +++ b/backend/prompts.py @@ -2,7 +2,7 @@ SYSTEM_PROMPT = """ You are an expert Tailwind developer You take screenshots of a reference web page from the user, and then build single page apps using Tailwind, HTML and JS. -You might also be given a screenshot of a web page that you have already built, and asked to +You might also be given a screenshot(The second image) of a web page that you have already built, and asked to update it to look more like the reference image. - Make sure the app looks exactly like the screenshot. @@ -27,21 +27,26 @@ USER_PROMPT = """ Generate code for a web page that looks exactly like this. """ - -def assemble_prompt(image_data_url): +def assemble_prompt(image_data_url, result_image_data_url=None): + content = [ + { + "type": "image_url", + "image_url": {"url": image_data_url, "detail": "high"}, + }, + { + "type": "text", + "text": USER_PROMPT, + }, + ] + if result_image_data_url: + content.insert(1, { + "type": "image_url", + "image_url": {"url": result_image_data_url, "detail": "high"}, + }) return [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", - "content": [ - { - "type": "image_url", - "image_url": {"url": image_data_url, "detail": "high"}, - }, - { - "type": "text", - "text": USER_PROMPT, - }, - ], + "content": content, }, ] From 02c3178b6bd3e413bd3b7f7791b9e64fa153d125 Mon Sep 17 00:00:00 2001 From: clean99 Date: Tue, 21 Nov 2023 21:14:45 +0800 Subject: [PATCH 02/10] feat: update api to support image comparison --- backend/main.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/backend/main.py b/backend/main.py index 4eb09fa..2fa9e9f 100644 --- a/backend/main.py +++ b/backend/main.py @@ -96,7 +96,10 @@ async def stream_code_test(websocket: WebSocket): async def process_chunk(content): await websocket.send_json({"type": "chunk", "value": content}) - prompt_messages = assemble_prompt(params["image"]) + if params.get("resultImg") and params["resultImg"]: + prompt_messages = assemble_prompt(params["image"], params["resultImg"]) + else: + prompt_messages = assemble_prompt(params["image"]) # Image cache for updates so that we don't have to regenerate images image_cache = {} From bc73613fb160954fafd4e3c6e13b1afb20ef162c Mon Sep 17 00:00:00 2001 From: clean99 Date: Tue, 21 Nov 2023 21:15:16 +0800 Subject: [PATCH 03/10] feat: update interface and fix close bug --- frontend/src/generateCode.ts | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/frontend/src/generateCode.ts b/frontend/src/generateCode.ts index 2083958..b848b78 100644 --- a/frontend/src/generateCode.ts +++ b/frontend/src/generateCode.ts @@ -10,6 +10,7 @@ const STOP_MESSAGE = "Code generation stopped"; export interface CodeGenerationParams { generationType: "create" | "update"; image: string; + resultImg?: string; history?: string[]; // isImageGenerationEnabled: boolean; // TODO: Merge with Settings type in types.ts } @@ -49,13 +50,11 @@ export function generateCode( console.log("Connection closed", event.code, event.reason); if (event.code === USER_CLOSE_WEB_SOCKET_CODE) { toast.success(STOP_MESSAGE); - onComplete(); - } else if (event.code === 1000) { - onComplete(); - } else { + } else if (event.code !== 1000) { console.error("WebSocket error code", event); toast.error(ERROR_MESSAGE); } + onComplete(); }); ws.addEventListener("error", (error) => { From cb1e7a39b93aa8ac4604b78df5f874b729ada5c8 Mon Sep 17 00:00:00 2001 From: clean99 Date: Tue, 21 Nov 2023 21:16:26 +0800 Subject: [PATCH 04/10] feat: add switch and screenshot feature in fe --- frontend/package.json | 1 + frontend/src/App.tsx | 50 +++++++++++++++++++++++++---- frontend/src/components/Preview.tsx | 1 + frontend/yarn.lock | 34 ++++++++++++++++++++ 4 files changed, 79 insertions(+), 7 deletions(-) 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( @@ -43,8 +46,20 @@ function App() { }, "setting" ); + const [isImgCompare, setIsImgCompare] = useState(false); const wsRef = useRef(null); + const takeScreenshot = async (): Promise => { + 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 && (
+
+
+ Auto Image Comparison +
+ +