diff --git a/src/app/api/flux/route.ts b/src/app/api/flux/route.ts
index f0b1e67..8749c4f 100644
--- a/src/app/api/flux/route.ts
+++ b/src/app/api/flux/route.ts
@@ -1,7 +1,8 @@
 import { NextRequest, NextResponse } from 'next/server'
 import { fal } from "@fal-ai/client"
+import { FLUX_MODELS } from '../../../services/fluxService'
 
-import { verifyPayment } from '../../../services/paymentService'
+import { verifyPayment } from '../../../utils/verifyPayment'
 
 if (!process.env.FAL_AI_KEY) {
   throw new Error('FAL_AI_KEY is not configured in environment variables')
@@ -18,7 +19,7 @@ const IMAGE_HEIGHT: number = 1024
 
 export async function POST(req: NextRequest): Promise<NextResponse> {
   try {
-    const { prompt, modelId, transactionSignature, expectedAmount } = await req.json()
+    const { prompt, modelId, transactionSignature } = await req.json()
 
     if (!prompt || !modelId) {
       return NextResponse.json(
@@ -27,15 +28,24 @@ export async function POST(req: NextRequest): Promise<NextResponse> {
       )
     }
 
-    if (!transactionSignature || !expectedAmount) {
+    if (!transactionSignature) {
       return NextResponse.json(
-      { error: 'Transaction signature and expected amount are required' },
+      { error: 'Transaction signature is required' },
       { status: 400 }
       )
     }
 
-    const isPaymentVerified = await verifyPayment(transactionSignature, expectedAmount);
-    console.log({isPaymentVerified});
+    const model = FLUX_MODELS.find((model) => model.modelId === modelId)
+    if (!model) {
+      return NextResponse.json(
+        { error: 'Invalid modelId' },
+        { status: 400 }
+      )
+    }
+
+    const expectedAmount = model.cost;
+    const isPaymentVerified = await verifyPayment(transactionSignature, expectedAmount)
+
     if (!isPaymentVerified) {
       throw new Error('Payment verification failed');
     }
diff --git a/src/app/page.tsx b/src/app/page.tsx
index 75c527f..287c7a2 100644
--- a/src/app/page.tsx
+++ b/src/app/page.tsx
@@ -49,10 +49,10 @@ const Page: React.FC = (): React.ReactElement => {
         return { error: paymentResult.error }
       }
 
-      const transactionSignature = paymentResult.signature;
+      const transactionSignature = paymentResult.transactionSignature;
 
       // Then generate image with specified model
-      return generateWithFlux(prompt, modelId, transactionSignature, cost)
+      return generateWithFlux(prompt, modelId, transactionSignature)
     }
   }
 
diff --git a/src/services/fluxService.ts b/src/services/fluxService.ts
index d20dfdf..7f22c2e 100644
--- a/src/services/fluxService.ts
+++ b/src/services/fluxService.ts
@@ -36,7 +36,6 @@ export async function generateWithFlux(
   prompt: string,
   modelId: string,
   transactionSignature: string,
-  expectedAmount: number,
 ): Promise<FluxGenerationResult> {
   try {
     const response = await fetch('/api/flux', {
@@ -48,7 +47,6 @@ export async function generateWithFlux(
         prompt,
         modelId,
         transactionSignature,
-        expectedAmount
       }),
     })
 
diff --git a/src/services/paymentService.ts b/src/services/paymentService.ts
index 9cd0f17..fd518e0 100644
--- a/src/services/paymentService.ts
+++ b/src/services/paymentService.ts
@@ -23,6 +23,7 @@ const connection = new Connection(
 
 export interface PaymentResult {
   success: boolean
+  transactionSignature?: string,
   error?: string
 }
 
@@ -48,7 +49,7 @@ export async function processMTMPayment(
   walletPublicKey: string,
   tokenAmount: number,
   walletType: WalletType
-): Promise<PaymentResult & { signature?: string }> {
+): Promise<PaymentResult> {
   try {
     let wallet: WalletAdapter | null = null;
 
@@ -146,7 +147,7 @@ export async function processMTMPayment(
       throw new Error(`Transaction failed: ${JSON.stringify(confirmation.value.err)}`)
     }
 
-    return { success: true, signature };
+    return { success: true, transactionSignature: signature };
   } catch (error) {
     console.error('Payment error:', error)
     return {
@@ -155,38 +156,3 @@ export async function processMTMPayment(
     }
   }
 }
-
-export async function verifyPayment(
-  transactionSignature: string,
-  expectedAmount: number,
-): Promise<boolean> {
-  try {
-    console.log('Verifying payment:');
-
-    const transaction = await connection.getParsedTransaction(transactionSignature, 'confirmed');
-    if (!transaction) {
-      throw new Error('Transaction not found');
-    }
-
-    const transferInstruction = transaction.transaction.message.instructions.find(
-      (instr) => 'parsed' in instr && instr.programId.equals(TOKEN_PROGRAM_ID)
-    );
-
-    if (!transferInstruction || !('parsed' in transferInstruction)) {
-      throw new Error('Transfer instruction not found');
-    }
-
-    const { parsed } = transferInstruction;
-    const { info } = parsed;
-    const { amount } = info;
-
-    if (BigInt(amount) === BigInt(expectedAmount * (10 ** 6))) {
-      return true;
-    }
-
-    return false;
-  } catch (error) {
-    console.error('Verification error:', error);
-    return false;
-  }
-}
diff --git a/src/utils/verifyPayment.ts b/src/utils/verifyPayment.ts
new file mode 100644
index 0000000..40cd0c9
--- /dev/null
+++ b/src/utils/verifyPayment.ts
@@ -0,0 +1,47 @@
+import { Connection } from '@solana/web3.js';
+import { TOKEN_PROGRAM_ID } from '@solana/spl-token';
+
+const SOLANA_RPC_URL = process.env.NEXT_PUBLIC_SOLANA_RPC_URL || 'https://young-radial-orb.solana-mainnet.quiknode.pro/67612b364664616c29514e551bf5de38447ca3d4';
+const SOLANA_WEBSOCKET_URL = process.env.NEXT_PUBLIC_SOLANA_WEBSOCKET_URL || 'wss://young-radial-orb.solana-mainnet.quiknode.pro/67612b364664616c29514e551bf5de38447ca3d4';
+
+const connection = new Connection(
+  SOLANA_RPC_URL,
+  {
+    commitment: 'confirmed',
+    wsEndpoint: SOLANA_WEBSOCKET_URL,
+    confirmTransactionInitialTimeout: 60000,
+  }
+);
+
+export async function verifyPayment(
+  transactionSignature: string,
+  expectedAmount: number,
+): Promise<boolean> {
+  try {
+    const transaction = await connection.getParsedTransaction(transactionSignature, 'finalized');
+    if (!transaction) {
+      throw new Error('Transaction not found');
+    }
+
+    const transferInstruction = transaction.transaction.message.instructions.find(
+      (instr) => 'parsed' in instr && instr.programId.equals(TOKEN_PROGRAM_ID)
+    );
+
+    if (!transferInstruction || !('parsed' in transferInstruction)) {
+      throw new Error('Transfer instruction not found');
+    }
+
+    const { parsed } = transferInstruction;
+    const { info } = parsed;
+    const { amount } = info;
+
+    if (BigInt(amount) === BigInt(expectedAmount * (10 ** 6))) {
+      return true;
+    }
+
+    return false;
+  } catch (error) {
+    console.error('Verification error:', error);
+    return false;
+  }
+}