diff --git a/backend/prompts/__init__.py b/backend/prompts/__init__.py
index 9aaf45a..4b6ae2b 100644
--- a/backend/prompts/__init__.py
+++ b/backend/prompts/__init__.py
@@ -43,12 +43,17 @@ async def create_prompt(
             prompt_messages.append(message)
     else:
         # Assemble the prompt for non-imported code
-        if params.get("resultImage"):
-            prompt_messages = assemble_prompt(
-                params["image"], stack, params["resultImage"]
-            )
+        if input_mode == "image":
+            if params.get("resultImage"):
+                prompt_messages = assemble_prompt(
+                    params["image"], stack, params["resultImage"]
+                )
+            else:
+                prompt_messages = assemble_prompt(params["image"], stack)
+        elif input_mode == "text":
+            prompt_messages = assemble_text_prompt(params["image"], stack)
         else:
-            prompt_messages = assemble_prompt(params["image"], stack)
+            raise Exception("Invalid input mode")
 
         if params["generationType"] == "update":
             # Transform the history tree into message format
diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index c51f783..e25fae9 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -343,7 +343,7 @@ function App({ navbarComponent }: Props) {
     doGenerateCode({
       generationType: "update",
       inputMode,
-      image: referenceImages[0],
+      image: inputMode === "text" ? initialPrompt : referenceImages[0],
       resultImage,
       history: updatedHistory,
       isImportedFromCode,