diff --git a/server/ai_http.go b/server/ai_http.go
index 34c359bc31..0ed9320ad1 100644
--- a/server/ai_http.go
+++ b/server/ai_http.go
@@ -302,7 +302,13 @@ func handleAIRequest(ctx context.Context, w http.ResponseWriter, r *http.Request
 			return orch.LlmGenerate(ctx, v)
 		}
 
-		// TODO: handle tokens for pricing
+		if v.MaxTokens == nil {
+			respondWithError(w, "MaxTokens not specified", http.StatusBadRequest)
+			return
+		}
+
+		// TODO: Improve pricing
+		outPixels = int64(*v.MaxTokens)
 	default:
 		respondWithError(w, "Unknown request type", http.StatusBadRequest)
 		return
diff --git a/server/ai_process.go b/server/ai_process.go
index 27d09a68da..ea145df4f0 100644
--- a/server/ai_process.go
+++ b/server/ai_process.go
@@ -728,8 +728,12 @@ func submitLlmGenerate(ctx context.Context, params aiRequestParams, sess *AISess
 		return nil, err
 	}
 
-	// TODO: calculate payment
-	setHeaders, balUpdate, err := prepareAIPayment(ctx, sess, 0)
+	// TODO: Improve pricing
+	if req.MaxTokens == nil {
+		req.MaxTokens = new(int)
+		*req.MaxTokens = 256
+	}
+	setHeaders, balUpdate, err := prepareAIPayment(ctx, sess, int64(*req.MaxTokens))
 	if err != nil {
 		if monitor.Enabled {
 			monitor.AIRequestError(err.Error(), "llm-generate", *req.ModelId, sess.OrchestratorInfo)