File size: 4,165 Bytes
27127dd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import { z } from "zod";

// Define schema for image generation request
export const imageGenerationSchema = z.object({
  prompt: z.string().min(1).max(1000),
  seed: z.number().optional().default(0),
  randomize_seed: z.boolean().optional().default(true),
  width: z.number().min(256).max(1024).optional().default(512),
  height: z.number().min(256).max(1024).optional().default(512),
  guidance_scale: z.number().min(0).max(20).optional().default(7.5),
  num_inference_steps: z.number().min(1).max(50).optional().default(20),
});

// Type for image generation parameters
export type ImageGenerationParams = z.infer<typeof imageGenerationSchema>;

/**
 * Generates an image using the FLUX.1-dev model via Replicate API
 * @param params Image generation parameters
 * @returns URL of the generated image
 */
export async function generateImage(params: ImageGenerationParams): Promise<string> {
  try {
    const apiKey = process.env.REPLICATE_API_KEY;
    if (!apiKey) {
      throw new Error("REPLICATE_API_KEY is not set in environment variables");
    }

    // Prepare the request data
    const inputData = {
      input: {
        prompt: params.prompt,
        width: params.width,
        height: params.height,
        seed: params.randomize_seed ? Math.floor(Math.random() * 1000000) : params.seed,
        guidance_scale: params.guidance_scale,
        num_inference_steps: params.num_inference_steps,
      },
    };

    // Make API request to create prediction
    const startResponse = await fetch(
      "https://api.replicate.com/v1/models/black-forest-labs/flux-dev/predictions",
      {
        method: "POST",
        headers: {
          "Authorization": `Bearer ${apiKey}`,
          "Content-Type": "application/json",
        },
        body: JSON.stringify(inputData),
      }
    );

    if (!startResponse.ok) {
      const errorData = await startResponse.json();
      throw new Error(`Replicate API error: ${JSON.stringify(errorData)}`);
    }

    const prediction = await startResponse.json();
    const predictionId = prediction.id;

    // Poll for result
    let imageUrl: string | null = null;
    let attempts = 0;
    const maxAttempts = 30;

    while (!imageUrl && attempts < maxAttempts) {
      await new Promise(resolve => setTimeout(resolve, 1000)); // Wait 1 second
      
      const statusResponse = await fetch(
        `https://api.replicate.com/v1/predictions/${predictionId}`,
        {
          headers: {
            "Authorization": `Bearer ${apiKey}`,
          },
        }
      );

      if (!statusResponse.ok) {
        const errorData = await statusResponse.json();
        throw new Error(`Replicate API status error: ${JSON.stringify(errorData)}`);
      }

      const status = await statusResponse.json();
      
      if (status.status === "succeeded") {
        if (status.output && typeof status.output === "string") {
          imageUrl = status.output;
        } else if (Array.isArray(status.output) && status.output.length > 0) {
          imageUrl = status.output[0];
        }
      } else if (status.status === "failed") {
        throw new Error(`Image generation failed: ${status.error || "Unknown error"}`);
      }
      
      attempts++;
    }

    if (!imageUrl) {
      throw new Error("Timed out waiting for image generation");
    }

    return imageUrl;
  } catch (error) {
    console.error("Error generating image:", error);
    throw new Error(`Failed to generate image: ${error instanceof Error ? error.message : String(error)}`);
  }
}

// Function to check if FLUX model is available via Replicate
export async function isFluxAvailable(): Promise<boolean> {
  try {
    const apiKey = process.env.REPLICATE_API_KEY;
    if (!apiKey) {
      return false;
    }

    // Check if the Replicate API is accessible by making a basic request
    const response = await fetch(
      "https://api.replicate.com/v1/models/black-forest-labs/flux-dev",
      {
        headers: {
          "Authorization": `Bearer ${apiKey}`,
        },
      }
    );

    return response.ok;
  } catch (error) {
    console.error("Error checking FLUX availability:", error);
    return false;
  }
}