diff --git a/frontend/src/lib/components/InputTransformForm.svelte b/frontend/src/lib/components/InputTransformForm.svelte
index 7e86b99fb3..b75d1dd443 100644
--- a/frontend/src/lib/components/InputTransformForm.svelte
+++ b/frontend/src/lib/components/InputTransformForm.svelte
@@ -212,5 +212,5 @@
 		<p>Not recognized arg type {arg.type}</p>
 	{/if}
 {:else}
-	<p>Arg at {argName} is undefined</p>
+	<p class="text-sm text-gray-700">Arg at {argName} is undefined</p>
 {/if}
diff --git a/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte b/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
index 631079d0d0..4735c2ab6f 100644
--- a/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
+++ b/frontend/src/lib/components/flows/content/FlowEditorPanel.svelte
@@ -23,7 +23,7 @@
 	{:else if $selectedId === 'failure'}
 		<FlowFailureModule />
 	{:else}
-		{#each $flowStore.value.modules as flowModule, index (index)}
+		{#each $flowStore.value.modules as flowModule, index (flowModule.id)}
 			<FlowModuleWrapper
 				bind:flowModule
 				previousModuleId={$flowStore.value.modules[index - 1]?.id}
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index 9f009ef9f3..de175ecd12 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -15,13 +15,13 @@
 	import { createScriptFromInlineScript, fork } from '$lib/components/flows/flowStateUtils'
 	import { flowStore } from '$lib/components/flows/flowStore'
 	import SchemaForm from '$lib/components/SchemaForm.svelte'
-	import { RawScript, type Flow, type FlowModule } from '$lib/gen'
+	import { RawScript, type FlowModule } from '$lib/gen'
 	import FlowCard from '../common/FlowCard.svelte'
 	import FlowModuleHeader from './FlowModuleHeader.svelte'
-	import { flowStateStore, type FlowModuleState } from '../flowState'
+	import { flowStateStore } from '../flowState'
 	import { scriptLangToEditorLang } from '$lib/utils'
 	import PropPickerWrapper from '../propPicker/PropPickerWrapper.svelte'
-	import { afterUpdate, getContext, onDestroy, setContext } from 'svelte'
+	import { afterUpdate, getContext, setContext } from 'svelte'
 	import type { FlowEditorContext } from '../types'
 	import { loadSchemaFromModule } from '../utils'
 	import { writable, type Writable } from 'svelte/store'
@@ -47,6 +47,15 @@
 	let panes: HTMLElement
 	let totalTopGap = 0
 
+	let inputTransforms: Record<string, any> =
+		flowModule.value.type === 'rawscript' || flowModule.value.type === 'script'
+			? flowModule.value.input_transforms
+			: {}
+
+	$: if (flowModule.value.type === 'rawscript' || flowModule.value.type === 'script') {
+		flowModule.value.input_transforms = inputTransforms
+	}
+
 	$: stepPropPicker = failureModule
 		? { pickableProperties: { previous_result: { error: 'the error message' } }, extraLib: '' }
 		: getStepPropPicker($flowStateStore, parentModule, previousModuleId, $flowStore, previewArgs)
@@ -62,23 +71,17 @@
 	async function reload(flowModule: FlowModule) {
 		const { input_transforms, schema } = await loadSchemaFromModule(flowModule)
 
-		let hasChanged = false
+		setTimeout(() => {
+			if (
+				(flowModule.value.type == 'script' || flowModule.value.type == 'rawscript') &&
+				JSON.stringify(flowModule.value.input_transforms) !== JSON.stringify(input_transforms)
+			) {
+				inputTransforms = input_transforms
+			}
+		})
+
 		if (JSON.stringify(schema) !== JSON.stringify($flowStateStore[flowModule.id].schema)) {
 			$flowStateStore[flowModule.id].schema = schema
-			hasChanged = true
-		}
-
-		if (
-			flowModule.value.type == 'script' ||
-			(flowModule.value.type == 'rawscript' &&
-				JSON.stringify(flowModule.value.input_transforms) !== JSON.stringify(input_transforms))
-		) {
-			flowModule.value.input_transforms = input_transforms
-			hasChanged = true
-		}
-
-		if (hasChanged) {
-			flowModule = flowModule
 		}
 	}
 
diff --git a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
index 49d8cf8934..be3114562d 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleWrapper.svelte
@@ -79,7 +79,7 @@
 			}}
 			failureModule={$selectedId === 'failure'}
 		/>
-	{:else}
+	{:else if flowModule.value.type === 'rawscript' || flowModule.value.type === 'script'}
 		<FlowModuleComponent bind:flowModule {parentModule} {previousModuleId} />
 	{/if}
 {:else if flowModule.value.type === 'forloopflow'}
diff --git a/frontend/src/lib/components/flows/flowStateUtils.ts b/frontend/src/lib/components/flows/flowStateUtils.ts
index e777adb132..e5dcadf403 100644
--- a/frontend/src/lib/components/flows/flowStateUtils.ts
+++ b/frontend/src/lib/components/flows/flowStateUtils.ts
@@ -1,5 +1,5 @@
 import type { Schema } from '$lib/common'
-import { Script, ScriptService, type FlowModule, type RawScript } from '$lib/gen'
+import { Script, ScriptService, type FlowModule, type PathScript, type RawScript } from '$lib/gen'
 import { initialCode } from '$lib/script_helpers'
 import { userStore, workspaceStore } from '$lib/stores'
 import { getScriptByPath } from '$lib/utils'
@@ -57,8 +57,8 @@ export async function pickScript(
 	path: string,
 	summary: string,
 	id: string
-): Promise<[FlowModule, FlowModuleState]> {
-	const flowModule: FlowModule = {
+): Promise<[FlowModule & { value: PathScript }, FlowModuleState]> {
+	const flowModule: FlowModule & { value: PathScript } = {
 		id,
 		value: { type: 'script', path, input_transforms: {} },
 		summary
@@ -129,7 +129,7 @@ export async function createBranchAll(id: string): Promise<[FlowModule, FlowModu
 	return [branchesFlowModules, flowModuleState]
 }
 
-export async function fork(flowModule: FlowModule): Promise<[FlowModule, FlowModuleState]> {
+export async function fork(flowModule: FlowModule): Promise<[FlowModule & { value: RawScript }, FlowModuleState]> {
 	if (flowModule.value.type !== 'script') {
 		throw new Error('Can only fork a script module')
 	}
@@ -141,7 +141,7 @@ export async function fork(flowModule: FlowModule): Promise<[FlowModule, FlowMod
 	return [forkedFlowModule, flowModuleState]
 }
 
-async function createInlineScriptModuleFromPath(path: string, id: string): Promise<FlowModule> {
+async function createInlineScriptModuleFromPath(path: string, id: string): Promise<FlowModule & { value: RawScript }> {
 	const { content, language } = await getScriptByPath(path)
 
 	return {
@@ -167,7 +167,7 @@ export async function createScriptFromInlineScript(
 	flowModule: FlowModule,
 	suffix: string,
 	schema: Schema
-): Promise<[FlowModule, FlowModuleState]> {
+): Promise<[FlowModule & { value: PathScript }, FlowModuleState]> {
 	const flow = get(flowStore)
 	const user = get(userStore)
 
diff --git a/frontend/src/lib/components/flows/utils.ts b/frontend/src/lib/components/flows/utils.ts
index fb75b125af..85d99892d4 100644
--- a/frontend/src/lib/components/flows/utils.ts
+++ b/frontend/src/lib/components/flows/utils.ts
@@ -19,7 +19,7 @@ export function cleanInputs(flow: Flow | any): Flow {
 	newFlow.value.modules.forEach((mod) => {
 		if (mod.value.type == 'rawscript' || mod.value.type == 'script') {
 			if (Object.keys(mod.input_transforms ?? {}).length > 0) {
-				mod.value.input_transforms = mod.input_transforms
+				mod.value.input_transforms = mod.input_transforms!
 				delete mod.input_transforms
 			}
 			Object.values(mod.input_transforms ?? {}).forEach((inp) => {
@@ -81,7 +81,7 @@ export async function loadSchemaFromModule(module: FlowModule): Promise<{
 		const keys = Object.keys(schema?.properties ?? {})
 
 		if (Object.keys(module.input_transforms ?? {}).length > 0) {
-			mod.input_transforms = module.input_transforms
+			mod.input_transforms = module.input_transforms!
 		}
 		let input_transforms = mod.input_transforms ?? module.input_transforms ?? {}
 
@@ -124,9 +124,8 @@ export function getDefaultExpr(
 	previousExpr?: string
 ) {
 	const expr = previousExpr ?? `previous_result.${key}`
-	return `import { previous_result, flow_input, step, variable, resource, params } from 'windmill${
-		importPath ? `@${importPath}` : ''
-	}'
+	return `import { previous_result, flow_input, step, variable, resource, params } from 'windmill${importPath ? `@${importPath}` : ''
+		}'
 
 ${expr}`
 }
diff --git a/openflow.openapi.yaml b/openflow.openapi.yaml
index ff8baf2281..77a660ee6d 100644
--- a/openflow.openapi.yaml
+++ b/openflow.openapi.yaml
@@ -185,6 +185,7 @@ components:
         - type
         - content
         - language
+        - input_transforms
 
     PathScript:
       type: object
@@ -203,6 +204,7 @@ components:
       required:
         - type
         - path
+        - input_transforms
 
     ForloopFlow:
       type: object