diff --git a/src/main/java/com/itfsw/mybatis/generator/plugins/ExampleEnhancedPlugin.java b/src/main/java/com/itfsw/mybatis/generator/plugins/ExampleEnhancedPlugin.java
index 80643c6..dfc88c1 100644
--- a/src/main/java/com/itfsw/mybatis/generator/plugins/ExampleEnhancedPlugin.java
+++ b/src/main/java/com/itfsw/mybatis/generator/plugins/ExampleEnhancedPlugin.java
@@ -106,7 +106,6 @@ public class ExampleEnhancedPlugin extends BasePlugin {
         // 添加接口CriteriaAdd
         InnerInterface criteriaAddInterface = new InnerInterface("ICriteriaAdd");
         criteriaAddInterface.setVisibility(JavaVisibility.PUBLIC);
-        commentGenerator.addInterfaceComment(criteriaAddInterface, introspectedTable);
         logger.debug("itfsw(Example增强插件):" + topLevelClass.getType().getShortName() + "." + innerClass.getType().getShortName() + "增加接口ICriteriaAdd");
 
         // ICriteriaAdd增加接口add
@@ -121,6 +120,8 @@ public class ExampleEnhancedPlugin extends BasePlugin {
         logger.debug("itfsw(Example增强插件):" + topLevelClass.getType().getShortName() + "." + innerClass.getType().getShortName() + "." + criteriaAddInterface.getType().getShortName() + "增加方法add");
 
         InnerClass innerClassWrapper = new InnerInterfaceWrapperToInnerClass(criteriaAddInterface);
+        // 添加注释
+        commentGenerator.addClassComment(innerClassWrapper, introspectedTable);
         innerClass.addInnerClass(innerClassWrapper);
 
         // 添加andIf方法
diff --git a/src/main/java/com/itfsw/mybatis/generator/plugins/utils/BasePlugin.java b/src/main/java/com/itfsw/mybatis/generator/plugins/utils/BasePlugin.java
index ad1efad..2efe648 100644
--- a/src/main/java/com/itfsw/mybatis/generator/plugins/utils/BasePlugin.java
+++ b/src/main/java/com/itfsw/mybatis/generator/plugins/utils/BasePlugin.java
@@ -16,8 +16,10 @@
 
 package com.itfsw.mybatis.generator.plugins.utils;
 
-import com.itfsw.mybatis.generator.plugins.utils.enhanced.IWCommentGenerator;
+import com.itfsw.mybatis.generator.plugins.utils.enhanced.DefaultCommentGenerator;
+import org.mybatis.generator.api.CommentGenerator;
 import org.mybatis.generator.api.PluginAdapter;
+import org.mybatis.generator.config.Context;
 import org.mybatis.generator.internal.util.StringUtility;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -34,7 +36,27 @@ import java.util.List;
  */
 public class BasePlugin extends PluginAdapter {
     protected static final Logger logger = LoggerFactory.getLogger(BasePlugin.class); // 日志
-    protected IWCommentGenerator commentGenerator = new IWCommentGenerator();  // 注释工具
+    protected CommentGenerator commentGenerator;  // 注释工具
+
+    /**
+     * Set the context under which this plugin is running.
+     *
+     * @param context
+     *            the new context
+     */
+    @Override
+    public void setContext(Context context) {
+        super.setContext(context);
+
+        // 配置插件使用的模板引擎
+        if (context.getCommentGenerator() instanceof org.mybatis.generator.internal.DefaultCommentGenerator){
+            // 使用默认模板引擎
+            commentGenerator = new DefaultCommentGenerator();
+        } else {
+            // 用户自定义
+            commentGenerator = context.getCommentGenerator();
+        }
+    }
 
     /**
      * {@inheritDoc}
diff --git a/src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/IWCommentGenerator.java b/src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/DefaultCommentGenerator.java
similarity index 89%
rename from src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/IWCommentGenerator.java
rename to src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/DefaultCommentGenerator.java
index 4e1a40b..9c826be 100644
--- a/src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/IWCommentGenerator.java
+++ b/src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/DefaultCommentGenerator.java
@@ -35,7 +35,7 @@ import java.util.Properties;
  * @time:2017/4/28 15:11
  * ---------------------------------------------------------------------------
  */
-public class IWCommentGenerator implements CommentGenerator {
+public class DefaultCommentGenerator implements CommentGenerator {
     /* (non-Javadoc)
      * @see org.mybatis.generator.api.CommentGenerator#addJavaFileComment(org.mybatis.generator.api.dom.java.CompilationUnit)
      */
@@ -97,39 +97,37 @@ public class IWCommentGenerator implements CommentGenerator {
      * @see org.mybatis.generator.api.CommentGenerator#addClassComment(org.mybatis.generator.api.dom.java.InnerClass, org.mybatis.generator.api.IntrospectedTable)
      */
     public void addClassComment(InnerClass innerClass, IntrospectedTable introspectedTable) {
-        StringBuilder sb = new StringBuilder();
+        if (innerClass instanceof InnerInterfaceWrapperToInnerClass){
+            InnerInterfaceWrapperToInnerClass innerInterfaceWrapperToInnerClass = (InnerInterfaceWrapperToInnerClass) innerClass;
+            InnerInterface innerInterface = innerInterfaceWrapperToInnerClass.getInnerInterface();
 
-        innerClass.addJavaDocLine("/**");
-        innerClass.addJavaDocLine(" * This class was generated by MyBatis Generator.");
+            StringBuilder sb = new StringBuilder();
 
-        sb.append(" * This class corresponds to the database table ");
-        sb.append(introspectedTable.getFullyQualifiedTable());
-        innerClass.addJavaDocLine(sb.toString());
+            innerInterface.addJavaDocLine("/**");
+            innerInterface.addJavaDocLine(" * This interface was generated by MyBatis Generator.");
 
-        addJavadocTag(innerClass, false);
+            sb.append(" * This interface corresponds to the database table ");
+            sb.append(introspectedTable.getFullyQualifiedTable());
+            innerInterface.addJavaDocLine(sb.toString());
 
-        innerClass.addJavaDocLine(" */");
-    }
+            addJavadocTag(innerInterface, false);
 
-    /**
-     * inner inter face 注释
-     *
-     * @param innerInterface
-     * @param introspectedTable
-     */
-    public void addInterfaceComment(Interface innerInterface, IntrospectedTable introspectedTable) {
-        StringBuilder sb = new StringBuilder();
+            innerInterface.addJavaDocLine(" */");
+        } else {
+            StringBuilder sb = new StringBuilder();
 
-        innerInterface.addJavaDocLine("/**");
-        innerInterface.addJavaDocLine(" * This interface was generated by MyBatis Generator.");
+            innerClass.addJavaDocLine("/**");
+            innerClass.addJavaDocLine(" * This class was generated by MyBatis Generator.");
 
-        sb.append(" * This interface corresponds to the database table ");
-        sb.append(introspectedTable.getFullyQualifiedTable());
-        innerInterface.addJavaDocLine(sb.toString());
+            sb.append(" * This class corresponds to the database table ");
+            sb.append(introspectedTable.getFullyQualifiedTable());
+            innerClass.addJavaDocLine(sb.toString());
 
-        addJavadocTag(innerInterface, false);
+            addJavadocTag(innerClass, false);
+
+            innerClass.addJavaDocLine(" */");
+        }
 
-        innerInterface.addJavaDocLine(" */");
     }
 
     /* (non-Javadoc)
diff --git a/src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/InnerInterfaceWrapperToInnerClass.java b/src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/InnerInterfaceWrapperToInnerClass.java
index 326d289..91a4652 100644
--- a/src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/InnerInterfaceWrapperToInnerClass.java
+++ b/src/main/java/com/itfsw/mybatis/generator/plugins/utils/enhanced/InnerInterfaceWrapperToInnerClass.java
@@ -56,4 +56,22 @@ public class InnerInterfaceWrapperToInnerClass extends InnerClass{
     public String getFormattedContent(int indentLevel, CompilationUnit compilationUnit) {
         return this.innerInterface.getFormattedContent(indentLevel, compilationUnit);
     }
+
+    /**
+     * Getter method for property <tt>innerInterface</tt>.
+     * @return property value of innerInterface
+     * @author hewei
+     */
+    public InnerInterface getInnerInterface() {
+        return innerInterface;
+    }
+
+    /**
+     * Setter method for property <tt>innerInterface</tt>.
+     * @param innerInterface value to be assigned to property innerInterface
+     * @author hewei
+     */
+    public void setInnerInterface(InnerInterface innerInterface) {
+        this.innerInterface = innerInterface;
+    }
 }