diff --git a/blocksuite/framework/store/src/model/block/draft.ts b/blocksuite/framework/store/src/model/block/draft.ts
index 538ee91b23..48baa54a95 100644
--- a/blocksuite/framework/store/src/model/block/draft.ts
+++ b/blocksuite/framework/store/src/model/block/draft.ts
@@ -16,27 +16,15 @@ export function toDraftModel<Model extends BlockModel = BlockModel>(
 ): DraftModel<Model> {
   const { id, version, flavour, role, keys, text, children } = origin;
 
-  // Process props data
-  let props: ModelProps<Model>;
-
-  // Check if it's a flat data structure (has _props property)
-  if ('_props' in origin && origin._props) {
-    // Flat data structure, get directly from _props
-    props = Object.entries(origin._props).reduce((acc, [key, value]) => {
-      return {
-        ...acc,
-        [key]: value,
-      };
-    }, {} as ModelProps<Model>);
-  } else {
-    // Non-flat data structure, get through keys array
-    props = origin.keys.reduce((acc, key) => {
-      return {
-        ...acc,
-        [key]: origin[key as keyof Model],
-      };
-    }, {} as ModelProps<Model>);
-  }
+  const isFlatData = origin.schema.model.isFlatData;
+  const props = origin.keys.reduce((acc, key) => {
+    const target = isFlatData ? origin.props : origin;
+    const value = target[key as keyof typeof target];
+    return {
+      ...acc,
+      [key]: value,
+    };
+  }, {} as ModelProps<Model>);
 
   return {
     id,