File size: 4,114 Bytes
dbc2c2a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
88b5cce
dbc2c2a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
88b5cce
 
dbc2c2a
 
88b5cce
dbc2c2a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
/**

 * API ROUTE: /api/improve-prompt

 * 

 * Uses Gemini 2.5 Flash to improve user prompts for better AI image generation.

 * Takes a basic prompt and enhances it with more detailed, descriptive language

 * that will produce better results from image generation models.

 */

import { NextRequest, NextResponse } from "next/server";
import { GoogleGenAI } from "@google/genai";
import { cookies } from "next/headers";

export const runtime = "nodejs";
export const maxDuration = 30;

export async function POST(req: NextRequest) {
  try {
    const body = await req.json() as {
      prompt: string;
      type?: string; // 'background', 'edit', etc.
      apiToken?: string; // User's Google AI API token
    };

    if (!body.prompt?.trim()) {
      return NextResponse.json(
        { error: "Prompt is required" },
        { status: 400 }
      );
    }

    // Check if user is logged in with HF Pro
    let isHfProUser = false;
    try {
      const cookieStore = await cookies();
      const hfToken = cookieStore.get('hf_token');
      isHfProUser = !!hfToken?.value;
    } catch (error) {
      console.error('Error reading HF token from cookies:', error);
    }

    // Validate and retrieve Google API key from user input or environment
    const apiKey = body.apiToken || process.env.GOOGLE_API_KEY;
    if (!apiKey || apiKey === 'your_actual_api_key_here') {
      return NextResponse.json(
        { error: `API key not provided. Please ${isHfProUser ? 'enter your Google Gemini API token in the top right' : 'login with HF Pro or enter your Google Gemini API token'}.` },
        { status: 500 }
      );
    }

    const ai = new GoogleGenAI({ apiKey });

    // Create context-specific improvement prompts
    const contextPrompts = {
      background: `You are an expert at writing prompts for AI image generation. Take the following simple background description and transform it into a detailed, vivid prompt that will generate stunning, realistic backgrounds.



Focus on:

- Visual details (lighting, colors, textures, atmosphere)

- Composition and depth

- Realistic environmental elements

- Photography/cinematic quality terms

- Maintaining the character while enhancing the background



Keep the character image and background realistic. Make the description rich and specific but not overly complex.



Original prompt: "${body.prompt}"



Write a short and concise improved background generation prompt and do not include anything unnecessary:`,

      edit: `You are an expert at writing prompts for AI image editing. Take the following simple editing request and transform it into a clear, detailed prompt that will produce precise, high-quality image modifications.

Original prompt: "${body.prompt}" Return a short and concise improved editing prompt and do not include anything unnecessary:`,

      default: `You are an expert at writing prompts for AI image generation and editing. Take the following simple prompt and transform it into a detailed, effective prompt that will produce better results.



Focus on:

- Clear, specific instructions

- Visual details and quality descriptors

- Professional terminology

- Realistic and natural-looking results



Original prompt: "${body.prompt}"



Write an improved prompt:`
    };

    const improvementPrompt = contextPrompts[body.type as keyof typeof contextPrompts] || contextPrompts.default;

    const response = await ai.models.generateContent({
      model: "gemini-2.5-flash",
      contents: [{ role: "user", parts: [{ text: improvementPrompt }] }],
    });

    const improvedPrompt = response?.text?.trim();

    if (!improvedPrompt) {
      return NextResponse.json(
        { error: "Failed to generate improved prompt" },
        { status: 500 }
      );
    }

    return NextResponse.json({ improvedPrompt });

  } catch (error: any) {
    console.error('[API] improve-prompt error:', error);
    
    return NextResponse.json(
      { error: "Failed to improve prompt. Please try again." },
      { status: 500 }
    );
  }
}