diff --git a/frontend/src/lib/components/InputTransformForm.svelte b/frontend/src/lib/components/InputTransformForm.svelte
index 7251efb1bf..1213ac4391 100644
--- a/frontend/src/lib/components/InputTransformForm.svelte
+++ b/frontend/src/lib/components/InputTransformForm.svelte
@@ -364,7 +364,7 @@
 </script>
 
 {#if arg != undefined}
-	<div class={twMerge('pl-2 pt-2 pb-2 ml-2 relative group', className)}>
+	<div class={twMerge('pt-2 pb-2 relative group', className)}>
 		<div class="flex flex-row justify-between gap-1 pb-1">
 			<div class="flex flex-wrap grow">
 				<FieldHeader
diff --git a/frontend/src/lib/components/InputTransformSchemaForm.svelte b/frontend/src/lib/components/InputTransformSchemaForm.svelte
index 5abf5a669f..56c5df0011 100644
--- a/frontend/src/lib/components/InputTransformSchemaForm.svelte
+++ b/frontend/src/lib/components/InputTransformSchemaForm.svelte
@@ -78,9 +78,9 @@
 	})
 </script>
 
-<div class="w-full {clazz}">
+<div class="w-full mb-6 {clazz}">
 	{#if enableAi}
-		<div class="px-0.5 pt-0.5">
+		<div class="mt-2">
 			<StepInputsGen
 				{pickableProperties}
 				argNames={keys
@@ -99,7 +99,7 @@
 	{#if keys.length > 0}
 		{#each keys as argName, index (argName)}
 			{#if (!filter || filter.includes(argName)) && Object.keys(schema.properties ?? {}).includes(argName)}
-				<div class="pt-2 relative">
+				<div class="mt-2 relative">
 					<InputTransformForm
 						{previousModuleId}
 						bind:arg={args[argName]}
diff --git a/frontend/src/lib/components/copilot/StepInputsGen.svelte b/frontend/src/lib/components/copilot/StepInputsGen.svelte
index 5d749382b9..e2a69da06c 100644
--- a/frontend/src/lib/components/copilot/StepInputsGen.svelte
+++ b/frontend/src/lib/components/copilot/StepInputsGen.svelte
@@ -178,6 +178,7 @@ input_name2: expression2
 		<Button
 			size="xs"
 			color="light"
+			wrapperClasses="flex-1 border rounded-md"
 			btnClasses={twMerge(
 				'text-violet-800 dark:text-violet-400',
 				!loading && Object.keys($generatedExprs || {}).length > 0
diff --git a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
index 67554c0efa..6288f75b4d 100644
--- a/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
+++ b/frontend/src/lib/components/flows/content/FlowModuleComponent.svelte
@@ -548,7 +548,7 @@
 														></div>
 													{/if}
 													<InputTransformSchemaForm
-														class="px-1 xl:px-2"
+														class="px-2 xl:px-4"
 														bind:this={inputTransformSchemaForm}
 														pickableProperties={stepPropPicker.pickableProperties}
 														schema={flowStateStore.val[$selectedId]?.schema ?? {}}
diff --git a/frontend/src/lib/components/text_input/TextInput.svelte b/frontend/src/lib/components/text_input/TextInput.svelte
index 6f85fe1a6b..e79fd72437 100644
--- a/frontend/src/lib/components/text_input/TextInput.svelte
+++ b/frontend/src/lib/components/text_input/TextInput.svelte
@@ -16,7 +16,7 @@
 	{...inputProps}
 	class={twMerge(
 		// TODO Make tailwind classes for every color
-		'!bg-surface-secondary !border-transparent !shadow-none !py-2.5 !px-4',
+		'!bg-surface-secondary !border-transparent !shadow-none !py-2 !px-4',
 		'focus:!ring-0',
 		'[appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none',
 		'!rounded-md',