1 line
11 KiB
JSON
1 line
11 KiB
JSON
{"repo": "aindywarhol/aindyagent", "files": {"packages/plugin-node/src/services/image.ts": {"42ff043f55f9a5031980e8a37829456c78d860e2": {"t": 1785595778.9513412, "n": 10868, "h": "c4a8cab41b9383a4", "c": "import { elizaLogger, models } from \"@ai16z/eliza\";\nimport { Service } from \"@ai16z/eliza\";\nimport {\n IAgentRuntime,\n ModelProviderName,\n ServiceType,\n IImageDescriptionService,\n} from \"@ai16z/eliza\";\nimport {\n AutoProcessor,\n AutoTokenizer,\n env,\n Florence2ForConditionalGeneration,\n Florence2Processor,\n PreTrainedModel,\n PreTrainedTokenizer,\n RawImage,\n type Tensor,\n} from \"@huggingface/transformers\";\nimport fs from \"fs\";\nimport gifFrames from \"gif-frames\";\nimport os from \"os\";\nimport path from \"path\";\n\nexport class ImageDescriptionService\n extends Service\n implements IImageDescriptionService\n{\n static serviceType: ServiceType = ServiceType.IMAGE_DESCRIPTION;\n\n private modelId: string = \"onnx-community/Florence-2-base-ft\";\n private device: string = \"gpu\";\n private model: PreTrainedModel | null = null;\n private processor: Florence2Processor | null = null;\n private tokenizer: PreTrainedTokenizer | null = null;\n private initialized: boolean = false;\n private runtime: IAgentRuntime | null = null;\n private queue: string[] = [];\n private processing: boolean = false;\n\n getInstance(): IImageDescriptionService {\n return ImageDescriptionService.getInstance();\n }\n\n async initialize(runtime: IAgentRuntime): Promise<void> {\n console.log(\"Initializing ImageDescriptionService\");\n this.runtime = runtime;\n }\n\n private async initializeLocalModel(): Promise<void> {\n env.allowLocalModels = false;\n env.allowRemoteModels = true;\n env.backends.onnx.logLevel = \"fatal\";\n env.backends.onnx.wasm.proxy = false;\n env.backends.onnx.wasm.numThreads = 1;\n\n elizaLogger.info(\"Downloading Florence model...\");\n\n this.model = await Florence2ForConditionalGeneration.from_pretrained(\n this.modelId,\n {\n device: \"gpu\",\n progress_callback: (progress) => {\n if (progress.status === \"downloading\") {\n const percent = (\n (progress.loaded / progress.total) *\n 100\n ).toFixed(1);\n const dots = \".\".repeat(\n Math.floor(Number(percent) / 5)\n );\n elizaLogger.info(\n `Downloading Florence model: [${dots.padEnd(20, \" \")}] ${percent}%`\n );\n }\n },\n }\n );\n\n elizaLogger.success(\"Florence model downloaded successfully\");\n\n elizaLogger.info(\"Downloading processor...\");\n this.processor = (await AutoProcessor.from_pretrained(\n this.modelId\n )) as Florence2Processor;\n\n elizaLogger.info(\"Downloading tokenizer...\");\n this.tokenizer = await AutoTokenizer.from_pretrained(this.modelId);\n elizaLogger.success(\"Image service initialization complete\");\n }\n\n async describeImage(\n imageUrl: string\n ): Promise<{ title: string; description: string }> {\n if (!this.initialized) {\n const model = models[this.runtime?.character?.modelProvider];\n\n if (model === models[ModelProviderName.LLAMALOCAL]) {\n await this.initializeLocalModel();\n } else {\n this.modelId = \"gpt-4o-mini\";\n this.device = \"cloud\";\n }\n\n this.initialized = true;\n }\n\n if (this.device === \"cloud\") {\n if (!this.runtime) {\n throw new Error(\n \"Runtime is required for OpenAI image recognition\"\n );\n }\n return this.recognizeWithOpenAI(imageUrl);\n }\n\n this.queue.push(imageUrl);\n this.processQueue();\n\n return new Promise((resolve, _reject) => {\n const checkQueue = () => {\n const index = this.queue.indexOf(imageUrl);\n if (index !== -1) {\n setTimeout(checkQueue, 100);\n } else {\n resolve(this.processImage(imageUrl));\n }\n };\n checkQueue();\n });\n }\n\n private async recognizeWithOpenAI(\n imageUrl: string\n ): Promise<{ title: string; description: string }> {\n const isGif = imageUrl.toLowerCase().endsWith(\".gif\");\n let imageData: Buffer | null = null;\n\n try {\n if (isGif) {\n const { filePath } =\n await this.extractFirstFrameFromGif(imageUrl);\n imageData = fs.readFileSync(filePath);\n } else {\n const response = await fetch(imageUrl);\n if (!response.ok) {\n throw new Error(\n `Failed to fetch image: ${response.statusText}`\n );\n }\n imageData = Buffer.from(await response.arrayBuffer());\n }\n\n if (!imageData || imageData.length === 0) {\n throw new Error(\"Failed to fetch image data\");\n }\n\n const prompt =\n \"Describe this image and give it a title. The first line should be the title, and then a line break, then a detailed description of the image. Respond with the format 'title\\ndescription'\";\n const text = await this.requestOpenAI(\n imageUrl,\n imageData,\n prompt,\n isGif\n );\n\n const [title, ...descriptionParts] = text.split(\"\\n\");\n return {\n title,\n description: descriptionParts.join(\"\\n\"),\n };\n } catch (error) {\n elizaLogger.error(\"Error in recognizeWithOpenAI:\", error);\n throw error;\n }\n }\n\n private async requestOpenAI(\n imageUrl: string,\n imageData: Buffer,\n prompt: string,\n isGif: boolean\n ): Promise<string> {\n for (let attempt = 0; attempt < 3; attempt++) {\n try {\n const content = [\n { type: \"text\", text: prompt },\n {\n type: \"image_url\",\n image_url: {\n url: isGif\n ? `data:image/png;base64,${imageData.toString(\"base64\")}`\n : imageUrl,\n },\n },\n ];\n\n const response = await fetch(\n \"https://api.openai.com/v1/chat/completions\",\n {\n method: \"POST\",\n headers: {\n \"Content-Type\": \"application/json\",\n Authorization: `Bearer ${this.runtime.getSetting(\"OPENAI_API_KEY\")}`,\n },\n body: JSON.stringify({\n model: \"gpt-4o-mini\",\n messages: [{ role: \"user\", content }],\n max_tokens: isGif ? 500 : 300,\n }),\n }\n );\n\n if (!response.ok) {\n throw new Error(`HTTP error! status: ${response.status}`);\n }\n\n const data = await response.json();\n return data.choices[0].message.content;\n } catch (error) {\n elizaLogger.error(\n `OpenAI request failed (attempt ${attempt + 1}):`,\n error\n );\n if (attempt === 2) throw error;\n }\n }\n throw new Error(\n \"Failed to recognize image with OpenAI after 3 attempts\"\n );\n }\n\n private async processQueue(): Promise<void> {\n if (this.processing || this.queue.length === 0) return;\n\n this.processing = true;\n while (this.queue.length > 0) {\n const imageUrl = this.queue.shift();\n await this.processImage(imageUrl);\n }\n this.processing = false;\n }\n\n private async processImage(\n imageUrl: string\n ): Promise<{ title: string; description: string }> {\n if (!this.model || !this.processor || !this.tokenizer) {\n throw new Error(\"Model components not initialized\");\n }\n\n elizaLogger.log(\"Processing image:\", imageUrl);\n const isGif = imageUrl.toLowerCase().endsWith(\".gif\");\n let imageToProcess = imageUrl;\n\n try {\n if (isGif) {\n elizaLogger.log(\"Extracting first frame from GIF\");\n const { filePath } =\n await this.extractFirstFrameFromGif(imageUrl);\n imageToProcess = filePath;\n }\n\n const image = await RawImage.fromURL(imageToProcess);\n const visionInputs = await this.processor(image);\n const prompts =\n this.processor.construct_prompts(\"<DETAILED_CAPTION>\");\n const textInputs = this.tokenizer(prompts);\n\n elizaLogger.log(\"Generating image description\");\n const generatedIds = (await this.model.generate({\n ...textInputs,\n ...visionInputs,\n max_new_tokens: 256,\n })) as Tensor;\n\n const generatedText = this.tokenizer.batch_decode(generatedIds, {\n skip_special_tokens: false,\n })[0];\n\n const result = this.processor.post_process_generation(\n generatedText,\n \"<DETAILED_CAPTION>\",\n image.size\n );\n\n const detailedCaption = result[\"<DETAILED_CAPTION>\"] as string;\n return { title: detailedCaption, description: detailedCaption };\n } catch (error) {\n elizaLogger.error(\"Error processing image:\", error);\n throw error;\n } finally {\n if (isGif && imageToProcess !== imageUrl) {\n fs.unlinkSync(imageToProcess);\n }\n }\n }\n\n private async extractFirstFrameFromGif(\n gifUrl: string\n ): Promise<{ filePath: string }> {\n const frameData = await gifFrames({\n url: gifUrl,\n frames: 1,\n outputType: \"png\",\n });\n\n const tempFilePath = path.join(\n os.tmpdir(),\n `gif_frame_${Date.now()}.png`\n );\n\n return new Promise((resolve, reject) => {\n const writeStream = fs.createWriteStream(tempFilePath);\n frameData[0].getImage().pipe(writeStream);\n writeStream.on(\"finish\", () => resolve({ filePath: tempFilePath }));\n writeStream.on(\"error\", reject);\n });\n }\n}\n\nexport default ImageDescriptionService;\n"}}}} |