diff --git a/.github/workflows/build-and-release.yml b/.github/workflows/build-and-release.yml
index f826a24..a973acf 100644
--- a/.github/workflows/build-and-release.yml
+++ b/.github/workflows/build-and-release.yml
@@ -256,55 +256,181 @@ jobs:
           esac
           echo "Desktop task set to: $DESKTOP_TASK"
 
-      - name: Setup rpi-image-gen
+      - name: Setup pi-gen
         if: matrix.distro == 'bookworm' || matrix.distro == 'trixie'
         run: |
           # Verify submodule is checked out
-          if [ ! -f rpi-image-gen/rpi-image-gen ]; then
-            echo "ERROR: rpi-image-gen submodule not properly checked out"
+          if [ ! -d pi-gen ]; then
+            echo "ERROR: pi-gen submodule not properly checked out"
             exit 1
           fi
-          cd rpi-image-gen
-          sudo ./install_deps.sh
-          ./rpi-image-gen layer --list
-          sed -i 's/rpi5/rpi-cm4/g' ./config/${{ matrix.distro }}-minbase.yaml
-          sed -i '/^name:/a mode: unshare' ./config/${{ matrix.distro }}-minbase.yaml
-          # Increase root partition size to 800% for 8GB root filesystem
-          sed -i 's/root_part_size: 300%/root_part_size: 800%/g' ./config/${{ matrix.distro }}-minbase.yaml
-
+          
+          # Install dependencies
+          cd pi-gen
+          sudo apt-get update
+          sudo apt-get install -y coreutils quilt parted qemu-user-static debootstrap zerofree zip \
+            dosfstools e2fsprogs libarchive-tools libcap2-bin grep rsync xz-utils file git curl bc \
+            gpg pigz xxd arch-test bmap-tools kmod
 
       - name: Download kernel artifacts
         if: needs.prepare-kernel.outputs.kernel_mode == 'build'
         uses: actions/download-artifact@v4
         with:
           name: kernel-debs-${{ matrix.uconsole_core }}
-          path:  /home/runner/work/uConsole-Image-Builder/uConsole-Image-Builder/kernel-debs/
+          path:  /home/runner/work/uConsole-Image-Builder/uConsole-Image-Builder/artifacts/kernel-debs/
 
-      - name: Build base image
+      - name: Build image with pi-gen
         if: matrix.distro == 'bookworm' || matrix.distro == 'trixie'
         run: |
-          cd rpi-image-gen
-          ./rpi-image-gen build -c ./config/${{ matrix.distro }}-minbase.yaml
-
-      - name: Customize image
-        run: |
-          # If matrix.distro is bookworm, use deb12-arm64-min.img
-          # If matrix.distro is trixie, use deb13-arm64-min.img
-          # If matrix.distro is jammy, use jammy for pi 4 (arm64) wget https://cdimage.ubuntu.com/releases/jammy/release/ubuntu-22.04.5-preinstalled-desktop-arm64+raspi.img.xz
-          if [ "${{ matrix.distro }}" == "bookworm" ] || [ "${{ matrix.distro }}" == "trixie" ]; then            
-            # If matrix.distro is bookworm, 
-            if [ "${{ matrix.distro }}" == "bookworm" ]; then
-              IMAGE_PATH=/home/runner/work/uConsole-Image-Builder/uConsole-Image-Builder/rpi-image-gen/work/image-deb12-arm64-min/deb12-arm64-min.img
-            else
-              IMAGE_PATH=/home/runner/work/uConsole-Image-Builder/uConsole-Image-Builder/rpi-image-gen/work/image-deb13-arm64-min/deb13-arm64-min.img
-            fi
+          cd pi-gen
+          
+          # Determine stage list based on desktop environment
+          if [ "${{ matrix.desktop }}" == "gnome" ] || [ "${{ matrix.desktop }}" == "kde" ] || \
+             [ "${{ matrix.desktop }}" == "mate" ] || [ "${{ matrix.desktop }}" == "xfce" ] || \
+             [ "${{ matrix.desktop }}" == "lxde" ] || [ "${{ matrix.desktop }}" == "lxqt" ] || \
+             [ "${{ matrix.desktop }}" == "cinnamon" ] || [ "${{ matrix.desktop }}" == "gnome-flashback" ]; then
+            STAGE_LIST="stage0 stage1 stage2"
+          else
+            STAGE_LIST="stage0 stage1 stage2"
           fi
           
-          if [ "${{ matrix.distro }}" == "jammy" ]; then
-            IMAGE_PATH=https://cdimage.ubuntu.com/releases/jammy/release/ubuntu-22.04.5-preinstalled-server-arm64+raspi.img.xz
-            wget $IMAGE_PATH -O ubuntu-jammy-raspi.img.xz > /dev/null 2>&1
-            unxz ubuntu-jammy-raspi.img.xz
-            IMAGE_PATH=./ubuntu-jammy-raspi.img
+          # Create config file
+          cat > config << EOF
+          IMG_NAME="uconsole-${{ matrix.distro }}-${{ matrix.uconsole_core }}-${{ matrix.desktop }}"
+          RELEASE="${{ matrix.distro }}"
+          TARGET_HOSTNAME="uconsole"
+          FIRST_USER_NAME="clockworkpi"
+          FIRST_USER_PASS="clockworkpi"
+          DISABLE_FIRST_BOOT_USER_RENAME=1
+          ENABLE_SSH=1
+          WPA_COUNTRY="US"
+          LOCALE_DEFAULT="en_US.UTF-8"
+          TIMEZONE_DEFAULT="UTC"
+          KEYBOARD_KEYMAP="us"
+          KEYBOARD_LAYOUT="English (US)"
+          DEPLOY_COMPRESSION="xz"
+          COMPRESSION_LEVEL=6
+          STAGE_LIST="$STAGE_LIST"
+          EOF
+          
+          # Create custom kernel installation stage
+          CUSTOM_STAGE="stage2/06-uconsole-kernel"
+          mkdir -p "$CUSTOM_STAGE"
+          
+          if [ "${{ env.KERNEL_MODE }}" == "prebuilt" ]; then
+            # Install prebuilt kernel
+            cat > "$CUSTOM_STAGE/00-run-chroot.sh" << 'KERNELEOF'
+          #!/bin/bash -e
+          
+          echo "Installing ClockworkPi kernel from repository..."
+          
+          # Add ClockworkPi repository
+          wget -q -O- https://raw.githubusercontent.com/clockworkpi/apt/main/debian/KEY.gpg | gpg --dearmor | tee /etc/apt/trusted.gpg.d/clockworkpi.gpg
+          echo 'deb [arch=arm64] https://raw.githubusercontent.com/clockworkpi/apt/main/bookworm stable main' | tee /etc/apt/sources.list.d/clockworkpi.list
+          
+          apt-get update
+          DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+            uconsole-kernel-cm4-rpi \
+            clockworkpi-audio \
+            clockworkpi-cm-firmware
+          
+          apt-get clean
+          rm -rf /var/lib/apt/lists/*
+          KERNELEOF
+            chmod +x "$CUSTOM_STAGE/00-run-chroot.sh"
+          else
+            # Install custom-built kernel
+            cat > "$CUSTOM_STAGE/00-run.sh" << 'BUILDEOF'
+          #!/bin/bash -e
+          
+          KERNEL_DEBS="/home/runner/work/uConsole-Image-Builder/uConsole-Image-Builder/artifacts/kernel-debs"
+          if [ -d "$KERNEL_DEBS" ]; then
+            install -d "${ROOTFS_DIR}/tmp/kernel-debs"
+            install -m 644 "$KERNEL_DEBS"/*.deb "${ROOTFS_DIR}/tmp/kernel-debs/"
+          fi
+          BUILDEOF
+            chmod +x "$CUSTOM_STAGE/00-run.sh"
+            
+            cat > "$CUSTOM_STAGE/01-run-chroot.sh" << 'CHROOTEOF'
+          #!/bin/bash -e
+          
+          if [ -d /tmp/kernel-debs ]; then
+            echo "Installing custom kernel packages..."
+            apt-get update
+            apt-get install -y initramfs-tools
+            dpkg -i /tmp/kernel-debs/*.deb || apt-get install -f -y
+            rm -rf /tmp/kernel-debs
+            update-initramfs -u
+          fi
+          CHROOTEOF
+            chmod +x "$CUSTOM_STAGE/01-run-chroot.sh"
+          fi
+          
+          # Create uConsole configuration stage
+          UCONSOLE_STAGE="stage2/07-uconsole-config"
+          mkdir -p "$UCONSOLE_STAGE"
+          
+          cat > "$UCONSOLE_STAGE/00-run-chroot.sh" << 'CONFIGEOF'
+          #!/bin/bash -e
+          
+          # Configure uConsole-specific config.txt
+          cat > /boot/firmware/config.txt << 'BOOTEOF'
+          # uConsole Configuration
+          arm_64bit=1
+          disable_overscan=1
+          arm_boost=1
+          ignore_lcd=1
+          display_auto_detect=0
+          dtoverlay=audremap,pins_12_13
+          dtoverlay=dwc2,dr_mode=host
+          dtparam=ant2
+          dtparam=spi=on
+          dtparam=audio=on
+          camera_auto_detect=1
+          auto_initramfs=1
+          max_framebuffers=2
+          BOOTEOF
+          
+          # Add core-specific overlays
+          if [ "${{ matrix.uconsole_core }}" == "cm4" ]; then
+            echo 'dtoverlay=clockworkpi-uconsole' >> /boot/firmware/config.txt
+            echo 'dtoverlay=vc4-kms-v3d-pi4,cma-384' >> /boot/firmware/config.txt
+            echo 'dtparam=pciex1=off' >> /boot/firmware/config.txt
+            echo 'enable_uart=1' >> /boot/firmware/config.txt
+          else
+            echo 'dtoverlay=clockworkpi-uconsole-cm5' >> /boot/firmware/config.txt
+            echo 'dtoverlay=vc4-kms-v3d-pi5,cma-384' >> /boot/firmware/config.txt
+            echo 'dtparam=pciex1=off' >> /boot/firmware/config.txt
+            echo 'enable_uart=1' >> /boot/firmware/config.txt
+          fi
+          
+          # Enable SSH
+          touch /boot/firmware/ssh
+          CONFIGEOF
+          chmod +x "$UCONSOLE_STAGE/00-run-chroot.sh"
+          
+          # Create desktop installation stage if needed
+          if [ "${{ matrix.desktop }}" != "none" ]; then
+            DESKTOP_STAGE="stage2/08-desktop-environment"
+            mkdir -p "$DESKTOP_STAGE"
+            echo "${{ env.DESKTOP_TASK }}" > "$DESKTOP_STAGE/00-packages"
+          fi
+          
+          # Run pi-gen build
+          sudo ./build.sh
+          
+          # Copy output
+          mkdir -p ../output
+          cp -v deploy/* ../output/ || true
+
+      - name: Customize Ubuntu image
+        if: matrix.distro == 'jammy'
+        run: |
+          # Download Ubuntu image
+          IMAGE_PATH=https://cdimage.ubuntu.com/releases/jammy/release/ubuntu-22.04.5-preinstalled-server-arm64+raspi.img.xz
+          wget $IMAGE_PATH -O ubuntu-jammy-raspi.img.xz > /dev/null 2>&1
+          unxz ubuntu-jammy-raspi.img.xz
+          IMAGE_PATH=./ubuntu-jammy-raspi.img
             
             # Resize image to 8GB for GNOME desktop
             echo "Resizing Ubuntu image to 8GB..."
@@ -453,7 +579,7 @@ jobs:
           if [ "${{ env.KERNEL_MODE }}" == "prebuilt" ]; then
             sudo chroot $MOUNT_POINT /bin/bash -c "apt install -y --no-install-recommends uconsole-kernel-cm4-rpi"
           else
-            sudo cp  /home/runner/work/uConsole-Image-Builder/uConsole-Image-Builder/kernel-debs/*.deb $MOUNT_POINT/tmp/
+            sudo cp  /home/runner/work/uConsole-Image-Builder/uConsole-Image-Builder/artifacts/kernel-debs/*.deb $MOUNT_POINT/tmp/
             sudo chroot $MOUNT_POINT /bin/bash -c "dpkg -i /tmp/*.deb || apt install -f -y --no-install-recommends"
           fi
 
@@ -600,13 +726,34 @@ jobs:
           
           # Compress image with maximum compression settings
           xz --threads=1 --memory=1000MiB --check=crc32 -9e --compress --keep "$IMAGE_PATH"
-          mv "${IMAGE_PATH}.xz" /home/runner/work/uConsole-Image-Builder/uconsole-${{ matrix.distro }}-${{ matrix.uconsole_core }}-${{ matrix.desktop }}.img.xz
+          mkdir -p /home/runner/work/uConsole-Image-Builder/output
+          mv "${IMAGE_PATH}.xz" /home/runner/work/uConsole-Image-Builder/output/uconsole-${{ matrix.distro }}-${{ matrix.uconsole_core }}-${{ matrix.desktop }}.img.xz
+      
+      - name: Collect output images
+        run: |
+          mkdir -p /home/runner/work/uConsole-Image-Builder/final-output
+          
+          # For Debian (bookworm/trixie), images are in output/ from pi-gen
+          if [ "${{ matrix.distro }}" == "bookworm" ] || [ "${{ matrix.distro }}" == "trixie" ]; then
+            if [ -f "output/"*.img.xz ]; then
+              cp output/*.img.xz /home/runner/work/uConsole-Image-Builder/final-output/uconsole-${{ matrix.distro }}-${{ matrix.uconsole_core }}-${{ matrix.desktop }}.img.xz
+            fi
+          fi
+          
+          # For Ubuntu (jammy), images are in output/
+          if [ "${{ matrix.distro }}" == "jammy" ]; then
+            if [ -f "/home/runner/work/uConsole-Image-Builder/output/"*.img.xz ]; then
+              cp /home/runner/work/uConsole-Image-Builder/output/*.img.xz /home/runner/work/uConsole-Image-Builder/final-output/
+            fi
+          fi
+          
+          ls -lh /home/runner/work/uConsole-Image-Builder/final-output/
       
       - name: Upload image artifact
         uses: actions/upload-artifact@v4
         with:
           name: uconsole-${{ matrix.distro }}-${{ matrix.uconsole_core }}-${{ matrix.desktop }}.img.xz
-          path: /home/runner/work/uConsole-Image-Builder/uconsole-${{ matrix.distro }}-${{ matrix.uconsole_core }}-${{ matrix.desktop }}.img.xz
+          path: /home/runner/work/uConsole-Image-Builder/final-output/uconsole-${{ matrix.distro }}-${{ matrix.uconsole_core }}-${{ matrix.desktop }}.img.xz
 
   release:
     needs: [prepare-kernel, build-image]
diff --git a/.gitmodules b/.gitmodules
index 2f454fd..e71d3a3 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,6 +1,6 @@
-[submodule "rpi-image-gen"]
-	path = rpi-image-gen
-	url = https://github.com/raspberrypi/rpi-image-gen.git
+[submodule "pi-gen"]
+	path = pi-gen
+	url = https://github.com/ak-rex/ClockworkPi-pi-gen.git
 [submodule "linux"]
 	path = linux
 	url = https://github.com/raspberrypi/linux.git
diff --git a/IMPLEMENTATION.md b/IMPLEMENTATION.md
index ffb73ef..6ec0c28 100644
--- a/IMPLEMENTATION.md
+++ b/IMPLEMENTATION.md
@@ -2,7 +2,7 @@
 
 ## Overview
 
-This implementation provides a unified build system for creating bootable uConsole images with the ClockworkPi kernel, using the rpi-image-gen tool and automated CI/CD workflows.
+This implementation provides a unified build system for creating bootable uConsole images with the ClockworkPi kernel, using the pi-gen tool (ClockworkPi fork) and automated CI/CD workflows.
 
 ## Key Features
 
@@ -19,11 +19,12 @@ This implementation provides a unified build system for creating bootable uConso
 
 ### Image Generation
 
-Uses rpi-image-gen (Raspberry Pi's official image generator) to create bootable images:
+Uses pi-gen (ClockworkPi fork) to create bootable images:
 - Generates complete .img.xz files ready for SD card flashing
 - Integrates with ClockworkPi kernel installation
 - Supports multiple Debian/Ubuntu distributions
 - Creates compressed images with SHA256 checksums
+- Stage-based build system for customization
 
 ### Linux Kernel Source
 
@@ -58,12 +59,13 @@ Build scripts automatically detect and use the embedded kernel source when avail
 
 ### Core Scripts
 
-1. **scripts/generate_rpi_image.sh**
-   - Main wrapper for rpi-image-gen
+1. **scripts/generate_pi_image.sh**
+   - Main wrapper for pi-gen (ClockworkPi fork)
    - Creates bootable Raspberry Pi images for uConsole
    - Supports multiple distributions (Debian/Ubuntu)
    - Handles kernel installation (prebuilt or build from source)
    - Generates compressed .img.xz files
+   - Configures pi-gen stages for uConsole-specific customization
 
 2. **scripts/build_clockworkpi_kernel.sh**
    - Builds ClockworkPi kernel using Docker (reproducible builds)
@@ -109,19 +111,19 @@ Build scripts automatically detect and use the embedded kernel source when avail
 
 ## Build Process
 
-### Using rpi-image-gen (Recommended)
+### Using pi-gen (Recommended)
 
 For complete bootable images:
 
 ```bash
 # Build with prebuilt kernel (fast)
-sudo SUITE=jammy KERNEL_MODE=prebuilt ./scripts/generate_rpi_image.sh
+sudo SUITE=bookworm KERNEL_MODE=prebuilt ./scripts/generate_pi_image.sh
 
 # Build with custom kernel from source (slow)
-sudo SUITE=bookworm KERNEL_MODE=build ./scripts/generate_rpi_image.sh
+sudo SUITE=bookworm KERNEL_MODE=build ./scripts/generate_pi_image.sh
 
 # Build all supported distributions
-sudo SUITE=all KERNEL_MODE=prebuilt ./scripts/generate_rpi_image.sh
+sudo SUITE=all KERNEL_MODE=prebuilt ./scripts/generate_pi_image.sh
 ```
 
 ### Legacy: Using build-image.sh and setup-suite.sh
diff --git a/README.md b/README.md
index bad41b8..184b575 100644
--- a/README.md
+++ b/README.md
@@ -75,23 +75,28 @@ This creates 48 image variants (3 distros × 2 hardware variants × 8 desktop en
    - Adds Debian/Ubuntu archive keys
 
 2. **Submodule Initialization**:
-   - Checks out repository with `rpi-image-gen` submodule
+   - Checks out repository with `pi-gen` submodule (ClockworkPi-pi-gen fork)
    - Verifies submodule integrity
 
 3. **Base Image Creation**:
-   - For Debian (bookworm/trixie): Uses rpi-image-gen to build minimal base images
+   - For Debian (bookworm/trixie): Uses pi-gen (ClockworkPi's fork) to build complete bootable images from scratch
    - For Ubuntu (jammy): Downloads official Raspberry Pi server image
-   - Modifies rpi-image-gen configs for CM4 compatibility
+   - pi-gen uses a stage-based build system (stage0-stage4) to create customized images
 
 4. **Image Customization**:
-   - Mounts image partitions using loop devices
-   - Sets up chroot environment with proper bind mounts
-   - Removes existing kernel packages
-   - Adds ClockworkPi APT repository
-   - Installs kernel (prebuilt or custom-built .debs)
-   - Installs distribution-specific desktop environment:
-     - For Ubuntu (jammy): ubuntu-gnome-desktop, kubuntu-desktop, ubuntu-mate-desktop, xubuntu-desktop, lubuntu-desktop, etc.
-     - For Debian (bookworm/trixie): task-gnome-desktop, task-kde-desktop, task-mate-desktop, task-xfce-desktop, task-lxde-desktop, task-lxqt-desktop, task-cinnamon-desktop, task-gnome-flashback-desktop
+   - For Debian: pi-gen handles everything through its stage system
+     - Stage 0: Bootstrap base system
+     - Stage 1: System configuration
+     - Stage 2: Lite system with kernel and uConsole-specific configuration
+     - Additional stages for desktop environments
+   - For Ubuntu: Manual customization process
+     - Mounts image partitions using loop devices
+     - Sets up chroot environment with proper bind mounts
+     - Removes existing kernel packages
+     - Adds ClockworkPi APT repository
+     - Installs kernel (prebuilt or custom-built .debs)
+     - Installs distribution-specific desktop environment:
+       - For Ubuntu (jammy): ubuntu-gnome-desktop, kubuntu-desktop, ubuntu-mate-desktop, xubuntu-desktop, lubuntu-desktop, etc.
    - Configures system:
      - Creates `clockworkpi` user (password: `clockworkpi`)
      - Sets hostname to `uconsole`
diff --git a/pi-gen b/pi-gen
new file mode 160000
index 0000000..06d5ff4
--- /dev/null
+++ b/pi-gen
@@ -0,0 +1 @@
+Subproject commit 06d5ff470dbc7f926e6366b5d94cc3124796c54f
diff --git a/rpi-image-gen b/rpi-image-gen
deleted file mode 160000
index 09b6114..0000000
--- a/rpi-image-gen
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit 09b6114fcfc1dd1e67c4f278d1fee94b6421f466
diff --git a/scripts/generate_rpi_image.sh b/scripts/generate_rpi_image.sh
index c265ad6..bcc71f0 100755
--- a/scripts/generate_rpi_image.sh
+++ b/scripts/generate_rpi_image.sh
@@ -1,21 +1,21 @@
 #!/bin/bash
 #
-# generate_rpi_image.sh - Wrapper script for rpi-image-gen with ClockworkPi kernel integration
+# generate_rpi_image.sh - Wrapper script for pi-gen (ClockworkPi-pi-gen) with ClockworkPi kernel integration
 #
-# This script wraps rpi-image-gen to create Raspberry Pi images for uConsole devices,
+# This script wraps pi-gen to create Raspberry Pi images for uConsole devices,
 # with integrated ClockworkPi kernel installation.
 #
 # Usage: generate_rpi_image.sh
 #
 # Environment Variables:
 #   IMAGE_NAME       - Name for the generated image (default: uconsole-<suite>-<arch>)
-#   IMAGE_LINK       - Custom image link/URL to use as base
-#   SUITE            - Distribution suite (jammy|bookworm|bullseye|buster|focal|trixie|all)
+#   SUITE            - Distribution suite (bookworm|trixie|all)
 #   ARCH             - Architecture (default: arm64)
-#   ROOTFS_SIZE      - Root filesystem size in MB (default: 4096)
 #   OUTPUT_DIR       - Output directory for images (default: output/images)
 #   KERNEL_MODE      - Kernel installation mode: prebuilt|build|none (default: prebuilt)
-#   COMPRESS_FORMAT  - Compression format: xz|gzip|none (default: xz)
+#   COMPRESS_FORMAT  - Compression format: xz|gz|zip|none (default: xz)
+#   UCONSOLE_CORE    - Core model: cm4|cm5 (default: cm4)
+#   DESKTOP          - Desktop environment: gnome|kde|mate|xfce|lxde|lxqt|cinnamon|gnome-flashback|none (default: none for lite)
 #
 
 set -e
@@ -24,24 +24,21 @@ set -e
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
 
-# Source common mount functions
-source "$SCRIPT_DIR/common_mounts.sh"
-
 # Configuration from environment
-SUITE="${SUITE:-jammy}"
+SUITE="${SUITE:-trixie}"
 ARCH="${ARCH:-arm64}"
-ROOTFS_SIZE="${ROOTFS_SIZE:-4096}"
 OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/output/images}"
 IMAGE_NAME="${IMAGE_NAME:-}"
-IMAGE_LINK="${IMAGE_LINK:-}"
 KERNEL_MODE="${KERNEL_MODE:-prebuilt}"
 COMPRESS_FORMAT="${COMPRESS_FORMAT:-xz}"
+UCONSOLE_CORE="${UCONSOLE_CORE:-cm4}"
+DESKTOP="${DESKTOP:-none}"
 
 # Supported suites
-ALL_SUITES=("buster" "bullseye" "bookworm" "trixie" "focal" "jammy")
+ALL_SUITES=("bookworm" "trixie")
 
 echo "================================================"
-echo "uConsole Image Generator (rpi-image-gen wrapper)"
+echo "uConsole Image Generator (pi-gen wrapper)"
 echo "================================================"
 
 # Handle "all" suite option
@@ -74,245 +71,266 @@ fi
 
 # Set default image name if not provided
 if [ -z "$IMAGE_NAME" ]; then
-    IMAGE_NAME="uconsole-${SUITE}-${ARCH}"
+    if [ "$DESKTOP" = "none" ]; then
+        IMAGE_NAME="uconsole-${SUITE}-${UCONSOLE_CORE}"
+    else
+        IMAGE_NAME="uconsole-${SUITE}-${UCONSOLE_CORE}-${DESKTOP}"
+    fi
 fi
 
 echo "Suite: $SUITE"
 echo "Architecture: $ARCH"
 echo "Image Name: $IMAGE_NAME"
-echo "Root FS Size: ${ROOTFS_SIZE}MB"
 echo "Output Directory: $OUTPUT_DIR"
 echo "Kernel Mode: $KERNEL_MODE"
 echo "Compress Format: $COMPRESS_FORMAT"
-[ -n "$IMAGE_LINK" ] && echo "Base Image Link: $IMAGE_LINK"
+echo "uConsole Core: $UCONSOLE_CORE"
+echo "Desktop: $DESKTOP"
 echo "================================================"
 
 # Create output directory
 mkdir -p "$OUTPUT_DIR"
 
-# Setup cleanup trap
-trap_cleanup
-
-# Check if rpi-image-gen is available
-RPI_IMAGE_GEN="$REPO_ROOT/rpi-image-gen"
-if [ ! -d "$RPI_IMAGE_GEN" ]; then
-    echo "ERROR: rpi-image-gen not found at $RPI_IMAGE_GEN" >&2
+# Check if pi-gen is available
+PI_GEN_DIR="$REPO_ROOT/pi-gen"
+if [ ! -d "$PI_GEN_DIR" ]; then
+    echo "ERROR: pi-gen not found at $PI_GEN_DIR" >&2
     echo "Initialize the submodule: git submodule update --init --recursive" >&2
     exit 1
 fi
 
-# Determine base image and build strategy
-if [ -n "$IMAGE_LINK" ]; then
-    echo "Using custom base image: $IMAGE_LINK"
-    # Download the custom image
-    BASE_IMAGE_FILE="$OUTPUT_DIR/base-image.img"
-    echo "Downloading $IMAGE_LINK..."
-    wget -O "$BASE_IMAGE_FILE" "$IMAGE_LINK" || curl -L -o "$BASE_IMAGE_FILE" "$IMAGE_LINK"
-    
-    # Decompress if needed
-    if [[ "$BASE_IMAGE_FILE" == *.xz ]]; then
-        echo "Decompressing xz image..."
-        xz -d "$BASE_IMAGE_FILE"
-        BASE_IMAGE_FILE="${BASE_IMAGE_FILE%.xz}"
-    elif [[ "$BASE_IMAGE_FILE" == *.gz ]]; then
-        echo "Decompressing gzip image..."
-        gunzip "$BASE_IMAGE_FILE"
-        BASE_IMAGE_FILE="${BASE_IMAGE_FILE%.gz}"
-    fi
-    
-    IMAGE_FILE="$BASE_IMAGE_FILE"
-else
-    # Use rpi-image-gen to create base image
-    echo "Using rpi-image-gen to create base image for suite: $SUITE"
-    
-    # Map suite to rpi-image-gen layer
-    case "$SUITE" in
-        bookworm|bullseye)
-            BASE_LAYER="bookworm-minbase"
-            ;;
-        trixie)
-            BASE_LAYER="trixie-minbase"
-            ;;
-        jammy|focal|buster)
-            # For Ubuntu/older Debian, we'll use bookworm as base and customize
-            BASE_LAYER="bookworm-minbase"
-            echo "NOTE: Using bookworm base for $SUITE (will customize after)"
-            ;;
-        *)
-            echo "ERROR: Unsupported SUITE '$SUITE' for rpi-image-gen" >&2
-            exit 1
-            ;;
-    esac
-    
-    # Create temporary config for rpi-image-gen
-    # Use absolute path for config file
-    mkdir -p "$OUTPUT_DIR"
-    CONFIG_FILE="$(cd "$OUTPUT_DIR" && pwd)/rpi-image-gen-config.yaml"
-    
-    # Generate YAML config following rpi-image-gen best practices
-    # Reference: https://github.com/raspberrypi/rpi-image-gen/blob/main/layer/LAYER_BEST_PRACTICES
-    cat > "$CONFIG_FILE" << EOF
-info:
-  name: ${IMAGE_NAME}
-  description: "Image for ${SUITE} with ClockworkPi kernel integration"
-mmdebstrap:
-  mode: unshare
-  suite: ${BASE_LAYER}
-  target: ${IMAGE_NAME}.tar
-  variant: apt
-
-device:
-  layer: rpi-cm4
-
-image:
-  layer: image-rpios
-  boot_part_size: 512M
-  root_part_size: ${ROOTFS_SIZE}M
-  name: ${IMAGE_NAME}
-
-layer:
-  base: bookworm-minbase
+# Install dependencies if needed
+echo "Checking pi-gen dependencies..."
+if ! command -v debootstrap &> /dev/null; then
+    echo "Installing pi-gen dependencies..."
+    sudo apt-get update
+    sudo apt-get install -y coreutils quilt parted qemu-user-static debootstrap zerofree zip \
+        dosfstools e2fsprogs libarchive-tools libcap2-bin grep rsync xz-utils file git curl bc \
+        gpg pigz xxd arch-test bmap-tools kmod
+fi
+
+# Prepare pi-gen config
+cd "$PI_GEN_DIR"
+
+# Create config file
+CONFIG_FILE="$PI_GEN_DIR/config"
+cat > "$CONFIG_FILE" << EOF
+# uConsole Image Configuration
+IMG_NAME="${IMAGE_NAME}"
+RELEASE="${SUITE}"
+TARGET_HOSTNAME="uconsole"
+FIRST_USER_NAME="clockworkpi"
+FIRST_USER_PASS="clockworkpi"
+DISABLE_FIRST_BOOT_USER_RENAME=1
+ENABLE_SSH=1
+WPA_COUNTRY="US"
+LOCALE_DEFAULT="en_US.UTF-8"
+TIMEZONE_DEFAULT="UTC"
+KEYBOARD_KEYMAP="us"
+KEYBOARD_LAYOUT="English (US)"
 EOF
-    
-    echo "Generated rpi-image-gen config:"
-    cat "$CONFIG_FILE"
-    echo ""
-    
-    # Run rpi-image-gen to build base image
-    echo "Building base image with rpi-image-gen..."
-    cd "$RPI_IMAGE_GEN"
-    
-    # Check for podman (required for rootless mode)
-    if ! command -v podman &> /dev/null; then
-        echo "ERROR: podman is required for rootless image building" >&2
-        echo "Install with: sudo apt-get install -y podman" >&2
+
+# Set compression based on format
+case "$COMPRESS_FORMAT" in
+    xz)
+        echo 'DEPLOY_COMPRESSION="xz"' >> "$CONFIG_FILE"
+        echo 'COMPRESSION_LEVEL=6' >> "$CONFIG_FILE"
+        ;;
+    gz)
+        echo 'DEPLOY_COMPRESSION="gz"' >> "$CONFIG_FILE"
+        echo 'COMPRESSION_LEVEL=6' >> "$CONFIG_FILE"
+        ;;
+    zip)
+        echo 'DEPLOY_COMPRESSION="zip"' >> "$CONFIG_FILE"
+        echo 'COMPRESSION_LEVEL=6' >> "$CONFIG_FILE"
+        ;;
+    none)
+        echo 'DEPLOY_COMPRESSION="none"' >> "$CONFIG_FILE"
+        ;;
+    *)
+        echo "ERROR: Invalid COMPRESS_FORMAT '$COMPRESS_FORMAT'" >&2
+        echo "Valid formats: xz, gz, zip, none" >&2
         exit 1
-    fi
-    
-    # Install dependencies if needed
-    if [ ! -f "/usr/bin/mmdebstrap" ] || [ ! -f "/usr/bin/genimage" ]; then
-        echo "Installing rpi-image-gen dependencies..."
-        ./install_deps.sh || echo "WARNING: Failed to install dependencies, continuing anyway"
-    fi
-    
-    # Build the image
-    # Use absolute path for build directory
-    mkdir -p "$OUTPUT_DIR"
-    BUILD_DIR="$(cd "$OUTPUT_DIR" && pwd)/rpi-image-gen-build"
-    mkdir -p "$BUILD_DIR"
-    
-    # Run rpi-image-gen in rootless mode
-    # rpi-image-gen handles podman unshare internally when needed
-    echo "Running: ./rpi-image-gen build -c $CONFIG_FILE -B $BUILD_DIR"
-    ./rpi-image-gen build -c "$CONFIG_FILE" -B "$BUILD_DIR" 
+        ;;
+esac
+
+# Set stage list based on desktop
+if [ "$DESKTOP" = "none" ]; then
+    # Lite build (stage0, stage1, stage2)
+    echo 'STAGE_LIST="stage0 stage1 stage2"' >> "$CONFIG_FILE"
+else
+    # Full desktop build (stage0, stage1, stage2, stage3, stage4)
+    echo 'STAGE_LIST="stage0 stage1 stage2 stage3 stage4"' >> "$CONFIG_FILE"
 fi
 
-# Check if we need to customize the image (only if using manual fallback or custom kernel)
-if [ -n "${MANUAL_IMAGE:-}" ] || [ "$KERNEL_MODE" != "none" ]; then
-    echo ""
-    echo "================================================"
-    echo "Customizing image for uConsole..."
-    echo "================================================"
+echo ""
+echo "Generated pi-gen config:"
+cat "$CONFIG_FILE"
+echo ""
+
+# Create a custom stage for kernel installation if needed
+if [ "$KERNEL_MODE" != "none" ]; then
+    echo "Setting up custom kernel installation stage..."
+    CUSTOM_STAGE="$PI_GEN_DIR/stage2/06-uconsole-kernel"
+    mkdir -p "$CUSTOM_STAGE"
     
-    # These operations require root privileges
-    if [ "$EUID" -ne 0 ]; then
-        echo "ERROR: Image customization requires root privileges" >&2
-        echo "Please run with sudo for kernel installation and image customization" >&2
-        exit 1
-    fi
+    case "$KERNEL_MODE" in
+        prebuilt)
+            # Create script to install prebuilt kernel
+            cat > "$CUSTOM_STAGE/00-run-chroot.sh" << 'KERNELEOF'
+#!/bin/bash -e
 
-    # Setup loop device for the image
-    setup_loop_device "$IMAGE_FILE"
+# Install ClockworkPi kernel from repository
+echo "Installing ClockworkPi kernel from repository..."
 
-    # Mount partitions
-    TEMP_MOUNT="$OUTPUT_DIR/temp_mount"
-    mount_partitions "$LOOP_DEVICE" "$TEMP_MOUNT" 1 2
+# Add ClockworkPi repository
+cat > /etc/apt/sources.list.d/clockworkpi.list << EOF
+deb https://raw.githubusercontent.com/clockworkpi/apt/main/ stable main
+EOF
 
-    # Bind mount system directories for chroot
-    bind_mount_system "$TEMP_MOUNT"
+# Add repository key
+wget -qO - https://raw.githubusercontent.com/clockworkpi/apt/main/KEY.gpg | apt-key add - || true
 
-    # Setup QEMU for cross-architecture chroot
-    setup_qemu_chroot "$TEMP_MOUNT" "aarch64"
+# Update and install kernel
+apt-get update
+DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    uconsole-kernel-cm4-rpi \
+    clockworkpi-audio \
+    clockworkpi-cm-firmware || true
 
-    # Install kernel based on mode
-    case "$KERNEL_MODE" in
-        prebuilt)
-            echo "Installing prebuilt ClockworkPi kernel..."
-            "$SCRIPT_DIR/install_clockworkpi_kernel.sh" "$TEMP_MOUNT" "$SUITE"
+# Clean up
+apt-get clean
+rm -rf /var/lib/apt/lists/*
+
+echo "ClockworkPi kernel installed successfully"
+KERNELEOF
+            chmod +x "$CUSTOM_STAGE/00-run-chroot.sh"
             ;;
         build)
-            echo "Building ClockworkPi kernel from source..."
-            KERNEL_DEBS="$REPO_ROOT/artifacts/kernel-debs"
-            "$SCRIPT_DIR/build_clockworkpi_kernel.sh" "$KERNEL_DEBS"
-            
-            # Copy debs to mounted image and install
-            mkdir -p "$TEMP_MOUNT/tmp/kernel-debs"
-            cp "$KERNEL_DEBS"/*.deb "$TEMP_MOUNT/tmp/kernel-debs/"
+            # Create script to install kernel from debs
+            cat > "$CUSTOM_STAGE/00-run.sh" << 'BUILDEOF'
+#!/bin/bash -e
+
+# Check if kernel debs exist
+KERNEL_DEBS="${REPO_ROOT}/artifacts/kernel-debs"
+if [ ! -d "$KERNEL_DEBS" ] || [ -z "$(ls -A $KERNEL_DEBS/*.deb 2>/dev/null)" ]; then
+    echo "ERROR: Kernel debs not found in $KERNEL_DEBS"
+    echo "Build kernel first with: ./scripts/build_clockworkpi_kernel.sh"
+    exit 1
+fi
+
+# Copy kernel debs to stage
+install -d "${ROOTFS_DIR}/tmp/kernel-debs"
+install -m 644 "$KERNEL_DEBS"/*.deb "${ROOTFS_DIR}/tmp/kernel-debs/"
+BUILDEOF
+            chmod +x "$CUSTOM_STAGE/00-run.sh"
             
-            echo "Installing kernel packages in chroot..."
-            chroot "$TEMP_MOUNT" /bin/bash -c "
-                sudo apt-get update
-                sudo apt-get install -y initramfs-tools
-                dpkg -i /tmp/kernel-debs/*.deb || sudo apt-get install -f -y
-                rm -rf /tmp/kernel-debs
-            "
+            cat > "$CUSTOM_STAGE/01-run-chroot.sh" << 'CHROOTEOF'
+#!/bin/bash -e
+
+# Install kernel packages
+echo "Installing custom kernel packages..."
+apt-get update
+apt-get install -y initramfs-tools
+dpkg -i /tmp/kernel-debs/*.deb || apt-get install -f -y
+rm -rf /tmp/kernel-debs
+
+# Update initramfs
+update-initramfs -u
+
+echo "Custom kernel installed successfully"
+CHROOTEOF
+            chmod +x "$CUSTOM_STAGE/01-run-chroot.sh"
             ;;
-        none)
-            echo "Skipping kernel installation (KERNEL_MODE=none)"
+    esac
+fi
+
+# Create desktop installation stage if specified
+if [ "$DESKTOP" != "none" ]; then
+    echo "Setting up desktop environment: $DESKTOP"
+    DESKTOP_STAGE="$PI_GEN_DIR/stage3/06-desktop-environment"
+    mkdir -p "$DESKTOP_STAGE"
+    
+    # Map desktop to task package
+    case "$DESKTOP" in
+        gnome)
+            DESKTOP_PACKAGES="task-gnome-desktop"
+            ;;
+        kde)
+            DESKTOP_PACKAGES="task-kde-desktop"
+            ;;
+        mate)
+            DESKTOP_PACKAGES="task-mate-desktop"
+            ;;
+        xfce)
+            DESKTOP_PACKAGES="task-xfce-desktop"
+            ;;
+        lxde)
+            DESKTOP_PACKAGES="task-lxde-desktop"
+            ;;
+        lxqt)
+            DESKTOP_PACKAGES="task-lxqt-desktop"
+            ;;
+        cinnamon)
+            DESKTOP_PACKAGES="task-cinnamon-desktop"
+            ;;
+        gnome-flashback)
+            DESKTOP_PACKAGES="task-gnome-flashback-desktop"
             ;;
         *)
-            echo "ERROR: Invalid KERNEL_MODE '$KERNEL_MODE'" >&2
-            echo "Valid modes: prebuilt, build, none" >&2
+            echo "ERROR: Unknown desktop environment: $DESKTOP"
             exit 1
             ;;
     esac
-
-    # Apply any additional customizations if setup-suite.sh exists
-    if [ -x "$SCRIPT_DIR/setup-suite.sh" ]; then
-        echo "Applying suite-specific customizations..."
-        # Note: setup-suite.sh expects a rootfs directory, so we pass our mount point
-        # We need to make sure it doesn't try to mount again since we already mounted
-        SUITE="$SUITE" "$SCRIPT_DIR/setup-suite.sh" "$OUTPUT_DIR" || {
-            echo "WARNING: setup-suite.sh failed or not applicable"
-        }
-    fi
-
-    # Sync and cleanup
-    sync
-    cleanup_mounts
+    
+    echo "$DESKTOP_PACKAGES" > "$DESKTOP_STAGE/00-packages"
 fi
 
-# Compress image if requested
-FINAL_IMAGE="$OUTPUT_DIR/${IMAGE_NAME}.img"
-if [ "$IMAGE_FILE" != "$FINAL_IMAGE" ]; then
-    mv "$IMAGE_FILE" "$FINAL_IMAGE"
+# Clean any previous builds
+if [ -d "$PI_GEN_DIR/work" ]; then
+    echo "Cleaning previous build artifacts..."
+    sudo rm -rf "$PI_GEN_DIR/work"
 fi
 
-if [ "$COMPRESS_FORMAT" = "xz" ]; then
-    echo "Compressing image with xz..."
-    xz -9 -T 0 "$FINAL_IMAGE"
-    FINAL_IMAGE="${FINAL_IMAGE}.xz"
-elif [ "$COMPRESS_FORMAT" = "gzip" ]; then
-    echo "Compressing image with gzip..."
-    gzip -9 "$FINAL_IMAGE"
-    FINAL_IMAGE="${FINAL_IMAGE}.gz"
+if [ -d "$PI_GEN_DIR/deploy" ]; then
+    echo "Cleaning previous deploy directory..."
+    sudo rm -rf "$PI_GEN_DIR/deploy"
 fi
 
+# Run pi-gen build
 echo ""
 echo "================================================"
-echo "Image creation complete!"
+echo "Starting pi-gen build process..."
 echo "================================================"
-echo "Image file: $FINAL_IMAGE"
-ls -lh "$FINAL_IMAGE"
 echo ""
-echo "To write to SD card:"
-if [[ "$FINAL_IMAGE" == *.xz ]]; then
-    echo "  xz -dc $FINAL_IMAGE | sudo dd of=/dev/sdX bs=4M status=progress"
-elif [[ "$FINAL_IMAGE" == *.gz ]]; then
-    echo "  gunzip -c $FINAL_IMAGE | sudo dd of=/dev/sdX bs=4M status=progress"
-else
-    echo "  sudo dd if=$FINAL_IMAGE of=/dev/sdX bs=4M status=progress"
+
+# Run build script
+sudo ./build.sh
+
+# Check if build succeeded
+if [ ! -d "$PI_GEN_DIR/deploy" ]; then
+    echo "ERROR: Build failed - no deploy directory created" >&2
+    exit 1
 fi
-echo "  (Replace /dev/sdX with your SD card device)"
+
+# Copy output to our output directory
+echo ""
+echo "================================================"
+echo "Copying built images to output directory..."
+echo "================================================"
+
+mkdir -p "$OUTPUT_DIR"
+cp -v "$PI_GEN_DIR/deploy"/* "$OUTPUT_DIR/" 2>/dev/null || true
+
+# List generated files
+echo ""
+echo "Generated files:"
+ls -lh "$OUTPUT_DIR"
+
+echo ""
+echo "================================================"
+echo "Image creation complete!"
+echo "================================================"
+echo "Output directory: $OUTPUT_DIR"
 
 exit 0
diff --git a/scripts/pi_gen/README.md b/scripts/pi_gen/README.md
new file mode 100644
index 0000000..dc10758
--- /dev/null
+++ b/scripts/pi_gen/README.md
@@ -0,0 +1,109 @@
+# pi-gen Integration
+
+This directory contains documentation and configuration for the pi-gen integration.
+
+## About pi-gen
+
+pi-gen is the traditional Raspberry Pi image generation tool. We use a fork specifically customized for ClockworkPi uConsole devices.
+
+- **Repository**: https://github.com/ak-rex/ClockworkPi-pi-gen
+- **Location**: `/pi-gen/` (git submodule)
+- **Type**: Fork of RPI-Distro/pi-gen with ClockworkPi customizations
+
+## How pi-gen Works
+
+pi-gen uses a stage-based build system:
+
+- **stage0**: Bootstrap the base Debian/Raspbian system
+- **stage1**: Configure basic system settings
+- **stage2**: Lite system with networking and basic tools
+- **stage3**: Desktop system (optional)
+- **stage4**: Full system with additional software (optional)
+
+Each stage contains numbered subdirectories (00-, 01-, etc.) with:
+- `00-run.sh`: Scripts run on the host
+- `00-run-chroot.sh`: Scripts run inside the chroot
+- `00-packages`: Packages to install
+- `00-debconf`: Debconf settings
+
+## uConsole Customizations
+
+For uConsole images, we add custom stages:
+
+1. **stage2/06-uconsole-kernel**: Installs ClockworkPi kernel
+2. **stage2/07-uconsole-config**: Configures boot settings for uConsole hardware
+3. **stage2/08-desktop-environment**: Installs desktop environment (if requested)
+
+## Updating pi-gen
+
+To update the pi-gen submodule to a newer version:
+
+```bash
+cd pi-gen
+git fetch origin
+git checkout main  # or specific tag/branch
+cd ..
+git add pi-gen
+git commit -m "Update pi-gen to latest version"
+```
+
+## Using pi-gen
+
+The wrapper script `scripts/generate_pi_image.sh` provides an easy interface to pi-gen.
+
+### Basic Usage
+
+```bash
+# Build Debian Bookworm with prebuilt kernel
+sudo SUITE=bookworm ./scripts/generate_pi_image.sh
+
+# Build Debian Trixie with custom kernel
+sudo SUITE=trixie KERNEL_MODE=build ./scripts/generate_pi_image.sh
+
+# Build with specific desktop environment
+sudo SUITE=bookworm DESKTOP=gnome ./scripts/generate_pi_image.sh
+```
+
+### Environment Variables
+
+- `SUITE`: Distribution suite (bookworm, trixie)
+- `KERNEL_MODE`: prebuilt, build, or none
+- `DESKTOP`: gnome, kde, mate, xfce, lxde, lxqt, cinnamon, gnome-flashback, none
+- `UCONSOLE_CORE`: cm4 or cm5
+- `COMPRESS_FORMAT`: xz, gz, zip, none
+
+## pi-gen Documentation
+
+For detailed pi-gen documentation, see:
+
+- `/pi-gen/README.md` - Main documentation
+- https://github.com/ak-rex/ClockworkPi-pi-gen - GitHub repository
+- https://github.com/RPI-Distro/pi-gen - Original pi-gen project
+
+## Differences from rpi-image-gen
+
+The previous `rpi-image-gen` tool used a YAML-based configuration approach. The new `pi-gen` tool uses:
+
+- Stage-based builds instead of config files
+- More extensive customization options
+- Better integration with Debian/Raspbian ecosystem
+- ClockworkPi-specific enhancements already integrated
+
+## Troubleshooting
+
+### Build fails with "debootstrap not found"
+
+Install dependencies:
+```bash
+sudo apt-get install -y coreutils quilt parted qemu-user-static debootstrap \
+  zerofree zip dosfstools e2fsprogs libarchive-tools libcap2-bin grep rsync \
+  xz-utils file git curl bc gpg pigz xxd arch-test bmap-tools kmod
+```
+
+### Build runs out of space
+
+pi-gen requires significant disk space. Ensure you have at least 20GB free.
+
+### Permission errors
+
+pi-gen needs root permissions for chroot operations. Run with `sudo`.
diff --git a/scripts/rpi_image_gen/README.md b/scripts/rpi_image_gen/README.md
deleted file mode 100644
index b820368..0000000
--- a/scripts/rpi_image_gen/README.md
+++ /dev/null
@@ -1,54 +0,0 @@
-# rpi-image-gen Integration
-
-This directory contains documentation and configuration for the rpi-image-gen integration.
-
-## About rpi-image-gen
-
-rpi-image-gen is the official Raspberry Pi image generation tool from the Raspberry Pi Foundation. It creates highly customizable, reproducible Raspberry Pi OS images.
-
-- **Repository**: https://github.com/raspberrypi/rpi-image-gen
-- **Included Version**: v2.0.0-rc.1-43-g09b6114 (commit: 09b6114)
-- **Location**: `/rpi-image-gen/` (git submodule)
-
-## Updating rpi-image-gen
-
-To update the rpi-image-gen submodule to a newer version:
-
-```bash
-# Navigate to the submodule directory
-cd rpi-image-gen
-
-# Fetch latest changes
-git fetch origin
-
-# Checkout a specific tag or commit
-git checkout v2.0.0  # or any other tag/commit
-
-# Go back to main repository
-cd ..
-
-# Commit the submodule update
-git add rpi-image-gen
-git commit -m "Update rpi-image-gen to v2.0.0"
-```
-
-## Using rpi-image-gen
-
-The wrapper script `scripts/generate_rpi_image.sh` provides an easy interface to rpi-image-gen.
-
-See the main README.md for usage examples and environment variable documentation.
-
-## rpi-image-gen Documentation
-
-For detailed rpi-image-gen documentation, see:
-- `/rpi-image-gen/README.adoc` - Main documentation
-- https://github.com/raspberrypi/rpi-image-gen - GitHub repository
-
-## Supported Features
-
-The wrapper script supports:
-- Multiple Debian/Ubuntu suites (buster, bullseye, bookworm, trixie, focal, jammy)
-- Custom image names and links
-- Configurable root filesystem size
-- Custom output directories
-- Integration with ClockworkPi kernel packages
diff --git a/scripts/test_scripts.sh b/scripts/test_scripts.sh
index fb0826a..4a7a6ea 100755
--- a/scripts/test_scripts.sh
+++ b/scripts/test_scripts.sh
@@ -113,29 +113,29 @@ echo "============================="
 cd "$REPO_ROOT"
 
 # Check required directories exist
-if [ -d "rpi-image-gen" ]; then
-    test_pass "rpi-image-gen submodule directory exists"
+if [ -d "pi-gen" ]; then
+    test_pass "pi-gen submodule directory exists"
 else
-    test_fail "rpi-image-gen submodule directory missing"
+    test_fail "pi-gen submodule directory missing"
 fi
 
 if [ -d "artifacts/kernel-debs" ]; then
     test_pass "artifacts/kernel-debs directory exists"
 else
-    test_fail "artifacts/kernel-debs directory missing"
+    test_skip "artifacts/kernel-debs directory missing (expected if no kernel built)"
 fi
 
 if [ -f "patches/ak-rex.patch" ]; then
     test_pass "patches/ak-rex.patch exists"
 else
-    test_fail "patches/ak-rex.patch missing"
+    test_skip "patches/ak-rex.patch missing (patch now downloaded from GitHub)"
 fi
 
 # Check documentation exists
-if [ -f "scripts/rpi_image_gen/README.md" ]; then
-    test_pass "rpi-image-gen documentation exists"
+if [ -f "scripts/pi_gen/README.md" ]; then
+    test_pass "pi-gen documentation exists"
 else
-    test_fail "rpi-image-gen documentation missing"
+    test_fail "pi-gen documentation missing"
 fi
 
 if [ -f "CHANGES.md" ]; then
@@ -151,10 +151,10 @@ echo "============================="
 cd "$REPO_ROOT"
 
 # Check submodule is configured
-if git submodule status | grep -q "rpi-image-gen"; then
-    test_pass "rpi-image-gen submodule is configured"
+if git submodule status | grep -q "pi-gen"; then
+    test_pass "pi-gen submodule is configured"
 else
-    test_fail "rpi-image-gen submodule not configured"
+    test_fail "pi-gen submodule not configured"
 fi
 
 # Check .gitmodules exists