diff --git a/src/module.ts b/src/module.ts
index 5be8c361..f23897c2 100644
--- a/src/module.ts
+++ b/src/module.ts
@@ -2,25 +2,6 @@ import { defu } from 'defu'
 import { createResolver, defineNuxtModule, addComponentsDir, addImportsDir, addVitePlugin, addPlugin, installModule } from '@nuxt/kit'
 import tailwindcss from '@tailwindcss/vite'
 import createTemplates from './templates'
-// import type { DeepPartial } from './runtime/types'
-// import * as theme from './runtime/theme'
-
-// type UI = {
-//   primary?: string
-//   gray?: string
-//   [key: string]: any
-// } & DeepPartial<typeof theme>
-
-// declare module 'nuxt/schema' {
-//   interface AppConfigInput {
-//     ui?: UI
-//   }
-// }
-// declare module '@nuxt/schema' {
-//   interface AppConfigInput {
-//     ui?: UI
-//   }
-// }
 
 export interface ModuleOptions {
   colors: string[]
diff --git a/src/runtime/components/Container.vue b/src/runtime/components/Container.vue
index 721f4393..6049447f 100644
--- a/src/runtime/components/Container.vue
+++ b/src/runtime/components/Container.vue
@@ -1,10 +1,9 @@
 <script lang="ts">
 import { tv } from 'tailwind-variants'
-// import appConfig from '#build/app.config'
+import appConfig from '#build/app.config'
 import theme from '#build/ui/container'
 
-const container = tv(theme)
-// const appContainer = tv({ extend: container, ...(appConfig.ui?.container || {}) })
+const container = tv({ extend: tv(theme), ...(appConfig.ui?.container || {}) })
 
 export interface ContainerProps {
   as?: string
diff --git a/src/runtime/components/Tooltip.vue b/src/runtime/components/Tooltip.vue
index a8377c8f..abed431b 100644
--- a/src/runtime/components/Tooltip.vue
+++ b/src/runtime/components/Tooltip.vue
@@ -5,7 +5,7 @@ import type { TooltipRootProps, TooltipRootEmits, TooltipContentProps } from 'ra
 import theme from '#build/ui/tooltip'
 
 const tooltip = tv(theme)
-// const appTooltip = tv({ extend: tooltip, ...(appConfig.ui?.tooltip || {}) })
+// const tooltip = tv({ extend: tv(theme), ...(appConfig.ui?.tooltip || {}) })
 
 export interface TooltipProps extends TooltipRootProps, Omit<TooltipContentProps, 'as' | 'asChild'> {
   text?: string
diff --git a/src/runtime/types/index.d.ts b/src/runtime/types/index.d.ts
index 129a7140..e69de29b 100644
--- a/src/runtime/types/index.d.ts
+++ b/src/runtime/types/index.d.ts
@@ -1,3 +0,0 @@
-export type DeepPartial<T> = Partial<{
-  [P in keyof T]: DeepPartial<T[P]> | { [key: string]: string | object }
-}>
\ No newline at end of file
diff --git a/src/templates.ts b/src/templates.ts
index 26de0e92..15394c0b 100644
--- a/src/templates.ts
+++ b/src/templates.ts
@@ -7,10 +7,9 @@ export default function createTemplates (options: ModuleOptions, nuxt: Nuxt) {
   const shades = [50, 100, 200, 300, 400, 500, 600, 700, 800, 900, 950]
 
   const template = addTemplate({
-    filename: 'ui.css',
+    filename: 'ui/index.css',
     write: true,
-    getContents: () => `
-      @import "tailwindcss";
+    getContents: () => `@import "tailwindcss";
 
       @theme {
         --color-gray-*: initial;
@@ -40,14 +39,39 @@ export default function createTemplates (options: ModuleOptions, nuxt: Nuxt) {
       write: true,
       getContents: () => `export default ${JSON.stringify((theme as any)[component]({ colors: options.colors }), null, 2)}`
     })
-
-    // addTypeTemplate({
-    //   filename: `ui/${component}.d.ts`,
-    //   write: true,
-    //   getContents: () => `import ${component} from './${component}'
-
-    //   type T = typeof ${component}
-    //   export default T`
-    // })
   }
+
+  addTemplate({
+    filename: 'ui/index.ts',
+    write: true,
+    getContents: () => Object.keys(theme).map(component => `export { default as ${component} } from './${component}'`).join('\n')
+  })
+
+  addTypeTemplate({
+    filename: 'types/ui.d.ts',
+    getContents: () => `
+      import * as ui from '#build/ui'
+
+      type DeepPartial<T> = Partial<{
+        [P in keyof T]: DeepPartial<T[P]> | { [key: string]: string | object }
+      }>
+
+      type UI = {
+        primary?: string
+        gray?: string
+        [key: string]: any
+      } & DeepPartial<typeof ui>
+
+      declare module 'nuxt/schema' {
+        interface AppConfigInput {
+          ui?: UI
+        }
+      }
+      declare module '@nuxt/schema' {
+        interface AppConfigInput {
+          ui?: UI
+        }
+      }
+      export {}`
+  })
 }
diff --git a/src/theme/button.ts b/src/theme/button.ts
index 1fab8f42..158267f2 100644
--- a/src/theme/button.ts
+++ b/src/theme/button.ts
@@ -5,7 +5,12 @@ export default (config: any) => ({
     icon: 'shrink-0'
   },
   variants: {
-    color: Object.fromEntries(config.colors.map((color: string) => [color, ''])),
+    color: {
+      ...Object.fromEntries(config.colors.map((color: string) => [color, ''])),
+      white: '',
+      gray: '',
+      black: ''
+    },
     variant: {
       solid: '',
       outline: '',