diff --git a/CloneToDisk.sh b/CloneToDisk.sh
index 815b990..5614a38 100755
--- a/CloneToDisk.sh
+++ b/CloneToDisk.sh
@@ -16,12 +16,17 @@ IMG_FILE="PBX-root.img.lz4"
 MOUNT_PATH="/mnt/debian"
 MOUNT_POINTS=("/dev" "/dev/pts" "/proc" "/sys" "/sys/firmware/efi/efivars" "/run")
 
+BOOT_PART="null"
+EFI_PART="null"
+ROOT_PART="null"
+
 VGROUP="PBX-vg"
 ROOT_LV="root"
 SWAP_LV="swap_1"
 VG_PATH="/dev/mapper/PBX--vg-"
 
-GRUB_CMD="grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian && update-grub"
+GRUB_CMD="export PATH=/usr/sbin:$PATH && update-initramfs -u -k all
+ && grub-install --target=x86_64-efi --efi-directory=/boot/efi --bootloader-id=debian && update-grub"
 
 for arg in "$@"; do
   case $arg in
@@ -33,10 +38,6 @@ for arg in "$@"; do
     NO_RETABLE=true
     shift
     ;;
-  --nousb)
-    NO_USB=true
-    shift
-    ;;
   --nocleanup)
     NO_CLEANUP=true
     shift
@@ -64,8 +65,12 @@ main() {
 
   trap 'echo "Status Code: $?"; cleanup;' EXIT
 
+  EFI_PART=$(get_partition $disk_name "1")
+  BOOT_PART=$(get_partition $disk_name "2")
+  ROOT_PART=$(get_partition $disk_name "3")
+
   if [[ $NO_RETABLE == true ]]; then
-    echo -e "${YELLOW}--noretable is specified. Tables will not be erased and no partitions will be created${RESET}"
+    echo -e "${YELLOW}Option --noretable is specified. Tables will not be erased and no partitions will be created${RESET}"
   else
     rewrite_gpt_table
     lvm_init
@@ -74,6 +79,8 @@ main() {
 
   write_img
 
+  mount_chroot
+
   gen_fstab
 
   install_grub
@@ -115,9 +122,9 @@ rewrite_gpt_table() {
 }
 
 lvm_init() {
-  pvcreate $(get_partition $disk_name "3")
+  pvcreate $ROOT_PART -ff
 
-  vgcreate $VGROUP $(get_partition $disk_name "3")
+  vgcreate $VGROUP $ROOT_PART -f
 
   vgchange -ay $VGROUP
 
@@ -128,15 +135,15 @@ lvm_init() {
 }
 
 make_filesystem() {
-  mkfs.vfat -F 32 $(get_partition $disk_name "1")
-  mkfs.ext4 $(get_partition $disk_name "2")
+  mkfs.vfat -F 32 $EFI_PART
+  mkfs.ext4 $BOOT_PART
   mkswap "${VG_PATH}${SWAP_LV}"
   echo -e "${GREEN}Make filesystems finished.${RESET}"
 }
 
 write_img() {
   if [[ $NO_WRITE == true ]]; then
-    echo -e "${YELLOW}--nowrite is specified. No need for Img no disks will be written to${RESET}"
+    echo -e "${YELLOW}Option --nowrite is specified. No need for Img no disks will be written to${RESET}"
     return 0
   fi
 
@@ -151,35 +158,15 @@ write_img() {
     exit 100
   fi
 
-  echo "Writing image to $(get_partition $disk_name "3")..."
+  echo "Writing image to ${VG_PATH}${ROOT_LV}..."
 
-  lz4 -dc ${img_file_dir} | partclone.ext4 -r -o $(get_partition $disk_name "3")
+  lz4 -dc ${img_file_dir} | partclone.ext4 -r -o "${VG_PATH}${ROOT_LV}"
 
   partprobe || true
 
   echo -e "${GREEN}Finished writing image to disk${RESET}"
 }
 
-cleanup() {
-  if [[ $NO_CLEANUP = true ]]; then
-    echo -e "${YELLOW}No cleanup specified. Make sure to manually clean mounts.${RESET}" 1>&2
-    exit 100
-  fi
-
-  echo "Cleaning Mounts"
-
-  umount -lf /mnt/ventoy 2>/dev/null || true
-
-  for i in "${MOUNT_POINTS[@]}"; do umount -lf "${MOUNT_PATH}${i}" 2>/dev/null || true; done
-
-  umount -lf /mnt/debian/boot/efi 2>/dev/null || true
-  umount -lf /mnt/debian/boot 2>/dev/null || true
-  umount -lf /mnt/debian 2>/dev/null || true
-
-  vgchange -an $VGROUP
-  echo -e "${GREEN}Clean.${RESET}"
-}
-
 get_partition() {
   local disk_name="$1"
   local part_num="$2"
@@ -194,11 +181,10 @@ get_partition() {
 }
 
 get_partuuid() {
-  local disk_name="$1"
-  local part_num="$2"
+  local target_part="$1"
 
   local part_uuid="$(
-    lsblk -ln -o PATH,PARTUUID $disk_name | awk '$1==$(get_partition $disk_name $part_num) {print $2}'
+    lsblk -ln -o PATH,PARTUUID $disk_name | awk '$1=="${target_part}" {print $2}'
   )"
 
   printf "$part_uuid"
@@ -209,33 +195,54 @@ mount_chroot() {
 
   echo "Mounting restored system..."
 
-  mkdir -p $MOUNT_PATH/boot
-  mount /dev/mapper/PBX--vg-root $MOUNT_PATH
-  mount $(get_partition $disk_name "2") ${MOUNT_PATH}/boot
+  mkdir -p "${MOUNT_PATH}/boot"
+  mount "${VG_PATH}${ROOT_LV}" $MOUNT_PATH
+  mount $BOOT_PART "${MOUNT_PATH}/boot"
 
-  mkdir -p $MOUNT_PATH/boot/efi
-  mount $(get_partition $disk_name "1") ${MOUNT_PATH}/boot/efi
+  mkdir -p "${MOUNT_PATH}/boot/efi"
+  mount $EFI_PART "${MOUNT_PATH}/boot/efi"
 
   for i in "${MOUNT_POINTS[@]}"; do mount -B $i ${MOUNT_PATH}${i}; done
+
   echo -e "${GREEN}Mounts prepared for chroot.${RESET}"
 }
 
 gen_fstab() {
   echo "Generating fstab..."
-  boot_uuid=$()
+
+  local boot_uuid=$(get_partuuid $BOOT_PART)
+  local efi_uuid=$(get_partuuid $BOOT_PART)
 
   cp ./fstab_template ${MOUNT_PATH}/etc/fstab
 
-  sed --in-place "s/boot_uuid/$(get_partuuid $disk_name "2")/g" ${MOUNT_PATH}/etc/fstab
-  sed --in-place "s/efi_uuid/$(get_partuuid $disk_name "1")/g" ${MOUNT_PATH}/etc/fstab
+  sed --in-place "s/BOOT_UUID/${boot_uuid}/g" ${MOUNT_PATH}/etc/fstab
+  sed --in-place "s/EFI_UUID/${efi_uuid}/g" ${MOUNT_PATH}/etc/fstab
 
   echo -e "${GREEN}fstab write complete!${RESET}"
 }
 
 install_grub() {
   echo "Entering chroot to reinstall GRUB..."
-  chroot $MOUNT_PATH /bin/bash -c $GRUB_CMD
+  chroot $MOUNT_PATH /bin/bash -c "$GRUB_CMD"
   echo -e "${GREEN}Grub installed!${RESET}"
 }
 
+cleanup() {
+  if [[ $NO_CLEANUP = true ]]; then
+    echo -e "${YELLOW}No cleanup specified. Make sure to manually clean mounts.${RESET}" 1>&2
+    exit 100
+  fi
+
+  echo "Cleaning Mounts"
+
+  for i in "${MOUNT_POINTS[@]}"; do umount -lf "${MOUNT_PATH}${i}" 2>/dev/null || true; done
+
+  umount -lf "${MOUNT_PATH}/boot/efi" 2>/dev/null || true
+  umount -lf "${MOUNT_PATH}/boot" 2>/dev/null || true
+  umount -lf "${MOUNT_PATH}" 2>/dev/null || true
+
+  vgchange -an $VGROUP
+  echo -e "${GREEN}Clean.${RESET}"
+}
+
 main