diff --git a/go.mod b/go.mod
index 75bd941..c2eefb4 100644
--- a/go.mod
+++ b/go.mod
@@ -5,7 +5,7 @@ go 1.23.7
 require (
 	github.com/containerd/containerd/v2 v2.0.4
 	github.com/containerd/platforms v1.0.0-rc.1
-	github.com/docker/model-distribution v0.0.0-20250410151231-bf9b59b512f7
+	github.com/docker/model-distribution v0.0.0-20250411163353-b33595b4e207
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.1.1
 	github.com/sirupsen/logrus v1.9.3
diff --git a/go.sum b/go.sum
index 91d2f9b..ade394a 100644
--- a/go.sum
+++ b/go.sum
@@ -51,6 +51,8 @@ github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4
 github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
 github.com/docker/model-distribution v0.0.0-20250410151231-bf9b59b512f7 h1:li7LReF/bddqOXbX7IfAnzRmYYAuMjG4C7xCMQLZPlI=
 github.com/docker/model-distribution v0.0.0-20250410151231-bf9b59b512f7/go.mod h1:/JWSwYc3pihCpHqFzDUyoiRKegA1srfYESxRh/vJE10=
+github.com/docker/model-distribution v0.0.0-20250411163353-b33595b4e207 h1:BZFQGpeo7H4JLeX1Gn+T9P7vSPwbhtH10QeSKdkzKKs=
+github.com/docker/model-distribution v0.0.0-20250411163353-b33595b4e207/go.mod h1:/JWSwYc3pihCpHqFzDUyoiRKegA1srfYESxRh/vJE10=
 github.com/ebitengine/purego v0.8.2 h1:jPPGWs2sZ1UgOSgD2bClL0MJIqu58nOmIcBuXr62z1I=
 github.com/ebitengine/purego v0.8.2/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=
 github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
diff --git a/pkg/inference/models/api.go b/pkg/inference/models/api.go
index c285754..68bb255 100644
--- a/pkg/inference/models/api.go
+++ b/pkg/inference/models/api.go
@@ -3,7 +3,7 @@ package models
 import (
 	"fmt"
 
-	"github.com/docker/model-distribution/pkg/types"
+	"github.com/docker/model-distribution/types"
 )
 
 // ModelCreateRequest represents a model create request. It is designed to
diff --git a/pkg/inference/models/manager.go b/pkg/inference/models/manager.go
index 8e0ae2b..d91dfd8 100644
--- a/pkg/inference/models/manager.go
+++ b/pkg/inference/models/manager.go
@@ -8,8 +8,8 @@ import (
 	"html"
 	"net/http"
 
-	"github.com/docker/model-distribution/pkg/distribution"
-	"github.com/docker/model-distribution/pkg/types"
+	"github.com/docker/model-distribution/distribution"
+	"github.com/docker/model-distribution/types"
 	"github.com/docker/model-runner/pkg/inference"
 	"github.com/docker/model-runner/pkg/logging"
 	"github.com/sirupsen/logrus"
@@ -219,6 +219,10 @@ func (m *Manager) handleDeleteModel(w http.ResponseWriter, r *http.Request) {
 
 	err := m.distributionClient.DeleteModel(r.PathValue("name"))
 	if err != nil {
+		if errors.Is(err, distribution.ErrModelNotFound) {
+			http.Error(w, err.Error(), http.StatusNotFound)
+			return
+		}
 		m.log.Warnln("Error while deleting model:", err)
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
diff --git a/pkg/inference/scheduling/scheduler.go b/pkg/inference/scheduling/scheduler.go
index 90041b1..2e11f11 100644
--- a/pkg/inference/scheduling/scheduler.go
+++ b/pkg/inference/scheduling/scheduler.go
@@ -9,7 +9,7 @@ import (
 	"io"
 	"net/http"
 
-	"github.com/docker/model-distribution/pkg/distribution"
+	"github.com/docker/model-distribution/distribution"
 	"github.com/docker/model-runner/pkg/inference"
 	"github.com/docker/model-runner/pkg/inference/models"
 	"github.com/docker/model-runner/pkg/logging"