diff --git a/src/ColorFunctionContext.tsx b/src/ColorFunctionContext.tsx
index eba7730..d76e966 100644
--- a/src/ColorFunctionContext.tsx
+++ b/src/ColorFunctionContext.tsx
@@ -1,20 +1,13 @@
-import {
-  createContext,
-  ReactNode,
-  useCallback,
-  useEffect,
-  useMemo,
-  useState,
-  useTransition,
-} from 'react';
-import type * as Wasm from './pkg/color_module';
-import init from './pkg/color_module';
+import { createContext, ReactNode, useEffect, useMemo, useState, useTransition } from 'react';
+import { ColorFunctionContext } from './ColorFunctionContext';
+import init, * as funcs from './pkg/color_module';
 
-export const ColorFunctionContext = createContext<{
-  colorFn: Wasm.InitOutput | null;
+export type ColorFunctionContextType = {
+  colorFn: typeof funcs | null;
   isLoading: boolean;
   error: Error | null;
-}>({
+};
+export const ColorFunctionContext = createContext<ColorFunctionContextType>({
   colorFn: null,
   isLoading: true,
   error: null,
@@ -29,17 +22,6 @@ export function ColorFunctionProvider({ children }: WasmProviderProps) {
   const [isPending, startTransition] = useTransition();
   const [error, setError] = useState<Error | null>(null);
 
-  const loadWasm = useCallback(() => {
-    startTransition(async () => {
-      try {
-        const wasm = await init();
-        setWasmInstance(wasm);
-      } catch (e) {
-        console.error('[Load WASM]', e);
-        setError(e);
-      }
-    });
-  }, []);
   const contextValue = useMemo(
     () => ({
       colorFn: wasmInstance,
@@ -49,7 +31,17 @@ export function ColorFunctionProvider({ children }: WasmProviderProps) {
     [wasmInstance, isPending, error],
   );
   useEffect(() => {
-    loadWasm();
+    (function () {
+      startTransition(async () => {
+        try {
+          await init();
+          setWasmInstance(funcs);
+        } catch (e) {
+          console.error('[Load WASM]', e);
+          setError(e);
+        }
+      });
+    })();
   }, []);
 
   return <ColorFunctionContext value={contextValue}>{children}</ColorFunctionContext>;