diff --git a/avm/res/virtual-machine-images/image-template/README.md b/avm/res/virtual-machine-images/image-template/README.md
index d8e9e5313d..022c56722c 100644
--- a/avm/res/virtual-machine-images/image-template/README.md
+++ b/avm/res/virtual-machine-images/image-template/README.md
@@ -45,12 +45,6 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
   name: 'imageTemplateDeployment'
   params: {
     // Required parameters
-    customizationSteps: [
-      {
-        restartTimeout: '30m'
-        type: 'WindowsRestart'
-      }
-    ]
     distributions: [
       {
         imageName: 'mi-vmiitmin-001'
@@ -89,14 +83,6 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
   "contentVersion": "1.0.0.0",
   "parameters": {
     // Required parameters
-    "customizationSteps": {
-      "value": [
-        {
-          "restartTimeout": "30m",
-          "type": "WindowsRestart"
-        }
-      ]
-    },
     "distributions": {
       "value": [
         {
@@ -149,26 +135,6 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
   name: 'imageTemplateDeployment'
   params: {
     // Required parameters
-    customizationSteps: [
-      {
-        name: 'PowerShell installation'
-        scriptUri: '<scriptUri>'
-        type: 'Shell'
-      }
-      {
-        destination: 'Initialize-LinuxSoftware.ps1'
-        name: 'Initialize-LinuxSoftware'
-        sourceUri: '<sourceUri>'
-        type: 'File'
-      }
-      {
-        inline: [
-          'pwsh \'Initialize-LinuxSoftware.ps1\''
-        ]
-        name: 'Software installation'
-        type: 'Shell'
-      }
-    ]
     distributions: [
       {
         imageName: 'mi-vmiitmax-001'
@@ -202,6 +168,26 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
     name: 'vmiitmax001'
     // Non-required parameters
     buildTimeoutInMinutes: 60
+    customizationSteps: [
+      {
+        name: 'PowerShell installation'
+        scriptUri: '<scriptUri>'
+        type: 'Shell'
+      }
+      {
+        destination: 'Initialize-LinuxSoftware.ps1'
+        name: 'Initialize-LinuxSoftware'
+        sourceUri: '<sourceUri>'
+        type: 'File'
+      }
+      {
+        inline: [
+          'pwsh \'Initialize-LinuxSoftware.ps1\''
+        ]
+        name: 'Software installation'
+        type: 'Shell'
+      }
+    ]
     location: '<location>'
     lock: {
       kind: 'CanNotDelete'
@@ -228,7 +214,7 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
         roleDefinitionIdOrName: '<roleDefinitionIdOrName>'
       }
     ]
-    stagingResourceGroup: '<stagingResourceGroup>'
+    stagingResourceGroupResourceId: '<stagingResourceGroupResourceId>'
     subnetResourceId: '<subnetResourceId>'
     tags: {
       Environment: 'Non-Prod'
@@ -269,28 +255,6 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
   "contentVersion": "1.0.0.0",
   "parameters": {
     // Required parameters
-    "customizationSteps": {
-      "value": [
-        {
-          "name": "PowerShell installation",
-          "scriptUri": "<scriptUri>",
-          "type": "Shell"
-        },
-        {
-          "destination": "Initialize-LinuxSoftware.ps1",
-          "name": "Initialize-LinuxSoftware",
-          "sourceUri": "<sourceUri>",
-          "type": "File"
-        },
-        {
-          "inline": [
-            "pwsh \"Initialize-LinuxSoftware.ps1\""
-          ],
-          "name": "Software installation",
-          "type": "Shell"
-        }
-      ]
-    },
     "distributions": {
       "value": [
         {
@@ -334,6 +298,28 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
     "buildTimeoutInMinutes": {
       "value": 60
     },
+    "customizationSteps": {
+      "value": [
+        {
+          "name": "PowerShell installation",
+          "scriptUri": "<scriptUri>",
+          "type": "Shell"
+        },
+        {
+          "destination": "Initialize-LinuxSoftware.ps1",
+          "name": "Initialize-LinuxSoftware",
+          "sourceUri": "<sourceUri>",
+          "type": "File"
+        },
+        {
+          "inline": [
+            "pwsh \"Initialize-LinuxSoftware.ps1\""
+          ],
+          "name": "Software installation",
+          "type": "Shell"
+        }
+      ]
+    },
     "location": {
       "value": "<location>"
     },
@@ -370,8 +356,8 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
         }
       ]
     },
-    "stagingResourceGroup": {
-      "value": "<stagingResourceGroup>"
+    "stagingResourceGroupResourceId": {
+      "value": "<stagingResourceGroupResourceId>"
     },
     "subnetResourceId": {
       "value": "<subnetResourceId>"
@@ -427,12 +413,6 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
   name: 'imageTemplateDeployment'
   params: {
     // Required parameters
-    customizationSteps: [
-      {
-        restartTimeout: '10m'
-        type: 'WindowsRestart'
-      }
-    ]
     distributions: [
       {
         sharedImageGalleryImageDefinitionResourceId: '<sharedImageGalleryImageDefinitionResourceId>'
@@ -453,6 +433,12 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
     }
     name: 'vmiitwaf001'
     // Non-required parameters
+    customizationSteps: [
+      {
+        restartTimeout: '10m'
+        type: 'WindowsRestart'
+      }
+    ]
     location: '<location>'
     subnetResourceId: '<subnetResourceId>'
     tags: {
@@ -477,14 +463,6 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
   "contentVersion": "1.0.0.0",
   "parameters": {
     // Required parameters
-    "customizationSteps": {
-      "value": [
-        {
-          "restartTimeout": "10m",
-          "type": "WindowsRestart"
-        }
-      ]
-    },
     "distributions": {
       "value": [
         {
@@ -513,6 +491,14 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
       "value": "vmiitwaf001"
     },
     // Non-required parameters
+    "customizationSteps": {
+      "value": [
+        {
+          "restartTimeout": "10m",
+          "type": "WindowsRestart"
+        }
+      ]
+    },
     "location": {
       "value": "<location>"
     },
@@ -539,7 +525,6 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
 
 | Parameter | Type | Description |
 | :-- | :-- | :-- |
-| [`customizationSteps`](#parameter-customizationsteps) | array | Customization steps to be run when building the VM image. |
 | [`distributions`](#parameter-distributions) | array | The distribution targets where the image output needs to go to. |
 | [`imageSource`](#parameter-imagesource) | object | Image source definition in object format. |
 | [`managedIdentities`](#parameter-managedidentities) | object | The managed identity definition for this resource. |
@@ -550,13 +535,14 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
 | Parameter | Type | Description |
 | :-- | :-- | :-- |
 | [`buildTimeoutInMinutes`](#parameter-buildtimeoutinminutes) | int | The image build timeout in minutes. 0 means the default 240 minutes. |
+| [`customizationSteps`](#parameter-customizationsteps) | array | Customization steps to be run when building the VM image. |
 | [`enableTelemetry`](#parameter-enabletelemetry) | bool | Enable/Disable usage telemetry for module. |
 | [`location`](#parameter-location) | string | Location for all resources. |
 | [`lock`](#parameter-lock) | object | The lock settings of the service. |
 | [`optimizeVmBoot`](#parameter-optimizevmboot) | string | The optimize property can be enabled while creating a VM image and allows VM optimization to improve image creation time. |
 | [`osDiskSizeGB`](#parameter-osdisksizegb) | int | Specifies the size of OS disk. |
 | [`roleAssignments`](#parameter-roleassignments) | array | Array of role assignments to create. |
-| [`stagingResourceGroup`](#parameter-stagingresourcegroup) | string | Resource ID of the staging resource group in the same subscription and location as the image template that will be used to build the image.</p>If this field is empty, a resource group with a random name will be created.</p>If the resource group specified in this field doesn't exist, it will be created with the same name.</p>If the resource group specified exists, it must be empty and in the same region as the image template.</p>The resource group created will be deleted during template deletion if this field is empty or the resource group specified doesn't exist,</p>but if the resource group specified exists the resources created in the resource group will be deleted during template deletion and the resource group itself will remain. |
+| [`stagingResourceGroupResourceId`](#parameter-stagingresourcegroupresourceid) | string | Resource ID of the staging resource group in the same subscription and location as the image template that will be used to build the image.</p>If this field is empty, a resource group with a random name will be created.</p>If the resource group specified in this field doesn't exist, it will be created with the same name.</p>If the resource group specified exists, it must be empty and in the same region as the image template.</p>The resource group created will be deleted during template deletion if this field is empty or the resource group specified doesn't exist,</p>but if the resource group specified exists the resources created in the resource group will be deleted during template deletion and the resource group itself will remain. |
 | [`subnetResourceId`](#parameter-subnetresourceid) | string | Resource ID of an already existing subnet, e.g.: /subscriptions/<subscriptionId>/resourceGroups/<resourceGroupName>/providers/Microsoft.Network/virtualNetworks/<vnetName>/subnets/<subnetName>.</p>If no value is provided, a new temporary VNET and subnet will be created in the staging resource group and will be deleted along with the remaining temporary resources. |
 | [`tags`](#parameter-tags) | object | Tags of the resource. |
 | [`validationProcess`](#parameter-validationprocess) | object | Configuration options and list of validations to be performed on the resulting image. |
@@ -569,13 +555,6 @@ module imageTemplate 'br/public:avm/res/virtual-machine-images/image-template:<v
 | :-- | :-- | :-- |
 | [`baseTime`](#parameter-basetime) | string | Do not provide a value! This date value is used to generate a unique image template name. |
 
-### Parameter: `customizationSteps`
-
-Customization steps to be run when building the VM image.
-
-- Required: Yes
-- Type: array
-
 ### Parameter: `distributions`
 
 The distribution targets where the image output needs to go to.
@@ -625,6 +604,13 @@ The image build timeout in minutes. 0 means the default 240 minutes.
 - Type: int
 - Default: `0`
 
+### Parameter: `customizationSteps`
+
+Customization steps to be run when building the VM image.
+
+- Required: No
+- Type: array
+
 ### Parameter: `enableTelemetry`
 
 Enable/Disable usage telemetry for module.
@@ -796,7 +782,7 @@ The principal type of the assigned principal ID.
   ]
   ```
 
-### Parameter: `stagingResourceGroup`
+### Parameter: `stagingResourceGroupResourceId`
 
 Resource ID of the staging resource group in the same subscription and location as the image template that will be used to build the image.</p>If this field is empty, a resource group with a random name will be created.</p>If the resource group specified in this field doesn't exist, it will be created with the same name.</p>If the resource group specified exists, it must be empty and in the same region as the image template.</p>The resource group created will be deleted during template deletion if this field is empty or the resource group specified doesn't exist,</p>but if the resource group specified exists the resources created in the resource group will be deleted during template deletion and the resource group itself will remain.
 
diff --git a/avm/res/virtual-machine-images/image-template/main.bicep b/avm/res/virtual-machine-images/image-template/main.bicep
index 9aefe03417..c32e5e36f0 100644
--- a/avm/res/virtual-machine-images/image-template/main.bicep
+++ b/avm/res/virtual-machine-images/image-template/main.bicep
@@ -25,11 +25,11 @@ param subnetResourceId string?
 @description('Required. Image source definition in object format.')
 param imageSource object
 
-@description('Required. Customization steps to be run when building the VM image.')
-param customizationSteps array
+@description('Optional. Customization steps to be run when building the VM image.')
+param customizationSteps array?
 
 @description('Optional. Resource ID of the staging resource group in the same subscription and location as the image template that will be used to build the image.</p>If this field is empty, a resource group with a random name will be created.</p>If the resource group specified in this field doesn\'t exist, it will be created with the same name.</p>If the resource group specified exists, it must be empty and in the same region as the image template.</p>The resource group created will be deleted during template deletion if this field is empty or the resource group specified doesn\'t exist,</p>but if the resource group specified exists the resources created in the resource group will be deleted during template deletion and the resource group itself will remain.')
-param stagingResourceGroup string?
+param stagingResourceGroupResourceId string?
 
 @description('Optional. The lock settings of the service.')
 param lock lockType
@@ -138,7 +138,7 @@ resource imageTemplate 'Microsoft.VirtualMachineImages/imageTemplates@2023-07-01
     }
     source: imageSource
     customize: customizationSteps
-    stagingResourceGroup: stagingResourceGroup
+    stagingResourceGroup: stagingResourceGroupResourceId
     distribute: [
       for distribution in distributions: union(
         {
diff --git a/avm/res/virtual-machine-images/image-template/main.json b/avm/res/virtual-machine-images/image-template/main.json
index 8af2add54f..df6d5ed0ea 100644
--- a/avm/res/virtual-machine-images/image-template/main.json
+++ b/avm/res/virtual-machine-images/image-template/main.json
@@ -6,7 +6,7 @@
     "_generator": {
       "name": "bicep",
       "version": "0.29.47.4906",
-      "templateHash": "12715542154180929791"
+      "templateHash": "495678619291242908"
     },
     "name": "Virtual Machine Image Templates",
     "description": "This module deploys a Virtual Machine Image Template that can be consumed by Azure Image Builder (AIB).",
@@ -458,11 +458,12 @@
     },
     "customizationSteps": {
       "type": "array",
+      "nullable": true,
       "metadata": {
-        "description": "Required. Customization steps to be run when building the VM image."
+        "description": "Optional. Customization steps to be run when building the VM image."
       }
     },
-    "stagingResourceGroup": {
+    "stagingResourceGroupResourceId": {
       "type": "string",
       "nullable": true,
       "metadata": {
@@ -607,7 +608,7 @@
         },
         "source": "[parameters('imageSource')]",
         "customize": "[parameters('customizationSteps')]",
-        "stagingResourceGroup": "[parameters('stagingResourceGroup')]",
+        "stagingResourceGroup": "[parameters('stagingResourceGroupResourceId')]",
         "validate": "[parameters('validationProcess')]",
         "optimize": "[if(not(equals(parameters('optimizeVmBoot'), null())), createObject('vmBoot', createObject('state', parameters('optimizeVmBoot'))), null())]"
       }
diff --git a/avm/res/virtual-machine-images/image-template/tests/e2e/defaults/main.test.bicep b/avm/res/virtual-machine-images/image-template/tests/e2e/defaults/main.test.bicep
index bb557a787c..246c9d8a1f 100644
--- a/avm/res/virtual-machine-images/image-template/tests/e2e/defaults/main.test.bicep
+++ b/avm/res/virtual-machine-images/image-template/tests/e2e/defaults/main.test.bicep
@@ -50,12 +50,6 @@ module testDeployment '../../../main.bicep' = {
   params: {
     name: '${namePrefix}${serviceShort}001'
     location: resourceLocation
-    customizationSteps: [
-      {
-        restartTimeout: '30m'
-        type: 'WindowsRestart'
-      }
-    ]
     imageSource: {
       offer: 'Windows-11'
       publisher: 'MicrosoftWindowsDesktop'
diff --git a/avm/res/virtual-machine-images/image-template/tests/e2e/max/dependencies.bicep b/avm/res/virtual-machine-images/image-template/tests/e2e/max/dependencies.bicep
index 036e72fb96..0581a62957 100644
--- a/avm/res/virtual-machine-images/image-template/tests/e2e/max/dependencies.bicep
+++ b/avm/res/virtual-machine-images/image-template/tests/e2e/max/dependencies.bicep
@@ -271,11 +271,11 @@ resource assetsStorageAccount_upload 'Microsoft.Resources/deploymentScripts@2023
     scriptContent: loadTextContent('../../../../../../utilities/e2e-template-assets/scripts/Set-StorageContainerContentByEnvVar.ps1')
     environmentVariables: [
       {
-        name: 'script_Install__LinuxPowerShell_sh' // May only be alphanumeric characters & underscores. The upload will replace '_' with '.' and '__' with '-'.
+        name: '__SCRIPT__Install__LinuxPowerShell_sh' // May only be alphanumeric characters & underscores. The upload will replace '_' with '.' and '__' with '-'.
         value: loadTextContent('src/Install-LinuxPowerShell.sh')
       }
       {
-        name: 'script_Initialize__LinuxSoftware_ps1' // May only be alphanumeric characters & underscores. The upload will replace '_' with '.' and '__' with '-'.
+        name: '__SCRIPT__Initialize__LinuxSoftware_ps1' // May only be alphanumeric characters & underscores. The upload will replace '_' with '.' and '__' with '-'.
         value: loadTextContent('src/Initialize-LinuxSoftware.ps1')
       }
     ]
diff --git a/avm/res/virtual-machine-images/image-template/tests/e2e/max/main.test.bicep b/avm/res/virtual-machine-images/image-template/tests/e2e/max/main.test.bicep
index ff20a8a137..f997482ca2 100644
--- a/avm/res/virtual-machine-images/image-template/tests/e2e/max/main.test.bicep
+++ b/avm/res/virtual-machine-images/image-template/tests/e2e/max/main.test.bicep
@@ -60,7 +60,7 @@ module testDeployment '../../../main.bicep' = {
   params: {
     name: '${namePrefix}${serviceShort}001'
     location: resourceLocation
-    stagingResourceGroup: '${subscription().id}/resourcegroups/${resourceGroupName}-staging'
+    stagingResourceGroupResourceId: '${subscription().id}/resourcegroups/${resourceGroupName}-staging'
     customizationSteps: [
       {
         type: 'Shell'
diff --git a/avm/utilities/e2e-template-assets/scripts/Set-StorageContainerContentByEnvVar.ps1 b/avm/utilities/e2e-template-assets/scripts/Set-StorageContainerContentByEnvVar.ps1
index d0af03f8c2..3d06cd5ab1 100644
--- a/avm/utilities/e2e-template-assets/scripts/Set-StorageContainerContentByEnvVar.ps1
+++ b/avm/utilities/e2e-template-assets/scripts/Set-StorageContainerContentByEnvVar.ps1
@@ -15,7 +15,7 @@ Required. The name of the Storage Account to upload to
 Required. The container to upload the files to
 
 .EXAMPLE
-Set-StorageContainerContentByEnvVar -StorageAccountName 'mystorage' -TargetContainer 'myContainer'
+. 'Set-StorageContainerContentByEnvVar.ps1' -StorageAccountName 'mystorage' -TargetContainer 'myContainer'
 
 Upload any required data to the storage account 'mystorage' and container 'myContainer'.
 #>
@@ -33,9 +33,9 @@ Write-Verbose 'Fetching & storing scripts' -Verbose
 $contentDirectoryName = 'scripts'
 $contentDirectory = (New-Item $contentDirectoryName -ItemType 'Directory' -Force).FullName
 $scriptPaths = @()
-foreach ($scriptEnvVar in (Get-ChildItem 'env:*').Name | Where-Object { $_ -like 'script_*' }) {
+foreach ($scriptEnvVar in (Get-ChildItem 'env:*').Name | Where-Object { $_ -like '__SCRIPT__*' }) {
     # Handle value like 'script_Initialize__LinuxSoftware_ps1'
-    $scriptName = $scriptEnvVar -replace 'script_', '' -replace '__', '-' -replace '_', '.'
+    $scriptName = $scriptEnvVar -replace '__SCRIPT__', '' -replace '__', '-' -replace '_', '.'
     $scriptContent = (Get-Item env:$scriptEnvVar).Value
 
     Write-Verbose ('Storing file [{0}] with length [{1}]' -f $scriptName, $scriptContent.Length) -Verbose