diff --git a/src/ubuntu-generic-diskimages/BUILDING.md b/src/ubuntu-generic-disk-images/BUILDING.md
similarity index 100%
rename from src/ubuntu-generic-diskimages/BUILDING.md
rename to src/ubuntu-generic-disk-images/BUILDING.md
diff --git a/src/ubuntu-generic-diskimages/README.md b/src/ubuntu-generic-disk-images/README.md
similarity index 100%
rename from src/ubuntu-generic-diskimages/README.md
rename to src/ubuntu-generic-disk-images/README.md
diff --git a/src/ubuntu-generic-diskimages/build-arm.sh b/src/ubuntu-generic-disk-images/build-arm.sh
similarity index 65%
rename from src/ubuntu-generic-diskimages/build-arm.sh
rename to src/ubuntu-generic-disk-images/build-arm.sh
index 2ddf6320d..3a5f9882a 100755
--- a/src/ubuntu-generic-diskimages/build-arm.sh
+++ b/src/ubuntu-generic-disk-images/build-arm.sh
@@ -29,7 +29,21 @@ if [[ "$ubuntu_version" != "22.04" && "$ubuntu_version" != "24.04" ]]; then
     exit 1
 fi
 
-# make the flash0.sh file
+# Create a dictionary to store the path to the kernel modules for each Ubuntu version
+declare -A kernel_modules_paths
+kernel_modules_paths["22.04"]="modules/u2204"
+kernel_modules_paths["24.04"]="modules/u2404/"
+
+# Get the path to the kernel modules for the specified Ubuntu version
+kernel_modules_path=${kernel_modules_paths[$ubuntu_version]}
+pushd $kernel_modules_path
+./copy_modules.sh
+popd
+
+# Store the image name from the second command line argument or default to "arm-ubuntu"
+image_name="${2:-arm-ubuntu}"
+
+# make the flash0.img file
 cd ./files
 dd if=/dev/zero of=flash0.img bs=1M count=64
 dd if=/usr/share/qemu-efi-aarch64/QEMU_EFI.fd of=flash0.img conv=notrunc
@@ -39,4 +53,4 @@ cd ..
 ./packer init ./packer-scripts/arm-ubuntu.pkr.hcl
 
 # Build the image with the specified Ubuntu version
-./packer build -var "ubuntu_version=${ubuntu_version}" ./packer-scripts/arm-ubuntu.pkr.hcl
+./packer build -var "ubuntu_version=${ubuntu_version}" -var "image_name=${image_name}" ./packer-scripts/arm-ubuntu.pkr.hcl
diff --git a/src/ubuntu-generic-diskimages/build-x86.sh b/src/ubuntu-generic-disk-images/build-x86.sh
similarity index 80%
rename from src/ubuntu-generic-diskimages/build-x86.sh
rename to src/ubuntu-generic-disk-images/build-x86.sh
index 52e0adb77..13bb892fa 100755
--- a/src/ubuntu-generic-diskimages/build-x86.sh
+++ b/src/ubuntu-generic-disk-images/build-x86.sh
@@ -27,8 +27,11 @@ if [[ "$ubuntu_version" != "22.04" && "$ubuntu_version" != "24.04" ]]; then
     exit 1
 fi
 
+# Store the image name from the second command line argument or default to "x86-ubuntu"
+image_name="${2:-x86-ubuntu}"
+
 # Install the needed plugins
 ./packer init ./packer-scripts/x86-ubuntu.pkr.hcl
 
 # Build the image with the specified Ubuntu version
-./packer build -var "ubuntu_version=${ubuntu_version}" ./packer-scripts/x86-ubuntu.pkr.hcl
+./packer build -var "ubuntu_version=${ubuntu_version}" -var "image_name=${image_name}" ./packer-scripts/x86-ubuntu.pkr.hcl
diff --git a/src/ubuntu-generic-diskimages/files/arm/after_boot.sh b/src/ubuntu-generic-disk-images/files/arm/after_boot.sh
similarity index 100%
rename from src/ubuntu-generic-diskimages/files/arm/after_boot.sh
rename to src/ubuntu-generic-disk-images/files/arm/after_boot.sh
diff --git a/src/ubuntu-generic-diskimages/files/arm/gem5_init.sh b/src/ubuntu-generic-disk-images/files/arm/gem5_init.sh
similarity index 63%
rename from src/ubuntu-generic-diskimages/files/arm/gem5_init.sh
rename to src/ubuntu-generic-disk-images/files/arm/gem5_init.sh
index 0c13e97ad..cf05f64c7 100755
--- a/src/ubuntu-generic-diskimages/files/arm/gem5_init.sh
+++ b/src/ubuntu-generic-disk-images/files/arm/gem5_init.sh
@@ -15,6 +15,25 @@ mount -t sysfs /sys /sys
 cmdline=$(cat /proc/cmdline)
 no_systemd=false
 
+# Load gem5_bridge driver
+## Default parameters (ARM64)
+gem5_bridge_baseaddr=0x10010000
+gem5_bridge_rangesize=0x10000
+## Try to read overloads from kernel arguments
+if [[ $cmdline =~ gem5_bridge_baseaddr=([[:alnum:]]+) ]]; then
+    gem5_bridge_baseaddr=${BASH_REMATCH[1]}
+fi
+if [[ $cmdline =~ gem5_bridge_rangesize=([[:alnum:]]+) ]]; then
+    gem5_bridge_rangesize=${BASH_REMATCH[1]}
+fi
+## Insert driver
+modprobe gem5_bridge \
+    gem5_bridge_baseaddr=$gem5_bridge_baseaddr \
+    gem5_bridge_rangesize=$gem5_bridge_rangesize
+
+# see if this modprode fails or not
+# print warning if it fails, gem5-bridge module is not going to work, you will need sudo for running exit events
+
 # gem5-bridge exit signifying that kernel is booted
 # This will cause the simulation to exit. Note that this will
 # cause qemu to fail.
diff --git a/src/ubuntu-generic-disk-images/files/serial-getty@.service-override.conf b/src/ubuntu-generic-disk-images/files/serial-getty@.service-override.conf
new file mode 100644
index 000000000..34ddc51c4
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/files/serial-getty@.service-override.conf
@@ -0,0 +1,3 @@
+[Service]
+ExecStart=
+ExecStart=-/sbin/agetty --autologin gem5 --keep-baud 115200,38400,9600 %I $TERM
\ No newline at end of file
diff --git a/src/ubuntu-generic-diskimages/files/x86/after_boot.sh b/src/ubuntu-generic-disk-images/files/x86/after_boot.sh
similarity index 100%
rename from src/ubuntu-generic-diskimages/files/x86/after_boot.sh
rename to src/ubuntu-generic-disk-images/files/x86/after_boot.sh
diff --git a/src/ubuntu-generic-diskimages/files/x86/gem5_init.sh b/src/ubuntu-generic-disk-images/files/x86/gem5_init.sh
similarity index 100%
rename from src/ubuntu-generic-diskimages/files/x86/gem5_init.sh
rename to src/ubuntu-generic-disk-images/files/x86/gem5_init.sh
diff --git a/src/ubuntu-generic-diskimages/http/arm-22-04/meta-data b/src/ubuntu-generic-disk-images/http/arm-2204/meta-data
similarity index 100%
rename from src/ubuntu-generic-diskimages/http/arm-22-04/meta-data
rename to src/ubuntu-generic-disk-images/http/arm-2204/meta-data
diff --git a/src/ubuntu-generic-diskimages/http/arm-22-04/user-data b/src/ubuntu-generic-disk-images/http/arm-2204/user-data
similarity index 98%
rename from src/ubuntu-generic-diskimages/http/arm-22-04/user-data
rename to src/ubuntu-generic-disk-images/http/arm-2204/user-data
index b82387329..47d90d512 100644
--- a/src/ubuntu-generic-diskimages/http/arm-22-04/user-data
+++ b/src/ubuntu-generic-disk-images/http/arm-2204/user-data
@@ -70,7 +70,7 @@ autoinstall:
       type: format
       id: format-0
     - device: disk-vda
-      size: 4257218560
+      size: 20909654016
       wipe: superblock
       flag: ''
       number: 2
@@ -92,4 +92,4 @@ autoinstall:
       type: mount
       id: mount-0
   updates: security
-  version: 1
\ No newline at end of file
+  version: 1
diff --git a/src/ubuntu-generic-diskimages/http/arm-24-04/meta-data b/src/ubuntu-generic-disk-images/http/arm-2404/meta-data
similarity index 100%
rename from src/ubuntu-generic-diskimages/http/arm-24-04/meta-data
rename to src/ubuntu-generic-disk-images/http/arm-2404/meta-data
diff --git a/src/ubuntu-generic-diskimages/http/arm-24-04/user-data b/src/ubuntu-generic-disk-images/http/arm-2404/user-data
similarity index 97%
rename from src/ubuntu-generic-diskimages/http/arm-24-04/user-data
rename to src/ubuntu-generic-disk-images/http/arm-2404/user-data
index 895062bbb..81a9a3649 100644
--- a/src/ubuntu-generic-diskimages/http/arm-24-04/user-data
+++ b/src/ubuntu-generic-disk-images/http/arm-2404/user-data
@@ -62,7 +62,7 @@ autoinstall:
       type: format
       id: format-0
     - device: disk-vda
-      size: 4257218560
+      size: 20909654016
       wipe: superblock
       flag: ''
       number: 2
@@ -84,4 +84,4 @@ autoinstall:
       type: mount
       id: mount-0
   updates: security
-  version: 1
\ No newline at end of file
+  version: 1
diff --git a/src/ubuntu-generic-diskimages/http/x86/meta-data b/src/ubuntu-generic-disk-images/http/x86/meta-data
similarity index 100%
rename from src/ubuntu-generic-diskimages/http/x86/meta-data
rename to src/ubuntu-generic-disk-images/http/x86/meta-data
diff --git a/src/ubuntu-generic-diskimages/http/x86/user-data b/src/ubuntu-generic-disk-images/http/x86/user-data
similarity index 100%
rename from src/ubuntu-generic-diskimages/http/x86/user-data
rename to src/ubuntu-generic-disk-images/http/x86/user-data
diff --git a/src/ubuntu-generic-disk-images/modules/u2204/Dockerfile b/src/ubuntu-generic-disk-images/modules/u2204/Dockerfile
new file mode 100644
index 000000000..cfd0bf93e
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/modules/u2204/Dockerfile
@@ -0,0 +1,46 @@
+FROM ubuntu:22.04
+
+# Install necessary packages for kernel build
+RUN apt update && apt install -y \
+    build-essential \
+    libncurses-dev \
+    bison \
+    flex \
+    libssl-dev \
+    libelf-dev \
+    bc \
+    wget \
+    git \
+    kmod \
+    apt-src \
+    vim \
+    curl \
+    file
+
+RUN sed -Ei 's/^# deb-src /deb-src /' /etc/apt/sources.list
+RUN apt update
+RUN mkdir /workspace
+RUN cd /workspace && apt source linux-image-unsigned-5.15.0-25-generic
+
+RUN cd /workspace/linux-5.15.0 && \
+    cd scripts && \
+    chmod +x pahole-version.sh && \
+    cd .. && \
+    make defconfig && \
+    make -j 32 && \
+    make INSTALL_MOD_PATH=/workspace/output modules_install
+
+RUN git clone https://github.com/nkrim/gem5.git --depth=1 --filter=blob:none --no-checkout --sparse --single-branch --branch=gem5-bridge && \
+    cd gem5 && \
+    git sparse-checkout add util/m5 && \
+    git sparse-checkout add util/gem5_bridge && \
+    git sparse-checkout add include && \
+    git checkout
+
+RUN cd gem5/util/gem5_bridge && \
+    make KMAKEDIR=/workspace/linux-5.15.0 INSTALL_MOD_PATH=/workspace/output build install
+
+RUN cd /workspace/output/lib/modules/5.15.167 && \
+    rm -rf build source
+
+CMD ["/bin/bash"]
diff --git a/src/ubuntu-generic-disk-images/modules/u2204/copy_modules.sh b/src/ubuntu-generic-disk-images/modules/u2204/copy_modules.sh
new file mode 100755
index 000000000..1532f82ca
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/modules/u2204/copy_modules.sh
@@ -0,0 +1,47 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+# Variables
+IMAGE_NAME=ms-u2204-kernel-modules-source-image
+CONTAINER_NAME=ms-u2204-kernel-modules-source-container
+
+# Check if the image exists
+if ! docker image inspect "$IMAGE_NAME" > /dev/null 2>&1; then
+    echo "Docker image $IMAGE_NAME not found. Building the image..."
+
+    # Replace with the actual Docker build command
+    docker build -t $IMAGE_NAME .
+
+    if [ $? -ne 0 ]; then
+        echo "Failed to build the Docker image."
+        exit 1
+    fi
+else
+    echo "Docker image $IMAGE_NAME exists."
+fi
+
+# Check if the container already exists and stop/remove it if necessary
+if docker ps -a --format '{{.Names}}' | grep -q "^$CONTAINER_NAME\$"; then
+    echo "Container $CONTAINER_NAME already exists. Removing it..."
+    docker rm -f $CONTAINER_NAME
+fi
+
+
+# Create and start the container
+echo "Creating and starting container $CONTAINER_NAME..."
+docker create --name $CONTAINER_NAME $IMAGE_NAME
+docker start -a $CONTAINER_NAME
+
+# Ensure the target directory exists
+mkdir -p files
+
+# Copy files from the container to the host
+docker cp $CONTAINER_NAME:/workspace/linux-5.15.0/vmlinux files/
+docker cp $CONTAINER_NAME:/workspace/output/lib/modules/5.15.167 files/
+
+# Stop and remove the container after the copy operation
+echo "Stopping and removing the container..."
+docker stop "$CONTAINER_NAME"
+docker rm "$CONTAINER_NAME"
diff --git a/src/ubuntu-generic-disk-images/modules/u2404/Dockerfile b/src/ubuntu-generic-disk-images/modules/u2404/Dockerfile
new file mode 100644
index 000000000..df925db00
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/modules/u2404/Dockerfile
@@ -0,0 +1,42 @@
+# Start from Ubuntu 24.04 base image
+FROM ubuntu:24.04
+
+# Install necessary packages for kernel and module build
+RUN apt update && apt install -y \
+    build-essential \
+    libncurses-dev \
+    bison \
+    flex \
+    libssl-dev \
+    libelf-dev \
+    bc \
+    wget \
+    git \
+    kmod
+
+
+RUN sed -i 's/^Types: deb$/Types: deb deb-src/' /etc/apt/sources.list.d/ubuntu.sources
+RUN apt update
+RUN mkdir /workspace
+RUN cd /workspace && apt source linux-image-unsigned-6.8.0-47-generic
+
+RUN cd /workspace/linux-6.8.0 && \
+    make defconfig && \
+    make -j 32 && \
+    make INSTALL_MOD_PATH=/workspace/output modules_install
+
+RUN git clone https://github.com/nkrim/gem5.git --depth=1 --filter=blob:none --no-checkout --sparse --single-branch --branch=gem5-bridge && \
+    cd gem5 && \
+    git sparse-checkout add util/m5 && \
+    git sparse-checkout add util/gem5_bridge && \
+    git sparse-checkout add include && \
+    git checkout
+
+RUN cd gem5/util/gem5_bridge && \
+    make KMAKEDIR=/workspace/linux-6.8.0 INSTALL_MOD_PATH=/workspace/output build install
+
+RUN cd /workspace/output/lib/modules/6.8.12 && \
+    rm -rf build
+
+# Use bash as the default shell
+CMD ["/bin/bash"]
diff --git a/src/ubuntu-generic-disk-images/modules/u2404/copy_modules.sh b/src/ubuntu-generic-disk-images/modules/u2404/copy_modules.sh
new file mode 100755
index 000000000..4f2293645
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/modules/u2404/copy_modules.sh
@@ -0,0 +1,47 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+# Variables
+IMAGE_NAME=ms-u2404-kernel-modules-source-image
+CONTAINER_NAME=ms-u2404-kernel-modules-source-container
+
+# Check if the image exists
+if ! docker image inspect "$IMAGE_NAME" > /dev/null 2>&1; then
+    echo "Docker image $IMAGE_NAME not found. Building the image..."
+
+    # Replace with the actual Docker build command
+    docker build -t $IMAGE_NAME .
+
+    if [ $? -ne 0 ]; then
+        echo "Failed to build the Docker image."
+        exit 1
+    fi
+else
+    echo "Docker image $IMAGE_NAME exists."
+fi
+
+# Check if the container already exists and stop/remove it if necessary
+if docker ps -a --format '{{.Names}}' | grep -q "^$CONTAINER_NAME\$"; then
+    echo "Container $CONTAINER_NAME already exists. Removing it..."
+    docker rm -f $CONTAINER_NAME
+fi
+
+
+# Create and start the container
+echo "Creating and starting container $CONTAINER_NAME..."
+docker create --name $CONTAINER_NAME $IMAGE_NAME
+docker start -a $CONTAINER_NAME
+
+# Ensure the target directory exists
+mkdir -p files
+
+# Copy files from the container to the host
+docker cp $CONTAINER_NAME:/workspace/linux-6.8.0/vmlinux files/
+docker cp $CONTAINER_NAME:/workspace/output/lib/modules/6.8.12 files/
+
+# Stop and remove the container after the copy operation
+echo "Stopping and removing the container..."
+docker stop "$CONTAINER_NAME"
+docker rm "$CONTAINER_NAME"
diff --git a/src/ubuntu-generic-diskimages/packer-scripts/arm-ubuntu.pkr.hcl b/src/ubuntu-generic-disk-images/packer-scripts/arm-ubuntu.pkr.hcl
similarity index 69%
rename from src/ubuntu-generic-diskimages/packer-scripts/arm-ubuntu.pkr.hcl
rename to src/ubuntu-generic-disk-images/packer-scripts/arm-ubuntu.pkr.hcl
index 6e81dd5e6..d237aa730 100644
--- a/src/ubuntu-generic-diskimages/packer-scripts/arm-ubuntu.pkr.hcl
+++ b/src/ubuntu-generic-disk-images/packer-scripts/arm-ubuntu.pkr.hcl
@@ -9,7 +9,6 @@ packer {
 
 variable "image_name" {
   type    = string
-  default = "arm-ubuntu"
 }
 
 variable "ssh_password" {
@@ -36,14 +35,18 @@ locals {
     "22.04" = {
       iso_url       = "https://old-releases.ubuntu.com/releases/jammy/ubuntu-22.04-live-server-arm64.iso"
       iso_checksum  = "sha256:c209ab013280d3cd26a344def60b7b19fbb427de904ea285057d94ca6ac82dd5"
-      output_dir    = "arm-disk-image-22-04"
-      http_directory = "http/arm-22-04"
+      output_dir    = "${var.image_name}-2204"
+      http_directory = "http/arm-2204"
+      kernel_version = "5.15.167"
+      modules_dir = "modules/u2204/files/"
     }
     "24.04" = {
-      iso_url       = "https://cdimage.ubuntu.com/releases/24.04/release/ubuntu-24.04-live-server-arm64.iso"
-      iso_checksum  = "sha256:d2d9986ada3864666e36a57634dfc97d17ad921fa44c56eeaca801e7dab08ad7"
-      output_dir    = "arm-disk-image-24-04"
-      http_directory = "http/arm-24-04"
+      iso_url       = "https://cdimage.ubuntu.com/releases/24.04/release/ubuntu-24.04.1-live-server-arm64.iso"
+      iso_checksum  = "sha256:5ceecb7ef5f976e8ab3fffee7871518c8e9927ec221a3bb548ee1193989e1773"
+      output_dir    = "${var.image_name}-2404"
+      http_directory = "http/arm-2404"
+      kernel_version = "6.8.12"
+      modules_dir = "modules/u2404/files/"
     }
   }
 }
@@ -91,14 +94,14 @@ source "qemu" "initialize" {
                       "<enter>",
                       "<wait>"
                       ]
-  cpus             = "4"
-  disk_size        = "4600"
+  cpus             = "32"
+  disk_size        = "21600"
   format           = "raw"
   headless         = "true"
   http_directory   = local.iso_data[var.ubuntu_version].http_directory
   iso_checksum     = local.iso_data[var.ubuntu_version].iso_checksum
   iso_urls         = [local.iso_data[var.ubuntu_version].iso_url]
-  memory           = "8192"
+  memory           = "65536"
   output_directory = local.iso_data[var.ubuntu_version].output_dir
   qemu_binary      = "/usr/bin/qemu-system-aarch64"
   qemuargs         = local.qemuargs
@@ -106,7 +109,7 @@ source "qemu" "initialize" {
   ssh_password     = "${var.ssh_password}"
   ssh_username     = "${var.ssh_username}"
   ssh_wait_timeout = "60m"
-  vm_name          = "${var.image_name}"
+  vm_name          = "disk-image"
   ssh_handshake_attempts = "1000"
 }
 
@@ -115,28 +118,39 @@ build {
 
   provisioner "file" {
     destination = "/home/gem5/"
-    source      = "files/exit.sh"
+    source      = "files/arm/gem5_init.sh"
   }
 
   provisioner "file" {
     destination = "/home/gem5/"
-    source      = "files/arm/gem5_init.sh"
+    source      = "files/arm/after_boot.sh"
   }
 
   provisioner "file" {
     destination = "/home/gem5/"
-    source      = "files/arm/after_boot.sh"
+    source      = "files/serial-getty@.service-override.conf"
   }
 
   provisioner "file" {
     destination = "/home/gem5/"
-    source      = "files/serial-getty@.service"
+    source      = "${local.iso_data[var.ubuntu_version].modules_dir}/${local.iso_data[var.ubuntu_version].kernel_version}"
   }
 
   provisioner "shell" {
     execute_command = "echo '${var.ssh_password}' | {{ .Vars }} sudo -E -S bash '{{ .Path }}'"
-    scripts         = ["scripts/post-installation.sh"]
-    environment_vars = ["ISA=arm64"]
+    scripts         = ["scripts/install-commons.sh", "scripts/update-modules.sh", "scripts/install-gem5-bridge-driver.sh", "scripts/install-user-packages.sh"]
+    environment_vars = ["ISA=arm64", "KERNEL_VERSION=${local.iso_data[var.ubuntu_version].kernel_version}", "DEBIAN_FRONTEND=noninteractive"]
     expect_disconnect = true
   }
+
+  provisioner "shell" {
+    scripts = ["scripts/install-benchmarks.sh"]
+  }
+
+  provisioner "shell" {
+    execute_command = "echo '${var.ssh_password}' | {{ .Vars }} sudo -E -S bash '{{ .Path }}'"
+    scripts         = ["scripts/install-gem5-init.sh", "scripts/disable-network.sh"]
+    expect_disconnect = true
+  }
+
 }
diff --git a/src/ubuntu-generic-diskimages/packer-scripts/x86-ubuntu.pkr.hcl b/src/ubuntu-generic-disk-images/packer-scripts/x86-ubuntu.pkr.hcl
similarity index 77%
rename from src/ubuntu-generic-diskimages/packer-scripts/x86-ubuntu.pkr.hcl
rename to src/ubuntu-generic-disk-images/packer-scripts/x86-ubuntu.pkr.hcl
index 3d88dce63..9f5c875ae 100644
--- a/src/ubuntu-generic-diskimages/packer-scripts/x86-ubuntu.pkr.hcl
+++ b/src/ubuntu-generic-disk-images/packer-scripts/x86-ubuntu.pkr.hcl
@@ -9,7 +9,6 @@ packer {
 
 variable "image_name" {
   type    = string
-  default = "x86-ubuntu"
 }
 
 variable "ssh_password" {
@@ -36,12 +35,12 @@ locals {
     "22.04" = {
       iso_url       = "https://old-releases.ubuntu.com/releases/jammy/ubuntu-22.04.2-live-server-amd64.iso"
       iso_checksum  = "sha256:5e38b55d57d94ff029719342357325ed3bda38fa80054f9330dc789cd2d43931"
-      output_dir    = "x86-disk-image-22-04"
+      output_dir    = "${var.image_name}-2204"
     }
     "24.04" = {
       iso_url       = "https://old-releases.ubuntu.com/releases/noble/ubuntu-24.04-live-server-amd64.iso"
       iso_checksum  = "sha256:8762f7e74e4d64d72fceb5f70682e6b069932deedb4949c6975d0f0fe0a91be3"
-      output_dir    = "x86-disk-image-24-04"
+      output_dir    = "${var.image_name}-2204"
     }
   }
 }
@@ -55,7 +54,7 @@ source "qemu" "initialize" {
                       "<f10><wait>"
                     ]
   cpus             = "4"
-  disk_size        = "5000"
+  disk_size        = "21600"
   format           = "raw"
   headless         = "true"
   http_directory   = "http/x86"
@@ -69,18 +68,13 @@ source "qemu" "initialize" {
   ssh_password     = "${var.ssh_password}"
   ssh_username     = "${var.ssh_username}"
   ssh_wait_timeout = "60m"
-  vm_name          = "${var.image_name}"
+  vm_name          = "disk-image"
   ssh_handshake_attempts = "1000"
 }
 
 build {
   sources = ["source.qemu.initialize"]
 
-  provisioner "file" {
-    destination = "/home/gem5/"
-    source      = "files/exit.sh"
-  }
-
   provisioner "file" {
     destination = "/home/gem5/"
     source      = "files/x86/gem5_init.sh"
@@ -93,15 +87,25 @@ build {
 
   provisioner "file" {
     destination = "/home/gem5/"
-    source      = "files/serial-getty@.service"
+    source      = "files/serial-getty@.service-override.conf"
+  }
+
+  provisioner "shell" {
+    execute_command = "echo '${var.ssh_password}' | {{ .Vars }} sudo -E -S bash '{{ .Path }}'"
+    scripts         = ["scripts/install-commons.sh", "scripts/extract-kernel.sh", "scripts/install-gem5-bridge-driver.sh", "scripts/install-user-packages.sh"]
+    environment_vars = ["ISA=x86", "DEBIAN_FRONTEND=noninteractive"]
+    expect_disconnect = true
+  }
+
+  provisioner "shell" {
+    scripts = ["scripts/install-benchmarks.sh"]
   }
 
   provisioner "shell" {
     execute_command = "echo '${var.ssh_password}' | {{ .Vars }} sudo -E -S bash '{{ .Path }}'"
-    scripts         = ["scripts/post-installation.sh"]
-    environment_vars = ["ISA=x86"]
+    scripts         = ["scripts/install-gem5-init.sh", "scripts/disable-network.sh"]
+    expect_disconnect = true
   }
-  
 
   provisioner "file" {
   source      = "/home/gem5/vmlinux-x86-ubuntu"
diff --git a/src/ubuntu-generic-disk-images/scripts/disable-network.sh b/src/ubuntu-generic-disk-images/scripts/disable-network.sh
new file mode 100755
index 000000000..314c71661
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/scripts/disable-network.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+# Disable network by default
+echo "Disabling network by default."
+echo "See README.md for instructions on how to enable network."
+if [ -f /etc/netplan/50-cloud-init.yaml ]; then
+    mv /etc/netplan/50-cloud-init.yaml /etc/netplan/50-cloud-init.yaml.bak
+elif [ -f /etc/netplan/00-installer-config.yaml ]; then
+    mv /etc/netplan/00-installer-config.yaml /etc/netplan/00-installer-config.yaml.bak
+    netplan apply
+fi
+# Disable systemd service that waits for network to be online
+systemctl disable systemd-networkd-wait-online.service
+systemctl mask systemd-networkd-wait-online.service
+
+echo "Network disabled!"
diff --git a/src/ubuntu-generic-disk-images/scripts/extract-kernel.sh b/src/ubuntu-generic-disk-images/scripts/extract-kernel.sh
new file mode 100755
index 000000000..f47e4a2fe
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/scripts/extract-kernel.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+# Check if ISA is set
+if [ -z "$ISA" ]; then
+    echo "Error: ISA is not set"
+    exit 1
+fi
+
+# Check if ISA is x86
+if [ "$ISA" != "x86" ]; then
+    echo "Error: ISA is not x86"
+    exit 1
+fi
+
+# Make sure the headers are installed to extract the kernel that DKMS
+# packages will be built against.
+sudo apt -y install "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)"
+
+echo "Extracting linux kernel $(uname -r) to /home/gem5/vmlinux-x86-ubuntu"
+sudo bash -c "/usr/src/linux-headers-$(uname -r)/scripts/extract-vmlinux /boot/vmlinuz-$(uname -r) > /home/gem5/vmlinux-x86-ubuntu"
diff --git a/src/ubuntu-generic-disk-images/scripts/install-benchmarks.sh b/src/ubuntu-generic-disk-images/scripts/install-benchmarks.sh
new file mode 100755
index 000000000..2e260d62c
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/scripts/install-benchmarks.sh
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+echo "Installing benchmarks..."
+
+# You should write this script with the following assumptions:
+# 1. common-packages and user-packages are installed; they are specified in
+#    install-common-packages.sh and install-user-packages.sh, respectively.
+# 2. This script is run without sudo access.
+# 3. All your files that you specify in your .pkr.hcl file are copied into
+#    the disk.
+
+echo "Done installing benchmarks."
\ No newline at end of file
diff --git a/src/ubuntu-generic-disk-images/scripts/install-commons.sh b/src/ubuntu-generic-disk-images/scripts/install-commons.sh
new file mode 100755
index 000000000..f77320984
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/scripts/install-commons.sh
@@ -0,0 +1,32 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+echo "Common Installation Started."
+
+
+echo "Installing serial service override for autologin after systemd."
+mkdir /etc/systemd/system/serial-getty@.service.d/
+mv serial-getty@.service-override.conf /etc/systemd/system/serial-getty@.service.d/override.conf
+
+# Installing the packages in this script instead of the user-data
+# file dueing ubuntu autoinstall. The reason is that sometimes
+# the package install failes. This method is more reliable.
+
+echo "Installing packages required for gem5-bridge (m5) and libm5."
+apt-get update
+apt-get install -y scons
+apt-get install -y git
+apt-get install -y vim
+apt-get install -y build-essential
+
+# Add after_boot.sh to bashrc in the gem5 user account
+# This will run the script after the user automatically logs in
+echo "Adding after_boot.sh to the gem5 user's .bashrc."
+echo -e "\nif [ -z \"\$AFTER_BOOT_EXECUTED\" ]; then\n   export AFTER_BOOT_EXECUTED=1\n    /home/gem5/after_boot.sh\nfi\n" >> /home/gem5/.bashrc
+
+# Remove the motd
+rm /etc/update-motd.d/*
+
+echo "Common Installation Done."
diff --git a/src/ubuntu-generic-disk-images/scripts/install-gem5-bridge-driver.sh b/src/ubuntu-generic-disk-images/scripts/install-gem5-bridge-driver.sh
new file mode 100755
index 000000000..09a7cda5b
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/scripts/install-gem5-bridge-driver.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+# Build and install the gem5-bridge (m5) binary, library, and headers
+echo "Building and installing gem5-bridge (m5) and libm5."
+
+# Ensure the ISA environment variable is set
+if [ -z "$ISA" ]; then
+  echo "Error: ISA environment variable is not set."
+  exit 1
+fi
+
+# Just get the files we need
+git clone https://github.com/nkrim/gem5.git --depth=1 --filter=blob:none --no-checkout --sparse --single-branch --branch=gem5-bridge
+pushd gem5
+# Checkout just the files we need
+git sparse-checkout add util/m5
+git sparse-checkout add util/gem5_bridge
+git sparse-checkout add include
+git checkout
+# Install the headers globally so that other benchmarks can use them
+cp -r include/gem5 /usr/local/include/
+
+# Build the library and binary
+pushd util/m5
+scons build/${ISA}/out/m5
+cp build/${ISA}/out/m5 /usr/local/bin/
+cp build/${ISA}/out/libm5.a /usr/local/lib/
+popd   # util/m5
+popd   # gem5
+
+# rename the m5 binary to gem5-bridge
+mv /usr/local/bin/m5 /usr/local/bin/gem5-bridge
+# Set the setuid bit on the m5 binary
+chmod 4755 /usr/local/bin/gem5-bridge
+chmod u+s /usr/local/bin/gem5-bridge
+
+#create a symbolic link to the gem5 binary for backward compatibility
+ln -s /usr/local/bin/gem5-bridge /usr/local/bin/m5
+
+# delete the git repo for gem5
+rm -rf gem5
+echo "Done building and installing gem5-bridge (m5) and libm5."
\ No newline at end of file
diff --git a/src/ubuntu-generic-disk-images/scripts/install-gem5-init.sh b/src/ubuntu-generic-disk-images/scripts/install-gem5-init.sh
new file mode 100755
index 000000000..38bf1c4e3
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/scripts/install-gem5-init.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+echo "Installing the gem5 init script in /sbin."
+
+mv /home/gem5/gem5_init.sh /sbin/
+mv /sbin/init /sbin/init.gem5.bak
+ln -s /sbin/gem5_init.sh /sbin/init
+
+echo "Done installing the gem5 init script as the init process."
diff --git a/src/ubuntu-generic-disk-images/scripts/install-user-packages.sh b/src/ubuntu-generic-disk-images/scripts/install-user-packages.sh
new file mode 100755
index 000000000..eb41d30a6
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/scripts/install-user-packages.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+echo "Installing user packages..."
+
+# Put your package installation commands here
+# NOTE: Since we use packer we can not interact with the process of apt-get
+# upgrade/install, I recommend that you add `Dpkg::Options::="--force-confnew"`
+# to your commands here. See example below on how to update, upgrade, and
+# install packages for cmake, python3-pip, gfortran, and openmpi, ...
+# sudo apt-get -y update
+# sudo apt-get -y -o Dpkg::Options::="--force-confnew" upgrade
+# sudo apt-get -y -o Dpkg::Options::="--force-confnew" install \
+# python3-pip gfortran cmake openmpi-bin libopenmpi-dev
+
+echo "Done installing user packages."
diff --git a/src/ubuntu-generic-disk-images/scripts/update-modules.sh b/src/ubuntu-generic-disk-images/scripts/update-modules.sh
new file mode 100755
index 000000000..7700c2e45
--- /dev/null
+++ b/src/ubuntu-generic-disk-images/scripts/update-modules.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+# Copyright (c) 2025 The Regents of the University of California.
+# SPDX-License-Identifier: BSD 3-Clause
+
+echo "Updating modules... This step is only necessary for arm64."
+
+# Ensure the ISA environment variable is set
+if [ -z "$ISA" ]; then
+    echo "Error: ISA environment variable is not set."
+    exit 1
+fi
+
+# Ensure the ISA environment variable is set to arm64
+if [ "$ISA" != "arm64" ]; then
+    echo "Error: This script should only be used for arm64 isa."
+    exit 1
+fi
+
+# Ensure the KERNEL_VERSION environment variable is set
+if [ -z "$KERNEL_VERSION" ]; then
+  echo "Error: KERNEL_VERSION environment variable is not set."
+  exit 1
+fi
+
+# moving modules to the correct location
+mv /home/gem5/$KERNEL_VERSION /lib/modules/$KERNEL_VERSION
+depmod --quick -a $KERNEL_VERSION
+update-initramfs -u -k $KERNEL_VERSION
+
+echo "Done updating modules."
diff --git a/src/ubuntu-generic-diskimages/.gitignore b/src/ubuntu-generic-diskimages/.gitignore
deleted file mode 100644
index 6fa003226..000000000
--- a/src/ubuntu-generic-diskimages/.gitignore
+++ /dev/null
@@ -1,7 +0,0 @@
-**packer
-**packer_*.zip
-packer_cache/
-x86-disk-image-22-04/
-x86-disk-image-24-04/
-arm-disk-image-22-04/
-arm-disk-image-24-04/
diff --git a/src/ubuntu-generic-diskimages/files/exit.sh b/src/ubuntu-generic-diskimages/files/exit.sh
deleted file mode 100755
index 4b147e01f..000000000
--- a/src/ubuntu-generic-diskimages/files/exit.sh
+++ /dev/null
@@ -1,6 +0,0 @@
-#!/bin/sh
-
-# Copyright (c) 2020 The Regents of the University of California.
-# SPDX-License-Identifier: BSD 3-Clause
-
-m5 exit
diff --git a/src/ubuntu-generic-diskimages/files/serial-getty@.service b/src/ubuntu-generic-diskimages/files/serial-getty@.service
deleted file mode 100644
index 121e8e355..000000000
--- a/src/ubuntu-generic-diskimages/files/serial-getty@.service
+++ /dev/null
@@ -1,46 +0,0 @@
-#  SPDX-License-Identifier: LGPL-2.1+
-#
-#  This file is part of systemd.
-#
-#  systemd is free software; you can redistribute it and/or modify it
-#  under the terms of the GNU Lesser General Public License as published by
-#  the Free Software Foundation; either version 2.1 of the License, or
-#  (at your option) any later version.
-
-[Unit]
-Description=Serial Getty on %I
-Documentation=man:agetty(8) man:systemd-getty-generator(8)
-Documentation=http://0pointer.de/blog/projects/serial-console.html
-BindsTo=dev-%i.device
-After=dev-%i.device systemd-user-sessions.service plymouth-quit-wait.service getty-pre.target
-After=rc-local.service
-
-# If additional gettys are spawned during boot then we should make
-# sure that this is synchronized before getty.target, even though
-# getty.target didn't actually pull it in.
-Before=getty.target
-IgnoreOnIsolate=yes
-
-# IgnoreOnIsolate causes issues with sulogin, if someone isolates
-# rescue.target or starts rescue.service from multi-user.target or
-# graphical.target.
-Conflicts=rescue.service
-Before=rescue.service
-
-[Service]
-# The '-o' option value tells agetty to replace 'login' arguments with an
-# option to preserve environment (-p), followed by '--' for safety, and then
-# the entered username.
-ExecStart=-/sbin/agetty --autologin gem5 --keep-baud 115200,38400,9600 %I $TERM
-Type=idle
-Restart=always
-UtmpIdentifier=%I
-TTYPath=/dev/%I
-TTYReset=yes
-TTYVHangup=yes
-KillMode=process
-IgnoreSIGPIPE=no
-SendSIGHUP=yes
-
-[Install]
-WantedBy=getty.target
diff --git a/src/ubuntu-generic-diskimages/scripts/post-installation.sh b/src/ubuntu-generic-diskimages/scripts/post-installation.sh
deleted file mode 100755
index 1f37ac98b..000000000
--- a/src/ubuntu-generic-diskimages/scripts/post-installation.sh
+++ /dev/null
@@ -1,126 +0,0 @@
-#!/bin/bash
-
-# Copyright (c) 2024 The Regents of the University of California.
-# SPDX-License-Identifier: BSD 3-Clause
-
-echo 'Post Installation Started'
-
-# Installing the packages in this script instead of the user-data
-# file dueing ubuntu autoinstall. The reason is that sometimes
-# the package install failes. This method is more reliable.
-echo 'installing packages'
-apt-get update
-apt-get install -y scons
-apt-get install -y git
-apt-get install -y vim
-apt-get install -y build-essential
-
-echo "Installing serial service for autologin after systemd"
-mv /home/gem5/serial-getty@.service /lib/systemd/system/
-
-# Make sure the headers are installed to extract the kernel that DKMS
-# packages will be built against.
-sudo apt -y install "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)"
-
-echo "Extracting linux kernel $(uname -r) to /home/gem5/vmlinux-x86-ubuntu"
-sudo bash -c "/usr/src/linux-headers-$(uname -r)/scripts/extract-vmlinux /boot/vmlinuz-$(uname -r) > /home/gem5/vmlinux-x86-ubuntu"
-
-echo "Installing the gem5 init script in /sbin"
-mv /home/gem5/gem5_init.sh /sbin
-mv /sbin/init /sbin/init.old
-ln -s /sbin/gem5_init.sh /sbin/init
-
-# Add after_boot.sh to bashrc in the gem5 user account
-# This will run the script after the user automatically logs in
-echo -e "\nif [ -z \"\$AFTER_BOOT_EXECUTED\" ]; then\n   export AFTER_BOOT_EXECUTED=1\n    /home/gem5/after_boot.sh\nfi\n" >> /home/gem5/.bashrc
-
-# Remove the motd
-rm /etc/update-motd.d/*
-
-# Build and install the gem5-bridge (m5) binary, library, and headers
-echo "Building and installing gem5-bridge (m5) and libm5"
-
-# Ensure the ISA environment variable is set
-if [ -z "$ISA" ]; then
-  echo "Error: ISA environment variable is not set."
-  exit 1
-fi
-
-# Just get the files we need
-git clone https://github.com/gem5/gem5.git --depth=1 --filter=blob:none --no-checkout --sparse --single-branch --branch=stable
-pushd gem5
-# Checkout just the files we need
-git sparse-checkout add util/m5
-git sparse-checkout add include
-git checkout
-# Install the headers globally so that other benchmarks can use them
-cp -r include/gem5 /usr/local/include/\
-
-# Build the library and binary
-pushd util/m5
-scons build/${ISA}/out/m5
-cp build/${ISA}/out/m5 /usr/local/bin/
-cp build/${ISA}/out/libm5.a /usr/local/lib/
-popd   # util/m5
-popd   # gem5
-
-# rename the m5 binary to gem5-bridge
-mv /usr/local/bin/m5 /usr/local/bin/gem5-bridge
-# Set the setuid bit on the m5 binary
-chmod 4755 /usr/local/bin/gem5-bridge
-chmod u+s /usr/local/bin/gem5-bridge
-
-#create a symbolic link to the gem5 binary for backward compatibility
-ln -s /usr/local/bin/gem5-bridge /usr/local/bin/m5
-
-# delete the git repo for gem5
-rm -rf gem5
-echo "Done building and installing gem5-bridge (m5) and libm5"
-
-# You can extend this script to install your own packages here or by modifying the `x86-ubuntu.pkr.hcl`
-# or `arm-ubuntu.pkr.hcl` file depending on the disk you are building.
-
-# Disable network by default
-echo "Disabling network by default"
-echo "See README.md for instructions on how to enable network"
-if [ -f /etc/netplan/50-cloud-init.yaml ]; then
-    mv /etc/netplan/50-cloud-init.yaml /etc/netplan/50-cloud-init.yaml.bak
-elif [ -f /etc/netplan/00-installer-config.yaml ]; then
-    mv /etc/netplan/00-installer-config.yaml /etc/netplan/00-installer-config.yaml.bak
-    netplan apply
-fi
-# Disable systemd service that waits for network to be online
-systemctl disable systemd-networkd-wait-online.service
-systemctl mask systemd-networkd-wait-online.service
-
-if [ "${ISA}" = "x86" ]; then
-    echo "Disabling systemd services for x86 architecture..."
-
-    # Disable multipathd service
-    systemctl disable multipathd.service
-
-    # Disable thermald service
-    systemctl disable thermald.service
-
-    # Disable snapd services and socket
-    systemctl disable snapd.service snapd.socket
-
-    # Disable unnecessary timers
-    systemctl disable apt-daily.timer apt-daily-upgrade.timer fstrim.timer
-
-    # Disable accounts-daemon
-    systemctl disable accounts-daemon.service
-
-    # Disable LVM monitoring service
-    systemctl disable lvm2-monitor.service
-
-    # Switch default target to multi-user (no GUI)
-    systemctl set-default multi-user.target
-
-    # Optionally disable AppArmor if not required
-    systemctl disable apparmor.service snapd.apparmor.service
-
-    echo "completed disabling systemd services for x86."
-fi
-
-echo "Post Installation Done"