diff --git a/Makefile b/Makefile
index e15c6e8..221c472 100644
--- a/Makefile
+++ b/Makefile
@@ -76,13 +76,13 @@ build-all: go.sum
 	rm -f $(CHECKSUM_FILE)
 	GOOS=linux GOARCH=amd64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-linux-amd64 $(SOURCE_DIR)
 	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-linux-amd64 >> $(CHECKSUM_FILE)
-	GOOS=linux GOARCH=arm64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 $(SOURCE_DIR)
-	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 >> $(CHECKSUM_FILE)
-	GOOS=darwin GOARCH=amd64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-amd64 $(SOURCE_DIR)
-	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-darwin-amd64 >> $(CHECKSUM_FILE)
-	GOOS=darwin GOARCH=arm64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 $(SOURCE_DIR)
-	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 >> $(CHECKSUM_FILE)
-	GOOS=windows GOARCH=amd64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe $(SOURCE_DIR)
-	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe >> $(CHECKSUM_FILE)
+#	GOOS=linux GOARCH=arm64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 $(SOURCE_DIR)
+#	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-linux-arm64 >> $(CHECKSUM_FILE)
+#	GOOS=darwin GOARCH=amd64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-amd64 $(SOURCE_DIR)
+#	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-darwin-amd64 >> $(CHECKSUM_FILE)
+#	GOOS=darwin GOARCH=arm64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 $(SOURCE_DIR)
+#	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-darwin-arm64 >> $(CHECKSUM_FILE)
+#	GOOS=windows GOARCH=amd64 $(GO_BUILD) -tags "$(build_tags)" -ldflags '$(ldflags)' -o $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe $(SOURCE_DIR)
+#	@shasum -a 256 $(BUILD_DIR)/$(BINARY_NAME)-windows-amd64.exe >> $(CHECKSUM_FILE)
 
 .PHONY: default build install test clean lint print-system build-all
diff --git a/app/app.go b/app/app.go
index a68ecd9..7f4add1 100644
--- a/app/app.go
+++ b/app/app.go
@@ -1,8 +1,11 @@
 package app
 
 import (
+	"context"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
+	trackgate "github.com/airchains-network/junction/x/trackgate/module"
 	trackgatemoduletypes "github.com/airchains-network/junction/x/trackgate/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"io"
 	"os"
 	"path/filepath"
@@ -325,6 +328,87 @@ func New(
 
 	app.App = appBuilder.Build(db, traceStore, baseAppOptions...)
 
+	upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
+	if err != nil {
+		panic(err)
+	}
+
+	if upgradeInfo.Name == "jip-2" && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
+		storeUpgrades := storetypes.StoreUpgrades{
+			Added: []string{trackgatemoduletypes.StoreKey},
+		}
+
+		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
+		app.UpgradeKeeper.SetUpgradeHandler(
+			"jip-2",
+			func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+				sdkCtx := sdk.UnwrapSDKContext(ctx)
+				//storeUpgrades := storetypes.StoreUpgrades{
+				//	Added: []string{trackgateTypes.StoreKey},
+				//}
+				//
+				//app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(plan.Height, &storeUpgrades))
+
+				// Skip the capability module migration to avoid setting the index again
+				// You can use the module manager to skip migrations for the capability module by adjusting the version map
+
+				// Check if the capability index is already set before attempting to initialize it
+				latestIndex := app.CapabilityKeeper.GetLatestIndex(sdkCtx)
+				if latestIndex == 0 {
+					// The index is not set, so we can safely initialize it
+					err := app.CapabilityKeeper.InitializeIndex(sdkCtx, 1) // Initialize with index 1 or a value > 0
+					if err != nil {
+						return nil, err
+					}
+				} else {
+					logger.Info("Capability index already initialized, skipping re-initialization")
+				}
+				//configurator := app.Configurator()
+				//versionMap, err := app.ModuleManager.RunMigrations(sdkCtx, configurator, fromVM)
+				//if err != nil {
+				//	return nil, err
+				//}
+				//// Convert the VersionMap to a string for logging
+				//versionMapString := fmt.Sprintf("%v", versionMap)
+				//logger.Info(versionMapString)
+				//
+				//// Ensure the capability module is not migrated again
+				//if version, exists := versionMap["capability"]; exists && version >= 1 {
+				//	logger.Info("Skipping capability module migration")
+				//}
+				versionMap := module.VersionMap{
+					"trackgate": 1,
+				}
+
+				authority := authtypes.NewModuleAddress(govtypes.ModuleName)
+
+				// Create the Trackgate Keeper
+				app.TrackgateKeeper = trackgatemodulekeeper.NewKeeper(
+					app.AppCodec(),
+					runtime.NewKVStoreService(app.GetKey(trackgatemoduletypes.StoreKey)),
+					logger,
+					authority.String(),
+					app.BankKeeper,
+				)
+
+				// Create the Trackgate AppModule
+				trackgateModule := trackgate.NewAppModule(
+					app.AppCodec(),
+					app.TrackgateKeeper,
+					app.AccountKeeper,
+					app.BankKeeper,
+				)
+
+				// Register the Trackgate module using app.RegisterModules
+				err = app.RegisterModules(trackgateModule)
+				if err != nil {
+					return nil, err
+				}
+
+				return versionMap, nil
+			}, // Upgrade handler function
+		)
+	}
 	// Register legacy modules
 	app.registerIBCModules()
 
@@ -356,31 +440,11 @@ func New(
 	// 	app.UpgradeKeeper.SetModuleVersionMap(ctx, app.ModuleManager.GetVersionMap())
 	// 	return app.App.InitChainer(ctx, req)
 	// })
-	// Handle store upgrades
-	upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
-	if err != nil {
-		panic(err)
-	}
-
-	if upgradeInfo.Name == "jip-2" && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
-		storeUpgrades := storetypes.StoreUpgrades{
-			Added: []string{trackgatemoduletypes.StoreKey},
-		}
-
-		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
-	}
 
 	if err := app.Load(loadLatest); err != nil {
 		return nil, err
 	}
 
-	configurator := app.Configurator()
-	UpgradeHandleFunc := CreateDefaultUpgradeHandler(app.ModuleManager, configurator, app)
-	app.UpgradeKeeper.SetUpgradeHandler(
-		"jip-2",
-		UpgradeHandleFunc, // Upgrade handler function
-	)
-
 	return app, nil
 }
 
diff --git a/app/upgrade.go b/app/upgrade.go
index 7082480..2851572 100644
--- a/app/upgrade.go
+++ b/app/upgrade.go
@@ -2,27 +2,79 @@ package app
 
 import (
 	"context"
-	storetypes "cosmossdk.io/store/types"
-	trackgatemoduletypes "github.com/airchains-network/junction/x/trackgate/types"
-
+	"cosmossdk.io/log"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
+	trackgateKeeper "github.com/airchains-network/junction/x/trackgate/keeper"
+	trackgate "github.com/airchains-network/junction/x/trackgate/module"
+	trackgateTypes "github.com/airchains-network/junction/x/trackgate/types"
+	"github.com/cosmos/cosmos-sdk/runtime"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
+	"strconv"
 )
 
-// for regular normal upgrades
-func CreateDefaultUpgradeHandler(
+func CreateUpgradeHandler(
 	mm *module.Manager,
 	configurator module.Configurator,
 	app *App,
+	logger log.Logger,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
-		storeUpgrades := storetypes.StoreUpgrades{
-			Added: []string{trackgatemoduletypes.StoreKey},
+		sdkCtx := sdk.UnwrapSDKContext(ctx)
+		//storeUpgrades := storetypes.StoreUpgrades{
+		//	Added: []string{trackgateTypes.StoreKey},
+		//}
+		//
+		//app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(plan.Height, &storeUpgrades))
+
+		// Check if the capability index is already set before attempting to initialize it
+		latestIndex := app.CapabilityKeeper.GetLatestIndex(sdkCtx)
+		indexString := strconv.FormatUint(latestIndex, 10)
+		logger.Debug(indexString)
+
+		// Run migrations for all modules
+		versionMap, err := mm.RunMigrations(sdkCtx, configurator, fromVM)
+		if err != nil {
+			return nil, err
+		}
+
+		if latestIndex == 0 {
+			// The index is not set, so we can safely initialize it
+			err := app.CapabilityKeeper.InitializeIndex(sdkCtx, 1) // Initialize with index 1 or a value > 0
+			if err != nil {
+				return nil, err
+			}
+		} else {
+			logger.Info("Capability index already initialized, skipping re-initialization")
 		}
 
-		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(plan.Height, &storeUpgrades))
+		authority := authtypes.NewModuleAddress(govtypes.ModuleName)
+
+		// Create the Trackgate Keeper
+		app.TrackgateKeeper = trackgateKeeper.NewKeeper(
+			app.AppCodec(),
+			runtime.NewKVStoreService(app.GetKey(trackgateTypes.StoreKey)),
+			logger,
+			authority.String(),
+			app.BankKeeper,
+		)
+
+		// Create the Trackgate AppModule
+		trackgateModule := trackgate.NewAppModule(
+			app.AppCodec(),
+			app.TrackgateKeeper,
+			app.AccountKeeper,
+			app.BankKeeper,
+		)
+
+		// Register the Trackgate module using app.RegisterModules
+		err = app.RegisterModules(trackgateModule)
+		if err != nil {
+			return nil, err
+		}
 
-		//return versionMap, nil
-		return app.ModuleManager.RunMigrations(ctx, configurator, fromVM)
+		return versionMap, nil
 	}
 }
diff --git a/x/trackgate/module/module.go b/x/trackgate/module/module.go
index bb7dda3..76d000c 100644
--- a/x/trackgate/module/module.go
+++ b/x/trackgate/module/module.go
@@ -32,7 +32,7 @@ var (
 	_ module.HasInvariants       = (*AppModule)(nil)
 	_ module.HasConsensusVersion = (*AppModule)(nil)
 
-	_ appmodule.AppModule       = (*AppModule)(nil)
+	_ module.AppModule          = (*AppModule)(nil)
 	_ appmodule.HasBeginBlocker = (*AppModule)(nil)
 	_ appmodule.HasEndBlocker   = (*AppModule)(nil)
 )
@@ -114,6 +114,11 @@ func NewAppModule(
 	}
 }
 
+// Name returns the name of the module as a string.
+func (AppModule) Name() string {
+	return types.ModuleName
+}
+
 // RegisterServices registers a gRPC query service to respond to the module-specific gRPC queries
 func (am AppModule) RegisterServices(cfg module.Configurator) {
 	types.RegisterMsgServer(cfg.MsgServer(), keeper.NewMsgServerImpl(am.keeper))
@@ -188,7 +193,7 @@ type ModuleOutputs struct {
 	depinject.Out
 
 	TrackgateKeeper keeper.Keeper
-	Module          appmodule.AppModule
+	Module          module.AppModule
 }
 
 func ProvideModule(in ModuleInputs) ModuleOutputs {