diff --git a/src/templates.ts b/src/templates.ts
index cf247f98b0..975c0fc0a4 100644
--- a/src/templates.ts
+++ b/src/templates.ts
@@ -29,24 +29,37 @@ export function getTemplates(options: ModuleOptions, uiConfig: Record<string, an
         let json = JSON.stringify(result, null, 2)
 
         for (const variant of variants) {
-          json = json.replace(new RegExp(`("${variant}": "[^"]+")`, 'g'), '$1 as const')
+          json = json.replace(new RegExp(`("${variant}": "[^"]+")`, 'g'), `$1 as typeof ${variant}[number]`)
           json = json.replace(new RegExp(`("${variant}": \\[\\s*)((?:"[^"]+",?\\s*)+)(\\])`, 'g'), (_, before, match, after) => {
-            const replaced = match.replace(/("[^"]+")/g, '$1 as const')
+            const replaced = match.replace(/("[^"]+")/g, `$1 as typeof ${variant}[number]`)
             return `${before}${replaced}${after}`
           })
         }
 
-        // For local development, directly import from theme
+        function generateVariantDeclarations(variants: string[]) {
+          return variants.map((variant) => {
+            const keys = Object.keys(result.variants[variant])
+            return `const ${variant} = ${JSON.stringify(keys, null, 2)} as const`
+          })
+        }
+
+        // For local development, import directly from theme
         if (process.env.DEV) {
+          const templatePath = fileURLToPath(new URL(`./theme/${kebabCase(component)}`, import.meta.url))
           return [
-            `import template from ${JSON.stringify(fileURLToPath(new URL(`./theme/${kebabCase(component)}`, import.meta.url)))}`,
-            `const result = typeof template === 'function' ? template(${JSON.stringify(options)}) : template`,
-            `const json = ${json}`,
-            `export default result as typeof json`
-          ].join('\n')
+            `import template from ${JSON.stringify(templatePath)}`,
+            ...generateVariantDeclarations(variants),
+            `const result = typeof template === 'function' ? template(${JSON.stringify(options, null, 2)}) : template`,
+            `const theme = ${json}`,
+            `export default result as typeof theme`
+          ].join('\n\n')
         }
 
-        return `export default ${json}`
+        // For production build
+        return [
+          ...generateVariantDeclarations(variants),
+          `export default ${json}`
+        ].join('\n\n')
       }
     })
   }