| 1 | "use server"; |
| 2 | |
| 3 | import { utapi } from "@/app/api/uploadthing/core"; |
| 4 | import { env } from "@/env"; |
| 5 | import { requireOptionalIntegration } from "@/lib/env/optional-integrations"; |
| 6 | import { auth } from "@/server/auth"; |
| 7 | import { db } from "@/server/db"; |
| 8 | import Together from "together-ai"; |
| 9 | import { UTFile } from "uploadthing/server"; |
| 10 | |
| 11 | export type ImageModelList = |
| 12 | | "black-forest-labs/FLUX1.1-pro" |
| 13 | | "black-forest-labs/FLUX.1-schnell" |
| 14 | | "black-forest-labs/FLUX.1-schnell-Free" |
| 15 | | "black-forest-labs/FLUX.1-pro" |
| 16 | | "black-forest-labs/FLUX.1-dev"; |
| 17 | |
| 18 | export async function generateImageAction( |
| 19 | prompt: string, |
| 20 | model: ImageModelList = "black-forest-labs/FLUX.1-schnell-Free", |
| 21 | ) { |
| 22 | // Get the current session |
| 23 | const session = await auth(); |
| 24 | |
| 25 | // Check if user is authenticated |
| 26 | if (!session?.user?.id) { |
| 27 | throw new Error("You must be logged in to generate images"); |
| 28 | } |
| 29 | |
| 30 | try { |
| 31 | const togetherConfig = requireOptionalIntegration({ |
| 32 | integration: "Together AI", |
| 33 | envVar: "TOGETHER_AI_API_KEY", |
| 34 | value: env.TOGETHER_AI_API_KEY, |
| 35 | feature: "AI image generation", |
| 36 | }); |
| 37 | |
| 38 | if (!togetherConfig.ok) { |
| 39 | return { |
| 40 | success: false, |
| 41 | error: togetherConfig.error, |
| 42 | }; |
| 43 | } |
| 44 | |
| 45 | const together = new Together({ apiKey: togetherConfig.value }); |
| 46 | |
| 47 | console.log(`Generating image with model: ${model}`); |
| 48 | |
| 49 | // Generate the image using Together AI |
| 50 | const response = (await together.images.create({ |
| 51 | model: model, |
| 52 | prompt: prompt, |
| 53 | width: 1024, |
| 54 | height: 768, |
| 55 | steps: model.includes("schnell") ? 4 : 28, // Fewer steps for schnell models |
| 56 | n: 1, |
| 57 | })) as unknown as { |
| 58 | id: string; |
| 59 | model: string; |
| 60 | object: string; |
| 61 | data: { |
| 62 | url: string; |
| 63 | }[]; |
| 64 | }; |
| 65 | |
| 66 | const imageUrl = response.data[0]?.url; |
| 67 | |
| 68 | if (!imageUrl) { |
| 69 | throw new Error("Failed to generate image"); |
| 70 | } |
| 71 | |
| 72 | console.log(`Generated image URL: ${imageUrl}`); |
| 73 | |
| 74 | // Download the image from Together AI URL |
| 75 | const imageResponse = await fetch(imageUrl); |
| 76 | if (!imageResponse.ok) { |
| 77 | throw new Error("Failed to download image from Together AI"); |
| 78 | } |
| 79 | |
| 80 | const imageBlob = await imageResponse.blob(); |
| 81 | const imageBuffer = await imageBlob.arrayBuffer(); |
| 82 | |
| 83 | // Generate a filename based on the prompt |
| 84 | const filename = `${prompt.substring(0, 20).replace(/[^a-z0-9]/gi, "_")}_${Date.now()}.png`; |
| 85 | |
| 86 | // Create a UTFile from the downloaded image |
| 87 | const utFile = new UTFile([new Uint8Array(imageBuffer)], filename); |
| 88 | |
| 89 | // Upload to UploadThing |
| 90 | const uploadResult = await utapi.uploadFiles([utFile]); |
| 91 | |
| 92 | if (!uploadResult[0]?.data?.ufsUrl) { |
| 93 | console.error("Upload error:", uploadResult[0]?.error); |
| 94 | throw new Error("Failed to upload image to UploadThing"); |
| 95 | } |
| 96 | |
| 97 | console.log(uploadResult); |
| 98 | const permanentUrl = uploadResult[0].data.ufsUrl; |
| 99 | console.log(`Uploaded to UploadThing URL: ${permanentUrl}`); |
| 100 | |
| 101 | // Store in database with the permanent URL |
| 102 | const generatedImage = await db.generatedImage.create({ |
| 103 | data: { |
| 104 | url: permanentUrl, // Store the UploadThing URL instead of the Together AI URL |
| 105 | prompt: prompt, |
| 106 | userId: session.user.id, |
| 107 | }, |
| 108 | }); |
| 109 | |
| 110 | return { |
| 111 | success: true, |
| 112 | image: generatedImage, |
| 113 | }; |
| 114 | } catch (error) { |
| 115 | console.error("Error generating image:", error); |
| 116 | return { |
| 117 | success: false, |
| 118 | error: |
| 119 | error instanceof Error ? error.message : "Failed to generate image", |
| 120 | }; |
| 121 | } |
| 122 | } |
| 123 |