feat: support custom dalle model

This commit is contained in:
Hk-Gosuto 2023-12-21 12:36:57 +08:00
parent 666e62d75a
commit 4044891f85
2 changed files with 13 additions and 3 deletions

View File

@ -54,6 +54,7 @@
- DALL-E 3 - DALL-E 3
- DALL-E 3 插件需要配置对象存储服务,请参考 [对象存储服务配置指南](./docs/s3-oss.md) 配置 - DALL-E 3 插件需要配置对象存储服务,请参考 [对象存储服务配置指南](./docs/s3-oss.md) 配置
- 如无需图像转存则可以配置 `DALLE_NO_IMAGE_STORAGE=1` ,此时将直接将 DALL-E 服务返回的临时 URL 用于图像显示,注意:该链接具有时效性 - 如无需图像转存则可以配置 `DALLE_NO_IMAGE_STORAGE=1` ,此时将直接将 DALL-E 服务返回的临时 URL 用于图像显示,注意:该链接具有时效性
- 默认使用 `dall-e-3` 模型,如果想使用 `dall-e-2` ,可以配置环境变量 `DALLE_MODEL=dall-e-2`
- StableDiffusion - StableDiffusion
- 本插件目前为测试版本,后续可能会有较大的变更,请谨慎使用 - 本插件目前为测试版本,后续可能会有较大的变更,请谨慎使用
- 使用本插件需要一定的专业知识Stable Diffusion 本身的相关问题不在本项目的解答范围内,如果您确定要使用本插件请参考 [Stable Diffusion 插件配置指南](./docs/stable-diffusion-plugin-cn.md) 文档进行配置 - 使用本插件需要一定的专业知识Stable Diffusion 本身的相关问题不在本项目的解答范围内,如果您确定要使用本插件请参考 [Stable Diffusion 插件配置指南](./docs/stable-diffusion-plugin-cn.md) 文档进行配置

View File

@ -7,6 +7,7 @@ export class DallEAPIWrapper extends StructuredTool {
n = 1; n = 1;
apiKey: string; apiKey: string;
baseURL?: string; baseURL?: string;
model: string;
noStorage: boolean; noStorage: boolean;
@ -26,6 +27,7 @@ export class DallEAPIWrapper extends StructuredTool {
this.callback = callback; this.callback = callback;
this.noStorage = !!process.env.DALLE_NO_IMAGE_STORAGE; this.noStorage = !!process.env.DALLE_NO_IMAGE_STORAGE;
this.model = process.env.DALLE_MODEL ?? "dall-e-3";
} }
async saveImageFromUrl(url: string) { async saveImageFromUrl(url: string) {
@ -59,18 +61,25 @@ export class DallEAPIWrapper extends StructuredTool {
Authorization: `Bearer ${this.apiKey}`, Authorization: `Bearer ${this.apiKey}`,
}, },
body: JSON.stringify({ body: JSON.stringify({
model: "dall-e-3", model: this.model,
prompt: prompt, prompt: prompt,
n: this.n, n: this.n,
size: size, size: size,
}), }),
}; };
console.log(requestOptions);
const response = await fetch(apiUrl, requestOptions); const response = await fetch(apiUrl, requestOptions);
const json = await response.json(); const json = await response.json();
console.log("[DALL-E]", json); try {
imageUrl = json.data[0].url; console.log("[DALL-E]", json);
imageUrl = json.data[0].url;
} catch (e) {
if (this.callback != null) await this.callback(JSON.stringify(json));
throw e;
}
} catch (e) { } catch (e) {
console.error("[DALL-E]", e); console.error("[DALL-E]", e);
return (e as Error).message;
} }
if (!imageUrl) return "No image was generated"; if (!imageUrl) return "No image was generated";
try { try {