diff --git a/src/ModelCustomization.ts b/src/ModelCustomization.ts
index d718222..c2f9bfe 100644
--- a/src/ModelCustomization.ts
+++ b/src/ModelCustomization.ts
@@ -128,10 +128,26 @@ function removeColumnDefaultProperties(
     });
     return dbModel;
 }
+
+function findFileImports(dbModel: Entity[]) {
+    dbModel.forEach(entity => {
+        entity.relations.forEach(relation => {
+            if (
+                relation.relatedTable !== entity.tscName &&
+                !entity.fileImports.some(v => v === relation.relatedTable)
+            ) {
+                entity.fileImports.push(relation.relatedTable);
+            }
+        });
+    });
+    return dbModel;
+}
+
 function addImportsAndGenerationOptions(
     dbModel: Entity[],
     generationOptions: IGenerationOptions
 ): Entity[] {
+    dbModel = findFileImports(dbModel);
     dbModel.forEach(entity => {
         entity.relations.forEach(relation => {
             if (generationOptions.lazy) {
diff --git a/src/drivers/AbstractDriver.ts b/src/drivers/AbstractDriver.ts
index 2b0072a..bb6ee21 100644
--- a/src/drivers/AbstractDriver.ts
+++ b/src/drivers/AbstractDriver.ts
@@ -207,20 +207,6 @@ export default abstract class AbstractDriver {
         );
         await this.DisconnectFromServer();
         dbModel = AbstractDriver.FindManyToManyRelations(dbModel);
-        dbModel = AbstractDriver.FindFileImports(dbModel);
-        return dbModel;
-    }
-
-    private static FindFileImports(dbModel: Entity[]) {
-        dbModel.forEach(entity => {
-            entity.relations.forEach(relation => {
-                if (
-                    !entity.fileImports.some(v => v === relation.relatedTable)
-                ) {
-                    entity.fileImports.push(relation.relatedTable);
-                }
-            });
-        });
         return dbModel;
     }