diff --git a/client.go b/client.go
index d13509b77..0301eeb0d 100644
--- a/client.go
+++ b/client.go
@@ -49,7 +49,6 @@ import (
 	"github.com/containerd/containerd/leases"
 	leasesproxy "github.com/containerd/containerd/leases/proxy"
 	"github.com/containerd/containerd/namespaces"
-	"github.com/containerd/containerd/oci"
 	"github.com/containerd/containerd/pkg/dialer"
 	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/containerd/plugin"
@@ -544,31 +543,16 @@ func (c *Client) Restore(ctx context.Context, id, ref string, opts ...RestoreOpt
 		return err
 	}
 
-	// get image from annotation
-	imageName, ok := index.Annotations["image.name"]
-	if !ok {
-		return ErrCheckpointIndexImageNameNotFound
-	}
-
-	image, err := c.Pull(ctx, imageName, WithPullUnpack)
-	if err != nil {
-		return err
-	}
-
 	ctx, done, err := c.WithLease(ctx)
 	if err != nil {
 		return err
 	}
 	defer done(ctx)
 
-	// container options
-	copts := []NewContainerOpts{
-		WithNewSpec(oci.WithImageConfig(image)),
-		WithNewSnapshot(id, image),
-	}
+	copts := []NewContainerOpts{}
 	topts := []NewTaskOpts{}
 	for _, o := range opts {
-		co, to, err := o(ctx, c, checkpoint, index)
+		co, to, err := o(ctx, id, c, checkpoint, index)
 		if err != nil {
 			return err
 		}
diff --git a/cmd/ctr/commands/containers/containers.go b/cmd/ctr/commands/containers/containers.go
index 0b1d4dbbd..72b0beae8 100644
--- a/cmd/ctr/commands/containers/containers.go
+++ b/cmd/ctr/commands/containers/containers.go
@@ -288,7 +288,7 @@ var infoCommand = cli.Command{
 var checkpointCommand = cli.Command{
 	Name:      "checkpoint",
 	Usage:     "checkpoint a container",
-	ArgsUsage: "CONTAINER REF [flags]",
+	ArgsUsage: "CONTAINER REF",
 	Flags: []cli.Flag{
 		cli.BoolFlag{
 			Name:  "rw",
@@ -342,7 +342,7 @@ var checkpointCommand = cli.Command{
 var restoreCommand = cli.Command{
 	Name:      "restore",
 	Usage:     "restore a container from checkpoint",
-	ArgsUsage: "CONTAINER REF [flags]",
+	ArgsUsage: "CONTAINER REF",
 	Flags: []cli.Flag{
 		cli.BoolFlag{
 			Name:  "live",
@@ -358,18 +358,20 @@ var restoreCommand = cli.Command{
 		if ref == "" {
 			return errors.New("ref must be provided")
 		}
-		opts := []containerd.RestoreOpts{
-			containerd.WithRestoreSpec,
-			containerd.WithRestoreRuntime,
-		}
-		if context.Bool("live") {
-			opts = append(opts, containerd.WithRestoreLive)
-		}
 		client, ctx, cancel, err := commands.NewClient(context)
 		if err != nil {
 			return err
 		}
 		defer cancel()
+
+		opts := []containerd.RestoreOpts{
+			containerd.WithRestoreSpec,
+			containerd.WithRestoreSnapshot,
+			containerd.WithRestoreRuntime,
+		}
+		if context.Bool("live") {
+			opts = append(opts, containerd.WithRestoreLive)
+		}
 		if err := client.Restore(ctx, id, ref, opts...); err != nil {
 			return err
 		}
diff --git a/container_checkpoint_opts.go b/container_checkpoint_opts.go
index 0af5f8ce2..b0d618070 100644
--- a/container_checkpoint_opts.go
+++ b/container_checkpoint_opts.go
@@ -36,10 +36,6 @@ var (
 	ErrCheckpointRWUnsupported = errors.New("rw checkpoint is only supported on v2 runtimes")
 	// ErrMediaTypeNotFound returns an error when a media type in the manifest is unknown
 	ErrMediaTypeNotFound = errors.New("media type not found")
-	// ErrCheckpointIndexImageNameNotFound is returned when the checkpoint image name is not present in the index
-	ErrCheckpointIndexImageNameNotFound = errors.New("image name not present in index")
-	// ErrCheckpointIndexRuntimeNameNotFound is returned when the checkpoint runtime name is not present in the index
-	ErrCheckpointIndexRuntimeNameNotFound = errors.New("runtime name not present in index")
 )
 
 // CheckpointOpts are options to manage the checkpoint operation
diff --git a/container_restore_opts.go b/container_restore_opts.go
index 52c343ec6..3b23bc57f 100644
--- a/container_restore_opts.go
+++ b/container_restore_opts.go
@@ -22,25 +22,34 @@ import (
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/images"
 	"github.com/containerd/containerd/oci"
+	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/typeurl"
 	"github.com/gogo/protobuf/proto"
 	ptypes "github.com/gogo/protobuf/types"
+	"github.com/opencontainers/image-spec/identity"
 	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
 )
 
+var (
+	// ErrCheckpointIndexImageNameNotFound is returned when the checkpoint image name is not present in the index
+	ErrCheckpointIndexImageNameNotFound = errors.New("image name not present in index")
+	// ErrCheckpointIndexRuntimeNameNotFound is returned when the checkpoint runtime name is not present in the index
+	ErrCheckpointIndexRuntimeNameNotFound = errors.New("runtime name not present in index")
+)
+
 // RestoreOpts are options to manage the restore operation
-type RestoreOpts func(context.Context, *Client, Image, *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error)
+type RestoreOpts func(context.Context, string, *Client, Image, *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error)
 
 // WithRestoreLive restores the runtime and memory data for the container
-func WithRestoreLive(ctx context.Context, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
+func WithRestoreLive(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
 	return nil, []NewTaskOpts{
 		WithTaskCheckpoint(checkpoint),
 	}, nil
 }
 
 // WithRestoreRuntime restores the runtime for the container
-func WithRestoreRuntime(ctx context.Context, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
+func WithRestoreRuntime(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
 	runtimeName, ok := index.Annotations["runtime.name"]
 	if !ok {
 		return nil, nil, ErrCheckpointIndexRuntimeNameNotFound
@@ -69,7 +78,7 @@ func WithRestoreRuntime(ctx context.Context, client *Client, checkpoint Image, i
 }
 
 // WithRestoreSpec restores the spec from the checkpoint for the container
-func WithRestoreSpec(ctx context.Context, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
+func WithRestoreSpec(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
 	m, err := GetIndexByMediaType(index, images.MediaTypeContainerd1CheckpointConfig)
 	if err != nil {
 		return nil, nil, err
@@ -93,3 +102,28 @@ func WithRestoreSpec(ctx context.Context, client *Client, checkpoint Image, inde
 		WithSpec(spec),
 	}, nil, nil
 }
+
+func WithRestoreSnapshot(ctx context.Context, id string, client *Client, checkpoint Image, index *imagespec.Index) ([]NewContainerOpts, []NewTaskOpts, error) {
+	// get image from annotation
+	imageName, ok := index.Annotations["image.name"]
+	if !ok {
+		return nil, nil, ErrCheckpointIndexImageNameNotFound
+	}
+	i, err := client.Pull(ctx, imageName, WithPullUnpack)
+	if err != nil {
+		return nil, nil, err
+	}
+	diffIDs, err := i.(*image).i.RootFS(ctx, client.ContentStore(), platforms.Default())
+	if err != nil {
+		return nil, nil, err
+	}
+	//setSnapshotterIfEmpty(client)
+	parent := identity.ChainID(diffIDs).String()
+	if _, err := client.SnapshotService(DefaultSnapshotter).Prepare(ctx, id, parent); err != nil {
+		return nil, nil, err
+	}
+	return []NewContainerOpts{
+		WithImage(i),
+		WithSnapshot(id),
+	}, nil, nil
+}