diff --git a/ci-operator/step-registry/upi/install/azure/upi-install-azure-commands.sh b/ci-operator/step-registry/upi/install/azure/upi-install-azure-commands.sh
index 9f05d919555d7..99d704ab5bc32 100644
--- a/ci-operator/step-registry/upi/install/azure/upi-install-azure-commands.sh
+++ b/ci-operator/step-registry/upi/install/azure/upi-install-azure-commands.sh
@@ -158,9 +158,11 @@ az storage account create -g $RESOURCE_GROUP --location $AZURE_REGION --name $AC
 ACCOUNT_KEY=$(az storage account keys list -g $RESOURCE_GROUP --account-name $ACCOUNT_NAME --query "[0].value" -o tsv)
 
 if openshift-install coreos print-stream-json 2>/tmp/err.txt >/tmp/coreos.json; then
-  VHD_URL="$(jq -r '.architectures.x86_64."rhel-coreos-extensions"."azure-disk".url' /tmp/coreos.json)"
+  VHD_URL="$(jq -r --arg arch "$OCP_ARCH" '.architectures[$arch]."rhel-coreos-extensions"."azure-disk".url' /tmp/coreos.json)"
+  RELEASE_VERSION="$(jq -r --arg arch "$OCP_ARCH" '.architectures[$arch]."rhel-coreos-extensions"."azure-disk".release' /tmp/coreos.json)"
 else
   VHD_URL="$(jq -r .azure.url /var/lib/openshift-install/rhcos.json)"
+  RELEASE_VERSION="$(jq -r .azure.release /var/lib/openshift-install/rhcos.json)"
 fi
 
 echo "Copying VHD image from ${VHD_URL}"
@@ -218,10 +220,23 @@ az network private-dns link vnet create -g $RESOURCE_GROUP -z ${CLUSTER_NAME}.${
 
 echo "Deploying 02_storage"
 VHD_BLOB_URL=$(az storage blob url --account-name $ACCOUNT_NAME --account-key $ACCOUNT_KEY -c vhd -n "rhcos.vhd" -o tsv)
-az deployment group create -g $RESOURCE_GROUP \
-  --template-file "02_storage.json" \
-  --parameters vhdBlobURL="${VHD_BLOB_URL}" \
-  --parameters baseName="$INFRA_ID"
+
+# Check if it's the new template using Image Galleries instead of Managed Images
+if grep -qs "Microsoft.Compute/galleries" 02_storage.json; then
+  AZ_ARCH=$(echo $OCP_ARCH | sed 's/x86_64/x64/;s/aarch64/Arm64/')
+  az deployment group create -g $RESOURCE_GROUP \
+    --template-file "02_storage.json" \
+    --parameters vhdBlobURL="${VHD_BLOB_URL}" \
+    --parameters baseName="$INFRA_ID" \
+    --parameters storageAccount="$ACCOUNT_NAME" \
+    --parameters imageRelease="${RELEASE_VERSION::15}" \
+    --parameters architecture="$AZ_ARCH"
+else
+  az deployment group create -g $RESOURCE_GROUP \
+    --template-file "02_storage.json" \
+    --parameters vhdBlobURL="${VHD_BLOB_URL}" \
+    --parameters baseName="$INFRA_ID"
+fi
 
 echo "Deploying 03_infra"
 az deployment group create -g $RESOURCE_GROUP \
diff --git a/ci-operator/step-registry/upi/install/azure/upi-install-azure-ref.yaml b/ci-operator/step-registry/upi/install/azure/upi-install-azure-ref.yaml
index f0bbf3356f8b9..d7e5d7bea873a 100644
--- a/ci-operator/step-registry/upi/install/azure/upi-install-azure-ref.yaml
+++ b/ci-operator/step-registry/upi/install/azure/upi-install-azure-ref.yaml
@@ -7,5 +7,10 @@ ref:
     requests:
       cpu: 10m
       memory: 100Mi
+  env:
+  - name: OCP_ARCH
+    default: "x86_64"
+    documentation: |-
+      The architecture of the control plane nodes (e.g., x86_64, aarch64)
   documentation: >-
     This step deploys a UPI cluster to the CI Azure project.