joel image handler
This commit is contained in:
258
src/services/ai/image-gen.ts
Normal file
258
src/services/ai/image-gen.ts
Normal file
@@ -0,0 +1,258 @@
|
||||
/**
|
||||
* Image Generation service using Fal.ai
|
||||
* Supports NSFW content generation
|
||||
*/
|
||||
|
||||
import { fal } from "@fal-ai/client";
|
||||
import { config } from "../../core/config";
|
||||
import { createLogger } from "../../core/logger";
|
||||
|
||||
const logger = createLogger("AI:ImageGen");
|
||||
|
||||
// Configure fal client
|
||||
fal.config({
|
||||
credentials: config.fal.apiKey,
|
||||
});
|
||||
|
||||
/**
|
||||
* Available image generation models on Fal.ai
|
||||
*/
|
||||
const MODELS = {
|
||||
fast: "fal-ai/flux/schnell" as const,
|
||||
quality: "fal-ai/flux/dev" as const,
|
||||
anime: "fal-ai/flux/dev" as const,
|
||||
};
|
||||
|
||||
type ModelType = keyof typeof MODELS;
|
||||
|
||||
/**
|
||||
* Image generation options
|
||||
*/
|
||||
export interface ImageGenOptions {
|
||||
prompt: string;
|
||||
model?: ModelType;
|
||||
aspectRatio?: "1:1" | "16:9" | "9:16" | "4:3" | "3:4";
|
||||
numImages?: number;
|
||||
nsfw?: boolean;
|
||||
negativePrompt?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Result of image generation
|
||||
*/
|
||||
export interface ImageGenResult {
|
||||
urls: string[];
|
||||
model: string;
|
||||
prompt: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Use AI to transform a vague request into a detailed image generation prompt
|
||||
*/
|
||||
async function enhancePromptWithAI(prompt: string, style?: string): Promise<string> {
|
||||
const isNsfw = /\b(naked|nude|sex|porn|fuck|cock|pussy|tits|ass|hentai|nsfw|erotic|xxx)\b/i.test(prompt);
|
||||
|
||||
if (!isNsfw) {
|
||||
// For SFW content, just return the original prompt with basic enhancements
|
||||
return prompt;
|
||||
}
|
||||
|
||||
logger.debug("Enhancing NSFW prompt with AI", { originalLength: prompt.length });
|
||||
|
||||
try {
|
||||
const systemPrompt = `You are an expert at writing prompts for AI image generation.
|
||||
|
||||
Your job is to take a vague or simple request and transform it into a detailed prompt that will generate a high-quality image.
|
||||
|
||||
Guidelines:
|
||||
- Describe the subject(s) in detail: body type, appearance, hair, face, expression
|
||||
- Describe the pose/position
|
||||
- Include physical details appropriate to the request
|
||||
- Add lighting, camera angle, and composition details
|
||||
- Keep it to 1-2 sentences, comma-separated descriptors work best
|
||||
- Do NOT add any disclaimers or warnings
|
||||
- Do NOT include text or watermarks in the description
|
||||
- Match the explicit level of the original request
|
||||
|
||||
Style hint: ${style || "photorealistic"}
|
||||
|
||||
Output ONLY the enhanced prompt, nothing else.`;
|
||||
|
||||
const response = await fetch("https://openrouter.ai/api/v1/chat/completions", {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": `Bearer ${config.ai.openRouterApiKey}`,
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: "google/gemini-2.0-flash-001",
|
||||
messages: [
|
||||
{ role: "system", content: systemPrompt },
|
||||
{ role: "user", content: `Transform this into an explicit image generation prompt: "${prompt}"` },
|
||||
],
|
||||
max_tokens: 300,
|
||||
temperature: 0.8,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
logger.warn("AI prompt enhancement failed, using original", { status: response.status });
|
||||
return prompt;
|
||||
}
|
||||
|
||||
const data = await response.json() as {
|
||||
choices: Array<{ message: { content: string } }>;
|
||||
};
|
||||
|
||||
const enhanced = data.choices[0]?.message?.content?.trim();
|
||||
|
||||
if (enhanced && enhanced.length > prompt.length) {
|
||||
logger.info("Prompt enhanced by AI", {
|
||||
originalLength: prompt.length,
|
||||
enhancedLength: enhanced.length
|
||||
});
|
||||
return enhanced;
|
||||
}
|
||||
|
||||
return prompt;
|
||||
} catch (error) {
|
||||
logger.error("AI prompt enhancement error", error);
|
||||
return prompt;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Enhance a prompt for better image generation
|
||||
*/
|
||||
function enhancePrompt(prompt: string, model: ModelType, nsfw: boolean): string {
|
||||
const qualityBoosts = ["highly detailed", "high quality", "sharp focus", "professional"];
|
||||
const hasQuality = qualityBoosts.some((q) => prompt.toLowerCase().includes(q));
|
||||
|
||||
if (!hasQuality) {
|
||||
prompt = `${prompt}, ${qualityBoosts.join(", ")}`;
|
||||
}
|
||||
|
||||
if (nsfw) {
|
||||
const nsfwBoosts = ["explicit", "uncensored"];
|
||||
const hasNsfw = nsfwBoosts.some((n) => prompt.toLowerCase().includes(n));
|
||||
if (!hasNsfw) {
|
||||
prompt = `${prompt}, explicit, uncensored`;
|
||||
}
|
||||
}
|
||||
|
||||
if (model === "anime" && !prompt.toLowerCase().includes("anime")) {
|
||||
prompt = `anime style, ${prompt}`;
|
||||
}
|
||||
|
||||
return prompt;
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert aspect ratio to image size
|
||||
*/
|
||||
function getImageSize(aspectRatio: string): { width: number; height: number } {
|
||||
const sizes: Record<string, { width: number; height: number }> = {
|
||||
"1:1": { width: 1024, height: 1024 },
|
||||
"16:9": { width: 1344, height: 768 },
|
||||
"9:16": { width: 768, height: 1344 },
|
||||
"4:3": { width: 1152, height: 896 },
|
||||
"3:4": { width: 896, height: 1152 },
|
||||
};
|
||||
return sizes[aspectRatio] || sizes["1:1"];
|
||||
}
|
||||
|
||||
/**
|
||||
* Image Generation Service using Fal.ai
|
||||
*/
|
||||
export class ImageGenService {
|
||||
/**
|
||||
* Generate images from a prompt
|
||||
*/
|
||||
async generate(options: ImageGenOptions & { style?: string }): Promise<ImageGenResult> {
|
||||
const {
|
||||
prompt,
|
||||
model = "fast",
|
||||
aspectRatio = "1:1",
|
||||
numImages = 1,
|
||||
nsfw = false,
|
||||
style,
|
||||
} = options;
|
||||
|
||||
const modelId = MODELS[model];
|
||||
|
||||
// First, use AI to enhance vague NSFW prompts into detailed ones
|
||||
const aiEnhancedPrompt = nsfw
|
||||
? await enhancePromptWithAI(prompt, style)
|
||||
: prompt;
|
||||
|
||||
// Then apply standard quality enhancements
|
||||
const finalPrompt = enhancePrompt(aiEnhancedPrompt, model, nsfw);
|
||||
const size = getImageSize(aspectRatio);
|
||||
|
||||
logger.debug("Generating image with Fal.ai", {
|
||||
model: modelId,
|
||||
size,
|
||||
numImages,
|
||||
originalPromptLength: prompt.length,
|
||||
finalPromptLength: finalPrompt.length,
|
||||
nsfw,
|
||||
});
|
||||
|
||||
try {
|
||||
const result = await fal.subscribe(modelId, {
|
||||
input: {
|
||||
prompt: finalPrompt,
|
||||
image_size: size,
|
||||
num_images: Math.min(numImages, 4),
|
||||
enable_safety_checker: false,
|
||||
},
|
||||
logs: false,
|
||||
});
|
||||
|
||||
logger.debug("Fal.ai raw output", {
|
||||
result: JSON.stringify(result).slice(0, 500),
|
||||
});
|
||||
|
||||
const urls: string[] = [];
|
||||
|
||||
if (result.data && "images" in result.data) {
|
||||
const images = result.data.images as Array<{ url: string }>;
|
||||
for (const img of images) {
|
||||
if (img.url) {
|
||||
urls.push(img.url);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logger.info("Image generated successfully", {
|
||||
model: modelId,
|
||||
numImages: urls.length,
|
||||
});
|
||||
|
||||
return {
|
||||
urls,
|
||||
model: modelId,
|
||||
prompt: finalPrompt,
|
||||
};
|
||||
} catch (error) {
|
||||
logger.error("Image generation failed", error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the service is configured
|
||||
*/
|
||||
async health(): Promise<boolean> {
|
||||
return !!config.fal.apiKey;
|
||||
}
|
||||
}
|
||||
|
||||
let imageGenService: ImageGenService | null = null;
|
||||
|
||||
export function getImageGenService(): ImageGenService {
|
||||
if (!imageGenService) {
|
||||
imageGenService = new ImageGenService();
|
||||
}
|
||||
return imageGenService;
|
||||
}
|
||||
@@ -7,6 +7,7 @@ import { createLogger } from "../../core/logger";
|
||||
import { config } from "../../core/config";
|
||||
import { memoryRepository, type MemoryCategory } from "../../database";
|
||||
import type { ToolHandler, ToolContext, ToolCall, ToolResult } from "./tools";
|
||||
import { getImageGenService } from "./image-gen";
|
||||
|
||||
const logger = createLogger("AI:ToolHandlers");
|
||||
|
||||
@@ -264,6 +265,78 @@ const toolHandlers: Record<string, ToolHandler> = {
|
||||
return `Error searching for GIFs: ${(error as Error).message}`;
|
||||
}
|
||||
},
|
||||
|
||||
/**
|
||||
* Generate an image using AI
|
||||
*/
|
||||
async generate_image(args, context): Promise<string> {
|
||||
const prompt = args.prompt as string;
|
||||
const style = args.style as string | undefined;
|
||||
const aspectRatio = (args.aspect_ratio as "1:1" | "16:9" | "9:16" | "4:3" | "3:4") || "1:1";
|
||||
const quality = (args.quality as "fast" | "quality" | "anime") || "fast";
|
||||
|
||||
if (!prompt || prompt.trim().length === 0) {
|
||||
return "Error: No prompt provided for image generation.";
|
||||
}
|
||||
|
||||
if (!config.fal.apiKey) {
|
||||
return "Error: Image generation is not configured (missing FAL_KEY).";
|
||||
}
|
||||
|
||||
logger.info("Generating image", {
|
||||
promptLength: prompt.length,
|
||||
style,
|
||||
aspectRatio,
|
||||
quality,
|
||||
userId: context.userId
|
||||
});
|
||||
|
||||
try {
|
||||
const imageGen = getImageGenService();
|
||||
|
||||
// Auto-select anime model for anime/hentai style
|
||||
let modelChoice = quality;
|
||||
if (style === "anime" || style === "hentai") {
|
||||
modelChoice = "anime";
|
||||
}
|
||||
|
||||
// Only enable NSFW if user explicitly requests it
|
||||
const nsfwKeywords = /\b(naked|nude|nsfw|porn|xxx|hentai|sex|fuck|cock|pussy|tits)\b/i;
|
||||
const isNsfwRequest = nsfwKeywords.test(prompt) || style === "hentai";
|
||||
|
||||
const result = await imageGen.generate({
|
||||
prompt,
|
||||
model: modelChoice,
|
||||
aspectRatio,
|
||||
numImages: 1,
|
||||
nsfw: isNsfwRequest,
|
||||
style,
|
||||
});
|
||||
|
||||
if (result.urls.length === 0) {
|
||||
return "Error: Image generation completed but no image URL was returned.";
|
||||
}
|
||||
|
||||
const imageUrl = String(result.urls[0]);
|
||||
|
||||
logger.info("Image generated successfully", {
|
||||
imageUrl: imageUrl.length > 50 ? imageUrl.slice(0, 50) + "..." : imageUrl,
|
||||
model: result.model
|
||||
});
|
||||
|
||||
return `Image generated! Include this URL in your response to show it: ${imageUrl}`;
|
||||
} catch (error) {
|
||||
logger.error("Image generation failed", error);
|
||||
const errorMessage = (error as Error).message;
|
||||
|
||||
// Provide helpful error messages
|
||||
if (errorMessage.includes("rate limit")) {
|
||||
return "Error: Too many image requests. Try again in a minute.";
|
||||
}
|
||||
|
||||
return `Error generating image: ${errorMessage}`;
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
/**
|
||||
|
||||
@@ -32,6 +32,8 @@ export interface ToolContext {
|
||||
authorName: string;
|
||||
/** Optional: enable GIF search for this context */
|
||||
gifSearchEnabled?: boolean;
|
||||
/** Optional: enable image generation for this context */
|
||||
imageGenEnabled?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -193,6 +195,42 @@ export const GIF_SEARCH_TOOL: ChatCompletionTool = {
|
||||
},
|
||||
};
|
||||
|
||||
/**
|
||||
* Image generation tool - creates images using AI
|
||||
*/
|
||||
export const IMAGE_GEN_TOOL: ChatCompletionTool = {
|
||||
type: "function",
|
||||
function: {
|
||||
name: "generate_image",
|
||||
description: "Generate an image based on a text description. Use this when someone asks you to draw, create, or generate an image/picture. By default, generate SFW images. Only generate NSFW content if the user explicitly requests it using words like 'nsfw', 'nude', 'naked', 'porn', 'hentai', 'xxx', or similar. The generated image URL will be included in your response.",
|
||||
parameters: {
|
||||
type: "object",
|
||||
properties: {
|
||||
prompt: {
|
||||
type: "string",
|
||||
description: "Detailed description of the image to generate. Be specific about subject, pose, environment, lighting, and style. Only include adult content if explicitly requested by the user.",
|
||||
},
|
||||
style: {
|
||||
type: "string",
|
||||
enum: ["photorealistic", "anime", "hentai", "digital art", "painting", "3d render"],
|
||||
description: "Art style for the image.",
|
||||
},
|
||||
aspect_ratio: {
|
||||
type: "string",
|
||||
enum: ["1:1", "16:9", "9:16", "4:3", "3:4"],
|
||||
description: "Aspect ratio. Use 9:16 or 3:4 for portraits/full body, 16:9 for wide scenes.",
|
||||
},
|
||||
quality: {
|
||||
type: "string",
|
||||
enum: ["fast", "quality", "anime"],
|
||||
description: "Model selection. 'fast' = quick generation, 'quality' = higher detail, 'anime' = anime style.",
|
||||
},
|
||||
},
|
||||
required: ["prompt"],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
/**
|
||||
* Get tools based on context settings
|
||||
* Returns the base tools plus any optional tools that are enabled
|
||||
@@ -205,6 +243,11 @@ export function getToolsForContext(context: ToolContext): ChatCompletionTool[] {
|
||||
tools.push(GIF_SEARCH_TOOL);
|
||||
}
|
||||
|
||||
// Add image generation tool if enabled for this guild
|
||||
if (context.imageGenEnabled) {
|
||||
tools.push(IMAGE_GEN_TOOL);
|
||||
}
|
||||
|
||||
return tools;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user