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 = {}