diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index 0e1146914c..2ea8f5b10d 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -1,2 +1,2 @@
-FROM mcr.microsoft.com/vscode/devcontainers/python:3.8
+FROM mcr.microsoft.com/vscode/devcontainers/python:3.9
 COPY setup.sh /setup.sh
diff --git a/.devcontainer/setup.sh b/.devcontainer/setup.sh
index 419a861781..cc8d8ef402 100644
--- a/.devcontainer/setup.sh
+++ b/.devcontainer/setup.sh
@@ -1,4 +1,5 @@
 sudo pip install --upgrade pip
+sudo pip install -r requirements.txt --progress-bar off
 sudo pip install -e ".[tests]"
 sudo apt update
 sudo apt install -y clang-format
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index 77da42b490..be5c0cb902 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -29,8 +29,8 @@ Fixes # (issue)
 Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
 members/contributors who may be interested in your PR.
 
-<!-- 
-Feel free to tag @ianstenbit and @jbischof in your reviews.
+<!--
+Feel free to tag @divyashreepathihalli and @sampathweb in your reviews.
 -->
 
 <!--
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000000..0df37b1230
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,23 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
+
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "monthly"
+    groups:
+      github-actions:
+        patterns:
+          - "*"
+  - package-ecosystem: "pip"
+    directory: "/"
+    schedule:
+      interval: "monthly"
+    groups:
+      python:
+        patterns:
+          - "*"
diff --git a/.github/workflows/actions.yml b/.github/workflows/actions.yml
index 9c218ce5cc..8be69b967b 100644
--- a/.github/workflows/actions.yml
+++ b/.github/workflows/actions.yml
@@ -3,25 +3,30 @@ name: Tests
 on:
   push:
   pull_request:
+  workflow_call:
   release:
     types: [created]
+
+permissions:
+  contents: read
+
 jobs:
-  test:
-    name: Test the code with tf.keras
+  keras_2:
+    name: Test the code with Keras 2
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v1
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.9
+      uses: actions/setup-python@v5
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: Get pip cache dir
       id: pip-cache
       run: |
         python -m pip install --upgrade pip setuptools
         echo "::set-output name=dir::$(pip cache dir)"
     - name: pip cache
-      uses: actions/cache@v2
+      uses: actions/cache@v4
       with:
         path: ${{ steps.pip-cache.outputs.dir }}
         key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
@@ -29,36 +34,38 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
-        pip install tensorflow>=2.13.0
+        pip install tensorflow~=2.14
         pip install torch>=2.0.1+cpu
         pip install "jax[cpu]"
         pip install keras-core
         pip install -e ".[tests]" --progress-bar off --upgrade
     - name: Test with pytest
       env:
+        KAGGLE_USERNAME:  ${{ secrets.KAGGLE_USERNAME }}
+        KAGGLE_KEY: ${{ secrets.KAGGLE_KEY }}
         TEST_CUSTOM_OPS: false
       run: |
         pytest keras_cv/ --ignore keras_cv/models/legacy/ --durations 0
-  multibackend:
-    name: Test the code with Keras Core
+  keras_3:
+    name: Test the code with Keras 3
     strategy:
       fail-fast: false
       matrix:
         backend: [tensorflow, jax, torch]
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v1
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.9
+      uses: actions/setup-python@v5
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: Get pip cache dir
       id: pip-cache
       run: |
         python -m pip install --upgrade pip setuptools
         echo "::set-output name=dir::$(pip cache dir)"
     - name: pip cache
-      uses: actions/cache@v2
+      uses: actions/cache@v4
       with:
         path: ${{ steps.pip-cache.outputs.dir }}
         key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
@@ -66,16 +73,14 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
-        pip install tensorflow>=2.13.0
-        pip install "jax[cpu]"
-        pip install torch>=2.0.1+cpu
-        pip install torchvision>=0.15.1
-        pip install keras-core
+        pip install -r requirements.txt
         pip install -e ".[tests]" --progress-bar off --upgrade
     - name: Test with pytest
       env:
         TEST_CUSTOM_OPS: false # TODO(ianstenbit): test custom ops, or figure out what our story is here
         KERAS_BACKEND: ${{ matrix.backend }}
+        KAGGLE_USERNAME:  ${{ secrets.KAGGLE_USERNAME}}
+        KAGGLE_KEY: ${{ secrets.KAGGLE_KEY}}
         JAX_ENABLE_X64: true
       run: |
         pytest keras_cv/bounding_box \
@@ -94,18 +99,18 @@ jobs:
     name: Check the code format
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v1
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.9
+      uses: actions/setup-python@v5
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: Get pip cache dir
       id: pip-cache
       run: |
         python -m pip install --upgrade pip setuptools
         echo "::set-output name=dir::$(pip cache dir)"
     - name: pip cache
-      uses: actions/cache@v2
+      uses: actions/cache@v4
       with:
         path: ${{ steps.pip-cache.outputs.dir }}
         key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
@@ -113,12 +118,12 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
-        pip install tensorflow>=2.13.0
+        pip install tensorflow~=2.14
         pip install -e ".[tests]" --progress-bar off --upgrade
     - name: Lint
       run: bash shell/lint.sh
     - name: Clang
-      uses: DoozyX/clang-format-lint-action@v0.14
+      uses: DoozyX/clang-format-lint-action@v0.17
       with:
         source: '.'
         extensions: 'h,c,cpp,hpp,cc'
diff --git a/.github/workflows/auto-assignment.yml b/.github/workflows/auto-assignment.yml
new file mode 100644
index 0000000000..de72da8ba2
--- /dev/null
+++ b/.github/workflows/auto-assignment.yml
@@ -0,0 +1,21 @@
+name: auto-assignment
+on:
+  issues:
+    types:
+      - opened
+
+permissions:
+  contents: read
+  issues: write
+  pull-requests: write
+
+jobs:
+  welcome:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/github-script@v7
+        with:
+          script: |
+            const script = require('./\.github/workflows/scripts/auto-assignment.js')
+            script({github, context})
diff --git a/.github/workflows/devcontainer.yml b/.github/workflows/devcontainer.yml
index 72f12471c7..ee0f49d549 100644
--- a/.github/workflows/devcontainer.yml
+++ b/.github/workflows/devcontainer.yml
@@ -11,6 +11,8 @@ on: # rebuild any PRs and main branch changes
       - '.devcontainer/**'
       - 'workflows/**'
 
+permissions:
+  contents: read
 
 jobs:
   build:
@@ -18,9 +20,9 @@ jobs:
     steps:
 
       - name: Checkout (GitHub)
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
 
       - name: Build and run dev container task
-        uses: devcontainers/ci@v0.2
+        uses: devcontainers/ci@v0.3
         with:
           runCmd: pytest keras_cv/ --ignore keras_cv/models
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
new file mode 100644
index 0000000000..ded0a461b2
--- /dev/null
+++ b/.github/workflows/nightly.yml
@@ -0,0 +1,50 @@
+name: Nightly
+
+on:
+  workflow_dispatch: # To Generate wheels on demand outside of schedule.
+  schedule:
+    - cron: '0 3 * * *' # run at 3 AM UTC / 8 PM PDT
+
+permissions:
+  contents: read
+
+jobs:
+  run-test-for-nightly:
+    uses: ./.github/workflows/actions.yml
+  nightly:
+    name: Build Wheel file and upload
+    needs: [run-test-for-nightly]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.9
+      - name: Get pip cache dir
+        id: pip-cache
+        run: |
+          python -m pip install --upgrade pip setuptools
+          echo "::set-output name=dir::$(pip cache dir)"
+      - name: pip cache
+        uses: actions/cache@v4
+        with:
+          path: ${{ steps.pip-cache.outputs.dir }}
+          key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip setuptools
+          pip install twine
+          pip install -r requirements.txt --progress-bar off
+      - name: Build wheel file
+        run: |
+          export BUILD_WITH_CUSTOM_OPS=false
+          python pip_build.py --nightly
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_NIGHTLY_API_TOKEN }}
+          packages-dir: dist/
+          verbose: true
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 2ee85c2055..7a471e938a 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -3,110 +3,43 @@ name: Release
 on:
   release:
     types: [published]
+
+permissions:
+  contents: read
+
 jobs:
-  deploy-with-custom-ops:
-    # This job is currently skipped until we cut a release with custom ops.
-    if: false
-    name: Build and deploy release wheels with custom ops
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        # To switch on windows-2022/latest, please verify the bazel version:
-        # https://github.com/bazelbuild/bazel/issues/14232#issuecomment-1011247429
-        os: ['macos-12', 'windows-2019', 'ubuntu-18.04']
-        py-version: ['3.8', '3.9', '3.10', '3.11']
-        tf-version: ['2.13.0']
-        use-macos-arm: [false]
-        include:
-          - os: 'macos-12'
-            tf-version: '2.13.0'
-            py-version: '3.8'
-            use-macos-arm: true
-          - os: 'macos-12'
-            tf-version: '2.13.0'
-            py-version: '3.9'
-            use-macos-arm: true
-          - os: 'macos-12'
-            tf-version: '2.13.0'
-            py-version: '3.10'
-            use-macos-arm: true
-      fail-fast: false
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v1
-      with:
-        python-version: ${{ matrix.py-version }}
-    - name: Get pip cache dir
-      id: pip-cache
-      run: |
-        python -m pip install --upgrade pip setuptools wheel auditwheel twine
-        echo "::set-output name=dir::$(pip cache dir)"
-    - name: pip cache
-      uses: actions/cache@v2
-      with:
-        path: ${{ steps.pip-cache.outputs.dir }}
-        key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
-        restore-keys: |
-          ${{ runner.os }}-pip-
-    - name: Install dependencies
-      run: |
-        python -m pip install tensorflow-cpu==${{ matrix.tf-version }}
-        python -m pip install -e ".[tests]" --progress-bar off --upgrade
-    - name: Configure Build Environment
-      run: |
-        python build_deps/configure.py
-    - name: Reinstall TensorFlow (MacOS ARM)
-      if: ${{ matrix.os == 'macos-12' && matrix.use-macos-arm}}
-      run: |
-        python -m pip uninstall -y tensorflow-cpu
-        python -m pip install --platform=macosx_12_0_arm64 --no-deps --target=$(python -c 'import site; print(site.getsitepackages()[0])') --upgrade tensorflow-macos==${{ matrix.tf-version }}
-    - name: Bazel Build
-      if: ${{ ! matrix.use-macos-arm }}
-      run: |
-        export BUILD_WITH_CUSTOM_OPS=true
-        bazel build build_pip_pkg
-    - name: Bazel Build (MacOS ARM)
-      if: ${{ matrix.use-macos-arm}}
-      run: |
-        bazel build --cpu=darwin_arm64 --copt -mmacosx-version-min=12.0 --linkopt -mmacosx-version-min=12.0 build_pip_pkg
-    - name: Build wheels
-      run: |
-        export BUILD_WITH_CUSTOM_OPS=true
-        bazel-bin/build_pip_pkg wheels
-    - name: Repair wheels (manylinux)
-      if: ${{ matrix.os == 'ubuntu-18.04' }}
-      run: |
-        python -m pip install --upgrade patchelf==0.14
-        bash build_deps/tf_auditwheel_patch.sh
-        python -m auditwheel repair --plat manylinux2014_x86_64 wheels/*.whl
-        rm wheels/*.whl
-        mv wheelhouse/* wheels/
-    - name: Upload wheels
-      env:
-        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
-        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
-      run: |
-        twine upload wheels/*
-  deploy-without-custom-ops:
-    name: Build and deploy release wheels without custom ops
+  run-test-for-release:
+    uses: ./.github/workflows/actions.yml
+  release:
+    name: Build and publish to PyPI
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v1
-      with:
-        python-version: 3.8
-    - name: Build wheels
-      run: |
-        pip install tensorflow==2.13.0
-        python -m pip install --upgrade setuptools wheel twine
-        python -m pip install --upgrade -r requirements.txt
-        export BUILD_WITH_CUSTOM_OPS=false
-        python pip_build.py
-    - name: Upload wheels
-      env:
-        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
-        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
-      run: |
-        twine upload dist/*.whl
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.9
+      - name: Get pip cache dir
+        id: pip-cache
+        run: |
+          python -m pip install --upgrade pip setuptools
+          echo "::set-output name=dir::$(pip cache dir)"
+      - name: pip cache
+        uses: actions/cache@v4
+        with:
+          path: ${{ steps.pip-cache.outputs.dir }}
+          key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-
+      - name: Install dependencies
+        run: |
+            pip install -r requirements.txt --progress-bar off
+      - name: Build a binary wheel and a source tarball
+        run: |
+          export BUILD_WITH_CUSTOM_OPS=false
+          python pip_build.py
+      - name: Publish distribution to PyPI
+        if: startsWith(github.ref, 'refs/tags')
+        uses: pypa/gh-action-pypi-publish@master
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.github/workflows/scorecard.yml b/.github/workflows/scorecard.yml
new file mode 100644
index 0000000000..98509aef93
--- /dev/null
+++ b/.github/workflows/scorecard.yml
@@ -0,0 +1,58 @@
+name: Scorecard supply-chain security
+on:
+  # For Branch-Protection check. Only the default branch is supported. See
+  # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection
+  branch_protection_rule:
+  # To guarantee Maintained check is occasionally updated. See
+  # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained
+  schedule:
+    - cron: "37 4 * * 6"
+  push:
+    branches: [ "master" ]
+
+permissions: read-all
+
+jobs:
+  analysis:
+    name: Scorecard analysis
+    runs-on: ubuntu-latest
+    permissions:
+      # Needed to upload the results to code-scanning dashboard.
+      security-events: write
+      # Needed to publish results and get a badge (see publish_results below).
+      id-token: write
+
+    steps:
+      - name: "Checkout code"
+        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+        with:
+          persist-credentials: false
+
+      - name: "Run analysis"
+        uses: ossf/scorecard-action@0864cf19026789058feabb7e87baa5f140aac736 # v2.3.1
+        with:
+          results_file: results.sarif
+          results_format: sarif
+          # (Optional) "write" PAT token. Uncomment the `repo_token` line below if
+          # you want to enable the Branch-Protection check on a *public* repository.
+          # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action#authentication-with-pat.
+          # repo_token: ${{ secrets.SCORECARD_TOKEN }}
+
+          # Publish results to OpenSSF REST API for easy access by consumers.
+          # See https://github.com/ossf/scorecard-action#publishing-results.
+          publish_results: true
+
+      # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF
+      # format to the repository Actions tab.
+      - name: "Upload artifact"
+        uses: actions/upload-artifact@26f96dfa697d77e81fd5907df203aa23a56210a8 # v4.3.0
+        with:
+          name: SARIF file
+          path: results.sarif
+          retention-days: 5
+
+      # Upload the results to GitHub's code scanning dashboard.
+      - name: "Upload to code-scanning"
+        uses: github/codeql-action/upload-sarif@b7bf0a3ed3ecfa44160715d7c442788f65f0f923 # v3.23.2
+        with:
+          sarif_file: results.sarif
diff --git a/.github/workflows/scripts/auto-assignment.js b/.github/workflows/scripts/auto-assignment.js
new file mode 100644
index 0000000000..d08b06a8b7
--- /dev/null
+++ b/.github/workflows/scripts/auto-assignment.js
@@ -0,0 +1,43 @@
+/** Automatically assign issues and PRs to users in the `assigneesList` 
+ *  on a rotating basis.
+
+  @param {!object}
+    GitHub objects can call GitHub APIs using their built-in library functions.
+    The context object contains issue and PR details.
+*/
+
+module.exports = async ({ github, context }) => {
+  let issueNumber;
+  let assigneesList;
+  // Is this an issue? If so, assign the issue number. Otherwise, assign the PR number.
+  if (context.payload.issue) {
+    //assignee List for issues. 
+    assigneesList = ["SuryanarayanaY", "sachinprasadhs"];
+    issueNumber = context.payload.issue.number;
+  } else {
+    //assignee List for PRs. 
+    assigneesList = [];
+    issueNumber = context.payload.number;
+  }
+  console.log("assignee list", assigneesList);
+  console.log("entered auto assignment for this issue:  ", issueNumber);
+  if (!assigneesList.length) {
+    console.log("No assignees found for this repo.");
+    return;
+  }
+  let noOfAssignees = assigneesList.length;
+  let selection = issueNumber % noOfAssignees;
+  let assigneeForIssue = assigneesList[selection];
+
+  console.log(
+    "issue Number = ",
+    issueNumber + " , assigning to: ",
+    assigneeForIssue
+  );
+  return github.rest.issues.addAssignees({
+    issue_number: context.issue.number,
+    owner: context.repo.owner,
+    repo: context.repo.repo,
+    assignees: [assigneeForIssue],
+  });
+};
diff --git a/.github/workflows/stale-issue-pr.yml b/.github/workflows/stale-issue-pr.yml
new file mode 100644
index 0000000000..034fb4c266
--- /dev/null
+++ b/.github/workflows/stale-issue-pr.yml
@@ -0,0 +1,50 @@
+name: Close inactive issues
+on:
+  schedule:
+    - cron: "30 1 * * *"
+jobs:
+  close-issues:
+    runs-on: ubuntu-latest
+    permissions:
+      issues: write
+      pull-requests: write
+    steps:
+      - name: Awaiting response issues
+        uses: actions/stale@v9
+        with:
+          days-before-issue-stale: 14
+          days-before-issue-close: 14
+          stale-issue-label: "stale"
+          # reason for closed the issue default value is not_planned
+          close-issue-reason: completed
+          only-labels: "stat:awaiting response from contributor"
+          stale-issue-message: > 
+            This issue is stale because it has been open for 14 days with no activity.
+            It will be closed if no further activity occurs. Thank you.
+          # List of labels to remove when issues/PRs unstale. 
+          labels-to-remove-when-unstale: "stat:awaiting response from contributor"
+          close-issue-message: >
+            This issue was closed because it has been inactive for 28 days.
+            Please reopen if you'd like to work on this further.
+          days-before-pr-stale: 14
+          days-before-pr-close: 14
+          stale-pr-message: "This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you."
+          close-pr-message: "This PR was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further."
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Contribution issues
+        uses: actions/stale@v9
+        with:
+          days-before-issue-stale: 180
+          days-before-issue-close: 365
+          stale-issue-label: "stale"
+          # reason for closed the issue default value is not_planned
+          close-issue-reason: not_planned
+          any-of-labels: "stat:contributions welcome,good first issue"
+          # List of labels to remove when issues/PRs unstale. 
+          labels-to-remove-when-unstale: "stat:contributions welcome,good first issue"
+          stale-issue-message: > 
+            This issue is stale because it has been open for 180 days with no activity.
+            It will be closed if no further activity occurs. Thank you.
+          close-issue-message: >
+            This issue was closed because it has been inactive for more than 1 year.
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.gitignore b/.gitignore
index 68d68189bd..566d9d109e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,6 +11,12 @@ build/
 *.egg-info
 __pycache__/
 *.so
+.bazelrc
+bazel-bin
+bazel-keras-cv
+bazel-out
+bazel-testlogs
+wheels
 
 #VS Code files and container
 .vscode/
diff --git a/.kokoro/README.md b/.kokoro/README.md
new file mode 100644
index 0000000000..2c7724d988
--- /dev/null
+++ b/.kokoro/README.md
@@ -0,0 +1 @@
+CI to run on PR and merge to Master.
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/build.sh b/.kokoro/github/ubuntu/gpu/build.sh
new file mode 100644
index 0000000000..9d07218317
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/build.sh
@@ -0,0 +1,81 @@
+set -e
+set -x
+
+cd "${KOKORO_ROOT}/"
+
+sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1
+
+PYTHON_BINARY="/usr/bin/python3.9"
+
+"${PYTHON_BINARY}" -m venv venv
+source venv/bin/activate
+# Check the python version
+python --version
+python3 --version
+
+export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:"
+# Check cuda
+nvidia-smi
+nvcc --version
+
+cd "src/github/keras-cv"
+pip install -U pip setuptools
+
+if [ "${KERAS2:-0}" == "1" ]
+then
+   echo "Keras2 detected."
+   pip install -r requirements-common.txt --progress-bar off
+   pip install tensorflow~=2.14
+   pip install --extra-index-url https://download.pytorch.org/whl/cpu torch==2.1.0+cpu
+   pip install torchvision~=0.16.0
+   pip install "jax[cpu]"
+
+elif [ "$KERAS_BACKEND" == "tensorflow" ]
+then
+   echo "TensorFlow backend detected."
+   pip install -r requirements-tensorflow-cuda.txt --progress-bar off
+
+elif [ "$KERAS_BACKEND" == "jax" ]
+then
+   echo "JAX backend detected."
+   pip install -r requirements-jax-cuda.txt --progress-bar off
+
+elif [ "$KERAS_BACKEND" == "torch" ]
+then
+   echo "PyTorch backend detected."
+   pip install -r requirements-torch-cuda.txt --progress-bar off
+fi
+
+pip install --no-deps -e "." --progress-bar off
+
+# Run Extra Large Tests for Continuous builds
+if [ "${RUN_XLARGE:-0}" == "1" ]
+then
+   pytest --cache-clear --check_gpu --run_large --run_extra_large --durations 0 \
+      keras_cv/bounding_box \
+      keras_cv/callbacks \
+      keras_cv/losses \
+      keras_cv/layers/object_detection \
+      keras_cv/layers/preprocessing \
+      keras_cv/models/backbones \
+      keras_cv/models/classification \
+      keras_cv/models/object_detection/retinanet \
+      keras_cv/models/object_detection/yolo_v8 \
+      keras_cv/models/object_detection_3d \
+      keras_cv/models/segmentation \
+      keras_cv/models/stable_diffusion
+else
+   pytest --cache-clear --check_gpu --run_large --durations 0 \
+      keras_cv/bounding_box \
+      keras_cv/callbacks \
+      keras_cv/losses \
+      keras_cv/layers/object_detection \
+      keras_cv/layers/preprocessing \
+      keras_cv/models/backbones \
+      keras_cv/models/classification \
+      keras_cv/models/object_detection/retinanet \
+      keras_cv/models/object_detection/yolo_v8 \
+      keras_cv/models/object_detection_3d \
+      keras_cv/models/segmentation \
+      keras_cv/models/stable_diffusion
+fi
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/jax/continuous.cfg b/.kokoro/github/ubuntu/gpu/jax/continuous.cfg
new file mode 100644
index 0000000000..2aca2e95ff
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/jax/continuous.cfg
@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "jax"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/jax/presubmit.cfg b/.kokoro/github/ubuntu/gpu/jax/presubmit.cfg
new file mode 100644
index 0000000000..2aca2e95ff
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/jax/presubmit.cfg
@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "jax"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/keras2/continuous.cfg b/.kokoro/github/ubuntu/gpu/keras2/continuous.cfg
new file mode 100644
index 0000000000..d5caba18f7
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/keras2/continuous.cfg
@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS2"
+   value: "1"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg b/.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg
new file mode 100644
index 0000000000..d5caba18f7
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/keras2/presubmit.cfg
@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS2"
+   value: "1"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/tensorflow/continuous.cfg b/.kokoro/github/ubuntu/gpu/tensorflow/continuous.cfg
new file mode 100644
index 0000000000..f7e02e6efa
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/tensorflow/continuous.cfg
@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "tensorflow"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/tensorflow/presubmit.cfg b/.kokoro/github/ubuntu/gpu/tensorflow/presubmit.cfg
new file mode 100644
index 0000000000..f7e02e6efa
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/tensorflow/presubmit.cfg
@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "tensorflow"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/torch/continuous.cfg b/.kokoro/github/ubuntu/gpu/torch/continuous.cfg
new file mode 100644
index 0000000000..a96e865152
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/torch/continuous.cfg
@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "torch"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60
\ No newline at end of file
diff --git a/.kokoro/github/ubuntu/gpu/torch/presubmit.cfg b/.kokoro/github/ubuntu/gpu/torch/presubmit.cfg
new file mode 100644
index 0000000000..a96e865152
--- /dev/null
+++ b/.kokoro/github/ubuntu/gpu/torch/presubmit.cfg
@@ -0,0 +1,16 @@
+build_file: "keras-cv/.kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "torch"
+}
+
+# Set timeout to 60 mins from default 180 mins
+timeout_mins: 60
\ No newline at end of file
diff --git a/.github/API_DESIGN.md b/API_DESIGN.md
similarity index 100%
rename from .github/API_DESIGN.md
rename to API_DESIGN.md
diff --git a/.github/CALL_FOR_CONTRIBUTIONS.md b/CALL_FOR_CONTRIBUTIONS.md
similarity index 100%
rename from .github/CALL_FOR_CONTRIBUTIONS.md
rename to CALL_FOR_CONTRIBUTIONS.md
diff --git a/.github/CONTRIBUTING.md b/CONTRIBUTING.md
similarity index 90%
rename from .github/CONTRIBUTING.md
rename to CONTRIBUTING.md
index 1708f542e2..6cd2487eed 100644
--- a/.github/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -40,7 +40,6 @@ the master branch in [keras-team/keras-cv](https://github.com/keras-team/keras-c
 After creating the pull request, you will need to sign the Google CLA agreement.
 The agreement can be found at [https://cla.developers.google.com/clas](https://cla.developers.google.com/clas).
 
-
 ### Step 5. Code review
 
 CI tests will automatically be run directly on your pull request. Their
@@ -57,6 +56,7 @@ approved by the reviewer.
 Once the pull request is approved, a team member will take care of merging.
 
 ## Contributing models
+
 When contributing new models, please validate model performance by providing training results. You can do this using our existing [ImageNet training script](https://github.com/keras-team/keras-cv/blob/master/examples/training/classification/imagenet/basic_training.py) or by contributing a custom training script of your own (see "Contributing training scripts" below). Training results can be added to the training history log with [this script](https://github.com/keras-team/keras-cv/blob/master/shell/weights/update_training_history.py), or shared with the team via Google Drive (we'll need TensorBoard logs as well as weights). Either way, the KerasCV team will need to upload the weights to our GCS bucket for distribution.
 
 For an initial submission, trained weights do not need to exactly match paper-claimed results. As a baseline, let's shoot for 90% of the paper-claimed ImageNet top-1 accuracy. However, we should strive to improve these weights quickly to at least match paper-claimed results.
@@ -74,6 +74,7 @@ You can also open an issue to add weights for a specific model using a pre-exist
 To contribute a new script, start by opening an issue and tagging @ianstenbit to discuss the task, dataset, and/or model for which you'd like to add a script. Once they've taken a look, you can prepare a PR to introduce the new training script.
 
 See [this example script](https://github.com/keras-team/keras-cv/blob/master/examples/training/classification/imagenet/basic_training.py) for training ImageNet classification. Please follow the structure of this training script in contributing your own script. New scripts should either:
+
 - Train a task for which we don't have a training script already
 - Include a meaningfully different training approach for a given task
 - Introduce a custom training method for a specific model or dataset, based on empirical evidence of efficacy.
@@ -103,7 +104,7 @@ You can achieve this by running the following commands:
 gh repo fork keras-team/keras-cv --clone --remote
 cd keras-cv
 pip install ".[tests]"
-python setup.py develop
+pip install -e .
 ```
 
 The first line relies on having an installation of [the GitHub CLI](https://github.com/cli/cli).
@@ -111,22 +112,13 @@ The first line relies on having an installation of [the GitHub CLI](https://gith
 Following these commands you should be able to run the tests using `pytest keras_cv`.
 Please report any issues running tests following these steps.
 
-Note that this will _not_ install custom ops. If you'd like to install custom ops from source, you can compile the binaries and add them to your local environment manually (requires Bazel):
-
-```shell
-python build_deps/configure.py
-
-bazel build keras_cv/custom_ops:all
-mv bazel-bin/keras_cv/custom_ops/*.so keras_cv/custom_ops
-```
-
 ## Run tests
 
 KerasCV is tested using [PyTest](https://docs.pytest.org/en/6.2.x/).
 
 ### Run a test file
 
-To run a test file, run `pytest path/to/file` from the root directory of keras\_cv.
+To run a test file, run `pytest path/to/file` from the root directory of keras_cv.
 
 ### Run a single test case
 
@@ -142,21 +134,13 @@ pytest keras_cv/layers/preprocessing/cut_mix_test.py -k="label"
 ### Run all tests
 
 You can run the unit tests for KerasCV by running:
+
 ```
 pytest keras_cv/
 ```
 
-### Tests that require custom ops
-For tests that require custom ops, you'll have to compile the custom ops and make them available to your local Python code:
-```shell
-python build_deps/configure.py
-bazel build keras_cv/custom_ops:all
-cp bazel-bin/keras_cv/custom_ops/*.so keras_cv/custom_ops/
-```
-
-Tests which use custom ops are disabled by default, but can be run by setting the environment variable `TEST_CUSTOM_OPS=true`.
-
 ## Formatting the Code
+
 We use `flake8`, `isort`, `black` and `clang-format` for code formatting. You can run
 the following commands manually every time you want to format your code:
 
@@ -168,7 +152,7 @@ This can be done by running `pip install --upgrade black`, `pip install --upgrad
 `pip install --upgrade isort` and `pip install --upgrade clang-format`
 
 Note: The linting checks could be automated activating  
-      pre-commit hooks with `git config core.hooksPath .github/.githooks`
+ pre-commit hooks with `git config core.hooksPath .github/.githooks`
 
 ## Community Guidelines
 
diff --git a/README.md b/README.md
index ad6bf5dc8c..3872931397 100644
--- a/README.md
+++ b/README.md
@@ -6,24 +6,24 @@
 ![Tensorflow](https://img.shields.io/badge/tensorflow-v2.9.0+-success.svg)
 [![Contributions Welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/keras-team/keras-cv/issues)
 
-KerasCV is a library of modular computer vision components that work natively 
-with TensorFlow, JAX, or PyTorch. Built on [Keras Core](https://keras.io/keras_core/announcement/),
-these models, layers, metrics, callbacks, etc., can be trained and serialized
-in any framework and re-used in another without costly migrations. See "Using 
-KerasCV with Keras Core" below for more details on multi-framework KerasCV.
+KerasCV is a library of modular computer vision components that work natively
+with TensorFlow, JAX, or PyTorch. Built on Keras 3, these models, layers, 
+metrics, callbacks, etc., can be trained and serialized in any framework and 
+re-used in another without costly migrations. See "Configuring your backend" 
+below for more details on multi-framework KerasCV.
 
 <img style="width: 440px; max-width: 90%;" src="https://storage.googleapis.com/keras-cv/guides/keras-cv-augmentations.gif">
 
-KerasCV can be understood as a horizontal extension of the Keras API: the 
-components are new first-party Keras objects that are too specialized to be 
-added to core Keras. They receive the same level of polish and backwards 
-compatibility guarantees as the core Keras API, and they are maintained by the 
+KerasCV can be understood as a horizontal extension of the Keras API: the
+components are new first-party Keras objects that are too specialized to be
+added to core Keras. They receive the same level of polish and backwards
+compatibility guarantees as the core Keras API, and they are maintained by the
 Keras team.
 
-Our APIs assist in common computer vision tasks such as data augmentation, 
+Our APIs assist in common computer vision tasks such as data augmentation,
 classification, object detection, segmentation, image generation, and more.
-Applied computer vision engineers can leverage KerasCV to quickly assemble 
-production-grade, state-of-the-art training and inference pipelines for all of 
+Applied computer vision engineers can leverage KerasCV to quickly assemble
+production-grade, state-of-the-art training and inference pipelines for all of
 these common tasks.
 
 ## Quick Links
@@ -34,39 +34,69 @@ these common tasks.
 - [API Design Guidelines](.github/API_DESIGN.md)
 
 ## Installation
+KerasCV supports both Keras 2 and Keras 3. We recommend Keras 3 for all new 
+users, as it enables using KerasCV models and layers with JAX, TensorFlow and 
+PyTorch.
+
+### Keras 2 Installation
+
+To install the latest KerasCV release with Keras 2, simply run:
+
+```
+pip install --upgrade keras-cv tensorflow
+```
+
+### Keras 3 Installation
+
+There are currently two ways to install Keras 3 with KerasCV. To install the
+latest changes for KerasCV and Keras, you can use our nightly package.
 
-To install the latest official release:
 
 ```
-pip install keras-cv tensorflow --upgrade
+pip install --upgrade keras-cv-nightly tf-nightly
 ```
 
-To install the latest unreleased changes to the library, we recommend using
-pip to install directly from the master branch on github:
+To install the stable versions of KerasCV and Keras 3, you should install Keras
+3 **after** installing KerasCV. This is a temporary step while TensorFlow is
+pinned to Keras 2, and will no longer be necessary after TensorFlow 2.16.
 
 ```
-pip install git+https://github.com/keras-team/keras-cv.git tensorflow --upgrade
+pip install --upgrade keras-cv tensorflow
+pip install --upgrade keras
+```
+> [!IMPORTANT]
+> Keras 3 will not function with TensorFlow 2.14 or earlier.
+
+## Configuring your backend
+
+If you have Keras 3 installed in your environment (see installation above),
+you can use KerasCV with any of JAX, TensorFlow and PyTorch. To do so, set the
+`KERAS_BACKEND` environment variable. For example:
+so by setting the `KERAS_BACKEND` environment variable. For example:
+
+```shell
+export KERAS_BACKEND=jax
 ```
 
-## Using KerasCV with Keras Core
+Or in Colab, with:
+
+```python
+import os
+os.environ["KERAS_BACKEND"] = "jax"
 
-As of version `0.6.0`, KerasCV supports multiple backends with Keras Core out of 
-the box. There are two ways to configure KerasCV to run with multi-backend 
-support:
+import keras_cv
+```
 
-1. Via the `KERAS_BACKEND` environment variable. If set, then KerasCV will be 
-using Keras Core with the backend specified (e.g., `KERAS_BACKEND=jax`).
-2. Via the `.keras/keras.json` and `.keras/keras_cv.json` config files (which 
-are automatically created the first time you import KerasCV):
-   - Set your backend of choice in `.keras/keras.json`; e.g., `"backend": "jax"`. 
-   - Set `"multi_backend": True` in `.keras/keras_cv.json`.
+> [!IMPORTANT]
+> Make sure to set the `KERAS_BACKEND` before import any Keras libraries, it
+> will be used to set up Keras when it is first imported.
 
-Once that configuration step is done, you can just import KerasCV and start 
+Once that configuration step is done, you can just import KerasCV and start
 using it on top of your backend of choice:
 
 ```python
 import keras_cv
-import keras_core as keras
+import keras
 
 filepath = keras.utils.get_file(origin="https://i.imgur.com/gCNcJJI.jpg")
 image = np.array(keras.utils.load_img(filepath))
@@ -79,19 +109,13 @@ model = keras_cv.models.YOLOV8Detector.from_preset(
 predictions = model.predict(image_resized)
 ```
 
-Until Keras Core is officially released as Keras 3.0, KerasCV will use 
-`tf.keras` as the default backend. To restore this default behavior, simply 
-`unset KERAS_BACKEND` and ensure that  `"multi_backend": False` or is unset in 
-`.keras/keras_cv.json`. You will need to restart the Python runtime for changes 
-to take effect.
-
 ## Quickstart
 
 ```python
 import tensorflow as tf
 import keras_cv
 import tensorflow_datasets as tfds
-import keras_core as keras
+import keras
 
 # Create a preprocessing pipeline with augmentations
 BATCH_SIZE = 16
@@ -157,9 +181,9 @@ We would like to leverage/outsource the Keras community not only for bug reporti
 but also for active development for feature delivery. To achieve this, here is the predefined
 process for how to contribute to this repository:
 
-1) Contributors are always welcome to help us fix an issue, add tests, better documentation. 
+1) Contributors are always welcome to help us fix an issue, add tests, better documentation.
 2) If contributors would like to create a backbone, we usually require a pre-trained weight set
-with the model for one dataset as the first PR, and a training script as a follow-up. The training script will preferably help us reproduce the results claimed from paper. The backbone should be generic but the training script can contain paper specific parameters such as learning rate schedules and weight decays. The training script will be used to produce leaderboard results. 
+with the model for one dataset as the first PR, and a training script as a follow-up. The training script will preferably help us reproduce the results claimed from paper. The backbone should be generic but the training script can contain paper specific parameters such as learning rate schedules and weight decays. The training script will be used to produce leaderboard results.
 Exceptions apply to large transformer-based models which are difficult to train. If this is the case,
 contributors should let us know so the team can help in training the model or providing GCP resources.
 3) If contributors would like to create a meta arch, please try to be aligned with our roadmap and create a PR for design review to make sure the meta arch is modular.
@@ -185,7 +209,7 @@ An example of this can be found in the ImageNet classification training
 All results are reproducible using the training scripts in this repository.
 
 Historically, many models have been trained on image datasets rescaled via manually
-crafted normalization schemes. 
+crafted normalization schemes.
 The most common variant of manually crafted normalization scheme is subtraction of the
 imagenet mean pixel followed by standard deviation normalization based on the imagenet
 pixel standard deviation.
@@ -233,7 +257,7 @@ These pre-trained models are provided on an "as is" basis, without warranties
 or conditions of any kind.
 The following underlying models are provided by third parties, and are subject to separate
 licenses:
-StableDiffusion, Vision Transfomer
+StableDiffusion, Vision Transformer
 
 ## Citing KerasCV
 
@@ -243,7 +267,7 @@ Here is the BibTeX entry:
 ```bibtex
 @misc{wood2022kerascv,
   title={KerasCV},
-  author={Wood, Luke and Tan, Zhenyu and Stenbit, Ian and Bischof, Jonathan and Zhu, Scott and Chollet, Fran\c{c}ois and others},
+  author={Wood, Luke and Tan, Zhenyu and Stenbit, Ian and Bischof, Jonathan and Zhu, Scott and Chollet, Fran\c{c}ois and Sreepathihalli, Divyashree and Sampath, Ramesh and others},
   year={2022},
   howpublished={\url{https://github.com/keras-team/keras-cv}},
 }
diff --git a/SECURITY.md b/SECURITY.md
new file mode 100644
index 0000000000..96015a7125
--- /dev/null
+++ b/SECURITY.md
@@ -0,0 +1,18 @@
+# Security Policy
+
+If you have discovered a security vulnerability in this project, please report it
+privately. **Do not disclose it as a public issue.** This gives us time to work with you
+to fix the issue before public exposure, reducing the chance that the exploit will be
+used before a patch is released.
+
+You may submit the report in the following ways:
+
+- send a [private vulnerability report](https://github.com/keras-team/keras-cv/security/advisories/new)
+
+Please provide the following information in your report:
+
+- A description of the vulnerability and its impact
+- How to reproduce the issue
+
+This project is maintained by volunteers on a reasonable-effort basis. As such,
+please give us 90 days to work on a fix before public exposure.
diff --git a/benchmarks/vectorized_jittered_resize.py b/benchmarks/vectorized_jittered_resize.py
index fe2eb26233..0faf1e4095 100644
--- a/benchmarks/vectorized_jittered_resize.py
+++ b/benchmarks/vectorized_jittered_resize.py
@@ -259,7 +259,7 @@ def test_consistency_with_old_impl(self):
         # makes offsets fixed to (0.5, 0.5)
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=tf.convert_to_tensor([[0.5, 0.5]]),
         ):
             output = layer(image)
diff --git a/benchmarks/vectorized_mosaic.py b/benchmarks/vectorized_mosaic.py
index 745022764f..54550358c1 100644
--- a/benchmarks/vectorized_mosaic.py
+++ b/benchmarks/vectorized_mosaic.py
@@ -101,7 +101,6 @@ def _batch_augment(self, inputs):
             minval=0,
             maxval=batch_size,
             dtype=tf.int32,
-            seed=self._random_generator.make_legacy_seed(),
         )
         # concatenate the batches with permutation order to get all 4 images of
         # the mosaic
diff --git a/benchmarks/vectorized_random_brightness.py b/benchmarks/vectorized_random_brightness.py
index 9babe0d692..115caa86ad 100644
--- a/benchmarks/vectorized_random_brightness.py
+++ b/benchmarks/vectorized_random_brightness.py
@@ -61,7 +61,7 @@ class OldRandomBrightness(BaseImageAugmentationLayer):
     """
 
     def __init__(self, factor, value_range=(0, 255), seed=None, **kwargs):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         if isinstance(factor, float) or isinstance(factor, int):
             factor = (-factor, factor)
         self.factor = preprocessing_utils.parse_factor(
diff --git a/benchmarks/vectorized_random_contrast.py b/benchmarks/vectorized_random_contrast.py
index 32ce5fbde6..7d5c7dbafe 100644
--- a/benchmarks/vectorized_random_contrast.py
+++ b/benchmarks/vectorized_random_contrast.py
@@ -60,7 +60,7 @@ class OldRandomContrast(BaseImageAugmentationLayer):
     """
 
     def __init__(self, factor, seed=None, **kwargs):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         if isinstance(factor, (tuple, list)):
             min = 1 - factor[0]
             max = 1 + factor[1]
diff --git a/benchmarks/vectorized_random_crop.py b/benchmarks/vectorized_random_crop.py
index bb26cc4cec..f932e7bda0 100644
--- a/benchmarks/vectorized_random_crop.py
+++ b/benchmarks/vectorized_random_crop.py
@@ -59,7 +59,9 @@ def __init__(
         self, height, width, seed=None, bounding_box_format=None, **kwargs
     ):
         super().__init__(
-            **kwargs, autocast=False, seed=seed, force_generator=True
+            **kwargs,
+            autocast=False,
+            seed=seed,
         )
         self.height = height
         self.width = width
@@ -72,7 +74,7 @@ def get_random_transformation(self, image=None, **kwargs):
         h_diff = image_shape[H_AXIS] - self.height
         w_diff = image_shape[W_AXIS] - self.width
         dtype = image_shape.dtype
-        rands = self._random_generator.random_uniform([2], 0, dtype.max, dtype)
+        rands = self._random_generator.uniform([2], 0, dtype.max, dtype)
         h_start = rands[0] % (h_diff + 1)
         w_start = rands[1] % (w_diff + 1)
         return {"top": h_start, "left": w_start}
diff --git a/benchmarks/vectorized_random_flip.py b/benchmarks/vectorized_random_flip.py
index 1a2da5e284..5fca5fd661 100644
--- a/benchmarks/vectorized_random_flip.py
+++ b/benchmarks/vectorized_random_flip.py
@@ -71,7 +71,7 @@ class OldRandomFlip(BaseImageAugmentationLayer):
     def __init__(
         self, mode=HORIZONTAL, seed=None, bounding_box_format=None, **kwargs
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.mode = mode
         self.seed = seed
         if mode == HORIZONTAL:
@@ -101,13 +101,9 @@ def get_random_transformation(self, **kwargs):
         flip_horizontal = False
         flip_vertical = False
         if self.horizontal:
-            flip_horizontal = (
-                self._random_generator.random_uniform(shape=[]) > 0.5
-            )
+            flip_horizontal = self._random_generator.uniform(shape=[]) > 0.5
         if self.vertical:
-            flip_vertical = (
-                self._random_generator.random_uniform(shape=[]) > 0.5
-            )
+            flip_vertical = self._random_generator.uniform(shape=[]) > 0.5
         return {
             "flip_horizontal": tf.cast(flip_horizontal, dtype=tf.bool),
             "flip_vertical": tf.cast(flip_vertical, dtype=tf.bool),
@@ -237,13 +233,13 @@ def test_consistency_with_old_impl(self):
 
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=tf.convert_to_tensor([[0.6]]),
         ):
             output = layer(image)
         with unittest.mock.patch.object(
             old_layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=tf.convert_to_tensor(0.6),
         ):
             old_output = old_layer(image)
diff --git a/benchmarks/vectorized_random_rotation.py b/benchmarks/vectorized_random_rotation.py
index 0c8821f632..ea852e0a66 100644
--- a/benchmarks/vectorized_random_rotation.py
+++ b/benchmarks/vectorized_random_rotation.py
@@ -97,7 +97,7 @@ def __init__(
         segmentation_classes=None,
         **kwargs,
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.factor = factor
         if isinstance(factor, (tuple, list)):
             self.lower = factor[0]
@@ -122,7 +122,7 @@ def __init__(
     def get_random_transformation(self, **kwargs):
         min_angle = self.lower * 2.0 * np.pi
         max_angle = self.upper * 2.0 * np.pi
-        angle = self._random_generator.random_uniform(
+        angle = self._random_generator.uniform(
             shape=[1], minval=min_angle, maxval=max_angle
         )
         return {"angle": angle}
diff --git a/benchmarks/vectorized_random_translation.py b/benchmarks/vectorized_random_translation.py
index 2146d7cae4..640f3c7933 100644
--- a/benchmarks/vectorized_random_translation.py
+++ b/benchmarks/vectorized_random_translation.py
@@ -141,7 +141,7 @@ def __init__(
         fill_value=0.0,
         **kwargs,
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.height_factor = height_factor
         if isinstance(height_factor, (tuple, list)):
             self.height_lower = height_factor[0]
@@ -217,13 +217,13 @@ def augment_image(self, image, transformation, **kwargs):
 
     def get_random_transformation(self, image=None, **kwargs):
         batch_size = 1
-        height_translation = self._random_generator.random_uniform(
+        height_translation = self._random_generator.uniform(
             shape=[batch_size, 1],
             minval=self.height_lower,
             maxval=self.height_upper,
             dtype=tf.float32,
         )
-        width_translation = self._random_generator.random_uniform(
+        width_translation = self._random_generator.uniform(
             shape=[batch_size, 1],
             minval=self.width_lower,
             maxval=self.width_upper,
diff --git a/benchmarks/vectorized_random_zoom.py b/benchmarks/vectorized_random_zoom.py
index b44919e3ab..0d3c84ce8d 100644
--- a/benchmarks/vectorized_random_zoom.py
+++ b/benchmarks/vectorized_random_zoom.py
@@ -103,7 +103,7 @@ def __init__(
         fill_value=0.0,
         **kwargs,
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.height_factor = height_factor
         if isinstance(height_factor, (tuple, list)):
             self.height_lower = height_factor[0]
@@ -143,13 +143,13 @@ def __init__(
         self.seed = seed
 
     def get_random_transformation(self, image=None, **kwargs):
-        height_zoom = self._random_generator.random_uniform(
+        height_zoom = self._random_generator.uniform(
             shape=[1, 1],
             minval=1.0 + self.height_lower,
             maxval=1.0 + self.height_upper,
         )
         if self.width_factor is not None:
-            width_zoom = self._random_generator.random_uniform(
+            width_zoom = self._random_generator.uniform(
                 shape=[1, 1],
                 minval=1.0 + self.width_lower,
                 maxval=1.0 + self.width_upper,
diff --git a/benchmarks/vectorized_randomly_zoomed_crop.py b/benchmarks/vectorized_randomly_zoomed_crop.py
index 66a12fec41..3a207ed2e3 100644
--- a/benchmarks/vectorized_randomly_zoomed_crop.py
+++ b/benchmarks/vectorized_randomly_zoomed_crop.py
@@ -109,14 +109,14 @@ def get_random_transformation(
 
         new_width = crop_size[1] * tf.sqrt(aspect_ratio)
 
-        height_offset = self._random_generator.random_uniform(
+        height_offset = self._random_generator.uniform(
             (),
             minval=tf.minimum(0.0, original_height - new_height),
             maxval=tf.maximum(0.0, original_height - new_height),
             dtype=tf.float32,
         )
 
-        width_offset = self._random_generator.random_uniform(
+        width_offset = self._random_generator.uniform(
             (),
             minval=tf.minimum(0.0, original_width - new_width),
             maxval=tf.maximum(0.0, original_width - new_width),
@@ -249,10 +249,10 @@ def from_config(cls, config):
                 config["zoom_factor"]
             )
         if isinstance(config["aspect_ratio_factor"], dict):
-            config[
-                "aspect_ratio_factor"
-            ] = keras.utils.deserialize_keras_object(
-                config["aspect_ratio_factor"]
+            config["aspect_ratio_factor"] = (
+                keras.utils.deserialize_keras_object(
+                    config["aspect_ratio_factor"]
+                )
             )
         return cls(**config)
 
diff --git a/cloudbuild/Dockerfile b/cloudbuild/Dockerfile
deleted file mode 100644
index 5aca71f5c2..0000000000
--- a/cloudbuild/Dockerfile
+++ /dev/null
@@ -1,4 +0,0 @@
-ARG IMAGE_NAME
-FROM $IMAGE_NAME
-COPY . /kerascv
-WORKDIR /kerascv
diff --git a/cloudbuild/README.md b/cloudbuild/README.md
deleted file mode 100644
index e6a8ba581d..0000000000
--- a/cloudbuild/README.md
+++ /dev/null
@@ -1,72 +0,0 @@
-# KerasCV Accelerators Testing
-
-This `cloudbuild/` directory contains configurations for accelerators (GPU/TPU)
-testing. Briefly, for each PR, it copies the PR's code to a base docker image
-which contains KerasCV dependencies to make a new docker image, and deploys the
-new image to Google Kubernetes Engine cluster, then run all tests in
-`keras_cv/` via Google Cloud Build.
-
-- `cloudbuild.yaml`: The cloud build configuration that specifies steps to run
-  by cloud build.
-- `Dockerfile`: The configuration to build the docker image for deployment.
-- `unit_test_jobs.jsonnet`: Jsonnet config that tells GKE cluster to run all
-  unit tests in `keras_cv/`.
-
-This test is powered by [ml-testing-accelerators](https://github.com/GoogleCloudPlatform/ml-testing-accelerators).
-
-
-### Adding Test Dependencies
-You must be authorized to run builds in the `keras-team-test` GCP project.
-If you are not, please open a GitHub issue and ping a team member.
-To authorize yourself with `keras-team-test`, run:
-
-```bash
-gcloud config set project keras-team-test
-```
-
-To add a dependency for GPU tests:
-- Create a PR adding the dependency to `requirements.txt`
-- Have a Keras team member update the Docker image for GPU tests by running the remaining steps
-- Create a `Dockerfile` with the following contents:
-```
-FROM tensorflow/tensorflow:2.13.0-gpu
-RUN \
-    apt-get -y update && \
-    apt-get -y install openjdk-8-jdk && \
-    echo "deb [arch=amd64] http://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list && \
-    curl https://bazel.build/bazel-release.pub.gpg | apt-key add
-RUN apt-get -y update
-RUN apt-get -y install bazel-5.4.0
-RUN apt-get -y install git
-RUN git clone https://github.com/{path_to_keras_cv_fork}.git
-RUN cd keras-cv && git checkout {branch_name}
-RUN pip install -r keras-cv/requirements.txt
-```
-- Run the following command from the directory with your `Dockerfile`:
-```
-gcloud builds submit --region=us-west1 --tag us-west1-docker.pkg.dev/keras-team-test/keras-cv-test/keras-cv-image-tensorflow:deps --timeout=20m
-```
-- Repeat the last two steps for Jax and Torch (replacing "tensorflow" with "jax"
- or "torch" in the docker image target name). `Dockerfile` for jax:
-```
-FROM nvidia/cuda:11.7.1-base-ubuntu20.04
-RUN apt-get update
-RUN apt-get install -y python3 python3-pip
-RUN apt-get install -y git
-RUN git clone https://github.com/{path_to_keras_cv_fork}.git
-RUN cd keras-cv && git checkout {branch_name}
-RUN pip install -r keras-cv/requirements.txt
-RUN pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-```
-  and for torch:
-```
-FROM nvidia/cuda:11.7.1-base-ubuntu20.04
-RUN apt-get update
-RUN apt-get install -y python3 python3-pip
-RUN apt-get install -y git
-RUN git clone https://github.com/{path_to_keras_cv_fork}.git
-RUN cd keras-cv && git checkout {branch_name}
-RUN pip install -r keras-cv/requirements.txt
-RUN pip install torch torchvision
-```
-- Merge the PR adding the dependency
diff --git a/cloudbuild/cloudbuild.yaml b/cloudbuild/cloudbuild.yaml
deleted file mode 100644
index 6fa3e1fd5f..0000000000
--- a/cloudbuild/cloudbuild.yaml
+++ /dev/null
@@ -1,77 +0,0 @@
-substitutions:
-  # GCS bucket name.
-  _GCS_BUCKET: 'gs://keras-cv-github-test'
-  # GKE cluster name.
-  _CLUSTER_NAME: 'keras-cv-test-cluster'
-  # Location of GKE cluster.
-  _CLUSTER_ZONE: 'us-west1-b'
-  # Image name.
-  _IMAGE_NAME: 'us-west1-docker.pkg.dev/keras-team-test/keras-cv-test/keras-cv-image-${_BACKEND}'
-steps:
-- name: 'gcr.io/cloud-builders/docker'
-  id: build-image
-  entrypoint: 'bash'
-  args:
-      ['-c', 'docker build -f cloudbuild/Dockerfile -t $_IMAGE_NAME:$BUILD_ID --build-arg IMAGE_NAME=$_IMAGE_NAME:deps .']
-- name: 'gcr.io/cloud-builders/docker'
-  id: push-image
-  waitFor:
-  - build-image
-  args: ['push', '$_IMAGE_NAME:$BUILD_ID']
-- name: 'golang'
-  id: download-jsonnet
-  waitFor: ['-']
-  entrypoint: 'go'
-  args: [
-    'install',
-    'github.com/google/go-jsonnet/cmd/jsonnet@latest',
-  ]
-- name: 'gcr.io/cloud-builders/gcloud'
-  id: clone-templates
-  waitFor: ['-']
-  entrypoint: 'git'
-  args: [
-    'clone',
-    'https://github.com/GoogleCloudPlatform/ml-testing-accelerators.git',
-  ]
-- name: 'golang'
-  id: build-templates
-  waitFor:
-  - download-jsonnet
-  - clone-templates
-  entrypoint: 'jsonnet'
-  args: [
-    'cloudbuild/unit_test_jobs.jsonnet',
-    '--string',
-    '-J', 'ml-testing-accelerators',
-    '--ext-str', 'image=$_IMAGE_NAME',
-    '--ext-str', 'tag_name=$BUILD_ID',
-    '--ext-str', 'gcs_bucket=$_GCS_BUCKET',
-    '--ext-str', 'backend=$_BACKEND',
-    '-o', 'output.yaml',
-  ]
-- name: 'gcr.io/cloud-builders/gcloud'
-  id: create-job
-  waitFor:
-  - push-image
-  - build-templates
-  entrypoint: bash
-  args:
-  - -c
-  - |
-    set -u
-    set -e
-    set -x
-    gcloud container clusters get-credentials $_CLUSTER_NAME --zone $_CLUSTER_ZONE --project keras-team-test
-    job_name=$(kubectl create -f output.yaml -o name)
-    sleep 5
-    pod_name=$(kubectl wait --for condition=ready --timeout=10m pod -l job-name=${job_name#job.batch/} -o name)
-    kubectl logs -f $pod_name --container=train
-    sleep 5
-    gcloud artifacts docker images delete $_IMAGE_NAME:$BUILD_ID
-    exit $(kubectl get $pod_name -o jsonpath={.status.containerStatuses[0].state.terminated.exitCode})
-timeout: 1800s  # 30 minutes
-options:
-  volumes:
-  - name: go-modules
-    path: /go
diff --git a/cloudbuild/unit_test_jobs.jsonnet b/cloudbuild/unit_test_jobs.jsonnet
deleted file mode 100644
index d688bf90b9..0000000000
--- a/cloudbuild/unit_test_jobs.jsonnet
+++ /dev/null
@@ -1,53 +0,0 @@
-local base = import 'templates/base.libsonnet';
-local gpus = import 'templates/gpus.libsonnet';
-
-local image = std.extVar('image');
-local tagName = std.extVar('tag_name');
-local gcsBucket = std.extVar('gcs_bucket');
-local backend = std.extVar('backend');
-
-local unittest = base.BaseTest {
-  // Configure job name.
-  frameworkPrefix: backend,
-  modelName: "keras-cv",
-  mode: "unit-tests",
-  timeout: 3600, # 1 hour, in seconds
-
-  // Set up runtime environment.
-  image: image,
-  imageTag: tagName,
-  accelerator: gpus.teslaT4,
-  outputBucket: gcsBucket,
-
-  entrypoint: [
-    'bash',
-    '-c',
-    std.format(
-      |||
-        export KERAS_BACKEND=%s
-        export JAX_ENABLE_X64=true
-
-        # Run whatever is in `command` here.
-        ${@:0}
-      |||,
-      backend
-    )
-  ],
-  command: [
-    'pytest --run_large --durations 0',
-    'keras_cv/bounding_box',
-    'keras_cv/callbacks',
-    'keras_cv/losses',
-    'keras_cv/layers/object_detection',
-    'keras_cv/layers/preprocessing',
-    'keras_cv/models/backbones',
-    'keras_cv/models/classification',
-    'keras_cv/models/object_detection/retinanet',
-    'keras_cv/models/object_detection/yolo_v8',
-    'keras_cv/models/object_detection_3d',
-    'keras_cv/models/segmentation',
-    'keras_cv/models/stable_diffusion',
-  ],
-};
-
-std.manifestYamlDoc(unittest.oneshotJob, quote_keys=false)
diff --git a/examples/layers/preprocessing/classification/random_crop_and_zoom.py b/examples/layers/preprocessing/classification/random_crop_and_resize_demo.py
similarity index 90%
rename from examples/layers/preprocessing/classification/random_crop_and_zoom.py
rename to examples/layers/preprocessing/classification/random_crop_and_resize_demo.py
index 3fe8aa1e00..cb8f9e5ffe 100644
--- a/examples/layers/preprocessing/classification/random_crop_and_zoom.py
+++ b/examples/layers/preprocessing/classification/random_crop_and_resize_demo.py
@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""random_resized_crop_demo.py.py shows how to use the RandomResizedCrop
-preprocessing layer. Operates on an image of elephant. In this script the image
+"""This demo example shows how to use the RandomCropAndResize preprocessing
+layer. Operates on an image of elephant. In this script the image
 is loaded, then are passed through the preprocessing layers.
 Finally, they are shown using matplotlib.
 """
diff --git a/keras_cv/__init__.py b/keras_cv/__init__.py
index db96e7ba4a..36c7d3511b 100644
--- a/keras_cv/__init__.py
+++ b/keras_cv/__init__.py
@@ -41,5 +41,5 @@
 from keras_cv.core import FactorSampler
 from keras_cv.core import NormalFactorSampler
 from keras_cv.core import UniformFactorSampler
-
-__version__ = "0.6.2"
+from keras_cv.version_utils import __version__
+from keras_cv.version_utils import version
diff --git a/keras_cv/backend/__init__.py b/keras_cv/backend/__init__.py
index 5375c97596..2190534ff8 100644
--- a/keras_cv/backend/__init__.py
+++ b/keras_cv/backend/__init__.py
@@ -14,81 +14,31 @@
 """
 Keras backend module.
 
-This module adds a temporarily Keras API surface that is fully under KerasCV
-control. This allows us to switch between `keras_core` and `tf.keras`, as well
-as add shims to support older version of `tf.keras`.
-- `config`: check which backend is being run.
-- `keras`: The full `keras` API (via `keras_core` or `tf.keras`).
-- `ops`: `keras_core.ops`, always tf-backed if using `tf.keras`.
+This module adds a temporary Keras API surface that is fully under KerasCV
+control. The goal is to allow us to write Keras 3-like code everywhere, while
+still supporting Keras 2. We do this by using the `keras_core` package with
+Keras 2 to backport Keras 3 numerics APIs (`keras.ops` and `keras.random`) into
+Keras 2. The sub-modules exposed are as follows:
+
+- `config`: check which version of Keras is being run.
+- `keras`: The full `keras` API with compat shims for older Keras versions.
+- `ops`: `keras.ops` for Keras 3 or `keras_core.ops` for Keras 2.
+- `random`: `keras.random` for Keras 3 or `keras_core.ops` for Keras 2.
 """
-
-import types
-
-from keras_cv.backend.config import multi_backend
-
-# Keys are of the form: "module.where.attr.exists->module.where.to.alias"
-# Value are of the form: ["attr1", "attr2", ...] or
-#                        [("attr1_original_name", "attr1_alias_name")]
-_KERAS_CORE_ALIASES = {
-    "utils->saving": [
-        "register_keras_serializable",
-        "deserialize_keras_object",
-        "serialize_keras_object",
-        "get_registered_object",
-    ],
-    "models->saving": ["load_model"],
-}
-
-
-if multi_backend():
-    import keras_core as keras
-
-    keras.backend.name_scope = keras.name_scope
-else:
-    from tensorflow import keras
-
-    if not hasattr(keras, "saving"):
-        keras.saving = types.SimpleNamespace()
-
-    # add aliases
-    for key, value in _KERAS_CORE_ALIASES.items():
-        src, _, dst = key.partition("->")
-        src = src.split(".")
-        dst = dst.split(".")
-
-        src_mod, dst_mod = keras, keras
-
-        # navigate to where we want to alias the attributes
-        for mod in src:
-            src_mod = getattr(src_mod, mod)
-        for mod in dst:
-            dst_mod = getattr(dst_mod, mod)
-
-        # add an alias for each attribute
-        for attr in value:
-            if isinstance(attr, tuple):
-                src_attr, dst_attr = attr
-            else:
-                src_attr, dst_attr = attr, attr
-            attr_val = getattr(src_mod, src_attr)
-            setattr(dst_mod, dst_attr, attr_val)
-
-    # TF Keras doesn't have this rename.
-    keras.activations.silu = keras.activations.swish
-
 from keras_cv.backend import config  # noqa: E402
+from keras_cv.backend import keras  # noqa: E402
 from keras_cv.backend import ops  # noqa: E402
 from keras_cv.backend import random  # noqa: E402
 from keras_cv.backend import tf_ops  # noqa: E402
 
 
 def assert_tf_keras(src):
-    if multi_backend():
+    if config.keras_3():
         raise NotImplementedError(
-            f"KerasCV component {src} does not yet support Keras Core, and can "
+            f"KerasCV component {src} does not yet support Keras 3, and can "
             "only be used in `tf.keras`."
         )
 
 
 def supports_ragged():
-    return not multi_backend()
+    return not config.keras_3()
diff --git a/keras_cv/backend/config.py b/keras_cv/backend/config.py
index b688b79291..00380c113b 100644
--- a/keras_cv/backend/config.py
+++ b/keras_cv/backend/config.py
@@ -11,54 +11,41 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import json
-import os
 
-_MULTI_BACKEND = False
 
-# Set Keras base dir path given KERAS_HOME env variable, if applicable.
-# Otherwise either ~/.keras or /tmp.
-if "KERAS_HOME" in os.environ:
-    _keras_dir = os.environ.get("KERAS_HOME")
-else:
-    _keras_base_dir = os.path.expanduser("~")
-    if not os.access(_keras_base_dir, os.W_OK):
-        _keras_base_dir = "/tmp"
-    _keras_dir = os.path.join(_keras_base_dir, ".keras")
-
-# Attempt to read KerasCV config file.
-_config_path = os.path.expanduser(os.path.join(_keras_dir, "keras-cv.json"))
-if os.path.exists(_config_path):
+def detect_if_tensorflow_uses_keras_3():
+    # We follow the version of keras that tensorflow is configured to use.
     try:
-        with open(_config_path) as f:
-            _config = json.load(f)
-    except ValueError:
-        _config = {}
-    _MULTI_BACKEND = _config.get("multi_backend", _MULTI_BACKEND)
+        from tensorflow import keras
 
-# Save config file, if possible.
-if not os.path.exists(_keras_dir):
-    try:
-        os.makedirs(_keras_dir)
-    except OSError:
-        # Except permission denied and potential race conditions
-        # in multi-threaded environments.
-        pass
+        # Note that only recent versions of keras have a `version()` function.
+        if hasattr(keras, "version") and keras.version().startswith("3."):
+            return True
+    except:
+        raise ValueError(
+            "Unable to import `keras` with `tensorflow`.  Please check your "
+            "Keras and Tensorflow version are compatible; Keras 3 requires "
+            "TensorFlow 2.15 or later. See keras.io/getting_started for more "
+            "information on installing Keras."
+        )
+
+    # No `keras.version()` means we are on an old version of keras.
+    return False
+
+
+_USE_KERAS_3 = detect_if_tensorflow_uses_keras_3()
+
+
+def keras_3():
+    """Check if Keras 3 is being used."""
+    return _USE_KERAS_3
 
-if not os.path.exists(_config_path):
-    _config = {
-        "multi_backend": _MULTI_BACKEND,
-    }
-    try:
-        with open(_config_path, "w") as f:
-            f.write(json.dumps(_config, indent=4))
-    except IOError:
-        # Except permission denied.
-        pass
 
-if "KERAS_BACKEND" in os.environ and os.environ["KERAS_BACKEND"]:
-    _MULTI_BACKEND = True
+def backend():
+    """Check the backend framework."""
+    if not keras_3():
+        return "tensorflow"
 
+    import keras
 
-def multi_backend():
-    return _MULTI_BACKEND
+    return keras.config.backend()
diff --git a/keras_cv/backend/keras.py b/keras_cv/backend/keras.py
new file mode 100644
index 0000000000..340efad3b1
--- /dev/null
+++ b/keras_cv/backend/keras.py
@@ -0,0 +1,65 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import types
+
+from keras_cv.backend import config
+
+_KERAS_CORE_ALIASES = {
+    "utils->saving": [
+        "register_keras_serializable",
+        "deserialize_keras_object",
+        "serialize_keras_object",
+        "get_registered_object",
+    ],
+    "models->saving": ["load_model"],
+}
+
+if config.keras_3():
+    import keras  # noqa: F403, F401
+    from keras import *  # noqa: F403, F401
+
+    keras.backend.name_scope = keras.name_scope
+else:
+    from tensorflow import keras  # noqa: F403, F401
+    from tensorflow.keras import *  # noqa: F403, F401
+
+    if not hasattr(keras, "saving"):
+        keras.saving = types.SimpleNamespace()
+
+    # add aliases
+    for key, value in _KERAS_CORE_ALIASES.items():
+        src, _, dst = key.partition("->")
+        src = src.split(".")
+        dst = dst.split(".")
+
+        src_mod, dst_mod = keras, keras
+
+        # navigate to where we want to alias the attributes
+        for mod in src:
+            src_mod = getattr(src_mod, mod)
+        for mod in dst:
+            dst_mod = getattr(dst_mod, mod)
+
+        # add an alias for each attribute
+        for attr in value:
+            if isinstance(attr, tuple):
+                src_attr, dst_attr = attr
+            else:
+                src_attr, dst_attr = attr, attr
+            attr_val = getattr(src_mod, src_attr)
+            setattr(dst_mod, dst_attr, attr_val)
+
+    # TF Keras doesn't have this rename.
+    keras.activations.silu = keras.activations.swish
diff --git a/keras_cv/backend/ops.py b/keras_cv/backend/ops.py
index 013fa3f124..c02dadfd6b 100644
--- a/keras_cv/backend/ops.py
+++ b/keras_cv/backend/ops.py
@@ -11,13 +11,25 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend import config
 
-if multi_backend():
-    from keras_core.src.backend import vectorized_map  # noqa: F403, F401
-    from keras_core.src.ops import *  # noqa: F403, F401
-    from keras_core.src.utils.image_utils import (  # noqa: F403, F401
-        smart_resize,
-    )
+if config.keras_3():
+    from keras.ops import *  # noqa: F403, F401
+    from keras.preprocessing.image import smart_resize  # noqa: F403, F401
+
+    from keras_cv.backend import keras
+
+    name_scope = keras.name_scope
 else:
-    from keras_cv.backend.tf_ops import *  # noqa: F403, F401
+    try:
+        from keras.src.ops import *  # noqa: F403, F401
+        from keras.src.utils.image_utils import smart_resize  # noqa: F403, F401
+    # Import error means Keras isn't installed, or is Keras 2.
+    except ImportError:
+        from keras_core.src.backend import vectorized_map  # noqa: F403, F401
+        from keras_core.src.ops import *  # noqa: F403, F401
+        from keras_core.src.utils.image_utils import (  # noqa: F403, F401
+            smart_resize,
+        )
+    if config.backend() == "tensorflow":
+        from keras_cv.backend.tf_ops import *  # noqa: F403, F401
diff --git a/keras_cv/backend/random.py b/keras_cv/backend/random.py
index 21d4b08c7d..049312a912 100644
--- a/keras_cv/backend/random.py
+++ b/keras_cv/backend/random.py
@@ -11,10 +11,136 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import random as python_random
 
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend import keras
+from keras_cv.backend.config import keras_3
 
-if multi_backend():
-    from keras_core.random import *  # noqa: F403, F401
+if keras_3():
+    from keras.random import *  # noqa: F403, F401
 else:
-    from keras_core.src.backend.tensorflow.random import *  # noqa: F403, F401
+    from keras_core.random import *  # noqa: F403, F401
+
+
+def _make_default_seed():
+    return python_random.randint(1, int(1e9))
+
+
+class SeedGenerator:
+    def __new__(cls, seed=None, **kwargs):
+        if keras_3():
+            return keras.random.SeedGenerator(seed=seed, **kwargs)
+        return super().__new__(cls)
+
+    def __init__(self, seed=None):
+        if seed is None:
+            seed = _make_default_seed()
+        self._initial_seed = seed
+        self._current_seed = [0, seed]
+
+    def next(self, ordered=True):
+        self._current_seed[0] += 1
+        return self._current_seed[:]
+
+    def get_config(self):
+        return {"seed": self._initial_seed}
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
+
+
+def _draw_seed(seed):
+    if keras_3():
+        # Keras 3 seed can be directly passed to random functions
+        return seed
+    if isinstance(seed, SeedGenerator):
+        init_seed = seed.next()
+    else:
+        if seed is None:
+            seed = _make_default_seed()
+        init_seed = [0, seed]
+    return init_seed
+
+
+def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
+    seed = _draw_seed(seed)
+    kwargs = {}
+    if dtype:
+        kwargs["dtype"] = dtype
+    if keras_3():
+        return keras.random.normal(
+            shape,
+            mean=mean,
+            stddev=stddev,
+            seed=seed,
+            **kwargs,
+        )
+    else:
+        import tensorflow as tf
+
+        return tf.random.stateless_normal(
+            shape,
+            mean=mean,
+            stddev=stddev,
+            seed=seed,
+            **kwargs,
+        )
+
+
+def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
+    init_seed = _draw_seed(seed)
+    kwargs = {}
+    if dtype:
+        kwargs["dtype"] = dtype
+    if keras_3():
+        return keras.random.uniform(
+            shape,
+            minval=minval,
+            maxval=maxval,
+            seed=init_seed,
+            **kwargs,
+        )
+    else:
+        import tensorflow as tf
+
+        return tf.random.stateless_uniform(
+            shape,
+            minval=minval,
+            maxval=maxval,
+            seed=init_seed,
+            **kwargs,
+        )
+
+
+def shuffle(x, axis=0, seed=None):
+    init_seed = _draw_seed(seed)
+    if keras_3():
+        return keras.random.shuffle(x=x, axis=axis, seed=init_seed)
+    else:
+        import tensorflow as tf
+
+        return tf.random.stateless_shuffle(x=x, axis=axis, seed=init_seed)
+
+
+def categorical(logits, num_samples, dtype=None, seed=None):
+    init_seed = _draw_seed(seed)
+    kwargs = {}
+    if dtype:
+        kwargs["dtype"] = dtype
+    if keras_3():
+        return keras.random.categorical(
+            logits=logits,
+            num_samples=num_samples,
+            seed=init_seed,
+            **kwargs,
+        )
+    else:
+        import tensorflow as tf
+
+        return tf.random.stateless_categorical(
+            logits=logits,
+            num_samples=num_samples,
+            seed=init_seed,
+            **kwargs,
+        )
diff --git a/keras_cv/backend/scope.py b/keras_cv/backend/scope.py
index d6bd1b40d2..b2eb094147 100644
--- a/keras_cv/backend/scope.py
+++ b/keras_cv/backend/scope.py
@@ -18,7 +18,7 @@
 from keras_cv.backend import keras
 from keras_cv.backend import ops
 from keras_cv.backend import tf_ops
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 _ORIGINAL_OPS = copy.copy(backend.ops.__dict__)
 _ORIGINAL_SUPPORTS_RAGGED = backend.supports_ragged
@@ -30,7 +30,7 @@
 def tf_data(function):
     @functools.wraps(function)
     def wrapper(*args, **kwargs):
-        if multi_backend() and keras.src.utils.backend_utils.in_tf_graph():
+        if keras_3() and keras.src.utils.backend_utils.in_tf_graph():
             with TFDataScope():
                 return function(*args, **kwargs)
         else:
diff --git a/keras_cv/backend/tf_ops.py b/keras_cv/backend/tf_ops.py
index dc93712095..1eb430d922 100644
--- a/keras_cv/backend/tf_ops.py
+++ b/keras_cv/backend/tf_ops.py
@@ -11,23 +11,103 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from keras_core.src.backend.tensorflow import *  # noqa: F403, F401
-from keras_core.src.backend.tensorflow import (  # noqa: F403, F401
-    convert_to_numpy,
-)
-from keras_core.src.backend.tensorflow.core import *  # noqa: F403, F401
-from keras_core.src.backend.tensorflow.math import *  # noqa: F403, F401
-from keras_core.src.backend.tensorflow.nn import *  # noqa: F403, F401
-from keras_core.src.backend.tensorflow.numpy import *  # noqa: F403, F401
+from keras_cv.backend import config
+
+if config.keras_3():
+    from keras.src.backend.tensorflow import *  # noqa: F403, F401
+    from keras.src.backend.tensorflow import (  # noqa: F403, F401
+        convert_to_numpy,
+    )
+    from keras.src.backend.tensorflow.core import *  # noqa: F403, F401
+    from keras.src.backend.tensorflow.math import *  # noqa: F403, F401
+    from keras.src.backend.tensorflow.nn import *  # noqa: F403, F401
+    from keras.src.backend.tensorflow.numpy import *  # noqa: F403, F401
+else:
+    # isort: off
+    from keras_core.src.backend.tensorflow import *  # noqa: F403, F401
+    from keras_core.src.backend.tensorflow import (  # noqa: F403, F401
+        convert_to_numpy,
+    )
+    from keras_core.src.backend.tensorflow.core import *  # noqa: F403, F401
+    from keras_core.src.backend.tensorflow.math import *  # noqa: F403, F401
+    from keras_core.src.backend.tensorflow.nn import *  # noqa: F403, F401
+    from keras_core.src.backend.tensorflow.numpy import *  # noqa: F403, F401, E501
+
 
 # Some TF APIs where the numpy API doesn't support raggeds that we need
+from tensorflow import broadcast_to  # noqa: F403, F401
 from tensorflow import concat as concatenate  # noqa: F403, F401
+from tensorflow import repeat  # noqa: F403, F401
+from tensorflow import reshape  # noqa: F403, F401
+
 from tensorflow import range as arange  # noqa: F403, F401
 from tensorflow import reduce_all as all  # noqa: F403, F401
 from tensorflow import reduce_max as max  # noqa: F403, F401
-from tensorflow import repeat  # noqa: F403, F401
-from tensorflow import reshape  # noqa: F403, F401
 from tensorflow import split  # noqa: F403, F401
-from tensorflow.keras.preprocessing.image import (  # noqa: F403, F401
-    smart_resize,
-)
+
+import numpy as np
+import tensorflow as tf
+
+
+def smart_resize(x, size, interpolation="bilinear"):
+    """Resize images to a target size without aspect ratio distortion.
+
+    Copied from `tf_keras` for Keras 3 and for use in `tf.data` pipeline.
+    """
+    if len(size) != 2:
+        raise ValueError(
+            f"Expected `size` to be a tuple of 2 integers, but got: {size}."
+        )
+    img = tf.convert_to_tensor(x)
+    if img.shape.rank is not None:
+        if img.shape.rank < 3 or img.shape.rank > 4:
+            raise ValueError(
+                "Expected an image array with shape `(height, width, "
+                "channels)`, or `(batch_size, height, width, channels)`, but "
+                f"got input with incorrect rank, of shape {img.shape}."
+            )
+    shape = tf.shape(img)
+    height, width = shape[-3], shape[-2]
+    target_height, target_width = size
+    if img.shape.rank is not None:
+        static_num_channels = img.shape[-1]
+    else:
+        static_num_channels = None
+
+    crop_height = tf.cast(
+        tf.cast(width * target_height, "float32") / target_width, "int32"
+    )
+    crop_width = tf.cast(
+        tf.cast(height * target_width, "float32") / target_height, "int32"
+    )
+
+    # Set back to input height / width if crop_height / crop_width is not
+    # smaller.
+    crop_height = tf.minimum(height, crop_height)
+    crop_width = tf.minimum(width, crop_width)
+
+    crop_box_hstart = tf.cast(
+        tf.cast(height - crop_height, "float32") / 2, "int32"
+    )
+    crop_box_wstart = tf.cast(
+        tf.cast(width - crop_width, "float32") / 2, "int32"
+    )
+
+    if img.shape.rank == 4:
+        crop_box_start = tf.stack([0, crop_box_hstart, crop_box_wstart, 0])
+        crop_box_size = tf.stack([-1, crop_height, crop_width, -1])
+    else:
+        crop_box_start = tf.stack([crop_box_hstart, crop_box_wstart, 0])
+        crop_box_size = tf.stack([crop_height, crop_width, -1])
+
+    img = tf.slice(img, crop_box_start, crop_box_size)
+    img = tf.image.resize(images=img, size=size, method=interpolation)
+    # Apparent bug in resize_images_v2 may cause shape to be lost
+    if img.shape.rank is not None:
+        if img.shape.rank == 4:
+            img.set_shape((None, None, None, static_num_channels))
+        if img.shape.rank == 3:
+            img.set_shape((None, None, static_num_channels))
+    if isinstance(x, np.ndarray):
+        return img.numpy()
+    return img
diff --git a/keras_cv/bounding_box/converters.py b/keras_cv/bounding_box/converters.py
index 8923527cbf..b671b6ff22 100644
--- a/keras_cv/bounding_box/converters.py
+++ b/keras_cv/bounding_box/converters.py
@@ -370,7 +370,8 @@ def convert_format(
             `"float32"`.
     """
     if isinstance(boxes, dict):
-        boxes["boxes"] = convert_format(
+        converted_boxes = boxes.copy()
+        converted_boxes["boxes"] = convert_format(
             boxes["boxes"],
             source=source,
             target=target,
@@ -378,7 +379,7 @@ def convert_format(
             image_shape=image_shape,
             dtype=dtype,
         )
-        return boxes
+        return converted_boxes
 
     if boxes.shape[-1] is not None and boxes.shape[-1] != 4:
         raise ValueError(
diff --git a/keras_cv/bounding_box/mask_invalid_detections.py b/keras_cv/bounding_box/mask_invalid_detections.py
index baf2fced6c..464b3fa247 100644
--- a/keras_cv/bounding_box/mask_invalid_detections.py
+++ b/keras_cv/bounding_box/mask_invalid_detections.py
@@ -23,8 +23,8 @@
 def mask_invalid_detections(bounding_boxes, output_ragged=False):
     """masks out invalid detections with -1s.
 
-    This utility is mainly used on the output of non-max supression operations.
-    The output of non-max-supression contains all the detections, even invalid
+    This utility is mainly used on the output of non-max suppression operations.
+    The output of non-max-suppression contains all the detections, even invalid
     ones. Users are expected to use `num_detections` to determine how many boxes
     are in each image.
 
diff --git a/keras_cv/bounding_box/mask_invalid_detections_test.py b/keras_cv/bounding_box/mask_invalid_detections_test.py
index ed7562ea7c..c2d9b8e710 100644
--- a/keras_cv/bounding_box/mask_invalid_detections_test.py
+++ b/keras_cv/bounding_box/mask_invalid_detections_test.py
@@ -18,15 +18,16 @@
 
 from keras_cv import bounding_box
 from keras_cv.backend import ops
+from keras_cv.backend import random
 from keras_cv.tests.test_case import TestCase
 
 
 class MaskInvalidDetectionsTest(TestCase):
     def test_correctly_masks_based_on_max_dets(self):
         bounding_boxes = {
-            "boxes": ops.random.uniform((4, 100, 4)),
+            "boxes": random.uniform((4, 100, 4)),
             "num_detections": ops.array([2, 3, 4, 2]),
-            "classes": ops.random.uniform((4, 100)),
+            "classes": random.uniform((4, 100)),
         }
 
         result = bounding_box.mask_invalid_detections(bounding_boxes)
diff --git a/keras_cv/callbacks/pycoco_callback.py b/keras_cv/callbacks/pycoco_callback.py
index 0de03fe799..f0f33ea2a6 100644
--- a/keras_cv/callbacks/pycoco_callback.py
+++ b/keras_cv/callbacks/pycoco_callback.py
@@ -46,7 +46,6 @@ def __init__(
                 operations and passing `cache=False`.
         """
         assert_pycocotools_installed("PyCOCOCallback")
-        self.model = None
         self.val_data = validation_data
         if cache:
             # We cache the dataset to preserve a consistent iteration order.
diff --git a/keras_cv/callbacks/waymo_evaluation_callback.py b/keras_cv/callbacks/waymo_evaluation_callback.py
index cc73eb9ae8..d2be6241f9 100644
--- a/keras_cv/callbacks/waymo_evaluation_callback.py
+++ b/keras_cv/callbacks/waymo_evaluation_callback.py
@@ -46,7 +46,6 @@ def __init__(self, validation_data, config=None, **kwargs):
         assert_waymo_open_dataset_installed(
             "keras_cv.callbacks.WaymoEvaluationCallback()"
         )
-        self.model = None
         self.val_data = validation_data
         self.evaluator = WODDetectionEvaluator(
             config=config or self._get_default_config()
diff --git a/keras_cv/conftest.py b/keras_cv/conftest.py
index 8e886dfd99..eaee5024b9 100644
--- a/keras_cv/conftest.py
+++ b/keras_cv/conftest.py
@@ -17,7 +17,8 @@
 import tensorflow as tf
 from packaging import version
 
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend import config as backend_config
+from keras_cv.backend.config import keras_3
 
 
 def pytest_addoption(parser):
@@ -33,9 +34,35 @@ def pytest_addoption(parser):
         default=False,
         help="run extra_large tests",
     )
+    parser.addoption(
+        "--check_gpu",
+        action="store_true",
+        default=False,
+        help="fail if a gpu is not present",
+    )
 
 
 def pytest_configure(config):
+    # Verify that device has GPU and detected by backend
+    if config.getoption("--check_gpu"):
+        found_gpu = False
+        backend = backend_config.backend()
+        if backend == "jax":
+            import jax
+
+            try:
+                found_gpu = bool(jax.devices("gpu"))
+            except RuntimeError:
+                found_gpu = False
+        elif backend == "tensorflow":
+            found_gpu = bool(tf.config.list_logical_devices("GPU"))
+        elif backend == "torch":
+            import torch
+
+            found_gpu = bool(torch.cuda.device_count())
+        if not found_gpu:
+            pytest.fail(f"No GPUs discovered on the {backend} backend.")
+
     config.addinivalue_line(
         "markers", "large: mark test as being slow or requiring a network"
     )
@@ -45,7 +72,7 @@ def pytest_configure(config):
     )
     config.addinivalue_line(
         "markers",
-        "tf_keras_only: mark test as a tf.keras-only test",
+        "tf_keras_only: mark test as a Keras 2-only test",
     )
     config.addinivalue_line(
         "markers",
@@ -69,12 +96,12 @@ def pytest_collection_modifyitems(config, items):
     skip_extra_large = pytest.mark.skipif(
         not run_extra_large_tests, reason="need --run_extra_large option to run"
     )
-    skip_tf_keras_only = pytest.mark.skipif(
-        multi_backend(),
-        reason="This test is only supported on tf.keras",
+    skip_keras_2_only = pytest.mark.skipif(
+        keras_3(),
+        reason="This test is only supported on Keras 2",
     )
     skip_tf_only = pytest.mark.skipif(
-        multi_backend() and keras_core.backend.backend() != "tensorflow",
+        keras_3() and keras_core.backend.backend() != "tensorflow",
         reason="This test is only supported on TensorFlow",
     )
     for item in items:
@@ -87,6 +114,6 @@ def pytest_collection_modifyitems(config, items):
         if "extra_large" in item.keywords:
             item.add_marker(skip_extra_large)
         if "tf_keras_only" in item.keywords:
-            item.add_marker(skip_tf_keras_only)
+            item.add_marker(skip_keras_2_only)
         if "tf_only" in item.keywords:
             item.add_marker(skip_tf_only)
diff --git a/keras_cv/datasets/imagenet/load.py b/keras_cv/datasets/imagenet/load.py
index 55438a474f..547bb51ab6 100644
--- a/keras_cv/datasets/imagenet/load.py
+++ b/keras_cv/datasets/imagenet/load.py
@@ -56,7 +56,9 @@ def apply(example):
     return apply
 
 
-@keras_cv_export("keras_cv.datasets.imagenet.load")
+@keras_cv_export(
+    "keras_cv.datasets.imagenet.load", package="keras_cv.datasets.imagenet"
+)
 def load(
     split,
     tfrecord_path,
diff --git a/keras_cv/datasets/load_test.py b/keras_cv/datasets/load_test.py
new file mode 100644
index 0000000000..bb12172b3b
--- /dev/null
+++ b/keras_cv/datasets/load_test.py
@@ -0,0 +1,24 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests ability to load keras_cv.datasets."""
+
+from keras_cv.tests.test_case import TestCase
+
+
+class LoadTest(TestCase):
+    def test_datasets_load(self):
+        import keras_cv.datasets  # noqa: F401
+
+    def test_imagenet_load(self):
+        import keras_cv.datasets.imagenet  # noqa: F401
diff --git a/keras_cv/datasets/pascal_voc/segmentation.py b/keras_cv/datasets/pascal_voc/segmentation.py
index 5e4ad1b93b..7e1acda281 100644
--- a/keras_cv/datasets/pascal_voc/segmentation.py
+++ b/keras_cv/datasets/pascal_voc/segmentation.py
@@ -58,7 +58,7 @@ class and instance segmentation masks.
     booktitle = "International Conference on Computer Vision (ICCV)",
     year = "2011"}}
 """  # noqa: E501
-SBD_URL = "https://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/semantic_contours/benchmark.tgz"  # noqa: E501
+SBD_URL = "https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/semantic_contours/benchmark.tgz"  # noqa: E501
 
 
 # Note that this list doesn't contain the background class. In the
@@ -465,7 +465,10 @@ def md_gen():
     return dataset
 
 
-@keras_cv_export("keras_cv.datasets.pascal_voc.segmentation.load")
+@keras_cv_export(
+    "keras_cv.datasets.pascal_voc.segmentation.load",
+    package="keras_cv.datasets.pascal_voc_segmentation",
+)
 def load(
     split="sbd_train",
     data_dir=None,
diff --git a/keras_cv/keypoint/converters.py b/keras_cv/keypoint/converters.py
index 4d4b290b63..453eafabc7 100644
--- a/keras_cv/keypoint/converters.py
+++ b/keras_cv/keypoint/converters.py
@@ -60,7 +60,9 @@ def _xy_noop(keypoints, images=None):
 }
 
 
-@keras_cv_export("keras_cv.keypoint.convert_format")
+@keras_cv_export(
+    "keras_cv.keypoint.convert_format", package="keras_cv.keypoint"
+)
 def convert_format(keypoints, source, target, images=None, dtype=None):
     """Converts keypoints from one format to another.
 
diff --git a/keras_cv/layers/__init__.py b/keras_cv/layers/__init__.py
index 342a942f64..0bfa2aa8ec 100644
--- a/keras_cv/layers/__init__.py
+++ b/keras_cv/layers/__init__.py
@@ -135,4 +135,9 @@
 )
 from keras_cv.layers.spatial_pyramid import SpatialPyramidPooling
 from keras_cv.layers.transformer_encoder import TransformerEncoder
+from keras_cv.layers.vit_det_layers import AddRelativePositionalEmbedding
+from keras_cv.layers.vit_det_layers import MultiHeadAttentionWithRelativePE
+from keras_cv.layers.vit_det_layers import ViTDetPatchingAndEmbedding
+from keras_cv.layers.vit_det_layers import WindowedTransformerEncoder
+from keras_cv.layers.vit_det_layers import WindowPartitioning
 from keras_cv.layers.vit_layers import PatchingAndEmbedding
diff --git a/keras_cv/layers/augmenter_test.py b/keras_cv/layers/augmenter_test.py
index a1886996f7..1afa7fffd9 100644
--- a/keras_cv/layers/augmenter_test.py
+++ b/keras_cv/layers/augmenter_test.py
@@ -37,7 +37,8 @@ def test_call_with_labels(self):
         augmenter = layers.Augmenter(
             [
                 layers.RandomFlip(),
-                layers.RandAugment(value_range=(0, 255)),
+                # TODO: RandAugment fails in Keras 3
+                # layers.RandAugment(value_range=(0, 255)),
                 layers.CutMix(),
             ]
         )
diff --git a/keras_cv/layers/feature_pyramid_test.py b/keras_cv/layers/feature_pyramid_test.py
index 8566db691d..b3118d8e2f 100644
--- a/keras_cv/layers/feature_pyramid_test.py
+++ b/keras_cv/layers/feature_pyramid_test.py
@@ -122,6 +122,9 @@ def test_invalid_input_features(self):
         c3 = np.ones([2, 32, 32, 3])
         c4 = np.ones([2, 16, 16, 3])
         c5 = np.ones([2, 8, 8, 3])
+        inputs = {2: c2, 3: c3, 4: c4, 5: c5}
+        # Build required for Keas 3
+        _ = layer(inputs)
         list_input = [c2, c3, c4, c5]
         with self.assertRaisesRegexp(
             ValueError, "expects input features to be a dict"
diff --git a/keras_cv/layers/object_detection/anchor_generator.py b/keras_cv/layers/object_detection/anchor_generator.py
index 30dd421afd..effc125143 100644
--- a/keras_cv/layers/object_detection/anchor_generator.py
+++ b/keras_cv/layers/object_detection/anchor_generator.py
@@ -172,7 +172,7 @@ def __call__(self, image=None, image_shape=None):
                     "Expected `image` to be a Tensor of rank 3. Got "
                     f"image.shape.rank={len(image.shape)}"
                 )
-            image_shape = image.shape
+            image_shape = tuple(image.shape)
 
         results = {}
         for key, generator in self.anchor_generators.items():
diff --git a/keras_cv/layers/object_detection/multi_class_non_max_suppression.py b/keras_cv/layers/object_detection/multi_class_non_max_suppression.py
index 7825268578..08bcd97528 100644
--- a/keras_cv/layers/object_detection/multi_class_non_max_suppression.py
+++ b/keras_cv/layers/object_detection/multi_class_non_max_suppression.py
@@ -18,7 +18,7 @@
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 
 @keras_cv_export("keras_cv.layers.MultiClassNonMaxSuppression")
@@ -73,7 +73,7 @@ def call(
                 `bounding_box_format` specified in the constructor.
             class_prediction: Dense Tensor of shape [batch, boxes, num_classes].
         """
-        if multi_backend() and keras.backend.backend() != "tensorflow":
+        if keras_3() and keras.backend.backend() != "tensorflow":
             raise NotImplementedError(
                 "MultiClassNonMaxSuppression does not support non-TensorFlow "
                 "backends. Consider using NonMaxSuppression instead."
diff --git a/keras_cv/layers/object_detection/non_max_suppression.py b/keras_cv/layers/object_detection/non_max_suppression.py
index ed988cf113..2c39fd0d12 100644
--- a/keras_cv/layers/object_detection/non_max_suppression.py
+++ b/keras_cv/layers/object_detection/non_max_suppression.py
@@ -20,7 +20,7 @@
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 EPSILON = 1e-8
 
@@ -89,7 +89,7 @@ def call(
 
         confidence_prediction = ops.max(class_prediction, axis=-1)
 
-        if not multi_backend() or keras.backend.backend() == "tensorflow":
+        if not keras_3() or keras.backend.backend() == "tensorflow":
             idx, valid_det = tf.image.non_max_suppression_padded(
                 box_prediction,
                 confidence_prediction,
@@ -318,7 +318,7 @@ def suppression_loop_body(boxes, iou_threshold, output_size, idx):
 
     # TODO(ianstenbit): Fix bug in tfnp.take_along_axis that causes this hack.
     # (This will be removed anyway when we use built-in NMS for TF.)
-    if multi_backend() and keras.backend.backend() != "tensorflow":
+    if keras_3() and keras.backend.backend() != "tensorflow":
         idx = ops.take_along_axis(
             ops.reshape(sorted_indices, [-1]), take_along_axis_idx
         )
diff --git a/keras_cv/layers/preprocessing/aug_mix.py b/keras_cv/layers/preprocessing/aug_mix.py
index b06cf5b1dd..13d7adf0d9 100644
--- a/keras_cv/layers/preprocessing/aug_mix.py
+++ b/keras_cv/layers/preprocessing/aug_mix.py
@@ -106,7 +106,6 @@ def _sample_from_dirichlet(self, alpha):
         gamma_sample = tf.random.gamma(
             shape=(),
             alpha=alpha,
-            seed=self._random_generator.make_legacy_seed(),
         )
         return gamma_sample / tf.reduce_sum(
             gamma_sample, axis=-1, keepdims=True
@@ -114,15 +113,17 @@ def _sample_from_dirichlet(self, alpha):
 
     def _sample_from_beta(self, alpha, beta):
         sample_alpha = tf.random.gamma(
-            (), alpha=alpha, seed=self._random_generator.make_legacy_seed()
+            (),
+            alpha=alpha,
         )
         sample_beta = tf.random.gamma(
-            (), alpha=beta, seed=self._random_generator.make_legacy_seed()
+            (),
+            alpha=beta,
         )
         return sample_alpha / (sample_alpha + sample_beta)
 
     def _sample_depth(self):
-        return self._random_generator.random_uniform(
+        return self._random_generator.uniform(
             shape=(),
             minval=self.chain_depth[0],
             maxval=self.chain_depth[1] + 1,
@@ -130,7 +131,7 @@ def _sample_depth(self):
         )
 
     def _loop_on_depth(self, depth_level, image_aug):
-        op_index = self._random_generator.random_uniform(
+        op_index = self._random_generator.uniform(
             shape=(), minval=0, maxval=8, dtype=tf.int32
         )
         image_aug = self._apply_op(image_aug, op_index)
diff --git a/keras_cv/layers/preprocessing/aug_mix_test.py b/keras_cv/layers/preprocessing/aug_mix_test.py
index a6b3a6d4ab..2adba4eab5 100644
--- a/keras_cv/layers/preprocessing/aug_mix_test.py
+++ b/keras_cv/layers/preprocessing/aug_mix_test.py
@@ -27,16 +27,16 @@ def test_return_shapes(self):
         xs = layer(xs)
         ys_segmentation_masks = tf.ones((2, 512, 512, 3))
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
 
         # greyscale
         xs = tf.ones((2, 512, 512, 1))
         xs = layer(xs)
         ys_segmentation_masks = tf.ones((2, 512, 512, 1))
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [2, 512, 512, 1])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 1])
+        self.assertEqual(xs.shape, (2, 512, 512, 1))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 1))
 
     def test_in_single_image_and_mask(self):
         layer = preprocessing.AugMix([0, 255])
@@ -54,8 +54,8 @@ def test_in_single_image_and_mask(self):
         )
 
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [512, 512, 3])
-        self.assertEqual(ys_segmentation_masks.shape, [512, 512, 3])
+        self.assertEqual(xs.shape, (512, 512, 3))
+        self.assertEqual(ys_segmentation_masks.shape, (512, 512, 3))
 
         # greyscale
         xs = tf.cast(
@@ -69,8 +69,8 @@ def test_in_single_image_and_mask(self):
             dtype=tf.float32,
         )
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [512, 512, 1])
-        self.assertEqual(ys_segmentation_masks.shape, [512, 512, 1])
+        self.assertEqual(xs.shape, (512, 512, 1))
+        self.assertEqual(ys_segmentation_masks.shape, (512, 512, 1))
 
     def test_non_square_images_and_masks(self):
         layer = preprocessing.AugMix([0, 255])
@@ -80,16 +80,16 @@ def test_non_square_images_and_masks(self):
         xs = layer(xs)
         ys_segmentation_masks = tf.ones((2, 256, 512, 3))
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [2, 256, 512, 3])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 256, 512, 3])
+        self.assertEqual(xs.shape, (2, 256, 512, 3))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 256, 512, 3))
 
         # greyscale
         xs = tf.ones((2, 256, 512, 1))
         xs = layer(xs)
         ys_segmentation_masks = tf.ones((2, 256, 512, 1))
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [2, 256, 512, 1])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 256, 512, 1])
+        self.assertEqual(xs.shape, (2, 256, 512, 1))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 256, 512, 1))
 
     def test_single_input_args(self):
         layer = preprocessing.AugMix([0, 255])
@@ -99,16 +99,16 @@ def test_single_input_args(self):
         xs = layer(xs)
         ys_segmentation_masks = tf.ones((2, 512, 512, 3))
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
 
         # greyscale
         xs = tf.ones((2, 512, 512, 1))
         xs = layer(xs)
         ys_segmentation_masks = tf.ones((2, 512, 512, 1))
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [2, 512, 512, 1])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 1])
+        self.assertEqual(xs.shape, (2, 512, 512, 1))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 1))
 
     def test_many_augmentations(self):
         layer = preprocessing.AugMix([0, 255], chain_depth=[25, 26])
@@ -118,13 +118,13 @@ def test_many_augmentations(self):
         xs = layer(xs)
         ys_segmentation_masks = tf.ones((2, 512, 512, 3))
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
 
         # greyscale
         xs = tf.ones((2, 512, 512, 1))
         xs = layer(xs)
         ys_segmentation_masks = tf.ones((2, 512, 512, 1))
         ys_segmentation_masks = layer(ys_segmentation_masks)
-        self.assertEqual(xs.shape, [2, 512, 512, 1])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 1])
+        self.assertEqual(xs.shape, (2, 512, 512, 1))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 1))
diff --git a/keras_cv/layers/preprocessing/auto_contrast_test.py b/keras_cv/layers/preprocessing/auto_contrast_test.py
index 9b87c2f100..bc34d5fae9 100644
--- a/keras_cv/layers/preprocessing/auto_contrast_test.py
+++ b/keras_cv/layers/preprocessing/auto_contrast_test.py
@@ -12,53 +12,53 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 
-import tensorflow as tf
-
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.tests.test_case import TestCase
 
 
 class AutoContrastTest(TestCase):
     def test_constant_channels_dont_get_nanned(self):
-        img = tf.constant([1, 1], dtype=tf.float32)
-        img = tf.expand_dims(img, axis=-1)
-        img = tf.expand_dims(img, axis=-1)
-        img = tf.expand_dims(img, axis=0)
+        img = np.array([1, 1], dtype=np.float32)
+        img = np.expand_dims(img, axis=-1)
+        img = np.expand_dims(img, axis=-1)
+        img = np.expand_dims(img, axis=0)
 
         layer = preprocessing.AutoContrast(value_range=(0, 255))
         ys = layer(img)
 
-        self.assertTrue(tf.math.reduce_any(ys[0] == 1.0))
-        self.assertTrue(tf.math.reduce_any(ys[0] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 1.0))
 
     def test_auto_contrast_expands_value_range(self):
-        img = tf.constant([0, 128], dtype=tf.float32)
-        img = tf.expand_dims(img, axis=-1)
-        img = tf.expand_dims(img, axis=-1)
-        img = tf.expand_dims(img, axis=0)
+        img = np.array([0, 128], dtype=np.float32)
+        img = np.expand_dims(img, axis=-1)
+        img = np.expand_dims(img, axis=-1)
+        img = np.expand_dims(img, axis=0)
 
         layer = preprocessing.AutoContrast(value_range=(0, 255))
         ys = layer(img)
 
-        self.assertTrue(tf.math.reduce_any(ys[0] == 0.0))
-        self.assertTrue(tf.math.reduce_any(ys[0] == 255.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 0.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 255.0))
 
     def test_auto_contrast_different_values_per_channel(self):
-        img = tf.constant(
+        img = np.array(
             [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]],
-            dtype=tf.float32,
+            dtype=np.float32,
         )
-        img = tf.expand_dims(img, axis=0)
+        img = np.expand_dims(img, axis=0)
 
         layer = preprocessing.AutoContrast(value_range=(0, 255))
         ys = layer(img)
 
-        self.assertTrue(tf.math.reduce_any(ys[0, ..., 0] == 0.0))
-        self.assertTrue(tf.math.reduce_any(ys[0, ..., 1] == 0.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 0]) == 0.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 1]) == 0.0))
 
-        self.assertTrue(tf.math.reduce_any(ys[0, ..., 0] == 255.0))
-        self.assertTrue(tf.math.reduce_any(ys[0, ..., 1] == 255.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 0]) == 255.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0, ..., 1]) == 255.0))
 
         self.assertAllClose(
             ys,
@@ -71,25 +71,25 @@ def test_auto_contrast_different_values_per_channel(self):
         )
 
     def test_auto_contrast_expands_value_range_uint8(self):
-        img = tf.constant([0, 128], dtype=tf.uint8)
-        img = tf.expand_dims(img, axis=-1)
-        img = tf.expand_dims(img, axis=-1)
-        img = tf.expand_dims(img, axis=0)
+        img = np.array([0, 128], dtype=np.uint8)
+        img = np.expand_dims(img, axis=-1)
+        img = np.expand_dims(img, axis=-1)
+        img = np.expand_dims(img, axis=0)
 
         layer = preprocessing.AutoContrast(value_range=(0, 255))
         ys = layer(img)
 
-        self.assertTrue(tf.math.reduce_any(ys[0] == 0.0))
-        self.assertTrue(tf.math.reduce_any(ys[0] == 255.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 0.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 255.0))
 
     def test_auto_contrast_properly_converts_value_range(self):
-        img = tf.constant([0, 0.5], dtype=tf.float32)
-        img = tf.expand_dims(img, axis=-1)
-        img = tf.expand_dims(img, axis=-1)
-        img = tf.expand_dims(img, axis=0)
+        img = np.array([0, 0.5], dtype=np.float32)
+        img = np.expand_dims(img, axis=-1)
+        img = np.expand_dims(img, axis=-1)
+        img = np.expand_dims(img, axis=0)
 
         layer = preprocessing.AutoContrast(value_range=(0, 1))
         ys = layer(img)
 
-        self.assertTrue(tf.math.reduce_any(ys[0] == 0.0))
-        self.assertTrue(tf.math.reduce_any(ys[0] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 0.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(ys[0]) == 1.0))
diff --git a/keras_cv/layers/preprocessing/base_image_augmentation_layer.py b/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
index 322dc5d92a..167da7ad0b 100644
--- a/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
+++ b/keras_cv/layers/preprocessing/base_image_augmentation_layer.py
@@ -14,6 +14,7 @@
 
 import keras
 import tensorflow as tf
+import tree
 
 if hasattr(keras, "src"):
     keras_backend = keras.src.backend
@@ -22,9 +23,10 @@
 
 from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import config
 from keras_cv.backend import keras
+from keras_cv.backend import ops
 from keras_cv.backend import scope
-from keras_cv.backend.config import multi_backend
 from keras_cv.utils import preprocessing
 
 # In order to support both unbatched and batched inputs, the horizontal
@@ -42,15 +44,8 @@
 USE_TARGETS = "use_targets"
 
 
-base_class = (
-    keras.src.layers.preprocessing.tf_data_layer.TFDataLayer
-    if multi_backend()
-    else keras.layers.Layer
-)
-
-
 @keras_cv_export("keras_cv.layers.BaseImageAugmentationLayer")
-class BaseImageAugmentationLayer(base_class):
+class BaseImageAugmentationLayer(keras.layers.Layer):
     """Abstract base layer for image augmentation.
 
     This layer contains base functionalities for preprocessing layers which
@@ -130,11 +125,14 @@ def augment_image(self, image, transformation):
     """
 
     def __init__(self, seed=None, **kwargs):
-        force_generator = kwargs.pop("force_generator", False)
-        self._random_generator = keras_backend.RandomGenerator(
-            seed=seed, force_generator=force_generator
-        )
+        if seed is not None:
+            self._random_generator = tf.random.Generator.from_seed(seed=seed)
+        else:
+            self._random_generator = tf.random.get_global_generator()
         super().__init__(**kwargs)
+
+        self._allow_non_tensor_positional_args = True
+        self.built = True
         self._convert_input_args = False
 
     @property
@@ -238,15 +236,15 @@ def _compute_output_signature(self, inputs):
         bounding_boxes = inputs.get(BOUNDING_BOXES, None)
 
         if bounding_boxes is not None:
-            fn_output_signature[
-                BOUNDING_BOXES
-            ] = self._compute_bounding_box_signature(bounding_boxes)
+            fn_output_signature[BOUNDING_BOXES] = (
+                self._compute_bounding_box_signature(bounding_boxes)
+            )
 
         segmentation_masks = inputs.get(SEGMENTATION_MASKS, None)
         if segmentation_masks is not None:
-            fn_output_signature[
-                SEGMENTATION_MASKS
-            ] = self.compute_image_signature(segmentation_masks)
+            fn_output_signature[SEGMENTATION_MASKS] = (
+                self.compute_image_signature(segmentation_masks)
+            )
 
         keypoints = inputs.get(KEYPOINTS, None)
         if keypoints is not None:
@@ -412,6 +410,21 @@ def get_random_transformation(
         return None
 
     def call(self, inputs):
+        # try to convert a given backend native tensor to TensorFlow tensor
+        # before passing it over to TFDataScope
+        is_tf_backend = config.backend() == "tensorflow"
+        is_in_tf_graph = not tf.executing_eagerly()
+        contains_ragged = lambda y: any(
+            tree.map_structure(
+                lambda x: isinstance(x, (tf.RaggedTensor, tf.SparseTensor)),
+                tree.flatten(y),
+            )
+        )
+        inputs_contain_ragged = contains_ragged(inputs)
+        if not is_tf_backend and not inputs_contain_ragged:
+            inputs = tree.map_structure(
+                lambda x: tf.convert_to_tensor(x), inputs
+            )
         with scope.TFDataScope():
             inputs = self._ensure_inputs_are_compute_dtype(inputs)
             inputs, metadata = self._format_inputs(inputs)
@@ -428,7 +441,22 @@ def call(self, inputs):
                     "rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
                     f"{images.shape}"
                 )
-            return outputs
+        # convert the outputs to backend native tensors if none of them
+        # contain RaggedTensors. Note that if the user passed in Raggeds
+        # but the outputs are dense, we still don't want to convert to
+        # backend native tensors. This is to avoid breaking TF data
+        # pipelines that can't easily be ported to become backend
+        # agnostic.
+        # Skip this step for TF backend or if in `tf.graph` like `tf.data`.
+        if not is_tf_backend and not is_in_tf_graph:
+            if not inputs_contain_ragged and not contains_ragged(outputs):
+                outputs = tree.map_structure(
+                    # some layers return None, handle that case when
+                    # converting to tensors
+                    lambda x: ops.convert_to_tensor(x) if x is not None else x,
+                    outputs,
+                )
+        return outputs
 
     def _augment(self, inputs):
         raw_image = inputs.get(IMAGES, None)
@@ -568,6 +596,8 @@ def _ensure_inputs_are_compute_dtype(self, inputs):
                 inputs,
                 self.compute_dtype,
             )
+        # Copy the input dict before we mutate it.
+        inputs = dict(inputs)
         inputs[IMAGES] = preprocessing.ensure_tensor(
             inputs[IMAGES],
             self.compute_dtype,
diff --git a/keras_cv/layers/preprocessing/base_image_augmentation_layer_test.py b/keras_cv/layers/preprocessing/base_image_augmentation_layer_test.py
index e3fcfeddb5..edc484e694 100644
--- a/keras_cv/layers/preprocessing/base_image_augmentation_layer_test.py
+++ b/keras_cv/layers/preprocessing/base_image_augmentation_layer_test.py
@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
+import pytest
 import tensorflow as tf
 
 from keras_cv import bounding_box
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
 )
@@ -30,7 +32,7 @@ def __init__(self, value_range=(0.0, 1.0), fixed_value=None, **kwargs):
     def get_random_transformation(self, **kwargs):
         if self.fixed_value:
             return self.fixed_value
-        return self._random_generator.random_uniform(
+        return self._random_generator.uniform(
             [], minval=self.value_range[0], maxval=self.value_range[1]
         )
 
@@ -78,17 +80,17 @@ def test_augment_dict_return_type(self):
 
     def test_augment_casts_dtypes(self):
         add_layer = RandomAddLayer(fixed_value=2.0)
-        images = tf.ones((2, 8, 8, 3), dtype="uint8")
+        images = np.ones((2, 8, 8, 3), dtype="uint8")
         output = add_layer(images)
 
         self.assertAllClose(
-            tf.ones((2, 8, 8, 3), dtype="float32") * 3.0, output
+            np.ones((2, 8, 8, 3), dtype="float32") * 3.0, output
         )
 
     def test_augment_batch_images(self):
         add_layer = RandomAddLayer()
         images = np.random.random(size=(2, 8, 8, 3)).astype("float32")
-        output = add_layer(images)
+        output = ops.convert_to_numpy(add_layer(images))
 
         diff = output - images
         # Make sure the first image and second image get different augmentation
@@ -118,8 +120,8 @@ def test_augment_batch_images_and_targets(self):
         targets = np.random.random(size=(2, 1)).astype("float32")
         output = add_layer({"images": images, "targets": targets})
 
-        image_diff = output["images"] - images
-        label_diff = output["targets"] - targets
+        image_diff = ops.convert_to_numpy(output["images"]) - images
+        label_diff = ops.convert_to_numpy(output["targets"]) - targets
         # Make sure the first image and second image get different augmentation
         self.assertNotAllClose(image_diff[0], image_diff[1])
         self.assertNotAllClose(label_diff[0], label_diff[1])
@@ -225,6 +227,7 @@ def test_augment_batch_image_and_localization_data(self):
             segmentation_mask_diff[0], segmentation_mask_diff[1]
         )
 
+    @pytest.mark.tf_only
     def test_augment_all_data_in_tf_function(self):
         add_layer = RandomAddLayer()
         images = np.random.random(size=(2, 8, 8, 3)).astype("float32")
@@ -262,3 +265,11 @@ def in_tf_function(inputs):
         self.assertNotAllClose(
             segmentation_mask_diff[0], segmentation_mask_diff[1]
         )
+
+    def test_augment_tf_data_pipeline(self):
+        image = np.random.random(size=(1, 8, 8, 3)).astype("float32")
+        tf_dataset = tf.data.Dataset.from_tensor_slices(image).map(
+            RandomAddLayer(fixed_value=2.0)
+        )
+        output = iter(tf_dataset).get_next()
+        self.assertAllClose(image[0] + 2.0, output)
diff --git a/keras_cv/layers/preprocessing/channel_shuffle.py b/keras_cv/layers/preprocessing/channel_shuffle.py
index 90f2100d6a..68e2088df0 100644
--- a/keras_cv/layers/preprocessing/channel_shuffle.py
+++ b/keras_cv/layers/preprocessing/channel_shuffle.py
@@ -55,7 +55,7 @@ def get_random_transformation_batch(self, batch_size, **kwargs):
         #     [0, 2, 3, 4, 1],
         #     [4, 1, 0, 2, 3]
         # ]
-        indices_distribution = self._random_generator.random_uniform(
+        indices_distribution = self._random_generator.uniform(
             (batch_size, self.groups)
         )
         indices = tf.argsort(indices_distribution, axis=-1)
diff --git a/keras_cv/layers/preprocessing/channel_shuffle_test.py b/keras_cv/layers/preprocessing/channel_shuffle_test.py
index 8063c446c4..e608e5584d 100644
--- a/keras_cv/layers/preprocessing/channel_shuffle_test.py
+++ b/keras_cv/layers/preprocessing/channel_shuffle_test.py
@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
+import pytest
 import tensorflow as tf
 
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.channel_shuffle import ChannelShuffle
 from keras_cv.tests.test_case import TestCase
 
@@ -24,7 +26,7 @@ def test_return_shapes(self):
 
         layer = ChannelShuffle(groups=3)
         xs = layer(xs, training=True)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
 
     def test_channel_shuffle_call_results_one_channel(self):
         xs = tf.cast(
@@ -37,8 +39,8 @@ def test_channel_shuffle_call_results_one_channel(self):
 
         layer = ChannelShuffle(groups=1)
         xs = layer(xs, training=True)
-        self.assertTrue(tf.math.reduce_any(xs[0] == 3.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 3.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
 
     def test_channel_shuffle_call_results_multi_channel(self):
         xs = tf.cast(
@@ -51,8 +53,8 @@ def test_channel_shuffle_call_results_multi_channel(self):
 
         layer = ChannelShuffle(groups=5)
         xs = layer(xs, training=True)
-        self.assertTrue(tf.math.reduce_any(xs[0] == 3.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 3.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
 
     def test_non_square_image(self):
         xs = tf.cast(
@@ -65,9 +67,10 @@ def test_non_square_image(self):
 
         layer = ChannelShuffle(groups=1)
         xs = layer(xs, training=True)
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
@@ -83,8 +86,8 @@ def augment(x):
             return layer(x, training=True)
 
         xs = augment(xs)
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
     def test_in_single_image(self):
         xs = tf.cast(
@@ -94,8 +97,9 @@ def test_in_single_image(self):
 
         layer = ChannelShuffle(groups=1)
         xs = layer(xs, training=True)
-        self.assertTrue(tf.math.reduce_any(xs == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs) == 1.0))
 
+    @pytest.mark.skip(reason="flaky")
     def test_channel_shuffle_on_batched_images_independently(self):
         image = tf.random.uniform((100, 100, 3))
         batched_images = tf.stack((image, image), axis=0)
@@ -114,9 +118,11 @@ def test_config_with_custom_name(self):
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = ChannelShuffle(groups=1)
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+        )
         layer = ChannelShuffle(groups=1, dtype="uint8")
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
 
     def test_config(self):
         layer = ChannelShuffle(groups=5)
diff --git a/keras_cv/layers/preprocessing/cut_mix.py b/keras_cv/layers/preprocessing/cut_mix.py
index 44d03c6134..7b76c38128 100644
--- a/keras_cv/layers/preprocessing/cut_mix.py
+++ b/keras_cv/layers/preprocessing/cut_mix.py
@@ -49,10 +49,12 @@ def __init__(
 
     def _sample_from_beta(self, alpha, beta, shape):
         sample_alpha = tf.random.gamma(
-            shape, alpha=alpha, seed=self._random_generator.make_legacy_seed()
+            shape,
+            alpha=alpha,
         )
         sample_beta = tf.random.gamma(
-            shape, alpha=beta, seed=self._random_generator.make_legacy_seed()
+            shape,
+            alpha=beta,
         )
         return sample_alpha / (sample_alpha + sample_beta)
 
diff --git a/keras_cv/layers/preprocessing/cut_mix_test.py b/keras_cv/layers/preprocessing/cut_mix_test.py
index 09f398328d..4dfa3120e8 100644
--- a/keras_cv/layers/preprocessing/cut_mix_test.py
+++ b/keras_cv/layers/preprocessing/cut_mix_test.py
@@ -11,8 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import numpy as np
+import pytest
 import tensorflow as tf
 
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.cut_mix import CutMix
 from keras_cv.tests.test_case import TestCase
 
@@ -51,9 +54,9 @@ def test_return_shapes(self):
             outputs["segmentation_masks"],
         )
 
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys_labels.shape, [2, 10])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertEqual(ys_labels.shape, (2, 10))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
 
     def test_cut_mix_call_results_with_labels(self):
         xs = tf.cast(
@@ -70,10 +73,10 @@ def test_cut_mix_call_results_with_labels(self):
         xs, ys = outputs["images"], outputs["labels"]
 
         # At least some pixels should be replaced in the CutMix operation
-        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
         # No labels should still be close to their original values
         self.assertNotAllClose(ys, 1.0)
         self.assertNotAllClose(ys, 0.0)
@@ -93,10 +96,10 @@ def test_cut_mix_call_results_one_channel_with_labels(self):
         xs, ys = outputs["images"], outputs["labels"]
 
         # At least some pixels should be replaced in the CutMix operation
-        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
         # No labels should still be close to their original values
         self.assertNotAllClose(ys, 1.0)
         self.assertNotAllClose(ys, 0.0)
@@ -128,15 +131,23 @@ def test_cut_mix_call_results_with_dense_encoded_segmentation_masks(self):
         )
 
         # At least some pixels should be replaced in the images
-        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
         # At least some pixels should be replaced in the segmentation_masks
-        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[0] == 1.0))
-        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[1] == 1.0))
-        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[1] == 2.0))
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(ys_segmentation_masks[0]) == 1.0)
+        )
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(ys_segmentation_masks[0]) == 2.0)
+        )
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(ys_segmentation_masks[1]) == 1.0)
+        )
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(ys_segmentation_masks[1]) == 2.0)
+        )
 
     def test_cut_mix_call_results_with_one_hot_encoded_segmentation_masks(self):
         xs = tf.cast(
@@ -166,24 +177,33 @@ def test_cut_mix_call_results_with_one_hot_encoded_segmentation_masks(self):
         )
 
         # At least some pixels should be replaced in the images
-        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
         # At least some pixels should be replaced in the segmentation_masks
         self.assertTrue(
-            tf.math.reduce_any(ys_segmentation_masks[0][:, :, 2] == 1.0)
+            np.any(
+                ops.convert_to_numpy(ys_segmentation_masks[0][:, :, 2]) == 1.0
+            )
         )
         self.assertTrue(
-            tf.math.reduce_any(ys_segmentation_masks[0][:, :, 2] == 0.0)
+            np.any(
+                ops.convert_to_numpy(ys_segmentation_masks[0][:, :, 2]) == 0.0
+            )
         )
         self.assertTrue(
-            tf.math.reduce_any(ys_segmentation_masks[1][:, :, 1] == 1.0)
+            np.any(
+                ops.convert_to_numpy(ys_segmentation_masks[1][:, :, 1]) == 1.0
+            )
         )
         self.assertTrue(
-            tf.math.reduce_any(ys_segmentation_masks[1][:, :, 1] == 0.0)
+            np.any(
+                ops.convert_to_numpy(ys_segmentation_masks[1][:, :, 1]) == 0.0
+            )
         )
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
@@ -203,10 +223,10 @@ def augment(x, y):
         xs, ys = outputs["images"], outputs["labels"]
 
         # At least some pixels should be replaced in the CutMix operation
-        self.assertTrue(tf.math.reduce_any(xs[0] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
         # No labels should still be close to their original values
         self.assertNotAllClose(ys, 1.0)
         self.assertNotAllClose(ys, 0.0)
diff --git a/keras_cv/layers/preprocessing/equalization.py b/keras_cv/layers/preprocessing/equalization.py
index cdc8808433..7dbeb636a6 100644
--- a/keras_cv/layers/preprocessing/equalization.py
+++ b/keras_cv/layers/preprocessing/equalization.py
@@ -12,17 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import partial
+
 import tensorflow as tf
 
 from keras_cv.api_export import keras_cv_export
-from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
-    BaseImageAugmentationLayer,
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
+    VectorizedBaseImageAugmentationLayer,
 )
 from keras_cv.utils import preprocessing
 
 
 @keras_cv_export("keras_cv.layers.Equalization")
-class Equalization(BaseImageAugmentationLayer):
+class Equalization(VectorizedBaseImageAugmentationLayer):
     """Equalization performs histogram equalization on a channel-wise basis.
 
     Args:
@@ -52,7 +54,7 @@ def __init__(self, value_range, bins=256, **kwargs):
         self.bins = bins
         self.value_range = value_range
 
-    def equalize_channel(self, image, channel_index):
+    def equalize_channel(self, images, channel_index):
         """equalize_channel performs histogram equalization on a single channel.
 
         Args:
@@ -60,9 +62,25 @@ def equalize_channel(self, image, channel_index):
                 with channels last
             channel_index: channel to equalize
         """
-        image = image[..., channel_index]
+        is_single_image = tf.rank(images) == 4 and tf.shape(images)[0] == 1
+
+        images = images[..., channel_index]
         # Compute the histogram of the image channel.
-        histogram = tf.histogram_fixed_width(image, [0, 255], nbins=self.bins)
+
+        # If the input is not a batch of images, directly using
+        # tf.histogram_fixed_width is much faster than using tf.vectorized_map
+        if is_single_image:
+            histogram = tf.histogram_fixed_width(
+                images, [0, 255], nbins=self.bins
+            )
+            histogram = tf.expand_dims(histogram, axis=0)
+        else:
+            partial_hist = partial(
+                tf.histogram_fixed_width, value_range=[0, 255], nbins=self.bins
+            )
+            histogram = tf.vectorized_map(
+                partial_hist, images, fallback_to_while_loop=True, warn=True
+            )
 
         # For the purposes of computing the step, filter out the non-zeros.
         # Zeroes are replaced by a big number while calculating min to keep
@@ -77,56 +95,94 @@ def equalize_channel(self, image, channel_index):
         )
 
         step = (
-            tf.reduce_sum(histogram) - tf.reduce_min(histogram_without_zeroes)
+            tf.reduce_sum(histogram, axis=-1)
+            - tf.reduce_min(histogram_without_zeroes, axis=-1)
         ) // (self.bins - 1)
 
         def build_mapping(histogram, step):
+            bacth_size = tf.shape(histogram)[0]
+
+            # Replace where step is 0 with 1 to avoid division by 0.
+            # This doesn't change the result, because where step==0 the
+            # original image is returned
+            _step = tf.where(
+                tf.equal(step, 0),
+                1,
+                step,
+            )
+            _step = tf.expand_dims(_step, -1)
+
             # Compute the cumulative sum, shifting by step // 2
             # and then normalization by step.
-            lookup_table = (tf.cumsum(histogram) + (step // 2)) // step
+            lookup_table = (
+                tf.cumsum(histogram, axis=-1) + (_step // 2)
+            ) // _step
+
             # Shift lookup_table, prepending with 0.
-            lookup_table = tf.concat([[0], lookup_table[:-1]], 0)
+            lookup_table = tf.concat(
+                [tf.tile([[0]], [bacth_size, 1]), lookup_table[..., :-1]],
+                axis=1,
+            )
+
             # Clip the counts to be in range. This is done
             # in the C code for image.point.
             return tf.clip_by_value(lookup_table, 0, 255)
 
         # If step is zero, return the original image. Otherwise, build
         # lookup table from the full histogram and step and then index from it.
-        result = tf.cond(
-            tf.equal(step, 0),
-            lambda: image,
-            lambda: tf.gather(build_mapping(histogram, step), image),
+        # The lookup table is built for all images,
+        # regardless of the corresponding value of step.
+        result = tf.where(
+            tf.reshape(tf.equal(step, 0), (-1, 1, 1)),
+            images,
+            tf.gather(
+                build_mapping(histogram, step), images, batch_dims=1, axis=1
+            ),
         )
 
         return result
 
-    def augment_image(self, image, **kwargs):
-        image = preprocessing.transform_value_range(
-            image, self.value_range, (0, 255), dtype=self.compute_dtype
+    def augment_images(self, images, transformations=None, **kwargs):
+        images = preprocessing.transform_value_range(
+            images, self.value_range, (0, 255), dtype=self.compute_dtype
         )
-        image = tf.cast(image, tf.int32)
-        image = tf.map_fn(
-            lambda channel: self.equalize_channel(image, channel),
-            tf.range(tf.shape(image)[-1]),
+        images = tf.cast(images, tf.int32)
+
+        images = tf.map_fn(
+            lambda channel: self.equalize_channel(images, channel),
+            tf.range(tf.shape(images)[-1]),
         )
+        images = tf.transpose(images, [1, 2, 3, 0])
 
-        image = tf.transpose(image, [1, 2, 0])
-        image = tf.cast(image, self.compute_dtype)
-        image = preprocessing.transform_value_range(
-            image, (0, 255), self.value_range, dtype=self.compute_dtype
+        images = tf.cast(images, self.compute_dtype)
+        images = preprocessing.transform_value_range(
+            images, (0, 255), self.value_range, dtype=self.compute_dtype
         )
-        return image
+        return images
 
     def augment_bounding_boxes(self, bounding_boxes, **kwargs):
         return bounding_boxes
 
-    def augment_label(self, label, transformation=None, **kwargs):
-        return label
+    def augment_labels(self, labels, transformations=None, **kwargs):
+        return labels
 
-    def augment_segmentation_mask(
-        self, segmentation_mask, transformation, **kwargs
+    def augment_segmentation_masks(
+        self, segmentation_masks, transformations, **kwargs
     ):
-        return segmentation_mask
+        return segmentation_masks
+
+    def augment_keypoints(self, keypoints, transformations, **kwargs):
+        return keypoints
+
+    def augment_targets(self, targets, transformations, **kwargs):
+        return targets
+
+    def augment_ragged_image(self, image, transformation, **kwargs):
+        image = tf.expand_dims(image, axis=0)
+        image = self.augment_images(
+            images=image, transformations=transformation, **kwargs
+        )
+        return tf.squeeze(image, axis=0)
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_cv/layers/preprocessing/equalization_test.py b/keras_cv/layers/preprocessing/equalization_test.py
index b6849ca142..937f19ca28 100644
--- a/keras_cv/layers/preprocessing/equalization_test.py
+++ b/keras_cv/layers/preprocessing/equalization_test.py
@@ -12,23 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import pytest
-import tensorflow as tf
 from absl.testing import parameterized
 
 from keras_cv.backend import keras
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.equalization import Equalization
 from keras_cv.tests.test_case import TestCase
 
 
 class EqualizationTest(TestCase):
     def test_return_shapes(self):
-        xs = 255 * tf.ones((2, 512, 512, 3), dtype=tf.int32)
+        xs = 255 * np.ones((2, 512, 512, 3), dtype=np.int32)
         layer = Equalization(value_range=(0, 255))
         xs = layer(xs)
 
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertAllEqual(xs, 255 * tf.ones((2, 512, 512, 3)))
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertAllEqual(xs, 255 * np.ones((2, 512, 512, 3)))
 
     @pytest.mark.tf_keras_only
     def test_return_shapes_inside_model(self):
@@ -40,28 +41,34 @@ def test_return_shapes_inside_model(self):
         self.assertEqual(model.output_shape, (None, 512, 512, 5))
 
     def test_equalizes_to_all_bins(self):
-        xs = tf.random.uniform((2, 512, 512, 3), 0, 255, dtype=tf.float32)
+        xs = np.random.uniform(size=(2, 512, 512, 3), low=0, high=255).astype(
+            np.float32
+        )
         layer = Equalization(value_range=(0, 255))
         xs = layer(xs)
 
         for i in range(0, 256):
-            self.assertTrue(tf.math.reduce_any(xs == i))
+            self.assertTrue(np.any(ops.convert_to_numpy(xs) == i))
 
     @parameterized.named_parameters(
-        ("float32", tf.float32), ("int32", tf.int32), ("int64", tf.int64)
+        ("float32", np.float32), ("int32", np.int32), ("int64", np.int64)
     )
     def test_input_dtypes(self, dtype):
-        xs = tf.random.uniform((2, 512, 512, 3), 0, 255, dtype=dtype)
+        xs = np.random.uniform(size=(2, 512, 512, 3), low=0, high=255).astype(
+            dtype
+        )
         layer = Equalization(value_range=(0, 255))
-        xs = layer(xs)
+        xs = ops.convert_to_numpy(layer(xs))
 
         for i in range(0, 256):
-            self.assertTrue(tf.math.reduce_any(xs == i))
+            self.assertTrue(np.any(xs == i))
         self.assertAllInRange(xs, 0, 255)
 
     @parameterized.named_parameters(("0_255", 0, 255), ("0_1", 0, 1))
     def test_output_range(self, lower, upper):
-        xs = tf.random.uniform((2, 512, 512, 3), lower, upper, dtype=tf.float32)
+        xs = np.random.uniform(
+            size=(2, 512, 512, 3), low=lower, high=upper
+        ).astype(np.float32)
         layer = Equalization(value_range=(lower, upper))
-        xs = layer(xs)
+        xs = ops.convert_to_numpy(layer(xs))
         self.assertAllInRange(xs, lower, upper)
diff --git a/keras_cv/layers/preprocessing/fourier_mix.py b/keras_cv/layers/preprocessing/fourier_mix.py
index b224fdc49e..46cf361da7 100644
--- a/keras_cv/layers/preprocessing/fourier_mix.py
+++ b/keras_cv/layers/preprocessing/fourier_mix.py
@@ -54,10 +54,12 @@ def __init__(self, alpha=0.5, decay_power=3, seed=None, **kwargs):
 
     def _sample_from_beta(self, alpha, beta, shape):
         sample_alpha = tf.random.gamma(
-            shape, alpha=alpha, seed=self._random_generator.make_legacy_seed()
+            shape,
+            alpha=alpha,
         )
         sample_beta = tf.random.gamma(
-            shape, alpha=beta, seed=self._random_generator.make_legacy_seed()
+            shape,
+            alpha=beta,
         )
         return sample_alpha / (sample_alpha + sample_beta)
 
@@ -100,7 +102,7 @@ def _get_spectrum(self, freqs, decay_power, channel, h, w):
         param_size = tf.concat(
             [tf.constant([channel]), tf.shape(freqs), tf.constant([2])], 0
         )
-        param = self._random_generator.random_normal(param_size)
+        param = self._random_generator.normal(param_size)
 
         scale = tf.expand_dims(scale, -1)[None, :]
 
diff --git a/keras_cv/layers/preprocessing/fourier_mix_test.py b/keras_cv/layers/preprocessing/fourier_mix_test.py
index d033aa83db..cc97c84017 100644
--- a/keras_cv/layers/preprocessing/fourier_mix_test.py
+++ b/keras_cv/layers/preprocessing/fourier_mix_test.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pytest
 import tensorflow as tf
 
 from keras_cv.layers.preprocessing.fourier_mix import FourierMix
@@ -50,9 +51,9 @@ def test_return_shapes(self):
             outputs["segmentation_masks"],
         )
 
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys.shape, [2, 10])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertEqual(ys.shape, (2, 10))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
 
     def test_fourier_mix_call_results_with_labels(self):
         xs = tf.cast(
@@ -110,6 +111,7 @@ def test_mix_up_call_results_with_masks(self):
         self.assertNotAllClose(ys_segmentation_masks, 1.0)
         self.assertNotAllClose(ys_segmentation_masks, 0.0)
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
diff --git a/keras_cv/layers/preprocessing/grayscale_test.py b/keras_cv/layers/preprocessing/grayscale_test.py
index 07072475e6..1a52d34019 100644
--- a/keras_cv/layers/preprocessing/grayscale_test.py
+++ b/keras_cv/layers/preprocessing/grayscale_test.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pytest
 import tensorflow as tf
 
 from keras_cv.layers import preprocessing
@@ -31,9 +32,10 @@ def test_return_shapes(self):
         )
         xs2 = layer(xs, training=True)
 
-        self.assertEqual(xs1.shape, [2, 52, 24, 1])
-        self.assertEqual(xs2.shape, [2, 52, 24, 3])
+        self.assertEqual(xs1.shape, (2, 52, 24, 1))
+        self.assertEqual(xs2.shape, (2, 52, 24, 3))
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack([2 * tf.ones((10, 10, 3)), tf.ones((10, 10, 3))], axis=0),
@@ -62,8 +64,8 @@ def augment(x):
 
         xs2 = augment(xs)
 
-        self.assertEqual(xs1.shape, [2, 10, 10, 1])
-        self.assertEqual(xs2.shape, [2, 10, 10, 3])
+        self.assertEqual(xs1.shape, (2, 10, 10, 1))
+        self.assertEqual(xs2.shape, (2, 10, 10, 3))
 
     def test_non_square_image(self):
         xs = tf.cast(
@@ -81,8 +83,8 @@ def test_non_square_image(self):
         )
         xs2 = layer(xs, training=True)
 
-        self.assertEqual(xs1.shape, [2, 52, 24, 1])
-        self.assertEqual(xs2.shape, [2, 52, 24, 3])
+        self.assertEqual(xs1.shape, (2, 52, 24, 1))
+        self.assertEqual(xs2.shape, (2, 52, 24, 3))
 
     def test_in_single_image(self):
         xs = tf.cast(
@@ -100,5 +102,5 @@ def test_in_single_image(self):
         )
         xs2 = layer(xs, training=True)
 
-        self.assertEqual(xs1.shape, [52, 24, 1])
-        self.assertEqual(xs2.shape, [52, 24, 3])
+        self.assertEqual(xs1.shape, (52, 24, 1))
+        self.assertEqual(xs2.shape, (52, 24, 3))
diff --git a/keras_cv/layers/preprocessing/grid_mask.py b/keras_cv/layers/preprocessing/grid_mask.py
index a7cea9f346..bdada47624 100644
--- a/keras_cv/layers/preprocessing/grid_mask.py
+++ b/keras_cv/layers/preprocessing/grid_mask.py
@@ -179,7 +179,7 @@ def _compute_grid_mask(self, input_shape, ratio):
         mask_side_len = tf.math.ceil(input_diagonal_len)
 
         # grid unit size
-        unit_size = self._random_generator.random_uniform(
+        unit_size = self._random_generator.uniform(
             shape=(),
             minval=tf.math.minimum(height * 0.5, width * 0.3),
             maxval=tf.math.maximum(height * 0.5, width * 0.3) + 1,
@@ -188,10 +188,10 @@ def _compute_grid_mask(self, input_shape, ratio):
         rectangle_side_len = tf.cast((ratio) * unit_size, tf.float32)
 
         # sample x and y offset for grid units randomly between 0 and unit_size
-        delta_x = self._random_generator.random_uniform(
+        delta_x = self._random_generator.uniform(
             shape=(), minval=0.0, maxval=unit_size, dtype=tf.float32
         )
-        delta_y = self._random_generator.random_uniform(
+        delta_y = self._random_generator.uniform(
             shape=(), minval=0.0, maxval=unit_size, dtype=tf.float32
         )
 
diff --git a/keras_cv/layers/preprocessing/grid_mask_test.py b/keras_cv/layers/preprocessing/grid_mask_test.py
index 3226a0ddd7..d06f9df9b8 100644
--- a/keras_cv/layers/preprocessing/grid_mask_test.py
+++ b/keras_cv/layers/preprocessing/grid_mask_test.py
@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import numpy as np
+import pytest
 import tensorflow as tf
 
 import keras_cv
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.grid_mask import GridMask
 from keras_cv.tests.test_case import TestCase
 
@@ -27,7 +29,7 @@ def test_return_shapes(self):
         layer = GridMask(ratio_factor=0.1, rotation_factor=(-0.2, 0.3))
         xs = layer(xs, training=True)
 
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
 
     def test_gridmask_call_results_one_channel(self):
         xs = tf.cast(
@@ -48,10 +50,14 @@ def test_gridmask_call_results_one_channel(self):
         xs = layer(xs, training=True)
 
         # Some pixels should be replaced with fill_value
-        self.assertTrue(tf.math.reduce_any(xs[0] == float(fill_value)))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 3.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == float(fill_value)))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 2.0))
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(xs[0]) == float(fill_value))
+        )
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 3.0))
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(xs[1]) == float(fill_value))
+        )
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 2.0))
 
     def test_non_square_image(self):
         xs = tf.cast(
@@ -72,11 +78,16 @@ def test_non_square_image(self):
         xs = layer(xs, training=True)
 
         # Some pixels should be replaced with fill_value
-        self.assertTrue(tf.math.reduce_any(xs[0] == float(fill_value)))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == float(fill_value)))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(xs[0]) == float(fill_value))
+        )
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(xs[1]) == float(fill_value))
+        )
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
@@ -100,10 +111,14 @@ def augment(x):
         xs = augment(xs)
 
         # Some pixels should be replaced with fill_value
-        self.assertTrue(tf.math.reduce_any(xs[0] == float(fill_value)))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == float(fill_value)))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(xs[0]) == float(fill_value))
+        )
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(xs[1]) == float(fill_value))
+        )
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
     def test_in_single_image(self):
         xs = tf.cast(
@@ -115,5 +130,5 @@ def test_in_single_image(self):
             ratio_factor=(0.5, 0.5), fill_mode="constant", fill_value=0.0
         )
         xs = layer(xs, training=True)
-        self.assertTrue(tf.math.reduce_any(xs == 0.0))
-        self.assertTrue(tf.math.reduce_any(xs == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs) == 0.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs) == 1.0))
diff --git a/keras_cv/layers/preprocessing/jittered_resize.py b/keras_cv/layers/preprocessing/jittered_resize.py
index d260c175ae..cb10ed77d4 100644
--- a/keras_cv/layers/preprocessing/jittered_resize.py
+++ b/keras_cv/layers/preprocessing/jittered_resize.py
@@ -166,7 +166,7 @@ def get_random_transformation_batch(
         max_offsets = tf.where(
             tf.less(max_offsets, 0), tf.zeros_like(max_offsets), max_offsets
         )
-        offsets = max_offsets * self._random_generator.random_uniform(
+        offsets = max_offsets * self._random_generator.uniform(
             shape=(batch_size, 2), minval=0, maxval=1, dtype=tf.float32
         )
         offsets = tf.cast(offsets, tf.int32)
diff --git a/keras_cv/layers/preprocessing/jittered_resize_test.py b/keras_cv/layers/preprocessing/jittered_resize_test.py
index 3bf2f1ad8f..dfa6a2ed69 100644
--- a/keras_cv/layers/preprocessing/jittered_resize_test.py
+++ b/keras_cv/layers/preprocessing/jittered_resize_test.py
@@ -17,6 +17,7 @@
 from keras_cv import bounding_box
 from keras_cv import core
 from keras_cv import layers
+from keras_cv.backend import ops
 from keras_cv.tests.test_case import TestCase
 
 
@@ -220,13 +221,15 @@ def test_output_dtypes(self):
             target_size=self.target_size,
             scale_factor=(3 / 4, 4 / 3),
         )
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+        )
         layer = layers.JitteredResize(
             target_size=self.target_size,
             scale_factor=(3 / 4, 4 / 3),
             dtype="uint8",
         )
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
 
     def test_config(self):
         layer = layers.JitteredResize(
diff --git a/keras_cv/layers/preprocessing/mix_up.py b/keras_cv/layers/preprocessing/mix_up.py
index d75fc9a33c..6cb0622193 100644
--- a/keras_cv/layers/preprocessing/mix_up.py
+++ b/keras_cv/layers/preprocessing/mix_up.py
@@ -57,10 +57,12 @@ def __init__(self, alpha=0.2, seed=None, **kwargs):
 
     def _sample_from_beta(self, alpha, beta, shape):
         sample_alpha = tf.random.gamma(
-            shape, alpha=alpha, seed=self._random_generator.make_legacy_seed()
+            shape,
+            alpha=alpha,
         )
         sample_beta = tf.random.gamma(
-            shape, alpha=beta, seed=self._random_generator.make_legacy_seed()
+            shape,
+            alpha=beta,
         )
         return sample_alpha / (sample_alpha + sample_beta)
 
diff --git a/keras_cv/layers/preprocessing/mix_up_test.py b/keras_cv/layers/preprocessing/mix_up_test.py
index ec1e0af1a6..332e0e226c 100644
--- a/keras_cv/layers/preprocessing/mix_up_test.py
+++ b/keras_cv/layers/preprocessing/mix_up_test.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pytest
 import tensorflow as tf
 
 from keras_cv.layers.preprocessing.mix_up import MixUp
@@ -60,11 +61,11 @@ def test_return_shapes(self):
             outputs["segmentation_masks"],
         )
 
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys_labels.shape, [2, 10])
-        self.assertEqual(ys_bounding_boxes["boxes"].shape, [2, 6, 4])
-        self.assertEqual(ys_bounding_boxes["classes"].shape, [2, 6])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertEqual(ys_labels.shape, (2, 10))
+        self.assertEqual(ys_bounding_boxes["boxes"].shape, (2, 6, 4))
+        self.assertEqual(ys_bounding_boxes["classes"].shape, (2, 6))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
 
     def test_mix_up_call_results_with_labels(self):
         xs = tf.cast(
@@ -122,6 +123,7 @@ def test_mix_up_call_results_with_masks(self):
         self.assertNotAllClose(ys_segmentation_masks, 1.0)
         self.assertNotAllClose(ys_segmentation_masks, 0.0)
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
diff --git a/keras_cv/layers/preprocessing/mosaic.py b/keras_cv/layers/preprocessing/mosaic.py
index 9b7580f19f..2ecf3421d6 100644
--- a/keras_cv/layers/preprocessing/mosaic.py
+++ b/keras_cv/layers/preprocessing/mosaic.py
@@ -96,7 +96,7 @@ def __init__(
 
     def get_random_transformation_batch(self, batch_size, **kwargs):
         # pick 3 indices for every batch to create the mosaic output with.
-        permutation_order = self._random_generator.random_uniform(
+        permutation_order = self._random_generator.uniform(
             (batch_size, 3),
             minval=0,
             maxval=batch_size,
diff --git a/keras_cv/layers/preprocessing/mosaic_test.py b/keras_cv/layers/preprocessing/mosaic_test.py
index c5f5c9d1c1..5d2561731c 100644
--- a/keras_cv/layers/preprocessing/mosaic_test.py
+++ b/keras_cv/layers/preprocessing/mosaic_test.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pytest
 import tensorflow as tf
 
 from keras_cv.layers.preprocessing.mosaic import Mosaic
@@ -59,6 +60,7 @@ def test_return_shapes(self):
         self.assertEqual(ys_bounding_boxes["classes"].shape, [2, None])
         self.assertEqual(ys_segmentation_masks.shape, input_shape)
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
diff --git a/keras_cv/layers/preprocessing/posterization_test.py b/keras_cv/layers/preprocessing/posterization_test.py
index 0825ffe3a4..8ea37b5565 100644
--- a/keras_cv/layers/preprocessing/posterization_test.py
+++ b/keras_cv/layers/preprocessing/posterization_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
+import pytest
 import tensorflow as tf
 
 from keras_cv.layers.preprocessing.posterization import Posterization
@@ -74,6 +75,7 @@ def test_batched_input(self):
 
         self.assertAllEqual(output, expected_output)
 
+    @pytest.mark.tf_only
     def test_works_with_xla(self):
         dummy_input = self.rng.uniform(shape=(2, 224, 224, 3))
         layer = Posterization(bits=4, value_range=[0, 1])
diff --git a/keras_cv/layers/preprocessing/rand_augment_test.py b/keras_cv/layers/preprocessing/rand_augment_test.py
index 739fd9c47a..5f93349c29 100644
--- a/keras_cv/layers/preprocessing/rand_augment_test.py
+++ b/keras_cv/layers/preprocessing/rand_augment_test.py
@@ -16,10 +16,20 @@
 from absl.testing import parameterized
 
 from keras_cv import layers
+from keras_cv.backend import ops
 from keras_cv.tests.test_case import TestCase
 
 
 class RandAugmentTest(TestCase):
+    def test_zero_rate_pass_through(self):
+        rand_augment = layers.RandAugment(
+            value_range=(0, 255),
+            rate=0.0,
+        )
+        xs = np.ones((2, 512, 512, 3))
+        ys = rand_augment(xs)
+        self.assertAllClose(ys, xs)
+
     @parameterized.named_parameters(
         ("0", 0),
         ("20", 0.2),
@@ -47,10 +57,8 @@ def test_runs_with_value_range(self, low, high):
             value_range=(low, high),
         )
         xs = tf.random.uniform((2, 512, 512, 3), low, high, dtype=tf.float32)
-        ys = rand_augment(xs)
-        self.assertTrue(
-            tf.math.reduce_all(tf.logical_and(ys >= low, ys <= high))
-        )
+        ys = ops.convert_to_numpy(rand_augment(xs))
+        self.assertTrue(np.all(np.logical_and(ys >= low, ys <= high)))
 
     @parameterized.named_parameters(
         ("float32", "float32"),
@@ -76,9 +84,9 @@ def test_standard_policy_respects_value_range(self, lower, upper):
             layers=my_layers, augmentations_per_image=3
         )
         xs = tf.random.uniform((2, 512, 512, 3), lower, upper, dtype=tf.float32)
-        ys = rand_augment(xs)
-        self.assertLessEqual(tf.math.reduce_max(ys), upper)
-        self.assertGreaterEqual(tf.math.reduce_min(ys), lower)
+        ys = ops.convert_to_numpy(rand_augment(xs))
+        self.assertLessEqual(np.max(ys), upper)
+        self.assertGreaterEqual(np.min(ys), lower)
 
     def test_runs_unbatched(self):
         rand_augment = layers.RandAugment(
diff --git a/keras_cv/layers/preprocessing/random_apply.py b/keras_cv/layers/preprocessing/random_apply.py
index be0aa2a408..166764bff9 100644
--- a/keras_cv/layers/preprocessing/random_apply.py
+++ b/keras_cv/layers/preprocessing/random_apply.py
@@ -112,9 +112,7 @@ def __init__(
         self.built = True
 
     def _should_augment(self):
-        return (
-            self._random_generator.random_uniform(shape=()) > 1.0 - self._rate
-        )
+        return self._random_generator.uniform(shape=()) > 1.0 - self._rate
 
     def _batch_augment(self, inputs):
         if self.batchwise:
diff --git a/keras_cv/layers/preprocessing/random_apply_test.py b/keras_cv/layers/preprocessing/random_apply_test.py
index 86a62eefe6..693e788be0 100644
--- a/keras_cv/layers/preprocessing/random_apply_test.py
+++ b/keras_cv/layers/preprocessing/random_apply_test.py
@@ -11,11 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 
 from keras_cv import layers
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
 )
@@ -49,7 +50,7 @@ def test_works_with_batched_input(self):
         dummy_inputs = self.rng.uniform(shape=(batch_size, 224, 224, 3))
         layer = RandomApply(rate=0.5, layer=ZeroOut(), seed=1234)
 
-        outputs = layer(dummy_inputs)
+        outputs = ops.convert_to_numpy(layer(dummy_inputs))
         num_zero_inputs = self._num_zero_batches(dummy_inputs)
         num_zero_outputs = self._num_zero_batches(outputs)
 
@@ -107,6 +108,7 @@ def test_can_modify_label(self):
 
         self.assertAllEqual(outputs["labels"], tf.zeros_like(dummy_labels))
 
+    @pytest.mark.tf_only
     def test_works_with_xla(self):
         dummy_inputs = self.rng.uniform(shape=(32, 224, 224, 3))
         # auto_vectorize=True will crash XLA
diff --git a/keras_cv/layers/preprocessing/random_aspect_ratio.py b/keras_cv/layers/preprocessing/random_aspect_ratio.py
index ee321812de..4839b2d7ba 100644
--- a/keras_cv/layers/preprocessing/random_aspect_ratio.py
+++ b/keras_cv/layers/preprocessing/random_aspect_ratio.py
@@ -14,12 +14,13 @@
 
 import tensorflow as tf
 
-import keras_cv
 from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
 )
+from keras_cv.utils import get_interpolation
+from keras_cv.utils import parse_factor
 
 
 @keras_cv_export("keras_cv.layers.RandomAspectRatio")
@@ -47,8 +48,8 @@ def __init__(
         **kwargs
     ):
         super().__init__(**kwargs)
-        self.interpolation = keras_cv.utils.get_interpolation(interpolation)
-        self.factor = keras_cv.utils.parse_factor(
+        self.interpolation = get_interpolation(interpolation)
+        self.factor = parse_factor(
             factor,
             min_value=0.0,
             max_value=None,
diff --git a/keras_cv/layers/preprocessing/random_augmentation_pipeline.py b/keras_cv/layers/preprocessing/random_augmentation_pipeline.py
index a593af9074..e437a90147 100644
--- a/keras_cv/layers/preprocessing/random_augmentation_pipeline.py
+++ b/keras_cv/layers/preprocessing/random_augmentation_pipeline.py
@@ -81,7 +81,7 @@ def __init__(
         seed=None,
         **kwargs,
     ):
-        super().__init__(**kwargs, seed=seed, force_generator=True)
+        super().__init__(**kwargs, seed=seed)
         self.augmentations_per_image = augmentations_per_image
         self.rate = rate
         self.layers = list(layers)
@@ -98,12 +98,12 @@ def _augment(self, inputs):
 
         result = inputs
         for _ in range(self.augmentations_per_image):
-            skip_augment = self._random_generator.random_uniform(
+            skip_augment = self._random_generator.uniform(
                 shape=(), minval=0.0, maxval=1.0, dtype=tf.float32
             )
             result = tf.cond(
                 skip_augment > self.rate,
-                lambda: inputs,
+                lambda: result,
                 lambda: self._random_choice(result),
             )
         return result
diff --git a/keras_cv/layers/preprocessing/random_brightness.py b/keras_cv/layers/preprocessing/random_brightness.py
index ed30a1ecc7..9c2a8fb8c5 100644
--- a/keras_cv/layers/preprocessing/random_brightness.py
+++ b/keras_cv/layers/preprocessing/random_brightness.py
@@ -62,7 +62,7 @@ class RandomBrightness(VectorizedBaseImageAugmentationLayer):
     """
 
     def __init__(self, factor, value_range=(0, 255), seed=None, **kwargs):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         if isinstance(factor, float) or isinstance(factor, int):
             factor = (-factor, factor)
         self.factor = preprocessing_utils.parse_factor(
diff --git a/keras_cv/layers/preprocessing/random_channel_shift_test.py b/keras_cv/layers/preprocessing/random_channel_shift_test.py
index 9f1f724eb5..0b2b8b1086 100644
--- a/keras_cv/layers/preprocessing/random_channel_shift_test.py
+++ b/keras_cv/layers/preprocessing/random_channel_shift_test.py
@@ -13,8 +13,10 @@
 # limitations under the License.
 
 import numpy as np
+import pytest
 import tensorflow as tf
 
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.tests.test_case import TestCase
 
@@ -27,7 +29,7 @@ def test_return_shapes(self):
         )
 
         xs = layer(xs, training=True)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
 
     def test_non_square_image(self):
         xs = tf.cast(
@@ -42,9 +44,10 @@ def test_non_square_image(self):
         )
 
         xs = layer(xs, training=True)
-        self.assertFalse(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertFalse(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertFalse(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertFalse(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
@@ -61,8 +64,8 @@ def augment(x):
             return layer(x, training=True)
 
         xs = augment(xs)
-        self.assertFalse(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertFalse(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertFalse(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertFalse(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
     def test_5_channels(self):
         xs = tf.cast(
@@ -73,7 +76,7 @@ def test_5_channels(self):
             factor=0.4, channels=5, value_range=(0, 255)
         )
         xs = layer(xs, training=True)
-        self.assertFalse(tf.math.reduce_any(xs == 1.0))
+        self.assertFalse(np.any(ops.convert_to_numpy(xs) == 1.0))
 
     def test_1_channel(self):
         xs = tf.cast(
@@ -84,7 +87,7 @@ def test_1_channel(self):
             factor=0.4, channels=1, value_range=(0, 255)
         )
         xs = layer(xs, training=True)
-        self.assertFalse(tf.math.reduce_any(xs == 1.0))
+        self.assertFalse(np.any(ops.convert_to_numpy(xs) == 1.0))
 
     def test_in_single_image(self):
         xs = tf.cast(
@@ -95,7 +98,7 @@ def test_in_single_image(self):
             factor=0.4, value_range=(0, 255)
         )
         xs = layer(xs, training=True)
-        self.assertFalse(tf.math.reduce_any(xs == 1.0))
+        self.assertFalse(np.any(ops.convert_to_numpy(xs) == 1.0))
 
     def test_config(self):
         layer = preprocessing.RandomChannelShift(
diff --git a/keras_cv/layers/preprocessing/random_choice.py b/keras_cv/layers/preprocessing/random_choice.py
index a2b1fbef81..eff23dc022 100644
--- a/keras_cv/layers/preprocessing/random_choice.py
+++ b/keras_cv/layers/preprocessing/random_choice.py
@@ -68,7 +68,7 @@ def __init__(
         seed=None,
         **kwargs,
     ):
-        super().__init__(**kwargs, seed=seed, force_generator=True)
+        super().__init__(**kwargs, seed=seed)
         self.layers = layers
         self.auto_vectorize = auto_vectorize
         self.batchwise = batchwise
@@ -87,7 +87,7 @@ def _batch_augment(self, inputs):
             return super()._batch_augment(inputs)
 
     def _augment(self, inputs, *args, **kwargs):
-        selected_op = self._random_generator.random_uniform(
+        selected_op = self._random_generator.uniform(
             (), minval=0, maxval=len(self.layers), dtype=tf.int32
         )
         # Warning:
diff --git a/keras_cv/layers/preprocessing/random_color_degeneration.py b/keras_cv/layers/preprocessing/random_color_degeneration.py
index f678a054b2..1d8206a3a2 100644
--- a/keras_cv/layers/preprocessing/random_color_degeneration.py
+++ b/keras_cv/layers/preprocessing/random_color_degeneration.py
@@ -16,14 +16,14 @@
 
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
-from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
-    BaseImageAugmentationLayer,
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
+    VectorizedBaseImageAugmentationLayer,
 )
 from keras_cv.utils import preprocessing
 
 
 @keras_cv_export("keras_cv.layers.RandomColorDegeneration")
-class RandomColorDegeneration(BaseImageAugmentationLayer):
+class RandomColorDegeneration(VectorizedBaseImageAugmentationLayer):
     """Randomly performs the color degeneration operation on given images.
 
     The sharpness operation first converts an image to gray scale, then back to
@@ -57,24 +57,39 @@ def __init__(
         )
         self.seed = seed
 
-    def get_random_transformation(self, **kwargs):
-        return self.factor(dtype=self.compute_dtype)
+    def get_random_transformation_batch(self, batch_size, **kwargs):
+        return self.factor(
+            shape=(batch_size, 1, 1, 1), dtype=self.compute_dtype
+        )
 
-    def augment_image(self, image, transformation=None, **kwargs):
-        degenerate = tf.image.grayscale_to_rgb(tf.image.rgb_to_grayscale(image))
-        result = preprocessing.blend(image, degenerate, transformation)
+    def augment_images(self, images, transformations=None, **kwargs):
+        degenerates = tf.image.grayscale_to_rgb(
+            tf.image.rgb_to_grayscale(images)
+        )
+        result = preprocessing.blend(images, degenerates, transformations)
         return result
 
     def augment_bounding_boxes(self, bounding_boxes, **kwargs):
         return bounding_boxes
 
-    def augment_label(self, label, transformation=None, **kwargs):
-        return label
+    def augment_labels(self, labels, transformations=None, **kwargs):
+        return labels
 
-    def augment_segmentation_mask(
-        self, segmentation_mask, transformation, **kwargs
+    def augment_segmentation_masks(
+        self, segmentation_masks, transformations, **kwargs
     ):
-        return segmentation_mask
+        return segmentation_masks
+
+    def augment_keypoints(self, keypoints, transformations, **kwargs):
+        return keypoints
+
+    def augment_targets(self, targets, transformations, **kwargs):
+        return targets
+
+    def augment_ragged_image(self, image, transformation, **kwargs):
+        return self.augment_images(
+            image, transformations=transformation, **kwargs
+        )
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_cv/layers/preprocessing/random_color_degeneration_test.py b/keras_cv/layers/preprocessing/random_color_degeneration_test.py
index 975a95b8ae..a27a7b67fd 100644
--- a/keras_cv/layers/preprocessing/random_color_degeneration_test.py
+++ b/keras_cv/layers/preprocessing/random_color_degeneration_test.py
@@ -14,6 +14,7 @@
 import numpy as np
 import tensorflow as tf
 
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.tests.test_case import TestCase
 
@@ -39,7 +40,7 @@ def test_color_degeneration_full_factor(self):
         xs = tf.concat([r, g, b], axis=-1)
 
         layer = preprocessing.RandomColorDegeneration(factor=(1, 1))
-        ys = layer(xs)
+        ys = ops.convert_to_numpy(layer(xs))
 
         # Color degeneration uses standard luma conversion for RGB->Grayscale.
         # The formula for luma is result= 0.2989*r + 0.5870*g + 0.1140*b
@@ -54,7 +55,7 @@ def test_color_degeneration_70p_factor(self):
         xs = tf.concat([r, g, b], axis=-1)
 
         layer = preprocessing.RandomColorDegeneration(factor=(0.7, 0.7))
-        ys = layer(xs)
+        ys = ops.convert_to_numpy(layer(xs))
 
         # Color degeneration uses standard luma conversion for RGB->Grayscale.
         # The formula for luma is result= 0.2989*r + 0.5870*g + 0.1140*b
diff --git a/keras_cv/layers/preprocessing/random_color_jitter_test.py b/keras_cv/layers/preprocessing/random_color_jitter_test.py
index 17fa2f8b60..24387d962a 100644
--- a/keras_cv/layers/preprocessing/random_color_jitter_test.py
+++ b/keras_cv/layers/preprocessing/random_color_jitter_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import numpy as np
+import pytest
 import tensorflow as tf
 
 from keras_cv.layers import preprocessing
@@ -37,9 +38,9 @@ def test_return_shapes(self):
         non_square_batch_output = layer(non_square_batch_input, training=True)
         unbatch_output = layer(unbatch_input, training=True)
 
-        self.assertEqual(batch_output.shape, [2, 512, 512, 3])
-        self.assertEqual(non_square_batch_output.shape, [2, 1024, 512, 3])
-        self.assertEqual(unbatch_output.shape, [512, 512, 3])
+        self.assertEqual(batch_output.shape, (2, 512, 512, 3))
+        self.assertEqual(non_square_batch_output.shape, (2, 1024, 512, 3))
+        self.assertEqual(unbatch_output.shape, (512, 512, 3))
 
     # Test 2: Check if the factor ranges are set properly.
     def test_factor_range(self):
@@ -57,6 +58,7 @@ def test_factor_range(self):
         self.assertEqual(layer.hue_factor, (0.5, 0.9))
 
     # Test 3: Test if it is OK to run on graph mode.
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         inputs = np.ones((2, 512, 512, 3))
 
diff --git a/keras_cv/layers/preprocessing/random_contrast.py b/keras_cv/layers/preprocessing/random_contrast.py
index d265001841..d66e7ea21d 100644
--- a/keras_cv/layers/preprocessing/random_contrast.py
+++ b/keras_cv/layers/preprocessing/random_contrast.py
@@ -67,7 +67,7 @@ class RandomContrast(VectorizedBaseImageAugmentationLayer):
     """
 
     def __init__(self, value_range, factor, seed=None, **kwargs):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         if isinstance(factor, (tuple, list)):
             min = 1 - factor[0]
             max = 1 + factor[1]
diff --git a/keras_cv/layers/preprocessing/random_crop.py b/keras_cv/layers/preprocessing/random_crop.py
index 05e0b24bfc..b64674d914 100644
--- a/keras_cv/layers/preprocessing/random_crop.py
+++ b/keras_cv/layers/preprocessing/random_crop.py
@@ -60,7 +60,9 @@ def __init__(
         self, height, width, seed=None, bounding_box_format=None, **kwargs
     ):
         super().__init__(
-            **kwargs, autocast=False, seed=seed, force_generator=True
+            **kwargs,
+            autocast=False,
+            seed=seed,
         )
         self.height = height
         self.width = width
@@ -79,13 +81,13 @@ def compute_ragged_image_signature(self, images):
 
     def get_random_transformation_batch(self, batch_size, **kwargs):
         tops = tf.cast(
-            self._random_generator.random_uniform(
+            self._random_generator.uniform(
                 shape=(batch_size, 1), minval=0, maxval=1
             ),
             self.compute_dtype,
         )
         lefts = tf.cast(
-            self._random_generator.random_uniform(
+            self._random_generator.uniform(
                 shape=(batch_size, 1), minval=0, maxval=1
             ),
             self.compute_dtype,
diff --git a/keras_cv/layers/preprocessing/random_crop_and_resize.py b/keras_cv/layers/preprocessing/random_crop_and_resize.py
index d660e982c5..cd947d5835 100644
--- a/keras_cv/layers/preprocessing/random_crop_and_resize.py
+++ b/keras_cv/layers/preprocessing/random_crop_and_resize.py
@@ -109,14 +109,14 @@ def get_random_transformation(
             tf.sqrt(crop_area_factor * aspect_ratio), 0.0, 1.0
         )
 
-        height_offset = self._random_generator.random_uniform(
+        height_offset = self._random_generator.uniform(
             (),
             minval=tf.minimum(0.0, 1.0 - new_height),
             maxval=tf.maximum(0.0, 1.0 - new_height),
             dtype=tf.float32,
         )
 
-        width_offset = self._random_generator.random_uniform(
+        width_offset = self._random_generator.uniform(
             (),
             minval=tf.minimum(0.0, 1.0 - new_width),
             maxval=tf.maximum(0.0, 1.0 - new_width),
@@ -272,10 +272,10 @@ def from_config(cls, config):
                 config["crop_area_factor"]
             )
         if isinstance(config["aspect_ratio_factor"], dict):
-            config[
-                "aspect_ratio_factor"
-            ] = keras.utils.deserialize_keras_object(
-                config["aspect_ratio_factor"]
+            config["aspect_ratio_factor"] = (
+                keras.utils.deserialize_keras_object(
+                    config["aspect_ratio_factor"]
+                )
             )
         return cls(**config)
 
diff --git a/keras_cv/layers/preprocessing/random_crop_and_resize_test.py b/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
index dc8c2b76b8..34d4613f93 100644
--- a/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
+++ b/keras_cv/layers/preprocessing/random_crop_and_resize_test.py
@@ -16,6 +16,7 @@
 from absl.testing import parameterized
 
 from keras_cv import bounding_box
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.tests.test_case import TestCase
 
@@ -147,7 +148,9 @@ def test_augment_sparse_segmentation_mask(self):
             seed=self.seed,
         )
         output = layer(inputs, training=True)
-        self.assertAllInSet(output["segmentation_masks"], [0, 7])
+        self.assertAllInSet(
+            ops.convert_to_numpy(output["segmentation_masks"]), [0, 7]
+        )
 
     def test_augment_one_hot_segmentation_mask(self):
         num_classes = 8
diff --git a/keras_cv/layers/preprocessing/random_crop_test.py b/keras_cv/layers/preprocessing/random_crop_test.py
index 43fff1e764..bd00b0e21b 100644
--- a/keras_cv/layers/preprocessing/random_crop_test.py
+++ b/keras_cv/layers/preprocessing/random_crop_test.py
@@ -15,10 +15,12 @@
 import unittest
 
 import numpy as np
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 
 from keras_cv import layers as cv_layers
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.random_crop import RandomCrop
 from keras_cv.tests.test_case import TestCase
 
@@ -63,6 +65,7 @@ def test_input_smaller_than_crop_box(self):
         expected_output = resizing_layer(inp)
         self.assertAllEqual(expected_output, actual_output)
 
+    @pytest.mark.skip(reason="need to update tests for keras 3")
     def test_training_with_mock(self):
         np.random.seed(1337)
         batch_size = 12
@@ -106,7 +109,7 @@ def test_unbatched_image(self):
         layer = RandomCrop(8, 8)
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_offset,
         ):
             actual_output = layer(inp, training=True)
@@ -120,7 +123,7 @@ def test_batched_input(self):
         layer = RandomCrop(8, 8)
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_offset,
         ):
             actual_output = layer(inp, training=True)
@@ -183,6 +186,7 @@ def test_augment_bounding_boxes_resize(self):
         )
         self.assertAllClose(expected_output, output["bounding_boxes"]["boxes"])
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         np.random.seed(1337)
         inp = np.random.random((20, 16, 16, 3))
@@ -195,7 +199,7 @@ def augment(x):
 
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_offset,
         ):
             actual_output = augment(inp)
@@ -250,9 +254,11 @@ def test_config_with_custom_name(self):
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = RandomCrop(2, 2)
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+        )
         layer = RandomCrop(2, 2, dtype="uint8")
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
 
     def test_config(self):
         layer = RandomCrop(height=2, width=3, bounding_box_format="xyxy")
diff --git a/keras_cv/layers/preprocessing/random_cutout.py b/keras_cv/layers/preprocessing/random_cutout.py
index 4eb4bb4a24..8760630f5c 100644
--- a/keras_cv/layers/preprocessing/random_cutout.py
+++ b/keras_cv/layers/preprocessing/random_cutout.py
@@ -15,15 +15,18 @@
 import tensorflow as tf
 
 from keras_cv.api_export import keras_cv_export
-from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
-    BaseImageAugmentationLayer,
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
+    VectorizedBaseImageAugmentationLayer,
 )
 from keras_cv.utils import fill_utils
 from keras_cv.utils import preprocessing
 
+H_AXIS = -3
+W_AXIS = -2
+
 
 @keras_cv_export("keras_cv.layers.RandomCutout")
-class RandomCutout(BaseImageAugmentationLayer):
+class RandomCutout(VectorizedBaseImageAugmentationLayer):
     """Randomly cut out rectangles from images and fill them.
 
     Args:
@@ -90,74 +93,132 @@ def __init__(
                 f'or "constant". Got `fill_mode`={fill_mode}'
             )
 
-    def _parse_bounds(self, factor):
-        if isinstance(factor, (tuple, list)):
-            return factor[0], factor[1]
-        else:
-            return type(factor)(0), factor
-
-    def get_random_transformation(self, image=None, **kwargs):
-        center_x, center_y = self._compute_rectangle_position(image)
-        rectangle_height, rectangle_width = self._compute_rectangle_size(image)
-        return center_x, center_y, rectangle_height, rectangle_width
+    def get_random_transformation_batch(self, batch_size, images, **kwargs):
+        centers_x, centers_y = self._compute_rectangle_position(images)
+        rectangles_height, rectangles_width = self._compute_rectangle_size(
+            images
+        )
+        return {
+            "centers_x": centers_x,
+            "centers_y": centers_y,
+            "rectangles_height": rectangles_height,
+            "rectangles_width": rectangles_width,
+        }
 
-    def augment_image(self, image, transformation=None, **kwargs):
+    def augment_images(self, images, transformations=None, **kwargs):
         """Apply random cutout."""
-        inputs = tf.expand_dims(image, 0)
-        center_x, center_y, rectangle_height, rectangle_width = transformation
-
-        rectangle_fill = self._compute_rectangle_fill(inputs)
-        inputs = fill_utils.fill_rectangle(
-            inputs,
-            center_x,
-            center_y,
-            rectangle_width,
-            rectangle_height,
-            rectangle_fill,
+        centers_x, centers_y = (
+            transformations["centers_x"],
+            transformations["centers_y"],
+        )
+        rectangles_height, rectangles_width = (
+            transformations["rectangles_height"],
+            transformations["rectangles_width"],
+        )
+
+        rectangles_fill = self._compute_rectangle_fill(images)
+        images = fill_utils.fill_rectangle(
+            images,
+            centers_x,
+            centers_y,
+            rectangles_width,
+            rectangles_height,
+            rectangles_fill,
         )
-        return inputs[0]
+        return images
 
-    def augment_label(self, label, transformation=None, **kwargs):
-        return label
+    def augment_bounding_boxes(self, bounding_boxes, **kwargs):
+        return bounding_boxes
 
-    def augment_segmentation_mask(
-        self, segmentation_masks, transformation=None, **kwargs
+    def augment_labels(self, labels, transformations=None, **kwargs):
+        return labels
+
+    def augment_segmentation_masks(
+        self, segmentation_masks, transformations, **kwargs
     ):
         return segmentation_masks
 
-    def _compute_rectangle_position(self, inputs):
-        input_shape = tf.shape(inputs)
-        image_height, image_width = (
-            input_shape[0],
-            input_shape[1],
+    def augment_keypoints(self, keypoints, transformations, **kwargs):
+        return keypoints
+
+    def augment_targets(self, targets, transformations, **kwargs):
+        return targets
+
+    def augment_ragged_image(self, image, transformation, **kwargs):
+        image = tf.expand_dims(image, axis=0)
+        centers_x, centers_y = (
+            transformation["centers_x"],
+            transformation["centers_y"],
+        )
+        rectangles_height, rectangles_width = (
+            transformation["rectangles_height"],
+            transformation["rectangles_width"],
+        )
+        transformation = {
+            "centers_x": tf.expand_dims(centers_x, axis=0),
+            "centers_y": tf.expand_dims(centers_y, axis=0),
+            "rectangles_height": tf.expand_dims(rectangles_height, axis=0),
+            "rectangles_width": tf.expand_dims(rectangles_width, axis=0),
+        }
+        image = self.augment_images(
+            images=image, transformations=transformation, **kwargs
         )
-        center_x = self._random_generator.random_uniform(
-            [1], 0, image_width, dtype=tf.int32
+        return tf.squeeze(image, axis=0)
+
+    def _get_image_shape(self, images):
+        if isinstance(images, tf.RaggedTensor):
+            heights = tf.reshape(images.row_lengths(), (-1,))
+            widths = tf.reshape(
+                tf.reduce_max(images.row_lengths(axis=2), 1), (-1,)
+            )
+        else:
+            batch_size = tf.shape(images)[0]
+            heights = tf.repeat(tf.shape(images)[H_AXIS], repeats=[batch_size])
+            heights = tf.reshape(heights, shape=(-1,))
+            widths = tf.repeat(tf.shape(images)[W_AXIS], repeats=[batch_size])
+            widths = tf.reshape(widths, shape=(-1,))
+        return tf.cast(heights, dtype=tf.int32), tf.cast(widths, dtype=tf.int32)
+
+    def _compute_rectangle_position(self, inputs):
+        batch_size = tf.shape(inputs)[0]
+        heights, widths = self._get_image_shape(inputs)
+
+        # generate values in float32 and then cast (i.e. round) to int32 because
+        # random.uniform do not support maxval broadcasting for integer types.
+        # Needed because maxval is a 1-D tensor to support ragged inputs.
+
+        heights = tf.cast(heights, dtype=tf.float32)
+        widths = tf.cast(widths, dtype=tf.float32)
+
+        center_x = self._random_generator.uniform(
+            (batch_size,), 0, widths, dtype=tf.float32
         )
-        center_y = self._random_generator.random_uniform(
-            [1], 0, image_height, dtype=tf.int32
+        center_y = self._random_generator.uniform(
+            (batch_size,), 0, heights, dtype=tf.float32
         )
+
+        center_x = tf.cast(center_x, tf.int32)
+        center_y = tf.cast(center_y, tf.int32)
+
         return center_x, center_y
 
     def _compute_rectangle_size(self, inputs):
-        input_shape = tf.shape(inputs)
-        image_height, image_width = (
-            input_shape[0],
-            input_shape[1],
-        )
-        height = self.height_factor()
-        width = self.width_factor()
+        batch_size = tf.shape(inputs)[0]
+        images_heights, images_widths = self._get_image_shape(inputs)
 
-        height = height * tf.cast(image_height, tf.float32)
-        width = width * tf.cast(image_width, tf.float32)
+        height = self.height_factor(shape=(batch_size,))
+        width = self.width_factor(shape=(batch_size,))
+
+        height = height * tf.cast(images_heights, tf.float32)
+        width = width * tf.cast(images_widths, tf.float32)
 
         height = tf.cast(tf.math.ceil(height), tf.int32)
         width = tf.cast(tf.math.ceil(width), tf.int32)
 
-        height = tf.minimum(height, image_height)
-        width = tf.minimum(width, image_width)
+        height = tf.minimum(height, images_heights)
+        width = tf.minimum(width, images_heights)
 
-        return tf.expand_dims(height, axis=0), tf.expand_dims(width, axis=0)
+        return height, width
 
     def _compute_rectangle_fill(self, inputs):
         input_shape = tf.shape(inputs)
@@ -167,16 +228,25 @@ def _compute_rectangle_fill(self, inputs):
         else:
             # gaussian noise
             fill_value = tf.random.normal(input_shape, dtype=self.compute_dtype)
-
+            # rescale the random noise to the original image range
+            image_max = tf.reduce_max(inputs)
+            image_min = tf.reduce_min(inputs)
+            fill_max = tf.reduce_max(fill_value)
+            fill_min = tf.reduce_min(fill_value)
+            fill_value = (image_max - image_min) * (fill_value - fill_min) / (
+                fill_max - fill_min
+            ) + image_min
         return fill_value
 
     def get_config(self):
-        config = {
-            "height_factor": self.height_factor,
-            "width_factor": self.width_factor,
-            "fill_mode": self.fill_mode,
-            "fill_value": self.fill_value,
-            "seed": self.seed,
-        }
-        base_config = super().get_config()
-        return dict(list(base_config.items()) + list(config.items()))
+        config = super().get_config()
+        config.update(
+            {
+                "height_factor": self.height_factor,
+                "width_factor": self.width_factor,
+                "fill_mode": self.fill_mode,
+                "fill_value": self.fill_value,
+                "seed": self.seed,
+            }
+        )
+        return config
diff --git a/keras_cv/layers/preprocessing/random_cutout_test.py b/keras_cv/layers/preprocessing/random_cutout_test.py
index 14930b6fd1..b8d549b0d0 100644
--- a/keras_cv/layers/preprocessing/random_cutout_test.py
+++ b/keras_cv/layers/preprocessing/random_cutout_test.py
@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
+import pytest
 import tensorflow as tf
 
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.tests.test_case import TestCase
 
@@ -38,10 +40,10 @@ def _run_test(self, height_factor, width_factor):
         xs = layer(xs)
 
         # Some pixels should be replaced with fill value
-        self.assertTrue(tf.math.reduce_any(xs[0] == fill_value))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == fill_value))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == fill_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == fill_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
     def test_return_shapes(self):
         xs = np.ones((2, 512, 512, 3))
@@ -53,8 +55,8 @@ def test_return_shapes(self):
         xs = layer(xs)
         ys_segmentation_masks = layer(ys_segmentation_masks)
 
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
 
     def test_return_shapes_single_element(self):
         xs = np.ones((512, 512, 3))
@@ -66,8 +68,8 @@ def test_return_shapes_single_element(self):
         xs = layer(xs)
         ys_segmentation_masks = layer(ys_segmentation_masks)
 
-        self.assertEqual(xs.shape, [512, 512, 3])
-        self.assertEqual(ys_segmentation_masks.shape, [512, 512, 3])
+        self.assertEqual(xs.shape, (512, 512, 3))
+        self.assertEqual(ys_segmentation_masks.shape, (512, 512, 3))
 
     def test_random_cutout_single_float(self):
         self._run_test(0.5, 0.5)
@@ -103,10 +105,10 @@ def test_random_cutout_call_results_one_channel(self):
         xs = layer(xs)
 
         # Some pixels should be replaced with fill value
-        self.assertTrue(tf.math.reduce_any(xs[0] == patch_value))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == patch_value))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == patch_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == patch_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
     def test_random_cutout_call_tiny_image(self):
         img_shape = (4, 4, 3)
@@ -127,11 +129,12 @@ def test_random_cutout_call_tiny_image(self):
         xs = layer(xs)
 
         # Some pixels should be replaced with fill value
-        self.assertTrue(tf.math.reduce_any(xs[0] == fill_value))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == fill_value))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == fill_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == fill_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         xs = tf.cast(
             tf.stack(
@@ -156,7 +159,7 @@ def augment(x):
         xs = augment(xs)
 
         # Some pixels should be replaced with fill value
-        self.assertTrue(tf.math.reduce_any(xs[0] == patch_value))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == patch_value))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == patch_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == patch_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
diff --git a/keras_cv/layers/preprocessing/random_flip.py b/keras_cv/layers/preprocessing/random_flip.py
index 927bab485f..fd53f47b12 100644
--- a/keras_cv/layers/preprocessing/random_flip.py
+++ b/keras_cv/layers/preprocessing/random_flip.py
@@ -69,7 +69,7 @@ def __init__(
         bounding_box_format=None,
         **kwargs,
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.mode = mode
         self.seed = seed
         if mode == HORIZONTAL:
@@ -98,12 +98,12 @@ def get_random_transformation_batch(self, batch_size, **kwargs):
         flip_verticals = tf.zeros(shape=(batch_size, 1))
 
         if self.horizontal:
-            flip_horizontals = self._random_generator.random_uniform(
+            flip_horizontals = self._random_generator.uniform(
                 shape=(batch_size, 1)
             )
 
         if self.vertical:
-            flip_verticals = self._random_generator.random_uniform(
+            flip_verticals = self._random_generator.uniform(
                 shape=(batch_size, 1)
             )
 
diff --git a/keras_cv/layers/preprocessing/random_flip_test.py b/keras_cv/layers/preprocessing/random_flip_test.py
index 33d4117997..5b130091e3 100644
--- a/keras_cv/layers/preprocessing/random_flip_test.py
+++ b/keras_cv/layers/preprocessing/random_flip_test.py
@@ -17,6 +17,7 @@
 import tensorflow as tf
 
 from keras_cv import bounding_box
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.random_flip import HORIZONTAL_AND_VERTICAL
 from keras_cv.layers.preprocessing.random_flip import RandomFlip
 from keras_cv.tests.test_case import TestCase
@@ -31,7 +32,7 @@ def test_horizontal_flip(self):
         layer = RandomFlip("horizontal")
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             actual_output = layer(inp)
@@ -57,7 +58,7 @@ def test_vertical_flip(self):
         layer = RandomFlip("vertical")
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             actual_output = layer(inp)
@@ -72,7 +73,7 @@ def test_flip_both(self):
         layer = RandomFlip("horizontal_and_vertical")
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             actual_output = layer(inp)
@@ -85,7 +86,7 @@ def test_random_flip_default(self):
         layer = RandomFlip()
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             actual_output = layer(input_images)
@@ -99,7 +100,7 @@ def test_random_flip_low_rate(self):
         layer = RandomFlip(rate=0.1)
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             actual_output = layer(input_images)
@@ -113,7 +114,7 @@ def test_random_flip_high_rate(self):
         layer = RandomFlip(rate=0.9)
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             actual_output = layer(input_images)
@@ -132,7 +133,7 @@ def test_random_flip_unbatched_image(self):
         layer = RandomFlip("vertical")
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             actual_output = layer(input_image)
@@ -141,9 +142,11 @@ def test_random_flip_unbatched_image(self):
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = RandomFlip()
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+        )
         layer = RandomFlip(dtype="uint8")
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
 
     def test_augment_bounding_box_batched_input(self):
         image = tf.zeros([20, 20, 3])
@@ -170,7 +173,7 @@ def test_augment_bounding_box_batched_input(self):
         )
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             output = layer(input)
@@ -216,7 +219,7 @@ def test_augment_boxes_ragged(self):
         )
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             output = layer(input)
@@ -253,7 +256,7 @@ def test_augment_segmentation_mask(self):
 
         with unittest.mock.patch.object(
             layer._random_generator,
-            "random_uniform",
+            "uniform",
             return_value=mock_random,
         ):
             output = layer(input)
diff --git a/keras_cv/layers/preprocessing/random_gaussian_blur_test.py b/keras_cv/layers/preprocessing/random_gaussian_blur_test.py
index be917133cf..698c65d885 100644
--- a/keras_cv/layers/preprocessing/random_gaussian_blur_test.py
+++ b/keras_cv/layers/preprocessing/random_gaussian_blur_test.py
@@ -28,12 +28,12 @@ def test_return_shapes(self):
         # RGB
         xs = np.ones((2, 512, 512, 3))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
 
         # greyscale
         xs = np.ones((2, 512, 512, 1))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 512, 512, 1])
+        self.assertEqual(xs.shape, (2, 512, 512, 1))
 
     def test_in_single_image(self):
         layer = preprocessing.RandomGaussianBlur(
@@ -47,7 +47,7 @@ def test_in_single_image(self):
         )
 
         xs = layer(xs)
-        self.assertEqual(xs.shape, [512, 512, 3])
+        self.assertEqual(xs.shape, (512, 512, 3))
 
         # greyscale
         xs = tf.cast(
@@ -56,7 +56,7 @@ def test_in_single_image(self):
         )
 
         xs = layer(xs)
-        self.assertEqual(xs.shape, [512, 512, 1])
+        self.assertEqual(xs.shape, (512, 512, 1))
 
     def test_non_square_images(self):
         layer = preprocessing.RandomGaussianBlur(
@@ -66,12 +66,12 @@ def test_non_square_images(self):
         # RGB
         xs = np.ones((2, 256, 512, 3))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 256, 512, 3])
+        self.assertEqual(xs.shape, (2, 256, 512, 3))
 
         # greyscale
         xs = np.ones((2, 256, 512, 1))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 256, 512, 1])
+        self.assertEqual(xs.shape, (2, 256, 512, 1))
 
     def test_single_input_args(self):
         layer = preprocessing.RandomGaussianBlur(kernel_size=7, factor=2)
@@ -79,12 +79,12 @@ def test_single_input_args(self):
         # RGB
         xs = np.ones((2, 512, 512, 3))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
 
         # greyscale
         xs = np.ones((2, 512, 512, 1))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 512, 512, 1])
+        self.assertEqual(xs.shape, (2, 512, 512, 1))
 
     def test_numerical(self):
         layer = preprocessing.RandomGaussianBlur(
diff --git a/keras_cv/layers/preprocessing/random_hue.py b/keras_cv/layers/preprocessing/random_hue.py
index f2f9027362..0161982414 100644
--- a/keras_cv/layers/preprocessing/random_hue.py
+++ b/keras_cv/layers/preprocessing/random_hue.py
@@ -65,9 +65,7 @@ def __init__(self, factor, value_range, seed=None, **kwargs):
         self.seed = seed
 
     def get_random_transformation_batch(self, batch_size, **kwargs):
-        invert = self._random_generator.random_uniform(
-            (batch_size,), 0, 1, tf.float32
-        )
+        invert = self._random_generator.uniform((batch_size,), 0, 1, tf.float32)
         invert = tf.where(
             invert > 0.5, -tf.ones_like(invert), tf.ones_like(invert)
         )
diff --git a/keras_cv/layers/preprocessing/random_hue_test.py b/keras_cv/layers/preprocessing/random_hue_test.py
index cad371f514..9977c7f2ae 100644
--- a/keras_cv/layers/preprocessing/random_hue_test.py
+++ b/keras_cv/layers/preprocessing/random_hue_test.py
@@ -11,10 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import tensorflow as tf
+import numpy as np
 from absl.testing import parameterized
 
 from keras_cv import core
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.tests.test_case import TestCase
 
@@ -22,7 +23,7 @@
 class RandomHueTest(TestCase):
     def test_preserves_output_shape(self):
         image_shape = (4, 8, 8, 3)
-        image = tf.random.uniform(shape=image_shape) * 255.0
+        image = np.random.uniform(size=image_shape) * 255.0
 
         layer = preprocessing.RandomHue(factor=(0.3, 0.8), value_range=(0, 255))
         output = layer(image)
@@ -32,7 +33,7 @@ def test_preserves_output_shape(self):
 
     def test_adjust_no_op(self):
         image_shape = (4, 8, 8, 3)
-        image = tf.random.uniform(shape=image_shape) * 255.0
+        image = np.random.uniform(size=image_shape) * 255.0
 
         layer = preprocessing.RandomHue(factor=(0.0, 0.0), value_range=(0, 255))
         output = layer(image)
@@ -40,24 +41,24 @@ def test_adjust_no_op(self):
 
     def test_adjust_full_opposite_hue(self):
         image_shape = (4, 8, 8, 3)
-        image = tf.random.uniform(shape=image_shape) * 255.0
+        image = np.random.uniform(size=image_shape) * 255.0
 
         layer = preprocessing.RandomHue(factor=(1.0, 1.0), value_range=(0, 255))
-        output = layer(image)
+        output = ops.convert_to_numpy(layer(image))
 
-        channel_max = tf.math.reduce_max(output, axis=-1)
-        channel_min = tf.math.reduce_min(output, axis=-1)
+        channel_max = np.max(output, axis=-1)
+        channel_min = np.min(output, axis=-1)
         # Make sure the max and min channel are the same between input and
         # output. In the meantime, and channel will swap between each other.
         self.assertAllClose(
             channel_max,
-            tf.math.reduce_max(image, axis=-1),
+            np.max(image, axis=-1),
             atol=1e-5,
             rtol=1e-5,
         )
         self.assertAllClose(
             channel_min,
-            tf.math.reduce_min(image, axis=-1),
+            np.min(image, axis=-1),
             atol=1e-5,
             rtol=1e-5,
         )
@@ -68,7 +69,7 @@ def test_adjust_full_opposite_hue(self):
     def test_adjusts_all_values_for_factor(self, factor):
         image_shape = (4, 8, 8, 3)
         # Value range (0, 100)
-        image = tf.random.uniform(shape=image_shape) * 100.0
+        image = np.random.uniform(size=image_shape) * 100.0
 
         layer = preprocessing.RandomHue(
             factor=(factor, factor), value_range=(0, 255)
@@ -79,7 +80,7 @@ def test_adjusts_all_values_for_factor(self, factor):
     def test_adjustment_for_non_rgb_value_range(self):
         image_shape = (4, 8, 8, 3)
         # Value range (0, 100)
-        image = tf.random.uniform(shape=image_shape) * 100.0
+        image = np.random.uniform(size=image_shape) * 100.0
 
         layer = preprocessing.RandomHue(factor=(0.0, 0.0), value_range=(0, 255))
         output = layer(image)
@@ -91,9 +92,7 @@ def test_adjustment_for_non_rgb_value_range(self):
 
     def test_with_uint8(self):
         image_shape = (4, 8, 8, 3)
-        image = tf.cast(
-            tf.random.uniform(shape=image_shape) * 255.0, dtype=tf.uint8
-        )
+        image = (np.random.uniform(size=image_shape) * 255.0).astype(np.uint8)
 
         layer = preprocessing.RandomHue(factor=(0.0, 0.0), value_range=(0, 255))
         output = layer(image)
diff --git a/keras_cv/layers/preprocessing/random_jpeg_quality_test.py b/keras_cv/layers/preprocessing/random_jpeg_quality_test.py
index 69e9b42ebf..c52a6f2e52 100644
--- a/keras_cv/layers/preprocessing/random_jpeg_quality_test.py
+++ b/keras_cv/layers/preprocessing/random_jpeg_quality_test.py
@@ -26,12 +26,12 @@ def test_return_shapes(self):
         # RGB
         xs = np.ones((2, 512, 512, 3))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
 
         # greyscale
         xs = np.ones((2, 512, 512, 1))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 512, 512, 1])
+        self.assertEqual(xs.shape, (2, 512, 512, 1))
 
     def test_in_single_image(self):
         layer = preprocessing.RandomJpegQuality(factor=[0, 100])
@@ -43,7 +43,7 @@ def test_in_single_image(self):
         )
 
         xs = layer(xs)
-        self.assertEqual(xs.shape, [512, 512, 3])
+        self.assertEqual(xs.shape, (512, 512, 3))
 
         # greyscale
         xs = tf.cast(
@@ -52,7 +52,7 @@ def test_in_single_image(self):
         )
 
         xs = layer(xs)
-        self.assertEqual(xs.shape, [512, 512, 1])
+        self.assertEqual(xs.shape, (512, 512, 1))
 
     def test_non_square_images(self):
         layer = preprocessing.RandomJpegQuality(factor=[0, 100])
@@ -60,9 +60,9 @@ def test_non_square_images(self):
         # RGB
         xs = np.ones((2, 256, 512, 3))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 256, 512, 3])
+        self.assertEqual(xs.shape, (2, 256, 512, 3))
 
         # greyscale
         xs = np.ones((2, 256, 512, 1))
         xs = layer(xs)
-        self.assertEqual(xs.shape, [2, 256, 512, 1])
+        self.assertEqual(xs.shape, (2, 256, 512, 1))
diff --git a/keras_cv/layers/preprocessing/random_rotation.py b/keras_cv/layers/preprocessing/random_rotation.py
index 56d30ef59d..15304b7a0a 100644
--- a/keras_cv/layers/preprocessing/random_rotation.py
+++ b/keras_cv/layers/preprocessing/random_rotation.py
@@ -92,7 +92,7 @@ def __init__(
         segmentation_classes=None,
         **kwargs,
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.factor = factor
         if isinstance(factor, (tuple, list)):
             self.lower = factor[0]
@@ -117,7 +117,7 @@ def __init__(
     def get_random_transformation_batch(self, batch_size, **kwargs):
         min_angle = self.lower * 2.0 * np.pi
         max_angle = self.upper * 2.0 * np.pi
-        angles = self._random_generator.random_uniform(
+        angles = self._random_generator.uniform(
             shape=[batch_size], minval=min_angle, maxval=max_angle
         )
         return {"angles": angles}
diff --git a/keras_cv/layers/preprocessing/random_rotation_test.py b/keras_cv/layers/preprocessing/random_rotation_test.py
index 4581a1b5f7..0fa6e59704 100644
--- a/keras_cv/layers/preprocessing/random_rotation_test.py
+++ b/keras_cv/layers/preprocessing/random_rotation_test.py
@@ -15,6 +15,7 @@
 import tensorflow as tf
 
 from keras_cv import bounding_box
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.random_rotation import RandomRotation
 from keras_cv.tests.test_case import TestCase
 
@@ -83,9 +84,11 @@ def test_augment_bounding_boxes(self):
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = RandomRotation(0.5)
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+        )
         layer = RandomRotation(0.5, dtype="uint8")
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
 
     def test_ragged_bounding_boxes(self):
         input_image = tf.random.uniform((2, 512, 512, 3))
@@ -174,7 +177,9 @@ def test_augment_sparse_segmentation_mask(self):
             factor=(0.125, 0.125), segmentation_classes=num_classes
         )
         outputs = layer(inputs)
-        self.assertAllInSet(outputs["segmentation_masks"], [0, 7])
+        self.assertAllInSet(
+            ops.convert_to_numpy(outputs["segmentation_masks"]), [0, 7]
+        )
 
     def test_augment_one_hot_segmentation_mask(self):
         num_classes = 8
diff --git a/keras_cv/layers/preprocessing/random_saturation_test.py b/keras_cv/layers/preprocessing/random_saturation_test.py
index 441a2665f3..652ad7c6f9 100644
--- a/keras_cv/layers/preprocessing/random_saturation_test.py
+++ b/keras_cv/layers/preprocessing/random_saturation_test.py
@@ -17,6 +17,7 @@
 
 from keras_cv import core
 from keras_cv.backend import keras
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
@@ -133,9 +134,9 @@ def test_adjust_to_grayscale(self):
         image = tf.random.uniform(shape=image_shape) * 255.0
 
         layer = preprocessing.RandomSaturation(factor=(0.0, 0.0))
-        output = layer(image)
+        output = ops.convert_to_numpy(layer(image))
 
-        channel_mean = tf.math.reduce_mean(output, axis=-1)
+        channel_mean = np.mean(output, axis=-1)
         channel_values = tf.unstack(output, axis=-1)
         # Make sure all the pixel has the same value among the channel dim,
         # which is a fully gray RGB.
@@ -149,9 +150,9 @@ def test_adjust_to_full_saturation(self):
         image = tf.random.uniform(shape=image_shape) * 255.0
 
         layer = preprocessing.RandomSaturation(factor=(1.0, 1.0))
-        output = layer(image)
+        output = ops.convert_to_numpy(layer(image))
 
-        channel_mean = tf.math.reduce_min(output, axis=-1)
+        channel_mean = np.min(output, axis=-1)
         # Make sure at least one of the channel is 0.0 (fully saturated image)
         self.assertAllClose(channel_mean, np.zeros((4, 8, 8)))
 
diff --git a/keras_cv/layers/preprocessing/random_shear.py b/keras_cv/layers/preprocessing/random_shear.py
index dc69288d70..55635da7e5 100644
--- a/keras_cv/layers/preprocessing/random_shear.py
+++ b/keras_cv/layers/preprocessing/random_shear.py
@@ -16,7 +16,6 @@
 
 import tensorflow as tf
 
-import keras_cv
 from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
 from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
@@ -280,7 +279,7 @@ def augment_bounding_boxes(
                 bounding_boxes, default_value=0
             )
 
-        bounding_boxes = keras_cv.bounding_box.convert_format(
+        bounding_boxes = bounding_box.convert_format(
             bounding_boxes,
             source=self.bounding_box_format,
             target="rel_xyxy",
@@ -337,7 +336,7 @@ def augment_bounding_boxes(
         bounding_boxes = bounding_box.clip_to_image(
             bounding_boxes, images=images, bounding_box_format="rel_xyxy"
         )
-        bounding_boxes = keras_cv.bounding_box.convert_format(
+        bounding_boxes = bounding_box.convert_format(
             bounding_boxes,
             source="rel_xyxy",
             target=self.bounding_box_format,
diff --git a/keras_cv/layers/preprocessing/random_shear_test.py b/keras_cv/layers/preprocessing/random_shear_test.py
index 51933b7f0b..0d2ef32ad8 100644
--- a/keras_cv/layers/preprocessing/random_shear_test.py
+++ b/keras_cv/layers/preprocessing/random_shear_test.py
@@ -11,9 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import numpy as np
+import pytest
 import tensorflow as tf
 
 from keras_cv import bounding_box
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.tests.test_case import TestCase
 
@@ -42,18 +45,22 @@ def test_aggressive_shear_fills_at_least_some_pixels(self):
         ys_segmentation_masks = layer(ys_segmentation_masks)
 
         # Some pixels should be replaced with fill value
-        self.assertTrue(tf.math.reduce_any(xs[0] == fill_value))
-        self.assertTrue(tf.math.reduce_any(xs[0] == 2.0))
-        self.assertTrue(tf.math.reduce_any(xs[1] == fill_value))
-        self.assertTrue(tf.math.reduce_any(xs[1] == 1.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == fill_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[0]) == 2.0))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == fill_value))
+        self.assertTrue(np.any(ops.convert_to_numpy(xs[1]) == 1.0))
         self.assertTrue(
-            tf.math.reduce_any(ys_segmentation_masks[0] == fill_value)
+            np.any(ops.convert_to_numpy(ys_segmentation_masks[0]) == fill_value)
         )
-        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[0] == 2.0))
         self.assertTrue(
-            tf.math.reduce_any(ys_segmentation_masks[1] == fill_value)
+            np.any(ops.convert_to_numpy(ys_segmentation_masks[0]) == 2.0)
+        )
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(ys_segmentation_masks[1]) == fill_value)
+        )
+        self.assertTrue(
+            np.any(ops.convert_to_numpy(ys_segmentation_masks[1]) == 1.0)
         )
-        self.assertTrue(tf.math.reduce_any(ys_segmentation_masks[1] == 1.0))
 
     def test_return_shapes(self):
         """test return dict keys and value pairs"""
@@ -95,11 +102,11 @@ def test_return_shapes(self):
             outputs["segmentation_masks"],
         )
         ys_bounding_boxes = bounding_box.to_dense(ys_bounding_boxes)
-        self.assertEqual(xs.shape, [2, 512, 512, 3])
-        self.assertEqual(ys_labels.shape, [2, 10])
-        self.assertEqual(ys_bounding_boxes["boxes"].shape, [2, 3, 4])
-        self.assertEqual(ys_bounding_boxes["classes"].shape, [2, 3])
-        self.assertEqual(ys_segmentation_masks.shape, [2, 512, 512, 3])
+        self.assertEqual(xs.shape, (2, 512, 512, 3))
+        self.assertEqual(ys_labels.shape, (2, 10))
+        self.assertEqual(ys_bounding_boxes["boxes"].shape, (2, 3, 4))
+        self.assertEqual(ys_bounding_boxes["classes"].shape, (2, 3))
+        self.assertEqual(ys_segmentation_masks.shape, (2, 512, 512, 3))
 
     def test_single_image_input(self):
         """test for single image input"""
@@ -111,8 +118,9 @@ def test_single_image_input(self):
             fill_mode="constant",
         )
         outputs = layer(inputs)
-        self.assertEqual(outputs["images"].shape, [512, 512, 3])
+        self.assertEqual(outputs["images"].shape, (512, 512, 3))
 
+    @pytest.mark.skip(reason="Flaky")
     def test_area(self):
         xs = tf.ones((1, 512, 512, 3))
         ys = {
@@ -153,6 +161,7 @@ def test_area(self):
         )
         self.assertTrue(tf.math.reduce_all(new_area > old_area))
 
+    @pytest.mark.tf_only
     def test_in_tf_function(self):
         """test for class works with tf function"""
         xs = tf.cast(
diff --git a/keras_cv/layers/preprocessing/random_translation.py b/keras_cv/layers/preprocessing/random_translation.py
index 3fcb6d7daf..4340d86064 100644
--- a/keras_cv/layers/preprocessing/random_translation.py
+++ b/keras_cv/layers/preprocessing/random_translation.py
@@ -14,7 +14,7 @@
 
 import tensorflow as tf
 
-import keras_cv
+from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
 from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
     VectorizedBaseImageAugmentationLayer,
@@ -96,7 +96,7 @@ def __init__(
         bounding_box_format=None,
         **kwargs,
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.height_factor = height_factor
         if isinstance(height_factor, (tuple, list)):
             self.height_lower = height_factor[0]
@@ -144,13 +144,13 @@ def __init__(
         self.bounding_box_format = bounding_box_format
 
     def get_random_transformation_batch(self, batch_size, **kwargs):
-        height_translations = self._random_generator.random_uniform(
+        height_translations = self._random_generator.uniform(
             shape=[batch_size, 1],
             minval=self.height_lower,
             maxval=self.height_upper,
             dtype=tf.float32,
         )
-        width_translations = self._random_generator.random_uniform(
+        width_translations = self._random_generator.uniform(
             shape=[batch_size, 1],
             minval=self.width_lower,
             maxval=self.width_upper,
@@ -239,7 +239,7 @@ def augment_bounding_boxes(
                 "Please specify a bounding box format in the constructor. i.e."
                 "`RandomTranslation(bounding_box_format='xyxy')`"
             )
-        bounding_boxes = keras_cv.bounding_box.convert_format(
+        bounding_boxes = bounding_box.convert_format(
             bounding_boxes,
             source=self.bounding_box_format,
             target="rel_xyxy",
@@ -255,16 +255,16 @@ def augment_bounding_boxes(
         y2 += tf.expand_dims(transformations["height_translations"], axis=1)
 
         bounding_boxes["boxes"] = tf.concat([x1, y1, x2, y2], axis=-1)
-        bounding_boxes = keras_cv.bounding_box.to_dense(bounding_boxes)
+        bounding_boxes = bounding_box.to_dense(bounding_boxes)
 
-        bounding_boxes = keras_cv.bounding_box.clip_to_image(
+        bounding_boxes = bounding_box.clip_to_image(
             bounding_boxes,
             bounding_box_format="rel_xyxy",
             images=images,
         )
-        bounding_boxes = keras_cv.bounding_box.to_ragged(bounding_boxes)
+        bounding_boxes = bounding_box.to_ragged(bounding_boxes)
 
-        bounding_boxes = keras_cv.bounding_box.convert_format(
+        bounding_boxes = bounding_box.convert_format(
             bounding_boxes,
             source="rel_xyxy",
             target=self.bounding_box_format,
diff --git a/keras_cv/layers/preprocessing/random_translation_test.py b/keras_cv/layers/preprocessing/random_translation_test.py
index ed98920560..6a2f36a556 100644
--- a/keras_cv/layers/preprocessing/random_translation_test.py
+++ b/keras_cv/layers/preprocessing/random_translation_test.py
@@ -15,6 +15,7 @@
 import numpy as np
 import tensorflow as tf
 
+from keras_cv.backend import ops
 from keras_cv.layers import preprocessing
 from keras_cv.tests.test_case import TestCase
 
@@ -224,6 +225,8 @@ def test_unbatched_image(self):
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = preprocessing.RandomTranslation(0.5, 0.6)
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+        )
         layer = preprocessing.RandomTranslation(0.5, 0.6, dtype="uint8")
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
diff --git a/keras_cv/layers/preprocessing/random_zoom.py b/keras_cv/layers/preprocessing/random_zoom.py
index 7af69a2bbf..089a9922dc 100644
--- a/keras_cv/layers/preprocessing/random_zoom.py
+++ b/keras_cv/layers/preprocessing/random_zoom.py
@@ -100,7 +100,7 @@ def __init__(
         fill_value=0.0,
         **kwargs,
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.height_factor = height_factor
         if isinstance(height_factor, (tuple, list)):
             self.height_lower = height_factor[0]
@@ -140,13 +140,13 @@ def __init__(
         self.seed = seed
 
     def get_random_transformation_batch(self, batch_size, **kwargs):
-        height_zooms = self._random_generator.random_uniform(
+        height_zooms = self._random_generator.uniform(
             shape=[batch_size, 1],
             minval=1.0 + self.height_lower,
             maxval=1.0 + self.height_upper,
         )
         if self.width_factor is not None:
-            width_zooms = self._random_generator.random_uniform(
+            width_zooms = self._random_generator.uniform(
                 shape=[batch_size, 1],
                 minval=1.0 + self.width_lower,
                 maxval=1.0 + self.width_upper,
diff --git a/keras_cv/layers/preprocessing/random_zoom_test.py b/keras_cv/layers/preprocessing/random_zoom_test.py
index 0fdcf6eec3..875402750f 100644
--- a/keras_cv/layers/preprocessing/random_zoom_test.py
+++ b/keras_cv/layers/preprocessing/random_zoom_test.py
@@ -16,6 +16,7 @@
 import tensorflow as tf
 from absl.testing import parameterized
 
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.random_zoom import RandomZoom
 from keras_cv.tests.test_case import TestCase
 
@@ -163,6 +164,8 @@ def test_unbatched_image(self):
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = RandomZoom(0.5, 0.5)
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+        )
         layer = RandomZoom(0.5, 0.5, dtype="uint8")
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
diff --git a/keras_cv/layers/preprocessing/repeated_augmentation_test.py b/keras_cv/layers/preprocessing/repeated_augmentation_test.py
index bb3a30b8d5..6a8da0828d 100644
--- a/keras_cv/layers/preprocessing/repeated_augmentation_test.py
+++ b/keras_cv/layers/preprocessing/repeated_augmentation_test.py
@@ -11,13 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pytest
 import tensorflow as tf
 
 import keras_cv.layers as cv_layers
+from keras_cv.backend.config import keras_3
 from keras_cv.tests.test_case import TestCase
 
 
 class RepeatedAugmentationTest(TestCase):
+    @pytest.mark.skipif(keras_3(), reason="Disabled for Keras 3")
     def test_output_shapes(self):
         repeated_augment = cv_layers.RepeatedAugmentation(
             augmenters=[
@@ -34,6 +37,7 @@ def test_output_shapes(self):
         self.assertEqual(outputs["images"].shape, (16, 512, 512, 3))
         self.assertEqual(outputs["labels"].shape, (16,))
 
+    @pytest.mark.skipif(keras_3(), reason="disabling test for Keras 3")
     def test_with_mix_up(self):
         repeated_augment = cv_layers.RepeatedAugmentation(
             augmenters=[
diff --git a/keras_cv/layers/preprocessing/rescaling.py b/keras_cv/layers/preprocessing/rescaling.py
index 938f23dc75..faeae51f37 100644
--- a/keras_cv/layers/preprocessing/rescaling.py
+++ b/keras_cv/layers/preprocessing/rescaling.py
@@ -55,7 +55,8 @@ class Rescaling(BaseImageAugmentationLayer):
     """
 
     def __init__(self, scale, offset=0.0, **kwargs):
-        super().__init__(**kwargs, autocast=False, force_generator=True)
+        super().__init__(**kwargs)
+
         self.scale = scale
         self.offset = offset
 
diff --git a/keras_cv/layers/preprocessing/rescaling_test.py b/keras_cv/layers/preprocessing/rescaling_test.py
index 0801c227d0..1ef2ed6c2f 100644
--- a/keras_cv/layers/preprocessing/rescaling_test.py
+++ b/keras_cv/layers/preprocessing/rescaling_test.py
@@ -14,6 +14,7 @@
 import numpy as np
 import tensorflow as tf
 
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.rescaling import Rescaling
 from keras_cv.tests.test_case import TestCase
 
@@ -23,14 +24,17 @@ def test_rescaling_correctness_float(self):
         layer = Rescaling(scale=1.0 / 127.5, offset=-1.0)
         inputs = tf.random.uniform((2, 4, 5, 3))
         outputs = layer(inputs)
-        self.assertAllClose(outputs.numpy(), inputs.numpy() * (1.0 / 127.5) - 1)
+        self.assertAllClose(outputs, inputs * (1.0 / 127.5) - 1)
 
     def test_rescaling_correctness_int(self):
         layer = Rescaling(scale=1.0 / 127.5, offset=-1)
         inputs = tf.random.uniform((2, 4, 5, 3), 0, 100, dtype="int32")
         outputs = layer(inputs)
+        outputs = ops.convert_to_numpy(outputs)
         self.assertEqual(outputs.dtype.name, "float32")
-        self.assertAllClose(outputs.numpy(), inputs.numpy() * (1.0 / 127.5) - 1)
+        self.assertAllClose(
+            outputs, ops.convert_to_numpy(inputs) * (1.0 / 127.5) - 1
+        )
 
     def test_config_with_custom_name(self):
         layer = Rescaling(0.5, name="rescaling")
@@ -42,11 +46,15 @@ def test_unbatched_image(self):
         layer = Rescaling(scale=1.0 / 127.5, offset=-1)
         inputs = tf.random.uniform((4, 5, 3))
         outputs = layer(inputs)
-        self.assertAllClose(outputs.numpy(), inputs.numpy() * (1.0 / 127.5) - 1)
+        self.assertAllClose(outputs, inputs * (1.0 / 127.5) - 1)
 
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = Rescaling(0.5)
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype.name, "float32"
+        )
         layer = Rescaling(0.5, dtype="uint8")
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype.name, "uint8"
+        )
diff --git a/keras_cv/layers/preprocessing/resizing.py b/keras_cv/layers/preprocessing/resizing.py
index e52e6ed22e..12804bab02 100644
--- a/keras_cv/layers/preprocessing/resizing.py
+++ b/keras_cv/layers/preprocessing/resizing.py
@@ -14,13 +14,13 @@
 
 import tensorflow as tf
 
-import keras_cv.utils
 from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
-from keras_cv.backend import ops
+from keras_cv.backend import tf_ops
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
 )
+from keras_cv.utils import get_interpolation
 
 H_AXIS = -3
 W_AXIS = -2
@@ -91,9 +91,7 @@ def __init__(
         self.interpolation = interpolation
         self.crop_to_aspect_ratio = crop_to_aspect_ratio
         self.pad_to_aspect_ratio = pad_to_aspect_ratio
-        self._interpolation_method = keras_cv.utils.get_interpolation(
-            interpolation
-        )
+        self._interpolation_method = get_interpolation(interpolation)
         self.bounding_box_format = bounding_box_format
         self.force_output_dense_images = True
 
@@ -199,7 +197,7 @@ def resize_single_with_pad_to_aspect(x):
             img_width = tf.cast(img_size[W_AXIS], self.compute_dtype)
             if bounding_boxes is not None:
                 bounding_boxes = bounding_box.to_dense(bounding_boxes)
-                bounding_boxes = keras_cv.bounding_box.convert_format(
+                bounding_boxes = bounding_box.convert_format(
                     bounding_boxes,
                     image_shape=img_size,
                     source=self.bounding_box_format,
@@ -220,7 +218,7 @@ def resize_single_with_pad_to_aspect(x):
                 method=self._interpolation_method,
             )
             if bounding_boxes is not None:
-                bounding_boxes = keras_cv.bounding_box.convert_format(
+                bounding_boxes = bounding_box.convert_format(
                     bounding_boxes,
                     images=image,
                     source="rel_xyxy",
@@ -230,10 +228,10 @@ def resize_single_with_pad_to_aspect(x):
                 image, 0, 0, self.height, self.width
             )
             if bounding_boxes is not None:
-                bounding_boxes = keras_cv.bounding_box.clip_to_image(
+                bounding_boxes = bounding_box.clip_to_image(
                     bounding_boxes, images=image, bounding_box_format="xyxy"
                 )
-                bounding_boxes = keras_cv.bounding_box.convert_format(
+                bounding_boxes = bounding_box.convert_format(
                     bounding_boxes,
                     images=image,
                     source="xyxy",
@@ -242,7 +240,7 @@ def resize_single_with_pad_to_aspect(x):
             inputs["images"] = image
 
             if bounding_boxes is not None:
-                inputs["bounding_boxes"] = keras_cv.bounding_box.to_ragged(
+                inputs["bounding_boxes"] = bounding_box.to_ragged(
                     bounding_boxes
                 )
 
@@ -311,7 +309,7 @@ def _resize_with_crop(self, inputs):
         def resize_with_crop_to_aspect(x, interpolation_method):
             if isinstance(x, tf.RaggedTensor):
                 x = x.to_tensor()
-            return ops.smart_resize(
+            return tf_ops.smart_resize(
                 x,
                 size=size,
                 interpolation=interpolation_method,
diff --git a/keras_cv/layers/preprocessing/resizing_test.py b/keras_cv/layers/preprocessing/resizing_test.py
index ad9fec9790..c094dee55a 100644
--- a/keras_cv/layers/preprocessing/resizing_test.py
+++ b/keras_cv/layers/preprocessing/resizing_test.py
@@ -17,6 +17,8 @@
 from absl.testing import parameterized
 
 from keras_cv import layers as cv_layers
+from keras_cv.backend import ops
+from keras_cv.backend.config import keras_3
 from keras_cv.tests.test_case import TestCase
 
 
@@ -182,9 +184,11 @@ def test_ragged_image(self, crop_to_aspect_ratio):
     def test_output_dtypes(self):
         inputs = np.array([[[1], [2]], [[3], [4]]], dtype="float64")
         layer = cv_layers.Resizing(2, 2)
-        self.assertAllEqual(layer(inputs).dtype, "float32")
+        self.assertAllEqual(
+            ops.convert_to_numpy(layer(inputs)).dtype, "float32"
+        )
         layer = cv_layers.Resizing(2, 2, dtype="uint8")
-        self.assertAllEqual(layer(inputs).dtype, "uint8")
+        self.assertAllEqual(ops.convert_to_numpy(layer(inputs)).dtype, "uint8")
 
     @parameterized.named_parameters(
         ("batch_crop_to_aspect_ratio", True, False, True),
@@ -194,6 +198,9 @@ def test_output_dtypes(self):
         ("batch_pad_to_aspect_ratio", False, True, True),
         ("single_sample_pad_to_aspect_ratio", False, True, False),
     )
+    @pytest.mark.skipif(
+        keras_3(), reason="ragged tests not yet enabled for keras 3"
+    )
     def test_static_shape_inference(
         self, crop_to_aspect_ratio, pad_to_aspect_ratio, batch
     ):
diff --git a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
index b327ab9785..fd36e22065 100644
--- a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
+++ b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer.py
@@ -12,19 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import keras
 import tensorflow as tf
-
-if hasattr(keras, "src"):
-    keras_backend = keras.src.backend
-else:
-    keras_backend = keras.backend
+import tree
 
 from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import config
 from keras_cv.backend import keras
+from keras_cv.backend import ops
 from keras_cv.backend import scope
-from keras_cv.backend.config import multi_backend
 from keras_cv.utils import preprocessing
 
 H_AXIS = -3
@@ -42,15 +38,8 @@
 USE_TARGETS = "use_targets"
 
 
-base_class = (
-    keras.src.layers.preprocessing.tf_data_layer.TFDataLayer
-    if multi_backend()
-    else keras.layers.Layer
-)
-
-
 @keras_cv_export("keras_cv.layers.VectorizedBaseImageAugmentationLayer")
-class VectorizedBaseImageAugmentationLayer(base_class):
+class VectorizedBaseImageAugmentationLayer(keras.layers.Layer):
     """Abstract base layer for vectorized image augmentation.
 
     This layer contains base functionalities for preprocessing layers which
@@ -109,12 +98,13 @@ def __init__(self):
     """
 
     def __init__(self, seed=None, **kwargs):
-        force_generator = kwargs.pop("force_generator", False)
-        self._random_generator = keras_backend.RandomGenerator(
-            seed=seed, force_generator=force_generator
-        )
         super().__init__(**kwargs)
+        if seed:
+            self._random_generator = tf.random.Generator.from_seed(seed=seed)
+        else:
+            self._random_generator = tf.random.get_global_generator()
         self._convert_input_args = False
+        self._allow_non_tensor_positional_args = True
 
     @property
     def force_output_dense_images(self):
@@ -421,6 +411,21 @@ def _batch_augment(self, inputs):
         return result
 
     def call(self, inputs):
+        # try to convert a given backend native tensor to TensorFlow tensor
+        # before passing it over to TFDataScope
+        is_tf_backend = config.backend() == "tensorflow"
+        is_in_tf_graph = not tf.executing_eagerly()
+        contains_ragged = lambda y: any(
+            tree.map_structure(
+                lambda x: isinstance(x, (tf.RaggedTensor, tf.SparseTensor)),
+                tree.flatten(y),
+            )
+        )
+        inputs_contain_ragged = contains_ragged(inputs)
+        if not is_tf_backend and not inputs_contain_ragged:
+            inputs = tree.map_structure(
+                lambda x: tf.convert_to_tensor(x), inputs
+            )
         with scope.TFDataScope():
             inputs = self._ensure_inputs_are_compute_dtype(inputs)
             inputs, metadata = self._format_inputs(inputs)
@@ -435,7 +440,21 @@ def call(self, inputs):
                     "rank 3 (HWC) or 4D (NHWC) tensors. Got shape: "
                     f"{images.shape}"
                 )
-            return outputs
+        # convert the outputs to backend native tensors if none of them
+        # contain RaggedTensors. Note that if the user passed in Raggeds
+        # but the outputs are dense, we still don't want to convert to
+        # backend native tensors. This is to avoid breaking TF data
+        # pipelines that can't easily be ported to become backend
+        # agnostic.
+        if not is_tf_backend and not is_in_tf_graph:
+            if not inputs_contain_ragged and not contains_ragged(outputs):
+                outputs = tree.map_structure(
+                    # some layers return None, handle that case when
+                    # converting to tensors
+                    lambda x: ops.convert_to_tensor(x) if x is not None else x,
+                    outputs,
+                )
+        return outputs
 
     def _format_inputs(self, inputs):
         metadata = {IS_DICT: True, USE_TARGETS: False}
@@ -443,6 +462,9 @@ def _format_inputs(self, inputs):
             # single image input tensor
             metadata[IS_DICT] = False
             inputs = {IMAGES: inputs}
+        else:
+            # Copy the input dict before we mutate it.
+            inputs = dict(inputs)
 
         metadata[BATCHED] = inputs["images"].shape.rank == 4
         if inputs["images"].shape.rank == 3:
@@ -503,6 +525,8 @@ def _ensure_inputs_are_compute_dtype(self, inputs):
                 inputs,
                 self.compute_dtype,
             )
+        # Copy the input dict before we mutate it.
+        inputs = dict(inputs)
         inputs[IMAGES] = preprocessing.ensure_tensor(
             inputs[IMAGES],
             self.compute_dtype,
diff --git a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py
index 38cff72940..c2d0daa840 100644
--- a/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py
+++ b/keras_cv/layers/preprocessing/vectorized_base_image_augmentation_layer_test.py
@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
+import pytest
 import tensorflow as tf
 
 from keras_cv import bounding_box
+from keras_cv.backend import keras
+from keras_cv.backend import ops
 from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
     VectorizedBaseImageAugmentationLayer,
 )
@@ -33,7 +36,7 @@ def augment_ragged_image(self, image, transformation, **kwargs):
     def get_random_transformation_batch(self, batch_size, **kwargs):
         if self.fixed_value:
             return tf.ones((batch_size,)) * self.fixed_value
-        return self._random_generator.random_uniform(
+        return self._random_generator.uniform(
             (batch_size,), minval=self.add_range[0], maxval=self.add_range[1]
         )
 
@@ -100,7 +103,7 @@ def get_random_transformation_batch(
         assert isinstance(bounding_boxes["classes"], TF_ALL_TENSOR_TYPES)
         assert isinstance(keypoints, TF_ALL_TENSOR_TYPES)
         assert isinstance(segmentation_masks, TF_ALL_TENSOR_TYPES)
-        return self._random_generator.random_uniform((batch_size,))
+        return self._random_generator.uniform((batch_size,))
 
     def augment_images(
         self,
@@ -207,11 +210,11 @@ def test_augment_dict_return_type(self):
 
     def test_augment_casts_dtypes(self):
         add_layer = VectorizedRandomAddLayer(fixed_value=2.0)
-        images = tf.ones((2, 8, 8, 3), dtype="uint8")
+        images = np.ones((2, 8, 8, 3), dtype="uint8")
         output = add_layer(images)
 
         self.assertAllClose(
-            tf.ones((2, 8, 8, 3), dtype="float32") * 3.0, output
+            np.ones((2, 8, 8, 3), dtype="float32") * 3.0, output
         )
 
     def test_augment_batch_images(self):
@@ -219,7 +222,7 @@ def test_augment_batch_images(self):
         images = np.random.random(size=(2, 8, 8, 3)).astype("float32")
         output = add_layer(images)
 
-        diff = output - images
+        diff = ops.convert_to_numpy(output) - images
         # Make sure the first image and second image get different augmentation
         self.assertNotAllClose(diff[0], diff[1])
 
@@ -247,8 +250,8 @@ def test_augment_batch_images_and_targets(self):
         targets = np.random.random(size=(2, 1)).astype("float32")
         output = add_layer({"images": images, "targets": targets})
 
-        image_diff = output["images"] - images
-        label_diff = output["targets"] - targets
+        image_diff = ops.convert_to_numpy(output["images"]) - images
+        label_diff = ops.convert_to_numpy(output["targets"]) - targets
         # Make sure the first image and second image get different augmentation
         self.assertNotAllClose(image_diff[0], image_diff[1])
         self.assertNotAllClose(label_diff[0], label_diff[1])
@@ -356,6 +359,13 @@ def test_augment_batch_image_and_localization_data(self):
             segmentation_mask_diff[0], segmentation_mask_diff[1]
         )
 
+        # the test finishes here for the non-tensorflow backends.
+        if (
+            getattr(keras.config, "backend", lambda: "tensorflow")()
+            != "tensorflow"
+        ):
+            return
+
         @tf.function
         def in_tf_function(inputs):
             return add_layer(inputs)
@@ -382,6 +392,7 @@ def in_tf_function(inputs):
             segmentation_mask_diff[0], segmentation_mask_diff[1]
         )
 
+    @pytest.mark.tf_only
     def test_augment_all_data_in_tf_function(self):
         add_layer = VectorizedRandomAddLayer()
         images = np.random.random(size=(2, 8, 8, 3)).astype("float32")
@@ -442,11 +453,11 @@ def test_augment_unbatched_all_data(self):
         self.assertAllClose(output["keypoints"], keypoints + 2.0)
         self.assertAllClose(
             output["bounding_boxes"]["boxes"],
-            tf.squeeze(bounding_boxes["boxes"]) + 2.0,
+            np.squeeze(bounding_boxes["boxes"]) + 2.0,
         )
         self.assertAllClose(
             output["bounding_boxes"]["classes"],
-            tf.squeeze(bounding_boxes["classes"]) + 2.0,
+            np.squeeze(bounding_boxes["classes"]) + 2.0,
         )
         self.assertAllClose(
             output["segmentation_masks"], segmentation_masks + 2.0
@@ -495,15 +506,6 @@ def test_augment_all_data_with_ragged_images_for_assertion(self):
         segmentation_masks = tf.random.uniform(shape=(2, 8, 8, 1))
         assertion_layer = VectorizedAssertionLayer()
 
-        print(
-            {
-                "images": type(images),
-                "labels": type(labels),
-                "bounding_boxes": type(bounding_boxes),
-                "keypoints": type(keypoints),
-                "segmentation_masks": type(segmentation_masks),
-            }
-        )
         _ = assertion_layer(
             {
                 "images": images,
@@ -547,3 +549,15 @@ def test_converts_ragged_to_dense_segmentation_masks(self):
             {"images": images, "segmentation_masks": segmentation_masks}
         )
         self.assertTrue(isinstance(result["segmentation_masks"], tf.Tensor))
+
+    def test_in_tf_data_pipeline(self):
+        images = np.random.randn(4, 100, 100, 3).astype("float32")
+        train_ds = tf.data.Dataset.from_tensor_slices(images)
+        train_ds = train_ds.map(lambda x: {"images": x})
+        train_ds = train_ds.map(
+            VectorizedRandomAddLayer(fixed_value=2.0)
+        ).batch(4)
+        for output in train_ds.take(1):
+            pass
+        self.assertTrue(isinstance(output["images"], tf.Tensor))
+        self.assertAllClose(output["images"], images + 2.0)
diff --git a/keras_cv/layers/preprocessing_3d/base_augmentation_layer_3d.py b/keras_cv/layers/preprocessing_3d/base_augmentation_layer_3d.py
index 49fa2772d1..073088863c 100644
--- a/keras_cv/layers/preprocessing_3d/base_augmentation_layer_3d.py
+++ b/keras_cv/layers/preprocessing_3d/base_augmentation_layer_3d.py
@@ -13,9 +13,10 @@
 # limitations under the License.
 
 import tensorflow as tf
-from tensorflow import keras
 
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import random
 
 POINT_CLOUDS = "point_clouds"
 BOUNDING_BOXES = "bounding_boxes"
@@ -29,7 +30,7 @@
 
 
 @keras_cv_export("keras_cv.layers.BaseAugmentationLayer3D")
-class BaseAugmentationLayer3D(keras.__internal__.layers.BaseRandomLayer):
+class BaseAugmentationLayer3D(keras.layers.Layer):
     """Abstract base layer for data augmentation for 3D perception.
 
     This layer contains base functionalities for preprocessing layers which
@@ -91,16 +92,13 @@ def augment_pointclouds(self, point_clouds, transformation):
         pointcloud = tf.concat([pointcloud_xyz, pointcloud[..., 3:]], axis=-1)
         return pointcloud, boxes
     ```
-
-    Note that since the randomness is also a common functionality, this layer
-    also includes a keras.backend.RandomGenerator, which can be used to
-    produce the random numbers. The random number generator is stored in the
-    `self._random_generator` attribute.
     """
 
     def __init__(self, seed=None, **kwargs):
-        super().__init__(seed=seed, **kwargs)
+        super().__init__(**kwargs)
         self.auto_vectorize = False
+        self.seed = seed
+        self._random_generator = random.SeedGenerator(seed=self.seed)
 
     @property
     def auto_vectorize(self):
diff --git a/keras_cv/layers/preprocessing_3d/base_augmentation_layer_3d_test.py b/keras_cv/layers/preprocessing_3d/base_augmentation_layer_3d_test.py
index e31e2c7ebc..ee7fe5b177 100644
--- a/keras_cv/layers/preprocessing_3d/base_augmentation_layer_3d_test.py
+++ b/keras_cv/layers/preprocessing_3d/base_augmentation_layer_3d_test.py
@@ -14,6 +14,7 @@
 import numpy as np
 import tensorflow as tf
 
+from keras_cv.backend import random
 from keras_cv.layers.preprocessing_3d import base_augmentation_layer_3d
 from keras_cv.tests.test_case import TestCase
 
@@ -27,14 +28,23 @@ def __init__(self, translate_noise=(0.0, 0.0, 0.0), **kwargs):
         self._translate_noise = translate_noise
 
     def get_random_transformation(self, **kwargs):
-        random_x = self._random_generator.random_normal(
-            (), mean=0.0, stddev=self._translate_noise[0]
+        random_x = random.normal(
+            (),
+            mean=0.0,
+            stddev=self._translate_noise[0],
+            seed=self._random_generator,
         )
-        random_y = self._random_generator.random_normal(
-            (), mean=0.0, stddev=self._translate_noise[1]
+        random_y = random.normal(
+            (),
+            mean=0.0,
+            stddev=self._translate_noise[1],
+            seed=self._random_generator,
         )
-        random_z = self._random_generator.random_normal(
-            (), mean=0.0, stddev=self._translate_noise[2]
+        random_z = random.normal(
+            (),
+            mean=0.0,
+            stddev=self._translate_noise[2],
+            seed=self._random_generator,
         )
 
         return {
diff --git a/keras_cv/layers/preprocessing_3d/input_format_test.py b/keras_cv/layers/preprocessing_3d/input_format_test.py
index 0fe4caf1ab..1eb6744241 100644
--- a/keras_cv/layers/preprocessing_3d/input_format_test.py
+++ b/keras_cv/layers/preprocessing_3d/input_format_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
+import pytest
 import tensorflow as tf
 from absl.testing import parameterized
 
@@ -95,12 +96,18 @@ def convert_to_model_format(inputs):
     }
 
 
+@pytest.mark.skip(
+    reason="values are not matching because of changes to random.py"
+)
 class InputFormatTest(TestCase):
     @parameterized.named_parameters(*TEST_CONFIGURATIONS)
     def test_equivalent_results_with_model_format(self, layer):
         point_clouds = np.random.random(size=(3, 2, 50, 10)).astype("float32")
         bounding_boxes = np.random.random(size=(3, 2, 10, 9)).astype("float32")
-        inputs = {POINT_CLOUDS: point_clouds, BOUNDING_BOXES: bounding_boxes}
+        inputs = {
+            POINT_CLOUDS: point_clouds,
+            BOUNDING_BOXES: bounding_boxes,
+        }
 
         tf.random.set_seed(123)
         outputs_with_legacy_format = convert_to_model_format(layer(inputs))
diff --git a/keras_cv/layers/preprocessing_3d/waymo/frustum_random_dropping_points.py b/keras_cv/layers/preprocessing_3d/waymo/frustum_random_dropping_points.py
index cd2f926fda..cf0a00856b 100644
--- a/keras_cv/layers/preprocessing_3d/waymo/frustum_random_dropping_points.py
+++ b/keras_cv/layers/preprocessing_3d/waymo/frustum_random_dropping_points.py
@@ -6,6 +6,7 @@
 
 from keras_cv import point_cloud
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import random
 from keras_cv.layers.preprocessing_3d import base_augmentation_layer_3d
 
 POINT_CLOUDS = base_augmentation_layer_3d.POINT_CLOUDS
@@ -122,8 +123,11 @@ def get_random_transformation(self, point_clouds, **kwargs):
         frustum_mask = tf.concat(frustum_mask, axis=0)
         # Generate mask along point dimension.
         random_point_mask = (
-            self._random_generator.random_uniform(
-                [1, num_points, 1], minval=0.0, maxval=1
+            random.uniform(
+                [1, num_points, 1],
+                minval=0.0,
+                maxval=1,
+                seed=self._random_generator,
             )
             < self._keep_probability
         )
diff --git a/keras_cv/layers/preprocessing_3d/waymo/global_random_dropping_points.py b/keras_cv/layers/preprocessing_3d/waymo/global_random_dropping_points.py
index 6702789275..5a6c14748b 100644
--- a/keras_cv/layers/preprocessing_3d/waymo/global_random_dropping_points.py
+++ b/keras_cv/layers/preprocessing_3d/waymo/global_random_dropping_points.py
@@ -5,6 +5,7 @@
 import tensorflow as tf
 
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import random
 from keras_cv.layers.preprocessing_3d import base_augmentation_layer_3d
 
 POINT_CLOUDS = base_augmentation_layer_3d.POINT_CLOUDS
@@ -63,9 +64,7 @@ def get_random_transformation(self, point_clouds, **kwargs):
         num_points = point_clouds.get_shape().as_list()[-2]
         # Generate mask along point dimension.
         random_point_mask = (
-            self._random_generator.random_uniform(
-                [1, num_points, 1], minval=0.0, maxval=1
-            )
+            random.uniform([1, num_points, 1], minval=0.0, maxval=1)
             < self._keep_probability
         )
 
diff --git a/keras_cv/layers/preprocessing_3d/waymo/global_random_rotation.py b/keras_cv/layers/preprocessing_3d/waymo/global_random_rotation.py
index 3cdf451cde..83ca1fd4f5 100644
--- a/keras_cv/layers/preprocessing_3d/waymo/global_random_rotation.py
+++ b/keras_cv/layers/preprocessing_3d/waymo/global_random_rotation.py
@@ -5,6 +5,7 @@
 import tensorflow as tf
 
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import random
 from keras_cv.bounding_box_3d import CENTER_XYZ_DXDYDZ_PHI
 from keras_cv.layers.preprocessing_3d import base_augmentation_layer_3d
 from keras_cv.point_cloud import coordinate_transform
@@ -84,23 +85,26 @@ def get_config(self):
         }
 
     def get_random_transformation(self, **kwargs):
-        random_rotation_x = self._random_generator.random_uniform(
+        random_rotation_x = random.uniform(
             (),
             minval=-self._max_rotation_angle_x,
             maxval=self._max_rotation_angle_x,
             dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
-        random_rotation_y = self._random_generator.random_uniform(
+        random_rotation_y = random.uniform(
             (),
             minval=-self._max_rotation_angle_y,
             maxval=self._max_rotation_angle_y,
             dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
-        random_rotation_z = self._random_generator.random_uniform(
+        random_rotation_z = random.uniform(
             (),
             minval=-self._max_rotation_angle_z,
             maxval=self._max_rotation_angle_z,
             dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
         return {
             "pose": tf.stack(
diff --git a/keras_cv/layers/preprocessing_3d/waymo/global_random_scaling.py b/keras_cv/layers/preprocessing_3d/waymo/global_random_scaling.py
index 34ff7315f4..0ddf8edc58 100644
--- a/keras_cv/layers/preprocessing_3d/waymo/global_random_scaling.py
+++ b/keras_cv/layers/preprocessing_3d/waymo/global_random_scaling.py
@@ -5,6 +5,7 @@
 import tensorflow as tf
 
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import random
 from keras_cv.bounding_box_3d import CENTER_XYZ_DXDYDZ_PHI
 from keras_cv.layers.preprocessing_3d import base_augmentation_layer_3d
 
@@ -133,23 +134,26 @@ def get_config(self):
         }
 
     def get_random_transformation(self, **kwargs):
-        random_scaling_x = self._random_generator.random_uniform(
+        random_scaling_x = random.uniform(
             (),
             minval=self._min_x_factor,
             maxval=self._max_x_factor,
             dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
-        random_scaling_y = self._random_generator.random_uniform(
+        random_scaling_y = random.uniform(
             (),
             minval=self._min_y_factor,
             maxval=self._max_y_factor,
             dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
-        random_scaling_z = self._random_generator.random_uniform(
+        random_scaling_z = random.uniform(
             (),
             minval=self._min_z_factor,
             maxval=self._max_z_factor,
             dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
         if not self._preserve_aspect_ratio:
             return {
diff --git a/keras_cv/layers/preprocessing_3d/waymo/global_random_translation.py b/keras_cv/layers/preprocessing_3d/waymo/global_random_translation.py
index 4ece16ec5e..44214f592a 100644
--- a/keras_cv/layers/preprocessing_3d/waymo/global_random_translation.py
+++ b/keras_cv/layers/preprocessing_3d/waymo/global_random_translation.py
@@ -5,6 +5,7 @@
 import tensorflow as tf
 
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import random
 from keras_cv.bounding_box_3d import CENTER_XYZ_DXDYDZ_PHI
 from keras_cv.layers.preprocessing_3d import base_augmentation_layer_3d
 from keras_cv.point_cloud import coordinate_transform
@@ -67,17 +68,26 @@ def get_config(self):
         }
 
     def get_random_transformation(self, **kwargs):
-        random_x_translation = self._random_generator.random_normal(
-            (), mean=0.0, stddev=self._x_stddev, dtype=self.compute_dtype
+        random_x_translation = random.normal(
+            (),
+            mean=0.0,
+            stddev=self._x_stddev,
+            dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
-        random_y_translation = self._random_generator.random_normal(
-            (), mean=0.0, stddev=self._y_stddev, dtype=self.compute_dtype
+        random_y_translation = random.normal(
+            (),
+            mean=0.0,
+            stddev=self._y_stddev,
+            dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
-        random_z_translation = self._random_generator.random_normal(
+        random_z_translation = random.normal(
             (),
             mean=0.0,
             stddev=self._z_stddev,
             dtype=self.compute_dtype,
+            seed=self._random_generator,
         )
         return {
             "pose": tf.stack(
diff --git a/keras_cv/layers/preprocessing_3d/waymo/random_copy_paste.py b/keras_cv/layers/preprocessing_3d/waymo/random_copy_paste.py
index c9b87ec58d..67dc989cc2 100644
--- a/keras_cv/layers/preprocessing_3d/waymo/random_copy_paste.py
+++ b/keras_cv/layers/preprocessing_3d/waymo/random_copy_paste.py
@@ -5,6 +5,7 @@
 import tensorflow as tf
 
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import random
 from keras_cv.bounding_box_3d import CENTER_XYZ_DXDYDZ_PHI
 from keras_cv.layers.preprocessing_3d import base_augmentation_layer_3d
 from keras_cv.ops import iou_3d
@@ -98,10 +99,11 @@ def get_random_transformation(
         **kwargs
     ):
         del point_clouds
-        num_paste_bounding_boxes = self._random_generator.random_uniform(
+        num_paste_bounding_boxes = random.uniform(
             (),
             minval=self._min_paste_bounding_boxes,
             maxval=self._max_paste_bounding_boxes,
+            seed=self._random_generator,
         )
         num_paste_bounding_boxes = tf.cast(
             num_paste_bounding_boxes, dtype=tf.int32
diff --git a/keras_cv/layers/regularization/drop_path.py b/keras_cv/layers/regularization/drop_path.py
index 4475e2365f..c860b22a79 100644
--- a/keras_cv/layers/regularization/drop_path.py
+++ b/keras_cv/layers/regularization/drop_path.py
@@ -50,7 +50,8 @@ class DropPath(keras.layers.Layer):
     def __init__(self, rate=0.5, seed=None, **kwargs):
         super().__init__(**kwargs)
         self.rate = rate
-        self.seed = seed
+        self._seed_val = seed
+        self.seed = random.SeedGenerator(seed=seed)
 
     def call(self, x, training=None):
         if self.rate == 0.0 or not training:
@@ -67,6 +68,6 @@ def call(self, x, training=None):
             return x
 
     def get_config(self):
-        config = {"rate": self.rate, "seed": self.seed}
-        base_config = super().get_config()
-        return dict(list(base_config.items()) + list(config.items()))
+        config = super().get_config()
+        config.update({"rate": self.rate, "seed": self._seed_val})
+        return config
diff --git a/keras_cv/layers/regularization/dropblock_2d.py b/keras_cv/layers/regularization/dropblock_2d.py
index 72c5835fd4..cabd2411fd 100644
--- a/keras_cv/layers/regularization/dropblock_2d.py
+++ b/keras_cv/layers/regularization/dropblock_2d.py
@@ -12,15 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import tensorflow as tf
-from tensorflow.keras.__internal__.layers import BaseRandomLayer
-
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.backend import random
 from keras_cv.utils import conv_utils
 
 
 @keras_cv_export("keras_cv.layers.DropBlock2D")
-class DropBlock2D(BaseRandomLayer):
+class DropBlock2D(keras.layers.Layer):
     """Applies DropBlock regularization to input features.
 
     DropBlock is a form of structured dropout, where units in a contiguous
@@ -145,7 +145,7 @@ def __init__(
         seed=None,
         **kwargs,
     ):
-        super().__init__(seed=seed, **kwargs)
+        super().__init__(**kwargs)
         if not 0.0 <= rate <= 1.0:
             raise ValueError(
                 f"rate must be a number between 0 and 1. " f"Received: {rate}"
@@ -159,77 +159,82 @@ def __init__(
             value=block_size, n=2, name="block_size", allow_zero=False
         )
         self.seed = seed
+        self._random_generator = random.SeedGenerator(self.seed)
 
     def call(self, x, training=None):
         if not training or self._rate == 0.0:
             return x
 
-        _, height, width, _ = tf.split(tf.shape(x), 4)
+        _, height, width, _ = ops.split(ops.shape(x), 4)
 
         # Unnest scalar values
-        height = tf.squeeze(height)
-        width = tf.squeeze(width)
+        height = ops.squeeze(height)
+        width = ops.squeeze(width)
 
-        dropblock_height = tf.math.minimum(self._dropblock_height, height)
-        dropblock_width = tf.math.minimum(self._dropblock_width, width)
+        dropblock_height = ops.minimum(self._dropblock_height, height)
+        dropblock_width = ops.minimum(self._dropblock_width, width)
 
         gamma = (
             self._rate
-            * tf.cast(width * height, dtype=tf.float32)
-            / tf.cast(dropblock_height * dropblock_width, dtype=tf.float32)
-            / tf.cast(
+            * ops.cast(width * height, dtype="float32")
+            / ops.cast(dropblock_height * dropblock_width, dtype="float32")
+            / ops.cast(
                 (width - self._dropblock_width + 1)
                 * (height - self._dropblock_height + 1),
-                tf.float32,
+                "float32",
             )
         )
 
         # Forces the block to be inside the feature map.
-        w_i, h_i = tf.meshgrid(tf.range(width), tf.range(height))
-        valid_block = tf.logical_and(
-            tf.logical_and(
+        w_i, h_i = ops.meshgrid(ops.arange(width), ops.arange(height))
+        valid_block = ops.logical_and(
+            ops.logical_and(
                 w_i >= int(dropblock_width // 2),
                 w_i < width - (dropblock_width - 1) // 2,
             ),
-            tf.logical_and(
+            ops.logical_and(
                 h_i >= int(dropblock_height // 2),
                 h_i < width - (dropblock_height - 1) // 2,
             ),
         )
 
-        valid_block = tf.reshape(valid_block, [1, height, width, 1])
+        valid_block = ops.reshape(valid_block, [1, height, width, 1])
 
-        random_noise = self._random_generator.random_uniform(
-            tf.shape(x), dtype=tf.float32
+        random_noise = random.uniform(
+            ops.shape(x), seed=self._random_generator, dtype="float32"
         )
-        valid_block = tf.cast(valid_block, dtype=tf.float32)
-        seed_keep_rate = tf.cast(1 - gamma, dtype=tf.float32)
+        valid_block = ops.cast(valid_block, dtype="float32")
+        seed_keep_rate = ops.cast(1 - gamma, dtype="float32")
         block_pattern = (1 - valid_block + seed_keep_rate + random_noise) >= 1
-        block_pattern = tf.cast(block_pattern, dtype=tf.float32)
+        block_pattern = ops.cast(block_pattern, dtype="float32")
 
         window_size = [1, self._dropblock_height, self._dropblock_width, 1]
 
         # Double negative and max_pool is essentially min_pooling
-        block_pattern = -tf.nn.max_pool(
+        block_pattern = -ops.max_pool(
             -block_pattern,
-            ksize=window_size,
+            pool_size=window_size,
             strides=[1, 1, 1, 1],
             padding="SAME",
         )
 
         # Slightly scale the values, to account for magnitude change
-        percent_ones = tf.cast(
-            tf.reduce_sum(block_pattern), tf.float32
-        ) / tf.cast(tf.size(block_pattern), tf.float32)
+        percent_ones = ops.cast(ops.sum(block_pattern), "float32") / ops.cast(
+            ops.size(block_pattern), "float32"
+        )
         return (
-            x / tf.cast(percent_ones, x.dtype) * tf.cast(block_pattern, x.dtype)
+            x
+            / ops.cast(percent_ones, x.dtype)
+            * ops.cast(block_pattern, x.dtype)
         )
 
     def get_config(self):
-        config = {
-            "rate": self._rate,
-            "block_size": (self._dropblock_height, self._dropblock_width),
-            "seed": self.seed,
-        }
-        base_config = super().get_config()
-        return dict(list(base_config.items()) + list(config.items()))
+        config = super().get_config()
+        config.update(
+            {
+                "rate": self._rate,
+                "block_size": (self._dropblock_height, self._dropblock_width),
+                "seed": self.seed,
+            }
+        )
+        return config
diff --git a/keras_cv/layers/regularization/dropblock_2d_test.py b/keras_cv/layers/regularization/dropblock_2d_test.py
index edf1a8d72f..122eded38f 100644
--- a/keras_cv/layers/regularization/dropblock_2d_test.py
+++ b/keras_cv/layers/regularization/dropblock_2d_test.py
@@ -84,13 +84,3 @@ def test_input_gets_partially_zeroed_out_with_non_square_block_size(self):
     @staticmethod
     def _count_zeros(tensor: tf.Tensor) -> tf.Tensor:
         return tf.size(tensor) - tf.math.count_nonzero(tensor, dtype=tf.int32)
-
-    def test_works_with_xla(self):
-        dummy_inputs = self.rng.uniform(shape=self.FEATURE_SHAPE)
-        layer = DropBlock2D(rate=0.1, block_size=7)
-
-        @tf.function(jit_compile=True)
-        def apply(x):
-            return layer(x, training=True)
-
-        apply(dummy_inputs)
diff --git a/keras_cv/layers/regularization/squeeze_excite.py b/keras_cv/layers/regularization/squeeze_excite.py
index cb03cc6942..8cbcc5bd94 100644
--- a/keras_cv/layers/regularization/squeeze_excite.py
+++ b/keras_cv/layers/regularization/squeeze_excite.py
@@ -118,10 +118,10 @@ def get_config(self):
     @classmethod
     def from_config(cls, config):
         if isinstance(config["squeeze_activation"], dict):
-            config[
-                "squeeze_activation"
-            ] = keras.saving.deserialize_keras_object(
-                config["squeeze_activation"]
+            config["squeeze_activation"] = (
+                keras.saving.deserialize_keras_object(
+                    config["squeeze_activation"]
+                )
             )
         if isinstance(config["excite_activation"], dict):
             config["excite_activation"] = keras.saving.deserialize_keras_object(
diff --git a/keras_cv/layers/serialization_test.py b/keras_cv/layers/serialization_test.py
index 356e17a89b..6896aeab79 100644
--- a/keras_cv/layers/serialization_test.py
+++ b/keras_cv/layers/serialization_test.py
@@ -18,6 +18,7 @@
 from tensorflow import keras
 
 from keras_cv import layers as cv_layers
+from keras_cv.backend.config import keras_3
 from keras_cv.layers.vit_layers import PatchingAndEmbedding
 from keras_cv.tests.test_case import TestCase
 from keras_cv.utils import test_utils
@@ -376,13 +377,40 @@ class SerializationTest(TestCase):
                 "from_logits": True,
             },
         ),
+        (
+            "NonMaxSuppression",
+            cv_layers.NonMaxSuppression,
+            {
+                "bounding_box_format": "yxyx",
+                "from_logits": True,
+            },
+        ),
     )
     def test_layer_serialization(self, layer_cls, init_args):
+        # TODO: Some layers are not yet compatible with Keras 3.
+        if keras_3:
+            skip_layers = [
+                cv_layers.DropBlock2D,
+                cv_layers.FrustumRandomDroppingPoints,
+                cv_layers.FrustumRandomPointFeatureNoise,
+                cv_layers.GlobalRandomDroppingPoints,
+                cv_layers.GlobalRandomFlip,
+                cv_layers.GlobalRandomRotation,
+                cv_layers.GlobalRandomScaling,
+                cv_layers.GlobalRandomTranslation,
+                cv_layers.GroupPointsByBoundingBoxes,
+                cv_layers.RandomCopyPaste,
+                cv_layers.RandomDropBox,
+                cv_layers.SwapBackground,
+                cv_layers.SqueezeAndExcite2D,  # TODO: Fails in Keras 3
+            ]
+            if layer_cls in skip_layers:
+                self.skipTest("Not supported on Keras 3")
         layer = layer_cls(**init_args)
         config = layer.get_config()
         self.assertAllInitParametersAreInConfig(layer_cls, config)
 
-        model = keras.models.Sequential(layer)
+        model = keras.models.Sequential([layer])
         model_config = model.get_config()
 
         reconstructed_model = keras.Sequential().from_config(model_config)
diff --git a/keras_cv/layers/spatial_pyramid.py b/keras_cv/layers/spatial_pyramid.py
index b45ee7bda3..3d114f7826 100644
--- a/keras_cv/layers/spatial_pyramid.py
+++ b/keras_cv/layers/spatial_pyramid.py
@@ -19,7 +19,7 @@
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 
 @keras_cv_export("keras_cv.layers.SpatialPyramidPooling")
@@ -68,7 +68,7 @@ def __init__(
         self.dropout = dropout
         # TODO(ianstenbit): Remove this once TF 2.14 is released which adds
         # XLA support for resizing with bilinear interpolation.
-        if multi_backend() and keras.backend.backend() == "tensorflow":
+        if keras_3() and keras.backend.backend() == "tensorflow":
             self.supports_jit = False
 
     def build(self, input_shape):
@@ -164,8 +164,12 @@ def call(self, inputs, training=None):
             temp = ops.cast(channel(inputs, training=training), inputs.dtype)
             result.append(temp)
 
+        image_shape = ops.shape(inputs)
+        height, width = image_shape[1], image_shape[2]
         result[-1] = keras.layers.Resizing(
-            inputs.shape[1], inputs.shape[2], interpolation="bilinear"
+            height,
+            width,
+            interpolation="bilinear",
         )(result[-1])
 
         result = ops.concatenate(result, axis=-1)
diff --git a/keras_cv/layers/vit_det_layers.py b/keras_cv/layers/vit_det_layers.py
new file mode 100644
index 0000000000..2e053db4cb
--- /dev/null
+++ b/keras_cv/layers/vit_det_layers.py
@@ -0,0 +1,588 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+
+
+class MLP(keras.layers.Layer):
+    """A MLP block with architecture
+    `input_dim -> [hidden_dim] * (num_layers - 1) -> output_dim`.
+
+    Args:
+        hidden_dim (int): The number of units in the hidden layers.
+        output_dim (int): The number of units in the output layer.
+        num_layers (int): The total number of dense layers to use.
+        activation (str): Activation to use in the hidden layers.
+            Default is `"relu"`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+        - [Detectron2](https://github.com/facebookresearch/detectron2)
+    """  # noqa: E501
+
+    def __init__(
+        self, hidden_dim, output_dim, num_layers, activation="relu", **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.num_layers = num_layers
+        self.activation = activation
+        h = [hidden_dim] * (num_layers - 1)
+        self.dense_net = []
+        for hidden_dim in h:
+            self.dense_net.append(keras.layers.Dense(hidden_dim))
+            self.dense_net.append(keras.layers.Activation(activation))
+        self.dense_net.append(keras.layers.Dense(output_dim))
+        self.dense_net = keras.models.Sequential(self.dense_net)
+
+    def build(self, input_shape):
+        self.dense_net.build(input_shape)
+        self.built = True
+
+    def call(self, x):
+        return self.dense_net(x)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "hidden_dim": self.hidden_dim,
+                "output_dim": self.output_dim,
+                "num_layers": self.num_layers,
+                "activation": self.activation,
+            }
+        )
+        return config
+
+
+@keras_cv_export(
+    "keras_cv.layers.AddRelativePositionalEmbedding", package="keras_cv.layers"
+)
+class AddRelativePositionalEmbedding(keras.layers.Layer):
+    def __init__(self, input_size, key_dim, **kwargs):
+        super().__init__(**kwargs)
+        self.input_size = input_size
+        self.key_dim = key_dim
+        self.rel_pos_h = self.add_weight(
+            name="rel_pos_h",
+            shape=(2 * self.input_size[0] - 1, self.key_dim),
+            initializer="zeros",
+            trainable=True,
+        )
+        self.rel_pos_w = self.add_weight(
+            name="rel_pos_w",
+            shape=(2 * self.input_size[1] - 1, self.key_dim),
+            initializer="zeros",
+            trainable=True,
+        )
+        self.built = True
+
+    def _get_rel_pos(self, query_size, key_size, rel_pos):
+        """
+        Get relative positional embeddings according to the relative positions
+        of query and key sizes.
+
+        Args:
+            query_size (int): The number of features of the queries.
+            key_size (int): The number of features of the keys.
+            rel_pos (tensor): Relative positional embedding tensor.
+
+        Returns:
+            tensor: Extracted positional embeddings according to relative
+                positions.
+        """
+        max_rel_dist = 2 * max(query_size, key_size) - 1
+
+        if ops.shape(rel_pos)[0] != max_rel_dist:
+            rel_pos_resized = ops.image.resize(
+                image=ops.reshape(
+                    rel_pos,
+                    (1, ops.shape(rel_pos)[0], ops.shape(rel_pos)[1], 1),
+                ),
+                size=(max_rel_dist, ops.shape(rel_pos)[1]),
+                interpolation="bilinear",
+            )
+            rel_pos_resized = ops.squeeze(rel_pos_resized, axis=(0, -1))
+            return rel_pos_resized
+        else:
+            rel_pos_resized = rel_pos
+        query_coordinates = ops.cast(
+            ops.arange(query_size), dtype=self.compute_dtype
+        )[:, None] * (max(key_size / query_size, 1.0))
+        key_coordinates = ops.cast(
+            ops.arange(key_size), dtype=self.compute_dtype
+        )[None, :] * (max(query_size / key_size, 1.0))
+        relative_coordinates = (query_coordinates - key_coordinates) + (
+            key_size - 1
+        ) * max(query_size / key_size, 1.0)
+        relative_coordinates = ops.cast(relative_coordinates, dtype="int32")
+        return ops.take(rel_pos_resized, relative_coordinates, 0)
+
+    def call(self, attention_map, queries, query_size, key_size):
+        """
+        Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+
+        Args:
+            attention_map (tensor): Attention map.
+            queries (tensor): Queries in the attention layer with shape
+                `(B, q_h * q_w, C)`.
+            query_size (tuple[int, int]): Spatial sequence size of queries with
+                `(q_h, q_w)`.
+            key_size (tuple[int, int]): Spatial sequence size of keys with
+                `(k_h, k_w)`.
+
+        Returns:
+            tensor: attention map with added relative positional embeddings.
+
+        References:
+            - https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py  # noqa: E501
+        """
+        query_height, query_width = query_size[0], query_size[1]
+        key_height, key_width = key_size[0], key_size[1]
+        rel_heights = self._get_rel_pos(
+            query_height, key_height, self.rel_pos_h
+        )
+        rel_widths = self._get_rel_pos(query_width, key_width, self.rel_pos_w)
+
+        shape = ops.shape(queries)
+        B, C = shape[0], shape[2]
+        rel_queries = ops.reshape(queries, (B, query_height, query_width, C))
+        rel_heights = ops.einsum("bhwc,hkc->bhwk", rel_queries, rel_heights)
+        rel_widths = ops.einsum("bhwc,wkc->bhwk", rel_queries, rel_widths)
+
+        attention_map = ops.reshape(
+            attention_map, (B, query_height, query_width, key_height, key_width)
+        )
+        attention_map = attention_map + rel_heights[..., :, None]
+        attention_map = attention_map + rel_widths[..., None, :]
+        attention_map = ops.reshape(
+            attention_map,
+            (B, query_height * query_width, key_height * key_width),
+        )
+        return attention_map
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"input_size": self.input_size, "key_dim": self.key_dim})
+        return config
+
+
+@keras_cv_export(
+    "keras_cv.layers.MultiHeadAttentionWithRelativePE",
+    package="keras_cv.layers",
+)
+class MultiHeadAttentionWithRelativePE(keras.layers.Layer):
+    """Multi-head Attention block with relative position embeddings.
+
+    Args:
+        num_heads (int): Number of attention heads.
+        key_dim (int): Size of each attention head for query, key, and
+            value.
+        use_bias (bool, optional): Whether to use bias when projecting
+            the queries, keys, and values. Defaults to `True`.
+        use_rel_pos (bool, optional): Whether to use relative positional
+            embeddings or not. Defaults to `False`.
+        input_size (tuple[int, int], optional): Size of the input image.
+            Must be provided when using relative positional embeddings.
+            Defaults to `None`.
+
+    Raises:
+        ValueError: When `input_size = None` with `use_rel_pos = True`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+        - [Detectron2](https://github.com/facebookresearch/detectron2)
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        num_heads,
+        key_dim,
+        use_bias=True,
+        use_rel_pos=False,
+        input_size=None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+        self.scale = self.key_dim**-0.5
+        self.use_bias = use_bias
+        self.input_size = input_size
+        self.use_rel_pos = use_rel_pos
+
+        self.qkv = keras.layers.Dense(
+            key_dim * self.num_heads * 3, use_bias=self.use_bias
+        )
+        self.projection = keras.layers.Dense(key_dim * self.num_heads)
+
+        if self.use_rel_pos:
+            if input_size is None:
+                raise ValueError(
+                    "Input size must be provided if using relative "
+                    "positional encoding."
+                )
+            self.add_decomposed_reative_pe = AddRelativePositionalEmbedding(
+                self.input_size, self.key_dim
+            )
+
+    def build(self, input_shape=None):
+        self.qkv.build([self.key_dim * self.num_heads])
+        self.projection.build([self.key_dim * self.num_heads])
+        self.built = True
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    def call(self, x):
+        shape = ops.shape(x)
+        B, H, W, C = shape[0], shape[1], shape[2], shape[3]
+        qkv = ops.transpose(
+            ops.reshape(
+                self.qkv(x), (B, H * W, 3, self.num_heads, self.key_dim)
+            ),
+            axes=(2, 0, 3, 1, 4),
+        )
+        qkv = ops.reshape(qkv, (3, B * self.num_heads, H * W, self.key_dim))
+        queries, keys, values = ops.unstack(qkv, axis=0)
+        attention_map = (queries * self.scale) @ ops.transpose(
+            keys, axes=(0, 2, 1)
+        )
+
+        if self.use_rel_pos:
+            attention_map = self.add_decomposed_reative_pe(
+                attention_map,
+                queries=queries,
+                query_size=(H, W),
+                key_size=(H, W),
+            )
+        attention_map = ops.softmax(attention_map, axis=-1)
+        x = ops.reshape(
+            attention_map @ values, (B, self.num_heads, H, W, self.key_dim)
+        )
+        x = ops.transpose(x, axes=(0, 2, 3, 1, 4))
+        x = ops.reshape(x, (B, H, W, C))
+        x = self.projection(x)
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "key_dim": self.key_dim,
+                "use_bias": self.use_bias,
+                "use_rel_pos": self.use_rel_pos,
+                "input_size": self.input_size,
+            }
+        )
+        return config
+
+
+@keras_cv_export(
+    "keras_cv.layers.WindowPartitioning", package="keras_cv.layers"
+)
+class WindowPartitioning(keras.layers.Layer):
+    def __init__(self, window_size, **kwargs):
+        super().__init__(**kwargs)
+        self.window_size = window_size
+        self.built = True
+
+    def partition(self, x):
+        shape = ops.shape(x)
+        B, H, W, C = shape[0], shape[1], shape[2], shape[3]
+        pad_height = (
+            self.window_size - H % self.window_size
+        ) % self.window_size
+        pad_width = (self.window_size - W % self.window_size) % self.window_size
+        if pad_height > 0 or pad_width > 0:
+            x = ops.pad(x, ((0, 0), (0, pad_height), (0, pad_width), (0, 0)))
+        H_padded, W_padded = H + pad_height, W + pad_width
+        x = ops.reshape(
+            x,
+            (
+                B,
+                H_padded // self.window_size,
+                self.window_size,
+                W_padded // self.window_size,
+                self.window_size,
+                C,
+            ),
+        )
+        windows = ops.reshape(
+            ops.transpose(x, axes=(0, 1, 3, 2, 4, 5)),
+            (-1, self.window_size, self.window_size, C),
+        )
+        return windows, (H_padded, W_padded)
+
+    def unpartition(self, windows, HW_padded, HW):
+        H_padded, W_padded = HW_padded
+        H, W = HW
+        B = ops.shape(windows)[0] // (
+            (H_padded // self.window_size) * (W_padded // self.window_size)
+        )
+        x = ops.reshape(
+            windows,
+            (
+                B,
+                H_padded // self.window_size,
+                W_padded // self.window_size,
+                self.window_size,
+                self.window_size,
+                -1,
+            ),
+        )
+        x = ops.reshape(
+            ops.transpose(x, axes=(0, 1, 3, 2, 4, 5)),
+            (B, H_padded, W_padded, -1),
+        )
+        return x[:, :H, :W, :]
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"window_size": self.window_size})
+        return config
+
+
+@keras_cv_export(
+    "keras_cv.layers.WindowedTransformerEncoder", package="keras_cv.layers"
+)
+class WindowedTransformerEncoder(keras.layers.Layer):
+    """Transformer blocks with support of window attention and residual
+    propagation blocks.
+
+    Args:
+        project_dim (int): the dimensionality of the projection of the
+            encoder, and output of the `MultiHeadAttention`.
+        mlp_dim (int): the intermediate dimensionality of the MLP head before
+            projecting to `project_dim`.
+        num_heads (int): the number of heads for the `MultiHeadAttention`
+            layer.
+        use_bias (bool, optional): Whether to use bias to project the keys,
+            queries, and values in the attention layer. Defaults to `True`.
+        use_rel_pos (bool, optional): Whether to use relative positional
+            emcodings in the attention layer. Defaults to `False`.
+        window_size (int, optional): Window size for windowed attention.
+            Defaults to `0`.
+        input_size (tuple[int, int], optional): Height and width of the input
+            image as a tuple of integers. Must be provided when using relative
+            positional embeddings. Defaults to `None`.
+        activation (str, optional): the activation function to apply in the
+            MLP head - should be a function. Defaults to `"gelu"`.
+        layer_norm_epsilon (float, optional): The epsilon to use in the layer
+            normalization layers. Defaults to `1e-6`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+        - [Detectron2](https://github.com/facebookresearch/detectron2)
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        project_dim,
+        mlp_dim,
+        num_heads,
+        use_bias=True,
+        use_rel_pos=False,
+        window_size=0,
+        input_size=None,
+        activation="gelu",
+        layer_norm_epsilon=1e-6,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.project_dim = project_dim
+        self.mlp_dim = mlp_dim
+        self.num_heads = num_heads
+        self.use_bias = use_bias
+        self.input_size = input_size
+        self.activation = activation
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.window_size = window_size
+        self.use_rel_pos = use_rel_pos
+
+        self.layer_norm1 = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon
+        )
+        self.layer_norm2 = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon
+        )
+        self.attention = MultiHeadAttentionWithRelativePE(
+            num_heads=self.num_heads,
+            key_dim=self.project_dim // self.num_heads,
+            use_bias=use_bias,
+            use_rel_pos=use_rel_pos,
+            input_size=(
+                input_size if window_size == 0 else (window_size, window_size)
+            ),
+        )
+        self.mlp_block = MLP(
+            mlp_dim,
+            project_dim,
+            num_layers=2,
+            activation="gelu",
+        )
+        self.window_partitioning = WindowPartitioning(window_size)
+
+    def build(self, input_shape=None):
+        self.layer_norm1.build([None, None, None, self.project_dim])
+        self.layer_norm2.build([None, None, None, self.project_dim])
+        self.attention.build()
+        self.mlp_block.build([None, None, None, self.project_dim])
+        self.built = True
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    def call(self, x):
+        shortcut = x
+        x = self.layer_norm1(x)
+        # Window Partition
+        if self.window_size > 0:
+            H, W = ops.shape(x)[1], ops.shape(x)[2]
+            x, HW_padded = self.window_partitioning.partition(x)
+
+        x = self.attention(x)
+        # Reverse Window Partition
+        if self.window_size > 0:
+            x = self.window_partitioning.unpartition(
+                x, HW_padded=HW_padded, HW=(H, W)
+            )
+
+        x = shortcut + x
+        x = x + self.mlp_block(self.layer_norm2(x))
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "project_dim": self.project_dim,
+                "mlp_dim": self.mlp_dim,
+                "num_heads": self.num_heads,
+                "use_bias": self.use_bias,
+                "use_rel_pos": self.use_rel_pos,
+                "window_size": self.window_size,
+                "input_size": self.input_size,
+                "activation": self.activation,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+
+@keras_cv_export(
+    "keras_cv.layers.ViTDetPatchingAndEmbedding", package="keras_cv.layers"
+)
+class ViTDetPatchingAndEmbedding(keras.layers.Layer):
+    """Image to Patch Embedding using only a conv layer (without
+    layer normalization).
+
+    Args:
+        kernel_size (tuple[int, int], optional): Kernel size of the
+            projection layer. Defaults to `(16, 16)`.
+        strides (tuple, optional): Strides of the projection layer.
+            Defaults to `(16, 16)`.
+        embed_dim (int, optional): Number of filters to use in the
+            projection layer i.e. projection size. Defaults to `768`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+        - [Detectron2](https://github.com/facebookresearch/detectron2)
+    """  # noqa: E501
+
+    def __init__(
+        self, kernel_size=(16, 16), strides=(16, 16), embed_dim=768, **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        self.projection = keras.layers.Conv2D(
+            embed_dim, kernel_size=kernel_size, strides=strides
+        )
+
+        self.kernel_size = kernel_size
+        self.strides = strides
+        self.embed_dim = embed_dim
+
+    def build(self, input_shape):
+        self.projection.build(input_shape)
+        self.built = True
+
+    def compute_output_shape(self, input_shape):
+        return self.projection.compute_output_shape(input_shape)
+
+    def call(self, x):
+        x = self.projection(x)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "kernel_size": self.kernel_size,
+                "strides": self.strides,
+                "embed_dim": self.embed_dim,
+            }
+        )
+        return config
+
+
+# TODO: Merge this with the `keras_cv.layers.PatchingAndEmbedding` class once
+# it has been ported to Keras Core.
+@keras_cv_export(
+    "keras_cv.layers.AddPositionalEmbedding", package="keras_cv.layers"
+)
+class AddPositionalEmbedding(keras.layers.Layer):
+    def __init__(self, img_size, patch_size, embed_dim, **kwargs):
+        super().__init__(**kwargs)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.pos_embed = self.add_weight(
+            name="pos_embed",
+            shape=(
+                1,
+                img_size // patch_size,
+                img_size // patch_size,
+                embed_dim,
+            ),
+            initializer="zeros",
+            trainable=True,
+        )
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+    def call(self, x):
+        return x + self.pos_embed
+
+    def get_confg(self):
+        config = super().get_config()
+        config.update(
+            {
+                "img_size": self.img_size,
+                "patch_size": self.patch_size,
+                "embed_dim": self.embed_dim,
+            }
+        )
+        return config
diff --git a/keras_cv/layers/vit_det_layers_test.py b/keras_cv/layers/vit_det_layers_test.py
new file mode 100644
index 0000000000..05c698730e
--- /dev/null
+++ b/keras_cv/layers/vit_det_layers_test.py
@@ -0,0 +1,64 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+from keras_cv.backend import ops
+from keras_cv.layers.vit_det_layers import AddPositionalEmbedding
+from keras_cv.layers.vit_det_layers import MultiHeadAttentionWithRelativePE
+from keras_cv.layers.vit_det_layers import ViTDetPatchingAndEmbedding
+from keras_cv.layers.vit_det_layers import WindowedTransformerEncoder
+from keras_cv.tests.test_case import TestCase
+
+
+class TestViTDetLayers(TestCase):
+    def test_multi_head_attention_with_relative_pe(self):
+        attention_with_rel_pe = MultiHeadAttentionWithRelativePE(
+            num_heads=16,
+            key_dim=1280 // 16,
+            use_bias=True,
+            input_size=(64, 64),
+        )
+        x = np.ones(shape=(1, 64, 64, 1280))
+        x_out = ops.convert_to_numpy(attention_with_rel_pe(x))
+        self.assertEqual(x_out.shape, (1, 64, 64, 1280))
+
+    def test_windowed_transformer_encoder(self):
+        windowed_transformer_encoder = WindowedTransformerEncoder(
+            project_dim=1280,
+            mlp_dim=1280 * 4,
+            num_heads=16,
+            use_bias=True,
+            use_rel_pos=True,
+            window_size=14,
+            input_size=(64, 64),
+        )
+        x = np.ones((1, 64, 64, 1280))
+        x_out = ops.convert_to_numpy(windowed_transformer_encoder(x))
+        self.assertEqual(x_out.shape, (1, 64, 64, 1280))
+        self.assertAllClose(x_out, np.ones_like(x_out))
+
+    def test_vit_patching_and_embedding(self):
+        vit_patching_and_embedding = ViTDetPatchingAndEmbedding()
+        x = np.ones((1, 1024, 1024, 3))
+        x_out = vit_patching_and_embedding(x)
+        self.assertEqual(x_out.shape, (1, 64, 64, 768))
+
+    def test_add_positional_embedding(self):
+        add_positional_embedding = AddPositionalEmbedding(
+            img_size=1024, patch_size=16, embed_dim=256
+        )
+        x = np.ones((1, 64, 64, 256))
+        x_out = add_positional_embedding(x)
+        self.assertEqual(x_out.shape, (1, 64, 64, 256))
diff --git a/keras_cv/losses/focal_test.py b/keras_cv/losses/focal_test.py
index 5a95fc04a9..a8f4463fba 100644
--- a/keras_cv/losses/focal_test.py
+++ b/keras_cv/losses/focal_test.py
@@ -16,14 +16,14 @@
 
 from keras_cv.backend import keras
 from keras_cv.backend import ops
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 from keras_cv.losses import FocalLoss
 from keras_cv.tests.test_case import TestCase
 
 
 class FocalTest(TestCase):
     def test_output_shape(self):
-        y_true = np.random.uniform(size=[2, 5], low=0, high=2)
+        y_true = np.random.uniform(size=[2, 5], low=0, high=1)
         y_pred = np.random.uniform(size=[2, 5], low=0, high=1)
 
         focal_loss = FocalLoss(reduction="sum")
@@ -31,7 +31,7 @@ def test_output_shape(self):
         self.assertAllEqual(focal_loss(y_true, y_pred).shape, [])
 
     def test_output_shape_reduction_none(self):
-        y_true = np.random.uniform(size=[2, 5], low=0, high=2)
+        y_true = np.random.uniform(size=[2, 5], low=0, high=1)
         y_pred = np.random.uniform(size=[2, 5], low=0, high=1)
 
         focal_loss = FocalLoss(reduction="none")
@@ -44,7 +44,7 @@ def test_output_shape_reduction_none(self):
         )
 
     def test_output_shape_from_logits(self):
-        y_true = np.random.uniform(size=[2, 5], low=0, high=2)
+        y_true = np.random.uniform(size=[2, 5], low=0, high=1)
         y_pred = np.random.uniform(size=[2, 5], low=-10, high=10)
 
         focal_loss = FocalLoss(reduction="none", from_logits=True)
@@ -73,7 +73,7 @@ def test_from_logits_argument(self):
         # TF backend somehow has different numerics.
         expected_loss = (
             31.11176
-            if multi_backend() and keras.backend.backend() != "tensorflow"
+            if keras_3() and keras.backend.backend() != "tensorflow"
             else 925.28081
         )
         self.assertAllClose(
diff --git a/keras_cv/losses/numerical_tests/focal_loss_numerical_test.py b/keras_cv/losses/numerical_tests/focal_loss_numerical_test.py
index ded529fedf..e9f1088ece 100644
--- a/keras_cv/losses/numerical_tests/focal_loss_numerical_test.py
+++ b/keras_cv/losses/numerical_tests/focal_loss_numerical_test.py
@@ -17,6 +17,7 @@
 from absl.testing import parameterized
 from tensorflow import keras
 
+from keras_cv.backend import ops
 from keras_cv.losses import FocalLoss
 from keras_cv.tests.test_case import TestCase
 
@@ -31,8 +32,8 @@ def __init__(
 
     def call(self, y_true, y_pred):
         with tf.name_scope("focal_loss"):
-            y_true = tf.cast(y_true, dtype=tf.float32)
-            y_pred = tf.cast(y_pred, dtype=tf.float32)
+            y_true = tf.cast(ops.convert_to_numpy(y_true), dtype=tf.float32)
+            y_pred = tf.cast(ops.convert_to_numpy(y_pred), dtype=tf.float32)
             positive_label_mask = tf.equal(y_true, 1.0)
             cross_entropy = tf.nn.sigmoid_cross_entropy_with_logits(
                 labels=y_true, logits=y_pred
diff --git a/keras_cv/metrics/coco/pycoco_wrapper.py b/keras_cv/metrics/coco/pycoco_wrapper.py
index 3c09784388..659cdef0a0 100644
--- a/keras_cv/metrics/coco/pycoco_wrapper.py
+++ b/keras_cv/metrics/coco/pycoco_wrapper.py
@@ -125,6 +125,9 @@ def _convert_predictions_to_coco_annotations(predictions):
     num_batches = len(predictions["source_id"])
     for i in range(num_batches):
         batch_size = predictions["source_id"][i].shape[0]
+        predictions["detection_boxes"][i] = predictions["detection_boxes"][
+            i
+        ].copy()
         for j in range(batch_size):
             max_num_detections = predictions["num_detections"][i][j]
             predictions["detection_boxes"][i][j] = _yxyx_to_xywh(
diff --git a/keras_cv/metrics/object_detection/box_coco_metrics.py b/keras_cv/metrics/object_detection/box_coco_metrics.py
index b293ae4fbb..47d86ba1c2 100644
--- a/keras_cv/metrics/object_detection/box_coco_metrics.py
+++ b/keras_cv/metrics/object_detection/box_coco_metrics.py
@@ -19,10 +19,10 @@
 import tensorflow as tf
 import tensorflow.keras as keras
 
-import keras_cv
 from keras_cv import bounding_box
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import ops
+from keras_cv.metrics import coco
 
 
 class HidePrints:
@@ -139,6 +139,8 @@ class BoxCOCOMetrics(keras.metrics.Metric):
     """
 
     def __init__(self, bounding_box_format, evaluate_freq, name=None, **kwargs):
+        if "dtype" not in kwargs:
+            kwargs["dtype"] = "float32"
         super().__init__(name=name, **kwargs)
         self.ground_truths = []
         self.predictions = []
@@ -210,9 +212,9 @@ def result_fn(self, force=False):
             )
             result = {}
             for i, key in enumerate(METRIC_NAMES):
-                result[
-                    self.name_prefix() + METRIC_MAPPING[key]
-                ] = py_func_result[i]
+                result[self.name_prefix() + METRIC_MAPPING[key]] = (
+                    py_func_result[i]
+                )
             return result
 
         obj.result = types.MethodType(result_fn, obj)
@@ -310,6 +312,4 @@ def compute_pycocotools_metric(y_true, y_pred, bounding_box_format):
         ops.sum(ops.cast(confidence_pred > 0, "int32"), axis=-1)
     ]
 
-    return keras_cv.metrics.coco.compute_pycoco_metrics(
-        ground_truth, predictions
-    )
+    return coco.compute_pycoco_metrics(ground_truth, predictions)
diff --git a/keras_cv/models/__init__.py b/keras_cv/models/__init__.py
index 9c83a3891a..b9b90b946a 100644
--- a/keras_cv/models/__init__.py
+++ b/keras_cv/models/__init__.py
@@ -178,6 +178,10 @@
 from keras_cv.models.backbones.resnet_v2.resnet_v2_backbone import (
     ResNetV2Backbone,
 )
+from keras_cv.models.backbones.vit_det.vit_det_aliases import ViTDetBBackbone
+from keras_cv.models.backbones.vit_det.vit_det_aliases import ViTDetHBackbone
+from keras_cv.models.backbones.vit_det.vit_det_aliases import ViTDetLBackbone
+from keras_cv.models.backbones.vit_det.vit_det_backbone import ViTDetBackbone
 from keras_cv.models.classification.image_classifier import ImageClassifier
 from keras_cv.models.object_detection.retinanet.retinanet import RetinaNet
 from keras_cv.models.object_detection.yolo_v8.yolo_v8_backbone import (
@@ -186,7 +190,12 @@
 from keras_cv.models.object_detection.yolo_v8.yolo_v8_detector import (
     YOLOV8Detector,
 )
+from keras_cv.models.segmentation import BASNet
 from keras_cv.models.segmentation import DeepLabV3Plus
+from keras_cv.models.segmentation import SAMMaskDecoder
+from keras_cv.models.segmentation import SAMPromptEncoder
+from keras_cv.models.segmentation import SegmentAnythingModel
+from keras_cv.models.segmentation import TwoWayTransformer
 from keras_cv.models.segmentation.segformer.segformer_aliases import SegFormer
 from keras_cv.models.segmentation.segformer.segformer_aliases import SegFormerB0
 from keras_cv.models.segmentation.segformer.segformer_aliases import SegFormerB1
diff --git a/keras_cv/models/backbones/backbone.py b/keras_cv/models/backbones/backbone.py
index 5e302b78b3..6130945308 100644
--- a/keras_cv/models/backbones/backbone.py
+++ b/keras_cv/models/backbones/backbone.py
@@ -13,10 +13,10 @@
 # limitations under the License.
 """Base class for Backbone models."""
 
-import os
-
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
+from keras_cv.utils.preset_utils import check_preset_class
+from keras_cv.utils.preset_utils import load_from_preset
 from keras_cv.utils.python_utils import classproperty
 from keras_cv.utils.python_utils import format_docstring
 
@@ -32,6 +32,20 @@ class Backbone(keras.Model):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._pyramid_level_inputs = {}
+        self._functional_layer_ids = set(
+            id(layer) for layer in self._flatten_layers()
+        )
+
+    def __dir__(self):
+        # Temporary fixes for weight saving. This mimics the following PR for
+        # older version of Keras: https://github.com/keras-team/keras/pull/18982
+        def filter_fn(attr):
+            try:
+                return id(getattr(self, attr)) not in self._functional_layer_ids
+            except:
+                return True
+
+        return filter(filter_fn, super().__dir__())
 
     def get_config(self):
         # Don't chain to super here. The default `get_config()` for functional
@@ -95,43 +109,18 @@ def from_preset(
             load_weights=False,
         ```
         """
-
-        if not cls.presets:
-            raise NotImplementedError(
-                "No presets have been created for this class."
-            )
-
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
-            )
-
-        if load_weights and preset not in cls.presets_with_weights:
-            raise ValueError(
-                f"""Pretrained weights not available for preset "{preset}". """
-                "Set `load_weights=False` to use this preset or choose one of "
-                "the following presets with weights:"
-                f""" "{'", "'.join(cls.presets_with_weights)}"."""
-            )
-
-        metadata = cls.presets[preset]
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if preset not in cls.presets_with_weights or load_weights is False:
-            return model
-
-        weights = keras.utils.get_file(
-            "model.h5",
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
+        # We support short IDs for official presets, e.g. `"bert_base_en"`.
+        # Map these to a Kaggle Models handle.
+        if preset in cls.presets:
+            preset = cls.presets[preset]["kaggle_handle"]
+
+        check_preset_class(preset, cls)
+        return load_from_preset(
+            preset,
+            load_weights=load_weights,
+            config_overrides=kwargs,
         )
 
-        model.load_weights(weights)
-        return model
-
     def __init_subclass__(cls, **kwargs):
         # Use __init_subclass__ to set up a correct docstring for from_preset.
         super().__init_subclass__(**kwargs)
diff --git a/keras_cv/models/backbones/backbone_presets.py b/keras_cv/models/backbones/backbone_presets.py
index 614f85cd24..95d3ccd522 100644
--- a/keras_cv/models/backbones/backbone_presets.py
+++ b/keras_cv/models/backbones/backbone_presets.py
@@ -28,6 +28,7 @@
 from keras_cv.models.backbones.mobilenet_v3 import mobilenet_v3_backbone_presets
 from keras_cv.models.backbones.resnet_v1 import resnet_v1_backbone_presets
 from keras_cv.models.backbones.resnet_v2 import resnet_v2_backbone_presets
+from keras_cv.models.backbones.vit_det import vit_det_backbone_presets
 from keras_cv.models.object_detection.yolo_v8 import yolo_v8_backbone_presets
 
 backbone_presets_no_weights = {
@@ -40,6 +41,7 @@
     **densenet_backbone_presets.backbone_presets_no_weights,
     **efficientnet_lite_backbone_presets.backbone_presets_no_weights,
     **yolo_v8_backbone_presets.backbone_presets_no_weights,
+    **vit_det_backbone_presets.backbone_presets_no_weights,
 }
 
 backbone_presets_with_weights = {
@@ -52,6 +54,7 @@
     **densenet_backbone_presets.backbone_presets_with_weights,
     **efficientnet_lite_backbone_presets.backbone_presets_with_weights,
     **yolo_v8_backbone_presets.backbone_presets_with_weights,
+    **vit_det_backbone_presets.backbone_presets_with_weights,
 }
 
 backbone_presets = {
diff --git a/keras_cv/models/backbones/csp_darknet/csp_darknet_backbone.py b/keras_cv/models/backbones/csp_darknet/csp_darknet_backbone.py
index fdb0870771..70fb629eca 100644
--- a/keras_cv/models/backbones/csp_darknet/csp_darknet_backbone.py
+++ b/keras_cv/models/backbones/csp_darknet/csp_darknet_backbone.py
@@ -41,7 +41,7 @@
 from keras_cv.utils.python_utils import classproperty
 
 
-@keras_cv_export("keras_cv.losses.CSPDarkNetBackbone")
+@keras_cv_export("keras_cv.models.CSPDarkNetBackbone")
 class CSPDarkNetBackbone(Backbone):
     """This class represents the CSPDarkNet architecture.
 
diff --git a/keras_cv/models/backbones/csp_darknet/csp_darknet_backbone_presets.py b/keras_cv/models/backbones/csp_darknet/csp_darknet_backbone_presets.py
index c4d4b0f2ef..6edfd02d4d 100644
--- a/keras_cv/models/backbones/csp_darknet/csp_darknet_backbone_presets.py
+++ b/keras_cv/models/backbones/csp_darknet/csp_darknet_backbone_presets.py
@@ -26,15 +26,7 @@
             "official_name": "CSPDarkNet",
             "path": "csp_darknet",
         },
-        "class_name": "keras_cv>CSPDarkNetBackbone",
-        "config": {
-            "stackwise_channels": [48, 96, 192, 384],
-            "stackwise_depth": [1, 3, 3, 1],
-            "include_rescaling": True,
-            "use_depthwise": False,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/cspdarknet/keras/csp_darknet_tiny/2",
     },
     "csp_darknet_s": {
         "metadata": {
@@ -47,15 +39,7 @@
             "official_name": "CSPDarkNet",
             "path": "csp_darknet",
         },
-        "class_name": "keras_cv>CSPDarkNetBackbone",
-        "config": {
-            "stackwise_channels": [64, 128, 256, 512],
-            "stackwise_depth": [1, 3, 3, 1],
-            "include_rescaling": True,
-            "use_depthwise": False,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/cspdarknet/keras/csp_darknet_s/2",
     },
     "csp_darknet_m": {
         "metadata": {
@@ -68,15 +52,7 @@
             "official_name": "CSPDarkNet",
             "path": "csp_darknet",
         },
-        "class_name": "keras_cv>CSPDarkNetBackbone",
-        "config": {
-            "stackwise_channels": [96, 192, 384, 768],
-            "stackwise_depth": [2, 6, 6, 2],
-            "include_rescaling": True,
-            "use_depthwise": False,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/cspdarknet/keras/csp_darknet_m/2",
     },
     "csp_darknet_l": {
         "metadata": {
@@ -89,15 +65,7 @@
             "official_name": "CSPDarkNet",
             "path": "csp_darknet",
         },
-        "class_name": "keras_cv>CSPDarkNetBackbone",
-        "config": {
-            "stackwise_channels": [128, 256, 512, 1024],
-            "stackwise_depth": [3, 9, 9, 3],
-            "include_rescaling": True,
-            "use_depthwise": False,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/cspdarknet/keras/csp_darknet_l/2",
     },
     "csp_darknet_xl": {
         "metadata": {
@@ -110,15 +78,7 @@
             "official_name": "CSPDarkNet",
             "path": "csp_darknet",
         },
-        "class_name": "keras_cv>CSPDarkNetBackbone",
-        "config": {
-            "stackwise_channels": [170, 340, 680, 1360],
-            "stackwise_depth": [4, 12, 12, 4],
-            "include_rescaling": True,
-            "use_depthwise": False,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/cspdarknet/keras/csp_darknet_xl/2",
     },
 }
 
@@ -135,10 +95,7 @@
             "official_name": "CSPDarkNet",
             "path": "csp_darknet",
         },
-        "class_name": "keras_cv>CSPDarkNetBackbone",
-        "config": backbone_presets_no_weights["csp_darknet_tiny"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/cspdarknettiny/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "0007ae82c95be4d4aef06368a7c38e006381324d77e5df029b04890e18a8ad19",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/cspdarknet/keras/csp_darknet_tiny_imagenet/2",  # noqa: E501
     },
     "csp_darknet_l_imagenet": {
         "metadata": {
@@ -152,10 +109,7 @@
             "official_name": "CSPDarkNet",
             "path": "csp_darknet",
         },
-        "class_name": "keras_cv>CSPDarkNetBackbone",
-        "config": backbone_presets_no_weights["csp_darknet_l"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/cspdarknetl/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "9303aabfadffbff8447171fce1e941f96d230d8f3cef30d3f05a9c85097f8f1e",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/cspdarknet/keras/csp_darknet_l_imagenet/2",  # noqa: E501
     },
 }
 
diff --git a/keras_cv/models/backbones/csp_darknet/csp_darknet_utils.py b/keras_cv/models/backbones/csp_darknet/csp_darknet_utils.py
index 304557cd7b..1be85aa976 100644
--- a/keras_cv/models/backbones/csp_darknet/csp_darknet_utils.py
+++ b/keras_cv/models/backbones/csp_darknet/csp_darknet_utils.py
@@ -310,7 +310,8 @@ def __init__(
                 )
             )
 
-        self.add = keras.layers.Add()
+        if self.residual:
+            self.add = keras.layers.Add()
         self.concatenate = keras.layers.Concatenate()
 
         self.darknet_conv3 = DarknetConvBlock(
diff --git a/keras_cv/models/backbones/densenet/densenet_backbone.py b/keras_cv/models/backbones/densenet/densenet_backbone.py
index 28109b64fa..251f3601ec 100644
--- a/keras_cv/models/backbones/densenet/densenet_backbone.py
+++ b/keras_cv/models/backbones/densenet/densenet_backbone.py
@@ -119,9 +119,9 @@ def __init__(
             name=f"conv{len(stackwise_num_repeats) + 1}",
         )
 
-        pyramid_level_inputs[
-            f"P{len(stackwise_num_repeats) + 1}"
-        ] = utils.get_tensor_input_name(x)
+        pyramid_level_inputs[f"P{len(stackwise_num_repeats) + 1}"] = (
+            utils.get_tensor_input_name(x)
+        )
         x = keras.layers.BatchNormalization(
             axis=BN_AXIS, epsilon=BN_EPSILON, name="bn"
         )(x)
diff --git a/keras_cv/models/backbones/densenet/densenet_backbone_presets.py b/keras_cv/models/backbones/densenet/densenet_backbone_presets.py
index 50230b8d89..7668016ac0 100644
--- a/keras_cv/models/backbones/densenet/densenet_backbone_presets.py
+++ b/keras_cv/models/backbones/densenet/densenet_backbone_presets.py
@@ -18,43 +18,19 @@
         "metadata": {
             "description": "DenseNet model with 121 layers.",
         },
-        "class_name": "keras_cv>DenseNetBackbone",
-        "config": {
-            "stackwise_num_repeats": [6, 12, 24, 16],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "compression_ratio": 0.5,
-            "growth_rate": 32,
-        },
+        "kaggle_handle": "kaggle://keras/densenet/keras/densenet121/2",
     },
     "densenet169": {
         "metadata": {
             "description": "DenseNet model with 169 layers.",
         },
-        "class_name": "keras_cv>DenseNetBackbone",
-        "config": {
-            "stackwise_num_repeats": [6, 12, 32, 32],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "compression_ratio": 0.5,
-            "growth_rate": 32,
-        },
+        "kaggle_handle": "kaggle://keras/densenet/keras/densenet169/2",
     },
     "densenet201": {
         "metadata": {
             "description": "DenseNet model with 201 layers.",
         },
-        "class_name": "keras_cv>DenseNetBackbone",
-        "config": {
-            "stackwise_num_repeats": [6, 12, 48, 32],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "compression_ratio": 0.5,
-            "growth_rate": 32,
-        },
+        "kaggle_handle": "kaggle://keras/densenet/keras/densenet201/2",
     },
 }
 
@@ -66,10 +42,7 @@
                 "classification task."
             ),
         },
-        "class_name": "keras_cv>DenseNetBackbone",
-        "config": backbone_presets_no_weights["densenet121"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/densenet121/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "709afe0321d9f2b2562e562ff9d0dc44cca10ed09e0e2cfba08d783ff4dab6bf",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/densenet/keras/densenet121_imagenet/2",
     },
     "densenet169_imagenet": {
         "metadata": {
@@ -78,10 +51,7 @@
                 "classification task."
             ),
         },
-        "class_name": "keras_cv>DenseNetBackbone",
-        "config": backbone_presets_no_weights["densenet169"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/densenet169/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "a99d1bb2cbe1a59a1cdd1f435fb265453a97c2a7b723d26f4ebee96e5fb49d62",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/densenet/keras/densenet169_imagenet/2",
     },
     "densenet201_imagenet": {
         "metadata": {
@@ -90,10 +60,7 @@
                 "classification task."
             ),
         },
-        "class_name": "keras_cv>DenseNetBackbone",
-        "config": backbone_presets_no_weights["densenet201"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/densenet201/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "c1189a934f12c1a676a9cf52238e5994401af925e2adfc0365bad8133c052060",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/densenet/keras/densenet201_imagenet/2",
     },
 }
 
diff --git a/keras_cv/models/backbones/efficientnet_lite/efficientnet_lite_backbone_presets.py b/keras_cv/models/backbones/efficientnet_lite/efficientnet_lite_backbone_presets.py
index db9838e3de..5405742151 100644
--- a/keras_cv/models/backbones/efficientnet_lite/efficientnet_lite_backbone_presets.py
+++ b/keras_cv/models/backbones/efficientnet_lite/efficientnet_lite_backbone_presets.py
@@ -26,24 +26,7 @@
             "official_name": "EfficientNetLite",
             "path": "EfficientNetLite",
         },
-        "class_name": "keras_cv.models>EfficientNetLiteBackbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.0,
-            "dropout_rate": 0.2,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "relu6",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetlite_b0",
     },
     "efficientnetlite_b1": {
         "metadata": {
@@ -56,24 +39,7 @@
             "official_name": "EfficientNetLite",
             "path": "EfficientNetLite",
         },
-        "class_name": "keras_cv.models>EfficientNetLiteBackbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.1,
-            "dropout_rate": 0.2,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "relu6",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetlite_b1",
     },
     "efficientnetlite_b2": {
         "metadata": {
@@ -86,24 +52,7 @@
             "official_name": "EfficientNetLite",
             "path": "EfficientNetLite",
         },
-        "class_name": "keras_cv.models>EfficientNetLiteBackbone",
-        "config": {
-            "width_coefficient": 1.1,
-            "depth_coefficient": 1.2,
-            "dropout_rate": 0.3,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "relu6",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetlite_b2",
     },
     "efficientnetlite_b3": {
         "metadata": {
@@ -116,24 +65,7 @@
             "official_name": "EfficientNetLite",
             "path": "EfficientNetLite",
         },
-        "class_name": "keras_cv.models>EfficientNetLiteBackbone",
-        "config": {
-            "width_coefficient": 1.2,
-            "depth_coefficient": 1.4,
-            "dropout_rate": 0.3,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "relu6",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetlite_b3",
     },
     "efficientnetlite_b4": {
         "metadata": {
@@ -146,24 +78,7 @@
             "official_name": "EfficientNetLite",
             "path": "EfficientNetLite",
         },
-        "class_name": "keras_cv.models>EfficientNetLiteBackbone",
-        "config": {
-            "width_coefficient": 1.4,
-            "depth_coefficient": 1.8,
-            "dropout_rate": 0.3,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "relu6",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetlite_b4",
     },
 }
 
diff --git a/keras_cv/models/backbones/efficientnet_v1/efficientnet_v1_backbone_presets.py b/keras_cv/models/backbones/efficientnet_v1/efficientnet_v1_backbone_presets.py
index a2aac81d26..8d0e0c6f7d 100644
--- a/keras_cv/models/backbones/efficientnet_v1/efficientnet_v1_backbone_presets.py
+++ b/keras_cv/models/backbones/efficientnet_v1/efficientnet_v1_backbone_presets.py
@@ -26,33 +26,7 @@
             "official_name": "EfficientNetV1",
             "path": "efficientnetv1",
         },
-        "class_name": "keras_cv.models>EfficientNetV1Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.0,
-            "dropout_rate": 0.2,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "swish",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetv1_b0",
     },
     "efficientnetv1_b1": {
         "metadata": {
@@ -65,33 +39,7 @@
             "official_name": "EfficientNetV1",
             "path": "efficientnetv1",
         },
-        "class_name": "keras_cv.models>EfficientNetV1Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.1,
-            "dropout_rate": 0.2,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "swish",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetv1_b1",
     },
     "efficientnetv1_b2": {
         "metadata": {
@@ -104,33 +52,7 @@
             "official_name": "EfficientNetV1",
             "path": "efficientnetv1",
         },
-        "class_name": "keras_cv.models>EfficientNetV1Backbone",
-        "config": {
-            "width_coefficient": 1.1,
-            "depth_coefficient": 1.2,
-            "dropout_rate": 0.3,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "swish",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetv1_b2",
     },
     "efficientnetv1_b3": {
         "metadata": {
@@ -143,33 +65,7 @@
             "official_name": "EfficientNetV1",
             "path": "efficientnetv1",
         },
-        "class_name": "keras_cv.models>EfficientNetV1Backbone",
-        "config": {
-            "width_coefficient": 1.2,
-            "depth_coefficient": 1.4,
-            "dropout_rate": 0.3,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "swish",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetv1_b3",
     },
     "efficientnetv1_b4": {
         "metadata": {
@@ -182,33 +78,7 @@
             "official_name": "EfficientNetV1",
             "path": "efficientnetv1",
         },
-        "class_name": "keras_cv.models>EfficientNetV1Backbone",
-        "config": {
-            "width_coefficient": 1.4,
-            "depth_coefficient": 1.8,
-            "dropout_rate": 0.4,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "swish",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetv1_b4",
     },
     "efficientnetv1_b5": {
         "metadata": {
@@ -221,33 +91,7 @@
             "official_name": "EfficientNetV1",
             "path": "efficientnetv1",
         },
-        "class_name": "keras_cv.models>EfficientNetV1Backbone",
-        "config": {
-            "width_coefficient": 1.6,
-            "depth_coefficient": 2.2,
-            "dropout_rate": 0.4,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "swish",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetv1_b5",
     },
     "efficientnetv1_b6": {
         "metadata": {
@@ -260,33 +104,7 @@
             "official_name": "EfficientNetV1",
             "path": "efficientnetv1",
         },
-        "class_name": "keras_cv.models>EfficientNetV1Backbone",
-        "config": {
-            "width_coefficient": 1.8,
-            "depth_coefficient": 2.6,
-            "dropout_rate": 0.5,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "swish",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetv1_b6",
     },
     "efficientnetv1_b7": {
         "metadata": {
@@ -299,33 +117,7 @@
             "official_name": "EfficientNetV1",
             "path": "efficientnetv1",
         },
-        "class_name": "keras_cv.models>EfficientNetV1Backbone",
-        "config": {
-            "width_coefficient": 2.0,
-            "depth_coefficient": 3.1,
-            "dropout_rate": 0.5,
-            "drop_connect_rate": 0.2,
-            "depth_divisor": 8,
-            "stackwise_kernel_sizes": [3, 3, 5, 3, 5, 5, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 3, 4, 1],
-            "stackwise_input_filters": [32, 16, 24, 40, 80, 112, 192],
-            "stackwise_output_filters": [16, 24, 40, 80, 112, 192, 320],
-            "stackwise_expansion_ratios": [1, 6, 6, 6, 6, 6, 6],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "activation": "swish",
-        },
+        "kaggle_handle": "gs://keras-cv-kaggle/efficientnetv1_b7",
     },
 }
 
diff --git a/keras_cv/models/backbones/efficientnet_v2/efficientnet_v2_backbone.py b/keras_cv/models/backbones/efficientnet_v2/efficientnet_v2_backbone.py
index cb836f88d4..1160184789 100644
--- a/keras_cv/models/backbones/efficientnet_v2/efficientnet_v2_backbone.py
+++ b/keras_cv/models/backbones/efficientnet_v2/efficientnet_v2_backbone.py
@@ -69,7 +69,7 @@ class EfficientNetV2Backbone(Backbone):
         input_tensor: optional Keras tensor (i.e. output of `keras.layers.Input()`)
             to use as image input for the model.
 
-    Usage:
+    Example:
     ```python
     # Construct an EfficientNetV2 from a preset:
     efficientnet = keras_cv.models.EfficientNetV2Backbone.from_preset(
diff --git a/keras_cv/models/backbones/efficientnet_v2/efficientnet_v2_backbone_presets.py b/keras_cv/models/backbones/efficientnet_v2/efficientnet_v2_backbone_presets.py
index 4bdc099c13..7052d2b258 100644
--- a/keras_cv/models/backbones/efficientnet_v2/efficientnet_v2_backbone_presets.py
+++ b/keras_cv/models/backbones/efficientnet_v2/efficientnet_v2_backbone_presets.py
@@ -24,40 +24,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.0,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [2, 4, 4, 6, 9, 15],
-            "stackwise_input_filters": [24, 24, 48, 64, 128, 160],
-            "stackwise_output_filters": [24, 48, 64, 128, 160, 256],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.0,
-                0.0,
-                0,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_s/2",  # noqa: E501
     },
     "efficientnetv2_m": {
         "metadata": {
@@ -68,42 +35,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.0,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [3, 5, 5, 7, 14, 18, 5],
-            "stackwise_input_filters": [24, 24, 48, 80, 160, 176, 304],
-            "stackwise_output_filters": [24, 48, 80, 160, 176, 304, 512],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [
-                0,
-                0,
-                0,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_m/2",  # noqa: E501
     },
     "efficientnetv2_l": {
         "metadata": {
@@ -115,42 +47,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.0,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [4, 7, 7, 10, 19, 25, 7],
-            "stackwise_input_filters": [32, 32, 64, 96, 192, 224, 384],
-            "stackwise_output_filters": [32, 64, 96, 192, 224, 384, 640],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [
-                0,
-                0,
-                0,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2, 1],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_l/2",  # noqa: E501
     },
     "efficientnetv2_b0": {
         "metadata": {
@@ -163,33 +60,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.0,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 5, 8],
-            "stackwise_input_filters": [32, 16, 32, 48, 96, 112],
-            "stackwise_output_filters": [16, 32, 48, 96, 112, 192],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [0, 0, 0, 0.25, 0.25, 0.25],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b0/2",  # noqa: E501
     },
     "efficientnetv2_b1": {
         "metadata": {
@@ -202,33 +73,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.1,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 5, 8],
-            "stackwise_input_filters": [32, 16, 32, 48, 96, 112],
-            "stackwise_output_filters": [16, 32, 48, 96, 112, 192],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [0, 0, 0, 0.25, 0.25, 0.25],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b1/2",  # noqa: E501
     },
     "efficientnetv2_b2": {
         "metadata": {
@@ -241,33 +86,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.1,
-            "depth_coefficient": 1.2,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 5, 8],
-            "stackwise_input_filters": [32, 16, 32, 48, 96, 112],
-            "stackwise_output_filters": [16, 32, 48, 96, 112, 192],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [0, 0, 0, 0.25, 0.25, 0.25],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b2/2",  # noqa: E501
     },
     "efficientnetv2_b3": {
         "metadata": {
@@ -280,33 +99,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.2,
-            "depth_coefficient": 1.4,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 5, 8],
-            "stackwise_input_filters": [32, 16, 32, 48, 96, 112],
-            "stackwise_output_filters": [16, 32, 48, 96, 112, 192],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [0, 0, 0, 0.25, 0.25, 0.25],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b3/2",  # noqa: E501
     },
 }
 
@@ -324,42 +117,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.0,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [2, 4, 4, 6, 9, 15],
-            "stackwise_input_filters": [24, 24, 48, 64, 128, 160],
-            "stackwise_output_filters": [24, 48, 64, 128, 160, 256],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [
-                0.0,
-                0.0,
-                0,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/efficientnetv2s/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "80555436ea49100893552614b4dce98de461fa3b6c14f8132673817d28c83654",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_s_imagenet/2",  # noqa: E501
     },
     "efficientnetv2_b0_imagenet": {
         "metadata": {
@@ -376,35 +134,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.0,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 5, 8],
-            "stackwise_input_filters": [32, 16, 32, 48, 96, 112],
-            "stackwise_output_filters": [16, 32, 48, 96, 112, 192],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [0, 0, 0, 0.25, 0.25, 0.25],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/efficientnetv2b0/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "ac95f13a8ad1cee41184fc16fd0eb769f7c5b3131151c6abf7fcee5cc3d09bc8",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b0_imagenet/2",  # noqa: E501
     },
     "efficientnetv2_b1_imagenet": {
         "metadata": {
@@ -421,35 +151,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.0,
-            "depth_coefficient": 1.1,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 5, 8],
-            "stackwise_input_filters": [32, 16, 32, 48, 96, 112],
-            "stackwise_output_filters": [16, 32, 48, 96, 112, 192],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [0, 0, 0, 0.25, 0.25, 0.25],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/efficientnetv2b1/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "82da111f8411f47e3f5eef090da76340f38e222f90a08bead53662f2ebafb01c",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b1_imagenet/2",  # noqa: E501
     },
     "efficientnetv2_b2_imagenet": {
         "metadata": {
@@ -466,35 +168,7 @@
             "official_name": "EfficientNetV2",
             "path": "efficientnetv2",
         },
-        "class_name": "keras_cv>EfficientNetV2Backbone",
-        "config": {
-            "width_coefficient": 1.1,
-            "depth_coefficient": 1.2,
-            "skip_connection_dropout": 0.2,
-            "depth_divisor": 8,
-            "min_depth": 8,
-            "activation": "swish",
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "stackwise_kernel_sizes": [3, 3, 3, 3, 3, 3],
-            "stackwise_num_repeats": [1, 2, 2, 3, 5, 8],
-            "stackwise_input_filters": [32, 16, 32, 48, 96, 112],
-            "stackwise_output_filters": [16, 32, 48, 96, 112, 192],
-            "stackwise_expansion_ratios": [1, 4, 4, 4, 6, 6],
-            "stackwise_squeeze_and_excite_ratios": [0, 0, 0, 0.25, 0.25, 0.25],
-            "stackwise_strides": [1, 2, 2, 2, 1, 2],
-            "stackwise_conv_types": [
-                "fused",
-                "fused",
-                "fused",
-                "unfused",
-                "unfused",
-                "unfused",
-            ],
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/efficientnetv2b2/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "02d12c9d1589b540b4e84ffdb54ff30c96099bd59e311a85ddc7180efc65e955",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b2_imagenet/2",  # noqa: E501
     },
 }
 
diff --git a/keras_cv/models/backbones/mix_transformer/mix_transformer_backbone_presets.py b/keras_cv/models/backbones/mix_transformer/mix_transformer_backbone_presets.py
index a4c1c2a3e1..542d6662d6 100644
--- a/keras_cv/models/backbones/mix_transformer/mix_transformer_backbone_presets.py
+++ b/keras_cv/models/backbones/mix_transformer/mix_transformer_backbone_presets.py
@@ -23,14 +23,7 @@
             "official_name": "MiT",
             "path": "mit",
         },
-        "class_name": "keras_cv>MiTBackbone",
-        "config": {
-            "embedding_dims": [32, 64, 160, 256],
-            "depths": [2, 2, 2, 2],
-            "include_rescaling": True,
-            "input_shape": (224, 224, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/mit/keras/mit_b0/2",
     },
     "mit_b1": {
         "metadata": {
@@ -41,14 +34,7 @@
             "official_name": "MiT",
             "path": "mit",
         },
-        "class_name": "keras_cv>MiTBackbone",
-        "config": {
-            "embedding_dims": [64, 128, 320, 512],
-            "depths": [2, 2, 2, 2],
-            "include_rescaling": True,
-            "input_shape": (224, 224, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/mit/keras/mit_b1/2",
     },
     "mit_b2": {
         "metadata": {
@@ -59,14 +45,7 @@
             "official_name": "MiT",
             "path": "mit",
         },
-        "class_name": "keras_cv>MiTBackbone",
-        "config": {
-            "embedding_dims": [64, 128, 320, 512],
-            "depths": [3, 4, 6, 3],
-            "include_rescaling": True,
-            "input_shape": (224, 224, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/mit/keras/mit_b2/2",
     },
     "mit_b3": {
         "metadata": {
@@ -77,14 +56,7 @@
             "official_name": "MiT",
             "path": "mit",
         },
-        "class_name": "keras_cv>MiTBackbone",
-        "config": {
-            "embedding_dims": [64, 128, 320, 512],
-            "depths": [3, 4, 18, 3],
-            "include_rescaling": True,
-            "input_shape": (224, 224, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/mit/keras/mit_b3/2",
     },
     "mit_b4": {
         "metadata": {
@@ -95,14 +67,7 @@
             "official_name": "MiT",
             "path": "mit",
         },
-        "class_name": "keras_cv>MiTBackbone",
-        "config": {
-            "embedding_dims": [64, 128, 320, 512],
-            "depths": [3, 8, 27, 3],
-            "include_rescaling": True,
-            "input_shape": (224, 224, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/mit/keras/mit_b4/2",
     },
     "mit_b5": {
         "metadata": {
@@ -113,14 +78,7 @@
             "official_name": "MiT",
             "path": "mit",
         },
-        "class_name": "keras_cv>MiTBackbone",
-        "config": {
-            "embedding_dims": [64, 128, 320, 512],
-            "depths": [3, 6, 40, 3],
-            "include_rescaling": True,
-            "input_shape": (224, 224, 3),
-            "input_tensor": None,
-        },
+        "kaggle_handle": "kaggle://keras/mit/keras/mit_b5/2",
     },
 }
 
@@ -134,16 +92,7 @@
             "official_name": "MiT",
             "path": "mit",
         },
-        "class_name": "keras_cv>MiTBackbone",
-        "config": {
-            "embedding_dims": [32, 64, 160, 256],
-            "depths": [2, 2, 2, 2],
-            "include_rescaling": True,
-            "input_shape": (224, 224, 3),
-            "input_tensor": None,
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/mitb0/imagenet/classification-v0.h5",  # noqa: E501
-        "weights_hash": "8e0c416cd330b6fa0bcfb3a5ccc43edcbcabf6a463aee3c2a9b6a1398c207d10",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/mit/keras/mit_b0_imagenet/2",
     },
 }
 
diff --git a/keras_cv/models/backbones/mobilenet_v3/mobilenet_v3_backbone_presets.py b/keras_cv/models/backbones/mobilenet_v3/mobilenet_v3_backbone_presets.py
index 75ad436fe1..460a1875a7 100644
--- a/keras_cv/models/backbones/mobilenet_v3/mobilenet_v3_backbone_presets.py
+++ b/keras_cv/models/backbones/mobilenet_v3/mobilenet_v3_backbone_presets.py
@@ -25,55 +25,7 @@
             "official_name": "MobileNetV3",
             "path": "mobilenetv3",
         },
-        "class_name": "keras_cv>MobileNetV3Backbone",
-        "config": {
-            "stackwise_expansion": [
-                1,
-                72.0 / 16,
-                88.0 / 24,
-                4,
-                6,
-                6,
-                3,
-                3,
-                6,
-                6,
-                6,
-            ],
-            "stackwise_filters": [16, 24, 24, 40, 40, 40, 48, 48, 96, 96, 96],
-            "stackwise_kernel_size": [3, 3, 3, 5, 5, 5, 5, 5, 5, 5, 5],
-            "stackwise_stride": [2, 2, 1, 2, 1, 1, 1, 1, 2, 1, 1],
-            "stackwise_se_ratio": [
-                0.25,
-                None,
-                None,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "stackwise_activation": [
-                "relu",
-                "relu",
-                "relu",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "alpha": 1.0,
-        },
+        "kaggle_handle": "kaggle://keras/mobilenetv3/keras/mobilenet_v3_small/2",  # noqa: E501
     },
     "mobilenet_v3_large": {
         "metadata": {
@@ -86,99 +38,7 @@
             "official_name": "MobileNetV3",
             "path": "mobilenetv3",
         },
-        "class_name": "keras_cv>MobileNetV3Backbone",
-        "config": {
-            "stackwise_expansion": [
-                1,
-                4,
-                3,
-                3,
-                3,
-                3,
-                6,
-                2.5,
-                2.3,
-                2.3,
-                6,
-                6,
-                6,
-                6,
-                6,
-            ],
-            "stackwise_filters": [
-                16,
-                24,
-                24,
-                40,
-                40,
-                40,
-                80,
-                80,
-                80,
-                80,
-                112,
-                112,
-                160,
-                160,
-                160,
-            ],
-            "stackwise_kernel_size": [
-                3,
-                3,
-                3,
-                5,
-                5,
-                5,
-                3,
-                3,
-                3,
-                3,
-                3,
-                3,
-                5,
-                5,
-                5,
-            ],
-            "stackwise_stride": [1, 2, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 2, 1, 1],
-            "stackwise_se_ratio": [
-                None,
-                None,
-                None,
-                0.25,
-                0.25,
-                0.25,
-                None,
-                None,
-                None,
-                None,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-                0.25,
-            ],
-            "stackwise_activation": [
-                "relu",
-                "relu",
-                "relu",
-                "relu",
-                "relu",
-                "relu",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-                "hard_swish",
-            ],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "alpha": 1.0,
-        },
+        "kaggle_handle": "kaggle://keras/mobilenetv3/keras/mobilenet_v3_large/2",  # noqa: E501
     },
 }
 
@@ -195,10 +55,7 @@
             "official_name": "MobileNetV3",
             "path": "mobilenetv3",
         },
-        "class_name": "keras_cv>MobileNetV3Backbone",
-        "config": backbone_presets_no_weights["mobilenet_v3_large"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/mobilenetv3/mobilenetv3_large_imagenet_backbone.h5",  # noqa: E501
-        "weights_hash": "ec55ea2f4f4ee9a2ddf3ee8e2dd784e9d5732690c1fc5afc7e1b2a66703f3337",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/mobilenetv3/keras/mobilenet_v3_large_imagenet/2",  # noqa: E501
     },
     "mobilenet_v3_small_imagenet": {
         "metadata": {
@@ -212,10 +69,7 @@
             "official_name": "MobileNetV3",
             "path": "mobilenetv3",
         },
-        "class_name": "keras_cv>MobileNetV3Backbone",
-        "config": backbone_presets_no_weights["mobilenet_v3_small"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/mobilenetv3/mobilenetv3_small_imagenet_backbone.h5",  # noqa: E501
-        "weights_hash": "592c2707edfc6c673a3b2d9aaf76dee678557f4a32d573c74f96c8122effa503",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/mobilenetv3/keras/mobilenet_v3_small_imagenet/2",  # noqa: E501
     },
 }
 
diff --git a/keras_cv/models/backbones/resnet_v1/resnet_v1_backbone.py b/keras_cv/models/backbones/resnet_v1/resnet_v1_backbone.py
index 61046234d3..07c896613c 100644
--- a/keras_cv/models/backbones/resnet_v1/resnet_v1_backbone.py
+++ b/keras_cv/models/backbones/resnet_v1/resnet_v1_backbone.py
@@ -130,9 +130,9 @@ def __init__(
                 first_shortcut=(block_type == "block" or stack_index > 0),
                 name=f"v2_stack_{stack_index}",
             )
-            pyramid_level_inputs[
-                f"P{stack_index + 2}"
-            ] = utils.get_tensor_input_name(x)
+            pyramid_level_inputs[f"P{stack_index + 2}"] = (
+                utils.get_tensor_input_name(x)
+            )
 
         # Create model.
         super().__init__(inputs=inputs, outputs=x, **kwargs)
diff --git a/keras_cv/models/backbones/resnet_v1/resnet_v1_backbone_presets.py b/keras_cv/models/backbones/resnet_v1/resnet_v1_backbone_presets.py
index 370a31cf03..eb312ebbbc 100644
--- a/keras_cv/models/backbones/resnet_v1/resnet_v1_backbone_presets.py
+++ b/keras_cv/models/backbones/resnet_v1/resnet_v1_backbone_presets.py
@@ -25,16 +25,7 @@
             "official_name": "ResNetV1",
             "path": "resnet_v1",
         },
-        "class_name": "keras_cv>ResNetBackbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [2, 2, 2, 2],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "basic_block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv1/keras/resnet18/2",
     },
     "resnet34": {
         "metadata": {
@@ -47,16 +38,7 @@
             "official_name": "ResNetV1",
             "path": "resnet_v1",
         },
-        "class_name": "keras_cv>ResNetBackbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [3, 4, 6, 3],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "basic_block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv1/keras/resnet34/2",
     },
     "resnet50": {
         "metadata": {
@@ -69,16 +51,7 @@
             "official_name": "ResNetV1",
             "path": "resnet_v1",
         },
-        "class_name": "keras_cv>ResNetBackbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [3, 4, 6, 3],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv1/keras/resnet50/2",
     },
     "resnet101": {
         "metadata": {
@@ -91,16 +64,7 @@
             "official_name": "ResNetV1",
             "path": "resnet_v1",
         },
-        "class_name": "keras_cv>ResNetBackbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [3, 4, 23, 3],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv1/keras/resnet101/2",
     },
     "resnet152": {
         "metadata": {
@@ -113,16 +77,7 @@
             "official_name": "ResNetV1",
             "path": "resnet_v1",
         },
-        "class_name": "keras_cv>ResNetBackbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [3, 8, 36, 3],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv1/keras/resnet152/2",
     },
 }
 
@@ -139,10 +94,7 @@
             "official_name": "ResNetV1",
             "path": "resnet_v1",
         },
-        "class_name": "keras_cv>ResNetBackbone",
-        "config": backbone_presets_no_weights["resnet50"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/resnet50/imagenet/classification-v0-notop.h5",  # noqa: E501
-        "weights_hash": "dc5f6d8f929c78d0fc192afecc67b11ac2166e9d8b9ef945742368ae254c07af",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/resnetv1/keras/resnet50_imagenet/2",
     },
 }
 
diff --git a/keras_cv/models/backbones/resnet_v2/resnet_v2_backbone.py b/keras_cv/models/backbones/resnet_v2/resnet_v2_backbone.py
index a31841f7fc..6a0cc74740 100644
--- a/keras_cv/models/backbones/resnet_v2/resnet_v2_backbone.py
+++ b/keras_cv/models/backbones/resnet_v2/resnet_v2_backbone.py
@@ -136,9 +136,9 @@ def __init__(
                 first_shortcut=(block_type == "block" or stack_index > 0),
                 name=f"v2_stack_{stack_index}",
             )
-            pyramid_level_inputs[
-                f"P{stack_index + 2}"
-            ] = utils.get_tensor_input_name(x)
+            pyramid_level_inputs[f"P{stack_index + 2}"] = (
+                utils.get_tensor_input_name(x)
+            )
 
         x = keras.layers.BatchNormalization(
             axis=BN_AXIS, epsilon=BN_EPSILON, name="post_bn"
diff --git a/keras_cv/models/backbones/resnet_v2/resnet_v2_backbone_presets.py b/keras_cv/models/backbones/resnet_v2/resnet_v2_backbone_presets.py
index 815ff5ee39..a351f027b8 100644
--- a/keras_cv/models/backbones/resnet_v2/resnet_v2_backbone_presets.py
+++ b/keras_cv/models/backbones/resnet_v2/resnet_v2_backbone_presets.py
@@ -24,17 +24,7 @@
             "official_name": "ResNetV2",
             "path": "resnet_v2",
         },
-        "class_name": "keras_cv>ResNetV2Backbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [2, 2, 2, 2],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "stackwise_dilations": None,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "basic_block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv2/keras/resnet18_v2/2",
     },
     "resnet34_v2": {
         "metadata": {
@@ -46,17 +36,7 @@
             "official_name": "ResNetV2",
             "path": "resnet_v2",
         },
-        "class_name": "keras_cv>ResNetV2Backbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [3, 4, 6, 3],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "stackwise_dilations": None,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "basic_block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv2/keras/resnet34_v2/2",
     },
     "resnet50_v2": {
         "metadata": {
@@ -68,17 +48,7 @@
             "official_name": "ResNetV2",
             "path": "resnet_v2",
         },
-        "class_name": "keras_cv>ResNetV2Backbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [3, 4, 6, 3],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "stackwise_dilations": None,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv2/keras/resnet50_v2/2",
     },
     "resnet101_v2": {
         "metadata": {
@@ -90,17 +60,7 @@
             "official_name": "ResNetV2",
             "path": "resnet_v2",
         },
-        "class_name": "keras_cv>ResNetV2Backbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [3, 4, 23, 3],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "stackwise_dilations": None,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv2/keras/resnet101_v2/2",
     },
     "resnet152_v2": {
         "metadata": {
@@ -112,17 +72,7 @@
             "official_name": "ResNetV2",
             "path": "resnet_v2",
         },
-        "class_name": "keras_cv>ResNetV2Backbone",
-        "config": {
-            "stackwise_filters": [64, 128, 256, 512],
-            "stackwise_blocks": [3, 8, 36, 3],
-            "stackwise_strides": [1, 2, 2, 2],
-            "include_rescaling": True,
-            "stackwise_dilations": None,
-            "input_shape": (None, None, 3),
-            "input_tensor": None,
-            "block_type": "block",
-        },
+        "kaggle_handle": "kaggle://keras/resnetv2/keras/resnet152_v2/2",
     },
 }
 
@@ -138,10 +88,7 @@
             "official_name": "ResNetV2",
             "path": "resnet_v2",
         },
-        "class_name": "keras_cv>ResNetV2Backbone",
-        "config": backbone_presets_no_weights["resnet50_v2"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/resnet50v2/imagenet/classification-v2-notop.h5",  # noqa: E501
-        "weights_hash": "e711c83d6db7034871f6d345a476c8184eab99dbf3ffcec0c1d8445684890ad9",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/resnetv2/keras/resnet50_v2_imagenet/2",
     },
 }
 
diff --git a/keras_cv/models/stable_diffusion/__internal__/__init__.py b/keras_cv/models/backbones/vit_det/__init__.py
similarity index 93%
rename from keras_cv/models/stable_diffusion/__internal__/__init__.py
rename to keras_cv/models/backbones/vit_det/__init__.py
index 65be099991..3992ffb59a 100644
--- a/keras_cv/models/stable_diffusion/__internal__/__init__.py
+++ b/keras_cv/models/backbones/vit_det/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The KerasCV Authors
+# Copyright 2023 The KerasCV Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/keras_cv/models/backbones/vit_det/data/vitdet_base_out.npz b/keras_cv/models/backbones/vit_det/data/vitdet_base_out.npz
new file mode 100644
index 0000000000..2063e4fa2a
Binary files /dev/null and b/keras_cv/models/backbones/vit_det/data/vitdet_base_out.npz differ
diff --git a/keras_cv/models/backbones/vit_det/vit_det_aliases.py b/keras_cv/models/backbones/vit_det/vit_det_aliases.py
new file mode 100644
index 0000000000..dec51207e3
--- /dev/null
+++ b/keras_cv/models/backbones/vit_det/vit_det_aliases.py
@@ -0,0 +1,115 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from keras_cv.models.backbones.vit_det.vit_det_backbone import ViTDetBackbone
+from keras_cv.models.backbones.vit_det.vit_det_backbone_presets import (
+    backbone_presets,
+)
+from keras_cv.utils.python_utils import classproperty
+
+ALIAS_DOCSTRING = """VitDet{size}Backbone model.
+
+    Reference:
+        - [Detectron2](https://github.com/facebookresearch/detectron2)
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+
+    For transfer learning use cases, make sure to read the
+    [guide to transfer learning & fine-tuning](https://keras.io/guides/transfer_learning/).
+
+    Examples:
+    ```python
+    input_data = np.ones(shape=(1, 1024, 1024, 3))
+
+    # Randomly initialized backbone
+    model = VitDet{size}Backbone()
+    output = model(input_data)
+    ```
+"""  # noqa: E501
+
+
+class ViTDetBBackbone(ViTDetBackbone):
+    def __new__(
+        cls,
+        **kwargs,
+    ):
+        return ViTDetBackbone.from_preset("vitdet_base", **kwargs)
+
+    @classproperty
+    def presets(cls):
+        """Dictionary of preset names and configurations."""
+        return {
+            "vitdet_base_sa1b": copy.deepcopy(
+                backbone_presets["vitdet_base_sa1b"]
+            ),
+        }
+
+    @classproperty
+    def presets_with_weights(cls):
+        """Dictionary of preset names and configurations that include
+        weights."""
+        return cls.presets
+
+
+class ViTDetLBackbone(ViTDetBackbone):
+    def __new__(
+        cls,
+        **kwargs,
+    ):
+        return ViTDetBackbone.from_preset("vitdet_large", **kwargs)
+
+    @classproperty
+    def presets(cls):
+        """Dictionary of preset names and configurations."""
+        return {
+            "vitdet_large_sa1b": copy.deepcopy(
+                backbone_presets["vitdet_large_sa1b"]
+            ),
+        }
+
+    @classproperty
+    def presets_with_weights(cls):
+        """Dictionary of preset names and configurations that include
+        weights."""
+        return cls.presets
+
+
+class ViTDetHBackbone(ViTDetBackbone):
+    def __new__(
+        cls,
+        **kwargs,
+    ):
+        return ViTDetBackbone.from_preset("vitdet_huge", **kwargs)
+
+    @classproperty
+    def presets(cls):
+        """Dictionary of preset names and configurations."""
+        return {
+            "vitdet_huge_sa1b": copy.deepcopy(
+                backbone_presets["vitdet_huge_sa1b"]
+            ),
+        }
+
+    @classproperty
+    def presets_with_weights(cls):
+        """Dictionary of preset names and configurations that include
+        weights."""
+        return cls.presets
+
+
+setattr(ViTDetBBackbone, "__doc__", ALIAS_DOCSTRING.format(size="B"))
+setattr(ViTDetLBackbone, "__doc__", ALIAS_DOCSTRING.format(size="L"))
+setattr(ViTDetHBackbone, "__doc__", ALIAS_DOCSTRING.format(size="H"))
diff --git a/keras_cv/models/backbones/vit_det/vit_det_backbone.py b/keras_cv/models/backbones/vit_det/vit_det_backbone.py
new file mode 100644
index 0000000000..beb730f4df
--- /dev/null
+++ b/keras_cv/models/backbones/vit_det/vit_det_backbone.py
@@ -0,0 +1,224 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.layers.vit_det_layers import AddPositionalEmbedding
+from keras_cv.layers.vit_det_layers import ViTDetPatchingAndEmbedding
+from keras_cv.layers.vit_det_layers import WindowedTransformerEncoder
+from keras_cv.models import utils
+from keras_cv.models.backbones.backbone import Backbone
+from keras_cv.models.backbones.vit_det.vit_det_backbone_presets import (
+    backbone_presets,
+)
+from keras_cv.models.backbones.vit_det.vit_det_backbone_presets import (
+    backbone_presets_with_weights,
+)
+from keras_cv.utils.python_utils import classproperty
+
+
+@keras_cv_export("keras_cv.models.ViTDetBackbone", package="keras_cv.models")
+class ViTDetBackbone(Backbone):
+    """A ViT image encoder that uses a windowed transformer encoder and
+    relative positional encodings.
+
+    Args:
+        input_shape (tuple[int], optional): The size of the input image in
+            `(H, W, C)` format. Defaults to `(1024, 1024, 3)`.
+        input_tensor (KerasTensor, optional): Output of
+            `keras.layers.Input()`) to use as image input for the model.
+            Defaults to `None`.
+        include_rescaling (bool, optional): Whether to rescale the inputs. If
+            set to `True`, inputs will be passed through a
+            `Rescaling(1/255.0)` layer. Defaults to `False`.
+        patch_size (int, optional): the patch size to be supplied to the
+            Patching layer to turn input images into a flattened sequence of
+            patches. Defaults to `16`.
+        embed_dim (int, optional): The latent dimensionality to be projected
+            into in the output of each stacked windowed transformer encoder.
+            Defaults to `768`.
+        depth (int, optional): The number of transformer encoder layers to
+            stack in the Vision Transformer. Defaults to `12`.
+        mlp_dim (int, optional): The dimensionality of the hidden Dense
+            layer in the transformer MLP head. Defaults to `768*4`.
+        num_heads (int, optional): the number of heads to use in the
+            `MultiHeadAttentionWithRelativePE` layer of each transformer
+            encoder. Defaults to `12`.
+        out_chans (int, optional): The number of channels (features) in the
+            output (image encodings). Defaults to `256`.
+        use_bias (bool, optional): Whether to use bias to project the keys,
+            queries, and values in the attention layer. Defaults to `True`.
+        use_abs_pos (bool, optional): Whether to add absolute positional
+            embeddings to the output patches. Defaults to `True`.
+        use_rel_pos (bool, optional): Whether to use relative positional
+            emcodings in the attention layer. Defaults to `True`.
+        window_size (int, optional): The size of the window for windowed
+            attention in the transformer encoder blocks. Defaults to `14`.
+        global_attention_indices (list, optional): Indexes for blocks using
+            global attention. Defaults to `[2, 5, 8, 11]`.
+        layer_norm_epsilon (int, optional): The epsilon to use in the layer
+            normalization blocks in transformer encoder. Defaults to `1e-6`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+        - [Detectron2](https://github.com/facebookresearch/detectron2)
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        *,
+        include_rescaling,
+        input_shape=(1024, 1024, 3),
+        input_tensor=None,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        mlp_dim=768 * 4,
+        num_heads=12,
+        out_chans=256,
+        use_bias=True,
+        use_abs_pos=True,
+        use_rel_pos=True,
+        window_size=14,
+        global_attention_indices=[2, 5, 8, 11],
+        layer_norm_epsilon=1e-6,
+        **kwargs
+    ):
+        img_input = utils.parse_model_inputs(
+            input_shape, input_tensor, name="images"
+        )
+
+        # Check that the input image is well specified.
+        if img_input.shape[-3] is None or img_input.shape[-2] is None:
+            raise ValueError(
+                "Height and width of the image must be specified"
+                " in `input_shape`."
+            )
+        if img_input.shape[-3] != img_input.shape[-2]:
+            raise ValueError(
+                "Input image must be square i.e. the height must"
+                " be equal to the width in the `input_shape`"
+                " tuple/tensor."
+            )
+
+        img_size = img_input.shape[-3]
+
+        x = img_input
+
+        if include_rescaling:
+            # Use common rescaling strategy across keras_cv
+            x = keras.layers.Rescaling(1.0 / 255.0)(x)
+
+        # VITDet scales inputs based on the standard ImageNet mean/stddev.
+        x = (x - ops.array([0.485, 0.456, 0.406], dtype=x.dtype)) / (
+            ops.array([0.229, 0.224, 0.225], dtype=x.dtype)
+        )
+
+        x = ViTDetPatchingAndEmbedding(
+            kernel_size=(patch_size, patch_size),
+            strides=(patch_size, patch_size),
+            embed_dim=embed_dim,
+        )(x)
+        if use_abs_pos:
+            x = AddPositionalEmbedding(img_size, patch_size, embed_dim)(x)
+
+        for i in range(depth):
+            x = WindowedTransformerEncoder(
+                project_dim=embed_dim,
+                mlp_dim=mlp_dim,
+                num_heads=num_heads,
+                use_bias=use_bias,
+                use_rel_pos=use_rel_pos,
+                window_size=(
+                    window_size if i not in global_attention_indices else 0
+                ),
+                input_size=(img_size // patch_size, img_size // patch_size),
+            )(x)
+        x = keras.models.Sequential(
+            [
+                keras.layers.Conv2D(
+                    filters=out_chans, kernel_size=1, use_bias=False
+                ),
+                keras.layers.LayerNormalization(epsilon=1e-6),
+                keras.layers.Conv2D(
+                    filters=out_chans,
+                    kernel_size=3,
+                    padding="same",
+                    use_bias=False,
+                ),
+                keras.layers.LayerNormalization(epsilon=1e-6),
+            ]
+        )(x)
+
+        super().__init__(inputs=img_input, outputs=x, **kwargs)
+
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.mlp_dim = mlp_dim
+        self.num_heads = num_heads
+        self.out_chans = out_chans
+        self.use_bias = use_bias
+        self.use_rel_pos = use_rel_pos
+        self.use_abs_pos = use_abs_pos
+        self.window_size = window_size
+        self.global_attention_indices = global_attention_indices
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.input_tensor = input_tensor
+        self.include_rescaling = include_rescaling
+
+    @property
+    def pyramid_level_inputs(self):
+        raise NotImplementedError(
+            "The `ViTDetBackbone` model doesn't compute"
+            " pyramid level features."
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "input_shape": self.input_shape[1:],
+                "input_tensor": self.input_tensor,
+                "include_rescaling": self.include_rescaling,
+                "patch_size": self.patch_size,
+                "embed_dim": self.embed_dim,
+                "depth": self.depth,
+                "mlp_dim": self.mlp_dim,
+                "num_heads": self.num_heads,
+                "out_chans": self.out_chans,
+                "use_bias": self.use_bias,
+                "use_abs_pos": self.use_abs_pos,
+                "use_rel_pos": self.use_rel_pos,
+                "window_size": self.window_size,
+                "global_attention_indices": self.global_attention_indices,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+            }
+        )
+        return config
+
+    @classproperty
+    def presets(cls):
+        """Dictionary of preset names and configurations."""
+        return copy.deepcopy(backbone_presets)
+
+    @classproperty
+    def presets_with_weights(cls):
+        """Dictionary of preset names and configurations that include
+        weights."""
+        return copy.deepcopy(backbone_presets_with_weights)
diff --git a/keras_cv/models/backbones/vit_det/vit_det_backbone_presets.py b/keras_cv/models/backbones/vit_det/vit_det_backbone_presets.py
new file mode 100644
index 0000000000..3c8fd3651b
--- /dev/null
+++ b/keras_cv/models/backbones/vit_det/vit_det_backbone_presets.py
@@ -0,0 +1,102 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""VitDet model preset configurations."""
+
+backbone_presets_no_weights = {
+    "vitdet_base": {
+        "metadata": {
+            "description": (
+                "Detectron2 ViT basebone with 12 "
+                "transformer encoders with embed dim 768 and attention layers"
+                " with 12 heads with global attention on encoders 2, 5, 8, "
+                "and 11."
+            ),
+            "params": 89_670_912,
+            "official_name": "VitDet",
+            "path": "vit_det",
+        },
+        "kaggle_handle": "kaggle://keras/vitdet/keras/vitdet_base/2",
+    },
+    "vitdet_large": {
+        "metadata": {
+            "description": (
+                "Detectron2 ViT basebone with 24 "
+                "transformer encoders with embed dim "
+                "1024 and attention layers with 16 heads with global "
+                "attention on encoders 5, 11, 17, and 23."
+            ),
+            "params": 308_278_272,
+            "official_name": "VitDet",
+            "path": "vit_det",
+        },
+        "kaggle_handle": "kaggle://keras/vitdet/keras/vitdet_large/2",
+    },
+    "vitdet_huge": {
+        "metadata": {
+            "description": (
+                "Detectron2 ViT basebone model "
+                "with 32 transformer encoders with embed dim "
+                "1280 and attention layers with 16 heads with global "
+                "attention on encoders 7, 15, 23, and 31."
+            ),
+            "params": 637_026_048,
+            "official_name": "VitDet",
+            "path": "vit_det",
+        },
+        "kaggle_handle": "kaggle://keras/vitdet/keras/vitdet_huge/2",
+    },
+}
+
+
+backbone_presets_with_weights = {
+    "vitdet_base_sa1b": {
+        "metadata": {
+            "description": (
+                "A base Detectron2 ViT backbone trained on the SA1B dataset."
+            ),
+            "params": 89_670_912,
+            "official_name": "VitDet",
+            "path": "vit_det",
+        },
+        "kaggle_handle": "kaggle://keras/vitdet/keras/vitdet_base_sa1b/2",
+    },
+    "vitdet_large_sa1b": {
+        "metadata": {
+            "description": (
+                "A large Detectron2 ViT backbone trained on the SA1B dataset."
+            ),
+            "params": 308_278_272,
+            "official_name": "VitDet",
+            "path": "vit_det",
+        },
+        "kaggle_handle": "kaggle://keras/vitdet/keras/vitdet_large_sa1b/2",
+    },
+    "vitdet_huge_sa1b": {
+        "metadata": {
+            "description": (
+                "A huge Detectron2 ViT backbone trained on the SA1B dataset."
+            ),
+            "params": 637_026_048,
+            "official_name": "VitDet",
+            "path": "vit_det",
+        },
+        "kaggle_handle": "kaggle://keras/vitdet/keras/vitdet_huge_sa1b/2",
+    },
+}
+
+
+backbone_presets = {
+    **backbone_presets_no_weights,
+    **backbone_presets_with_weights,
+}
diff --git a/keras_cv/models/backbones/vit_det/vit_det_backbone_presets_test.py b/keras_cv/models/backbones/vit_det/vit_det_backbone_presets_test.py
new file mode 100644
index 0000000000..7575ecd1d8
--- /dev/null
+++ b/keras_cv/models/backbones/vit_det/vit_det_backbone_presets_test.py
@@ -0,0 +1,96 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for loading pretrained model presets."""
+
+import pathlib
+
+import numpy as np
+import pytest
+
+from keras_cv.backend import ops
+from keras_cv.models.backbones.vit_det.vit_det_aliases import ViTDetBBackbone
+from keras_cv.models.backbones.vit_det.vit_det_backbone import ViTDetBackbone
+from keras_cv.tests.test_case import TestCase
+
+
+@pytest.mark.large
+class ViTDetPresetSmokeTest(TestCase):
+    """
+    A smoke test for ViTDet presets we run continuously.
+    This only tests the smallest weights we have available. Run with:
+    `pytest keras_cv/models/backbones/detectron2/detectron2_backbone_presets_test.py --run_large`  # noqa: E501
+    """
+
+    def setUp(self):
+        self.input_batch = np.ones(shape=(1, 1024, 1024, 3))
+
+    def test_backbone_output(self):
+        model = ViTDetBackbone.from_preset("vitdet_base_sa1b")
+        outputs = model(self.input_batch)
+
+        # The forward pass from a preset should be stable!
+        # This test should catch cases where we unintentionally change our
+        # network code in a way that would invalidate our preset weights.
+        # We should only update these numbers if we are updating a weights
+        # file, or have found a discrepancy with the upstream source.
+
+        expected = np.load(
+            pathlib.Path(__file__).parent / "data" / "vitdet_base_out.npz"
+        )
+        # Keep a high tolerance, so we are robust to different hardware.
+        self.assertAllClose(
+            ops.convert_to_numpy(outputs),
+            expected,
+            atol=1e-5,
+            rtol=1e-5,
+        )
+
+    def test_applications_model_output(self):
+        model = ViTDetBBackbone()
+        model(self.input_batch)
+
+    def test_applications_model_output_with_preset(self):
+        model = ViTDetBackbone.from_preset("vitdet_base")
+        model(self.input_batch)
+
+    def test_applications_model_predict(self):
+        model = ViTDetBBackbone()
+        # Test that the model XLA compiles
+        model.predict(self.input_batch)
+
+    def test_preset_docstring(self):
+        """Check we did our docstring formatting correctly."""
+        for name in ViTDetBackbone.presets:
+            self.assertRegex(ViTDetBackbone.from_preset.__doc__, name)
+
+    def test_unknown_preset_error(self):
+        # Not a preset name
+        with self.assertRaises(ValueError):
+            ViTDetBackbone.from_preset("vitdet_nonexistant")
+
+
+@pytest.mark.extra_large
+class ViTDetPresetFullTest(TestCase):
+    """
+    Test the full enumeration of our preset.
+    This tests every preset for ViTDet and is only run manually.
+    Run with:
+    `pytest keras_cv/models/backbones/detectron2/detectron2_backbone_presets_test.py --run_extra_large`  # noqa: E501
+    """
+
+    def test_load_ViTDet(self):
+        input_data = np.ones(shape=(1, 1024, 1024, 3))
+        for preset in ViTDetBackbone.presets:
+            model = ViTDetBackbone.from_preset(preset)
+            model(input_data)
diff --git a/keras_cv/models/backbones/vit_det/vit_det_backbone_test.py b/keras_cv/models/backbones/vit_det/vit_det_backbone_test.py
new file mode 100644
index 0000000000..fb12438119
--- /dev/null
+++ b/keras_cv/models/backbones/vit_det/vit_det_backbone_test.py
@@ -0,0 +1,61 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+import numpy as np
+import pytest
+
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.models.backbones.vit_det.vit_det_aliases import ViTDetBBackbone
+from keras_cv.tests.test_case import TestCase
+
+
+class TestViTDetBackbone(TestCase):
+    @pytest.mark.large
+    def test_call(self):
+        model = ViTDetBBackbone()
+        x = np.ones((1, 1024, 1024, 3))
+        x_out = ops.convert_to_numpy(model(x))
+        num_parameters = sum(
+            np.prod(tuple(x.shape)) for x in model.trainable_variables
+        )
+        self.assertEqual(x_out.shape, (1, 64, 64, 256))
+        self.assertEqual(num_parameters, 89_670_912)
+
+    @pytest.mark.extra_large
+    def teat_save(self):
+        # saving test
+        model = ViTDetBBackbone()
+        x = np.ones((1, 1024, 1024, 3))
+        x_out = ops.convert_to_numpy(model(x))
+        path = os.path.join(self.get_temp_dir(), "model.keras")
+        model.save(path)
+        loaded_model = keras.saving.load_model(path)
+        x_out_loaded = ops.convert_to_numpy(loaded_model(x))
+        self.assertAllClose(x_out, x_out_loaded)
+
+    @pytest.mark.extra_large
+    def test_fit(self):
+        model = ViTDetBBackbone()
+        x = np.ones((1, 1024, 1024, 3))
+        y = np.zeros((1, 64, 64, 256))
+        model.compile(optimizer="adam", loss="mse", metrics=["mse"])
+        model.fit(x, y, epochs=1)
+
+    def test_pyramid_level_inputs_error(self):
+        model = ViTDetBBackbone()
+        with self.assertRaises(NotImplementedError, msg="doesn't compute"):
+            model.pyramid_level_inputs
diff --git a/keras_cv/models/classification/image_classifier.py b/keras_cv/models/classification/image_classifier.py
index ce1cc2d23e..a2d6a16674 100644
--- a/keras_cv/models/classification/image_classifier.py
+++ b/keras_cv/models/classification/image_classifier.py
@@ -28,7 +28,12 @@
 from keras_cv.utils.python_utils import classproperty
 
 
-@keras_cv_export("keras_cv.models.ImageClassifier")
+@keras_cv_export(
+    [
+        "keras_cv.models.ImageClassifier",
+        "keras_cv.models.classification.ImageClassifier",
+    ]
+)
 class ImageClassifier(Task):
     """Image classifier with pooling and dense layer prediction head.
 
diff --git a/keras_cv/models/classification/image_classifier_presets.py b/keras_cv/models/classification/image_classifier_presets.py
index c210997951..2c11696b18 100644
--- a/keras_cv/models/classification/image_classifier_presets.py
+++ b/keras_cv/models/classification/image_classifier_presets.py
@@ -13,11 +13,6 @@
 # limitations under the License.
 """ImageClassifier Task presets."""
 
-from keras_cv.models.backbones.efficientnet_v2 import (
-    efficientnet_v2_backbone_presets,
-)
-from keras_cv.models.backbones.mobilenet_v3 import mobilenet_v3_backbone_presets
-from keras_cv.models.backbones.resnet_v2 import resnet_v2_backbone_presets
 
 classifier_presets = {
     "resnet50_v2_imagenet_classifier": {
@@ -32,16 +27,7 @@
             "official_name": "ImageClassifier",
             "path": "image_classifier",
         },
-        "config": {
-            "backbone": resnet_v2_backbone_presets.backbone_presets[
-                "resnet50_v2"
-            ],
-            "num_classes": 1000,
-            "pooling": "avg",
-            "activation": "softmax",
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/resnet50v2/imagenet-classifier-v0.h5",  # noqa: E501
-        "weights_hash": "77fa9f1cd1de0e202309e51d4e598e441d1111dacb6c41a182b6c63f76ff26cd",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/resnetv2/keras/resnet50_v2_imagenet_classifier/2",  # noqa: E501
     },
     "efficientnetv2_s_imagenet_classifier": {
         "metadata": {
@@ -58,16 +44,7 @@
             "official_name": "ImageClassifier",
             "path": "image_classifier",
         },
-        "config": {
-            "backbone": efficientnet_v2_backbone_presets.backbone_presets[
-                "efficientnetv2_s"
-            ],
-            "num_classes": 1000,
-            "pooling": "avg",
-            "activation": "softmax",
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/efficientnet_v2/efficientnetv2-s_imagenet_classifier.h5",  # noqa: E501
-        "weights_hash": "4da57ade035e11aff7e1c5cb04c4235e",
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_s_imagenet_classifier/2",  # noqa: E501
     },
     "efficientnetv2_b0_imagenet_classifier": {
         "metadata": {
@@ -87,16 +64,7 @@
             "official_name": "ImageClassifier",
             "path": "image_classifier",
         },
-        "config": {
-            "backbone": efficientnet_v2_backbone_presets.backbone_presets[
-                "efficientnetv2_b0"
-            ],
-            "num_classes": 1000,
-            "pooling": "avg",
-            "activation": "softmax",
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/efficientnet_v2/efficientnetv2-b0_imagenet_classifier.h5",  # noqa: E501
-        "weights_hash": "4b739de648c346b7e4156e11a223c338",
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b0_imagenet_classifier/2",  # noqa: E501
     },
     "efficientnetv2_b1_imagenet_classifier": {
         "metadata": {
@@ -116,16 +84,7 @@
             "official_name": "ImageClassifier",
             "path": "image_classifier",
         },
-        "config": {
-            "backbone": efficientnet_v2_backbone_presets.backbone_presets[
-                "efficientnetv2_b1"
-            ],
-            "num_classes": 1000,
-            "pooling": "avg",
-            "activation": "softmax",
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/efficientnet_v2/efficientnetv2-b1_imagenet_classifier.h5",  # noqa: E501
-        "weights_hash": "78c1c879143dbd8f74e6ffc4d3180197",
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b1_imagenet_classifier/2",  # noqa: E501
     },
     "efficientnetv2_b2_imagenet_classifier": {
         "metadata": {
@@ -145,16 +104,7 @@
             "official_name": "ImageClassifier",
             "path": "image_classifier",
         },
-        "config": {
-            "backbone": efficientnet_v2_backbone_presets.backbone_presets[
-                "efficientnetv2_b2"
-            ],
-            "num_classes": 1000,
-            "pooling": "avg",
-            "activation": "softmax",
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/efficientnet_v2/efficientnetv2-b2_imagenet_classifier.h5",  # noqa: E501
-        "weights_hash": "07eda1c48aee27e12a3fe2545e6c65ed",
+        "kaggle_handle": "kaggle://keras/efficientnetv2/keras/efficientnetv2_b2_imagenet_classifier/2",  # noqa: E501
     },
     "mobilenet_v3_large_imagenet_classifier": {
         "metadata": {
@@ -171,15 +121,6 @@
             "official_name": "ImageClassifier",
             "path": "image_classifier",
         },
-        "config": {
-            "backbone": mobilenet_v3_backbone_presets.backbone_presets[
-                "mobilenet_v3_large"
-            ],
-            "num_classes": 1000,
-            "pooling": "avg",
-            "activation": "softmax",
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/mobilenetv3/mobilenetv3_large_imagenet_classifier.h5",  # noqa: E501
-        "weights_hash": "4130fbc249d3812719396439031d89042dce440dd256cda941a109bce6d4e305",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/mobilenetv3/keras/mobilenet_v3_large_imagenet_classifier/2",  # noqa: E501
     },
 }
diff --git a/keras_cv/models/classification/image_classifier_test.py b/keras_cv/models/classification/image_classifier_test.py
index 67c4b2a31c..8d9e48cc6f 100644
--- a/keras_cv/models/classification/image_classifier_test.py
+++ b/keras_cv/models/classification/image_classifier_test.py
@@ -23,6 +23,7 @@
 
 from keras_cv.backend import keras
 from keras_cv.backend import ops
+from keras_cv.backend.config import keras_3
 from keras_cv.models.backbones.resnet_v2.resnet_v2_aliases import (
     ResNet18V2Backbone,
 )
@@ -50,6 +51,9 @@ def test_valid_call(self):
     @pytest.mark.large  # Fit is slow, so mark these large.
     @pytest.mark.filterwarnings("ignore::UserWarning")  # Torch + jit_compile
     def test_classifier_fit(self, jit_compile):
+        if keras_3() and jit_compile and keras.backend.backend() == "torch":
+            self.skipTest("TODO: Torch Backend `jit_compile` fails on GPU.")
+            self.supports_jit = False
         model = ImageClassifier(
             backbone=ResNet18V2Backbone(),
             num_classes=2,
diff --git a/keras_cv/models/legacy/darknet.py b/keras_cv/models/legacy/darknet.py
index ea7fd429f2..2dc14d499d 100644
--- a/keras_cv/models/legacy/darknet.py
+++ b/keras_cv/models/legacy/darknet.py
@@ -76,7 +76,6 @@
 
 @keras.utils.register_keras_serializable(package="keras_cv.models")
 class DarkNet(keras.Model):
-
     """Represents the DarkNet architecture.
 
     The DarkNet architecture is commonly used for detection tasks. It is
diff --git a/keras_cv/models/legacy/mlp_mixer.py b/keras_cv/models/legacy/mlp_mixer.py
index a48544f905..170d0a4c6f 100644
--- a/keras_cv/models/legacy/mlp_mixer.py
+++ b/keras_cv/models/legacy/mlp_mixer.py
@@ -143,7 +143,6 @@ def apply_mixer_block(x, tokens_mlp_dim, channels_mlp_dim, name=None):
 
 @keras.utils.register_keras_serializable(package="keras_cv.models")
 class MLPMixer(keras.Model):
-
     """Instantiates the MLP Mixer architecture.
 
     Args:
diff --git a/keras_cv/models/legacy/object_detection/faster_rcnn/faster_rcnn.py b/keras_cv/models/legacy/object_detection/faster_rcnn/faster_rcnn.py
index f380e3c1ba..f4c2543e75 100644
--- a/keras_cv/models/legacy/object_detection/faster_rcnn/faster_rcnn.py
+++ b/keras_cv/models/legacy/object_detection/faster_rcnn/faster_rcnn.py
@@ -16,9 +16,9 @@
 from absl import logging
 from tensorflow import keras
 
-import keras_cv
 from keras_cv import bounding_box
 from keras_cv import layers as cv_layers
+from keras_cv import models
 from keras_cv.bounding_box.converters import _decode_deltas_to_boxes
 from keras_cv.bounding_box.utils import _clip_boxes
 from keras_cv.layers.object_detection.anchor_generator import AnchorGenerator
@@ -321,7 +321,7 @@ def __init__(
         )
         self.roi_pooler = _ROIAligner(bounding_box_format="yxyx")
         self.rcnn_head = rcnn_head or RCNNHead(num_classes)
-        self.backbone = backbone or keras_cv.models.ResNet50Backbone()
+        self.backbone = backbone or models.ResNet50Backbone()
         extractor_levels = ["P2", "P3", "P4", "P5"]
         extractor_layer_names = [
             self.backbone.pyramid_level_inputs[i] for i in extractor_levels
diff --git a/keras_cv/models/object_detection/predict_utils.py b/keras_cv/models/object_detection/predict_utils.py
index 8eb6e10e62..9cd7a1bd44 100644
--- a/keras_cv/models/object_detection/predict_utils.py
+++ b/keras_cv/models/object_detection/predict_utils.py
@@ -15,15 +15,23 @@
 import tensorflow as tf
 
 try:
-    from keras.src.engine.training import _minimum_control_deps
-    from keras.src.engine.training import reduce_per_replica
     from keras.src.utils import tf_utils
 except ImportError:
-    from keras.engine.training import _minimum_control_deps
-    from keras.engine.training import reduce_per_replica
     from keras.utils import tf_utils
 
 
+def _minimum_control_deps(outputs):
+    """Returns the minimum control dependencies to ensure step succeeded."""
+    if tf.executing_eagerly():
+        return []  # Control dependencies not needed.
+    outputs = tf.nest.flatten(outputs, expand_composites=True)
+    for out in outputs:
+        # Variables can't be control dependencies.
+        if not isinstance(out, tf.Variable):
+            return [out]  # Return first Tensor or Op from outputs.
+    return []  # No viable Tensor or Op to use for control deps.
+
+
 def make_predict_function(model, force=False):
     if model.predict_function is not None and not force:
         return model.predict_function
@@ -45,9 +53,7 @@ def run_step(data):
 
         data = next(iterator)
         outputs = model.distribute_strategy.run(run_step, args=(data,))
-        outputs = reduce_per_replica(
-            outputs, model.distribute_strategy, reduction="concat"
-        )
+        outputs = model.distribute_strategy.gather(outputs, axis=0)
         # Note that this is the only deviation from the base keras.Model
         # implementation. We add the decode_step inside of the computation
         # graph but outside of the distribute_strategy (i.e on host CPU).
diff --git a/keras_cv/models/object_detection/retinanet/retinanet.py b/keras_cv/models/object_detection/retinanet/retinanet.py
index 5f047897ed..2244da108c 100644
--- a/keras_cv/models/object_detection/retinanet/retinanet.py
+++ b/keras_cv/models/object_detection/retinanet/retinanet.py
@@ -16,9 +16,9 @@
 
 import numpy as np
 
-import keras_cv
 from keras_cv import bounding_box
 from keras_cv import layers as cv_layers
+from keras_cv import losses
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
@@ -41,7 +41,9 @@
 BOX_VARIANCE = [0.1, 0.1, 0.2, 0.2]
 
 
-@keras_cv_export("keras_cv.models.RetinaNet")
+@keras_cv_export(
+    ["keras_cv.models.RetinaNet", "keras_cv.models.object_detection.RetinaNet"]
+)
 class RetinaNet(Task):
     """A Keras model implementing the RetinaNet meta-architecture.
 
@@ -522,6 +524,12 @@ def from_config(cls, config):
             config["label_encoder"] = keras.layers.deserialize(
                 config["label_encoder"]
             )
+        if "prediction_decoder" in config and isinstance(
+            config["prediction_decoder"], dict
+        ):
+            config["prediction_decoder"] = keras.layers.deserialize(
+                config["prediction_decoder"]
+            )
         return super().from_config(config)
 
     @classproperty
@@ -551,7 +559,7 @@ def _parse_box_loss(loss):
 
     # case insensitive comparison
     if loss.lower() == "smoothl1":
-        return keras_cv.losses.SmoothL1Loss(l1_cutoff=1.0, reduction="sum")
+        return losses.SmoothL1Loss(l1_cutoff=1.0, reduction="sum")
     if loss.lower() == "huber":
         return keras.losses.Huber(reduction="sum")
 
@@ -568,7 +576,7 @@ def _parse_classification_loss(loss):
 
     # case insensitive comparison
     if loss.lower() == "focal":
-        return keras_cv.losses.FocalLoss(from_logits=True, reduction="sum")
+        return losses.FocalLoss(from_logits=True, reduction="sum")
 
     raise ValueError(
         "Expected `classification_loss` to be either a Keras Loss, "
diff --git a/keras_cv/models/object_detection/retinanet/retinanet_label_encoder.py b/keras_cv/models/object_detection/retinanet/retinanet_label_encoder.py
index 00718d4377..8b08bbc87a 100644
--- a/keras_cv/models/object_detection/retinanet/retinanet_label_encoder.py
+++ b/keras_cv/models/object_detection/retinanet/retinanet_label_encoder.py
@@ -68,7 +68,7 @@ def __init__(
         self.box_variance = ops.array(box_variance, "float32")
         self.background_class = background_class
         self.ignore_class = ignore_class
-        self.matched_boxes_metric = keras.metrics.BinaryAccuracy(
+        self.matched_boxes_metric = MatchedBoxesMetric(
             name="percent_boxes_matched_with_anchor"
         )
         self.positive_threshold = positive_threshold
@@ -124,7 +124,7 @@ def _encode_sample(self, box_labels, anchor_boxes, image_shape):
             gt_boxes, matched_gt_idx
         )
         matched_gt_boxes = ops.reshape(
-            matched_gt_boxes, (-1, matched_gt_boxes.shape[1], 4)
+            matched_gt_boxes, (-1, ops.shape(matched_gt_boxes)[1], 4)
         )
 
         box_target = bounding_box._encode_box_to_deltas(
@@ -166,7 +166,7 @@ def _encode_sample(self, box_labels, anchor_boxes, image_shape):
         box_shape = ops.shape(gt_boxes)
         batch_size = box_shape[0]
         n_boxes = box_shape[1]
-        box_ids = ops.arange(gt_boxes.shape[1], dtype=matched_gt_idx.dtype)
+        box_ids = ops.arange(n_boxes, dtype=matched_gt_idx.dtype)
         matched_ids = ops.expand_dims(matched_gt_idx, axis=-1)
         matches = box_ids == matched_ids
         matches = ops.any(matches, axis=1)
@@ -197,7 +197,8 @@ def call(self, images, box_labels):
                 f"Received `type(images)={type(images)}`."
             )
 
-        image_shape = tuple(images[0].shape)
+        image_shape = ops.shape(images)
+        image_shape = (image_shape[1], image_shape[2], image_shape[3])
         box_labels = bounding_box.to_dense(box_labels)
         if len(box_labels["classes"].shape) == 2:
             box_labels["classes"] = ops.expand_dims(
@@ -215,7 +216,7 @@ def call(self, images, box_labels):
         result = self._encode_sample(box_labels, anchor_boxes, image_shape)
         encoded_box_targets = result["boxes"]
         encoded_box_targets = ops.reshape(
-            encoded_box_targets, (-1, encoded_box_targets.shape[1], 4)
+            encoded_box_targets, (-1, ops.shape(encoded_box_targets)[1], 4)
         )
         class_targets = result["classes"]
         return encoded_box_targets, class_targets
@@ -232,3 +233,20 @@ def get_config(self):
         }
         base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
+
+    @classmethod
+    def from_config(cls, config):
+        if "anchor_generator" in config and isinstance(
+            config["anchor_generator"], dict
+        ):
+            config["anchor_generator"] = keras.layers.deserialize(
+                config["anchor_generator"]
+            )
+
+        return super().from_config(config)
+
+
+class MatchedBoxesMetric(keras.metrics.BinaryAccuracy):
+    # Prevent `load_weights` from accessing metric
+    def load_own_variables(self, store):
+        return
diff --git a/keras_cv/models/object_detection/retinanet/retinanet_presets.py b/keras_cv/models/object_detection/retinanet/retinanet_presets.py
index 6bcd6d616b..bc7f747c38 100644
--- a/keras_cv/models/object_detection/retinanet/retinanet_presets.py
+++ b/keras_cv/models/object_detection/retinanet/retinanet_presets.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """RetinaNet Task presets."""
-from keras_cv.models.backbones.resnet_v1 import resnet_v1_backbone_presets
 
 retinanet_presets = {
     "retinanet_resnet50_pascalvoc": {
@@ -27,13 +26,6 @@
             "official_name": "RetinaNet",
             "path": "retinanet",
         },
-        "config": {
-            "backbone": resnet_v1_backbone_presets.backbone_presets["resnet50"],
-            # 21 used as an implicit background class marginally improves
-            # performance.
-            "num_classes": 20,
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/retinanet/pascal_voc/resnet50-v4.h5",  # noqa: E501
-        "weights_hash": "961f6f9a03c869900d61ae74a9f0c31c6439b4f81f13593f5c16c4733061fbac",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/retinanet/keras/retinanet_resnet50_pascalvoc/2",  # noqa: E501
     },
 }
diff --git a/keras_cv/models/object_detection/retinanet/retinanet_test.py b/keras_cv/models/object_detection/retinanet/retinanet_test.py
index cc0f7c9131..bbede75943 100644
--- a/keras_cv/models/object_detection/retinanet/retinanet_test.py
+++ b/keras_cv/models/object_detection/retinanet/retinanet_test.py
@@ -20,6 +20,7 @@
 from absl.testing import parameterized
 
 import keras_cv
+from keras_cv.backend import config
 from keras_cv.backend import keras
 from keras_cv.backend import ops
 from keras_cv.models.backbones.test_backbone_presets import (
@@ -249,6 +250,49 @@ def test_call_with_custom_label_encoder(self):
         )
         model(ops.ones(shape=(2, 224, 224, 3)))
 
+    def test_tf_dataset_data_generator(self):
+        if config.backend() != "tensorflow":
+            pytest.skip("TensorFlow required for `tf.data.Dataset` test.")
+
+        def data_generator():
+            image = tf.ones((512, 512, 3), dtype=tf.float32)
+
+            bounding_boxes = {
+                "boxes": tf.ones((3, 4), dtype=tf.float32),
+                "classes": tf.ones((3,), dtype=tf.float32),
+            }
+
+            yield {"images": image, "bounding_boxes": bounding_boxes}
+
+        data = tf.data.Dataset.from_generator(
+            generator=data_generator,
+            output_signature={
+                "images": tf.TensorSpec(shape=(512, 512, 3), dtype=tf.float32),
+                "bounding_boxes": {
+                    "boxes": tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
+                    "classes": tf.TensorSpec(shape=(None,), dtype=tf.float32),
+                },
+            },
+        ).batch(1)
+
+        model = keras_cv.models.RetinaNet(
+            num_classes=2,
+            bounding_box_format="xyxy",
+            backbone=keras_cv.models.ResNet50Backbone.from_preset(
+                "resnet50_imagenet",
+                load_weights=False,
+            ),
+        )
+
+        model.compile(
+            classification_loss="focal",
+            box_loss="smoothl1",
+            optimizer="adam",
+            jit_compile=False,
+        )
+
+        model.fit(data, epochs=1, batch_size=1, steps_per_epoch=1)
+
 
 @pytest.mark.large
 class RetinaNetSmokeTest(TestCase):
diff --git a/keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone.py b/keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone.py
index f4bd99fafa..a2bf4bdd3b 100644
--- a/keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone.py
+++ b/keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone.py
@@ -178,9 +178,9 @@ def __init__(
                     activation=activation,
                     name=f"{stack_name}_spp_fast",
                 )
-            pyramid_level_inputs[
-                f"P{stack_id + 2}"
-            ] = utils.get_tensor_input_name(x)
+            pyramid_level_inputs[f"P{stack_id + 2}"] = (
+                utils.get_tensor_input_name(x)
+            )
 
         super().__init__(inputs=inputs, outputs=x, **kwargs)
         self.pyramid_level_inputs = pyramid_level_inputs
diff --git a/keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone_presets.py b/keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone_presets.py
index 6113ea2d4b..2c5cf4baa5 100644
--- a/keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone_presets.py
+++ b/keras_cv/models/object_detection/yolo_v8/yolo_v8_backbone_presets.py
@@ -22,14 +22,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": {
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "stackwise_channels": [32, 64, 128, 256],
-            "stackwise_depth": [1, 2, 2, 1],
-            "activation": "swish",
-        },
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_xs_backbone/2",
     },
     "yolo_v8_s_backbone": {
         "metadata": {
@@ -38,14 +31,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": {
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "stackwise_channels": [64, 128, 256, 512],
-            "stackwise_depth": [1, 2, 2, 1],
-            "activation": "swish",
-        },
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_s_backbone/2",
     },
     "yolo_v8_m_backbone": {
         "metadata": {
@@ -54,14 +40,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": {
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "stackwise_channels": [96, 192, 384, 576],
-            "stackwise_depth": [2, 4, 4, 2],
-            "activation": "swish",
-        },
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_m_backbone/2",
     },
     "yolo_v8_l_backbone": {
         "metadata": {
@@ -70,14 +49,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": {
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "stackwise_channels": [128, 256, 512, 512],
-            "stackwise_depth": [3, 6, 6, 3],
-            "activation": "swish",
-        },
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_l_backbone/2",
     },
     "yolo_v8_xl_backbone": {
         "metadata": {
@@ -86,14 +58,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": {
-            "include_rescaling": True,
-            "input_shape": (None, None, 3),
-            "stackwise_channels": [160, 320, 640, 640],
-            "stackwise_depth": [3, 6, 6, 3],
-            "activation": "swish",
-        },
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_xl_backbone/2",
     },
 }
 
@@ -107,10 +72,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": backbone_presets_no_weights["yolo_v8_xs_backbone"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/yolov8/coco/yolov8_n_backbone.h5",  # noqa: E501
-        "weights_hash": "636ba3cba064c7c457e30a0e4759716006c305c30876df1c1caf2e56b99eab6c",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_xs_backbone_coco/2",  # noqa: E501
     },
     "yolo_v8_s_backbone_coco": {
         "metadata": {
@@ -119,10 +81,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": backbone_presets_no_weights["yolo_v8_s_backbone"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/yolov8/coco/yolov8_s_backbone.h5",  # noqa: E501
-        "weights_hash": "49ab5da87d6b36a1943e7f111a1960355171332c25312b6cc01526baaecf1b69",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_s_backbone_coco/2",  # noqa: E501
     },
     "yolo_v8_m_backbone_coco": {
         "metadata": {
@@ -131,10 +90,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": backbone_presets_no_weights["yolo_v8_m_backbone"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/yolov8/coco/yolov8_m_backbone.h5",  # noqa: E501
-        "weights_hash": "a9719807699a2540da14aa7f9a0dda272d400d30c40a956298a63a2805aa6436",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_m_backbone_coco/2",  # noqa: E501
     },
     "yolo_v8_l_backbone_coco": {
         "metadata": {
@@ -143,10 +99,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": backbone_presets_no_weights["yolo_v8_l_backbone"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/yolov8/coco/yolov8_l_backbone.h5",  # noqa: E501
-        "weights_hash": "2c94ffe75492491974c6d7347d5c1d1aa209d8f6d78c63ab62df0f5dd51680b9",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_l_backbone_coco/2",  # noqa: E501
     },
     "yolo_v8_xl_backbone_coco": {
         "metadata": {
@@ -157,10 +110,7 @@
             "official_name": "YOLOV8",
             "path": "yolo_v8",
         },
-        "class_name": "keras_cv>YOLOV8Backbone",
-        "config": backbone_presets_no_weights["yolo_v8_xl_backbone"]["config"],
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/yolov8/coco/yolov8_x_backbone.h5",  # noqa: E501
-        "weights_hash": "ce0cc3235eacaffc4a9824e28b2366e674b6d42befc4c7b77f3be7d1d39960bd",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_xl_backbone_coco/2",  # noqa: E501
     },
 }
 
diff --git a/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector.py b/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector.py
index 584ed1af79..6c17c71a72 100644
--- a/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector.py
+++ b/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector.py
@@ -14,8 +14,8 @@
 import copy
 import warnings
 
-import keras_cv
 from keras_cv import bounding_box
+from keras_cv import layers
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
@@ -320,7 +320,12 @@ def dist2bbox(distance, anchor_points):
     return ops.concatenate((x1y1, x2y2), axis=-1)  # xyxy bbox
 
 
-@keras_cv_export("keras_cv.models.YOLOV8Detector")
+@keras_cv_export(
+    [
+        "keras_cv.models.YOLOV8Detector",
+        "keras_cv.models.object_detection.YOLOV8Detector",
+    ]
+)
 class YOLOV8Detector(Task):
     """Implements the YOLOV8 architecture for object detection.
 
@@ -353,22 +358,23 @@ class YOLOV8Detector(Task):
     ```python
     images = tf.ones(shape=(1, 512, 512, 3))
     labels = {
-        "boxes": [
+        "boxes": tf.constant([
             [
                 [0, 0, 100, 100],
                 [100, 100, 200, 200],
                 [300, 300, 100, 100],
             ]
-        ],
-        "classes": [[1, 1, 1]],
+        ], dtype=tf.float32),
+        "classes": tf.constant([[1, 1, 1]], dtype=tf.int64),
     }
+
     model = keras_cv.models.YOLOV8Detector(
         num_classes=20,
         bounding_box_format="xywh",
         backbone=keras_cv.models.YOLOV8Backbone.from_preset(
-            "yolo_v8_m_coco"
+            "yolo_v8_m_backbone_coco"
         ),
-        fpn_depth=2.
+        fpn_depth=2
     )
 
     # Evaluate model without box decoding and NMS
@@ -430,7 +436,7 @@ def __init__(
         self.bounding_box_format = bounding_box_format
         self._prediction_decoder = (
             prediction_decoder
-            or keras_cv.layers.NonMaxSuppression(
+            or layers.NonMaxSuppression(
                 bounding_box_format=bounding_box_format,
                 from_logits=False,
                 confidence_threshold=0.2,
@@ -635,10 +641,33 @@ def get_config(self):
             "bounding_box_format": self.bounding_box_format,
             "fpn_depth": self.fpn_depth,
             "backbone": keras.saving.serialize_keras_object(self.backbone),
-            "label_encoder": self.label_encoder,
-            "prediction_decoder": self._prediction_decoder,
+            "label_encoder": keras.saving.serialize_keras_object(
+                self.label_encoder
+            ),
+            "prediction_decoder": keras.saving.serialize_keras_object(
+                self._prediction_decoder
+            ),
         }
 
+    @classmethod
+    def from_config(cls, config):
+        config["backbone"] = keras.saving.deserialize_keras_object(
+            config["backbone"]
+        )
+        label_encoder = config.get("label_encoder")
+        if label_encoder is not None and isinstance(label_encoder, dict):
+            config["label_encoder"] = keras.saving.deserialize_keras_object(
+                label_encoder
+            )
+        prediction_decoder = config.get("prediction_decoder")
+        if prediction_decoder is not None and isinstance(
+            prediction_decoder, dict
+        ):
+            config["prediction_decoder"] = (
+                keras.saving.deserialize_keras_object(prediction_decoder)
+            )
+        return cls(**config)
+
     @classproperty
     def presets(cls):
         """Dictionary of preset names and configurations."""
diff --git a/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector_presets.py b/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector_presets.py
index dfbbe512c1..d7fa5d27e1 100644
--- a/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector_presets.py
+++ b/keras_cv/models/object_detection/yolo_v8/yolo_v8_detector_presets.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """YOLOv8 Task presets."""
-from keras_cv.models.object_detection.yolo_v8 import yolo_v8_backbone_presets
 
 yolo_v8_detector_presets = {
     "yolo_v8_m_pascalvoc": {
@@ -26,14 +25,6 @@
             "official_name": "YOLOV8Detector",
             "path": "yolo_v8_detector",
         },
-        "config": {
-            "backbone": yolo_v8_backbone_presets.backbone_presets[
-                "yolo_v8_m_backbone"
-            ],
-            "num_classes": 20,
-            "fpn_depth": 2,
-        },
-        "weights_url": "https://storage.googleapis.com/keras-cv/models/yolov8/pascal_voc/yolov8_m_v1.h5",  # noqa: E501
-        "weights_hash": "2891fbd66f71e0b9da0cb02ef3afbccb819e1b8f18204157f643f4ec058a71a8",  # noqa: E501
+        "kaggle_handle": "kaggle://keras/yolov8/keras/yolo_v8_m_pascalvoc/2",
     },
 }
diff --git a/keras_cv/models/object_detection_3d/center_pillar_backbone.py b/keras_cv/models/object_detection_3d/center_pillar_backbone.py
index f7b1f731bb..01554f260b 100644
--- a/keras_cv/models/object_detection_3d/center_pillar_backbone.py
+++ b/keras_cv/models/object_detection_3d/center_pillar_backbone.py
@@ -52,6 +52,10 @@ def __init__(
         input_shape=(None, None, 128),
         **kwargs
     ):
+        self.stackwise_down_blocks = stackwise_down_blocks
+        self.stackwise_down_filters = stackwise_down_filters
+        self.stackwise_up_filters = stackwise_up_filters
+
         input = keras.layers.Input(shape=input_shape)
         x = input
 
@@ -88,6 +92,17 @@ def __init__(
             **kwargs,
         )
 
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "stackwise_down_blocks": self.stackwise_down_blocks,
+                "stackwise_down_filters": self.stackwise_down_filters,
+                "stackwise_up_filters": self.stackwise_up_filters,
+            }
+        )
+        return config
+
     @classproperty
     def presets(cls):
         """Dictionary of preset names and configurations."""
diff --git a/keras_cv/models/object_detection_3d/center_pillar_backbone_presets.py b/keras_cv/models/object_detection_3d/center_pillar_backbone_presets.py
index cc063f43ac..8f875e638a 100644
--- a/keras_cv/models/object_detection_3d/center_pillar_backbone_presets.py
+++ b/keras_cv/models/object_detection_3d/center_pillar_backbone_presets.py
@@ -19,16 +19,8 @@
         "metadata": {
             "description": "An example CenterPillar backbone for WOD.",
             "params": 1277680,
-            "official_name": "WaymoOpenDataset",
-        },
-        "class_name": (
-            "keras_cv.models.object_detection_3d>CenterPillarBackbone"
-        ),
-        "config": {
-            "stackwise_down_blocks": [6, 2, 1],
-            "stackwise_down_filters": [128, 256, 512],
-            "stackwise_up_filters": [512, 256, 256],
-            "input_shape": (None, None, 128),
+            "official_name": "CenterPillar",
         },
+        "kaggle_handle": "gs://keras-cv-kaggle/center_pillar_waymo_open_dataset",  # noqa: E501
     },
 }
diff --git a/keras_cv/models/object_detection_3d/center_pillar_backbone_test.py b/keras_cv/models/object_detection_3d/center_pillar_backbone_test.py
index cc89cb1895..c02046b5e1 100644
--- a/keras_cv/models/object_detection_3d/center_pillar_backbone_test.py
+++ b/keras_cv/models/object_detection_3d/center_pillar_backbone_test.py
@@ -16,13 +16,13 @@
 import tensorflow as tf
 
 from keras_cv.backend import keras
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 from keras_cv.models.object_detection_3d import CenterPillarBackbone
 from keras_cv.tests.test_case import TestCase
 
 
 @pytest.mark.skipif(
-    multi_backend() and keras.backend.backend() == "torch",
+    keras_3() and keras.backend.backend() == "torch",
     reason="CenterPillar does not yet support PyTorch.",
 )
 class CenterPillarBackboneTest(TestCase):
diff --git a/keras_cv/models/object_detection_3d/center_pillar_test.py b/keras_cv/models/object_detection_3d/center_pillar_test.py
index 352d33fbbd..b4be9e6a5e 100644
--- a/keras_cv/models/object_detection_3d/center_pillar_test.py
+++ b/keras_cv/models/object_detection_3d/center_pillar_test.py
@@ -16,7 +16,7 @@
 import tensorflow as tf
 
 from keras_cv.backend import keras
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 from keras_cv.layers.object_detection_3d.voxelization import DynamicVoxelization
 from keras_cv.models.object_detection_3d.center_pillar import (
     MultiClassDetectionHead,
@@ -34,7 +34,7 @@
 
 
 @pytest.mark.skipif(
-    multi_backend() and keras.backend.backend() == "torch",
+    keras_3() and keras.backend.backend() == "torch",
     reason="CenterPillar does not yet support PyTorch.",
 )
 class CenterPillarTest(TestCase):
diff --git a/keras_cv/models/segmentation/__init__.py b/keras_cv/models/segmentation/__init__.py
index f25ee4ea7c..13a9795dda 100644
--- a/keras_cv/models/segmentation/__init__.py
+++ b/keras_cv/models/segmentation/__init__.py
@@ -12,5 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_cv.models.segmentation.basnet import BASNet
 from keras_cv.models.segmentation.deeplab_v3_plus import DeepLabV3Plus
 from keras_cv.models.segmentation.segformer import SegFormer
+from keras_cv.models.segmentation.segment_anything import SAMMaskDecoder
+from keras_cv.models.segmentation.segment_anything import SAMPromptEncoder
+from keras_cv.models.segmentation.segment_anything import SegmentAnythingModel
+from keras_cv.models.segmentation.segment_anything import TwoWayTransformer
diff --git a/keras_cv/models/segmentation/basnet/__init__.py b/keras_cv/models/segmentation/basnet/__init__.py
new file mode 100644
index 0000000000..b51fd6c004
--- /dev/null
+++ b/keras_cv/models/segmentation/basnet/__init__.py
@@ -0,0 +1,15 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_cv.models.segmentation.basnet.basnet import BASNet
diff --git a/keras_cv/models/segmentation/basnet/basnet.py b/keras_cv/models/segmentation/basnet/basnet.py
new file mode 100644
index 0000000000..2803d4425c
--- /dev/null
+++ b/keras_cv/models/segmentation/basnet/basnet.py
@@ -0,0 +1,454 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.models import utils
+from keras_cv.models.backbones.backbone_presets import backbone_presets
+from keras_cv.models.backbones.resnet_v1.resnet_v1_backbone import (
+    apply_basic_block as resnet_basic_block,
+)
+from keras_cv.models.segmentation.basnet.basnet_presets import basnet_presets
+from keras_cv.models.segmentation.basnet.basnet_presets import (
+    presets_no_weights,
+)
+from keras_cv.models.segmentation.basnet.basnet_presets import (
+    presets_with_weights,
+)
+from keras_cv.models.task import Task
+from keras_cv.utils.python_utils import classproperty
+
+
+@keras_cv_export(
+    [
+        "keras_cv.models.BASNet",
+        "keras_cv.models.segmentation.BASNet",
+    ]
+)
+class BASNet(Task):
+    """
+    A Keras model implementing the BASNet architecture for semantic
+    segmentation.
+
+    References:
+        - [BASNet: Boundary-Aware Segmentation Network for Mobile and Web Applications](https://arxiv.org/abs/2101.04704)
+
+    Args:
+        backbone: `keras.Model`. The backbone network for the model that is
+            used as a feature extractor for BASNet prediction encoder. Currently
+            supported backbones are ResNet18 and ResNet34. Default backbone is
+            `keras_cv.models.ResNet34Backbone()`
+            (Note: Do not specify 'input_shape', 'input_tensor', or 'include_rescaling'
+            within the backbone. Please provide these while initializing the
+            'BASNet' model.)
+        num_classes: int, the number of classes for the segmentation model.
+        input_shape: optional shape tuple, defaults to (None, None, 3).
+        input_tensor: optional Keras tensor (i.e., output of `layers.Input()`)
+            to use as image input for the model.
+        include_rescaling: bool, whether to rescale the inputs. If set
+            to `True`, inputs will be passed through a `Rescaling(1/255.0)`
+            layer.
+        projection_filters: int, number of filters in the convolution layer
+            projecting low-level features from the `backbone`.
+        prediction_heads: (Optional) List of `keras.layers.Layer` defining
+            the prediction module head for the model. If not provided, a
+            default head is created with a Conv2D layer followed by resizing.
+        refinement_head: (Optional) a `keras.layers.Layer` defining the
+            refinement module head for the model. If not provided, a default
+            head is created with a Conv2D layer.
+
+    Examples:
+    ```python
+
+    import keras_cv
+
+    images = np.ones(shape=(1, 288, 288, 3))
+    labels = np.zeros(shape=(1, 288, 288, 1))
+
+    # Note: Do not specify 'input_shape', 'input_tensor', or
+    # 'include_rescaling' within the backbone.
+    backbone = keras_cv.models.ResNet34Backbone()
+    model = keras_cv.models.segmentation.BASNet(
+        backbone=backbone,
+        num_classes=1,
+        input_shape=[288, 288, 3],
+        include_rescaling=False
+    )
+
+    # Evaluate model
+    output = model(images)
+    pred_labels = output[0]
+
+    # Train model
+    model.compile(
+        optimizer="adam",
+        loss=keras.losses.BinaryCrossentropy(from_logits=False),
+        metrics=["accuracy"],
+    )
+    model.fit(images, labels, epochs=3)
+        ```
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        backbone,
+        num_classes,
+        input_shape=(None, None, 3),
+        input_tensor=None,
+        include_rescaling=False,
+        projection_filters=64,
+        prediction_heads=None,
+        refinement_head=None,
+        **kwargs,
+    ):
+        if not isinstance(backbone, keras.layers.Layer) or not isinstance(
+            backbone, keras.Model
+        ):
+            raise ValueError(
+                "Argument `backbone` must be a `keras.layers.Layer` instance"
+                f" or `keras.Model`. Received instead"
+                f" backbone={backbone} (of type {type(backbone)})."
+            )
+
+        if backbone.input_shape != (None, None, None, 3):
+            raise ValueError(
+                "Do not specify 'input_shape' or 'input_tensor' within the"
+                " 'BASNet' backbone. \nPlease provide 'input_shape' or"
+                " 'input_tensor' while initializing the 'BASNet' model."
+            )
+
+        inputs = utils.parse_model_inputs(input_shape, input_tensor)
+        x = inputs
+
+        if include_rescaling:
+            x = keras.layers.Rescaling(1 / 255.0)(x)
+
+        if prediction_heads is None:
+            prediction_heads = []
+            for size in (1, 2, 4, 8, 16, 32, 32):
+                head_layers = [
+                    keras.layers.Conv2D(
+                        num_classes, kernel_size=(3, 3), padding="same"
+                    )
+                ]
+                if size != 1:
+                    head_layers.append(
+                        keras.layers.UpSampling2D(
+                            size=size, interpolation="bilinear"
+                        )
+                    )
+                prediction_heads.append(keras.Sequential(head_layers))
+
+        if refinement_head is None:
+            refinement_head = keras.Sequential(
+                [
+                    keras.layers.Conv2D(
+                        num_classes, kernel_size=(3, 3), padding="same"
+                    ),
+                ]
+            )
+
+        # Prediction model.
+        predict_model = basnet_predict(
+            x, backbone, projection_filters, prediction_heads
+        )
+
+        # Refinement model.
+        refine_model = basnet_rrm(
+            predict_model, projection_filters, refinement_head
+        )
+
+        outputs = refine_model.outputs  # Combine outputs.
+        outputs.extend(predict_model.outputs)
+
+        outputs = [
+            keras.layers.Activation("sigmoid", dtype="float32")(_)
+            for _ in outputs
+        ]  # Activations.
+
+        super().__init__(inputs=inputs, outputs=outputs, **kwargs)
+
+        self.backbone = backbone
+        self.num_classes = num_classes
+        self.input_tensor = input_tensor
+        self.include_rescaling = include_rescaling
+        self.projection_filters = projection_filters
+        self.prediction_heads = prediction_heads
+        self.refinement_head = refinement_head
+
+    def get_config(self):
+        return {
+            "backbone": keras.saving.serialize_keras_object(self.backbone),
+            "num_classes": self.num_classes,
+            "input_shape": self.input_shape[1:],
+            "input_tensor": keras.saving.serialize_keras_object(
+                self.input_tensor
+            ),
+            "include_rescaling": self.include_rescaling,
+            "projection_filters": self.projection_filters,
+            "prediction_heads": [
+                keras.saving.serialize_keras_object(prediction_head)
+                for prediction_head in self.prediction_heads
+            ],
+            "refinement_head": keras.saving.serialize_keras_object(
+                self.refinement_head
+            ),
+        }
+
+    @classmethod
+    def from_config(cls, config):
+        if "backbone" in config and isinstance(config["backbone"], dict):
+            input_shape = (None, None, 3)
+            if isinstance(config["backbone"]["config"]["input_shape"], list):
+                input_shape = list(input_shape)
+            if config["backbone"]["config"]["input_shape"] != input_shape:
+                config["input_shape"] = config["backbone"]["config"][
+                    "input_shape"
+                ]
+                config["backbone"]["config"]["input_shape"] = input_shape
+            config["backbone"] = keras.layers.deserialize(config["backbone"])
+
+        if "input_tensor" in config and isinstance(
+            config["input_tensor"], dict
+        ):
+            config["input_tensor"] = keras.layers.deserialize(
+                config["input_tensor"]
+            )
+
+        if "prediction_heads" in config and isinstance(
+            config["prediction_heads"], list
+        ):
+            for i in range(len(config["prediction_heads"])):
+                if isinstance(config["prediction_heads"][i], dict):
+                    config["prediction_heads"][i] = keras.layers.deserialize(
+                        config["prediction_heads"][i]
+                    )
+
+        if "refinement_head" in config and isinstance(
+            config["refinement_head"], dict
+        ):
+            config["refinement_head"] = keras.layers.deserialize(
+                config["refinement_head"]
+            )
+        return super().from_config(config)
+
+    @classproperty
+    def presets(cls):
+        """Dictionary of preset names and configurations."""
+        filtered_backbone_presets = copy.deepcopy(
+            {
+                k: v
+                for k, v in backbone_presets.items()
+                if k in ("resnet18", "resnet34")
+            }
+        )
+
+        return copy.deepcopy({**filtered_backbone_presets, **basnet_presets})
+
+    @classproperty
+    def presets_with_weights(cls):
+        """
+        Dictionary of preset names and configurations that include weights.
+        """
+        return copy.deepcopy(presets_with_weights)
+
+    @classproperty
+    def presets_without_weights(cls):
+        """
+        Dictionary of preset names and configurations that has no weights.
+        """
+        return copy.deepcopy(presets_no_weights)
+
+    @classproperty
+    def backbone_presets(cls):
+        """
+        Dictionary of preset names and configurations of compatible backbones.
+        """
+        filtered_backbone_presets = copy.deepcopy(
+            {
+                k: v
+                for k, v in backbone_presets.items()
+                if k in ("resnet18", "resnet34")
+            }
+        )
+        filtered_presets = copy.deepcopy(filtered_backbone_presets)
+        return filtered_presets
+
+
+def convolution_block(x_input, filters, dilation=1):
+    """
+    Apply convolution + batch normalization + ReLU activation.
+
+    Args:
+        x_input: Input keras tensor.
+        filters: int, number of output filters in the convolution.
+        dilation: int, dilation rate for the convolution operation.
+            Defaults to 1.
+
+    Returns:
+        A tensor with convolution, batch normalization, and ReLU
+        activation applied.
+    """
+    x = keras.layers.Conv2D(
+        filters, (3, 3), padding="same", dilation_rate=dilation
+    )(x_input)
+    x = keras.layers.BatchNormalization()(x)
+    return keras.layers.Activation("relu")(x)
+
+
+def get_resnet_block(_resnet, block_num):
+    """
+    Extract and return a specific ResNet block.
+
+    Args:
+        _resnet: `keras.Model`. ResNet model instance.
+        block_num: int, block number to extract.
+
+    Returns:
+        A Keras Model representing the specified ResNet block.
+    """
+
+    extractor_levels = ["P2", "P3", "P4", "P5"]
+    return keras.models.Model(
+        inputs=_resnet.get_layer(f"v2_stack_{block_num}_block1_1_conv").input,
+        outputs=_resnet.get_layer(
+            _resnet.pyramid_level_inputs[extractor_levels[block_num]]
+        ).output,
+        name=f"resnet_block{block_num + 1}",
+    )
+
+
+def basnet_predict(x_input, backbone, filters, segmentation_heads):
+    """
+    BASNet Prediction Module.
+
+    This module outputs a coarse label map by integrating heavy
+    encoder, bridge, and decoder blocks.
+
+    Args:
+        x_input: Input keras tensor.
+        backbone: `keras.Model`. The backbone network used as a feature
+            extractor for BASNet prediction encoder.
+        filters: int, the number of filters.
+        segmentation_heads: List of `keras.layers.Layer`, A list of Keras
+            layers serving as the segmentation head for prediction module.
+
+
+    Returns:
+        A Keras Model that integrates the encoder, bridge, and decoder
+        blocks for coarse label map prediction.
+    """
+    num_stages = 6
+
+    x = x_input
+
+    # -------------Encoder--------------
+    x = keras.layers.Conv2D(filters, kernel_size=(3, 3), padding="same")(x)
+
+    encoder_blocks = []
+    for i in range(num_stages):
+        if i < 4:  # First four stages are adopted from ResNet backbone.
+            x = get_resnet_block(backbone, i)(x)
+            encoder_blocks.append(x)
+        else:  # Last 2 stages consist of three basic resnet blocks.
+            x = keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))(x)
+            for j in range(3):
+                x = resnet_basic_block(
+                    x,
+                    filters=x.shape[3],
+                    conv_shortcut=False,
+                    name=f"v1_basic_block_{i + 1}_{j + 1}",
+                )
+            encoder_blocks.append(x)
+
+    # -------------Bridge-------------
+    x = convolution_block(x, filters=filters * 8, dilation=2)
+    x = convolution_block(x, filters=filters * 8, dilation=2)
+    x = convolution_block(x, filters=filters * 8, dilation=2)
+    encoder_blocks.append(x)
+
+    # -------------Decoder-------------
+    decoder_blocks = []
+    for i in reversed(range(num_stages)):
+        if i != (num_stages - 1):  # Except first, scale other decoder stages.
+            x = keras.layers.UpSampling2D(size=2, interpolation="bilinear")(x)
+
+        x = keras.layers.concatenate([encoder_blocks[i], x], axis=-1)
+        x = convolution_block(x, filters=filters * 8)
+        x = convolution_block(x, filters=filters * 8)
+        x = convolution_block(x, filters=filters * 8)
+        decoder_blocks.append(x)
+
+    decoder_blocks.reverse()  # Change order from last to first decoder stage.
+    decoder_blocks.append(encoder_blocks[-1])  # Copy bridge to decoder.
+
+    # -------------Side Outputs--------------
+    decoder_blocks = [
+        segmentation_head(decoder_block)  # Prediction segmentation head.
+        for segmentation_head, decoder_block in zip(
+            segmentation_heads, decoder_blocks
+        )
+    ]
+
+    return keras.models.Model(inputs=[x_input], outputs=decoder_blocks)
+
+
+def basnet_rrm(base_model, filters, segmentation_head):
+    """
+    BASNet Residual Refinement Module (RRM).
+
+    This module outputs a fine label map by integrating light encoder,
+    bridge, and decoder blocks.
+
+    Args:
+        base_model: Keras model used as the base or coarse label map.
+        filters: int, the number of filters.
+        segmentation_head: a `keras.layers.Layer`, A Keras layer serving
+            as the segmentation head for refinement module.
+
+    Returns:
+        A Keras Model that constructs the Residual Refinement Module (RRM).
+    """
+    num_stages = 4
+
+    x_input = base_model.output[0]
+
+    # -------------Encoder--------------
+    x = keras.layers.Conv2D(filters, kernel_size=(3, 3), padding="same")(
+        x_input
+    )
+
+    encoder_blocks = []
+    for _ in range(num_stages):
+        x = convolution_block(x, filters=filters)
+        encoder_blocks.append(x)
+        x = keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))(x)
+
+    # -------------Bridge--------------
+    x = convolution_block(x, filters=filters)
+
+    # -------------Decoder--------------
+    for i in reversed(range(num_stages)):
+        x = keras.layers.UpSampling2D(size=2, interpolation="bilinear")(x)
+        x = keras.layers.concatenate([encoder_blocks[i], x], axis=-1)
+        x = convolution_block(x, filters=filters)
+
+    x = segmentation_head(x)  # Refinement segmentation head.
+
+    # ------------- refined = coarse + residual
+    x = keras.layers.Add()([x_input, x])  # Add prediction + refinement output
+
+    return keras.models.Model(inputs=base_model.input, outputs=[x])
diff --git a/keras_cv/models/segmentation/basnet/basnet_presets.py b/keras_cv/models/segmentation/basnet/basnet_presets.py
new file mode 100644
index 0000000000..69d323fd0f
--- /dev/null
+++ b/keras_cv/models/segmentation/basnet/basnet_presets.py
@@ -0,0 +1,51 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""BASNet model preset configurations."""
+
+from keras_cv.models.backbones.resnet_v1 import resnet_v1_backbone_presets
+
+presets_no_weights = {
+    "basnet_resnet18": {
+        "metadata": {
+            "description": "BASNet with a ResNet18 v1 backbone.",
+            "params": 98780872,
+            "official_name": "BASNet",
+            "path": "basnet_resnet18",
+        },
+        "config": {
+            "backbone": resnet_v1_backbone_presets.backbone_presets["resnet18"],
+            "num_classes": 1,
+            "input_shape": (288, 288, 3),
+        },
+    },
+    "basnet_resnet34": {
+        "metadata": {
+            "description": "BASNet with a ResNet34 v1 backbone.",
+            "params": 108896456,
+            "official_name": "BASNet",
+            "path": "basnet_resnet34",
+        },
+        "config": {
+            "backbone": resnet_v1_backbone_presets.backbone_presets["resnet34"],
+            "num_classes": 1,
+            "input_shape": (288, 288, 3),
+        },
+    },
+}
+
+presets_with_weights = {
+    # TODO: Add BASNet preset with weights
+}
+
+basnet_presets = {**presets_no_weights, **presets_with_weights}
diff --git a/keras_cv/models/segmentation/basnet/basnet_test.py b/keras_cv/models/segmentation/basnet/basnet_test.py
new file mode 100644
index 0000000000..88408c134c
--- /dev/null
+++ b/keras_cv/models/segmentation/basnet/basnet_test.py
@@ -0,0 +1,142 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gc
+import os
+
+import numpy as np
+import pytest
+import tensorflow as tf
+from absl.testing import parameterized
+
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.backend.config import keras_3
+from keras_cv.models import BASNet
+from keras_cv.models import ResNet18Backbone
+from keras_cv.tests.test_case import TestCase
+
+
+class BASNetTest(TestCase):
+    def test_basnet_construction(self):
+        backbone = ResNet18Backbone()
+        model = BASNet(
+            input_shape=[64, 64, 3], backbone=backbone, num_classes=1
+        )
+        model.compile(
+            optimizer="adam",
+            loss=keras.losses.BinaryCrossentropy(),
+            metrics=["accuracy"],
+        )
+
+    @pytest.mark.large
+    def test_basnet_call(self):
+        backbone = ResNet18Backbone()
+        model = BASNet(
+            input_shape=[64, 64, 3], backbone=backbone, num_classes=1
+        )
+        images = np.random.uniform(size=(2, 64, 64, 3))
+        _ = model(images)
+        _ = model.predict(images)
+
+    @pytest.mark.large
+    @pytest.mark.filterwarnings("ignore::UserWarning")
+    def test_weights_change(self):
+        input_size = [64, 64, 3]
+        target_size = [64, 64, 1]
+
+        images = np.ones([1] + input_size)
+        labels = np.random.uniform(size=[1] + target_size)
+        ds = tf.data.Dataset.from_tensor_slices((images, labels))
+        ds = ds.repeat(2)
+        ds = ds.batch(2)
+
+        backbone = ResNet18Backbone()
+        model = BASNet(
+            input_shape=[64, 64, 3], backbone=backbone, num_classes=1
+        )
+        model_metrics = ["accuracy"]
+        if keras_3():
+            model_metrics = ["accuracy" for _ in range(8)]
+
+        model.compile(
+            optimizer="adam",
+            loss=keras.losses.BinaryCrossentropy(),
+            metrics=model_metrics,
+        )
+
+        original_weights = model.refinement_head.get_weights()
+        model.fit(ds, epochs=1, batch_size=1)
+        updated_weights = model.refinement_head.get_weights()
+
+        for w1, w2 in zip(original_weights, updated_weights):
+            self.assertNotAllEqual(w1, w2)
+            self.assertFalse(ops.any(ops.isnan(w2)))
+
+    @pytest.mark.large
+    def test_with_model_preset_forward_pass(self):
+        self.skipTest("Skipping preset test until BASNet weights are added.")
+        model = BASNet.from_preset(
+            "basnet_resnet34",
+        )
+        image = np.ones((1, 288, 288, 3))
+        output = ops.expand_dims(ops.argmax(model(image), axis=-1), axis=-1)
+        output = output[0]
+        expected_output = np.zeros((1, 288, 288, 1))
+        self.assertAllClose(output, expected_output)
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        target_size = [64, 64, 3]
+
+        backbone = ResNet18Backbone()
+        model = BASNet(
+            input_shape=[64, 64, 3], backbone=backbone, num_classes=1
+        )
+
+        input_batch = np.ones(shape=[2] + target_size)
+        model_output = model(input_batch)
+
+        save_path = os.path.join(self.get_temp_dir(), "model.keras")
+        if keras_3():
+            model.save(save_path)
+        else:
+            model.save(save_path, save_format="keras_v3")
+        # Free up model memory
+        del model
+        gc.collect()
+        restored_model = keras.models.load_model(save_path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, BASNet)
+
+        # Check that output matches.
+        restored_output = restored_model(input_batch)
+        self.assertAllClose(model_output, restored_output)
+
+
+@pytest.mark.large
+class BASNetSmokeTest(TestCase):
+    @parameterized.named_parameters(
+        *[(preset, preset) for preset in ["resnet18", "resnet34"]]
+    )
+    def test_backbone_preset(self, preset):
+        model = BASNet.from_preset(
+            preset,
+            num_classes=1,
+        )
+        xs = np.random.uniform(size=(1, 128, 128, 3))
+        output = model(xs)[0]
+
+        self.assertEqual(output.shape, (1, 128, 128, 1))
diff --git a/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus.py b/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus.py
index 61d6166a69..35a5dce8f8 100644
--- a/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus.py
+++ b/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus.py
@@ -16,17 +16,26 @@
 
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
+from keras_cv.backend.config import keras_3
 from keras_cv.layers.spatial_pyramid import SpatialPyramidPooling
 from keras_cv.models.backbones.backbone_presets import backbone_presets
 from keras_cv.models.backbones.backbone_presets import (
     backbone_presets_with_weights,
 )
+from keras_cv.models.segmentation.deeplab_v3_plus.deeplab_v3_plus_presets import (  # noqa: E501
+    deeplab_v3_plus_presets,
+)
 from keras_cv.models.task import Task
 from keras_cv.utils.python_utils import classproperty
 from keras_cv.utils.train import get_feature_extractor
 
 
-@keras_cv_export("keras_cv.models.DeepLabV3Plus")
+@keras_cv_export(
+    [
+        "keras_cv.models.DeepLabV3Plus",
+        "keras_cv.models.segmentation.DeepLabV3Plus",
+    ]
+)
 class DeepLabV3Plus(Task):
     """A Keras model implementing the DeepLabV3+ architecture for semantic
     segmentation.
@@ -46,7 +55,7 @@ class DeepLabV3Plus(Task):
             somewhat sensible backbone to use in many cases is the
             `keras_cv.models.ResNet50V2Backbone.from_preset("resnet50_v2_imagenet")`.
         num_classes: int, the number of classes for the detection model. Note
-            that the `num_classes` doesn't contain the background class, and the
+            that the `num_classes` contains the background class, and the
             classes from the data should be represented by integers with range
             [0, `num_classes`).
         projection_filters: int, number of filters in the convolution layer
@@ -229,13 +238,21 @@ def from_config(cls, config):
     @classproperty
     def presets(cls):
         """Dictionary of preset names and configurations."""
-        return copy.deepcopy(backbone_presets)
+        if keras_3():
+            return copy.deepcopy(
+                {**backbone_presets, **deeplab_v3_plus_presets}
+            )
+        else:
+            # TODO: #2246 Deeplab V3 presets don't work in Keras 2
+            return copy.deepcopy({**backbone_presets})
 
     @classproperty
     def presets_with_weights(cls):
         """Dictionary of preset names and configurations that include
         weights."""
-        return copy.deepcopy(backbone_presets_with_weights)
+        return copy.deepcopy(
+            {**backbone_presets_with_weights, **deeplab_v3_plus_presets}
+        )
 
     @classproperty
     def backbone_presets(cls):
diff --git a/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_presets.py b/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_presets.py
new file mode 100644
index 0000000000..13416b0697
--- /dev/null
+++ b/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_presets.py
@@ -0,0 +1,34 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""DeepLabV3Plus presets."""
+
+
+deeplab_v3_plus_presets = {
+    "deeplab_v3_plus_resnet50_pascalvoc": {
+        "metadata": {
+            "description": (
+                "DeeplabV3Plus with a ResNet50 v2 backbone. "
+                "Trained on PascalVOC 2012 Semantic segmentation task, which "
+                "consists of 20 classes and one background class. This model "
+                "achieves a final categorical accuracy of 89.34% and mIoU of "
+                "0.6391 on evaluation dataset. "
+                "This preset is only comptabile with Keras 3."
+            ),
+            "params": 39191488,
+            "official_name": "DeepLabV3Plus",
+            "path": "deeplab_v3_plus",
+        },
+        "kaggle_handle": "kaggle://keras/deeplabv3plus/keras/deeplab_v3_plus_resnet50_pascalvoc/2",  # noqa: E501
+    },
+}
diff --git a/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_test.py b/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_test.py
index b6d8ad0f19..90ec3406d9 100644
--- a/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_test.py
+++ b/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_test.py
@@ -21,6 +21,7 @@
 
 from keras_cv.backend import keras
 from keras_cv.backend import ops
+from keras_cv.backend.config import keras_3
 from keras_cv.models import DeepLabV3Plus
 from keras_cv.models import ResNet18V2Backbone
 from keras_cv.models.backbones.test_backbone_presets import (
@@ -31,8 +32,8 @@
 
 class DeepLabV3PlusTest(TestCase):
     def test_deeplab_v3_plus_construction(self):
-        backbone = ResNet18V2Backbone(input_shape=[512, 512, 3])
-        model = DeepLabV3Plus(backbone=backbone, num_classes=1)
+        backbone = ResNet18V2Backbone(input_shape=[256, 256, 3])
+        model = DeepLabV3Plus(backbone=backbone, num_classes=2)
         model.compile(
             optimizer="adam",
             loss=keras.losses.BinaryCrossentropy(),
@@ -41,15 +42,15 @@ def test_deeplab_v3_plus_construction(self):
 
     @pytest.mark.large
     def test_deeplab_v3_plus_call(self):
-        backbone = ResNet18V2Backbone(input_shape=[512, 512, 3])
-        model = DeepLabV3Plus(backbone=backbone, num_classes=1)
-        images = np.random.uniform(size=(2, 512, 512, 3))
+        backbone = ResNet18V2Backbone(input_shape=[256, 256, 3])
+        model = DeepLabV3Plus(backbone=backbone, num_classes=2)
+        images = np.random.uniform(size=(2, 256, 256, 3))
         _ = model(images)
         _ = model.predict(images)
 
     @pytest.mark.large
     def test_weights_change(self):
-        target_size = [512, 512, 3]
+        target_size = [256, 256, 3]
 
         images = np.ones([1] + target_size)
         labels = np.random.uniform(size=[1] + target_size)
@@ -74,22 +75,35 @@ def test_weights_change(self):
             self.assertNotAllEqual(w1, w2)
             self.assertFalse(ops.any(ops.isnan(w2)))
 
-    @parameterized.named_parameters(
-        ("tf_format", "tf", "model"),
-        ("keras_format", "keras_v3", "model.keras"),
-    )
+    @pytest.mark.large
+    def test_with_model_preset_forward_pass(self):
+        if not keras_3():
+            self.skipTest("TODO: #2246 Not supported for Keras 2")
+        model = DeepLabV3Plus.from_preset(
+            "deeplab_v3_plus_resnet50_pascalvoc",
+            num_classes=21,
+            input_shape=[256, 256, 3],
+        )
+        image = np.ones((1, 256, 256, 3))
+        output = ops.expand_dims(ops.argmax(model(image), axis=-1), axis=-1)
+        expected_output = np.zeros((1, 256, 256, 1))
+        self.assertAllClose(output, expected_output)
+
     @pytest.mark.large  # Saving is slow, so mark these large.
-    def test_saved_model(self, save_format, filename):
-        target_size = [512, 512, 3]
+    def test_saved_model(self):
+        target_size = [256, 256, 3]
 
         backbone = ResNet18V2Backbone(input_shape=target_size)
-        model = DeepLabV3Plus(backbone=backbone, num_classes=1)
+        model = DeepLabV3Plus(backbone=backbone, num_classes=2)
 
         input_batch = np.ones(shape=[2] + target_size)
         model_output = model(input_batch)
 
-        save_path = os.path.join(self.get_temp_dir(), filename)
-        model.save(save_path, save_format=save_format)
+        save_path = os.path.join(self.get_temp_dir(), "model.keras")
+        if keras_3():
+            model.save(save_path)
+        else:
+            model.save(save_path, save_format="keras_v3")
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_cv/models/segmentation/segformer/segformer.py b/keras_cv/models/segmentation/segformer/segformer.py
index 0985b13749..59e4e38f56 100644
--- a/keras_cv/models/segmentation/segformer/segformer.py
+++ b/keras_cv/models/segmentation/segformer/segformer.py
@@ -16,6 +16,7 @@
 
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
+from keras_cv.models import MiTBackbone
 from keras_cv.models.segmentation.segformer.segformer_presets import (  # noqa: E501
     presets,
 )
@@ -27,7 +28,9 @@
 from keras_cv.utils.train import get_feature_extractor
 
 
-@keras_cv_export("keras_cv.models.segmentation.SegFormer")
+@keras_cv_export(
+    ["keras_cv.models.SegFormer", "keras_cv.models.segmentation.SegFormer"]
+)
 class SegFormer(Task):
     """A Keras model implementing the SegFormer architecture for semantic
     segmentation.
@@ -163,6 +166,33 @@ def get_config(self):
         )
         return config
 
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        num_classes,
+        load_weights=None,
+        input_shape=None,
+        **kwargs,
+    ):
+        aliases = {
+            "segformer_b0": "mit_b0",
+            "segformer_b1": "mit_b1",
+            "segformer_b2": "mit_b2",
+            "segformer_b3": "mit_b3",
+            "segformer_b4": "mit_b4",
+            "segformer_b5": "mit_b5",
+        }
+        if preset in aliases:
+            preset = aliases[preset]
+        return super().from_preset(
+            preset,
+            load_weights=load_weights,
+            num_classes=num_classes,
+            input_shape=input_shape,
+            **kwargs,
+        )
+
     @classproperty
     def presets(cls):
         """Dictionary of preset names and configurations."""
@@ -173,3 +203,7 @@ def presets_with_weights(cls):
         """Dictionary of preset names and configurations that include
         weights."""
         return copy.deepcopy(presets_with_weights)
+
+    @classproperty
+    def backbone_presets(cls):
+        return copy.deepcopy(MiTBackbone.presets)
diff --git a/keras_cv/models/segmentation/segformer/segformer_presets.py b/keras_cv/models/segmentation/segformer/segformer_presets.py
index e19e2ec9ba..fde49da7b2 100644
--- a/keras_cv/models/segmentation/segformer/segformer_presets.py
+++ b/keras_cv/models/segmentation/segformer/segformer_presets.py
@@ -25,10 +25,7 @@
             "official_name": "SegFormerB0",
             "path": "segformer_b0",
         },
-        "class_name": "keras_cv>SegFormer",
-        "config": {
-            "backbone": backbone_presets["mit_b0"],
-        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0/2",
     },
     "segformer_b1": {
         "metadata": {
@@ -37,8 +34,7 @@
             "official_name": "SegFormerB1",
             "path": "segformer_b1",
         },
-        "class_name": "keras_cv>SegFormer",
-        "config": {"backbone": backbone_presets["mit_b1"]},
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b1/2",
     },
     "segformer_b2": {
         "metadata": {
@@ -47,8 +43,7 @@
             "official_name": "SegFormerB2",
             "path": "segformer_b2",
         },
-        "class_name": "keras_cv>SegFormer",
-        "config": {"backbone": backbone_presets["mit_b2"]},
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b2/2",
     },
     "segformer_b3": {
         "metadata": {
@@ -57,8 +52,7 @@
             "official_name": "SegFormerB3",
             "path": "segformer_b3",
         },
-        "class_name": "keras_cv>SegFormer",
-        "config": {"backbone": backbone_presets["mit_b3"]},
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b3/2",
     },
     "segformer_b4": {
         "metadata": {
@@ -67,8 +61,7 @@
             "official_name": "SegFormerB4",
             "path": "segformer_b4",
         },
-        "class_name": "keras_cv>SegFormer",
-        "config": {"backbone": backbone_presets["mit_b4"]},
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b4/2",
     },
     "segformer_b5": {
         "metadata": {
@@ -77,8 +70,7 @@
             "official_name": "SegFormerB5",
             "path": "segformer_b5",
         },
-        "class_name": "keras_cv>SegFormer",
-        "config": {"backbone": backbone_presets["mit_b5"]},
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b5/2",
     },
 }
 
@@ -92,14 +84,12 @@
             "official_name": "SegFormerB0",
             "path": "segformer_b0",
         },
-        "class_name": "keras_cv>SegFormer",
-        "config": {
-            "backbone": backbone_presets["mit_b0_imagenet"],
-        },
+        "kaggle_handle": "kaggle://keras/segformer/keras/segformer_b0_imagenet/2",  # noqa: E501
     },
 }
 
 presets = {
+    **backbone_presets,  # Add MiTBackbone presets
     **presets_no_weights,
     **presets_with_weights,
 }
diff --git a/keras_cv/models/segmentation/segformer/segformer_test.py b/keras_cv/models/segmentation/segformer/segformer_test.py
index 0990e0e88f..5f33b033da 100644
--- a/keras_cv/models/segmentation/segformer/segformer_test.py
+++ b/keras_cv/models/segmentation/segformer/segformer_test.py
@@ -20,6 +20,7 @@
 
 from keras_cv.backend import keras
 from keras_cv.backend import ops
+from keras_cv.backend.config import keras_3
 from keras_cv.models import MiTBackbone
 from keras_cv.models import SegFormer
 from keras_cv.tests.test_case import TestCase
@@ -35,13 +36,35 @@ def test_segformer_construction(self):
             metrics=["accuracy"],
         )
 
+    def test_segformer_preset_construction(self):
+        model = SegFormer.from_preset(
+            "segformer_b0", num_classes=1, input_shape=[512, 512, 3]
+        )
+        model.compile(
+            optimizer="adam",
+            loss=keras.losses.BinaryCrossentropy(),
+            metrics=["accuracy"],
+        )
+
+    def test_segformer_preset_error(self):
+        with self.assertRaises(TypeError):
+            _ = SegFormer.from_preset("segformer_b0")
+
     @pytest.mark.large
     def test_segformer_call(self):
-        backbone = MiTBackbone.from_preset("mit_b0", input_shape=[512, 512, 3])
-        model = SegFormer(backbone=backbone, num_classes=1)
-        images = np.random.uniform(size=(2, 512, 512, 3))
-        _ = model(images)
-        _ = model.predict(images)
+        backbone = MiTBackbone.from_preset("mit_b0")
+        mit_model = SegFormer(backbone=backbone, num_classes=1)
+
+        images = np.random.uniform(size=(2, 224, 224, 3))
+        mit_output = mit_model(images)
+        mit_pred = mit_model.predict(images)
+
+        seg_model = SegFormer.from_preset("segformer_b0", num_classes=1)
+        seg_output = seg_model(images)
+        seg_pred = seg_model.predict(images)
+
+        self.assertAllClose(mit_output, seg_output)
+        self.assertAllClose(mit_pred, seg_pred)
 
     @pytest.mark.large
     def test_weights_change(self):
@@ -81,7 +104,33 @@ def test_saved_model(self):
         model_output = model(input_batch)
 
         save_path = os.path.join(self.get_temp_dir(), "model.keras")
-        model.save(save_path, save_format="keras_v3")
+        if keras_3():
+            model.save(save_path)
+        else:
+            model.save(save_path, save_format="keras_v3")
+        restored_model = keras.models.load_model(save_path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, SegFormer)
+
+        # Check that output matches.
+        restored_output = restored_model(input_batch)
+        self.assertAllClose(model_output, restored_output)
+
+    @pytest.mark.large  # Saving is slow, so mark these large.
+    def test_preset_saved_model(self):
+        target_size = [224, 224, 3]
+
+        model = SegFormer.from_preset("segformer_b0", num_classes=1)
+
+        input_batch = np.ones(shape=[2] + target_size)
+        model_output = model(input_batch)
+
+        save_path = os.path.join(self.get_temp_dir(), "model.keras")
+        if keras_3():
+            model.save(save_path)
+        else:
+            model.save(save_path, save_format="keras_v3")
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
diff --git a/keras_cv/models/segmentation/segment_anything/__init__.py b/keras_cv/models/segmentation/segment_anything/__init__.py
new file mode 100644
index 0000000000..945955261f
--- /dev/null
+++ b/keras_cv/models/segmentation/segment_anything/__init__.py
@@ -0,0 +1,26 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_cv.models.segmentation.segment_anything.sam import (
+    SegmentAnythingModel,
+)
+from keras_cv.models.segmentation.segment_anything.sam_mask_decoder import (
+    SAMMaskDecoder,
+)
+from keras_cv.models.segmentation.segment_anything.sam_prompt_encoder import (
+    SAMPromptEncoder,
+)
+from keras_cv.models.segmentation.segment_anything.sam_transformer import (
+    TwoWayTransformer,
+)
diff --git a/keras_cv/models/segmentation/segment_anything/data/sam_base_out_iou_pred.npy b/keras_cv/models/segmentation/segment_anything/data/sam_base_out_iou_pred.npy
new file mode 100644
index 0000000000..2f9f2114c0
Binary files /dev/null and b/keras_cv/models/segmentation/segment_anything/data/sam_base_out_iou_pred.npy differ
diff --git a/keras_cv/models/segmentation/segment_anything/data/sam_base_out_masks.npy b/keras_cv/models/segmentation/segment_anything/data/sam_base_out_masks.npy
new file mode 100644
index 0000000000..d191e2ca70
Binary files /dev/null and b/keras_cv/models/segmentation/segment_anything/data/sam_base_out_masks.npy differ
diff --git a/keras_cv/models/segmentation/segment_anything/sam.py b/keras_cv/models/segmentation/segment_anything/sam.py
new file mode 100644
index 0000000000..84b808419a
--- /dev/null
+++ b/keras_cv/models/segmentation/segment_anything/sam.py
@@ -0,0 +1,290 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+import numpy as np
+
+from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.models.backbones.backbone_presets import backbone_presets
+from keras_cv.models.backbones.backbone_presets import (
+    backbone_presets_with_weights,
+)
+from keras_cv.models.segmentation.segment_anything.sam_presets import (
+    sam_presets,
+)
+from keras_cv.models.task import Task
+from keras_cv.utils.python_utils import classproperty
+
+
+@keras_cv_export(
+    [
+        "keras_cv.models.SegmentAnythingModel",
+        "keras_cv.models.segmentation.SegmentAnythingModel",
+    ],
+    package="keras_cv.models",
+)
+class SegmentAnythingModel(Task):
+    """
+    The Segment Anything (SAM) Model.
+
+    Args:
+        backbone (keras_cv.models.Backbone): A feature extractor for the input
+            images.
+        prompt_encoder (keras_cv.models.SAMPromptEncoder): A Keras layer to
+            compute embeddings for points, box, and mask prompt.
+        mask_decoder (keras_cv.models.SAMMaskDecoder): A Keras layer to
+            generate segmentation masks given the embeddings generated by the
+            backbone and the prompt encoder.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+
+    Examples:
+
+    >>> import numpy as np
+    >>> from keras_cv.models import ViTDetBBackbone
+    >>> from keras_cv.models import SAMPromptEncoder
+    >>> from keras_cv.models import SAMMaskDecoder
+
+    Create all the components of the SAM model:
+
+    >>> backbone = ViTDetBBackbone()
+    >>> prompt_encoder = SAMPromptEncoder()
+    >>> mask_decoder = SAMMaskDecoder()
+
+    Instantiate the model:
+
+    >>> sam = SegmentAnythingModel(
+    ...     backbone=backbone,
+    ...     prompt_encoder=prompt_encoder,
+    ...     mask_decoder=mask_decoder
+    ... )
+
+    Define the input of the backbone. This must be a batch of images of shape
+    `(1024, 1024, 3)` for the ViT backbone we are using:
+
+    >>> image = np.ones((1, 1024, 1024, 3))
+
+    SAM works by prompting the input images. There are three ways to prompt:
+
+    (1) Labelled Points: Foreground points (points with label 1) are encoded
+        such that the output masks generated by the mask decoder contain them
+        and background points (points with label 0) are encoded such that the
+        generated masks don't contain them.
+    (2) Box: A box tells the model which part/crop of the image to segment.
+    (3) Mask: An input mask can be used to refine the output of the mask
+        decoder.
+
+    These prompts can be mixed and matched but at least one of the prompts
+    must be present. To turn off a particular prompt, simply exclude it from
+    the inputs to the model.
+
+    # TODO(ianstenbit): Remove the need for the `1` axes, and fix the box shape.
+    (1) For points prompts, the expected shape is `(batch, num_points, 2)`.
+        The labels must have a corresponding shape of `(batch, num_points)`.
+    (2) For box prompt, the expected shape is `(batch, 1, 2, 2)`.
+    (3) Similarly, mask prompts have shape `(batch, 1, H, W, 1)`.
+
+    For example, to pass in all the prompts, do:
+
+    >>> points = np.array([[[512., 512.], [100., 100.]]])
+    >>> # For labels: 1 means foreground point, 0 means background
+    >>> labels = np.array([[1., 0.]])
+    >>> box = np.array([[[[384., 384.], [640., 640.]]]])
+    >>> input_mask = np.ones((1, 1, 256, 256, 1))
+
+    Prepare an input dictionary:
+
+    >>> inputs = {
+    ...     "images": image,
+    ...     "points": points,
+    ...     "labels": labels,
+    ...     "boxes": box,
+    ...     "masks": input_mask
+    ... }
+    ...
+    >>> outputs = sam.predict(inputs)
+    >>> masks, iou_pred = outputs["masks"], outputs["iou_pred"]
+
+    The first mask in the output `masks` (i.e. `masks[:, 0, ...]`) is the best
+    mask predicted by the model based on the prompts. Other `masks`
+    (i.e. `masks[:, 1:, ...]`) are alternate predictions that can be used if
+    they are desired over the first one.
+
+    Now, in case of only points and box prompts, simply exclude the masks:
+
+    >>> inputs = {
+    ...     "images": image,
+    ...     "points": points,
+    ...     "labels": labels,
+    ...     "boxes": box,
+    ... }
+    ...
+    >>> outputs = sam.predict(inputs)
+    >>> masks, iou_pred = outputs["masks"], outputs["iou_pred"]
+
+    # TODO(ianstenbit): Remove the need for this padding.
+    Another example is that only points prompts are present.
+    Note that if point prompts are present but no box prompt is present, the
+    points must be padded using a zero point and -1 label:
+
+    >>> padded_points = np.concatenate(
+    ...     [points, np.zeros((1, 1, 2))], axis=1
+    ... )
+    ...
+    >>> padded_labels = np.concatenate(
+    ...     [labels, -np.ones((1, 1))], axis=1
+    ... )
+    >>> inputs = {
+    ...     "images": image,
+    ...     "points": padded_points,
+    ...     "labels": padded_labels,
+    ... }
+    ...
+    >>> outputs = sam.predict(inputs)
+    >>> masks, iou_pred = outputs["masks"], outputs["iou_pred"]
+
+    Note that the segment anything model only supports inference and training
+    isn't support yet. So, calling the `fit` method will fail for now.
+    """  # noqa: E501
+
+    def __init__(self, *, backbone, prompt_encoder, mask_decoder, **kwargs):
+        # Get the image encoder input -- Images
+        backbone_input = backbone.input
+
+        # Define the prompt encoder inputs -- Prompts
+        prompt_inputs = {
+            "points": keras.Input(shape=[None, 2], name="points"),
+            "labels": keras.Input(shape=[None], name="labels"),
+            "boxes": keras.Input(shape=[None, 2, 2], name="boxes"),
+            "masks": keras.Input(shape=[None, None, None, 1], name="masks"),
+        }
+
+        # All Inputs -- Images + Prompts
+        all_inputs = {"images": backbone_input}
+        all_inputs.update(prompt_inputs)
+
+        # Build the prompt encoder
+        prompt_embeddings = prompt_encoder(prompt_inputs)
+
+        # Define the mask decoder inputs
+        mask_decoder_inputs = {
+            "image_embeddings": backbone.output,
+            "image_pe": prompt_embeddings["dense_positional_embeddings"],
+            "sparse_prompt_embeddings": prompt_embeddings["sparse_embeddings"],
+            "dense_prompt_embeddings": prompt_embeddings["dense_embeddings"],
+        }
+
+        # Build the mask decoder
+        outputs = mask_decoder(mask_decoder_inputs)
+
+        super().__init__(inputs=all_inputs, outputs=outputs, **kwargs)
+
+        self.backbone = backbone
+        self.prompt_encoder = prompt_encoder
+        self.mask_decoder = mask_decoder
+
+    # TODO(ianstenbit): Do something more elegant to handle empty prompts.
+    def predict_step(self, *args, **kwargs):
+        if len(args) == 2:
+            args = (args[0], _add_placeholder_prompts(args[-1]))
+        else:
+            args = (_add_placeholder_prompts(args[0]),)
+
+        return super().predict_step(*args, **kwargs)
+
+    def fit(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Segment Anything Model only supports inference for now. Training"
+            " the model isn't supported yet."
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "backbone": keras.saving.serialize_keras_object(self.backbone),
+                "prompt_encoder": keras.saving.serialize_keras_object(
+                    self.prompt_encoder
+                ),
+                "mask_decoder": keras.saving.serialize_keras_object(
+                    self.mask_decoder
+                ),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        config.update(
+            {
+                "prompt_encoder": keras.layers.deserialize(
+                    config["prompt_encoder"]
+                ),
+                "mask_decoder": keras.layers.deserialize(
+                    config["mask_decoder"]
+                ),
+            }
+        )
+        return super().from_config(config)
+
+    @classproperty
+    def presets(cls):
+        """Dictionary of preset names and configurations."""
+        return copy.deepcopy({**backbone_presets, **sam_presets})
+
+    @classproperty
+    def presets_with_weights(cls):
+        """Dictionary of preset names and configurations that include
+        weights."""
+        return copy.deepcopy({**backbone_presets_with_weights, **sam_presets})
+
+    @classproperty
+    def backbone_presets(cls):
+        """Dictionary of preset names and configurations of compatible
+        backbones."""
+        return copy.deepcopy(backbone_presets)
+
+
+def _add_placeholder_prompts(inputs):
+    """Adds placeholder prompt inputs for a call to SAM.
+
+    Because SAM is a functional subclass model, all inputs must be specified in
+    calls to the model. However, prompt inputs are all optional, so we have to
+    add placeholders when they're not specified by the user.
+    """
+    inputs = inputs.copy()
+
+    # Get the batch shape based on the image input
+    B = ops.shape(inputs["images"])[0]
+
+    # The type of the placeholders must match the existing inputs with respect
+    # to whether or not they are tensors (as opposed to Numpy arrays).
+    zeros = ops.zeros if ops.is_tensor(inputs["images"]) else np.zeros
+
+    # Fill in missing inputs.
+    if "points" not in inputs:
+        inputs["points"] = zeros((B, 0, 2))
+    if "labels" not in inputs:
+        inputs["labels"] = zeros((B, 0))
+    if "boxes" not in inputs:
+        inputs["boxes"] = zeros((B, 0, 2, 2))
+    if "masks" not in inputs:
+        inputs["masks"] = zeros((B, 0, 256, 256, 1))
+
+    return inputs
diff --git a/keras_cv/models/segmentation/segment_anything/sam_layers.py b/keras_cv/models/segmentation/segment_anything/sam_layers.py
new file mode 100644
index 0000000000..fffc4faee5
--- /dev/null
+++ b/keras_cv/models/segmentation/segment_anything/sam_layers.py
@@ -0,0 +1,347 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.layers.vit_det_layers import MLP
+
+
+@keras_cv_export(
+    "keras_cv.layers.MultiHeadAttentionWithDownsampling",
+    package="keras_cv.layers",
+)
+class MultiHeadAttentionWithDownsampling(keras.layers.Layer):
+    """Multi-Head Attention with downsampling.
+
+    An attention layer that allows for downscaling the size of the embedding
+    after projection to queries, keys, and values.
+
+    This layer first downscales the features of input queries, keys, and
+    values using a dense layer. Multi-head attention is then performed
+    and the attention map is projected back (upscaled) to the number of
+    input features.
+
+    Args:
+        num_heads (int): Number of attention heads.
+        key_dim (int): Size of each attention head for query, key, and
+            value.
+        downsample_rate (int, optional): The factor by which to downscale the
+            input features i.e. the input features of size `key_dim` are
+            projected down to `key_dim // downsample_rate`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+    """  # noqa: E501
+
+    def __init__(self, num_heads, key_dim, downsample_rate=1, **kwargs):
+        super().__init__(**kwargs)
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+        self.downsample_rate = downsample_rate
+        self.internal_dims = key_dim // downsample_rate
+
+        # Downsample
+        self.query_proj = keras.layers.Dense(
+            self.internal_dims * self.num_heads
+        )
+        self.key_proj = keras.layers.Dense(self.internal_dims * self.num_heads)
+        self.value_proj = keras.layers.Dense(
+            self.internal_dims * self.num_heads
+        )
+
+        # Upsample
+        self.out_proj = keras.layers.Dense(self.key_dim * self.num_heads)
+
+    def build(self, input_shape=None):
+        self.query_proj.build([None, None, self.num_heads * self.key_dim])
+        self.key_proj.build([None, None, self.num_heads * self.key_dim])
+        self.value_proj.build([None, None, self.num_heads * self.key_dim])
+        self.out_proj.build([None, None, self.internal_dims * self.num_heads])
+        self.built = True
+
+    def __separate_heads(self, x):
+        shape = ops.shape(x)
+        B, N, C = shape[0], shape[1], shape[2]
+        x = ops.reshape(x, (B, N, self.num_heads, C // self.num_heads))
+        return ops.transpose(x, axes=(0, 2, 1, 3))
+
+    def __recombine_heads(self, x):
+        shape = ops.shape(x)
+        B, N_H, N_T, C_PH = shape[0], shape[1], shape[2], shape[3]
+        x = ops.transpose(x, axes=(0, 2, 1, 3))
+        return ops.reshape(x, (B, N_T, N_H * C_PH))
+
+    def call(self, query, value, key):
+        query = self.query_proj(query)
+        key = self.key_proj(key)
+        value = self.value_proj(value)
+
+        # Separate into heads
+        query = self.__separate_heads(query)
+        key = self.__separate_heads(key)
+        value = self.__separate_heads(value)
+
+        # Attention
+        C_PH = ops.shape(query)[-1]
+        out = query @ ops.transpose(key, (0, 1, 3, 2))
+        out = out / ops.sqrt(ops.cast(C_PH, dtype=self.compute_dtype))
+        out = ops.softmax(out, axis=-1)
+
+        # Get output
+        attention_map = out @ value
+        attention_map = self.__recombine_heads(attention_map)
+        return self.out_proj(attention_map)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "key_dim": self.key_dim,
+                "downsample_rate": self.downsample_rate,
+            }
+        )
+        return config
+
+
+@keras_cv_export(
+    "keras_cv.layers.TwoWayMultiHeadAttention", package="keras_cv.layers"
+)
+class TwoWayMultiHeadAttention(keras.layers.Layer):
+    """Two-way multi-head attention layer.
+
+    Args:
+        num_heads (int): Number of attention heads.
+        key_dim (int): Size of each attention head for query, key, and
+            value.
+        mlp_dim (int): Number of hidden dims to use in the mlp block.
+        skip_first_layer_pe (bool): A boolean indicating whether to skip the
+            first layer positional embeddings.
+        attention_downsample_rate (int, optional): The downsample rate to use
+            in the attention layers. Defaults to 2.
+        activation (str, optional): The activation for the mlp block's output
+            layer. Defaults to "relu".
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        num_heads,
+        key_dim,
+        mlp_dim,
+        skip_first_layer_pe,
+        attention_downsample_rate=2,
+        activation="relu",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_heads = num_heads
+        self.key_dim = key_dim
+        self.mlp_dim = mlp_dim
+        self.skip_first_layer_pe = skip_first_layer_pe
+        self.attention_downsample_rate = attention_downsample_rate
+        self.activation = activation
+
+        self.self_attention = MultiHeadAttentionWithDownsampling(
+            num_heads=num_heads, key_dim=key_dim
+        )
+        self.layer_norm1 = keras.layers.LayerNormalization(epsilon=1e-5)
+        self.cross_attention_token_to_image = (
+            MultiHeadAttentionWithDownsampling(
+                num_heads=num_heads,
+                key_dim=key_dim,
+                downsample_rate=attention_downsample_rate,
+            )
+        )
+        self.layer_norm2 = keras.layers.LayerNormalization(epsilon=1e-5)
+
+        self.mlp_block = MLP(
+            mlp_dim,
+            key_dim * num_heads,
+            num_layers=2,
+            activation=activation,
+        )
+
+        self.layer_norm3 = keras.layers.LayerNormalization(epsilon=1e-5)
+        self.cross_attention_image_to_token = (
+            MultiHeadAttentionWithDownsampling(
+                num_heads=num_heads,
+                key_dim=key_dim,
+                downsample_rate=attention_downsample_rate,
+            )
+        )
+        self.layer_norm4 = keras.layers.LayerNormalization(epsilon=1e-5)
+
+    def build(self, input_shape=None):
+        self.self_attention.build()
+        self.layer_norm1.build([None, None, self.num_heads * self.key_dim])
+        self.cross_attention_token_to_image.build()
+        self.layer_norm2.build([None, None, self.num_heads * self.key_dim])
+        self.mlp_block.build([None, None, self.num_heads * self.key_dim])
+        self.layer_norm3.build([None, None, self.num_heads * self.key_dim])
+        self.cross_attention_image_to_token.build()
+        self.layer_norm4.build([None, None, self.num_heads * self.key_dim])
+        self.built = True
+
+    def call(self, queries, keys, query_pe, key_pe):
+        if self.skip_first_layer_pe:
+            queries = self.self_attention(
+                query=queries, value=queries, key=queries
+            )
+        else:
+            queries_with_pe = queries + query_pe
+            attention_map = self.self_attention(
+                query=queries_with_pe, key=queries_with_pe, value=queries
+            )
+            queries = queries + attention_map
+        queries = self.layer_norm1(queries)
+
+        queries_with_pe = queries + query_pe
+        keys_with_pe = keys + key_pe
+        attention_map = self.cross_attention_token_to_image(
+            query=queries_with_pe, key=keys_with_pe, value=keys
+        )
+        queries = queries + attention_map
+        queries = self.layer_norm2(queries)
+
+        mlp_out = self.mlp_block(queries)
+        queries = queries + mlp_out
+        queries = self.layer_norm3(queries)
+
+        queries_with_pe = queries + query_pe
+        keys_with_pe = keys + key_pe
+        attention_map = self.cross_attention_image_to_token(
+            query=keys_with_pe, key=queries_with_pe, value=queries
+        )
+        keys = keys + attention_map
+        keys = self.layer_norm4(keys)
+
+        return queries, keys
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "key_dim": self.key_dim,
+                "mlp_dim": self.mlp_dim,
+                "skip_first_layer_pe": self.skip_first_layer_pe,
+                "attention_downsample_rate": self.attention_downsample_rate,
+                "activation": self.activation,
+            }
+        )
+        return config
+
+
+@keras_cv_export(
+    "keras_cv.layers.RandomFrequencyPositionalEmbeddings",
+    package="keras_cv.layers",
+)
+class RandomFrequencyPositionalEmbeddings(keras.layers.Layer):
+    """Positional encoding using random spatial frequencies.
+
+    This layer maps coordinates/points in 2D space to positional
+    encodings using random spatial frequencies.
+
+    Args:
+        num_positional_features (int): Number of positional features
+            in the output.
+        scale (float): The standard deviation of the random frequencies.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+    """  # noqa: E501
+
+    def __init__(self, num_positional_features, scale, **kwargs):
+        super().__init__(**kwargs)
+        self.num_positional_features = num_positional_features
+        self.scale = scale
+        self.positional_encoding_gaussian_matrix = self.add_weight(
+            name="positional_encoding_gaussian_matrix",
+            shape=(2, self.num_positional_features),
+            dtype=self.variable_dtype,
+            trainable=False,
+            initializer=keras.initializers.get("normal"),
+        )
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def __positional_encodings(self, coords):
+        coords = coords * 2 - 1
+        coords = coords @ ops.cast(
+            self.positional_encoding_gaussian_matrix, dtype=self.compute_dtype
+        )
+        coords = coords * (2 * math.pi)
+        return ops.concatenate([ops.sin(coords), ops.cos(coords)], axis=-1)
+
+    def call(self, size):
+        return self.encode_image(size)
+
+    def encode_image(self, size):
+        """Generate a positional encoding for an image of any given size.
+
+        Args:
+            size (tuple[int, int]): The size of the image.
+
+        Returns:
+            tensor: Positional encoding of the image.
+        """
+        H, W = size
+        grid = ops.ones(shape=(H, W), dtype=self.compute_dtype)
+        y_embed = ops.cumsum(grid, axis=0) - 0.5
+        x_embed = ops.cumsum(grid, axis=1) - 0.5
+        y_embed = y_embed / ops.cast(H, self.compute_dtype)
+        x_embed = x_embed / ops.cast(W, self.compute_dtype)
+        return self.__positional_encodings(
+            ops.stack([x_embed, y_embed], axis=-1)
+        )
+
+    def encode_coordinates(self, coords_input, image_size):
+        """Positionally encode points that are not normalized to `[0, 1]`.
+
+        Args:
+            coords_input (tensor): 2D coordinates/points to map.
+            image_size (tuple[int, int]): Height and width of the image
+                being prompted.
+
+        Returns:
+            tensor: Positional encodings of the normalized coordinates.
+        """
+        coords_normalized = ops.stack(
+            [
+                coords_input[..., 0] / image_size[1],
+                coords_input[..., 1] / image_size[0],
+            ],
+            axis=-1,
+        )
+        return self.__positional_encodings(coords_normalized)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_positional_features": self.num_positional_features,
+                "scale": self.scale,
+            }
+        )
+        return config
diff --git a/keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py b/keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py
new file mode 100644
index 0000000000..141f90addc
--- /dev/null
+++ b/keras_cv/models/segmentation/segment_anything/sam_mask_decoder.py
@@ -0,0 +1,240 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.layers.vit_det_layers import MLP
+from keras_cv.models.segmentation.segment_anything.sam_transformer import (
+    TwoWayTransformer,
+)
+
+
+@keras_cv_export("keras_cv.models.SAMMaskDecoder", package="keras_cv.models")
+class SAMMaskDecoder(keras.layers.Layer):
+    """Mask decoder for the Segment Anything Model (SAM).
+
+    This lightweight module efficiently maps the image embedding and a set of
+    prompt embeddings to an output mask. Before applying the transformer
+    decoder, the layer first inserts into the set of prompt embeddings a
+    learned output token embedding that will be used at the decoder's output.
+    For simplicity, these embeddings (not including the image embedding) are
+    collectively called "tokens".
+
+    The image embeddings, positional image embeddings, and tokens are passed
+    through a transformer decoder. After running the decoder, the layer
+    upsamples the updated image embedding by 4x with two transposed
+    convolutional layers (now it's downscaled 4x relative to the input
+    image). Then, the tokens attend once more to the image embedding and
+    the updated output token embedding are passed to a small 3-layer MLP that
+    outputs a vector matching the channel dimension of the upscaled image
+    embedding. Finally, a mask is predicted with a spatially point-wise
+    product between the upscaled image embedding and the MLP's output.
+
+    Args:
+        transformer_dim (int, optional): The number of input features to the
+            transformer decoder. Defaults to `256`.
+        transformer (keras.layers.Layer, optional): A transformer decoder.
+            Defaults to `None`. When `None`, a
+            `keras_cv.models.TwoWayTransformer` layer is used.
+        num_multimask_outputs (int, optional): Number of multimask outputs.
+            The model would generate these many extra masks. The total masks
+            generated by the model are `1 + num_multimask_outputs`. Defaults
+            to `3`.
+        iou_head_depth (int, optional): The depth of the dense net used to
+            predict the IoU confidence score. Defaults to `3`.
+        iou_head_hidden_dim (int, optional): The number of units in the hidden
+            layers used in the dense net to predict the IoU confidence score.
+            Defaults to `256`.
+        activation (str, optional): Activation to use in the mask upscaler
+            network. Defaults to `"gelu"`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        *,
+        transformer_dim=256,
+        transformer=None,
+        num_multimask_outputs=3,
+        iou_head_depth=3,
+        iou_head_hidden_dim=256,
+        activation="gelu",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.transformer_dim = transformer_dim
+        if transformer is None:
+            transformer = TwoWayTransformer()
+        self.transformer = transformer
+        self.num_multimask_outputs = num_multimask_outputs
+        self.iou_head_depth = iou_head_depth
+        self.iou_head_hidden_dim = iou_head_hidden_dim
+        self.activation = activation
+
+        self.iou_token = keras.layers.Embedding(1, transformer_dim)
+        self.num_mask_tokens = num_multimask_outputs + 1
+        self.mask_tokens = keras.layers.Embedding(
+            self.num_mask_tokens, transformer_dim
+        )
+
+        self.output_upscaling = keras.models.Sequential(
+            [
+                keras.layers.Conv2DTranspose(
+                    transformer_dim // 4, kernel_size=2, strides=2
+                ),
+                keras.layers.LayerNormalization(epsilon=1e-6),
+                keras.layers.Activation(activation),
+                keras.layers.Conv2DTranspose(
+                    transformer_dim // 8, kernel_size=2, strides=2
+                ),
+                keras.layers.Activation(activation),
+            ]
+        )
+
+        self.output_hypernetworks_mlps = [
+            MLP(transformer_dim, transformer_dim // 8, 3)
+            for _ in range(self.num_mask_tokens)
+        ]
+
+        self.iou_prediction_head = MLP(
+            iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth
+        )
+
+    def build(self, input_shape=None):
+        self.transformer.build()
+        self.iou_token.build([None])
+        self.mask_tokens.build([None])
+        self.output_upscaling.build([None, None, None, self.transformer_dim])
+        for mlp in self.output_hypernetworks_mlps:
+            mlp.build([None, self.transformer_dim])
+        self.iou_prediction_head.build([None, self.transformer_dim])
+        self.built = True
+
+    def call(self, inputs):
+        image_embeddings = inputs["image_embeddings"]
+        image_pe = inputs["image_pe"]
+        sparse_prompt_embeddings = inputs["sparse_prompt_embeddings"]
+        dense_prompt_embeddings = inputs["dense_prompt_embeddings"]
+
+        masks, iou_pred = self._predict_masks(
+            image_embeddings=image_embeddings,
+            image_pe=image_pe,
+            sparse_prompt_embeddings=sparse_prompt_embeddings,
+            dense_prompt_embeddings=dense_prompt_embeddings,
+        )
+
+        return {"masks": masks, "iou_pred": iou_pred}
+
+    def _predict_masks(
+        self,
+        image_embeddings,
+        image_pe,
+        sparse_prompt_embeddings,
+        dense_prompt_embeddings,
+    ):
+        indices_iou = ops.arange(1, dtype="int32")
+        indices_mask = ops.arange(self.num_mask_tokens, dtype="int32")
+
+        output_tokens = ops.concatenate(
+            [self.iou_token(indices_iou), self.mask_tokens(indices_mask)],
+            axis=0,
+        )
+        output_tokens = ops.broadcast_to(
+            output_tokens[None, ...],
+            shape=(
+                ops.shape(sparse_prompt_embeddings)[0],
+                ops.shape(output_tokens)[0],
+                ops.shape(output_tokens)[1],
+            ),
+        )
+        tokens = ops.concatenate(
+            [output_tokens, sparse_prompt_embeddings], axis=1
+        )
+
+        source = ops.broadcast_to(
+            image_embeddings,
+            shape=(
+                ops.shape(tokens)[0],
+                ops.shape(image_embeddings)[1],
+                ops.shape(image_embeddings)[2],
+                ops.shape(image_embeddings)[3],
+            ),
+        )
+        source = source + dense_prompt_embeddings
+        positional_source = ops.broadcast_to(
+            image_pe,
+            shape=(
+                ops.shape(tokens)[0],
+                ops.shape(image_embeddings)[1],
+                ops.shape(image_embeddings)[2],
+                ops.shape(image_embeddings)[3],
+            ),
+        )
+        shape = ops.shape(source)
+        B, H, W, C = shape[0], shape[1], shape[2], shape[3]
+
+        hidden_state, source = self.transformer(
+            source, positional_source, tokens
+        )
+        iou_token_out = hidden_state[:, 0, :]
+        mask_tokens_out = hidden_state[:, 1 : (1 + self.num_mask_tokens), :]
+
+        source = ops.reshape(source, (B, H, W, C))
+        upscaled_embeddings = self.output_upscaling(source)
+        hyper_in_list = []
+        for i in range(self.num_mask_tokens):
+            hyper_in_list.append(
+                self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])
+            )
+        hyper_in = ops.stack(hyper_in_list, axis=1)
+        shape = ops.shape(upscaled_embeddings)
+        B, H, W, C = shape[0], shape[1], shape[2], shape[3]
+        upscaled_embeddings = ops.reshape(
+            ops.transpose(upscaled_embeddings, axes=(0, 3, 1, 2)),
+            (B, C, H * W),
+        )
+        masks = ops.reshape(
+            hyper_in @ upscaled_embeddings, (B, self.num_mask_tokens, H, W)
+        )
+
+        iou_pred = self.iou_prediction_head(iou_token_out)
+
+        return masks, iou_pred
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "transformer_dim": self.transformer_dim,
+                "transformer": keras.saving.serialize_keras_object(
+                    self.transformer
+                ),
+                "num_multimask_outputs": self.num_multimask_outputs,
+                "iou_head_depth": self.iou_head_depth,
+                "iou_head_hidden_dim": self.iou_head_hidden_dim,
+                "activation": self.activation,
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        config.update(
+            {"transformer": keras.layers.deserialize(config["transformer"])}
+        )
+        return super().from_config(config)
diff --git a/keras_cv/models/segmentation/segment_anything/sam_presets.py b/keras_cv/models/segmentation/segment_anything/sam_presets.py
new file mode 100644
index 0000000000..74368dfe10
--- /dev/null
+++ b/keras_cv/models/segmentation/segment_anything/sam_presets.py
@@ -0,0 +1,45 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""SAM model preset configurations."""
+
+
+sam_presets = {
+    "sam_base_sa1b": {
+        "metadata": {
+            "description": "The base SAM model trained on the SA1B dataset.",
+            "params": 93_735_728,
+            "official_name": "SAM",
+            "path": "segment_anything",
+        },
+        "kaggle_handle": "kaggle://keras/sam/keras/sam_base_sa1b/2",
+    },
+    "sam_large_sa1b": {
+        "metadata": {
+            "description": "The large SAM model trained on the SA1B dataset.",
+            "params": 312_343_088,
+            "official_name": "SAM",
+            "path": "segment_anything",
+        },
+        "kaggle_handle": "kaggle://keras/sam/keras/sam_large_sa1b/2",
+    },
+    "sam_huge_sa1b": {
+        "metadata": {
+            "description": "The huge SAM model trained on the SA1B dataset.",
+            "params": 641_090_864,
+            "official_name": "SAM",
+            "path": "segment_anything",
+        },
+        "kaggle_handle": "kaggle://keras/sam/keras/sam_huge_sa1b/2",
+    },
+}
diff --git a/keras_cv/models/segmentation/segment_anything/sam_prompt_encoder.py b/keras_cv/models/segmentation/segment_anything/sam_prompt_encoder.py
new file mode 100644
index 0000000000..b792246f06
--- /dev/null
+++ b/keras_cv/models/segmentation/segment_anything/sam_prompt_encoder.py
@@ -0,0 +1,301 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.models.segmentation.segment_anything.sam_layers import (
+    RandomFrequencyPositionalEmbeddings,
+)
+
+
+@keras_cv_export("keras_cv.models.SAMPromptEncoder", package="keras_cv.models")
+class SAMPromptEncoder(keras.layers.Layer):
+    """Prompt Encoder for the Segment Anything Model (SAM).
+
+    The prompt encoder generates encodings for three types of prompts:
+
+    - Point prompts: Points on the image along with a label indicating whether
+        the point is in the foreground (part of the mask) or in the background
+        (not a part of the mask).
+    - Box prompts: A batch of bounding boxes with format [(x1, y1), (x2, y2)]
+        used to determine the location of the masks in the image.
+    - Masks: An input mask can be passed to refine the positional embeddings
+        for the output mask.
+
+    First, the point prompts and box prompts are concatenated and positional
+    encodings are generated using random spatial frequencies. A point is
+    represented as the sum of a positional encoding of the point's location
+    and one of two learned embeddings that indicate if the point is either in
+    the foreground or background. A box is represented by an embedding pair:
+
+    (1) the positional encoding of its top-left corner summed with a learned
+    embedding representing "top-left corner" and
+    (2) the same structure but using a learned embedding indicating
+    "bottom-right corner".
+
+    The box and point encodings are referred to as "sparse encodings"
+
+    If a mask prompt is passed, a convolutional neural net is used to
+    downscale it to generate "dense encodings". If no mask prompt is passed,
+    an embedding layer is used instead to generate a "no mask" embedding.
+
+    Args:
+        embed_dim (int, optional): The number of features in the output
+            embeddings. Defaults to `256`.
+        image_embedding_size (int, optional): The number of features in the
+            image embeddings generated by an image encoder. Defaults to
+            `(64, 64)`.
+        input_image_size (tuple[int], optional): A tuple of the height and
+            width of the image being prompted. Defaults to `(1024, 1024)`.
+        mask_in_chans (int, optional): The number of channels of the mask
+            prompt. Defaults to `16`.
+        activation (str, optional): The activation to use in the mask
+            downscaler neural net. Defaults to `"gelu"`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        *,
+        embed_dim=256,
+        image_embedding_size=(64, 64),
+        input_image_size=(1024, 1024),
+        mask_in_chans=16,
+        activation="gelu",
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.embed_dim = embed_dim
+        self.image_embedding_size = image_embedding_size
+        self.input_image_size = input_image_size
+        self.mask_in_chans = mask_in_chans
+        self.activation = activation
+
+        self.positional_embedding_layer = RandomFrequencyPositionalEmbeddings(
+            num_positional_features=self.embed_dim // 2, scale=1
+        )
+
+        self.foreground_point_embed = keras.layers.Embedding(
+            1, embed_dim, name="foreground_point_embed"
+        )
+        self.background_point_embed = keras.layers.Embedding(
+            1, embed_dim, name="background_point_embed"
+        )
+        self.top_left_corner_embed = keras.layers.Embedding(
+            1, embed_dim, name="top_left_corner_embed"
+        )
+        self.bottom_right_corner_embed = keras.layers.Embedding(
+            1, embed_dim, name="bottom_right_corner_embed"
+        )
+        self.not_a_point_embed = keras.layers.Embedding(
+            1, embed_dim, name="not_a_point_embed"
+        )
+
+        self.mask_downscaler = keras.models.Sequential(
+            [
+                keras.layers.Conv2D(
+                    mask_in_chans // 4, kernel_size=2, strides=2
+                ),
+                keras.layers.LayerNormalization(epsilon=1e-6),
+                keras.layers.Activation(activation),
+                keras.layers.Conv2D(mask_in_chans, kernel_size=2, strides=2),
+                keras.layers.LayerNormalization(epsilon=1e-6),
+                keras.layers.Activation(activation),
+                keras.layers.Conv2D(embed_dim, kernel_size=1),
+            ],
+            name="mask_downscaler",
+        )
+        self.no_mask_embed = keras.layers.Embedding(
+            1, embed_dim, name="no_mask_embed"
+        )
+
+    def build(self, input_shape=None):
+        self.positional_embedding_layer.build()
+        for layer in [
+            self.foreground_point_embed,
+            self.background_point_embed,
+            self.top_left_corner_embed,
+            self.bottom_right_corner_embed,
+            self.not_a_point_embed,
+            self.no_mask_embed,
+        ]:
+            layer.build([None])
+        self.mask_downscaler.build(
+            [
+                None,
+                4 * self.image_embedding_size[0],
+                4 * self.image_embedding_size[1],
+                1,
+            ]
+        )
+        self.built = True
+
+    def compute_output_shape(self, input_shape):
+        return {
+            "sparse_embeddings": [None, None, self.embed_dim],
+            "dense_embeddings": [
+                None,
+                self.image_embedding_size[0],
+                self.image_embedding_size[1],
+                self.embed_dim,
+            ],
+            "dense_positional_embeddings": [
+                None,
+                self.image_embedding_size[0],
+                self.image_embedding_size[1],
+                self.embed_dim,
+            ],
+        }
+
+    def __embed_points(self, points, labels):
+        points = points + 0.5
+        indices = ops.arange(1, dtype="int32")
+
+        point_embeddings = self.positional_embedding_layer.encode_coordinates(
+            points, self.input_image_size
+        )
+        labels = ops.broadcast_to(
+            labels[..., None], ops.shape(point_embeddings)
+        )
+        point_embeddings = ops.where(
+            labels == 0,
+            point_embeddings + self.background_point_embed(indices),
+            point_embeddings + self.foreground_point_embed(indices),
+        )
+        point_embeddings = ops.where(
+            labels == -1,
+            self.not_a_point_embed(indices),
+            point_embeddings,
+        )
+        return point_embeddings
+
+    def __embed_box(self, box):
+        shape = ops.shape(box)
+        B, N = shape[0], shape[1]
+        box = box + 0.5
+        indices = ops.arange(1, dtype="int32")
+        corner_embedding = self.positional_embedding_layer.encode_coordinates(
+            box, self.input_image_size
+        )
+        top_left_embedding = corner_embedding[
+            :, :, 0, :
+        ] + self.top_left_corner_embed(indices)
+        bottom_right_embedding = corner_embedding[
+            :, :, 1, :
+        ] + self.bottom_right_corner_embed(indices)
+        corner_embedding = ops.stack(
+            [top_left_embedding, bottom_right_embedding], axis=2
+        )
+        return ops.reshape(corner_embedding, (B, N * 2, self.embed_dim))
+
+    def __embed_mask(self, mask):
+        mask_embedding = self.mask_downscaler(mask)
+        return mask_embedding
+
+    def call(self, inputs):
+        # Get the batch shape based on any arbitrary input, because batch
+        # shapes must all match.
+        B = ops.shape(next(iter(inputs.values())))[0]
+
+        points = inputs.get("points", ops.zeros((B, 0, 2)))
+        labels = inputs.get("labels", ops.zeros((B, 0)))
+        box = inputs.get("boxes", ops.zeros((B, 0, 2, 2)))
+        mask = inputs.get("masks", ops.zeros((B, 0, 256, 256, 1)))
+
+        # Compute point embeddings
+        point_embeddings = self.__embed_points(points, labels)
+
+        # Compute box embeddings
+        box_embeddings = self.__embed_box(box)
+
+        # Concatenate both into a sparse embeddings tensor
+        sparse_embeddings = ops.concatenate(
+            [point_embeddings, box_embeddings], axis=1
+        )
+
+        # Compute the mask embeddings
+        _no_mask_embed = lambda: (
+            ops.broadcast_to(
+                ops.reshape(
+                    self.no_mask_embed(ops.arange(1, dtype="int32")),
+                    (1, 1, 1, self.embed_dim),
+                ),
+                shape=(
+                    B,
+                    self.image_embedding_size[0],
+                    self.image_embedding_size[1],
+                    self.embed_dim,
+                ),
+            )
+        )
+
+        def _maybe_input_mask_embed():
+            # Keras Core passes the masks as concrete tensors for both the
+            # true and false functions to build the output shape. So, we
+            # need to handle the case when 0 size mask is passed and
+            # dispatch the call to `_no_mask_embed`. Note that we can't call
+            # the lambda directly since the inputs are bound to different
+            # values when called with concrete values.
+            if mask.shape[1] == 0:
+                return ops.broadcast_to(
+                    ops.reshape(
+                        self.no_mask_embed(ops.arange(1, dtype="int32")),
+                        (1, 1, 1, self.embed_dim),
+                    ),
+                    shape=(
+                        B,
+                        self.image_embedding_size[0],
+                        self.image_embedding_size[1],
+                        self.embed_dim,
+                    ),
+                )
+            shape = ops.shape(mask)
+            BM, N, H, W, C = shape[0], shape[1], shape[2], shape[3], shape[4]
+            return self.__embed_mask(ops.reshape(mask, (BM * N, H, W, C)))
+
+        dense_embeddings = ops.cond(
+            ops.equal(ops.size(mask), 0),
+            _no_mask_embed,
+            _maybe_input_mask_embed,
+        )
+
+        # Compute the dense positional embeddings
+        dense_positional_embeddings = (
+            self.positional_embedding_layer.encode_image(
+                self.image_embedding_size
+            )[None, ...]
+        )
+
+        return {
+            "sparse_embeddings": sparse_embeddings,
+            "dense_embeddings": dense_embeddings,
+            "dense_positional_embeddings": dense_positional_embeddings,
+        }
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "embed_dim": self.embed_dim,
+                "image_embedding_size": self.image_embedding_size,
+                "input_image_size": self.input_image_size,
+                "mask_in_chans": self.mask_in_chans,
+                "activation": self.activation,
+            }
+        )
+        return config
diff --git a/keras_cv/models/segmentation/segment_anything/sam_test.py b/keras_cv/models/segmentation/segment_anything/sam_test.py
new file mode 100644
index 0000000000..295355a716
--- /dev/null
+++ b/keras_cv/models/segmentation/segment_anything/sam_test.py
@@ -0,0 +1,378 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import os
+import pathlib
+
+import numpy as np
+import pytest
+from absl.testing import parameterized
+
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.backend.config import keras_3
+from keras_cv.models.backbones.vit_det.vit_det_aliases import ViTDetBBackbone
+from keras_cv.models.segmentation.segment_anything.sam import (
+    SegmentAnythingModel,
+)
+from keras_cv.models.segmentation.segment_anything.sam_layers import (
+    TwoWayMultiHeadAttention,
+)
+from keras_cv.models.segmentation.segment_anything.sam_mask_decoder import (
+    SAMMaskDecoder,
+)
+from keras_cv.models.segmentation.segment_anything.sam_prompt_encoder import (
+    SAMPromptEncoder,
+)
+from keras_cv.models.segmentation.segment_anything.sam_transformer import (
+    TwoWayTransformer,
+)
+from keras_cv.tests.test_case import TestCase
+
+
+class SAMTest(TestCase):
+    def setUp(self):
+        self.image_encoder = ViTDetBBackbone()
+        self.prompt_encoder = SAMPromptEncoder(
+            embed_dim=256,
+            image_embedding_size=(64, 64),
+            input_image_size=(1024, 1024),
+            mask_in_chans=16,
+        )
+        self.mask_decoder = SAMMaskDecoder(
+            transformer_dim=256,
+            transformer=TwoWayTransformer(
+                depth=2, embed_dim=256, mlp_dim=2048, num_heads=8
+            ),
+            num_multimask_outputs=3,
+            iou_head_depth=3,
+            iou_head_hidden_dim=256,
+        )
+
+    def get_prompts(self, B, prompts="all"):
+        rng = np.random.default_rng(0)
+
+        prompts_dict = {}
+
+        if "all" in prompts or "points" in prompts:
+            prompts_dict["points"] = ops.convert_to_tensor(
+                rng.integers(0, 1023, (B, 10, 2)), dtype="float32"
+            )
+            prompts_dict["labels"] = ops.convert_to_tensor(
+                1 * (rng.random((B, 10)) > 0.5), dtype="int32"
+            )
+
+        if "all" in prompts or "boxes" in prompts:
+            x1y1 = rng.integers(0, 1022, (B, 2))
+            x2y2 = rng.integers(x1y1, 1023, (B, 2))
+            box = np.stack([x1y1, x2y2], axis=1)
+            prompts_dict["boxes"] = ops.convert_to_tensor(
+                box[:, None, ...], dtype="float32"
+            )
+        if "all" in prompts or "masks" in prompts:
+            prompts_dict["masks"] = ops.convert_to_tensor(
+                1.0 * (rng.random((B, 1, 256, 256, 1)) > 0.5), dtype="float32"
+            )
+
+        return prompts_dict
+
+    def test_prompt_encoder_simple(self):
+        outputs = self.prompt_encoder(self.get_prompts(7))
+        sparse_embeddings, dense_embeddings, dense_positional_embeddings = (
+            outputs["sparse_embeddings"],
+            outputs["dense_embeddings"],
+            outputs["dense_positional_embeddings"],
+        )
+
+        trainable_parameters = np.sum(
+            [np.prod(x.shape) for x in self.prompt_encoder.trainable_weights]
+        )
+        num_parameters = np.sum(
+            [np.prod(x.shape) for x in self.prompt_encoder.weights]
+        )
+
+        sparse_embeddings = ops.convert_to_numpy(sparse_embeddings)
+        dense_embeddings = ops.convert_to_numpy(dense_embeddings)
+        dense_positional_embeddings = ops.convert_to_numpy(
+            dense_positional_embeddings
+        )
+
+        self.assertEqual(sparse_embeddings.shape, (7, 12, 256))
+        self.assertEqual(dense_embeddings.shape, (7, 64, 64, 256))
+        self.assertEqual(dense_positional_embeddings.shape, (1, 64, 64, 256))
+        self.assertEqual(trainable_parameters, 6_220)
+        self.assertEqual(num_parameters, 6_476)
+
+    @parameterized.named_parameters(
+        [
+            ("_".join(x), x)
+            for x in itertools.chain(
+                itertools.combinations(["points", "boxes", "masks"], 1),
+                itertools.combinations(["points", "boxes", "masks"], 2),
+            )
+        ]
+    )
+    def test_prompt_encoder_partial_prompts(self, prompts):
+        prompts_dict = self.get_prompts(7, prompts)
+        outputs = self.prompt_encoder(prompts_dict)
+        sparse_embeddings, dense_embeddings = (
+            outputs["sparse_embeddings"],
+            outputs["dense_embeddings"],
+        )
+
+        sparse_embeddings_dim = 0
+        if "points" in prompts:
+            sparse_embeddings_dim += prompts_dict["points"].shape[1]
+        if "boxes" in prompts:
+            sparse_embeddings_dim += prompts_dict["boxes"].shape[1] * 2
+        self.assertAllEqual(
+            sparse_embeddings.shape,
+            (7, sparse_embeddings_dim, 256),
+        )
+        self.assertAllEqual(dense_embeddings.shape, (7, 64, 64, 256))
+        if "masks" not in prompts:
+            no_mask_embed = ops.broadcast_to(
+                self.prompt_encoder.no_mask_embed(ops.arange(1)),
+                (7, 64, 64, 256),
+            )
+            self.assertAllClose(dense_embeddings, no_mask_embed)
+
+    def test_two_way_multi_head_attention(self):
+        image_embeddings = np.random.randn(1, 64, 64, 256).astype(np.float32)
+
+        prompt_encoder_outputs = self.prompt_encoder(self.get_prompts(1))
+        sparse_embeddings = prompt_encoder_outputs["sparse_embeddings"]
+
+        two_way_attention = TwoWayMultiHeadAttention(
+            num_heads=8,
+            key_dim=256 // 8,
+            mlp_dim=2048,
+            skip_first_layer_pe=False,
+        )
+        queries, keys = two_way_attention(
+            queries=sparse_embeddings,
+            keys=ops.reshape(image_embeddings, (1, 64 * 64, 256)),
+            query_pe=sparse_embeddings,
+            key_pe=ops.reshape(
+                prompt_encoder_outputs["dense_positional_embeddings"],
+                (1, 64 * 64, 256),
+            ),
+        )
+
+        queries, keys = map(ops.convert_to_numpy, [queries, keys])
+
+        self.assertEqual(queries.shape, (1, 12, 256))
+        self.assertEqual(keys.shape, (1, 64 * 64, 256))
+
+    def test_two_way_transformer(self):
+        prompt_encoder_outputs = self.prompt_encoder(self.get_prompts(1))
+        sparse_embeddings = prompt_encoder_outputs["sparse_embeddings"]
+        image_embeddings = np.random.randn(1, 64, 64, 256)
+        two_way_transformer = TwoWayTransformer(
+            depth=2, embed_dim=256, num_heads=8, mlp_dim=2048
+        )
+        queries, keys = two_way_transformer(
+            image_embedding=image_embeddings,
+            image_pe=prompt_encoder_outputs["dense_positional_embeddings"],
+            point_embedding=sparse_embeddings,
+        )
+
+        queries, keys = map(ops.convert_to_numpy, [queries, keys])
+
+        self.assertEqual(queries.shape, (1, 12, 256))
+        self.assertEqual(keys.shape, (1, 64 * 64, 256))
+
+    def test_mask_decoder(self):
+        prompt_encoder_outputs = self.prompt_encoder(self.get_prompts(1))
+        sparse_embeddings, dense_embeddings, dense_positional_embeddings = (
+            prompt_encoder_outputs["sparse_embeddings"],
+            prompt_encoder_outputs["dense_embeddings"],
+            prompt_encoder_outputs["dense_positional_embeddings"],
+        )
+        image_embeddings = np.random.randn(1, 64, 64, 256)
+        outputs = self.mask_decoder(
+            dict(
+                image_embeddings=image_embeddings,
+                image_pe=dense_positional_embeddings,
+                sparse_prompt_embeddings=sparse_embeddings,
+                dense_prompt_embeddings=dense_embeddings,
+            )
+        )
+        masks, iou_pred = outputs["masks"], outputs["iou_pred"]
+        num_parameters = np.sum(
+            [np.prod(x.shape) for x in self.mask_decoder.weights]
+        )
+        masks, iou_pred = map(ops.convert_to_numpy, [masks, iou_pred])
+        self.assertEqual(masks.shape, (1, 4, 256, 256))
+        self.assertEqual(iou_pred.shape, (1, 4))
+        self.assertEqual(num_parameters, 4_058_340)
+
+    @pytest.mark.large
+    @parameterized.named_parameters(
+        [
+            ("float32", "float32"),
+            ("mixed_float16", "mixed_float16"),
+            ("bfloat16", "bfloat16"),
+        ]
+    )
+    def test_end_to_end_model_predict(self, dtype_policy):
+        import threading
+
+        with threading.Lock():
+            # We are changing the global dtype policy here but don't want any
+            # other tests to use that policy, so compute under a lock until
+            # we reset the global policy.
+            old_policy = getattr(
+                keras.mixed_precision, "dtype_policy", lambda: "float32"
+            )()
+            keras.mixed_precision.set_global_policy(dtype_policy)
+            model = SegmentAnythingModel(
+                backbone=self.image_encoder,
+                prompt_encoder=self.prompt_encoder,
+                mask_decoder=self.mask_decoder,
+            )
+
+            # We use box-only prompting for this test.
+            mask_prompts = self.get_prompts(1, "boxes")
+            inputs = {
+                "images": np.ones((1, 1024, 1024, 3)),
+            }
+            inputs.update(mask_prompts)
+
+            # Check the number of parameters
+            num_parameters = np.sum([np.prod(x.shape) for x in model.weights])
+            self.assertEqual(num_parameters, 89_670_912 + 6_476 + 4_058_340)
+
+            # Forward pass through the model
+            outputs = model.predict(inputs)
+            masks, iou_pred = outputs["masks"], outputs["iou_pred"]
+
+            # Check the output is equal to the one we expect if we
+            # run each component separately. This is to confirm that
+            # the graph is getting compiled correctly i.e. the jitted
+            # execution is equivalent to the eager execution.
+            features = self.image_encoder(inputs["images"])
+            outputs_ex = self.prompt_encoder(
+                {k: v for k, v in inputs.items() if k != "images"}
+            )
+            outputs_ex = self.mask_decoder(
+                {
+                    "image_embeddings": features,
+                    "image_pe": outputs_ex["dense_positional_embeddings"],
+                    "sparse_prompt_embeddings": outputs_ex["sparse_embeddings"],
+                    "dense_prompt_embeddings": outputs_ex["dense_embeddings"],
+                },
+            )
+            masks_ex, iou_pred_ex = outputs_ex["masks"], outputs_ex["iou_pred"]
+
+            self.assertAllClose(masks, masks_ex, atol=1e-4)
+            self.assertAllClose(iou_pred, iou_pred_ex, atol=1e-4)
+
+            # Reset the global policy
+            keras.mixed_precision.set_global_policy(old_policy)
+
+    @pytest.mark.extra_large
+    def test_end_to_end_model_save(self):
+        # Build the model
+        model = SegmentAnythingModel(
+            backbone=self.image_encoder,
+            prompt_encoder=self.prompt_encoder,
+            mask_decoder=self.mask_decoder,
+        )
+
+        mask_prompts = self.get_prompts(1)
+        inputs = {
+            "images": np.ones((1, 1024, 1024, 3)),
+        }
+        inputs.update(mask_prompts)
+
+        # Forward pass
+        outputs = model.predict(inputs)
+
+        # Save the model
+        save_path = os.path.join(self.get_temp_dir(), "model.keras")
+        if keras_3():
+            model.save(save_path)
+        else:
+            model.save(save_path, save_format="keras_v3")
+        restored_model = keras.models.load_model(save_path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, SegmentAnythingModel)
+
+        # Check that output matches.
+        restored_outputs = restored_model.predict(inputs)
+        self.assertAllClose(outputs, restored_outputs)
+
+    @pytest.mark.large
+    def test_end_to_end_model_preset(self):
+        # Define the RNG. Don't change the seed. This seed
+        # was used to generate the inputs for the reference
+        # values.
+        rng = np.random.default_rng(0)
+
+        # Generate the inputs
+        inputs = {
+            "images": 255.0 * rng.random((1, 1024, 1024, 3), dtype=np.float32),
+            "points": np.array(
+                [[[10, 10], [100, 100], [500, 500]]], dtype=np.float32
+            ),
+            "labels": np.array([[0, 1, 0]], dtype=np.float32),
+            "boxes": np.array(
+                [[[[10.0, 10.0], [100.0, 100.0]]]], dtype=np.float32
+            ),
+            "masks": (rng.random((1, 1, 256, 256, 1)) > 0.5).astype(np.float32),
+        }
+
+        # Run the model
+        model = SegmentAnythingModel.from_preset("sam_base_sa1b")
+        outs = model.predict(inputs)
+
+        # Make sure the weights have been loaded correctly.
+        masks_expected = np.load(
+            pathlib.Path(__file__).parent / "data" / "sam_base_out_masks.npy"
+        )
+        iou_pred_expected = np.load(
+            pathlib.Path(__file__).parent / "data" / "sam_base_out_iou_pred.npy"
+        )
+        self.assertAllClose(outs["masks"], masks_expected, atol=1e-2, rtol=1e-2)
+        self.assertAllClose(
+            outs["iou_pred"], iou_pred_expected, atol=1e-2, rtol=1e-2
+        )
+
+    def test_end_to_end_model_fit_error(self):
+        # Build the model
+        model = SegmentAnythingModel(
+            backbone=self.image_encoder,
+            prompt_encoder=self.prompt_encoder,
+            mask_decoder=self.mask_decoder,
+        )
+
+        mask_prompts = self.get_prompts(1)
+        inputs = {
+            "images": np.ones((1, 1024, 1024, 3)),
+        }
+        inputs.update(mask_prompts)
+
+        # Compile the model
+        model.compile(
+            optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"]
+        )
+
+        # Check that calling fit raises a NotImplementedError.
+        with self.assertRaises(
+            NotImplementedError, msg=r"only supports inference"
+        ):
+            model.fit(inputs)
diff --git a/keras_cv/models/segmentation/segment_anything/sam_transformer.py b/keras_cv/models/segmentation/segment_anything/sam_transformer.py
new file mode 100644
index 0000000000..64d1de8575
--- /dev/null
+++ b/keras_cv/models/segmentation/segment_anything/sam_transformer.py
@@ -0,0 +1,155 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import keras
+from keras_cv.backend import ops
+from keras_cv.models.segmentation.segment_anything.sam_layers import (
+    MultiHeadAttentionWithDownsampling,
+)
+from keras_cv.models.segmentation.segment_anything.sam_layers import (
+    TwoWayMultiHeadAttention,
+)
+
+
+@keras_cv_export("keras_cv.models.TwoWayTransformer", package="keras_cv.models")
+class TwoWayTransformer(keras.layers.Layer):
+    """A two-way cross-attention transformer decoder.
+
+    A transformer decoder that attends to an input image using
+    queries whose positional embedding is supplied.
+
+    The transformer decoder design is shown in [1]_. Each decoder layer
+    performs 4 steps: (1) self-attention on the tokens, (2) cross-attention
+    from tokens (as queries) to the image embedding, (3) a point-wise MLP
+    updates each token, and (4) cross-attention from the image embedding (as
+    queries) to tokens. This last step updates the image embedding with prompt
+    information. Each self/cross-attention and MLP has a residual connection
+    and layer normalization.
+
+    To ensure the decoder has access to critical geometric information the
+    positional encodings are added to the image embedding whenever they
+    participate in an attention layer. Additionally, the entire original
+    prompt tokens (including their positional encodings) are re-added to the
+    updated tokens whenever they participate in an attention layer. This
+    allows for a strong dependence on both the prompt token's geometric
+    location and type.
+
+    Args:
+        depth (int, optional): The depth of the attention blocks (the number
+            of attention blocks to use). Defaults to `2`.
+        embed_dim (int, optional): The number of features of the input image
+            and point embeddings. Defaults to `256`.
+        num_heads (int, optional): Number of heads to use in the attention
+            layers. Defaults to `8`.
+        mlp_dim (int, optional): The number of units in the hidden layer of
+            the MLP block used in the attention layers. Defaults to `2048`.
+        activation (str, optional): The activation of the MLP block's output
+            layer used in the attention layers. Defaults to `"relu"`.
+        attention_downsample_rate (int, optional): The downsample rate of the
+            attention layers. Defaults to `2`.
+
+    References:
+        - [Segment Anything paper](https://arxiv.org/abs/2304.02643)
+        - [Segment Anything GitHub](https://github.com/facebookresearch/segment-anything)
+    """  # noqa: E501
+
+    def __init__(
+        self,
+        *,
+        depth=2,
+        embed_dim=256,
+        num_heads=8,
+        mlp_dim=2048,
+        activation="relu",
+        attention_downsample_rate=2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.depth = depth
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.mlp_dim = mlp_dim
+        self.activation = activation
+        self.attention_downsample_rate = attention_downsample_rate
+        self.layers = []
+        for i in range(depth):
+            self.layers.append(
+                TwoWayMultiHeadAttention(
+                    num_heads=num_heads,
+                    key_dim=embed_dim // num_heads,
+                    mlp_dim=mlp_dim,
+                    skip_first_layer_pe=(i == 0),
+                    attention_downsample_rate=attention_downsample_rate,
+                    activation=activation,
+                )
+            )
+        self.final_attention_token_to_image = (
+            MultiHeadAttentionWithDownsampling(
+                num_heads=num_heads,
+                key_dim=embed_dim // num_heads,
+                downsample_rate=attention_downsample_rate,
+            )
+        )
+        self.final_layer_norm = keras.layers.LayerNormalization(epsilon=1e-5)
+
+    def build(self, input_shape=None):
+        for layer in self.layers:
+            layer.build()
+        self.final_attention_token_to_image.build()
+        self.final_layer_norm.build([None, None, self.embed_dim])
+        self.built = True
+
+    def call(self, image_embedding, image_pe, point_embedding):
+        shape = ops.shape(image_embedding)
+        B, H, W, C = shape[0], shape[1], shape[2], shape[3]
+        image_embedding = ops.reshape(image_embedding, (B, H * W, C))
+
+        shape = ops.shape(image_pe)
+        B, H, W, C = shape[0], shape[1], shape[2], shape[3]
+        image_pe = ops.reshape(image_pe, (B, H * W, C))
+        queries = point_embedding
+        keys = image_embedding
+
+        for layer in self.layers:
+            queries, keys = layer(
+                queries=queries,
+                keys=keys,
+                query_pe=point_embedding,
+                key_pe=image_pe,
+            )
+
+        queries_with_pe = queries + point_embedding
+        keys_with_pe = keys + image_pe
+        attention_map = self.final_attention_token_to_image(
+            query=queries_with_pe, key=keys_with_pe, value=keys
+        )
+        queries = queries + attention_map
+        queries = self.final_layer_norm(queries)
+
+        return queries, keys
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "depth": self.depth,
+                "embed_dim": self.embed_dim,
+                "num_heads": self.num_heads,
+                "mlp_dim": self.mlp_dim,
+                "activation": self.activation,
+                "attention_downsample_rate": self.attention_downsample_rate,
+            }
+        )
+        return config
diff --git a/keras_cv/models/stable_diffusion/clip_tokenizer.py b/keras_cv/models/stable_diffusion/clip_tokenizer.py
index 13e73508c4..cf2f186314 100644
--- a/keras_cv/models/stable_diffusion/clip_tokenizer.py
+++ b/keras_cv/models/stable_diffusion/clip_tokenizer.py
@@ -20,6 +20,7 @@
 
 import regex as re
 
+from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 
 
@@ -76,6 +77,7 @@ def whitespace_clean(text):
     return text
 
 
+@keras_cv_export("keras_cv.models.stable_diffusion.SimpleTokenizer")
 class SimpleTokenizer:
     def __init__(self, bpe_path=None):
         bpe_path = bpe_path or keras.utils.get_file(
diff --git a/keras_cv/models/stable_diffusion/decoder.py b/keras_cv/models/stable_diffusion/decoder.py
index 318d64366c..e97f3a0089 100644
--- a/keras_cv/models/stable_diffusion/decoder.py
+++ b/keras_cv/models/stable_diffusion/decoder.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.models.stable_diffusion.attention_block import (  # noqa: E501
     AttentionBlock,
@@ -20,6 +21,7 @@
 from keras_cv.models.stable_diffusion.resnet_block import ResnetBlock
 
 
+@keras_cv_export("keras_cv.models.stable_diffusion.Decoder")
 class Decoder(keras.Sequential):
     def __init__(self, img_height, img_width, name=None, download_weights=True):
         super().__init__(
diff --git a/keras_cv/models/stable_diffusion/diffusion_model.py b/keras_cv/models/stable_diffusion/diffusion_model.py
index ba9e849737..7c8fd30f43 100644
--- a/keras_cv/models/stable_diffusion/diffusion_model.py
+++ b/keras_cv/models/stable_diffusion/diffusion_model.py
@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
 from keras_cv.models.stable_diffusion.padded_conv2d import PaddedConv2D
 
 
+@keras_cv_export("keras_cv.models.stable_diffusion.DiffusionModel")
 class DiffusionModel(keras.Model):
     def __init__(
         self,
diff --git a/keras_cv/models/stable_diffusion/image_encoder.py b/keras_cv/models/stable_diffusion/image_encoder.py
index 13128f0412..9c32ab6912 100644
--- a/keras_cv/models/stable_diffusion/image_encoder.py
+++ b/keras_cv/models/stable_diffusion/image_encoder.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.models.stable_diffusion.attention_block import (  # noqa: E501
     AttentionBlock,
@@ -20,6 +21,7 @@
 from keras_cv.models.stable_diffusion.resnet_block import ResnetBlock
 
 
+@keras_cv_export("keras_cv.models.stable_diffusion.ImageEncoder")
 class ImageEncoder(keras.Sequential):
     """ImageEncoder is the VAE Encoder for StableDiffusion."""
 
diff --git a/keras_cv/models/stable_diffusion/noise_scheduler.py b/keras_cv/models/stable_diffusion/noise_scheduler.py
index 069b9c6284..17bad6468f 100644
--- a/keras_cv/models/stable_diffusion/noise_scheduler.py
+++ b/keras_cv/models/stable_diffusion/noise_scheduler.py
@@ -16,10 +16,12 @@
 Adapted from https://github.com/huggingface/diffusers/blob/v0.3.0/src/diffusers/schedulers/scheduling_ddpm.py#L56
 """  # noqa: E501
 
+from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import ops
 from keras_cv.backend import random
 
 
+@keras_cv_export("keras_cv.models.stable_diffusion.NoiseScheduler")
 class NoiseScheduler:
     """
     Args:
@@ -29,8 +31,6 @@ class NoiseScheduler:
         beta_schedule: the beta schedule, a mapping from a beta range to a
             sequence of betas for stepping the model. Choose from `linear` or
             `quadratic`.
-        betas: a complete set of betas, in lieu of using one of the existing
-            schedules.
         variance_type: options to clip the variance used when adding noise to
             the de-noised sample. Choose from `fixed_small`, `fixed_small_log`,
             `fixed_large`, `fixed_large_log`, `learned` or `learned_range`.
@@ -41,11 +41,11 @@ class NoiseScheduler:
     def __init__(
         self,
         train_timesteps=1000,
-        beta_start=0.0001,
-        beta_end=0.02,
-        beta_schedule="linear",
+        beta_start=0.00085,
+        beta_end=0.012,
+        beta_schedule="scaled_linear",
         variance_type="fixed_small",
-        clip_sample=True,
+        clip_sample=False,
     ):
         self.train_timesteps = train_timesteps
 
@@ -54,16 +54,14 @@ def __init__(
         elif beta_schedule == "scaled_linear":
             # this schedule is very specific to the latent diffusion model.
             self.betas = (
-                ops.linspace(
-                    beta_start**0.5, beta_end**0.5, train_timesteps
-                )
+                ops.linspace(beta_start**0.5, beta_end**0.5, train_timesteps)
                 ** 2
             )
         else:
             raise ValueError(f"Invalid beta schedule: {beta_schedule}.")
 
         self.alphas = 1.0 - self.betas
-        self.alphas_cumprod = ops.math.cumprod(self.alphas)
+        self.alphas_cumprod = ops.cumprod(self.alphas)
 
         self.variance_type = variance_type
         self.clip_sample = clip_sample
@@ -205,7 +203,12 @@ def add_noise(
             sqrt_one_minus_alpha_prod = ops.expand_dims(
                 sqrt_one_minus_alpha_prod, axis=-1
             )
-
+        sqrt_alpha_prod = ops.cast(
+            sqrt_alpha_prod, dtype=original_samples.dtype
+        )
+        sqrt_one_minus_alpha_prod = ops.cast(
+            sqrt_one_minus_alpha_prod, dtype=noise.dtype
+        )
         noisy_samples = (
             sqrt_alpha_prod * original_samples
             + sqrt_one_minus_alpha_prod * noise
diff --git a/keras_cv/models/stable_diffusion/stable_diffusion.py b/keras_cv/models/stable_diffusion/stable_diffusion.py
index 2ec798b40a..a68923dc78 100644
--- a/keras_cv/models/stable_diffusion/stable_diffusion.py
+++ b/keras_cv/models/stable_diffusion/stable_diffusion.py
@@ -200,7 +200,7 @@ def generate_image(
 
         if diffusion_noise is not None:
             diffusion_noise = ops.squeeze(diffusion_noise)
-            if diffusion_noise.shape.rank == 3:
+            if len(ops.shape(diffusion_noise)) == 3:
                 diffusion_noise = ops.repeat(
                     ops.expand_dims(diffusion_noise, axis=0), batch_size, axis=0
                 )
@@ -209,7 +209,10 @@ def generate_image(
             latent = self._get_initial_diffusion_noise(batch_size, seed)
 
         # Iterative reverse diffusion stage
-        timesteps = np.arange(1, 1000, 1000 // num_steps)
+        num_timesteps = 1000
+        ratio = (num_timesteps - 1) / (num_steps - 1)
+        timesteps = (np.arange(0, num_steps) * ratio).round().astype(np.int64)
+
         alphas, alphas_prev = self._get_initial_alphas(timesteps)
         progbar = keras.utils.Progbar(len(timesteps))
         iteration = 0
@@ -235,6 +238,9 @@ def generate_image(
                 + unconditional_guidance_scale * (latent - unconditional_latent)
             )
             a_t, a_prev = alphas[index], alphas_prev[index]
+            # Keras backend array need to cast explicitly
+            target_dtype = latent_prev.dtype
+            latent = ops.cast(latent, target_dtype)
             pred_x0 = (latent_prev - math.sqrt(1 - a_t) * latent) / math.sqrt(
                 a_t
             )
diff --git a/keras_cv/models/stable_diffusion/stable_diffusion_test.py b/keras_cv/models/stable_diffusion/stable_diffusion_test.py
index 895b7b341d..edd8681483 100644
--- a/keras_cv/models/stable_diffusion/stable_diffusion_test.py
+++ b/keras_cv/models/stable_diffusion/stable_diffusion_test.py
@@ -25,6 +25,7 @@
 class StableDiffusionTest(TestCase):
     @pytest.mark.large
     def test_end_to_end_golden_value(self):
+        self.skipTest("TODO: #2246 values differ for Keras2 and Keras3 TF")
         prompt = "a caterpillar smoking a hookah while sitting on a mushroom"
         stablediff = StableDiffusion(128, 128)
 
diff --git a/keras_cv/models/stable_diffusion/text_encoder.py b/keras_cv/models/stable_diffusion/text_encoder.py
index 68a61bbd7e..766650d6c4 100644
--- a/keras_cv/models/stable_diffusion/text_encoder.py
+++ b/keras_cv/models/stable_diffusion/text_encoder.py
@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
 from keras_cv.backend import ops
 
 
+@keras_cv_export("keras_cv.models.stable_diffusion.TextEncoder")
 class TextEncoder(keras.Model):
     def __init__(
         self, max_length, vocab_size=49408, name=None, download_weights=True
@@ -40,6 +42,7 @@ def __init__(
             self.load_weights(text_encoder_weights_fpath)
 
 
+@keras_cv_export("keras_cv.models.stable_diffusion.TextEncoderV2")
 class TextEncoderV2(keras.Model):
     def __init__(
         self, max_length, vocab_size=49408, name=None, download_weights=True
diff --git a/keras_cv/models/task.py b/keras_cv/models/task.py
index 4bbac5531c..79a80d584e 100644
--- a/keras_cv/models/task.py
+++ b/keras_cv/models/task.py
@@ -13,10 +13,11 @@
 # limitations under the License.
 """Base class for Task models."""
 
-import os
-
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
+from keras_cv.models.backbones.backbone import Backbone
+from keras_cv.utils.preset_utils import check_preset_class
+from keras_cv.utils.preset_utils import load_from_preset
 from keras_cv.utils.python_utils import classproperty
 from keras_cv.utils.python_utils import format_docstring
 
@@ -28,6 +29,22 @@ class Task(keras.Model):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._backbone = None
+        self._functional_layer_ids = set(
+            id(layer) for layer in self._flatten_layers()
+        )
+
+    def __dir__(self):
+        # Temporary fixes for weight saving. This mimics the following PR for
+        # older version of Keras: https://github.com/keras-team/keras/pull/18982
+        def filter_fn(attr):
+            if attr in ["backbone", "_backbone"]:
+                return False
+            try:
+                return id(getattr(self, attr)) not in self._functional_layer_ids
+            except:
+                return True
+
+        return filter(filter_fn, super().__dir__())
 
     @property
     def backbone(self):
@@ -82,6 +99,7 @@ def from_preset(
         cls,
         preset,
         load_weights=None,
+        input_shape=None,
         **kwargs,
     ):
         """Instantiate {{model_name}} model from preset config and weights.
@@ -93,6 +111,9 @@ def from_preset(
             load_weights: Whether to load pre-trained weights into model.
                 Defaults to `None`, which follows whether the preset has
                 pretrained weights available.
+            input_shape : input shape that will be passed to backbone
+                initialization, Defaults to `None`.If `None`, the preset
+                value will be used.
 
         Examples:
         ```python
@@ -107,56 +128,29 @@ def from_preset(
             load_weights=False,
         ```
         """
-
-        if not cls.presets:
-            raise NotImplementedError(
-                "No presets have been created for this class."
-            )
-
-        if preset not in cls.presets:
-            raise ValueError(
-                "`preset` must be one of "
-                f"""{", ".join(cls.presets)}. Received: {preset}."""
+        # We support short IDs for official presets, e.g. `"bert_base_en"`.
+        # Map these to a Kaggle Models handle.
+        if preset in cls.presets:
+            preset = cls.presets[preset]["kaggle_handle"]
+
+        preset_cls = check_preset_class(preset, (cls, Backbone))
+
+        # Backbone case.
+        if issubclass(preset_cls, Backbone):
+            backbone = load_from_preset(
+                preset,
+                load_weights=load_weights,
             )
-
-        if load_weights and preset not in cls.presets_with_weights:
-            raise ValueError(
-                f"""Pretrained weights not available for preset "{preset}". """
-                "Set `load_weights=False` to use this preset or choose one of "
-                "the following presets with weights:"
-                f""" "{'", "'.join(cls.presets_with_weights)}"."""
-            )
-
-        metadata = cls.presets[preset]
-        # Check if preset is backbone-only model
-        if preset in cls.backbone_presets:
-            backbone_cls = keras.saving.get_registered_object(
-                metadata["class_name"]
-            )
-            backbone = backbone_cls.from_preset(preset, load_weights)
-            return cls(backbone, **kwargs)
-
-        # Otherwise must be one of class presets
-        config = metadata["config"]
-        model = cls.from_config({**config, **kwargs})
-
-        if preset not in cls.presets_with_weights or load_weights is False:
-            return model
-
-        local_weights_path = "model.h5"
-        if metadata["weights_url"].endswith(".weights.h5"):
-            local_weights_path = "model.weights.h5"
-
-        weights = keras.utils.get_file(
-            local_weights_path,
-            metadata["weights_url"],
-            cache_subdir=os.path.join("models", preset),
-            file_hash=metadata["weights_hash"],
+            return cls(backbone=backbone, **kwargs)
+
+        # Task case.
+        return load_from_preset(
+            preset,
+            load_weights=load_weights,
+            input_shape=input_shape,
+            config_overrides=kwargs,
         )
 
-        model.load_weights(weights)
-        return model
-
     @property
     def layers(self):
         # Some of our task models don't use the Backbone directly, but create
diff --git a/keras_cv/models/utils.py b/keras_cv/models/utils.py
index cb573b63c3..9d2d5b5efc 100644
--- a/keras_cv/models/utils.py
+++ b/keras_cv/models/utils.py
@@ -15,22 +15,24 @@
 """Utility functions for models"""
 
 from keras_cv.backend import keras
-from keras_cv.backend.config import multi_backend
+from keras_cv.backend.config import keras_3
 
 
 def get_tensor_input_name(tensor):
-    if multi_backend():
+    if keras_3():
         return tensor._keras_history.operation.name
     else:
         return tensor.node.layer.name
 
 
-def parse_model_inputs(input_shape, input_tensor):
+def parse_model_inputs(input_shape, input_tensor, **kwargs):
     if input_tensor is None:
-        return keras.layers.Input(shape=input_shape)
+        return keras.layers.Input(shape=input_shape, **kwargs)
     else:
         if not keras.backend.is_keras_tensor(input_tensor):
-            return keras.layers.Input(tensor=input_tensor, shape=input_shape)
+            return keras.layers.Input(
+                tensor=input_tensor, shape=input_shape, **kwargs
+            )
         else:
             return input_tensor
 
diff --git a/keras_cv/models/utils_test.py b/keras_cv/models/utils_test.py
index cc6c1ac596..4fd55dec4d 100644
--- a/keras_cv/models/utils_test.py
+++ b/keras_cv/models/utils_test.py
@@ -24,7 +24,7 @@ def test_parse_model_inputs(self):
         input_shape = (224, 244, 3)
 
         inputs = utils.parse_model_inputs(input_shape, None)
-        self.assertEqual(inputs.shape.as_list(), list((None,) + input_shape))
+        self.assertEqual(list(inputs.shape), list((None,) + input_shape))
 
         input_tensor = layers.Input(shape=input_shape)
         self.assertIs(
diff --git a/keras_cv/tools/convert_presets.py b/keras_cv/tools/convert_presets.py
new file mode 100644
index 0000000000..e4f5886ff7
--- /dev/null
+++ b/keras_cv/tools/convert_presets.py
@@ -0,0 +1,139 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import re
+
+import keras_cv  # noqa: E402
+
+BUCKET = "keras-cv-kaggle"
+
+
+def to_snake_case(name):
+    name = re.sub(r"\W+", "", name)
+    name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
+    name = re.sub("([a-z])([A-Z])", r"\1_\2", name).lower()
+    return name
+
+
+def convert_backbone_presets():
+    # Save and upload Backbone presets
+
+    backbone_models = [
+        keras_cv.models.ResNetBackbone,
+        keras_cv.models.ResNet18Backbone,
+        keras_cv.models.ResNet34Backbone,
+        keras_cv.models.ResNet50Backbone,
+        keras_cv.models.ResNet101Backbone,
+        keras_cv.models.ResNet152Backbone,
+        keras_cv.models.ResNetV2Backbone,
+        keras_cv.models.ResNet18V2Backbone,
+        keras_cv.models.ResNet34V2Backbone,
+        keras_cv.models.ResNet50V2Backbone,
+        keras_cv.models.ResNet101V2Backbone,
+        keras_cv.models.ResNet152V2Backbone,
+        keras_cv.models.YOLOV8Backbone,
+        keras_cv.models.MobileNetV3Backbone,
+        keras_cv.models.MobileNetV3SmallBackbone,
+        keras_cv.models.MobileNetV3LargeBackbone,
+        keras_cv.models.EfficientNetV2Backbone,
+        keras_cv.models.EfficientNetV2B0Backbone,
+        keras_cv.models.EfficientNetV2B1Backbone,
+        keras_cv.models.EfficientNetV2B2Backbone,
+        keras_cv.models.EfficientNetV2B3Backbone,
+        keras_cv.models.EfficientNetV2SBackbone,
+        keras_cv.models.EfficientNetV2MBackbone,
+        keras_cv.models.EfficientNetV2LBackbone,
+        keras_cv.models.CSPDarkNetBackbone,
+        keras_cv.models.DenseNetBackbone,
+        keras_cv.src.models.EfficientNetV1Backbone,
+        keras_cv.src.models.EfficientNetLiteBackbone,
+        keras_cv.models.MiTBackbone,
+        keras_cv.models.ViTDetBackbone,
+        keras_cv.models.CenterPillarBackbone,
+    ]
+    for backbone_cls in backbone_models:
+        for preset in backbone_cls.presets:
+            backbone = backbone_cls.from_preset(
+                preset, name=to_snake_case(backbone_cls.__name__)
+            )
+            save_weights = preset in backbone_cls.presets_with_weights
+            save_to_preset(
+                backbone,
+                preset,
+                save_weights=save_weights,
+                config_filename="config.json",
+            )
+            # Delete first to clean up any exising version.
+            os.system(f"gsutil rm -rf gs://{BUCKET}/{preset}")
+            os.system(f"gsutil cp -r {preset} gs://{BUCKET}/{preset}")
+            for root, _, files in os.walk(preset):
+                for file in files:
+                    path = os.path.join(BUCKET, root, file)
+                    os.system(
+                        f"gcloud storage objects update gs://{path} "
+                        "--add-acl-grant=entity=AllUsers,role=READER"
+                    )
+
+
+def convert_task_presets():
+    # Save and upload task presets
+
+    task_models = [
+        keras_cv.models.RetinaNet,
+        keras_cv.models.YOLOV8Detector,
+        keras_cv.models.ImageClassifier,
+        keras_cv.models.DeepLabV3Plus,
+        # keras_cv.models.SegFormer,
+        keras_cv.models.SegmentAnythingModel,
+    ]
+    for task_cls in task_models:
+        # Remove backbone-specific keys
+        task_preset_keys = set(task_cls.presets) ^ set(
+            task_cls.backbone_presets
+        )
+        for preset in task_preset_keys:
+            save_weights = preset in task_cls.presets_with_weights
+            kwargs = {"name": to_snake_case(task_cls.__name__)}
+            if task_cls in [
+                keras_cv.models.RetinaNet,
+                keras_cv.models.YOLOV8Detector,
+            ]:
+                kwargs.update({"bounding_box_format": "xywh"})
+                task = task_cls.from_preset(preset, **kwargs)
+            else:
+                task = task_cls.from_preset(preset, **kwargs)
+            save_to_preset(
+                task,
+                preset,
+                save_weights=save_weights,
+                config_filename="config.json",
+            )
+            # Delete first to clean up any exising version.
+            os.system(f"gsutil rm -rf gs://{BUCKET}/{preset}")
+            os.system(f"gsutil cp -r {preset} gs://{BUCKET}/{preset}")
+            for root, _, files in os.walk(preset):
+                for file in files:
+                    path = os.path.join(BUCKET, root, file)
+                    os.system(
+                        f"gcloud storage objects update gs://{path} "
+                        "--add-acl-grant=entity=AllUsers,role=READER"
+                    )
+
+
+if __name__ == "__main__":
+    from keras_cv.src.utils.preset_utils import save_to_preset  # noqa: E402
+
+    convert_backbone_presets()
+    convert_task_presets()
diff --git a/keras_cv/utils/__init__.py b/keras_cv/utils/__init__.py
index d2bfafd52f..525c99737d 100644
--- a/keras_cv/utils/__init__.py
+++ b/keras_cv/utils/__init__.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_cv.utils import test_utils
 from keras_cv.utils.conditional_imports import assert_cv2_installed
 from keras_cv.utils.conditional_imports import assert_matplotlib_installed
 from keras_cv.utils.conditional_imports import (
diff --git a/keras_cv/utils/preprocessing.py b/keras_cv/utils/preprocessing.py
index 09607dafd5..8b2b1075c1 100644
--- a/keras_cv/utils/preprocessing.py
+++ b/keras_cv/utils/preprocessing.py
@@ -184,16 +184,14 @@ def random_inversion(random_generator):
     Returns:
       either -1, or -1.
     """
-    negate = random_generator.random_uniform((), 0, 1, dtype=tf.float32) > 0.5
+    negate = random_generator.uniform((), 0, 1, dtype=tf.float32) > 0.5
     negate = tf.cond(negate, lambda: -1.0, lambda: 1.0)
     return negate
 
 
 def batch_random_inversion(random_generator, batch_size):
     """Same as `random_inversion` but for batched inputs."""
-    negate = random_generator.random_uniform(
-        (batch_size, 1), 0, 1, dtype=tf.float32
-    )
+    negate = random_generator.uniform((batch_size, 1), 0, 1, dtype=tf.float32)
     negate = tf.where(negate > 0.5, -1.0, 1.0)
     return negate
 
diff --git a/keras_cv/utils/preprocessing_test.py b/keras_cv/utils/preprocessing_test.py
index 96ad303d80..73e301ea35 100644
--- a/keras_cv/utils/preprocessing_test.py
+++ b/keras_cv/utils/preprocessing_test.py
@@ -22,7 +22,7 @@ class MockRandomGenerator:
     def __init__(self, value):
         self.value = value
 
-    def random_uniform(self, shape, minval, maxval, dtype=None):
+    def uniform(self, shape, minval, maxval, dtype=None):
         del minval, maxval
         return tf.constant(self.value, dtype=dtype)
 
diff --git a/keras_cv/utils/preset_utils.py b/keras_cv/utils/preset_utils.py
new file mode 100644
index 0000000000..bb0b6b0a0b
--- /dev/null
+++ b/keras_cv/utils/preset_utils.py
@@ -0,0 +1,218 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import inspect
+import json
+import os
+
+from keras_cv.backend import keras
+
+try:
+    import kagglehub
+except ImportError:
+    kagglehub = None
+
+KAGGLE_PREFIX = "kaggle://"
+GS_PREFIX = "gs://"
+
+
+def get_file(preset, path):
+    """Download a preset file in necessary and return the local path."""
+    if not isinstance(preset, str):
+        raise ValueError(
+            f"A preset identifier must be a string. Received: preset={preset}"
+        )
+    if preset.startswith(KAGGLE_PREFIX):
+        if kagglehub is None:
+            raise ImportError(
+                "`from_preset()` requires the `kagglehub` package. "
+                "Please install with `pip install kagglehub`."
+            )
+        # Insert the kaggle framework into the handle.
+        kaggle_handle = preset.removeprefix(KAGGLE_PREFIX)
+        num_segments = len(kaggle_handle.split("/"))
+        if num_segments not in (4, 5):
+            raise ValueError(
+                "Unexpected kaggle preset handle. Kaggle model handles "
+                "should have the form "
+                "kaggle://{org}/{model}/keras/{variant}[/{version}]. "
+                "For example, "
+                "'kaggle://keras/retinanet/keras/retinanet_base_en'. "
+                f"Received: preset={preset}"
+            )
+        return kagglehub.model_download(kaggle_handle, path)
+    elif preset.startswith(GS_PREFIX):
+        url = os.path.join(preset, path)
+        url = url.replace(GS_PREFIX, "https://storage.googleapis.com/")
+        subdir = preset.replace(GS_PREFIX, "gs_")
+        subdir = subdir.replace("/", "_").replace("-", "_")
+        filename = os.path.basename(path)
+        subdir = os.path.join(subdir, os.path.dirname(path))
+        return keras.utils.get_file(
+            filename,
+            url,
+            cache_subdir=os.path.join("models", subdir),
+        )
+    elif os.path.exists(preset):
+        # Assume a local filepath.
+        return os.path.join(preset, path)
+    else:
+        raise ValueError(
+            "Unknown preset identifier. A preset must be a one of:\n"
+            "1) a built in preset identifier like `'mobilenet_v3_small'`\n"
+            "2) a Kaggle Models handle like `'kaggle://keras/mobilenetv3/keras/mobilenet_v3_small'`\n"  # noqa: E501
+            "3) a path to a local preset directory like `'./mobilenet_v3_small`\n"  # noqa: E501
+            "Use `print(cls.presets.keys())` to view all built-in presets for "
+            "API symbol `cls`.\n"
+            f"Received: preset='{preset}'"
+        )
+
+
+def recursive_pop(config, key):
+    """Remove a key from a nested config object"""
+    config.pop(key, None)
+    for value in config.values():
+        if isinstance(value, dict):
+            recursive_pop(value, key)
+        if isinstance(value, list):
+            for v in value:
+                if isinstance(v, dict):
+                    recursive_pop(v, key)
+
+
+def save_to_preset(
+    layer,
+    preset,
+    save_weights=True,
+    config_filename="config.json",
+    weights_filename="model.weights.h5",
+):
+    """Save a KerasCV layer to a preset directory."""
+    os.makedirs(preset, exist_ok=True)
+
+    # Optionally save weights.
+    save_weights = save_weights and hasattr(layer, "save_weights")
+    if save_weights:
+        weights_path = os.path.join(preset, weights_filename)
+        layer.save_weights(weights_path)
+
+    # Save a serialized Keras object.
+    config_path = os.path.join(preset, config_filename)
+    config = keras.saving.serialize_keras_object(layer)
+    # Include references to weights.
+    config["weights"] = weights_filename if save_weights else None
+    recursive_pop(config, "compile_config")
+    recursive_pop(config, "build_config")
+    with open(config_path, "w") as config_file:
+        config_file.write(json.dumps(config, indent=4))
+
+    from keras_cv import __version__ as keras_cv_version
+
+    keras_version = keras.version() if hasattr(keras, "version") else None
+
+    # Save any associated metadata.
+    if config_filename == "config.json":
+        metadata = {
+            "keras_version": keras_version,
+            "keras_cv_version": keras_cv_version,
+            "parameter_count": layer.count_params(),
+            "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
+        }
+        metadata_path = os.path.join(preset, "metadata.json")
+        with open(metadata_path, "w") as metadata_file:
+            metadata_file.write(json.dumps(metadata, indent=4))
+
+
+def load_from_preset(
+    preset,
+    load_weights=None,
+    input_shape=None,
+    config_file="config.json",
+    config_overrides={},
+):
+    """Load a KerasCV layer to a preset directory."""
+    # Load a serialized Keras object.
+    config_path = get_file(preset, config_file)
+    with open(config_path) as config_file:
+        config = json.load(config_file)
+    config["config"] = {**config["config"], **config_overrides}
+    layer = keras.saving.deserialize_keras_object(config)
+    if input_shape is not None:
+        layer.build(input_shape)
+
+    # Check load_weights flag does not violate preset config.
+    if load_weights is True and config["weights"] is None:
+        raise ValueError(
+            f"The specified preset `{preset}` does not include weights. "
+            "Please remove the `load_weights` flag when calling "
+            "`from_preset()` on this preset."
+        )
+    # Default to loading weights if available.
+    if load_weights is not False and config["weights"] is not None:
+        weights_path = get_file(preset, config["weights"])
+        if hasattr(layer, "_layer_checkpoint_dependencies"):
+            legacy_load_weights(layer, weights_path)
+        else:
+            layer.load_weights(weights_path)
+
+    return layer
+
+
+def check_preset_class(
+    preset,
+    classes,
+    config_file="config.json",
+):
+    """Validate a preset is being loaded on the correct class."""
+    config_path = get_file(preset, config_file)
+    try:
+        with open(config_path) as config_file:
+            config = json.load(config_file)
+    except:
+        raise ValueError(
+            f"The specified preset  `{preset}` is unknown. "
+            "Please check documentation to ensure the correct preset "
+            "handle is being used."
+        )
+    cls = keras.saving.get_registered_object(config["registered_name"])
+    if not isinstance(classes, (tuple, list)):
+        classes = (classes,)
+
+    # Subclass checking and alias checking
+    if not any(issubclass(cls, obj) for obj in classes) and not any(
+        issubclass(alias, cls) for alias in classes
+    ):
+        raise ValueError(
+            f"Unexpected class in preset `'{preset}'`. "
+            "When calling `from_preset()` on a class object, the preset class "
+            f"much match allowed classes. Allowed classes are `{classes}`. "
+            f"Received: `{cls}`."
+        )
+    return cls
+
+
+def legacy_load_weights(layer, weights_path):
+    # Hacky fix for TensorFlow 2.13 and 2.14 when loading a `.weights.h5` file.
+    # We find the `Functional` class, and temporarily remove the
+    # `_layer_checkpoint_dependencies` property, which on older version of
+    # TensorFlow complete broke the variable paths for functional models.
+    functional_cls = None
+    for cls in inspect.getmro(layer.__class__):
+        if cls.__name__ == "Functional":
+            functional_cls = cls
+    property = functional_cls._layer_checkpoint_dependencies
+    functional_cls._layer_checkpoint_dependencies = {}
+    layer.load_weights(weights_path)
+    functional_cls._layer_checkpoint_dependencies = property
diff --git a/keras_cv/utils/preset_utils_test.py b/keras_cv/utils/preset_utils_test.py
new file mode 100644
index 0000000000..53546dd0d0
--- /dev/null
+++ b/keras_cv/utils/preset_utils_test.py
@@ -0,0 +1,87 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+
+import numpy as np
+import pytest
+from absl.testing import parameterized
+
+from keras_cv.models import DeepLabV3Plus
+from keras_cv.models import ImageClassifier
+from keras_cv.models import RetinaNet
+from keras_cv.models import YOLOV8Detector
+from keras_cv.tests.test_case import TestCase
+from keras_cv.utils import preset_utils
+
+
+class PresetUtilsTest(TestCase):
+    @parameterized.parameters(
+        (ImageClassifier, "resnet50_v2_imagenet_classifier", "classification"),
+        (
+            ImageClassifier,
+            "efficientnetv2_s_imagenet_classifier",
+            "classification",
+        ),
+        (
+            ImageClassifier,
+            "mobilenet_v3_large_imagenet_classifier",
+            "classification",
+        ),
+        (YOLOV8Detector, "yolo_v8_m_pascalvoc", "detection"),
+        (RetinaNet, "retinanet_resnet50_pascalvoc", "detection"),
+        (DeepLabV3Plus, "deeplab_v3_plus_resnet50_pascalvoc", "segmentation"),
+    )
+    @pytest.mark.large
+    def test_preset_saving(self, cls, preset_name, task_type):
+        save_dir = self.get_temp_dir()
+        if task_type == "detection":
+            model = cls.from_preset(preset_name, bounding_box_format="xywh")
+        else:
+            model = cls.from_preset(preset_name)
+        preset_utils.save_to_preset(model, save_dir)
+
+        # Check existence of files
+        self.assertTrue(os.path.exists(os.path.join(save_dir, "config.json")))
+        self.assertTrue(
+            os.path.exists(os.path.join(save_dir, "model.weights.h5"))
+        )
+        self.assertTrue(os.path.exists(os.path.join(save_dir, "metadata.json")))
+
+        # Check the model config (`config.json`)
+        with open(os.path.join(save_dir, "config.json"), "r") as f:
+            config_json = f.read()
+        self.assertTrue(
+            "build_config" not in config_json
+        )  # Test on raw json to include nested keys
+        self.assertTrue(
+            "compile_config" not in config_json
+        )  # Test on raw json to include nested keys
+        config = json.loads(config_json)
+        self.assertEqual(config["weights"], "model.weights.h5")
+
+        # Try loading the model from preset directory
+        restored_model = preset_utils.load_from_preset(save_dir)
+
+        input_batch = np.ones(shape=(2, 224, 224, 3))
+        expected_output = model(input_batch)
+        restored_output = restored_model(input_batch)
+        self.assertAllClose(expected_output, restored_output)
+
+    def test_preset_errors(self):
+        with self.assertRaisesRegex(ValueError, "must be a string"):
+            ImageClassifier.from_preset(ImageClassifier)
+        with self.assertRaisesRegex(ValueError, "Unknown preset identifier"):
+            ImageClassifier.from_preset("taggle://rednet/rednet/rednet")
diff --git a/keras_cv/version_check.py b/keras_cv/version_check.py
index d419594ff8..c3d32a62d8 100644
--- a/keras_cv/version_check.py
+++ b/keras_cv/version_check.py
@@ -14,8 +14,16 @@
 
 """KerasCV Version check."""
 
-
-import tensorflow as tf
+try:
+    import tensorflow as tf
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "KerasCV uses TensorFlow for its "
+        "preprocessing layers. While this dependency "
+        "will be dropped in the future, please install "
+        "TensorFlow with `pip install tensorflow` to "
+        "use KerasCV"
+    )
 from packaging.version import parse
 
 MIN_VERSION = "2.11.0"
diff --git a/keras_cv/version_utils.py b/keras_cv/version_utils.py
new file mode 100644
index 0000000000..7b0568bdd7
--- /dev/null
+++ b/keras_cv/version_utils.py
@@ -0,0 +1,23 @@
+# Copyright 2023 The KerasCV Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from keras_cv.api_export import keras_cv_export
+
+# Unique source of truth for the version number.
+__version__ = "0.8.3"
+
+
+@keras_cv_export("keras_cv.version")
+def version():
+    return __version__
diff --git a/keras_cv/visualization/plot_bounding_box_gallery.py b/keras_cv/visualization/plot_bounding_box_gallery.py
index 73112458bd..6a46dcc739 100644
--- a/keras_cv/visualization/plot_bounding_box_gallery.py
+++ b/keras_cv/visualization/plot_bounding_box_gallery.py
@@ -16,11 +16,11 @@
 
 import numpy as np
 
-import keras_cv
-import keras_cv.visualization.draw_bounding_boxes
+from keras_cv import bounding_box
 from keras_cv import utils
 from keras_cv.utils import assert_cv2_installed
 from keras_cv.utils import assert_matplotlib_installed
+from keras_cv.visualization.draw_bounding_boxes import draw_bounding_boxes
 from keras_cv.visualization.plot_image_gallery import plot_image_gallery
 
 try:
@@ -125,7 +125,7 @@ def unpackage_tfds_inputs(inputs):
     plotted_images = utils.to_numpy(images)
 
     draw_fn = functools.partial(
-        keras_cv.visualization.draw_bounding_boxes.draw_bounding_boxes,
+        draw_bounding_boxes,
         bounding_box_format="xyxy",
         line_thickness=line_thickness,
         text_thickness=text_thickness,
@@ -136,7 +136,7 @@ def unpackage_tfds_inputs(inputs):
         y_true = y_true.copy()
         y_true["boxes"] = utils.to_numpy(y_true["boxes"])
         y_true["classes"] = utils.to_numpy(y_true["classes"])
-        y_true = keras_cv.bounding_box.convert_format(
+        y_true = bounding_box.convert_format(
             y_true, images=images, source=bounding_box_format, target="xyxy"
         )
         plotted_images = draw_fn(
@@ -150,7 +150,7 @@ def unpackage_tfds_inputs(inputs):
         y_pred = y_pred.copy()
         y_pred["boxes"] = utils.to_numpy(y_pred["boxes"])
         y_pred["classes"] = utils.to_numpy(y_pred["classes"])
-        y_pred = keras_cv.bounding_box.convert_format(
+        y_pred = bounding_box.convert_format(
             y_pred, images=images, source=bounding_box_format, target="xyxy"
         )
         plotted_images = draw_fn(
diff --git a/keras_cv/visualization/plot_image_gallery.py b/keras_cv/visualization/plot_image_gallery.py
index 05cbbad796..4f64330892 100644
--- a/keras_cv/visualization/plot_image_gallery.py
+++ b/keras_cv/visualization/plot_image_gallery.py
@@ -17,9 +17,9 @@
 import numpy as np
 import tensorflow as tf
 
-import keras_cv
 from keras_cv import utils
 from keras_cv.api_export import keras_cv_export
+from keras_cv.backend import ops
 from keras_cv.utils import assert_matplotlib_installed
 
 try:
@@ -46,7 +46,7 @@ def unpack_images(inputs):
         return sample
 
     else:
-        if len(images.shape) != 4:
+        if len(ops.shape(images)) != 4:
             raise ValueError(
                 "`plot_images_gallery()` requires you to "
                 "batch your `np.array` samples together."
@@ -130,7 +130,7 @@ def plot_image_gallery(
         )  # batch_size from within passed `tf.data.Dataset`
     else:
         batch_size = (
-            images.shape[0] if len(images.shape) == 4 else 1
+            ops.shape(images)[0] if len(ops.shape(images)) == 4 else 1
         )  # batch_size from np.array or single image
 
     rows = rows or int(math.ceil(math.sqrt(batch_size)))
@@ -160,7 +160,7 @@ def plot_image_gallery(
         fig.legend(handles=legend_handles, loc="lower center")
 
     # Perform image range transform
-    images = keras_cv.utils.transform_value_range(
+    images = utils.transform_value_range(
         images, original_range=value_range, target_range=(0, 255)
     )
     images = utils.to_numpy(images)
diff --git a/pip_build.py b/pip_build.py
index c1c895e187..ba61963697 100644
--- a/pip_build.py
+++ b/pip_build.py
@@ -28,6 +28,7 @@
 ```
 """
 import argparse
+import datetime
 import glob
 import os
 import pathlib
@@ -45,56 +46,103 @@
 ]
 
 
-def build():
-    if os.path.exists(build_directory):
-        raise ValueError(f"Directory already exists: {build_directory}")
+def ignore_files(_, filenames):
+    return [f for f in filenames if "_test" in f]
+
+
+def export_version_string(version, is_nightly=False):
+    """Export Version and Package Name."""
+    if is_nightly:
+        date = datetime.datetime.now()
+        version += f".dev{date.strftime('%Y%m%d%H')}"
+        # Replaces `name="keras-cv"` in `setup.py` with `keras-cv-nightly`
+        with open("setup.py") as f:
+            setup_contents = f.read()
+        with open("setup.py", "w") as f:
+            setup_contents = setup_contents.replace(
+                'name="keras-cv"', 'name="keras-cv-nightly"'
+            )
+            f.write(setup_contents)
+
+    # Overwrite the version string with our package version.
+    with open(os.path.join(package, "src", "version_utils.py")) as f:
+        version_contents = f.readlines()
+    with open(os.path.join(package, "src", "version_utils.py"), "w") as f:
+        for line in version_contents:
+            if line.startswith("__version__"):
+                f.write(f'__version__ = "{version}"\n')
+            else:
+                f.write(line)
+
+    # Make sure to export the __version__ string
+    with open(os.path.join(package, "__init__.py")) as f:
+        init_contents = f.read()
+    with open(os.path.join(package, "__init__.py"), "w") as f:
+        f.write(init_contents)
+        f.write("from keras_cv.src.version_utils import __version__\n")
+
+
+def copy_source_to_build_directory(root_path):
+    # Copy sources (`keras_cv/` directory and setup files) to build
+    # directory
+    os.chdir(root_path)
+    os.mkdir(build_directory)
+    shutil.copytree(
+        package, os.path.join(build_directory, package), ignore=ignore_files
+    )
+    for fname in to_copy:
+        shutil.copy(fname, os.path.join(f"{build_directory}", fname))
+    os.chdir(build_directory)
+
+
+def run_namex_conversion():
+    # Restructure the codebase so that source files live in `keras_cv/src`
+    namex.convert_codebase(package, code_directory="src")
 
+    # Generate API __init__.py files in `keras_cv/`
+    namex.generate_api_files(package, code_directory="src", verbose=True)
+
+
+def build_and_save_output(root_path, __version__):
+    """Build the package."""
+    os.system("python3 -m build")
+
+    # Save the dist files generated by the build process
+    os.chdir(root_path)
+    if not os.path.exists(dist_directory):
+        os.mkdir(dist_directory)
+    for fpath in glob.glob(
+        os.path.join(build_directory, dist_directory, "*.*")
+    ):
+        shutil.copy(fpath, dist_directory)
+
+    # Find the .whl file path
     whl_path = None
-    try:
-        # Copy sources (`keras_cv/` directory and setup files) to build dir
-        root_path = pathlib.Path(__file__).parent.resolve()
-        os.chdir(root_path)
-        os.mkdir(build_directory)
-        shutil.copytree(package, os.path.join(build_directory, package))
-        for fname in to_copy:
-            shutil.copy(fname, os.path.join(f"{build_directory}", fname))
-        os.chdir(build_directory)
+    for fname in os.listdir(dist_directory):
+        if __version__ in fname and fname.endswith(".whl"):
+            whl_path = os.path.abspath(os.path.join(dist_directory, fname))
+    if whl_path:
+        print(f"Build successful. Wheel file available at {whl_path}")
+    else:
+        print("Build failed.")
+    return whl_path
 
-        # Restructure the codebase so that source files live in `keras_cv/src`
-        namex.convert_codebase(package, code_directory="src")
 
-        # Generate API __init__.py files in `keras_cv/`
-        namex.generate_api_files(package, code_directory="src", verbose=True)
+def build(root_path, is_nightly=False):
+    if os.path.exists(build_directory):
+        raise ValueError(f"Directory already exists: {build_directory}")
 
+    try:
+        copy_source_to_build_directory(root_path)
+        run_namex_conversion()
         # Make sure to export the __version__ string
         from keras_cv.src import __version__  # noqa: E402
 
-        with open(os.path.join(package, "__init__.py")) as f:
-            init_contents = f.read()
-        with open(os.path.join(package, "__init__.py"), "w") as f:
-            f.write(init_contents + "\n\n" + f'__version__ = "{__version__}"\n')
-
-        # Build the package
-        os.system("python3 -m build")
-
-        # Save the dist files generated by the build process
-        os.chdir(root_path)
-        if not os.path.exists(dist_directory):
-            os.mkdir(dist_directory)
-        for fpath in glob.glob(
-            os.path.join(build_directory, dist_directory, "*.*")
-        ):
-            shutil.copy(fpath, dist_directory)
-
-        # Find the .whl file path
-        for fname in os.listdir(dist_directory):
-            if __version__ in fname and fname.endswith(".whl"):
-                whl_path = os.path.abspath(os.path.join(dist_directory, fname))
-        print(f"Build successful. Wheel file available at {whl_path}")
+        export_version_string(__version__, is_nightly)
+        return build_and_save_output(root_path, __version__)
     finally:
         # Clean up: remove the build directory (no longer needed)
         shutil.rmtree(build_directory)
-    return whl_path
 
 
 def install_whl(whl_fpath):
@@ -109,7 +157,13 @@ def install_whl(whl_fpath):
         action="store_true",
         help="Whether to install the generated wheel file.",
     )
+    parser.add_argument(
+        "--nightly",
+        action="store_true",
+        help="Whether to generate nightly wheel file.",
+    )
     args = parser.parse_args()
-    whl_path = build()
+    root_path = pathlib.Path(__file__).parent.resolve()
+    whl_path = build(root_path, args.nightly)
     if whl_path and args.install:
         install_whl(whl_path)
diff --git a/requirements-common.txt b/requirements-common.txt
new file mode 100644
index 0000000000..fc21cc5f96
--- /dev/null
+++ b/requirements-common.txt
@@ -0,0 +1,16 @@
+# Library deps.
+absl-py
+regex
+pandas
+kagglehub
+keras-core>=0.1.6
+tensorflow-datasets
+pycocotools
+# Tooling deps.
+packaging
+flake8
+isort
+black
+pytest
+build
+namex
\ No newline at end of file
diff --git a/requirements-jax-cuda.txt b/requirements-jax-cuda.txt
new file mode 100644
index 0000000000..6f5ff522e0
--- /dev/null
+++ b/requirements-jax-cuda.txt
@@ -0,0 +1,13 @@
+# Tensorflow cpu-only version.
+tf-nightly-cpu==2.16.0.dev20240104  # Pin a working nightly until rc0.
+
+# Torch cpu-only version.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax with cuda support.
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jax[cuda12_pip]
+
+-r requirements-common.txt
\ No newline at end of file
diff --git a/requirements-tensorflow-cuda.txt b/requirements-tensorflow-cuda.txt
new file mode 100644
index 0000000000..3889e83934
--- /dev/null
+++ b/requirements-tensorflow-cuda.txt
@@ -0,0 +1,12 @@
+# Tensorflow with cuda support.
+tf-nightly[and-cuda]==2.16.0.dev20240104  # Pin a working nightly until rc0.
+
+# Torch cpu-only version.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax cpu-only version.
+jax[cpu]
+
+-r requirements-common.txt
\ No newline at end of file
diff --git a/requirements-torch-cuda.txt b/requirements-torch-cuda.txt
new file mode 100644
index 0000000000..5bfbc37546
--- /dev/null
+++ b/requirements-torch-cuda.txt
@@ -0,0 +1,12 @@
+# Tensorflow cpu-only version.
+tf-nightly-cpu==2.16.0.dev20240104  # Pin a working nightly until rc0.
+
+# Torch with cuda support.
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.1.2+cu121
+torchvision==0.16.2+cu121
+
+# Jax cpu-only version.
+jax[cpu]
+
+-r requirements-common.txt
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index eb83dc9544..02792b29b7 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,14 +1,12 @@
-absl-py
-packaging
-pandas
-tensorflow
-tensorflow-datasets
-flake8
-regex
-isort
-black
-pytest
-pycocotools
-keras-core
-build
-namex
+# Tensorflow.
+tf-nightly-cpu==2.16.0.dev20240104  # Pin a working nightly until rc0.
+
+# Torch.
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch>=2.1.0
+torchvision>=0.16.0
+
+# Jax.
+jax[cpu]
+
+-r requirements-common.txt
diff --git a/setup.cfg b/setup.cfg
index 7b52664a31..9e2ba6cfa5 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
-license_file = LICENSE
-description-file = README.md
+license_files = LICENSE
+description_file = README.md
 version = attr: keras_cv.__version__
 
 [tool:pytest]
diff --git a/setup.py b/setup.py
index 513519b0a2..19dc42248c 100644
--- a/setup.py
+++ b/setup.py
@@ -21,6 +21,21 @@
 from setuptools import setup
 from setuptools.dist import Distribution
 
+
+def read(rel_path):
+    here = os.path.abspath(os.path.dirname(__file__))
+    with open(os.path.join(here, rel_path)) as fp:
+        return fp.read()
+
+
+def get_version(rel_path):
+    for line in read(rel_path).splitlines():
+        if line.startswith("__version__"):
+            delim = '"' if '"' in line else "'"
+            return line.split(delim)[1]
+    raise RuntimeError("Unable to find version string.")
+
+
 BUILD_WITH_CUSTOM_OPS = (
     "BUILD_WITH_CUSTOM_OPS" in os.environ
     and os.environ["BUILD_WITH_CUSTOM_OPS"] == "true"
@@ -28,6 +43,10 @@
 
 HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
+if os.path.exists("keras_cv/version_utils.py"):
+    VERSION = get_version("keras_cv/version_utils.py")
+else:
+    VERSION = get_version("keras_cv/src/version_utils.py")
 
 
 class BinaryDistribution(Distribution):
@@ -45,6 +64,7 @@ def is_pure(self):
     description="Industry-strength computer Vision extensions for Keras.",
     long_description=README,
     long_description_content_type="text/markdown",
+    version=VERSION,
     url="https://github.com/keras-team/keras-cv",
     author="Keras team",
     author_email="keras-cv@google.com",
@@ -55,6 +75,7 @@ def is_pure(self):
         "regex",
         "tensorflow-datasets",
         "keras-core",
+        "kagglehub",
     ],
     extras_require={
         "tests": [
@@ -63,18 +84,15 @@ def is_pure(self):
             "black[jupyter]",
             "pytest",
             "pycocotools",
-            "tensorflow",
-            "keras-core",
         ],
         "examples": ["tensorflow_datasets", "matplotlib"],
     },
     distclass=BinaryDistribution,
     # Supported Python versions
-    python_requires=">=3.8",
+    python_requires=">=3.9",
     classifiers=[
         "Development Status :: 3 - Alpha",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",