mirror of
https://github.com/lobehub/lobehub
synced 2026-04-21 17:47:27 +00:00
✨ feat: add ComfyUI integration Phase1(RFC-128) (#9043)
Co-authored-by: YuTengjing <ytj2713151713@gmail.com>
This commit is contained in:
Maple Gao
GitHub
parent
2606f93146
commit
15ffe289f5
130 changed files with 22066 additions and 32 deletions
|
|
@ -165,6 +165,9 @@ ENV \
|
|||
CLOUDFLARE_API_KEY="" CLOUDFLARE_BASE_URL_OR_ACCOUNT_ID="" CLOUDFLARE_MODEL_LIST="" \
|
||||
# Cohere
|
||||
COHERE_API_KEY="" COHERE_MODEL_LIST="" COHERE_PROXY_URL="" \
|
||||
# ComfyUI
|
||||
COMFYUI_BASE_URL="" COMFYUI_AUTH_TYPE="" \
|
||||
COMFYUI_API_KEY="" COMFYUI_USERNAME="" COMFYUI_PASSWORD="" COMFYUI_CUSTOM_HEADERS="" \
|
||||
# DeepSeek
|
||||
DEEPSEEK_API_KEY="" DEEPSEEK_MODEL_LIST="" \
|
||||
# Fireworks AI
|
||||
|
|
|
|||
|
|
@ -218,6 +218,9 @@ ENV \
|
|||
CLOUDFLARE_API_KEY="" CLOUDFLARE_BASE_URL_OR_ACCOUNT_ID="" CLOUDFLARE_MODEL_LIST="" \
|
||||
# Cohere
|
||||
COHERE_API_KEY="" COHERE_MODEL_LIST="" COHERE_PROXY_URL="" \
|
||||
# ComfyUI
|
||||
COMFYUI_BASE_URL="" COMFYUI_AUTH_TYPE="" \
|
||||
COMFYUI_API_KEY="" COMFYUI_USERNAME="" COMFYUI_PASSWORD="" COMFYUI_CUSTOM_HEADERS="" \
|
||||
# DeepSeek
|
||||
DEEPSEEK_API_KEY="" DEEPSEEK_MODEL_LIST="" \
|
||||
# Fireworks AI
|
||||
|
|
|
|||
|
|
@ -167,6 +167,9 @@ ENV \
|
|||
CLOUDFLARE_API_KEY="" CLOUDFLARE_BASE_URL_OR_ACCOUNT_ID="" CLOUDFLARE_MODEL_LIST="" \
|
||||
# Cohere
|
||||
COHERE_API_KEY="" COHERE_MODEL_LIST="" COHERE_PROXY_URL="" \
|
||||
# ComfyUI
|
||||
COMFYUI_BASE_URL="" COMFYUI_AUTH_TYPE="" \
|
||||
COMFYUI_API_KEY="" COMFYUI_USERNAME="" COMFYUI_PASSWORD="" COMFYUI_CUSTOM_HEADERS="" \
|
||||
# DeepSeek
|
||||
DEEPSEEK_API_KEY="" DEEPSEEK_MODEL_LIST="" \
|
||||
# Fireworks AI
|
||||
|
|
|
|||
1009
docs/development/basic/comfyui-development.mdx
Normal file
1009
docs/development/basic/comfyui-development.mdx
Normal file
File diff suppressed because it is too large
Load diff
998
docs/development/basic/comfyui-development.zh-CN.mdx
Normal file
998
docs/development/basic/comfyui-development.zh-CN.mdx
Normal file
|
|
@ -0,0 +1,998 @@
|
|||
---
|
||||
title: ComfyUI 扩展开发指南
|
||||
description: 学习如何为 LobeChat ComfyUI 集成添加新模型、工作流和功能扩展
|
||||
tags:
|
||||
- ComfyUI
|
||||
- 开发指南
|
||||
- 模型扩展
|
||||
- 工作流开发
|
||||
---
|
||||
|
||||
# ComfyUI 扩展开发指南
|
||||
|
||||
本指南基于实际代码实现，帮助开发者扩展 LobeChat 的 ComfyUI 集成功能。
|
||||
|
||||
## 架构概览
|
||||
|
||||
LobeChat ComfyUI 集成采用四层服务架构，围绕 `LobeComfyUI` 主类构建：
|
||||
|
||||
```plaintext
|
||||
packages/model-runtime/src/providers/comfyui/
|
||||
├── index.ts # LobeComfyUI 主类入口
|
||||
├── services/ # 四大核心服务
|
||||
│ ├── comfyuiClient.ts # ComfyUIClientService - 客户端和认证
|
||||
│ ├── modelResolver.ts # ModelResolverService - 模型解析
|
||||
│ ├── workflowBuilder.ts # WorkflowBuilderService - 工作流构建
|
||||
│ └── imageService.ts # ImageService - 图像生成
|
||||
├── config/ # 配置系统
|
||||
│ ├── modelRegistry.ts # 主模型注册表（222个模型）
|
||||
│ ├── fluxModelRegistry.ts # 130个FLUX模型配置
|
||||
│ ├── sdModelRegistry.ts # 92个SD系列模型配置
|
||||
│ ├── systemComponents.ts # VAE/CLIP/T5/LoRA/ControlNet组件
|
||||
│ └── workflowRegistry.ts # 工作流路由配置
|
||||
├── workflows/ # 工作流实现
|
||||
│ ├── flux-dev.ts # FLUX Dev 20步工作流
|
||||
│ ├── flux-schnell.ts # FLUX Schnell 4步快速工作流
|
||||
│ ├── flux-kontext.ts # FLUX Kontext 填充工作流
|
||||
│ ├── sd35.ts # SD3.5 外部编码器工作流
|
||||
│ ├── simple-sd.ts # 通用SD工作流
|
||||
│ └── index.ts # 工作流导出
|
||||
├── utils/ # 工具层
|
||||
│ ├── staticModelLookup.ts # 模型查找函数
|
||||
│ ├── workflowDetector.ts # 模型架构检测
|
||||
│ ├── promptSplitter.ts # FLUX双提示词分割
|
||||
│ ├── seedGenerator.ts # 随机种子生成
|
||||
│ ├── cacheManager.ts # TTL缓存管理
|
||||
│ └── workflowUtils.ts # 工作流工具函数
|
||||
└── errors/ # 错误处理
|
||||
├── base.ts # 基础错误类
|
||||
├── modelResolverError.ts # 模型解析错误
|
||||
├── workflowError.ts # 工作流错误
|
||||
└── servicesError.ts # 服务错误
|
||||
|
||||
src/server/services/comfyui/ # 服务端实现
|
||||
├── core/ # 核心服务器服务
|
||||
│ ├── comfyUIAuthService.ts # 认证服务
|
||||
│ ├── comfyUIClientService.ts # 客户端服务
|
||||
│ ├── comfyUIConnectionService.ts # 连接服务
|
||||
│ ├── errorHandlerService.ts # 错误处理服务
|
||||
│ ├── imageService.ts # 图像生成服务
|
||||
│ ├── modelResolverService.ts # 模型解析服务
|
||||
│ └── workflowBuilderService.ts # 工作流构建服务
|
||||
├── config/ # 服务器端配置
|
||||
│ ├── constants.ts # 常量和默认值
|
||||
│ ├── modelRegistry.ts # 模型注册表
|
||||
│ ├── fluxModelRegistry.ts # FLUX模型
|
||||
│ ├── sdModelRegistry.ts # SD模型
|
||||
│ ├── systemComponents.ts # 系统组件
|
||||
│ └── workflowRegistry.ts # 工作流注册表
|
||||
├── workflows/ # 服务端工作流实现
|
||||
│ ├── flux-dev.ts # FLUX Dev 工作流
|
||||
│ ├── flux-schnell.ts # FLUX Schnell 工作流
|
||||
│ ├── flux-kontext.ts # FLUX Kontext 工作流
|
||||
│ ├── sd35.ts # SD3.5 工作流
|
||||
│ └── simple-sd.ts # Simple SD 工作流
|
||||
├── utils/ # 服务器工具
|
||||
│ ├── cacheManager.ts # 缓存管理
|
||||
│ ├── componentInfo.ts # 组件信息
|
||||
│ ├── imageResizer.ts # 图像调整
|
||||
│ ├── promptSplitter.ts # 提示词分割
|
||||
│ ├── staticModelLookup.ts # 模型查找
|
||||
│ ├── weightDType.ts # 权重数据类型工具
|
||||
│ ├── workflowDetector.ts # 工作流检测
|
||||
│ └── workflowUtils.ts # 工作流工具
|
||||
└── errors/ # 服务器错误处理
|
||||
├── base.ts # 基础错误类
|
||||
├── configError.ts # 配置错误
|
||||
├── modelResolverError.ts # 模型解析器错误
|
||||
├── servicesError.ts # 服务错误
|
||||
├── utilsError.ts # 工具错误
|
||||
└── workflowError.ts # 工作流错误
|
||||
|
||||
packages/model-runtime/src/utils/ # 共享工具
|
||||
└── comfyuiErrorParser.ts # 客户端/服务器统一错误解析器
|
||||
```
|
||||
|
||||
### 核心服务架构
|
||||
|
||||
`LobeComfyUI` 主类初始化四个核心服务：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/index.ts
|
||||
export class LobeComfyUI implements LobeRuntimeAI, AuthenticatedImageRuntime {
|
||||
constructor(options: ComfyUIKeyVault = {}) {
|
||||
// 1. 客户端服务 - 处理认证和API调用
|
||||
this.clientService = new ComfyUIClientService(options);
|
||||
|
||||
// 2. 模型解析服务 - 模型查找和组件选择
|
||||
const modelResolverService = new ModelResolverService(this.clientService);
|
||||
|
||||
// 3. 工作流构建服务 - 路由和构建工作流
|
||||
const workflowBuilderService = new WorkflowBuilderService({
|
||||
clientService: this.clientService,
|
||||
modelResolverService: modelResolverService,
|
||||
});
|
||||
|
||||
// 4. 图像服务 - 统一的图像生成入口
|
||||
this.imageService = new ImageService(
|
||||
this.clientService,
|
||||
modelResolverService,
|
||||
workflowBuilderService,
|
||||
);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## 认证系统
|
||||
|
||||
ComfyUI 集成支持四种认证方式，由 `ComfyUIClientService` 内的 `AuthManager` 处理：
|
||||
|
||||
### 支持的认证类型
|
||||
|
||||
```typescript
|
||||
interface ComfyUIKeyVault {
|
||||
baseURL: string;
|
||||
authType?: 'none' | 'basic' | 'bearer' | 'custom';
|
||||
// Basic Auth
|
||||
username?: string;
|
||||
password?: string;
|
||||
// Bearer Token
|
||||
apiKey?: string;
|
||||
// Custom Headers
|
||||
customHeaders?: Record<string, string>;
|
||||
}
|
||||
```
|
||||
|
||||
### 认证配置示例
|
||||
|
||||
```typescript
|
||||
// 无认证
|
||||
const comfyUI = new LobeComfyUI({
|
||||
baseURL: 'http://localhost:8000',
|
||||
authType: 'none'
|
||||
});
|
||||
|
||||
// 基础认证
|
||||
const comfyUI = new LobeComfyUI({
|
||||
baseURL: 'https://your-comfyui-server.com',
|
||||
authType: 'basic',
|
||||
username: 'your-username',
|
||||
password: 'your-password'
|
||||
});
|
||||
|
||||
// Bearer Token
|
||||
const comfyUI = new LobeComfyUI({
|
||||
baseURL: 'https://your-comfyui-server.com',
|
||||
authType: 'bearer',
|
||||
apiKey: 'your-api-key'
|
||||
});
|
||||
|
||||
// 自定义头部
|
||||
const comfyUI = new LobeComfyUI({
|
||||
baseURL: 'https://your-comfyui-server.com',
|
||||
authType: 'custom',
|
||||
customHeaders: {
|
||||
'X-API-Key': 'your-custom-key',
|
||||
'Authorization': 'Custom your-token'
|
||||
}
|
||||
});
|
||||
```
|
||||
|
||||
## WebAPI 路由
|
||||
|
||||
ComfyUI 提供了用于图像生成的 REST WebAPI 路由，支持常规认证和内部服务认证：
|
||||
|
||||
### 路由详情
|
||||
|
||||
```typescript
|
||||
// src/app/(backend)/webapi/create-image/comfyui/route.ts
|
||||
export const runtime = 'nodejs';
|
||||
export const maxDuration = 300; // 最长5分钟
|
||||
|
||||
// POST /api/create-image/comfyui
|
||||
{
|
||||
model: string; // 模型标识符
|
||||
params: { // 生成参数
|
||||
prompt: string;
|
||||
width?: number;
|
||||
height?: number;
|
||||
// ... 其他参数
|
||||
};
|
||||
options?: { // 可选生成选项
|
||||
// ... 额外选项
|
||||
};
|
||||
}
|
||||
```
|
||||
|
||||
### 认证中间件
|
||||
|
||||
WebAPI 路由使用 `checkAuth` 中间件进行认证：
|
||||
|
||||
```typescript
|
||||
import { checkAuth } from '@/app/(backend)/middleware/auth';
|
||||
|
||||
// 路由自动验证 JWT 令牌
|
||||
// 并将认证上下文传递给 tRPC 调用器
|
||||
```
|
||||
|
||||
### 错误处理
|
||||
|
||||
WebAPI 路由提供结构化的错误响应：
|
||||
|
||||
```typescript
|
||||
// 从 TRPCError 的 cause 中提取 AgentRuntimeError
|
||||
if (agentError && 'errorType' in agentError) {
|
||||
// 将 errorType 转换为适当的 HTTP 状态码
|
||||
// 401 对应 InvalidProviderAPIKey
|
||||
// 403 对应 PermissionDenied
|
||||
// 404 对应 NotFound
|
||||
// 500+ 对应服务器错误
|
||||
}
|
||||
```
|
||||
|
||||
## 添加新模型
|
||||
|
||||
### 1. 理解模型注册表结构
|
||||
|
||||
模型配置存储在配置文件中：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/config/modelRegistry.ts
|
||||
export interface ModelConfig {
|
||||
modelFamily: 'FLUX' | 'SD1' | 'SDXL' | 'SD3';
|
||||
priority: number; // 1=官方, 2=企业, 3=社区
|
||||
recommendedDtype?: 'default' | 'fp8_e4m3fn' | 'fp8_e5m2';
|
||||
variant: string; // 模型变体标识符
|
||||
}
|
||||
```
|
||||
|
||||
### 2. 添加 FLUX 模型
|
||||
|
||||
在 `fluxModelRegistry.ts` 中添加新模型：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/config/fluxModelRegistry.ts
|
||||
export const FLUX_MODEL_REGISTRY: Record<string, ModelConfig> = {
|
||||
// 现有模型...
|
||||
|
||||
// 添加新的FLUX Dev模型
|
||||
'your-custom-flux-dev.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2, // 企业级模型
|
||||
variant: 'dev',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 添加量化版本
|
||||
'your-custom-flux-dev-fp8.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
};
|
||||
```
|
||||
|
||||
### 3. 添加 SD 系列模型
|
||||
|
||||
在 `sdModelRegistry.ts` 中添加：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/config/sdModelRegistry.ts
|
||||
export const SD_MODEL_REGISTRY: Record<string, ModelConfig> = {
|
||||
// 现有模型...
|
||||
|
||||
// 添加新的SD3.5模型
|
||||
'your-custom-sd35.safetensors': {
|
||||
modelFamily: 'SD3',
|
||||
priority: 2,
|
||||
variant: 'sd35',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
};
|
||||
```
|
||||
|
||||
### 4. 更新模型 ID 映射（可选）
|
||||
|
||||
如果需要为前端提供友好的模型 ID，在 `modelRegistry.ts` 中添加映射：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/config/modelRegistry.ts
|
||||
export const MODEL_ID_VARIANT_MAP: Record<string, string> = {
|
||||
// 现有映射...
|
||||
|
||||
// 添加新模型的友好ID
|
||||
'my-custom-flux': 'dev', // 映射到dev变体
|
||||
'my-custom-sd35': 'sd35', // 映射到sd35变体
|
||||
};
|
||||
```
|
||||
|
||||
## 创建新工作流
|
||||
|
||||
### 工作流创建原理
|
||||
|
||||
**重要：工作流节点结构来自 ComfyUI 原生导出**
|
||||
|
||||
1. 在 ComfyUI 界面中设计工作流
|
||||
2. 使用 "Export (API Format)" 导出 JSON
|
||||
3. 将 JSON 结构复制到 TypeScript 文件
|
||||
4. 使用`PromptBuilder`包装并参数化
|
||||
|
||||
### 1. 从 ComfyUI 导出工作流
|
||||
|
||||
在 ComfyUI 界面中：
|
||||
|
||||
1. 拖拽节点构建所需工作流
|
||||
2. 连接各节点的输入输出
|
||||
3. 右键点击空白处 → "Export (API Format)"
|
||||
4. 复制生成的 JSON 结构
|
||||
|
||||
### 2. 工作流文件模板
|
||||
|
||||
创建新文件 `workflows/your-workflow.ts`：
|
||||
|
||||
```typescript
|
||||
import { PromptBuilder } from '@saintno/comfyui-sdk';
|
||||
|
||||
import type { WorkflowContext } from '../services/workflowBuilder';
|
||||
import { generateUniqueSeeds } from '../utils/seedGenerator';
|
||||
import { getWorkflowFilenamePrefix } from '../utils/workflowUtils';
|
||||
|
||||
/**
|
||||
* 构建自定义工作流
|
||||
* @param modelFileName - 模型文件名
|
||||
* @param params - 生成参数
|
||||
* @param context - 工作流上下文
|
||||
*/
|
||||
export async function buildYourCustomWorkflow(
|
||||
modelFileName: string,
|
||||
params: Record<string, any>,
|
||||
context: WorkflowContext,
|
||||
): Promise<PromptBuilder<any, any, any>> {
|
||||
|
||||
// 从ComfyUI "Export (API Format)" 获得的JSON结构
|
||||
const workflow = {
|
||||
'1': {
|
||||
_meta: { title: 'Load Checkpoint' },
|
||||
class_type: 'CheckpointLoaderSimple',
|
||||
inputs: {
|
||||
ckpt_name: modelFileName,
|
||||
},
|
||||
},
|
||||
'2': {
|
||||
_meta: { title: 'CLIP Text Encode' },
|
||||
class_type: 'CLIPTextEncode',
|
||||
inputs: {
|
||||
clip: ['1', 1], // 连接到节点1的CLIP输出
|
||||
text: params.prompt,
|
||||
},
|
||||
},
|
||||
'3': {
|
||||
_meta: { title: 'Empty Latent' },
|
||||
class_type: 'EmptyLatentImage',
|
||||
inputs: {
|
||||
width: params.width,
|
||||
height: params.height,
|
||||
batch_size: 1,
|
||||
},
|
||||
},
|
||||
'4': {
|
||||
_meta: { title: 'KSampler' },
|
||||
class_type: 'KSampler',
|
||||
inputs: {
|
||||
model: ['1', 0], // 连接到节点1的MODEL输出
|
||||
positive: ['2', 0], // 连接到节点2的CONDITIONING输出
|
||||
negative: ['2', 0], // 可以配置负面提示词
|
||||
latent_image: ['3', 0],
|
||||
seed: params.seed ?? generateUniqueSeeds(1)[0],
|
||||
steps: params.steps,
|
||||
cfg: params.cfg,
|
||||
sampler_name: 'euler',
|
||||
scheduler: 'normal',
|
||||
denoise: 1.0,
|
||||
},
|
||||
},
|
||||
'5': {
|
||||
_meta: { title: 'VAE Decode' },
|
||||
class_type: 'VAEDecode',
|
||||
inputs: {
|
||||
samples: ['4', 0],
|
||||
vae: ['1', 2], // 连接到节点1的VAE输出
|
||||
},
|
||||
},
|
||||
'6': {
|
||||
_meta: { title: 'Save Image' },
|
||||
class_type: 'SaveImage',
|
||||
inputs: {
|
||||
filename_prefix: getWorkflowFilenamePrefix('buildYourCustomWorkflow', context.variant),
|
||||
images: ['5', 0],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
// 使用PromptBuilder包装静态JSON
|
||||
const builder = new PromptBuilder(
|
||||
workflow,
|
||||
['width', 'height', 'steps', 'cfg', 'seed'], // 输入参数
|
||||
['images'], // 输出参数
|
||||
);
|
||||
|
||||
// 设置输出节点
|
||||
builder.setOutputNode('images', '6');
|
||||
|
||||
// 设置输入节点路径
|
||||
builder.setInputNode('width', '3.inputs.width');
|
||||
builder.setInputNode('height', '3.inputs.height');
|
||||
builder.setInputNode('steps', '4.inputs.steps');
|
||||
builder.setInputNode('cfg', '4.inputs.cfg');
|
||||
builder.setInputNode('seed', '4.inputs.seed');
|
||||
|
||||
// 设置参数值
|
||||
builder
|
||||
.input('width', params.width)
|
||||
.input('height', params.height)
|
||||
.input('steps', params.steps)
|
||||
.input('cfg', params.cfg)
|
||||
.input('seed', params.seed ?? generateUniqueSeeds(1)[0]);
|
||||
|
||||
return builder;
|
||||
}
|
||||
```
|
||||
|
||||
### 3. 注册新工作流
|
||||
|
||||
在 `workflowRegistry.ts` 中添加工作流映射：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/config/workflowRegistry.ts
|
||||
import { buildYourCustomWorkflow } from '../workflows/your-workflow';
|
||||
|
||||
export const VARIANT_WORKFLOW_MAP: Record<string, WorkflowBuilder> = {
|
||||
// 现有映射...
|
||||
|
||||
// 添加新工作流
|
||||
'your-variant': buildYourCustomWorkflow,
|
||||
};
|
||||
```
|
||||
|
||||
### 4. 实际工作流示例
|
||||
|
||||
参考 `flux-dev.ts` 的真实实现：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/workflows/flux-dev.ts (简化版)
|
||||
export async function buildFluxDevWorkflow(
|
||||
modelFileName: string,
|
||||
params: Record<string, any>,
|
||||
context: WorkflowContext,
|
||||
): Promise<PromptBuilder<any, any, any>> {
|
||||
// 获取所需组件
|
||||
const selectedT5Model = await context.modelResolverService.getOptimalComponent('t5', 'FLUX');
|
||||
const selectedVAE = await context.modelResolverService.getOptimalComponent('vae', 'FLUX');
|
||||
const selectedCLIP = await context.modelResolverService.getOptimalComponent('clip', 'FLUX');
|
||||
|
||||
// 处理双提示词分割
|
||||
const { t5xxlPrompt, clipLPrompt } = splitPromptForDualCLIP(params.prompt);
|
||||
|
||||
// 静态工作流定义（来自ComfyUI导出）
|
||||
const workflow = {
|
||||
'1': {
|
||||
class_type: 'DualCLIPLoader',
|
||||
inputs: {
|
||||
clip_name1: selectedT5Model,
|
||||
clip_name2: selectedCLIP,
|
||||
type: 'flux',
|
||||
},
|
||||
},
|
||||
// ... 更多节点
|
||||
};
|
||||
|
||||
// 参数注入（必须在workflow文件内完成）
|
||||
workflow['5'].inputs.clip_l = clipLPrompt;
|
||||
workflow['5'].inputs.t5xxl = t5xxlPrompt;
|
||||
workflow['4'].inputs.width = params.width;
|
||||
workflow['4'].inputs.height = params.height;
|
||||
|
||||
// 创建并配置PromptBuilder
|
||||
const builder = new PromptBuilder(workflow, inputs, outputs);
|
||||
// 配置输入输出映射...
|
||||
|
||||
return builder;
|
||||
}
|
||||
```
|
||||
|
||||
## 系统组件管理
|
||||
|
||||
### 组件配置结构
|
||||
|
||||
所有系统组件（VAE、CLIP、T5、LoRA、ControlNet）统一配置在 `systemComponents.ts`：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/config/systemComponents.ts
|
||||
export interface ComponentConfig {
|
||||
modelFamily: string; // 模型家族
|
||||
priority: number; // 1=必需, 2=标准, 3=可选
|
||||
type: string; // 组件类型
|
||||
compatibleVariants?: string[]; // 兼容变体（LoRA/ControlNet）
|
||||
controlnetType?: string; // ControlNet类型
|
||||
}
|
||||
|
||||
export const SYSTEM_COMPONENTS: Record<string, ComponentConfig> = {
|
||||
// VAE组件
|
||||
'ae.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'vae',
|
||||
},
|
||||
|
||||
// CLIP组件
|
||||
'clip_l.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'clip',
|
||||
},
|
||||
|
||||
// T5编码器
|
||||
't5xxl_fp16.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 't5',
|
||||
},
|
||||
|
||||
// LoRA适配器
|
||||
'realism_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
// ControlNet模型
|
||||
'flux-controlnet-canny-v3.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
controlnetType: 'canny',
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'controlnet',
|
||||
},
|
||||
};
|
||||
```
|
||||
|
||||
### 添加新组件
|
||||
|
||||
```typescript
|
||||
// 添加新的LoRA
|
||||
'your-custom-lora.safetensors': {
|
||||
compatibleVariants: ['dev', 'schnell'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
// 添加新的ControlNet
|
||||
'your-controlnet-pose.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
controlnetType: 'pose',
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'controlnet',
|
||||
},
|
||||
```
|
||||
|
||||
### 组件查询 API
|
||||
|
||||
```typescript
|
||||
import { getAllComponentsWithNames, getOptimalComponent } from '../config/systemComponents';
|
||||
|
||||
// 获取最优组件
|
||||
const bestVAE = getOptimalComponent('vae', 'FLUX');
|
||||
const bestT5 = getOptimalComponent('t5', 'FLUX');
|
||||
|
||||
// 查询特定类型的组件
|
||||
const availableLoras = getAllComponentsWithNames({
|
||||
type: 'lora',
|
||||
modelFamily: 'FLUX',
|
||||
compatibleVariant: 'dev'
|
||||
});
|
||||
|
||||
// 查询ControlNet
|
||||
const cannyControlNets = getAllComponentsWithNames({
|
||||
type: 'controlnet',
|
||||
controlnetType: 'canny',
|
||||
modelFamily: 'FLUX'
|
||||
});
|
||||
```
|
||||
|
||||
## 模型解析和查找
|
||||
|
||||
### ModelResolverService 工作原理
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/services/modelResolver.ts
|
||||
export class ModelResolverService {
|
||||
async resolveModelFileName(modelId: string): Promise<string | undefined> {
|
||||
// 1. 清理模型ID
|
||||
const cleanId = modelId.replace(/^comfyui\//, '');
|
||||
|
||||
// 2. 检查模型ID映射
|
||||
const mappedVariant = MODEL_ID_VARIANT_MAP[cleanId];
|
||||
if (mappedVariant) {
|
||||
const prioritizedModels = getModelsByVariant(mappedVariant);
|
||||
const serverModels = await this.getAvailableModelFiles();
|
||||
|
||||
// 按优先级查找第一个可用模型
|
||||
for (const filename of prioritizedModels) {
|
||||
if (serverModels.includes(filename)) {
|
||||
return filename;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 直接注册表查找
|
||||
if (MODEL_REGISTRY[cleanId]) {
|
||||
return cleanId;
|
||||
}
|
||||
|
||||
// 4. 检查服务器文件存在性
|
||||
if (isModelFile(cleanId)) {
|
||||
const serverModels = await this.getAvailableModelFiles();
|
||||
if (serverModels.includes(cleanId)) {
|
||||
return cleanId;
|
||||
}
|
||||
}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 模型查找示例
|
||||
|
||||
```typescript
|
||||
// 实际使用示例
|
||||
const resolver = new ModelResolverService(clientService);
|
||||
|
||||
// 友好ID查找
|
||||
const fluxDevFile = await resolver.resolveModelFileName('flux-dev');
|
||||
// 返回: 'flux1-dev.safetensors' (如果存在)
|
||||
|
||||
// 直接文件名查找
|
||||
const directFile = await resolver.resolveModelFileName('my-custom-model.safetensors');
|
||||
// 返回: 'my-custom-model.safetensors' (如果存在)
|
||||
|
||||
// 变体查找
|
||||
const devModels = getModelsByVariant('dev');
|
||||
console.log(devModels.slice(0, 3));
|
||||
// 输出: ['flux1-dev.safetensors', 'flux1-dev-fp8.safetensors', ...]
|
||||
```
|
||||
|
||||
## 错误处理
|
||||
|
||||
### 错误类型层次
|
||||
|
||||
```plaintext
|
||||
// packages/model-runtime/src/providers/comfyui/errors/
|
||||
ComfyUIInternalError // 基础错误
|
||||
├── ModelResolverError // 模型解析错误
|
||||
├── WorkflowError // 工作流错误
|
||||
├── ServicesError // 服务错误
|
||||
└── UtilsError // 工具错误
|
||||
```
|
||||
|
||||
### 错误处理示例
|
||||
|
||||
```typescript
|
||||
import { ModelResolverError, WorkflowError } from '../errors';
|
||||
|
||||
try {
|
||||
const result = await comfyUI.createImage({
|
||||
model: 'nonexistent-model',
|
||||
params: { prompt: '测试' }
|
||||
});
|
||||
} catch (error) {
|
||||
if (error instanceof ModelResolverError) {
|
||||
console.log('模型解析失败:', error.message);
|
||||
console.log('错误原因:', error.reason);
|
||||
console.log('错误详情:', error.details);
|
||||
} else if (error instanceof WorkflowError) {
|
||||
console.log('工作流错误:', error.message);
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 统一错误解析器
|
||||
|
||||
客户端和服务器端错误处理可使用共享的错误解析器：
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/utils/comfyuiErrorParser.ts
|
||||
import { parseComfyUIErrorMessage, cleanComfyUIErrorMessage } from '../utils/comfyuiErrorParser';
|
||||
|
||||
// 解析错误消息并确定错误类型
|
||||
const { error, errorType } = parseComfyUIErrorMessage(rawError);
|
||||
|
||||
// 清理 ComfyUI 格式的错误消息
|
||||
const cleanMessage = cleanComfyUIErrorMessage(errorMessage);
|
||||
```
|
||||
|
||||
错误解析器处理：
|
||||
|
||||
- HTTP 状态码映射到错误类型
|
||||
- 服务器端错误增强
|
||||
- 模型文件缺失检测
|
||||
- 网络错误识别
|
||||
- 工作流验证错误
|
||||
|
||||
## 测试架构与开发
|
||||
|
||||
### 测试架构概述
|
||||
|
||||
ComfyUI 集成使用了统一的测试架构，确保测试的可维护性和定制友好性。该架构包括：
|
||||
|
||||
- **统一 Mock 系统**：集中管理所有外部依赖的模拟
|
||||
- **参数化测试**：自动适应新模型，无需修改现有测试
|
||||
- **夹具系统**：从配置文件中获取测试数据，确保准确性
|
||||
- **覆盖率目标**：ComfyUI 模块维持 97%+ 覆盖率
|
||||
|
||||
### 测试文件结构
|
||||
|
||||
```plaintext
|
||||
packages/model-runtime/src/providers/comfyui/__tests__/
|
||||
├── setup/
|
||||
│ └── unifiedMocks.ts # 统一Mock配置
|
||||
├── fixtures/
|
||||
│ ├── parameters.fixture.ts # 参数测试夹具
|
||||
│ └── workflow.fixture.ts # 工作流测试夹具
|
||||
├── integration/
|
||||
│ ├── parameterMapping.test.ts # 参数映射集成测试
|
||||
│ └── workflowBuilder.test.ts # 工作流构建测试
|
||||
├── services/ # 各服务单元测试
|
||||
└── workflows/ # 工作流单元测试
|
||||
```
|
||||
|
||||
### 添加新模型测试
|
||||
|
||||
当添加新模型时，测试会自动识别并运行相应的参数映射测试。你只需要：
|
||||
|
||||
#### 1. 在模型配置中添加参数架构
|
||||
|
||||
```typescript
|
||||
// packages/model-bank/src/aiModels/comfyui.ts
|
||||
export const myNewModelParamsSchema = {
|
||||
prompt: { type: 'string', required: true },
|
||||
steps: { type: 'number', default: 20, min: 1, max: 150 },
|
||||
cfg: { type: 'number', default: 7.0, min: 1.0, max: 30.0 }
|
||||
};
|
||||
```
|
||||
|
||||
#### 2. 创建工作流构建器
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/workflows/myNewModel.ts
|
||||
export async function buildMyNewModelWorkflow(
|
||||
modelName: string,
|
||||
params: MyNewModelParams,
|
||||
context: ComfyUIContext
|
||||
) {
|
||||
const workflow = { /* 工作流定义 */ };
|
||||
|
||||
// 参数注入
|
||||
workflow['1'].inputs.prompt = params.prompt;
|
||||
workflow['2'].inputs.steps = params.steps;
|
||||
|
||||
return workflow;
|
||||
}
|
||||
```
|
||||
|
||||
#### 3. 在夹具中注册模型
|
||||
|
||||
```typescript
|
||||
// packages/model-runtime/src/providers/comfyui/__tests__/fixtures/parameters.fixture.ts
|
||||
import {
|
||||
myNewModelParamsSchema,
|
||||
// ... 其他架构
|
||||
} from '../../../../../model-bank/src/aiModels/comfyui';
|
||||
|
||||
export const parametersFixture = {
|
||||
models: {
|
||||
'my-new-model': {
|
||||
schema: myNewModelParamsSchema,
|
||||
defaults: {
|
||||
steps: myNewModelParamsSchema.steps.default,
|
||||
cfg: myNewModelParamsSchema.cfg.default,
|
||||
},
|
||||
boundaries: {
|
||||
min: { steps: myNewModelParamsSchema.steps.min },
|
||||
max: { steps: myNewModelParamsSchema.steps.max }
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
```
|
||||
|
||||
### 测试最佳实践
|
||||
|
||||
#### 使用统一 Mock 系统
|
||||
|
||||
```typescript
|
||||
import { setupAllMocks } from '../setup/unifiedMocks';
|
||||
|
||||
describe('MyTest', () => {
|
||||
const mocks = setupAllMocks();
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
});
|
||||
```
|
||||
|
||||
#### 编写参数映射测试
|
||||
|
||||
参数映射测试会自动运行，验证前端参数正确注入到工作流中：
|
||||
|
||||
```typescript
|
||||
// 测试会自动包含新注册的模型
|
||||
describe.each(
|
||||
Object.entries(models).filter(([name]) => workflowBuilders[name])
|
||||
)(
|
||||
'%s parameter mapping',
|
||||
(modelName, modelConfig) => {
|
||||
it('should map schema parameters to workflow', async () => {
|
||||
const params = {
|
||||
prompt: 'test prompt',
|
||||
...modelConfig.defaults,
|
||||
};
|
||||
|
||||
const workflow = await builder(`${modelName}.safetensors`, params, mockContext);
|
||||
expect(workflow).toBeDefined();
|
||||
});
|
||||
}
|
||||
);
|
||||
```
|
||||
|
||||
#### 定制友好的测试原则
|
||||
|
||||
- **不测试工作流结构**：工作流是 ComfyUI 官方格式，只测试参数映射
|
||||
- **使用配置驱动的数据**：测试数据来自模型配置文件，确保一致性
|
||||
- **避免脆性断言**：不检查具体的节点 ID 或内部结构
|
||||
- **支持扩展**：新增模型应该只影响覆盖率，不破坏现有测试
|
||||
|
||||
### 运行测试
|
||||
|
||||
```bash
|
||||
# 运行 ComfyUI 相关测试
|
||||
cd packages/model-runtime
|
||||
bunx vitest run --silent='passed-only' 'src/comfyui'
|
||||
|
||||
# 查看覆盖率
|
||||
bunx vitest run --coverage 'src/comfyui'
|
||||
|
||||
# 运行特定测试文件
|
||||
bunx vitest run 'src/comfyui/__tests__/integration/parameterMapping.test.ts'
|
||||
```
|
||||
|
||||
### 覆盖率目标
|
||||
|
||||
- **整体覆盖率**：ComfyUI 模块维持 97%+ 覆盖率
|
||||
- **核心功能**：100% 分支覆盖率
|
||||
- **新增功能**：保持或提升现有覆盖率水平
|
||||
|
||||
## 开发和测试
|
||||
|
||||
### 1. 本地开发设置
|
||||
|
||||
```bash
|
||||
# 启动ComfyUI调试模式
|
||||
DEBUG=lobe-image:* pnpm dev
|
||||
```
|
||||
|
||||
### 2. 测试新功能
|
||||
|
||||
```typescript
|
||||
// 创建测试文件
|
||||
import { buildYourCustomWorkflow } from './your-workflow';
|
||||
|
||||
describe('Custom Workflow', () => {
|
||||
test('should build workflow correctly', async () => {
|
||||
const mockContext = {
|
||||
clientService: mockClientService,
|
||||
modelResolverService: mockModelResolver,
|
||||
};
|
||||
|
||||
const workflow = await buildYourCustomWorkflow(
|
||||
'test-model.safetensors',
|
||||
{ prompt: '测试', width: 512, height: 512 },
|
||||
mockContext
|
||||
);
|
||||
|
||||
expect(workflow).toBeDefined();
|
||||
// 验证工作流结构...
|
||||
});
|
||||
});
|
||||
```
|
||||
|
||||
### 3. 模型配置测试
|
||||
|
||||
```typescript
|
||||
import { getModelConfig, getAllModelNames } from '../config/modelRegistry';
|
||||
|
||||
describe('Model Registry', () => {
|
||||
test('should find new model', () => {
|
||||
const config = getModelConfig('your-new-model.safetensors');
|
||||
expect(config).toBeDefined();
|
||||
expect(config?.variant).toBe('dev');
|
||||
expect(config?.modelFamily).toBe('FLUX');
|
||||
});
|
||||
});
|
||||
```
|
||||
|
||||
## 完整使用示例
|
||||
|
||||
### 基础图像生成
|
||||
|
||||
```typescript
|
||||
import { LobeComfyUI } from '@/libs/model-runtime/comfyui';
|
||||
|
||||
const comfyUI = new LobeComfyUI({
|
||||
baseURL: 'http://localhost:8000',
|
||||
authType: 'none'
|
||||
});
|
||||
|
||||
// FLUX Dev模型生成
|
||||
const result = await comfyUI.createImage({
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: '美丽的风景画，高质量，详细',
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
steps: 20,
|
||||
cfg: 3.5,
|
||||
seed: -1
|
||||
}
|
||||
});
|
||||
|
||||
console.log('生成图像URL:', result.imageUrl);
|
||||
```
|
||||
|
||||
### SD3.5 模型使用
|
||||
|
||||
```typescript
|
||||
// SD3.5会自动检测可用编码器
|
||||
const sd35Result = await comfyUI.createImage({
|
||||
model: 'stable-diffusion-35',
|
||||
params: {
|
||||
prompt: '未来主义城市景观',
|
||||
width: 1344,
|
||||
height: 768,
|
||||
steps: 28,
|
||||
cfg: 4.5
|
||||
}
|
||||
});
|
||||
```
|
||||
|
||||
### 企业优化模型
|
||||
|
||||
```typescript
|
||||
// 系统会自动选择最佳可用变体（如FP8量化版本）
|
||||
const optimizedResult = await comfyUI.createImage({
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: '专业商务肖像',
|
||||
width: 768,
|
||||
height: 1024,
|
||||
steps: 15 // FP8模型可以用更少步数
|
||||
}
|
||||
});
|
||||
```
|
||||
|
||||
## 注意事项
|
||||
|
||||
- 确保 ComfyUI 服务正常运行并可访问
|
||||
- 检查所有必需的模型文件是否已正确安装
|
||||
- 注意模型文件的命名规范和路径配置
|
||||
- 定期检查和更新工作流配置以支持新功能
|
||||
- 注意不同模型系列的参数差异和兼容性
|
||||
- 添加新模型时，请遵循测试架构指南确保测试完整性
|
||||
- 在提交代码前务必运行相关测试确保覆盖率达标
|
||||
|
||||
通过遵循这些指南，开发者可以有效地在 LobeChat 中使用和扩展 ComfyUI 功能，为用户提供强大的图像生成和处理能力。
|
||||
|
|
@ -651,6 +651,58 @@ If you need to use Azure OpenAI to provide model services, you can refer to the
|
|||
|
||||
The above example disables all models first, then enables `fal-ai/flux/schnell` and `fal-ai/flux-pro/kontext` (displayed as `FLUX.1 Kontext [pro]`).
|
||||
|
||||
## ComfyUI
|
||||
|
||||
### `COMFYUI_BASE_URL`
|
||||
|
||||
- Type: Optional
|
||||
- Description: The base URL address of the ComfyUI service
|
||||
- Default: `http://localhost:8188`
|
||||
- Example: `http://192.168.1.100:8188` or `https://my-comfyui-server.com`
|
||||
|
||||
### `COMFYUI_AUTH_TYPE`
|
||||
|
||||
- Type: Optional
|
||||
- Description: The authentication type for ComfyUI, supporting 4 authentication methods
|
||||
- `none`: No authentication (default)
|
||||
- `basic`: Basic authentication (username + password)
|
||||
- `bearer`: Bearer Token authentication (API key)
|
||||
- `custom`: Custom request header authentication
|
||||
- Default: `none`
|
||||
- Example: `basic`
|
||||
|
||||
### `COMFYUI_API_KEY`
|
||||
|
||||
- Type: Optional
|
||||
- Description: The API key used when the authentication type is `bearer`
|
||||
- Default: -
|
||||
- Example: `sk-xxxxxx...xxxxxx`
|
||||
|
||||
### `COMFYUI_USERNAME`
|
||||
|
||||
- Type: Optional
|
||||
- Description: The username used when the authentication type is `basic`
|
||||
- Default: -
|
||||
- Example: `admin`
|
||||
|
||||
### `COMFYUI_PASSWORD`
|
||||
|
||||
- Type: Optional
|
||||
- Description: The password used when the authentication type is `basic`
|
||||
- Default: -
|
||||
- Example: `password123`
|
||||
|
||||
### `COMFYUI_CUSTOM_HEADERS`
|
||||
|
||||
- Type: Optional
|
||||
- Description: Custom request headers used when the authentication type is `custom`, requires JSON format string
|
||||
- Default: -
|
||||
- Example: `{"X-Auth-Token": "your-token", "X-Custom-Header": "value"}`
|
||||
|
||||
<Callout type={'info'}>
|
||||
ComfyUI supports multiple authentication methods. Please choose the appropriate authentication type and corresponding authentication parameters according to your ComfyUI service configuration. If your ComfyUI service has no authentication set up, you can skip configuring authentication-related environment variables.
|
||||
</Callout>
|
||||
|
||||
## BFL
|
||||
|
||||
### `ENABLED_BFL`
|
||||
|
|
|
|||
|
|
@ -165,6 +165,58 @@ LobeChat 在部署时提供了丰富的模型服务商相关的环境变量，
|
|||
- 默认值：`us-east-1`
|
||||
- 示例：`us-east-1`
|
||||
|
||||
## ComfyUI
|
||||
|
||||
### `COMFYUI_BASE_URL`
|
||||
|
||||
- 类型：可选
|
||||
- 描述：ComfyUI 服务的基础 URL 地址
|
||||
- 默认值：`http://localhost:8000`
|
||||
- 示例：`http://192.168.1.100:8000` 或 `https://my-comfyui-server.com`
|
||||
|
||||
### `COMFYUI_AUTH_TYPE`
|
||||
|
||||
- 类型：可选
|
||||
- 描述：ComfyUI 的认证类型，支持 4 种认证方式
|
||||
- `none`: 无认证（默认）
|
||||
- `basic`: 基础认证（用户名 + 密码）
|
||||
- `bearer`: Bearer Token 认证（API 密钥）
|
||||
- `custom`: 自定义请求头认证
|
||||
- 默认值：`none`
|
||||
- 示例：`basic`
|
||||
|
||||
### `COMFYUI_API_KEY`
|
||||
|
||||
- 类型：可选
|
||||
- 描述：当认证类型为 `bearer` 时使用的 API 密钥
|
||||
- 默认值：-
|
||||
- 示例：`sk-xxxxxx...xxxxxx`
|
||||
|
||||
### `COMFYUI_USERNAME`
|
||||
|
||||
- 类型：可选
|
||||
- 描述：当认证类型为 `basic` 时使用的用户名
|
||||
- 默认值：-
|
||||
- 示例：`admin`
|
||||
|
||||
### `COMFYUI_PASSWORD`
|
||||
|
||||
- 类型：可选
|
||||
- 描述：当认证类型为 `basic` 时使用的密码
|
||||
- 默认值：-
|
||||
- 示例：`password123`
|
||||
|
||||
### `COMFYUI_CUSTOM_HEADERS`
|
||||
|
||||
- 类型：可选
|
||||
- 描述：当认证类型为 `custom` 时使用的自定义请求头，需要使用 JSON 格式字符串
|
||||
- 默认值：-
|
||||
- 示例：`{"X-Auth-Token": "your-token", "X-Custom-Header": "value"}`
|
||||
|
||||
<Callout type={'info'}>
|
||||
ComfyUI 支持多种认证方式，请根据您的 ComfyUI 服务配置选择合适的认证类型和相应的认证参数。如果您的 ComfyUI 服务没有设置认证，可以不配置认证相关的环境变量。
|
||||
</Callout>
|
||||
|
||||
## DeepSeek AI
|
||||
|
||||
### `DEEPSEEK_PROXY_URL`
|
||||
|
|
|
|||
816
docs/usage/providers/comfyui.mdx
Normal file
816
docs/usage/providers/comfyui.mdx
Normal file
|
|
@ -0,0 +1,816 @@
|
|||
---
|
||||
title: Using ComfyUI for Image Generation in LobeChat
|
||||
description: Learn how to configure and use ComfyUI service in LobeChat, supporting FLUX series models for high-quality image generation and editing features
|
||||
tags:
|
||||
- ComfyUI
|
||||
- FLUX
|
||||
- Text-to-Image
|
||||
- Image Editing
|
||||
- AI Image Generation
|
||||
---
|
||||
|
||||
# Using ComfyUI in LobeChat
|
||||
|
||||
<Image alt={'Using ComfyUI in LobeChat'} cover src={'https://github.com/lobehub/lobe-chat/assets/17870709/c9e5eafc-ca22-496b-a88d-cc0ae53bf720'} />
|
||||
|
||||
This documentation will guide you on how to use [ComfyUI](https://github.com/comfyanonymous/ComfyUI) in LobeChat for high-quality AI image generation and editing.
|
||||
|
||||
## ComfyUI Overview
|
||||
|
||||
ComfyUI is a powerful stable diffusion and flow diffusion GUI that provides a node-based workflow interface. LobeChat integrates with ComfyUI, supporting complete FLUX series models, including text-to-image generation and image editing capabilities.
|
||||
|
||||
### Key Features
|
||||
|
||||
- **Extensive Model Support**: Supports 223 models, including FLUX series (130) and SD series (93)
|
||||
- **Configuration-Driven Architecture**: Registry system provides intelligent model selection
|
||||
- **Multi-Format Support**: Supports .safetensors and .gguf formats with various quantization levels
|
||||
- **Dynamic Precision Selection**: Supports default, fp8\_e4m3fn, fp8\_e5m2, fp8\_e4m3fn\_fast precision
|
||||
- **Multiple Authentication Methods**: Supports no authentication, basic authentication, Bearer Token, and custom authentication
|
||||
- **Intelligent Component Selection**: Automatically selects optimal T5, CLIP, VAE encoder combinations
|
||||
- **Enterprise-Grade Optimization**: Includes NF4, SVDQuant, TorchAO, MFLUX optimization variants
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Step 1: Configure ComfyUI in LobeChat
|
||||
|
||||
#### 1. Open Settings Interface
|
||||
|
||||
- Access LobeChat's `Settings` interface
|
||||
- Find the `ComfyUI` setting item under `AI Providers`
|
||||
|
||||
<Image alt={'ComfyUI Settings Interface'} inStep src={'https://github.com/lobehub/lobe-chat/assets/17870709/3f31bc33-509f-4ad2-ba81-280c2a6ec5fa'} />
|
||||
|
||||
#### 2. Configure Connection Parameters
|
||||
|
||||
**Basic Configuration**:
|
||||
|
||||
- **Server Address**: Enter ComfyUI server address, e.g., `http://localhost:8188`
|
||||
- **Authentication Type**: Select appropriate authentication method (default: no authentication)
|
||||
|
||||
### Step 2: Select Model and Start Generating Images
|
||||
|
||||
#### 1. Select FLUX Model
|
||||
|
||||
In the conversation interface:
|
||||
|
||||
- Click the model selection button
|
||||
- Select the desired FLUX model from the ComfyUI category
|
||||
|
||||
<Image alt={'Select FLUX Model'} inStep src={'https://github.com/lobehub/lobe-chat/assets/17870709/ff7ebacf-27f0-42d7-810b-00314499a084'} />
|
||||
|
||||
#### 2. Text-to-Image Generation
|
||||
|
||||
**Using FLUX Schnell (Fast Generation)**:
|
||||
|
||||
```plaintext
|
||||
Generate an image: A cute orange cat sitting on a sunny windowsill, warm lighting, detailed fur texture
|
||||
```
|
||||
|
||||
**Using FLUX Dev (High Quality Generation)**:
|
||||
|
||||
```plaintext
|
||||
Generate high quality image: City skyline at sunset, cyberpunk style, neon lights, 4K high resolution, detailed architecture
|
||||
```
|
||||
|
||||
#### 3. Image Editing
|
||||
|
||||
**Using FLUX Kontext-dev for Image Editing**:
|
||||
|
||||
```plaintext
|
||||
Edit this image: Change the background to a starry night sky, keep the main subject, cosmic atmosphere
|
||||
```
|
||||
|
||||
Then upload the original image you want to edit.
|
||||
|
||||
<Callout type={'info'}>
|
||||
Image editing functionality requires uploading the original image first, then describing the modifications you want to make.
|
||||
</Callout>
|
||||
|
||||
## Authentication Configuration Guide
|
||||
|
||||
ComfyUI supports four authentication methods. Choose the appropriate method based on your server configuration and security requirements:
|
||||
|
||||
### No Authentication (none)
|
||||
|
||||
**Use Cases**:
|
||||
|
||||
- Local development environment (localhost)
|
||||
- Internal network with trusted users
|
||||
- Personal single-machine deployment
|
||||
|
||||
**Configuration**:
|
||||
|
||||
```yaml
|
||||
Authentication Type: None
|
||||
Server Address: http://localhost:8188
|
||||
```
|
||||
|
||||
### Basic Authentication (basic)
|
||||
|
||||
**Use Cases**:
|
||||
|
||||
- Deployments using Nginx reverse proxy
|
||||
- Team internal use requiring basic access control
|
||||
|
||||
**Configuration**:
|
||||
|
||||
1. **Create User Password**:
|
||||
|
||||
```bash
|
||||
# Install apache2-utils
|
||||
sudo apt-get install apache2-utils
|
||||
|
||||
# Create user 'admin'
|
||||
sudo htpasswd -c /etc/nginx/.htpasswd admin
|
||||
```
|
||||
|
||||
2. **LobeChat Configuration**:
|
||||
|
||||
```yaml
|
||||
Authentication Type: Basic Authentication
|
||||
Server Address: http://your-domain.com
|
||||
Username: admin
|
||||
Password: your_secure_password
|
||||
```
|
||||
|
||||
### Bearer Token (bearer)
|
||||
|
||||
**Use Cases**:
|
||||
|
||||
- API-driven application integration
|
||||
- Enterprise environments requiring Token authentication
|
||||
|
||||
**Generate Token**:
|
||||
|
||||
```python
|
||||
import jwt
|
||||
import datetime
|
||||
|
||||
payload = {
|
||||
'user': 'admin',
|
||||
'exp': datetime.datetime.utcnow() + datetime.timedelta(days=30)
|
||||
}
|
||||
|
||||
secret_key = "your-secret-key"
|
||||
token = jwt.encode(payload, secret_key, algorithm='HS256')
|
||||
print(f"Bearer Token: {token}")
|
||||
```
|
||||
|
||||
**LobeChat Configuration**:
|
||||
|
||||
```yaml
|
||||
Authentication Type: Bearer Token
|
||||
Server Address: http://your-server:8188
|
||||
API Key: eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...
|
||||
```
|
||||
|
||||
### Custom Authentication (custom)
|
||||
|
||||
**Use Cases**:
|
||||
|
||||
- Integration with existing enterprise authentication systems
|
||||
- Systems requiring multiple authentication headers
|
||||
|
||||
**LobeChat Configuration**:
|
||||
|
||||
```yaml
|
||||
Authentication Type: Custom
|
||||
Server Address: http://your-server:8188
|
||||
Custom Headers:
|
||||
{
|
||||
"X-API-Key": "your_api_key",
|
||||
"X-Client-ID": "lobechat"
|
||||
}
|
||||
```
|
||||
|
||||
## Common Issues Resolution
|
||||
|
||||
### 1. How to Install Comfy-Manager
|
||||
|
||||
Comfy-Manager is ComfyUI's extension manager that allows you to easily install and manage various nodes, models, and extensions.
|
||||
|
||||
<details>
|
||||
<summary><b>📦 Install Comfy-Manager Steps</b></summary>
|
||||
|
||||
#### Method 1: Manual Installation (Recommended)
|
||||
|
||||
```bash
|
||||
# Navigate to ComfyUI's custom_nodes directory
|
||||
cd ComfyUI/custom_nodes
|
||||
|
||||
# Clone Comfy-Manager repository
|
||||
git clone https://github.com/ltdrdata/ComfyUI-Manager.git
|
||||
|
||||
# Restart ComfyUI server
|
||||
# After restart, you'll see the Manager button in the UI
|
||||
```
|
||||
|
||||
#### Method 2: One-Click Installation Script
|
||||
|
||||
```bash
|
||||
# Execute in ComfyUI root directory
|
||||
curl -fsSL https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/install.sh | bash
|
||||
```
|
||||
|
||||
#### Verify Installation
|
||||
|
||||
1. Restart ComfyUI server
|
||||
2. Visit `http://localhost:8188`
|
||||
3. You should see the "Manager" button in the bottom-right corner
|
||||
|
||||
#### Using Comfy-Manager
|
||||
|
||||
**Install Models**:
|
||||
|
||||
1. Click "Manager" button
|
||||
2. Select "Install Models"
|
||||
3. Search for needed models (e.g., FLUX, SD3.5)
|
||||
4. Click "Install" to automatically download to correct directory
|
||||
|
||||
**Install Node Extensions**:
|
||||
|
||||
1. Click "Manager" button
|
||||
2. Select "Install Custom Nodes"
|
||||
3. Search for needed nodes (e.g., ControlNet, AnimateDiff)
|
||||
4. Click "Install" and restart server
|
||||
|
||||
**Manage Installed Content**:
|
||||
|
||||
1. Click "Manager" button
|
||||
2. Select "Installed" to view installed extensions
|
||||
3. Update, disable, or uninstall extensions
|
||||
</details>
|
||||
|
||||
### 2. How to Handle "Model not found" Errors
|
||||
|
||||
When you see errors like `Model not found: flux1-dev.safetensors, flux1-krea-dev.safetensors, flux1-schnell.safetensors`, it means the required model files are missing from the server.
|
||||
|
||||
<details>
|
||||
<summary><b>🔧 Resolve Model not found Errors</b></summary>
|
||||
|
||||
#### Error Example
|
||||
|
||||
```plaintext
|
||||
Model not found: flux1-dev.safetensors, flux1-krea-dev.safetensors, flux1-schnell.safetensors
|
||||
```
|
||||
|
||||
This error indicates the system expects to find these model files but couldn't locate them on the server.
|
||||
|
||||
#### Resolution Methods
|
||||
|
||||
**Method 1: Download using Comfy-Manager (Recommended)**
|
||||
|
||||
1. Open ComfyUI interface
|
||||
2. Click "Manager" → "Install Models"
|
||||
3. Search for the model name from the error (e.g., "flux1-dev")
|
||||
4. Click "Install" to automatically download
|
||||
|
||||
**Method 2: Manual Model Download**
|
||||
|
||||
1. **Download Model Files**:
|
||||
- Visit [Hugging Face](https://huggingface.co/black-forest-labs/FLUX.1-dev) or other model sources
|
||||
- Download the files mentioned in the error (e.g., `flux1-dev.safetensors`)
|
||||
|
||||
2. **Place in Correct Directory**:
|
||||
```bash
|
||||
# FLUX and SD3.5 main models go to
|
||||
ComfyUI/models/diffusion_models/flux1-dev.safetensors
|
||||
|
||||
# SD1.5 and SDXL models go to
|
||||
ComfyUI/models/checkpoints/
|
||||
```
|
||||
|
||||
3. **Verify Files**:
|
||||
```bash
|
||||
# Check if file exists
|
||||
ls -la ComfyUI/models/diffusion_models/flux1-dev.safetensors
|
||||
|
||||
# Check file integrity (optional)
|
||||
sha256sum flux1-dev.safetensors
|
||||
```
|
||||
|
||||
4. **Restart ComfyUI Server**
|
||||
|
||||
**Method 3: Direct Download with wget/curl**
|
||||
|
||||
```bash
|
||||
# Navigate to models directory
|
||||
cd ComfyUI/models/diffusion_models/
|
||||
|
||||
# Download using wget (replace with actual download link)
|
||||
wget https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors
|
||||
|
||||
# Or use curl
|
||||
curl -L -o flux1-dev.safetensors https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors
|
||||
```
|
||||
|
||||
#### Common Model Download Sources
|
||||
|
||||
- **Hugging Face**: [https://huggingface.co/models](https://huggingface.co/models)
|
||||
- **Civitai**: [https://civitai.com/models](https://civitai.com/models)
|
||||
- **Official Sources**:
|
||||
- FLUX: [https://huggingface.co/black-forest-labs](https://huggingface.co/black-forest-labs)
|
||||
- SD3.5: [https://huggingface.co/stabilityai](https://huggingface.co/stabilityai)
|
||||
|
||||
#### Prevention Measures
|
||||
|
||||
1. **Basic Model Package**: Download at least one base model
|
||||
- FLUX: `flux1-schnell.safetensors` (fast) or `flux1-dev.safetensors` (high quality)
|
||||
- SD3.5: `sd3.5_large.safetensors`
|
||||
|
||||
2. **Check Disk Space**:
|
||||
```bash
|
||||
# Check available space
|
||||
df -h ComfyUI/models/
|
||||
```
|
||||
|
||||
3. **Set Model Path** (optional):
|
||||
If your models are stored elsewhere, create symbolic links:
|
||||
```bash
|
||||
ln -s /path/to/your/models ComfyUI/models/diffusion_models/
|
||||
```
|
||||
</details>
|
||||
|
||||
### 3. How to Handle Missing System Component Errors
|
||||
|
||||
When you see errors like `Missing VAE encoder: ae.safetensors` or other component files missing, you need to download the corresponding system components.
|
||||
|
||||
<details>
|
||||
<summary><b>🛠️ Resolve Missing System Component Errors</b></summary>
|
||||
|
||||
#### Common Component Errors
|
||||
|
||||
```plaintext
|
||||
Missing VAE encoder: ae.safetensors. Please download and place it in the models/vae folder.
|
||||
Missing CLIP encoder: clip_l.safetensors. Please download and place it in the models/clip folder.
|
||||
Missing T5 encoder: t5xxl_fp16.safetensors. Please download and place it in the models/clip folder.
|
||||
```
|
||||
|
||||
#### Component Types Description
|
||||
|
||||
| Component Type | Example Filename | Purpose | Storage Directory |
|
||||
| -------------- | ------------------------------ | ----------------------- | ------------------ |
|
||||
| **VAE** | ae.safetensors | Image encoding/decoding | models/vae/ |
|
||||
| **CLIP** | clip\_l.safetensors | Text encoding (CLIP) | models/clip/ |
|
||||
| **T5** | t5xxl\_fp16.safetensors | Text encoding (T5) | models/clip/ |
|
||||
| **ControlNet** | flux-controlnet-\*.safetensors | Control networks | models/controlnet/ |
|
||||
|
||||
#### Resolution Methods
|
||||
|
||||
**Method 1: Use Comfy-Manager (Recommended)**
|
||||
|
||||
1. Click "Manager" → "Install Models"
|
||||
2. Select component type in "Filter" (VAE/CLIP/T5)
|
||||
3. Download corresponding component files
|
||||
|
||||
**Method 2: Manual Component Download**
|
||||
|
||||
##### FLUX Required Components
|
||||
|
||||
```bash
|
||||
# 1. VAE Encoder
|
||||
cd ComfyUI/models/vae/
|
||||
wget https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/ae.safetensors
|
||||
|
||||
# 2. CLIP-L Encoder
|
||||
cd ComfyUI/models/clip/
|
||||
wget https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors
|
||||
|
||||
# 3. T5-XXL Encoder (choose different precisions)
|
||||
# FP16 version (recommended, balanced performance)
|
||||
wget https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp16.safetensors
|
||||
|
||||
# Or FP8 version (saves VRAM)
|
||||
wget https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors
|
||||
```
|
||||
|
||||
##### SD3.5 Required Components
|
||||
|
||||
```bash
|
||||
# SD3.5 uses different encoders
|
||||
cd ComfyUI/models/clip/
|
||||
|
||||
# CLIP-G Encoder
|
||||
wget https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/text_encoders/clip_g.safetensors
|
||||
|
||||
# CLIP-L Encoder
|
||||
wget https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/text_encoders/clip_l.safetensors
|
||||
|
||||
# T5-XXL Encoder
|
||||
wget https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/text_encoders/t5xxl_fp16.safetensors
|
||||
```
|
||||
|
||||
##### SDXL Required Components
|
||||
|
||||
```bash
|
||||
# SDXL VAE
|
||||
cd ComfyUI/models/vae/
|
||||
wget https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors
|
||||
|
||||
# SDXL uses built-in CLIP encoders, usually no separate download needed
|
||||
```
|
||||
|
||||
#### Component Compatibility Matrix
|
||||
|
||||
| Model Series | Required VAE | Required CLIP | Required T5 | Optional Components |
|
||||
| ------------ | -------------- | ------------------- | ----------------------- | ------------------- |
|
||||
| **FLUX** | ae.safetensors | clip\_l.safetensors | t5xxl\_fp16.safetensors | ControlNet |
|
||||
| **SD3.5** | Built-in | clip\_g + clip\_l | t5xxl\_fp16 | - |
|
||||
| **SDXL** | sdxl\_vae | Built-in | - | Refiner |
|
||||
| **SD1.5** | vae-ft-mse | Built-in | - | ControlNet |
|
||||
|
||||
#### Precision Selection Recommendations
|
||||
|
||||
**T5 Encoder Precision Selection**:
|
||||
|
||||
| VRAM Capacity | Recommended Version | Filename |
|
||||
| ------------- | ------------------- | ------------------------------ |
|
||||
| \< 12GB | FP8 Quantized | t5xxl\_fp8\_e4m3fn.safetensors |
|
||||
| 12-16GB | FP16 | t5xxl\_fp16.safetensors |
|
||||
| > 16GB | FP32 | t5xxl.safetensors |
|
||||
|
||||
#### Verify Component Installation
|
||||
|
||||
```bash
|
||||
# Check all required components
|
||||
echo "=== VAE Components ==="
|
||||
ls -la ComfyUI/models/vae/
|
||||
|
||||
echo "=== CLIP/T5 Components ==="
|
||||
ls -la ComfyUI/models/clip/
|
||||
|
||||
echo "=== ControlNet Components ==="
|
||||
ls -la ComfyUI/models/controlnet/
|
||||
```
|
||||
|
||||
#### Troubleshooting
|
||||
|
||||
**Issue: Still getting errors after download**
|
||||
|
||||
1. **Check File Permissions**:
|
||||
```bash
|
||||
chmod 644 ComfyUI/models/vae/*.safetensors
|
||||
chmod 644 ComfyUI/models/clip/*.safetensors
|
||||
```
|
||||
|
||||
2. **Clear Cache**:
|
||||
```bash
|
||||
# Clear ComfyUI cache
|
||||
rm -rf ComfyUI/temp/*
|
||||
rm -rf ComfyUI/__pycache__/*
|
||||
```
|
||||
|
||||
3. **Restart Server**:
|
||||
```bash
|
||||
# Fully restart ComfyUI
|
||||
pkill -f "python main.py"
|
||||
python main.py --listen 0.0.0.0 --port 8188
|
||||
```
|
||||
|
||||
**Issue: Insufficient VRAM**
|
||||
|
||||
Use quantized component versions:
|
||||
|
||||
- T5: Use `t5xxl_fp8_e4m3fn.safetensors` instead of FP16/FP32
|
||||
- VAE: Some models support FP16 VAE versions
|
||||
|
||||
**Issue: Slow Downloads**
|
||||
|
||||
1. Use mirror sources (if applicable)
|
||||
2. Use download tools (like aria2c) with resume support:
|
||||
```bash
|
||||
aria2c -x 16 -s 16 -k 1M [download_link]
|
||||
```
|
||||
</details>
|
||||
|
||||
## ComfyUI Server Installation
|
||||
|
||||
<details>
|
||||
<summary><b>🚀 Install and Configure ComfyUI Server</b></summary>
|
||||
|
||||
### 1. Install ComfyUI
|
||||
|
||||
```bash
|
||||
# Clone ComfyUI repository
|
||||
git clone https://github.com/comfyanonymous/ComfyUI.git
|
||||
cd ComfyUI
|
||||
|
||||
# Install dependencies
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Optional: Install JWT support (for Token authentication)
|
||||
pip install PyJWT
|
||||
|
||||
# Start ComfyUI server
|
||||
python main.py --listen 0.0.0.0 --port 8188
|
||||
```
|
||||
|
||||
### 2. Download Model Files
|
||||
|
||||
**Recommended Basic Configuration** (Minimal installation):
|
||||
|
||||
**Main Models** (place in `models/diffusion_models/` directory):
|
||||
|
||||
- `flux1-schnell.safetensors` - Fast generation (4 steps)
|
||||
- `flux1-dev.safetensors` - High-quality creation (20 steps)
|
||||
|
||||
**Required Components** (place in respective directories):
|
||||
|
||||
- `models/vae/ae.safetensors` - VAE encoder
|
||||
- `models/clip/clip_l.safetensors` - CLIP text encoder
|
||||
- `models/clip/t5xxl_fp16.safetensors` - T5 text encoder
|
||||
|
||||
### 3. Verify Server Running
|
||||
|
||||
Visit `http://localhost:8188` to confirm ComfyUI interface loads properly.
|
||||
|
||||
<Callout type={'info'}>
|
||||
**Smart Model Selection**: LobeChat will automatically select the best model based on available model files on the server. You don't need to download all models; the system will automatically choose from available models by priority (Official > Enterprise > Community).
|
||||
</Callout>
|
||||
</details>
|
||||
|
||||
## Supported Models
|
||||
|
||||
LobeChat's ComfyUI integration uses a configuration-driven architecture, supporting **223 models**, providing complete coverage from official models to community-optimized versions.
|
||||
|
||||
### FLUX Series Recommended Parameters
|
||||
|
||||
| Model Type | Recommended Steps | CFG Scale | Resolution Range |
|
||||
| ----------- | ----------------- | --------- | -------------------- |
|
||||
| **Schnell** | 4 steps | - | 512×512 to 1536×1536 |
|
||||
| **Dev** | 20 steps | 3.5 | 512×512 to 2048×2048 |
|
||||
| **Kontext** | 20 steps | 3.5 | 512×512 to 2048×2048 |
|
||||
| **Krea** | 20 steps | 4.5 | 512×512 to 2048×2048 |
|
||||
|
||||
### SD3.5 Series Parameters
|
||||
|
||||
| Model Type | Recommended Steps | CFG Scale | Resolution Range |
|
||||
| --------------- | ----------------- | --------- | -------------------- |
|
||||
| **Large** | 25 steps | 7.0 | 512×512 to 2048×2048 |
|
||||
| **Large Turbo** | 8 steps | 3.5 | 512×512 to 1536×1536 |
|
||||
| **Medium** | 20 steps | 6.0 | 512×512 to 1536×1536 |
|
||||
|
||||
<details>
|
||||
<summary><b>📋 Complete Supported Model List</b></summary>
|
||||
|
||||
### Model Classification System
|
||||
|
||||
#### Priority 1: Official Core Models
|
||||
|
||||
**FLUX.1 Official Series**:
|
||||
|
||||
- `flux1-dev.safetensors` - High-quality creation model
|
||||
- `flux1-schnell.safetensors` - Fast generation model
|
||||
- `flux1-kontext-dev.safetensors` - Image editing model
|
||||
- `flux1-krea-dev.safetensors` - Safety-enhanced model
|
||||
|
||||
**SD3.5 Official Series**:
|
||||
|
||||
- `sd3.5_large.safetensors` - SD3.5 large base model
|
||||
- `sd3.5_large_turbo.safetensors` - Fast generation version
|
||||
- `sd3.5_medium.safetensors` - Medium-scale model
|
||||
|
||||
#### Priority 2: Enterprise Optimized Models (106 FLUX)
|
||||
|
||||
**Quantization Optimization Series**:
|
||||
|
||||
- **GGUF Quantization**: Each variant supports 11 quantization levels (F16, Q8\_0, Q6\_K, Q5\_K\_M, Q5\_K\_S, Q4\_K\_M, Q4\_K\_S, Q4\_0, Q3\_K\_M, Q3\_K\_S, Q2\_K)
|
||||
- **FP8 Precision**: fp8\_e4m3fn, fp8\_e5m2 optimized versions
|
||||
- **Enterprise Lightweight**: FLUX.1-lite-8B series
|
||||
- **Technical Experiments**: NF4, SVDQuant, TorchAO, optimum-quanto, MFLUX optimized versions
|
||||
|
||||
#### Priority 3: Community Fine-tuned Models (48 FLUX)
|
||||
|
||||
**Community Optimization Series**:
|
||||
|
||||
- **Jib Mix Flux** Series: High-quality mixed models
|
||||
- **Real Dream FLUX** Series: Realism style
|
||||
- **Vision Realistic** Series: Visual realism
|
||||
- **PixelWave FLUX** Series: Pixel art optimization
|
||||
- **Fluxmania** Series: Diverse style support
|
||||
|
||||
### SD Series Model Support (93 models)
|
||||
|
||||
**SD3.5 Series**: 5 models
|
||||
**SD1.5 Series**: 37 models (including official, quantized, and community versions)
|
||||
**SDXL Series**: 50 models (including base, Refiner, and Playground models)
|
||||
|
||||
### Workflow Support
|
||||
|
||||
System supports **6 workflows**:
|
||||
|
||||
- **flux-dev**: High-quality creation workflow
|
||||
- **flux-schnell**: Fast generation workflow
|
||||
- **flux-kontext**: Image editing workflow
|
||||
- **sd35**: SD3.5 dedicated workflow
|
||||
- **simple-sd**: Simple SD workflow
|
||||
- **index**: Workflow entry point
|
||||
</details>
|
||||
|
||||
## Performance Optimization Recommendations
|
||||
|
||||
### Hardware Requirements
|
||||
|
||||
**Minimum Configuration** (GGUF quantized models):
|
||||
|
||||
- GPU: 6GB VRAM (using Q4 quantization)
|
||||
- RAM: 12GB
|
||||
- Storage: 30GB available space
|
||||
|
||||
**Recommended Configuration** (standard models):
|
||||
|
||||
- GPU: 12GB+ VRAM (RTX 4070 Ti or higher)
|
||||
- RAM: 24GB+
|
||||
- Storage: SSD 100GB+ available space
|
||||
|
||||
### VRAM Optimization Strategy
|
||||
|
||||
| VRAM Capacity | Recommended Quantization | Model Example | Performance Characteristics |
|
||||
| ------------- | ------------------------ | ---------------------------------- | --------------------------- |
|
||||
| **6-8GB** | Q4\_0, Q4\_K\_S | `flux1-dev-Q4_0.gguf` | Minimal VRAM usage |
|
||||
| **10-12GB** | Q6\_K, Q8\_0 | `flux1-dev-Q6_K.gguf` | Balance performance/quality |
|
||||
| **16GB+** | FP8, FP16 | `flux1-dev-fp8-e4m3fn.safetensors` | Near-original quality |
|
||||
| **24GB+** | Full model | `flux1-dev.safetensors` | Best quality |
|
||||
|
||||
## Custom Model Usage
|
||||
|
||||
<details>
|
||||
<summary><b>🎨 Configure Custom SD Models</b></summary>
|
||||
|
||||
LobeChat supports using custom Stable Diffusion models. The system uses fixed filenames to identify custom models.
|
||||
|
||||
### 1. Model File Preparation
|
||||
|
||||
**Required Files**:
|
||||
|
||||
- **Main Model File**: `custom_sd_lobe.safetensors`
|
||||
- **VAE File (Optional)**: `custom_sd_vae_lobe.safetensors`
|
||||
|
||||
### 2. Add Custom Model
|
||||
|
||||
**Method 1: Rename Existing Model**
|
||||
|
||||
```bash
|
||||
# Rename your model to fixed filename
|
||||
mv your_custom_model.safetensors custom_sd_lobe.safetensors
|
||||
|
||||
# Move to correct directory
|
||||
mv custom_sd_lobe.safetensors ComfyUI/models/diffusion_models/
|
||||
```
|
||||
|
||||
**Method 2: Create Symbolic Link (Recommended)**
|
||||
|
||||
```bash
|
||||
# Create soft link for easy model switching
|
||||
ln -s /path/to/your_model.safetensors ComfyUI/models/diffusion_models/custom_sd_lobe.safetensors
|
||||
```
|
||||
|
||||
### 3. Use Custom Model
|
||||
|
||||
In LobeChat, custom models will appear as:
|
||||
|
||||
- **stable-diffusion-custom**: Standard custom model
|
||||
- **stable-diffusion-custom-refiner**: Refiner custom model
|
||||
|
||||
### Custom Model Parameter Recommendations
|
||||
|
||||
| Parameter | SD 1.5 Models | SDXL Models |
|
||||
| ---------- | ------------- | ----------- |
|
||||
| **steps** | 20-30 | 25-40 |
|
||||
| **cfg** | 7.0 | 6.0-8.0 |
|
||||
| **width** | 512 | 1024 |
|
||||
| **height** | 512 | 1024 |
|
||||
</details>
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Smart Error Diagnosis System
|
||||
|
||||
LobeChat integrates a smart error handling system that can automatically diagnose and provide targeted solutions.
|
||||
|
||||
#### Error Types and Solutions
|
||||
|
||||
| Error Type | User Prompt | Automatic Diagnosis |
|
||||
| ------------------ | ---------------------------------- | --------------------------------------------------- |
|
||||
| **Connection** | "Cannot connect to ComfyUI server" | Auto-detect server status and connectivity |
|
||||
| **Authentication** | "API key invalid or expired" | Auto-verify authentication credentials |
|
||||
| **Permissions** | "Access permissions insufficient" | Auto-check user permissions and file access |
|
||||
| **Model Issues** | "Cannot find specified model file" | Auto-scan available models and suggest alternatives |
|
||||
| **Configuration** | "Configuration file error" | Auto-verify config completeness and syntax |
|
||||
|
||||
<details>
|
||||
<summary><b>🔍 Traditional Troubleshooting Methods</b></summary>
|
||||
|
||||
#### 1. Connection Failure
|
||||
|
||||
**Issue**: Cannot connect to ComfyUI server
|
||||
|
||||
**Solution**:
|
||||
|
||||
```bash
|
||||
# Confirm server running
|
||||
curl http://localhost:8188/system_stats
|
||||
|
||||
# Check port
|
||||
netstat -tulpn | grep 8188
|
||||
```
|
||||
|
||||
#### 2. Out of Memory
|
||||
|
||||
**Issue**: Memory errors during generation
|
||||
|
||||
**Solution**:
|
||||
|
||||
- Lower image resolution
|
||||
- Reduce generation steps
|
||||
- Use quantized models
|
||||
|
||||
#### 3. Authentication Failure
|
||||
|
||||
**Issue**: 401 or 403 errors
|
||||
|
||||
**Solution**:
|
||||
|
||||
- Verify authentication configuration
|
||||
- Check if Token is expired
|
||||
- Confirm user permissions
|
||||
</details>
|
||||
|
||||
## Best Practices
|
||||
|
||||
### Prompt Writing
|
||||
|
||||
1. **Detailed Description**: Provide clear, detailed image descriptions
|
||||
2. **Style Specification**: Clearly specify artistic style, color style, etc.
|
||||
3. **Quality Keywords**: Add "4K", "high quality", "detailed" keywords
|
||||
4. **Avoid Contradictions**: Ensure description content is logically consistent
|
||||
|
||||
**Example**:
|
||||
|
||||
```plaintext
|
||||
A young woman with flowing long hair, wearing an elegant blue dress, standing in a cherry blossom park,
|
||||
sunlight filtering through leaves, warm atmosphere, cinematic lighting, 4K high resolution, detailed, photorealistic
|
||||
```
|
||||
|
||||
### Parameter Optimization
|
||||
|
||||
1. **FLUX Schnell**: Suitable for quick previews, use 4-step generation
|
||||
2. **FLUX Dev**: Balance quality and speed, CFG 3.5, 20 steps
|
||||
3. **FLUX Krea-dev**: Safe creation, CFG 4.5, note content filtering
|
||||
4. **FLUX Kontext-dev**: Image editing, strength 0.6-0.9
|
||||
|
||||
<Callout type={'warning'}>
|
||||
Please note during use:
|
||||
|
||||
- FLUX Dev, Krea-dev, Kontext-dev models are for non-commercial use only
|
||||
- Generated content must comply with relevant laws and platform policies
|
||||
- Large model generation may take considerable time, please be patient
|
||||
</Callout>
|
||||
|
||||
## API Reference
|
||||
|
||||
<details>
|
||||
<summary><b>📚 API Documentation</b></summary>
|
||||
|
||||
### Request Format
|
||||
|
||||
```typescript
|
||||
interface ComfyUIRequest {
|
||||
model: string; // Model ID, e.g., 'flux-schnell'
|
||||
prompt: string; // Text prompt
|
||||
width: number; // Image width
|
||||
height: number; // Image height
|
||||
steps: number; // Generation steps
|
||||
seed: number; // Random seed
|
||||
cfg?: number; // CFG Scale (Dev/Krea/Kontext specific)
|
||||
strength?: number; // Edit strength (Kontext specific)
|
||||
imageUrl?: string; // Input image (Kontext specific)
|
||||
}
|
||||
```
|
||||
|
||||
### Response Format
|
||||
|
||||
```typescript
|
||||
interface ComfyUIResponse {
|
||||
images: Array<{
|
||||
url: string; // Generated image URL
|
||||
filename: string; // Filename
|
||||
subfolder: string; // Subdirectory
|
||||
type: string; // File type
|
||||
}>;
|
||||
prompt_id: string; // Prompt ID
|
||||
}
|
||||
```
|
||||
|
||||
### Error Codes
|
||||
|
||||
| Error Code | Description | Resolution Suggestions |
|
||||
| ---------- | ------------------------ | -------------------------------- |
|
||||
| `400` | Invalid parameters | Check parameter format and range |
|
||||
| `401` | Authentication failed | Verify API key and auth config |
|
||||
| `403` | Insufficient permissions | Check user permissions |
|
||||
| `404` | Model not found | Confirm model file exists |
|
||||
| `500` | Server error | Check ComfyUI logs |
|
||||
</details>
|
||||
|
||||
You can now use ComfyUI in LobeChat for high-quality AI image generation and editing. If you encounter issues, please refer to the troubleshooting section or consult the [ComfyUI official documentation](https://github.com/comfyanonymous/ComfyUI).
|
||||
816
docs/usage/providers/comfyui.zh-CN.mdx
Normal file
816
docs/usage/providers/comfyui.zh-CN.mdx
Normal file
|
|
@ -0,0 +1,816 @@
|
|||
---
|
||||
title: 在 LobeChat 中使用 ComfyUI 生成图像
|
||||
description: 学习如何在 LobeChat 中配置和使用 ComfyUI 服务，支持 FLUX 系列模型的高质量图像生成和编辑功能
|
||||
tags:
|
||||
- ComfyUI
|
||||
- FLUX
|
||||
- 文生图
|
||||
- 图像编辑
|
||||
- AI 图像生成
|
||||
---
|
||||
|
||||
# 在 LobeChat 中使用 ComfyUI
|
||||
|
||||
<Image alt={'在 LobeChat 中使用 ComfyUI'} cover src={'https://github.com/lobehub/lobe-chat/assets/17870709/c9e5eafc-ca22-496b-a88d-cc0ae53bf720'} />
|
||||
|
||||
本文档将指导你如何在 LobeChat 中使用 [ComfyUI](https://github.com/comfyanonymous/ComfyUI) 进行高质量的 AI 图像生成和编辑。
|
||||
|
||||
## ComfyUI 简介
|
||||
|
||||
ComfyUI 是一个功能强大的稳定扩散和流扩散 GUI，提供基于节点的工作流界面。LobeChat 集成了 ComfyUI，支持完整的 FLUX 系列模型，包括文本生成图像和图像编辑功能。
|
||||
|
||||
### 主要特性
|
||||
|
||||
- **广泛模型支持**：支持 223 个模型，包含 FLUX 系列（130 个）和 SD 系列（93 个）
|
||||
- **配置驱动架构**：注册表系统提供智能模型选择
|
||||
- **多格式支持**：支持 .safetensors 和 .gguf 格式，包含多种量化级别
|
||||
- **动态精度选择**：支持 default、fp8\_e4m3fn、fp8\_e5m2、fp8\_e4m3fn\_fast 精度
|
||||
- **多种认证方式**：支持无认证、基本认证、Bearer Token 和自定义认证
|
||||
- **智能组件选择**：自动选择最优的 T5、CLIP、VAE 编码器组合
|
||||
- **企业级优化**：包含 NF4、SVDQuant、TorchAO、MFLUX 等优化变体
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 步骤一：在 LobeChat 中配置 ComfyUI
|
||||
|
||||
#### 1. 打开设置界面
|
||||
|
||||
- 访问 LobeChat 的 `设置` 界面
|
||||
- 在 `AI 服务商` 下找到 `ComfyUI` 的设置项
|
||||
|
||||
<Image alt={'ComfyUI 设置界面'} inStep src={'https://github.com/lobehub/lobe-chat/assets/17870709/3f31bc33-509f-4ad2-ba81-280c2a6ec5fa'} />
|
||||
|
||||
#### 2. 配置连接参数
|
||||
|
||||
**基本配置**：
|
||||
|
||||
- **服务器地址**：输入 ComfyUI 服务器地址，如 `http://localhost:8000`
|
||||
- **认证类型**：选择合适的认证方式（默认无认证）
|
||||
|
||||
### 步骤二：选择模型并开始生成图像
|
||||
|
||||
#### 1. 选择 FLUX 模型
|
||||
|
||||
在对话界面中：
|
||||
|
||||
- 点击模型选择按钮
|
||||
- 从 ComfyUI 分类中选择所需的 FLUX 模型
|
||||
|
||||
<Image alt={'选择 FLUX 模型'} inStep src={'https://github.com/lobehub/lobe-chat/assets/17870709/ff7ebacf-27f0-42d7-810b-00314499a084'} />
|
||||
|
||||
#### 2. 文本生成图像
|
||||
|
||||
**使用 FLUX Schnell（快速生成）**：
|
||||
|
||||
```plaintext
|
||||
Generate an image: A cute orange cat sitting on a sunny windowsill, warm lighting, detailed fur texture
|
||||
```
|
||||
|
||||
**使用 FLUX Dev（高质量生成）**：
|
||||
|
||||
```plaintext
|
||||
Generate high quality image: City skyline at sunset, cyberpunk style, neon lights, 4K high resolution, detailed architecture
|
||||
```
|
||||
|
||||
#### 3. 图像编辑
|
||||
|
||||
**使用 FLUX Kontext-dev 编辑图像**：
|
||||
|
||||
```plaintext
|
||||
Edit this image: Change the background to a starry night sky, keep the main subject, cosmic atmosphere
|
||||
```
|
||||
|
||||
然后上传需要编辑的原始图像。
|
||||
|
||||
<Callout type={'info'}>
|
||||
图像编辑功能需要先上传原始图像，然后描述你希望进行的修改。
|
||||
</Callout>
|
||||
|
||||
## 认证配置指南
|
||||
|
||||
ComfyUI 支持四种认证方式，请根据你的服务器配置和安全需求选择合适的认证方式：
|
||||
|
||||
### 无认证 (none)
|
||||
|
||||
**适用场景**：
|
||||
|
||||
- 本地开发环境（localhost）
|
||||
- 内网环境且信任所有用户
|
||||
- 个人使用的单机部署
|
||||
|
||||
**配置方法**：
|
||||
|
||||
```yaml
|
||||
认证类型：无认证
|
||||
服务器地址：http://localhost:8000
|
||||
```
|
||||
|
||||
### 基本认证 (basic)
|
||||
|
||||
**适用场景**：
|
||||
|
||||
- 使用 Nginx 反向代理的部署
|
||||
- 团队内部使用且需要基础访问控制
|
||||
|
||||
**配置方法**：
|
||||
|
||||
1. **创建用户密码**：
|
||||
|
||||
```bash
|
||||
# 安装 apache2-utils
|
||||
sudo apt-get install apache2-utils
|
||||
|
||||
# 创建用户 'admin'
|
||||
sudo htpasswd -c /etc/nginx/.htpasswd admin
|
||||
```
|
||||
|
||||
2. **LobeChat 配置**：
|
||||
|
||||
```yaml
|
||||
认证类型：基本认证
|
||||
服务器地址：https://your-domain.com
|
||||
用户名：admin
|
||||
密码：your_secure_password
|
||||
```
|
||||
|
||||
### Bearer Token (bearer)
|
||||
|
||||
**适用场景**：
|
||||
|
||||
- API 驱动的应用集成
|
||||
- 需要 Token 认证的企业环境
|
||||
|
||||
**生成 Token**：
|
||||
|
||||
```python
|
||||
import jwt
|
||||
import datetime
|
||||
|
||||
payload = {
|
||||
'user': 'admin',
|
||||
'exp': datetime.datetime.utcnow() + datetime.timedelta(days=30)
|
||||
}
|
||||
|
||||
secret_key = "your-secret-key"
|
||||
token = jwt.encode(payload, secret_key, algorithm='HS256')
|
||||
print(f"Bearer Token: {token}")
|
||||
```
|
||||
|
||||
**LobeChat 配置**：
|
||||
|
||||
```yaml
|
||||
认证类型：Bearer Token
|
||||
服务器地址：https://your-domain.com
|
||||
API 密钥：example-eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...
|
||||
```
|
||||
|
||||
### 自定义认证 (custom)
|
||||
|
||||
**适用场景**：
|
||||
|
||||
- 集成现有企业认证系统
|
||||
- 需要多重认证头的系统
|
||||
|
||||
**LobeChat 配置**：
|
||||
|
||||
```yaml
|
||||
认证类型：自定义
|
||||
服务器地址：https://your-domain.com
|
||||
自定义请求头：
|
||||
{
|
||||
"X-API-Key": "your_api_key",
|
||||
"X-Client-ID": "lobechat"
|
||||
}
|
||||
```
|
||||
|
||||
## 常见问题处理
|
||||
|
||||
### 1. 如何安装 Comfy-Manager
|
||||
|
||||
Comfy-Manager 是 ComfyUI 的扩展管理器，让你能够轻松安装和管理各种节点、模型和扩展。
|
||||
|
||||
<details>
|
||||
<summary><b>📦 安装 Comfy-Manager 步骤</b></summary>
|
||||
|
||||
#### 方法一：手动安装（推荐）
|
||||
|
||||
```bash
|
||||
# 进入 ComfyUI 的 custom_nodes 目录
|
||||
cd ComfyUI/custom_nodes
|
||||
|
||||
# 克隆 Comfy-Manager 仓库
|
||||
git clone https://github.com/ltdrdata/ComfyUI-Manager.git
|
||||
|
||||
# 重启 ComfyUI 服务器
|
||||
# 重新启动后，你会在 UI 中看到 Manager 按钮
|
||||
```
|
||||
|
||||
#### 方法二：使用一键安装脚本
|
||||
|
||||
```bash
|
||||
# 在 ComfyUI 根目录下执行
|
||||
curl -fsSL https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/install.sh | bash
|
||||
```
|
||||
|
||||
#### 验证安装
|
||||
|
||||
1. 重启 ComfyUI 服务器
|
||||
2. 访问 `http://localhost:8000`
|
||||
3. 你应该能在界面右下角看到 "Manager" 按钮
|
||||
|
||||
#### 使用 Comfy-Manager
|
||||
|
||||
**安装模型**：
|
||||
|
||||
1. 点击 "Manager" 按钮
|
||||
2. 选择 "Install Models"
|
||||
3. 搜索需要的模型（如 FLUX、SD3.5）
|
||||
4. 点击 "Install" 自动下载到正确目录
|
||||
|
||||
**安装节点扩展**：
|
||||
|
||||
1. 点击 "Manager" 按钮
|
||||
2. 选择 "Install Custom Nodes"
|
||||
3. 搜索需要的节点（如 ControlNet、AnimateDiff）
|
||||
4. 点击 "Install" 并重启服务器
|
||||
|
||||
**管理已安装内容**：
|
||||
|
||||
1. 点击 "Manager" 按钮
|
||||
2. 选择 "Installed" 查看已安装的扩展
|
||||
3. 可以更新、禁用或卸载扩展
|
||||
</details>
|
||||
|
||||
### 2. 如何处理 "Model not found" 错误
|
||||
|
||||
当你看到类似 `Model not found: flux1-dev.safetensors, please install one first.` 的错误时，说明服务器上缺少所需的模型文件。
|
||||
|
||||
<details>
|
||||
<summary><b>🔧 解决 Model not found 错误</b></summary>
|
||||
|
||||
#### 错误示例
|
||||
|
||||
```plaintext
|
||||
Model not found: flux1-dev.safetensors, please install one first.
|
||||
```
|
||||
|
||||
这个错误表示系统期望找到 `flux1-dev.safetensors` 模型文件，但在服务器上没有找到。
|
||||
|
||||
#### 解决方法
|
||||
|
||||
**方法一：使用 Comfy-Manager 下载（推荐）**
|
||||
|
||||
1. 打开 ComfyUI 界面
|
||||
2. 点击 "Manager" → "Install Models"
|
||||
3. 搜索错误提示中的模型名（如 "flux1-dev"）
|
||||
4. 点击 "Install" 自动下载
|
||||
|
||||
**方法二：手动下载模型**
|
||||
|
||||
1. **下载模型文件**：
|
||||
- 访问 [Hugging Face](https://huggingface.co/black-forest-labs/FLUX.1-dev) 或其他模型源
|
||||
- 下载错误提示中的文件（如 `flux1-dev.safetensors`）
|
||||
|
||||
2. **放置到正确目录**：
|
||||
```bash
|
||||
# FLUX 和 SD3.5 主模型放入
|
||||
ComfyUI/models/diffusion_models/flux1-dev.safetensors
|
||||
|
||||
# SD1.5 和 SDXL 模型放入
|
||||
ComfyUI/models/checkpoints/
|
||||
```
|
||||
|
||||
3. **验证文件**：
|
||||
```bash
|
||||
# 检查文件是否存在
|
||||
ls -la ComfyUI/models/diffusion_models/flux1-dev.safetensors
|
||||
|
||||
# 检查文件完整性（可选）
|
||||
sha256sum flux1-dev.safetensors
|
||||
```
|
||||
|
||||
4. **重启 ComfyUI 服务器**
|
||||
|
||||
**方法三：使用 wget/curl 直接下载**
|
||||
|
||||
```bash
|
||||
# 进入模型目录
|
||||
cd ComfyUI/models/diffusion_models/
|
||||
|
||||
# 使用 wget 下载（替换为实际下载链接）
|
||||
wget https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors
|
||||
|
||||
# 或使用 curl
|
||||
curl -L -o flux1-dev.safetensors https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors
|
||||
```
|
||||
|
||||
#### 常见模型下载源
|
||||
|
||||
- **Hugging Face**：[https://huggingface.co/models](https://huggingface.co/models)
|
||||
- **Civitai**：[https://civitai.com/models](https://civitai.com/models)
|
||||
- **官方源**：
|
||||
- FLUX: [https://huggingface.co/black-forest-labs](https://huggingface.co/black-forest-labs)
|
||||
- SD3.5: [https://huggingface.co/stabilityai](https://huggingface.co/stabilityai)
|
||||
|
||||
#### 预防措施
|
||||
|
||||
1. **基础模型包**：至少下载一个基础模型
|
||||
- FLUX: `flux1-schnell.safetensors`（快速）或 `flux1-dev.safetensors`（高质量）
|
||||
- SD3.5: `sd3.5_large.safetensors`
|
||||
|
||||
2. **检查磁盘空间**：
|
||||
```bash
|
||||
# 检查可用空间
|
||||
df -h ComfyUI/models/
|
||||
```
|
||||
|
||||
3. **设置模型路径**（可选）：
|
||||
如果你的模型存储在其他位置，可以创建符号链接：
|
||||
```bash
|
||||
ln -s /path/to/your/models ComfyUI/models/diffusion_models/
|
||||
```
|
||||
</details>
|
||||
|
||||
### 3. 如何处理缺少 System Component 错误
|
||||
|
||||
当你看到类似 `Missing VAE encoder: ae.safetensors` 或其他组件文件缺失的错误时，需要下载相应的系统组件。
|
||||
|
||||
<details>
|
||||
<summary><b>🛠️ 解决缺少 System Component 错误</b></summary>
|
||||
|
||||
#### 常见组件错误
|
||||
|
||||
```plaintext
|
||||
Missing VAE encoder: ae.safetensors. Please download and place it in the models/vae folder.
|
||||
Missing CLIP encoder: clip_l.safetensors. Please download and place it in the models/clip folder.
|
||||
Missing T5 encoder: t5xxl_fp16.safetensors. Please download and place it in the models/clip folder.
|
||||
```
|
||||
|
||||
#### 组件类型说明
|
||||
|
||||
| 组件类型 | 文件名示例 | 用途 | 存放目录 |
|
||||
| -------------- | ------------------------------ | ---------- | ------------------ |
|
||||
| **VAE** | ae.safetensors | 图像编码 / 解码 | models/vae/ |
|
||||
| **CLIP** | clip\_l.safetensors | 文本编码（CLIP） | models/clip/ |
|
||||
| **T5** | t5xxl\_fp16.safetensors | 文本编码（T5） | models/clip/ |
|
||||
| **ControlNet** | flux-controlnet-\*.safetensors | 控制网络 | models/controlnet/ |
|
||||
|
||||
#### 解决方法
|
||||
|
||||
**方法一：使用 Comfy-Manager（推荐）**
|
||||
|
||||
1. 点击 "Manager" → "Install Models"
|
||||
2. 在 "Filter" 中选择组件类型（VAE/CLIP/T5）
|
||||
3. 下载对应的组件文件
|
||||
|
||||
**方法二：手动下载必需组件**
|
||||
|
||||
##### FLUX 必需组件
|
||||
|
||||
```bash
|
||||
# 1. VAE 编码器
|
||||
cd ComfyUI/models/vae/
|
||||
wget https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/ae.safetensors
|
||||
|
||||
# 2. CLIP-L 编码器
|
||||
cd ComfyUI/models/clip/
|
||||
wget https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors
|
||||
|
||||
# 3. T5-XXL 编码器（可选择不同精度）
|
||||
# FP16 版本（推荐，平衡性能）
|
||||
wget https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp16.safetensors
|
||||
|
||||
# 或 FP8 版本（节省显存）
|
||||
wget https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors
|
||||
```
|
||||
|
||||
##### SD3.5 必需组件
|
||||
|
||||
```bash
|
||||
# SD3.5 使用不同的编码器
|
||||
cd ComfyUI/models/clip/
|
||||
|
||||
# CLIP-G 编码器
|
||||
wget https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/text_encoders/clip_g.safetensors
|
||||
|
||||
# CLIP-L 编码器
|
||||
wget https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/text_encoders/clip_l.safetensors
|
||||
|
||||
# T5-XXL 编码器
|
||||
wget https://huggingface.co/stabilityai/stable-diffusion-3.5-large/resolve/main/text_encoders/t5xxl_fp16.safetensors
|
||||
```
|
||||
|
||||
##### SDXL 必需组件
|
||||
|
||||
```bash
|
||||
# SDXL VAE
|
||||
cd ComfyUI/models/vae/
|
||||
wget https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors
|
||||
|
||||
# SDXL 使用内置的 CLIP 编码器，通常不需要单独下载
|
||||
```
|
||||
|
||||
#### 组件兼容性矩阵
|
||||
|
||||
| 模型系列 | 必需 VAE | 必需 CLIP | 必需 T5 | 可选组件 |
|
||||
| --------- | -------------- | ------------------- | ----------------------- | ---------- |
|
||||
| **FLUX** | ae.safetensors | clip\_l.safetensors | t5xxl\_fp16.safetensors | ControlNet |
|
||||
| **SD3.5** | 内置 | clip\_g + clip\_l | t5xxl\_fp16 | - |
|
||||
| **SDXL** | sdxl\_vae | 内置 | - | Refiner |
|
||||
| **SD1.5** | vae-ft-mse | 内置 | - | ControlNet |
|
||||
|
||||
#### 精度选择建议
|
||||
|
||||
**T5 编码器精度选择**：
|
||||
|
||||
| 显存容量 | 推荐版本 | 文件名 |
|
||||
| ------- | ------ | ------------------------------ |
|
||||
| \< 12GB | FP8 量化 | t5xxl\_fp8\_e4m3fn.safetensors |
|
||||
| 12-16GB | FP16 | t5xxl\_fp16.safetensors |
|
||||
| > 16GB | FP32 | t5xxl.safetensors |
|
||||
|
||||
#### 验证组件安装
|
||||
|
||||
```bash
|
||||
# 检查所有必需组件
|
||||
echo "=== VAE Components ==="
|
||||
ls -la ComfyUI/models/vae/
|
||||
|
||||
echo "=== CLIP/T5 Components ==="
|
||||
ls -la ComfyUI/models/clip/
|
||||
|
||||
echo "=== ControlNet Components ==="
|
||||
ls -la ComfyUI/models/controlnet/
|
||||
```
|
||||
|
||||
#### 故障排除
|
||||
|
||||
**问题：下载后仍然报错**
|
||||
|
||||
1. **检查文件权限**：
|
||||
```bash
|
||||
chmod 644 ComfyUI/models/vae/*.safetensors
|
||||
chmod 644 ComfyUI/models/clip/*.safetensors
|
||||
```
|
||||
|
||||
2. **清除缓存**：
|
||||
```bash
|
||||
# 清除 ComfyUI 缓存
|
||||
rm -rf ComfyUI/temp/*
|
||||
rm -rf ComfyUI/__pycache__/*
|
||||
```
|
||||
|
||||
3. **重启服务器**：
|
||||
```bash
|
||||
# 完全重启 ComfyUI
|
||||
pkill -f "python main.py"
|
||||
python main.py --listen 0.0.0.0 --port 8000
|
||||
```
|
||||
|
||||
**问题：显存不足**
|
||||
|
||||
使用量化版本的组件：
|
||||
|
||||
- T5: 使用 `t5xxl_fp8_e4m3fn.safetensors` 而不是 FP16/FP32
|
||||
- VAE: 某些模型支持 FP16 VAE 版本
|
||||
|
||||
**问题：下载速度慢**
|
||||
|
||||
1. 使用镜像源（如适用）
|
||||
2. 使用下载工具（如 aria2c）支持断点续传：
|
||||
```bash
|
||||
aria2c -x 16 -s 16 -k 1M [下载链接]
|
||||
```
|
||||
</details>
|
||||
|
||||
## ComfyUI 服务器安装
|
||||
|
||||
<details>
|
||||
<summary><b>🚀 安装和配置 ComfyUI 服务器</b></summary>
|
||||
|
||||
### 1. 安装 ComfyUI
|
||||
|
||||
```bash
|
||||
# 克隆 ComfyUI 仓库
|
||||
git clone https://github.com/comfyanonymous/ComfyUI.git
|
||||
cd ComfyUI
|
||||
|
||||
# 安装依赖
|
||||
pip install -r requirements.txt
|
||||
|
||||
# 可选：安装JWT支持（用于Token认证）
|
||||
pip install PyJWT
|
||||
|
||||
# 启动 ComfyUI 服务器
|
||||
python main.py --listen 0.0.0.0 --port 8000
|
||||
```
|
||||
|
||||
### 2. 下载模型文件
|
||||
|
||||
**推荐基础配置** (最小化安装)：
|
||||
|
||||
**主模型** (放入 `models/diffusion_models/` 目录)：
|
||||
|
||||
- `flux1-schnell.safetensors` - 快速生成（4 步）
|
||||
- `flux1-dev.safetensors` - 高质量创作（20 步）
|
||||
|
||||
**必需组件** (放入相应目录)：
|
||||
|
||||
- `models/vae/ae.safetensors` - VAE 编码器
|
||||
- `models/clip/clip_l.safetensors` - CLIP 文本编码器
|
||||
- `models/clip/t5xxl_fp16.safetensors` - T5 文本编码器
|
||||
|
||||
### 3. 验证服务器运行
|
||||
|
||||
访问 `http://localhost:8000` 确认 ComfyUI 界面正常加载。
|
||||
|
||||
<Callout type={'info'}>
|
||||
**智能模型选择**：LobeChat 会根据服务器上可用的模型文件自动选择最佳模型。您无需下载所有模型，系统会在可用模型中按优先级（官方 > 企业 > 社区）自动选择。
|
||||
</Callout>
|
||||
</details>
|
||||
|
||||
## 支持的模型
|
||||
|
||||
LobeChat ComfyUI 集成采用配置驱动的架构，支持 **223 个模型**，提供从官方模型到社区优化版本的全覆盖。
|
||||
|
||||
### FLUX 系列推荐参数
|
||||
|
||||
| 模型类型 | 推荐步数 | CFG Scale | 分辨率范围 |
|
||||
| ----------- | ---- | --------- | ------------------- |
|
||||
| **Schnell** | 4 步 | - | 512×512 至 1536×1536 |
|
||||
| **Dev** | 20 步 | 3.5 | 512×512 至 2048×2048 |
|
||||
| **Kontext** | 20 步 | 3.5 | 512×512 至 2048×2048 |
|
||||
| **Krea** | 20 步 | 4.5 | 512×512 至 2048×2048 |
|
||||
|
||||
### SD3.5 系列参数
|
||||
|
||||
| 模型类型 | 推荐步数 | CFG Scale | 分辨率范围 |
|
||||
| --------------- | ---- | --------- | ------------------- |
|
||||
| **Large** | 25 步 | 7.0 | 512×512 至 2048×2048 |
|
||||
| **Large Turbo** | 8 步 | 3.5 | 512×512 至 1536×1536 |
|
||||
| **Medium** | 20 步 | 6.0 | 512×512 至 1536×1536 |
|
||||
|
||||
<details>
|
||||
<summary><b>📋 当前完整支持的模型列表</b></summary>
|
||||
|
||||
### 模型分类体系
|
||||
|
||||
#### 优先级 1：官方核心模型
|
||||
|
||||
**FLUX.1 Official 系列**：
|
||||
|
||||
- `flux1-dev.safetensors` - 高质量创作模型
|
||||
- `flux1-schnell.safetensors` - 快速生成模型
|
||||
- `flux1-kontext-dev.safetensors` - 图像编辑模型
|
||||
- `flux1-krea-dev.safetensors` - 安全增强模型
|
||||
|
||||
**SD3.5 Official 系列**：
|
||||
|
||||
- `sd3.5_large.safetensors` - SD3.5 大型基础模型
|
||||
- `sd3.5_large_turbo.safetensors` - 快速生成版本
|
||||
- `sd3.5_medium.safetensors` - 中等规模模型
|
||||
|
||||
#### 优先级 2：企业优化模型（106 个 FLUX）
|
||||
|
||||
**量化优化系列**：
|
||||
|
||||
- **GGUF 量化**：每个变体支持 11 种量化级别（F16, Q8\_0, Q6\_K, Q5\_K\_M, Q5\_K\_S, Q4\_K\_M, Q4\_K\_S, Q4\_0, Q3\_K\_M, Q3\_K\_S, Q2\_K）
|
||||
- **FP8 精度**：fp8\_e4m3fn、fp8\_e5m2 优化版本
|
||||
- **企业轻量级**：FLUX.1-lite-8B 系列
|
||||
- **技术实验**：NF4、SVDQuant、TorchAO、optimum-quanto、MFLUX 优化版本
|
||||
|
||||
#### 优先级 3：社区精调模型（48 个 FLUX）
|
||||
|
||||
**社区优化系列**：
|
||||
|
||||
- **Jib Mix Flux** 系列：高质量混合模型
|
||||
- **Real Dream FLUX** 系列：现实主义风格
|
||||
- **Vision Realistic** 系列：视觉现实化
|
||||
- **PixelWave FLUX** 系列：像素艺术优化
|
||||
- **Fluxmania** 系列：多样化风格支持
|
||||
|
||||
### SD 系列模型支持（93 个）
|
||||
|
||||
**SD3.5 系列**：5 个模型
|
||||
**SD1.5 系列**：37 个模型（包括官方、量化和社区版本）
|
||||
**SDXL 系列**：50 个模型（包括基础、Refiner 和 Playground 模型）
|
||||
|
||||
### 工作流支持
|
||||
|
||||
系统支持 **6 种工作流**：
|
||||
|
||||
- **flux-dev**：高质量创作工作流
|
||||
- **flux-schnell**：快速生成工作流
|
||||
- **flux-kontext**：图像编辑工作流
|
||||
- **sd35**：SD3.5 专用工作流
|
||||
- **simple-sd**：简单 SD 工作流
|
||||
- **index**：工作流入口
|
||||
</details>
|
||||
|
||||
## 性能优化建议
|
||||
|
||||
### 硬件要求
|
||||
|
||||
**最低配置** (GGUF 量化模型)：
|
||||
|
||||
- GPU：6GB VRAM (使用 Q4 量化)
|
||||
- RAM：12GB
|
||||
- 存储：30GB 可用空间
|
||||
|
||||
**推荐配置** (标准模型)：
|
||||
|
||||
- GPU：12GB+ VRAM (RTX 4070 Ti 或更高)
|
||||
- RAM：24GB+
|
||||
- 存储：SSD 100GB+ 可用空间
|
||||
|
||||
### 显存优化策略
|
||||
|
||||
| 显存容量 | 推荐量化 | 模型示例 | 性能特点 |
|
||||
| ----------- | --------------- | ---------------------------------- | ------- |
|
||||
| **6-8GB** | Q4\_0, Q4\_K\_S | `flux1-dev-Q4_0.gguf` | 最小显存占用 |
|
||||
| **10-12GB** | Q6\_K, Q8\_0 | `flux1-dev-Q6_K.gguf` | 平衡性能与质量 |
|
||||
| **16GB+** | FP8, FP16 | `flux1-dev-fp8-e4m3fn.safetensors` | 接近原始质量 |
|
||||
| **24GB+** | 完整模型 | `flux1-dev.safetensors` | 最佳质量 |
|
||||
|
||||
## 自定义模型使用
|
||||
|
||||
<details>
|
||||
<summary><b>🎨 配置自定义 SD 模型</b></summary>
|
||||
|
||||
LobeChat 支持使用自定义的 Stable Diffusion 模型。系统使用固定的文件名来识别自定义模型。
|
||||
|
||||
### 1. 模型文件准备
|
||||
|
||||
**必需文件**：
|
||||
|
||||
- **主模型文件**：`custom_sd_lobe.safetensors`
|
||||
- **VAE 文件（可选）**：`custom_sd_vae_lobe.safetensors`
|
||||
|
||||
### 2. 添加自定义模型
|
||||
|
||||
**方法一：重命名现有模型**
|
||||
|
||||
```bash
|
||||
# 将您的模型重命名为固定文件名
|
||||
mv your_custom_model.safetensors custom_sd_lobe.safetensors
|
||||
|
||||
# 移动到正确目录
|
||||
mv custom_sd_lobe.safetensors ComfyUI/models/diffusion_models/
|
||||
```
|
||||
|
||||
**方法二：创建符号链接（推荐）**
|
||||
|
||||
```bash
|
||||
# 创建软链接，方便切换不同模型
|
||||
ln -s /path/to/your_model.safetensors ComfyUI/models/diffusion_models/custom_sd_lobe.safetensors
|
||||
```
|
||||
|
||||
### 3. 使用自定义模型
|
||||
|
||||
在 LobeChat 中，自定义模型会显示为：
|
||||
|
||||
- **stable-diffusion-custom**：标准自定义模型
|
||||
- **stable-diffusion-custom-refiner**：Refiner 自定义模型
|
||||
|
||||
### 自定义模型参数建议
|
||||
|
||||
| 参数 | SD 1.5 模型 | SDXL 模型 |
|
||||
| ---------- | --------- | ------- |
|
||||
| **steps** | 20-30 | 25-40 |
|
||||
| **cfg** | 7.0 | 6.0-8.0 |
|
||||
| **width** | 512 | 1024 |
|
||||
| **height** | 512 | 1024 |
|
||||
</details>
|
||||
|
||||
## 故障排除
|
||||
|
||||
### 智能错误诊断系统
|
||||
|
||||
LobeChat 集成了智能错误处理系统，能够自动诊断并提供针对性的解决方案。
|
||||
|
||||
#### 错误类型与解决方案
|
||||
|
||||
| 错误类型 | 用户提示 | 自动诊断 |
|
||||
| -------- | ------------------- | --------------- |
|
||||
| **连接问题** | "无法连接到 ComfyUI 服务器" | 自动检测服务器状态和网络连通性 |
|
||||
| **认证问题** | "API 密钥无效或已过期" | 自动验证认证凭据有效性 |
|
||||
| **权限问题** | "访问权限不足" | 自动检查用户权限和文件访问权限 |
|
||||
| **模型问题** | "找不到指定的模型文件" | 自动扫描可用模型并建议替代方案 |
|
||||
| **配置问题** | "配置文件存在错误" | 自动验证配置完整性和语法正确性 |
|
||||
|
||||
<details>
|
||||
<summary><b>🔍 传统故障排除方法</b></summary>
|
||||
|
||||
#### 1. 连接失败
|
||||
|
||||
**问题**：无法连接到 ComfyUI 服务器
|
||||
|
||||
**解决方案**：
|
||||
|
||||
```bash
|
||||
# 确认服务器运行
|
||||
curl http://localhost:8000/system_stats
|
||||
|
||||
# 检查端口
|
||||
netstat -tulpn | grep 8000
|
||||
```
|
||||
|
||||
#### 2. 内存不足
|
||||
|
||||
**问题**：生成过程中出现内存错误
|
||||
|
||||
**解决方案**：
|
||||
|
||||
- 降低图像分辨率
|
||||
- 减少生成步数
|
||||
- 使用量化模型
|
||||
|
||||
#### 3. 认证失败
|
||||
|
||||
**问题**：401 或 403 错误
|
||||
|
||||
**解决方案**：
|
||||
|
||||
- 验证认证配置
|
||||
- 检查 Token 是否过期
|
||||
- 确认用户权限
|
||||
</details>
|
||||
|
||||
## 最佳实践
|
||||
|
||||
### 提示词编写
|
||||
|
||||
1. **详细描述**：提供清晰、详细的图像描述
|
||||
2. **风格指定**：明确指定艺术风格、色彩风格等
|
||||
3. **质量关键词**：添加 "4K", "high quality", "detailed" 等关键词
|
||||
4. **避免矛盾**：确保描述内容逻辑一致
|
||||
|
||||
**示例**：
|
||||
|
||||
```plaintext
|
||||
A young woman with flowing long hair, wearing an elegant blue dress, standing in a cherry blossom park,
|
||||
sunlight filtering through leaves, warm atmosphere, cinematic lighting, 4K high resolution, detailed, photorealistic
|
||||
```
|
||||
|
||||
### 参数调优
|
||||
|
||||
1. **FLUX Schnell**：适合快速预览，使用 4 步生成
|
||||
2. **FLUX Dev**：平衡质量和速度，CFG 3.5，步数 20
|
||||
3. **FLUX Krea-dev**：安全创作，CFG 4.5，注意内容过滤
|
||||
4. **FLUX Kontext-dev**：图像编辑，strength 0.6-0.9
|
||||
|
||||
<Callout type={'warning'}>
|
||||
在使用过程中请注意：
|
||||
|
||||
- FLUX Dev、Krea-dev、Kontext-dev 模型仅限非商业使用
|
||||
- 生成内容请遵守相关法律法规和平台政策
|
||||
- 大型模型生成可能需要较长时间，请耐心等待
|
||||
</Callout>
|
||||
|
||||
## API 参考
|
||||
|
||||
<details>
|
||||
<summary><b>📚 API 文档</b></summary>
|
||||
|
||||
### 请求格式
|
||||
|
||||
```typescript
|
||||
interface ComfyUIRequest {
|
||||
model: string; // 模型 ID，如 'flux-schnell'
|
||||
prompt: string; // 文本提示词
|
||||
width: number; // 图像宽度
|
||||
height: number; // 图像高度
|
||||
steps: number; // 生成步数
|
||||
seed: number; // 随机种子
|
||||
cfg?: number; // CFG Scale (Dev/Krea/Kontext 专用)
|
||||
strength?: number; // 编辑强度 (Kontext 专用)
|
||||
imageUrl?: string; // 输入图像 (Kontext 专用)
|
||||
}
|
||||
```
|
||||
|
||||
### 响应格式
|
||||
|
||||
```typescript
|
||||
interface ComfyUIResponse {
|
||||
images: Array<{
|
||||
url: string; // 生成的图像 URL
|
||||
filename: string; // 文件名
|
||||
subfolder: string; // 子目录
|
||||
type: string; // 文件类型
|
||||
}>;
|
||||
prompt_id: string; // 提示 ID
|
||||
}
|
||||
```
|
||||
|
||||
### 错误代码
|
||||
|
||||
| 错误代码 | 描述 | 解决建议 |
|
||||
| ----- | ------ | -------------- |
|
||||
| `400` | 请求参数无效 | 检查参数格式和范围 |
|
||||
| `401` | 认证失败 | 验证 API 密钥和认证配置 |
|
||||
| `403` | 权限不足 | 检查用户权限 |
|
||||
| `404` | 模型未找到 | 确认模型文件存在 |
|
||||
| `500` | 服务器错误 | 检查 ComfyUI 日志 |
|
||||
</details>
|
||||
|
||||
至此你已经可以在 LobeChat 中使用 ComfyUI 进行高质量的 AI 图像生成和编辑了。如果遇到问题，请参考故障排除部分或查阅 [ComfyUI 官方文档](https://github.com/comfyanonymous/ComfyUI)。
|
||||
|
|
@ -82,6 +82,58 @@
|
|||
"title": "Cloudflare Account ID / API Address"
|
||||
}
|
||||
},
|
||||
"comfyui": {
|
||||
"apiKey": {
|
||||
"desc": "API key for Bearer Token authentication",
|
||||
"placeholder": "Enter API key",
|
||||
"title": "API Key"
|
||||
},
|
||||
"authType": {
|
||||
"desc": "Choose authentication method for ComfyUI server",
|
||||
"options": {
|
||||
"basic": "Basic Authentication",
|
||||
"bearer": "Bearer Token",
|
||||
"custom": "Custom Authentication",
|
||||
"none": "No Authentication"
|
||||
},
|
||||
"placeholder": "Select authentication method",
|
||||
"title": "Authentication Type"
|
||||
},
|
||||
"baseURL": {
|
||||
"desc": "ComfyUI server access address, e.g., http://localhost:8000",
|
||||
"placeholder": "http://localhost:8000",
|
||||
"title": "Server Address"
|
||||
},
|
||||
"checker": {
|
||||
"desc": "Test whether the ComfyUI server can connect normally",
|
||||
"title": "Connection Test"
|
||||
},
|
||||
"customHeaders": {
|
||||
"addButton": "Add Header",
|
||||
"deleteTooltip": "Delete this header",
|
||||
"desc": "Custom HTTP request headers for custom authentication, in key-value format",
|
||||
"duplicateKeyError": "Header names cannot be duplicated",
|
||||
"keyPlaceholder": "Header name",
|
||||
"title": "Custom Headers",
|
||||
"valuePlaceholder": "Header value"
|
||||
},
|
||||
"password": {
|
||||
"desc": "Password for basic authentication",
|
||||
"placeholder": "Enter password",
|
||||
"title": "Password"
|
||||
},
|
||||
"title": "ComfyUI",
|
||||
"unlock": {
|
||||
"customAuth": "Custom Authentication",
|
||||
"description": "Configure ComfyUI server connection information to start image generation",
|
||||
"title": "Use ComfyUI Image Generation"
|
||||
},
|
||||
"username": {
|
||||
"desc": "Username for basic authentication",
|
||||
"placeholder": "Enter username",
|
||||
"title": "Username"
|
||||
}
|
||||
},
|
||||
"createNewAiProvider": {
|
||||
"apiKey": {
|
||||
"placeholder": "Please enter your API Key",
|
||||
|
|
|
|||
|
|
@ -82,6 +82,58 @@
|
|||
"title": "Cloudflare 账户 ID / API 地址"
|
||||
}
|
||||
},
|
||||
"comfyui": {
|
||||
"apiKey": {
|
||||
"desc": "Bearer Token 认证所需的 API 密钥",
|
||||
"placeholder": "请输入 API 密钥",
|
||||
"title": "API 密钥"
|
||||
},
|
||||
"authType": {
|
||||
"desc": "选择 ComfyUI 服务器的认证方式",
|
||||
"options": {
|
||||
"basic": "账号/密码",
|
||||
"bearer": "Bearer (API 密钥)",
|
||||
"custom": "自定义请求头",
|
||||
"none": "无需认证"
|
||||
},
|
||||
"placeholder": "请选择认证方式",
|
||||
"title": "认证类型"
|
||||
},
|
||||
"baseURL": {
|
||||
"desc": "ComfyUI 网页访问地址",
|
||||
"placeholder": "http://localhost:8000",
|
||||
"title": "访问地址"
|
||||
},
|
||||
"checker": {
|
||||
"desc": "测试 ComfyUI 服务器是否可以正常连接",
|
||||
"title": "连接测试"
|
||||
},
|
||||
"customHeaders": {
|
||||
"addButton": "添加请求头",
|
||||
"deleteTooltip": "删除此请求头",
|
||||
"desc": "自定义认证方式下的HTTP请求头，格式为键值对",
|
||||
"duplicateKeyError": "请求头键名不能重复",
|
||||
"keyPlaceholder": "请求头键名",
|
||||
"title": "自定义请求头",
|
||||
"valuePlaceholder": "请求头值"
|
||||
},
|
||||
"password": {
|
||||
"desc": "基本认证所需的密码",
|
||||
"placeholder": "请输入密码",
|
||||
"title": "密码"
|
||||
},
|
||||
"title": "ComfyUI",
|
||||
"unlock": {
|
||||
"customAuth": "自定义认证",
|
||||
"description": "配置 ComfyUI 服务器连接信息即可开始图像生成",
|
||||
"title": "使用 ComfyUI 图像生成"
|
||||
},
|
||||
"username": {
|
||||
"desc": "基本认证所需的用户名",
|
||||
"placeholder": "请输入用户名",
|
||||
"title": "用户名"
|
||||
}
|
||||
},
|
||||
"createNewAiProvider": {
|
||||
"apiKey": {
|
||||
"placeholder": "请填写你的 API Key",
|
||||
|
|
|
|||
|
|
@ -866,6 +866,46 @@
|
|||
"cohere/embed-v4.0": {
|
||||
"description": "一个允许对文本、图像或混合内容进行分类或转换为嵌入的模型。"
|
||||
},
|
||||
"comfyui/flux-dev": {
|
||||
"description": "FLUX.1 Dev - 高质量文生图模型，支持 guidance scale 调节，10-50步生成，非商业许可，适合高质量创作和艺术作品生成"
|
||||
},
|
||||
"comfyui/flux-kontext-dev": {
|
||||
"description": "FLUX.1 Kontext-dev - 图像编辑模型，支持基于文本指令修改现有图像，支持局部修改和风格迁移，非商业许可"
|
||||
},
|
||||
"comfyui/flux-krea-dev": {
|
||||
"description": "FLUX.1 Krea-dev - 增强安全的文生图模型，与 Krea 合作开发，内置安全过滤，避免生成不当内容，非商业许可"
|
||||
},
|
||||
"comfyui/flux-schnell": {
|
||||
"description": "FLUX.1 Schnell - 超快速文生图模型，1-4步即可生成高质量图像，Apache 2.0开源许可，适合实时应用和快速原型制作"
|
||||
},
|
||||
"comfyui/stable-diffusion-15": {
|
||||
"displayName": "SD 1.5",
|
||||
"description": "Stable Diffusion 1.5 文生图模型，经典的512x512分辨率文本到图像生成，适合快速原型和创意实验。支持负向提示。"
|
||||
},
|
||||
"comfyui/stable-diffusion-35": {
|
||||
"displayName": "Stable Diffusion 3.5",
|
||||
"description": "Stable Diffusion 3.5 新一代文生图模型，支持 Large 和 Medium 两个版本，需要外部 CLIP 编码器文件，提供卓越的图像质量和提示词匹配度。"
|
||||
},
|
||||
"comfyui/stable-diffusion-35-inclclip": {
|
||||
"displayName": "Stable Diffusion 3.5 (内置编码器)",
|
||||
"description": "Stable Diffusion 3.5 内置 CLIP/T5 编码器版本，无需外部编码器文件，适用于 sd3.5_medium_incl_clips 等模型，资源占用更少。"
|
||||
},
|
||||
"comfyui/stable-diffusion-custom": {
|
||||
"displayName": "Custom SD",
|
||||
"description": "自定义 SD 文生图模型，模型文件名请使用 custom_sd_lobe.safetensors，如有 VAE 请使用 custom_sd_vae_lobe.safetensors，模型文件需要按照 Comfy 的要求放入对应文件夹。"
|
||||
},
|
||||
"comfyui/stable-diffusion-custom-refiner": {
|
||||
"displayName": "Custom SD Refiner",
|
||||
"description": "自定义 SD 图生图模型，模型文件名请使用 custom_sd_lobe.safetensors，如有 VAE 请使用 custom_sd_vae_lobe.safetensors，模型文件需要按照 Comfy 的要求放入对应文件夹。"
|
||||
},
|
||||
"comfyui/stable-diffusion-refiner": {
|
||||
"displayName": "SDXL Image-to-Image",
|
||||
"description": "SDXL 图生图模型，基于输入图像进行高质量的图像到图像转换，支持风格迁移、图像修复和创意变换。"
|
||||
},
|
||||
"comfyui/stable-diffusion-xl": {
|
||||
"displayName": "SDXL Text-to-Image",
|
||||
"description": "SDXL 文生图模型，支持1024x1024高分辨率文本到图像生成，提供更好的图像质量和细节表现。支持负向提示。"
|
||||
},
|
||||
"command": {
|
||||
"description": "一个遵循指令的对话模型，在语言任务中表现出高质量、更可靠，并且相比我们的基础生成模型具有更长的上下文长度。"
|
||||
},
|
||||
|
|
@ -3017,6 +3057,9 @@
|
|||
"sonar-reasoning-pro": {
|
||||
"description": "支持搜索上下文的高级搜索产品，支持高级查询和跟进。"
|
||||
},
|
||||
"stable-diffusion-15": {
|
||||
"description": "Stable Diffusion 1.5 文生图模型，经典的512x512分辨率文本到图像生成，适合快速原型和创意实验。支持负向提示。"
|
||||
},
|
||||
"stable-diffusion-3-medium": {
|
||||
"description": "由 Stability AI 推出的最新文生图大模型。这一版本在继承了前代的优点上，对图像质量、文本理解和风格多样性等方面进行了显著改进，能够更准确地解读复杂的自然语言提示，并生成更为精确和多样化的图像。"
|
||||
},
|
||||
|
|
@ -3026,6 +3069,15 @@
|
|||
"stable-diffusion-3.5-large-turbo": {
|
||||
"description": "stable-diffusion-3.5-large-turbo 是在 stable-diffusion-3.5-large 的基础上采用对抗性扩散蒸馏（ADD）技术的模型，具备更快的速度。"
|
||||
},
|
||||
"stable-diffusion-custom": {
|
||||
"description": "自定义 SD 文生图模型，支持社区和第三方训练的 Stable Diffusion 文本到图像模型，提供灵活的参数配置。"
|
||||
},
|
||||
"stable-diffusion-custom-refiner": {
|
||||
"description": "自定义 SD 图生图模型，支持社区和第三方训练的 Stable Diffusion 图像到图像模型，适合专业图像处理工作流。"
|
||||
},
|
||||
"stable-diffusion-refiner": {
|
||||
"description": "SDXL 图生图模型，基于输入图像进行高质量的图像到图像转换，支持风格迁移、图像修复和创意变换。"
|
||||
},
|
||||
"stable-diffusion-v1.5": {
|
||||
"description": "stable-diffusion-v1.5 是以 stable-diffusion-v1.2 检查点的权重进行初始化，并在 \"laion-aesthetics v2 5+\" 上以 512x512 的分辨率进行了595k步的微调，减少了 10% 的文本条件化，以提高无分类器的引导采样。"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -44,6 +44,9 @@
|
|||
"cometapi": {
|
||||
"description": "CometAPI 是一个提供多种前沿大模型接口的服务平台，支持 OpenAI、Anthropic、Google 及更多，适合多样化的开发和应用需求。用户可根据自身需求灵活选择最优的模型和价格，助力AI体验的提升。"
|
||||
},
|
||||
"comfyui": {
|
||||
"description": "强大的开源图像、视频、音频生成工作流引擎，支持 SD FLUX Qwen Hunyuan WAN 等先进模型，提供节点化工作流编辑和私有化部署能力"
|
||||
},
|
||||
"deepseek": {
|
||||
"description": "DeepSeek 是一家专注于人工智能技术研究和应用的公司，其最新模型 DeepSeek-V3 多项评测成绩超越 Qwen2.5-72B 和 Llama-3.1-405B 等开源模型，性能对齐领军闭源模型 GPT-4o 与 Claude-3.5-Sonnet。"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -175,6 +175,7 @@
|
|||
"@opentelemetry/exporter-jaeger": "^2.1.0",
|
||||
"@opentelemetry/winston-transport": "^0.17.0",
|
||||
"@react-spring/web": "^9.7.5",
|
||||
"@saintno/comfyui-sdk": "^0.2.48",
|
||||
"@serwist/next": "^9.2.1",
|
||||
"@t3-oss/env-nextjs": "^0.13.8",
|
||||
"@tanstack/react-query": "^5.90.2",
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@
|
|||
"./cloudflare": "./src/aiModels/cloudflare.ts",
|
||||
"./cohere": "./src/aiModels/cohere.ts",
|
||||
"./cometapi": "./src/aiModels/cometapi.ts",
|
||||
"./comfyui": "./src/aiModels/comfyui.ts",
|
||||
"./deepseek": "./src/aiModels/deepseek.ts",
|
||||
"./fal": "./src/aiModels/fal.ts",
|
||||
"./fireworksai": "./src/aiModels/fireworksai.ts",
|
||||
|
|
|
|||
335
packages/model-bank/src/aiModels/comfyui.ts
Normal file
335
packages/model-bank/src/aiModels/comfyui.ts
Normal file
|
|
@ -0,0 +1,335 @@
|
|||
import { ModelParamsSchema, PRESET_ASPECT_RATIOS } from '../standard-parameters';
|
||||
import { AIImageModelCard } from '../types';
|
||||
|
||||
/**
|
||||
* Aspect ratios supported by FLUX models
|
||||
* Support wide range ratios from 21:9 to 9:21, including foldable screen devices
|
||||
*/
|
||||
const FLUX_ASPECT_RATIOS = [
|
||||
'21:9', // Ultra-wide screen
|
||||
'16:9', // Widescreen
|
||||
'8:7', // Foldable screen (e.g. Galaxy Z Fold, unfolded state ~7.6 inch)
|
||||
'4:3', // Traditional landscape
|
||||
'3:2', // Classic landscape
|
||||
'1:1', // Square
|
||||
'2:3', // Classic portrait
|
||||
'3:4', // Traditional portrait
|
||||
'7:8', // Foldable screen portrait
|
||||
'9:16', // Portrait
|
||||
'9:21', // Ultra-tall portrait
|
||||
];
|
||||
|
||||
/**
|
||||
* Standard aspect ratios supported by SD models
|
||||
* Based on preset aspect ratios, suitable for traditional SD model use cases
|
||||
*/
|
||||
const SD_ASPECT_RATIOS = PRESET_ASPECT_RATIOS;
|
||||
|
||||
/**
|
||||
* Extended aspect ratios supported by SDXL models
|
||||
* Support more modern display ratios, similar to FLUX but more conservative
|
||||
*/
|
||||
const SDXL_ASPECT_RATIOS = [
|
||||
'16:9', // Modern widescreen
|
||||
'4:3', // Traditional landscape
|
||||
'3:2', // Classic landscape
|
||||
'1:1', // Square
|
||||
'2:3', // Classic portrait
|
||||
'3:4', // Traditional portrait
|
||||
'9:16', // Modern portrait
|
||||
];
|
||||
|
||||
/**
|
||||
* FLUX.1 Schnell model parameter configuration
|
||||
* Ultra-fast text-to-image mode, generates in 1-4 steps, Apache 2.0 license
|
||||
*/
|
||||
export const fluxSchnellParamsSchema: ModelParamsSchema = {
|
||||
aspectRatio: {
|
||||
default: '1:1',
|
||||
enum: FLUX_ASPECT_RATIOS,
|
||||
},
|
||||
cfg: { default: 1, max: 1, min: 1, step: 0 }, // Schnell uses fixed CFG of 1
|
||||
height: { default: 1024, max: 1536, min: 512, step: 8 },
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'euler' },
|
||||
scheduler: { default: 'simple' },
|
||||
seed: { default: null },
|
||||
steps: { default: 4, max: 4, min: 1, step: 1 },
|
||||
width: { default: 1024, max: 1536, min: 512, step: 8 },
|
||||
};
|
||||
|
||||
/**
|
||||
* FLUX.1 Dev model parameter configuration
|
||||
* High-quality text-to-image mode, supports guidance scale adjustment, non-commercial license
|
||||
*/
|
||||
export const fluxDevParamsSchema: ModelParamsSchema = {
|
||||
aspectRatio: {
|
||||
default: '1:1',
|
||||
enum: FLUX_ASPECT_RATIOS,
|
||||
},
|
||||
cfg: { default: 3.5, max: 10, min: 1, step: 0.5 },
|
||||
height: { default: 1024, max: 2048, min: 512, step: 8 },
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'euler' },
|
||||
scheduler: { default: 'simple' },
|
||||
seed: { default: null },
|
||||
steps: { default: 20, max: 50, min: 10, step: 1 },
|
||||
width: { default: 1024, max: 2048, min: 512, step: 8 },
|
||||
};
|
||||
|
||||
/**
|
||||
* FLUX.1 Krea-dev model parameter configuration
|
||||
* Enhanced safety text-to-image mode, developed in collaboration with Krea, non-commercial license
|
||||
*/
|
||||
export const fluxKreaDevParamsSchema: ModelParamsSchema = {
|
||||
aspectRatio: {
|
||||
default: '1:1',
|
||||
enum: FLUX_ASPECT_RATIOS,
|
||||
},
|
||||
cfg: { default: 3.5, max: 10, min: 1, step: 0.5 },
|
||||
height: { default: 1024, max: 2048, min: 512, step: 8 },
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'dpmpp_2m_sde' },
|
||||
scheduler: { default: 'karras' },
|
||||
seed: { default: null },
|
||||
steps: { default: 15, max: 50, min: 10, step: 1 },
|
||||
width: { default: 1024, max: 2048, min: 512, step: 8 },
|
||||
};
|
||||
|
||||
/**
|
||||
* FLUX.1 Kontext-dev model parameter configuration
|
||||
* Image editing mode, supports modifying existing images based on text instructions, non-commercial license
|
||||
*/
|
||||
export const fluxKontextDevParamsSchema: ModelParamsSchema = {
|
||||
cfg: { default: 3.5, max: 10, min: 1, step: 0.5 },
|
||||
imageUrl: { default: '' }, // Input image URL (supports text-to-image and image-to-image)
|
||||
prompt: { default: '' },
|
||||
seed: { default: null },
|
||||
steps: { default: 28, max: 50, min: 10, step: 1 }, // Kontext defaults to 28 steps
|
||||
strength: { default: 0.85, max: 1, min: 0, step: 0.05 }, // Image editing strength control (frontend parameter)
|
||||
};
|
||||
|
||||
/**
|
||||
* SD3.5 model parameter configuration
|
||||
* Stable Diffusion 3.5, supports Large and Medium versions, automatically selects by priority
|
||||
*/
|
||||
export const sd35ParamsSchema: ModelParamsSchema = {
|
||||
aspectRatio: {
|
||||
default: '1:1',
|
||||
enum: FLUX_ASPECT_RATIOS, // SD3.5 also supports multiple aspect ratios
|
||||
},
|
||||
cfg: { default: 4, max: 20, min: 1, step: 0.5 },
|
||||
height: { default: 1024, max: 2048, min: 512, step: 8 },
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'euler' },
|
||||
scheduler: { default: 'sgm_uniform' },
|
||||
seed: { default: null },
|
||||
steps: { default: 20, max: 50, min: 10, step: 1 },
|
||||
width: { default: 1024, max: 2048, min: 512, step: 8 },
|
||||
};
|
||||
|
||||
/**
|
||||
* SD1.5 text-to-image model parameter configuration
|
||||
* Stable Diffusion 1.5 text-to-image generation, suitable for 512x512 resolution
|
||||
*/
|
||||
export const sd15T2iParamsSchema: ModelParamsSchema = {
|
||||
aspectRatio: {
|
||||
default: '1:1',
|
||||
enum: SD_ASPECT_RATIOS,
|
||||
},
|
||||
cfg: { default: 7, max: 20, min: 1, step: 0.5 },
|
||||
height: { default: 512, max: 1024, min: 256, step: 8 },
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'euler' },
|
||||
scheduler: { default: 'normal' },
|
||||
seed: { default: null },
|
||||
steps: { default: 25, max: 50, min: 10, step: 1 },
|
||||
width: { default: 512, max: 1024, min: 256, step: 8 },
|
||||
};
|
||||
|
||||
/**
|
||||
* SDXL text-to-image model parameter configuration
|
||||
* SDXL text-to-image generation, suitable for 1024x1024 resolution
|
||||
*/
|
||||
export const sdxlT2iParamsSchema: ModelParamsSchema = {
|
||||
aspectRatio: {
|
||||
default: '1:1',
|
||||
enum: SDXL_ASPECT_RATIOS,
|
||||
},
|
||||
cfg: { default: 8, max: 20, min: 1, step: 0.5 },
|
||||
height: { default: 1024, max: 2048, min: 512, step: 8 },
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'euler' },
|
||||
scheduler: { default: 'normal' },
|
||||
seed: { default: null },
|
||||
steps: { default: 30, max: 50, min: 10, step: 1 },
|
||||
width: { default: 1024, max: 2048, min: 512, step: 8 },
|
||||
};
|
||||
|
||||
/**
|
||||
* SDXL image-to-image model parameter configuration
|
||||
* SDXL image-to-image generation, supports input image modification
|
||||
*/
|
||||
export const sdxlI2iParamsSchema: ModelParamsSchema = {
|
||||
cfg: { default: 8, max: 20, min: 1, step: 0.5 },
|
||||
imageUrl: { default: '' }, // Input image URL
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'euler' },
|
||||
scheduler: { default: 'normal' },
|
||||
seed: { default: null },
|
||||
steps: { default: 30, max: 50, min: 10, step: 1 },
|
||||
strength: { default: 0.75, max: 1, min: 0, step: 0.05 }, // Image modification strength (frontend parameter)
|
||||
};
|
||||
|
||||
/**
|
||||
* Custom SD text-to-image model parameter configuration
|
||||
* Custom Stable Diffusion text-to-image model with flexible parameter settings
|
||||
*/
|
||||
export const customSdT2iParamsSchema: ModelParamsSchema = {
|
||||
aspectRatio: {
|
||||
default: '1:1',
|
||||
enum: SDXL_ASPECT_RATIOS, // Use broader aspect ratio support
|
||||
},
|
||||
cfg: { default: 7, max: 30, min: 1, step: 0.5 },
|
||||
height: { default: 768, max: 2048, min: 256, step: 8 },
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'euler' }, // Use SDXL common parameters
|
||||
scheduler: { default: 'normal' }, // Use SDXL common parameters
|
||||
seed: { default: null },
|
||||
steps: { default: 25, max: 100, min: 5, step: 1 },
|
||||
width: { default: 768, max: 2048, min: 256, step: 8 },
|
||||
};
|
||||
|
||||
/**
|
||||
* Custom SD image-to-image model parameter configuration
|
||||
* Custom Stable Diffusion image-to-image model, supports image editing
|
||||
*/
|
||||
export const customSdI2iParamsSchema: ModelParamsSchema = {
|
||||
cfg: { default: 7, max: 30, min: 1, step: 0.5 },
|
||||
imageUrl: { default: '' }, // Input image URL
|
||||
prompt: { default: '' },
|
||||
samplerName: { default: 'euler' }, // Use SDXL common parameters
|
||||
scheduler: { default: 'normal' }, // Use SDXL common parameters
|
||||
seed: { default: null },
|
||||
steps: { default: 25, max: 100, min: 5, step: 1 },
|
||||
strength: { default: 0.75, max: 1, min: 0, step: 0.05 }, // Image modification strength (frontend parameter)
|
||||
};
|
||||
|
||||
/**
|
||||
* List of image generation models supported by ComfyUI
|
||||
* Supports FLUX series and Stable Diffusion 3.5 models
|
||||
*/
|
||||
const comfyuiImageModels: AIImageModelCard[] = [
|
||||
{
|
||||
description:
|
||||
'FLUX.1 Schnell - 超快速文生图模型，1-4步即可生成高质量图像，适合实时应用和快速原型制作',
|
||||
displayName: 'FLUX.1 Schnell',
|
||||
enabled: true,
|
||||
id: 'comfyui/flux-schnell',
|
||||
parameters: fluxSchnellParamsSchema,
|
||||
releasedAt: '2024-08-01',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description: 'FLUX.1 Dev - 高质量文生图模型，10-50步生成，适合高质量创作和艺术作品生成',
|
||||
displayName: 'FLUX.1 Dev',
|
||||
enabled: true,
|
||||
id: 'comfyui/flux-dev',
|
||||
parameters: fluxDevParamsSchema,
|
||||
releasedAt: '2024-08-01',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description: 'FLUX.1 Krea-dev - 增强安全的文生图模型，与 Krea 合作开发，内置安全过滤',
|
||||
displayName: 'FLUX.1 Krea-dev',
|
||||
enabled: false,
|
||||
id: 'comfyui/flux-krea-dev',
|
||||
parameters: fluxKreaDevParamsSchema,
|
||||
releasedAt: '2025-07-31',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description:
|
||||
'FLUX.1 Kontext-dev - 图像编辑模型，支持基于文本指令修改现有图像，支持局部修改和风格迁移',
|
||||
displayName: 'FLUX.1 Kontext-dev',
|
||||
enabled: true,
|
||||
id: 'comfyui/flux-kontext-dev',
|
||||
parameters: fluxKontextDevParamsSchema,
|
||||
releasedAt: '2025-05-29', // Aligned with BFL official Kontext series release date
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description:
|
||||
'Stable Diffusion 3.5 新一代文生图模型，支持 Large 和 Medium 两个版本，需要外部 CLIP 编码器文件，提供卓越的图像质量和提示词匹配度。',
|
||||
displayName: 'Stable Diffusion 3.5',
|
||||
enabled: true,
|
||||
id: 'comfyui/stable-diffusion-35',
|
||||
parameters: sd35ParamsSchema,
|
||||
releasedAt: '2024-10-22',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description:
|
||||
'Stable Diffusion 3.5 内置 CLIP/T5 编码器版本，无需外部编码器文件，适用于 sd3.5_medium_incl_clips 等模型，资源占用更少。',
|
||||
displayName: 'Stable Diffusion 3.5 (内置编码器)',
|
||||
enabled: false,
|
||||
id: 'comfyui/stable-diffusion-35-inclclip',
|
||||
parameters: sd35ParamsSchema,
|
||||
releasedAt: '2024-10-22',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description:
|
||||
'Stable Diffusion 1.5 文生图模型，经典的512x512分辨率文本到图像生成，适合快速原型和创意实验',
|
||||
displayName: 'SD 1.5',
|
||||
enabled: false,
|
||||
id: 'comfyui/stable-diffusion-15',
|
||||
parameters: sd15T2iParamsSchema,
|
||||
releasedAt: '2022-08-22',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description:
|
||||
'SDXL 文生图模型，支持1024x1024高分辨率文本到图像生成，提供更好的图像质量和细节表现',
|
||||
displayName: 'SDXL 文生图',
|
||||
enabled: true,
|
||||
id: 'comfyui/stable-diffusion-xl',
|
||||
parameters: sdxlT2iParamsSchema,
|
||||
releasedAt: '2023-07-26',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description:
|
||||
'SDXL 图生图模型，基于输入图像进行高质量的图像到图像转换，支持风格迁移、图像修复和创意变换。',
|
||||
displayName: 'SDXL Refiner',
|
||||
enabled: true,
|
||||
id: 'comfyui/stable-diffusion-refiner',
|
||||
parameters: sdxlI2iParamsSchema,
|
||||
releasedAt: '2023-07-26',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description:
|
||||
'自定义 SD 文生图模型，模型文件名请使用 custom_sd_lobe.safetensors，如有 VAE 请使用 custom_sd_vae_lobe.safetensors，模型文件需要按照 Comfy 的要求放入对应文件夹',
|
||||
displayName: '自定义 SD 文生图',
|
||||
enabled: false,
|
||||
id: 'comfyui/stable-diffusion-custom',
|
||||
parameters: customSdT2iParamsSchema,
|
||||
releasedAt: '2023-01-01',
|
||||
type: 'image',
|
||||
},
|
||||
{
|
||||
description:
|
||||
'自定义 SDXL 图生图模型，模型文件名请使用 custom_sd_lobe.safetensors，如有 VAE 请使用 custom_sd_vae_lobe.safetensors，模型文件需要按照 Comfy 的要求放入对应文件夹',
|
||||
displayName: '自定义 SDXL Refiner',
|
||||
enabled: false,
|
||||
id: 'comfyui/stable-diffusion-custom-refiner',
|
||||
parameters: customSdI2iParamsSchema,
|
||||
releasedAt: '2023-01-01',
|
||||
type: 'image',
|
||||
},
|
||||
];
|
||||
|
||||
export const allModels = [...comfyuiImageModels];
|
||||
|
||||
export default allModels;
|
||||
|
|
@ -14,6 +14,7 @@ import { default as cerebras } from './cerebras';
|
|||
import { default as cloudflare } from './cloudflare';
|
||||
import { default as cohere } from './cohere';
|
||||
import { default as cometapi } from './cometapi';
|
||||
import { default as comfyui } from './comfyui';
|
||||
import { default as deepseek } from './deepseek';
|
||||
import { default as fal } from './fal';
|
||||
import { default as fireworksai } from './fireworksai';
|
||||
|
|
@ -100,6 +101,7 @@ export const LOBE_DEFAULT_MODEL_LIST = buildDefaultModelList({
|
|||
cloudflare,
|
||||
cohere,
|
||||
cometapi,
|
||||
comfyui,
|
||||
deepseek,
|
||||
fal,
|
||||
fireworksai,
|
||||
|
|
@ -167,6 +169,7 @@ export { default as cerebras } from './cerebras';
|
|||
export { default as cloudflare } from './cloudflare';
|
||||
export { default as cohere } from './cohere';
|
||||
export { default as cometapi } from './cometapi';
|
||||
export { default as comfyui } from './comfyui';
|
||||
export { default as deepseek } from './deepseek';
|
||||
export { default as fal, fluxSchnellParamsSchema } from './fal';
|
||||
export { default as fireworksai } from './fireworksai';
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ export enum ModelProvider {
|
|||
Cloudflare = 'cloudflare',
|
||||
Cohere = 'cohere',
|
||||
CometAPI = 'cometapi',
|
||||
ComfyUI = 'comfyui',
|
||||
DeepSeek = 'deepseek',
|
||||
Fal = 'fal',
|
||||
FireworksAI = 'fireworksai',
|
||||
|
|
|
|||
|
|
@ -96,6 +96,30 @@ export const ModelParamsMetaSchema = z.object({
|
|||
})
|
||||
.optional(),
|
||||
|
||||
/**
|
||||
* samplerName is not requires by all i2i providers
|
||||
*/
|
||||
samplerName: z
|
||||
.object({
|
||||
default: z.string(),
|
||||
description: z.string().optional(),
|
||||
enum: z.array(z.string()).optional(),
|
||||
type: z.literal('string').optional(),
|
||||
})
|
||||
.optional(),
|
||||
|
||||
/**
|
||||
* scheduler is not requires by all i2i providers
|
||||
*/
|
||||
scheduler: z
|
||||
.object({
|
||||
default: z.string(),
|
||||
description: z.string().optional(),
|
||||
enum: z.array(z.string()).optional(),
|
||||
type: z.literal('string').optional(),
|
||||
})
|
||||
.optional(),
|
||||
|
||||
height: z
|
||||
.object({
|
||||
default: z.number(),
|
||||
|
|
@ -136,6 +160,20 @@ export const ModelParamsMetaSchema = z.object({
|
|||
})
|
||||
.optional(),
|
||||
|
||||
/**
|
||||
* strength/denoise is optional for t2i but must be used for i2i
|
||||
*/
|
||||
strength: z
|
||||
.object({
|
||||
default: z.number(),
|
||||
description: z.string().optional(),
|
||||
max: z.number().optional().default(1),
|
||||
min: z.number().optional().default(0),
|
||||
step: z.number().optional().default(0.05),
|
||||
type: z.literal('number').optional(),
|
||||
})
|
||||
.optional(),
|
||||
|
||||
steps: z
|
||||
.object({
|
||||
default: z.number(),
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ import {
|
|||
TextToImagePayload,
|
||||
TextToSpeechPayload,
|
||||
} from '../types';
|
||||
import { CreateImagePayload } from '../types/image';
|
||||
import { AuthenticatedImageRuntime, CreateImagePayload } from '../types/image';
|
||||
import { LobeRuntimeAI } from './BaseAI';
|
||||
|
||||
export interface AgentChatOptions {
|
||||
|
|
@ -92,6 +92,13 @@ export class ModelRuntime {
|
|||
return this._runtime.pullModel?.(params, options);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get authentication headers if runtime supports it
|
||||
*/
|
||||
getAuthHeaders(): Record<string, string> | undefined {
|
||||
return (this._runtime as AuthenticatedImageRuntime).getAuthHeaders?.();
|
||||
}
|
||||
|
||||
/**
|
||||
* @description Initialize the runtime with the provider and the options
|
||||
* @param provider choose a model provider
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ export { LobeBedrockAI } from './providers/bedrock';
|
|||
export { LobeBflAI } from './providers/bfl';
|
||||
export { LobeCerebrasAI } from './providers/cerebras';
|
||||
export { LobeCometAPIAI } from './providers/cometapi';
|
||||
export { LobeComfyUI } from './providers/comfyui';
|
||||
export { LobeDeepSeekAI } from './providers/deepseek';
|
||||
export { LobeGoogleAI } from './providers/google';
|
||||
export { LobeGroq } from './providers/groq';
|
||||
|
|
@ -39,3 +40,4 @@ export { AgentRuntimeError } from './utils/createError';
|
|||
export { getModelPropertyWithFallback } from './utils/getFallbackModelProperty';
|
||||
export { getModelPricing } from './utils/getModelPricing';
|
||||
export { parseDataUri } from './utils/uriParser';
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,521 @@
|
|||
// @vitest-environment node
|
||||
import { createBasicAuthCredentials } from '@lobechat/utils';
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import type { ComfyUIKeyVault } from '@/types/index';
|
||||
|
||||
import type { CreateImagePayload } from '../../../types/image';
|
||||
import { LobeComfyUI } from '../index';
|
||||
|
||||
// Mock debug
|
||||
vi.mock('debug', () => ({
|
||||
default: vi.fn(() => vi.fn()),
|
||||
}));
|
||||
|
||||
describe('LobeComfyUI Runtime', () => {
|
||||
let runtime: LobeComfyUI;
|
||||
let mockFetch: ReturnType<typeof vi.fn>;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
// Mock global fetch
|
||||
mockFetch = vi.fn();
|
||||
vi.stubGlobal('fetch', mockFetch);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.unstubAllGlobals();
|
||||
});
|
||||
|
||||
describe('constructor', () => {
|
||||
it('should initialize with default options', () => {
|
||||
runtime = new LobeComfyUI();
|
||||
|
||||
expect(runtime.baseURL).toBe('http://localhost:8188');
|
||||
});
|
||||
|
||||
it('should initialize with custom baseURL', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
baseURL: 'https://custom.comfyui.com',
|
||||
};
|
||||
|
||||
runtime = new LobeComfyUI(options);
|
||||
|
||||
expect(runtime.baseURL).toBe('https://custom.comfyui.com');
|
||||
});
|
||||
|
||||
it('should use environment variable if no baseURL provided', () => {
|
||||
const originalEnv = process.env.COMFYUI_DEFAULT_URL;
|
||||
process.env.COMFYUI_DEFAULT_URL = 'https://env.comfyui.com';
|
||||
|
||||
runtime = new LobeComfyUI();
|
||||
|
||||
expect(runtime.baseURL).toBe('https://env.comfyui.com');
|
||||
|
||||
// Restore environment
|
||||
if (originalEnv === undefined) {
|
||||
delete process.env.COMFYUI_DEFAULT_URL;
|
||||
} else {
|
||||
process.env.COMFYUI_DEFAULT_URL = originalEnv;
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAuthHeaders', () => {
|
||||
it('should return undefined for no auth', () => {
|
||||
runtime = new LobeComfyUI({ authType: 'none' });
|
||||
|
||||
const headers = runtime.getAuthHeaders();
|
||||
|
||||
expect(headers).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return undefined for default auth type', () => {
|
||||
runtime = new LobeComfyUI();
|
||||
|
||||
const headers = runtime.getAuthHeaders();
|
||||
|
||||
expect(headers).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return Basic auth headers when configured correctly', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'basic',
|
||||
username: 'testuser',
|
||||
password: 'testpass',
|
||||
};
|
||||
|
||||
runtime = new LobeComfyUI(options);
|
||||
|
||||
const headers = runtime.getAuthHeaders();
|
||||
|
||||
expect(headers).toEqual({
|
||||
Authorization: `Basic ${createBasicAuthCredentials('testuser', 'testpass')}`,
|
||||
});
|
||||
});
|
||||
|
||||
it('should return undefined for basic auth without credentials', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'basic',
|
||||
};
|
||||
|
||||
runtime = new LobeComfyUI(options);
|
||||
|
||||
const headers = runtime.getAuthHeaders();
|
||||
|
||||
expect(headers).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return Bearer auth headers when configured correctly', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'bearer',
|
||||
apiKey: 'test-api-key',
|
||||
};
|
||||
|
||||
runtime = new LobeComfyUI(options);
|
||||
|
||||
const headers = runtime.getAuthHeaders();
|
||||
|
||||
expect(headers).toEqual({
|
||||
Authorization: 'Bearer test-api-key',
|
||||
});
|
||||
});
|
||||
|
||||
it('should return undefined for bearer auth without apiKey', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'bearer',
|
||||
};
|
||||
|
||||
runtime = new LobeComfyUI(options);
|
||||
|
||||
const headers = runtime.getAuthHeaders();
|
||||
|
||||
expect(headers).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return custom headers when configured', () => {
|
||||
const customHeaders = {
|
||||
'X-Custom-Auth': 'custom-value',
|
||||
'Authorization': 'Custom auth-token',
|
||||
};
|
||||
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'custom',
|
||||
customHeaders,
|
||||
};
|
||||
|
||||
runtime = new LobeComfyUI(options);
|
||||
|
||||
const headers = runtime.getAuthHeaders();
|
||||
|
||||
expect(headers).toEqual(customHeaders);
|
||||
});
|
||||
|
||||
it('should return undefined for custom auth without headers', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'custom',
|
||||
};
|
||||
|
||||
runtime = new LobeComfyUI(options);
|
||||
|
||||
const headers = runtime.getAuthHeaders();
|
||||
|
||||
expect(headers).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('createImage', () => {
|
||||
beforeEach(() => {
|
||||
runtime = new LobeComfyUI({
|
||||
baseURL: 'https://test.comfyui.com',
|
||||
authType: 'bearer',
|
||||
apiKey: 'test-key',
|
||||
});
|
||||
});
|
||||
|
||||
it('should call WebAPI endpoint with correct URL and payload', async () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'flux1-dev.safetensors',
|
||||
params: {
|
||||
prompt: 'a beautiful landscape',
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
steps: 20,
|
||||
cfg: 7,
|
||||
},
|
||||
};
|
||||
|
||||
const mockResponse = {
|
||||
imageUrl: 'https://test.comfyui.com/image/output.png',
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
|
||||
// Mock successful response
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue(mockResponse),
|
||||
});
|
||||
|
||||
// Set APP_URL environment variable
|
||||
const originalAppUrl = process.env.APP_URL;
|
||||
process.env.APP_URL = 'http://localhost:3010';
|
||||
|
||||
const result = await runtime.createImage(mockPayload);
|
||||
|
||||
// Verify fetch was called with correct parameters
|
||||
expect(mockFetch).toHaveBeenCalledWith('http://localhost:3010/webapi/create-image/comfyui', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': 'Bearer test-key',
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: 'flux1-dev.safetensors',
|
||||
options: {
|
||||
baseURL: 'https://test.comfyui.com',
|
||||
authType: 'bearer',
|
||||
apiKey: 'test-key',
|
||||
},
|
||||
params: {
|
||||
prompt: 'a beautiful landscape',
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
steps: 20,
|
||||
cfg: 7,
|
||||
},
|
||||
}),
|
||||
});
|
||||
|
||||
expect(result).toEqual(mockResponse);
|
||||
|
||||
// Restore environment
|
||||
if (originalAppUrl === undefined) {
|
||||
delete process.env.APP_URL;
|
||||
} else {
|
||||
process.env.APP_URL = originalAppUrl;
|
||||
}
|
||||
});
|
||||
|
||||
it('should use default APP_URL when environment variable is not set', async () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'flux1-dev.safetensors',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Mock successful response
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ imageUrl: 'test.png' }),
|
||||
});
|
||||
|
||||
// Ensure APP_URL is not set
|
||||
const originalAppUrl = process.env.APP_URL;
|
||||
const originalPort = process.env.PORT;
|
||||
delete process.env.APP_URL;
|
||||
process.env.PORT = '3000';
|
||||
|
||||
await runtime.createImage(mockPayload);
|
||||
|
||||
// Should use default localhost:3000
|
||||
expect(mockFetch).toHaveBeenCalledWith(
|
||||
'http://localhost:3000/webapi/create-image/comfyui',
|
||||
expect.any(Object),
|
||||
);
|
||||
|
||||
// Restore environment
|
||||
if (originalAppUrl !== undefined) {
|
||||
process.env.APP_URL = originalAppUrl;
|
||||
}
|
||||
if (originalPort === undefined) {
|
||||
delete process.env.PORT;
|
||||
} else {
|
||||
process.env.PORT = originalPort;
|
||||
}
|
||||
});
|
||||
|
||||
it('should use default port 3010 when PORT is not set', async () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'flux1-dev.safetensors',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Mock successful response
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ imageUrl: 'test.png' }),
|
||||
});
|
||||
|
||||
// Ensure both APP_URL and PORT are not set
|
||||
const originalAppUrl = process.env.APP_URL;
|
||||
const originalPort = process.env.PORT;
|
||||
delete process.env.APP_URL;
|
||||
delete process.env.PORT;
|
||||
|
||||
await runtime.createImage(mockPayload);
|
||||
|
||||
// Should use default localhost:3010
|
||||
expect(mockFetch).toHaveBeenCalledWith(
|
||||
'http://localhost:3010/webapi/create-image/comfyui',
|
||||
expect.any(Object),
|
||||
);
|
||||
|
||||
// Restore environment
|
||||
if (originalAppUrl !== undefined) {
|
||||
process.env.APP_URL = originalAppUrl;
|
||||
}
|
||||
if (originalPort !== undefined) {
|
||||
process.env.PORT = originalPort;
|
||||
}
|
||||
});
|
||||
|
||||
it('should include auth headers when configured', async () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'test-model.safetensors',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Test with basic auth
|
||||
runtime = new LobeComfyUI({
|
||||
baseURL: 'https://test.comfyui.com',
|
||||
authType: 'basic',
|
||||
username: 'testuser',
|
||||
password: 'testpass',
|
||||
});
|
||||
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ imageUrl: 'test.png' }),
|
||||
});
|
||||
|
||||
await runtime.createImage(mockPayload);
|
||||
|
||||
expect(mockFetch).toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
expect.objectContaining({
|
||||
headers: expect.objectContaining({
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': `Basic ${createBasicAuthCredentials('testuser', 'testpass')}`,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should not include auth headers when auth is disabled', async () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'test-model.safetensors',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Test with no auth
|
||||
runtime = new LobeComfyUI({
|
||||
baseURL: 'https://test.comfyui.com',
|
||||
authType: 'none',
|
||||
});
|
||||
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ imageUrl: 'test.png' }),
|
||||
});
|
||||
|
||||
await runtime.createImage(mockPayload);
|
||||
|
||||
expect(mockFetch).toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
expect.objectContaining({
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw error when fetch response is not ok', async () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'flux1-dev.safetensors',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Mock error response
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: false,
|
||||
status: 500,
|
||||
text: vi.fn().mockResolvedValue('Internal server error'),
|
||||
});
|
||||
|
||||
await expect(runtime.createImage(mockPayload)).rejects.toMatchObject({
|
||||
errorType: 'ComfyUIServiceUnavailable',
|
||||
provider: 'comfyui',
|
||||
});
|
||||
|
||||
expect(mockFetch).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should throw error when fetch throws', async () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'flux1-dev.safetensors',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
const networkError = new Error('Network connection failed');
|
||||
mockFetch.mockRejectedValue(networkError);
|
||||
|
||||
await expect(runtime.createImage(mockPayload)).rejects.toMatchObject({
|
||||
errorType: 'ComfyUIBizError',
|
||||
provider: 'comfyui',
|
||||
error: {
|
||||
message: 'Network connection failed',
|
||||
},
|
||||
});
|
||||
|
||||
expect(mockFetch).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should handle complex payload with all parameters', async () => {
|
||||
const complexPayload: CreateImagePayload = {
|
||||
model: 'sd3.5-large.safetensors',
|
||||
params: {
|
||||
prompt: 'complex image generation with multiple parameters',
|
||||
width: 1152,
|
||||
height: 896,
|
||||
steps: 25,
|
||||
cfg: 8.5,
|
||||
seed: 12345,
|
||||
},
|
||||
};
|
||||
|
||||
const mockResponse = {
|
||||
imageUrl: 'https://test.comfyui.com/complex-image.png',
|
||||
width: 1152,
|
||||
height: 896,
|
||||
};
|
||||
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue(mockResponse),
|
||||
});
|
||||
|
||||
const result = await runtime.createImage(complexPayload);
|
||||
|
||||
expect(result).toEqual(mockResponse);
|
||||
|
||||
// Verify that complex payload was passed correctly
|
||||
expect(mockFetch).toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
expect.objectContaining({
|
||||
body: JSON.stringify({
|
||||
model: 'sd3.5-large.safetensors',
|
||||
options: {
|
||||
baseURL: 'https://test.comfyui.com',
|
||||
authType: 'bearer',
|
||||
apiKey: 'test-key',
|
||||
},
|
||||
params: {
|
||||
prompt: 'complex image generation with multiple parameters',
|
||||
width: 1152,
|
||||
height: 896,
|
||||
steps: 25,
|
||||
cfg: 8.5,
|
||||
seed: 12345,
|
||||
},
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle WebAPI error responses with JSON body', async () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'flux1-dev.safetensors',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Mock error response with JSON body
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: false,
|
||||
status: 400,
|
||||
text: vi
|
||||
.fn()
|
||||
.mockResolvedValue('{"message":"Invalid model specified","error":"Model not found"}'),
|
||||
});
|
||||
|
||||
await expect(runtime.createImage(mockPayload)).rejects.toMatchObject({
|
||||
errorType: 'ComfyUIBizError',
|
||||
provider: 'comfyui',
|
||||
error: {
|
||||
message: 'Invalid model specified',
|
||||
},
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('runtime interface compliance', () => {
|
||||
it('should implement LobeRuntimeAI interface', () => {
|
||||
runtime = new LobeComfyUI();
|
||||
|
||||
expect(runtime).toHaveProperty('baseURL');
|
||||
expect(typeof runtime.createImage).toBe('function');
|
||||
});
|
||||
|
||||
it('should implement AuthenticatedImageRuntime interface', () => {
|
||||
runtime = new LobeComfyUI();
|
||||
|
||||
expect(typeof runtime.getAuthHeaders).toBe('function');
|
||||
});
|
||||
});
|
||||
});
|
||||
116
packages/model-runtime/src/providers/comfyui/auth/AuthManager.ts
Normal file
116
packages/model-runtime/src/providers/comfyui/auth/AuthManager.ts
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
import { createBasicAuthCredentials } from '@lobechat/utils';
|
||||
|
||||
import type { ComfyUIKeyVault } from '@/types/index';
|
||||
|
||||
export interface BasicCredentials {
|
||||
password: string;
|
||||
type: 'basic';
|
||||
username: string;
|
||||
}
|
||||
|
||||
export interface BearerTokenCredentials {
|
||||
apiKey: string;
|
||||
type: 'bearer';
|
||||
}
|
||||
|
||||
export interface CustomCredentials {
|
||||
customHeaders: Record<string, string>;
|
||||
type: 'custom';
|
||||
}
|
||||
|
||||
/**
|
||||
* ComfyUI Authentication Manager
|
||||
* Handles authentication headers generation for ComfyUI requests
|
||||
*/
|
||||
export class AuthManager {
|
||||
private credentials: BasicCredentials | BearerTokenCredentials | CustomCredentials | undefined;
|
||||
private authHeaders: Record<string, string> | undefined;
|
||||
|
||||
constructor(options: ComfyUIKeyVault) {
|
||||
this.validateOptions(options);
|
||||
this.credentials = this.createCredentials(options);
|
||||
this.authHeaders = this.createAuthHeaders(options);
|
||||
}
|
||||
|
||||
getAuthHeaders(): Record<string, string> | undefined {
|
||||
return this.authHeaders;
|
||||
}
|
||||
|
||||
private validateOptions(options: ComfyUIKeyVault): void {
|
||||
const { authType = 'none', apiKey, username, password, customHeaders } = options;
|
||||
|
||||
switch (authType) {
|
||||
case 'basic': {
|
||||
if (!username || !password) {
|
||||
throw new TypeError('Username and password are required for basic authentication');
|
||||
}
|
||||
break;
|
||||
}
|
||||
case 'bearer': {
|
||||
if (!apiKey) {
|
||||
throw new TypeError('API key is required for bearer token authentication');
|
||||
}
|
||||
break;
|
||||
}
|
||||
case 'custom': {
|
||||
if (!customHeaders || Object.keys(customHeaders).length === 0) {
|
||||
throw new TypeError('Custom headers are required for custom authentication');
|
||||
}
|
||||
break;
|
||||
}
|
||||
case 'none': {
|
||||
// No validation needed for none authentication
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
throw new TypeError(`Unsupported authentication type: ${authType}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private createCredentials(
|
||||
options: ComfyUIKeyVault,
|
||||
): BasicCredentials | BearerTokenCredentials | CustomCredentials | undefined {
|
||||
const { authType = 'none', apiKey, username, password, customHeaders } = options;
|
||||
|
||||
switch (authType) {
|
||||
case 'basic': {
|
||||
return { password: password!, type: 'basic', username: username! };
|
||||
}
|
||||
case 'bearer': {
|
||||
return { apiKey: apiKey!, type: 'bearer' };
|
||||
}
|
||||
case 'custom': {
|
||||
return { customHeaders: customHeaders!, type: 'custom' };
|
||||
}
|
||||
case 'none': {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private createAuthHeaders(options: ComfyUIKeyVault): Record<string, string> | undefined {
|
||||
const { authType = 'none', apiKey, username, password, customHeaders } = options;
|
||||
|
||||
switch (authType) {
|
||||
case 'basic': {
|
||||
if (!username || !password) return undefined;
|
||||
const credentials = createBasicAuthCredentials(username, password);
|
||||
return { Authorization: `Basic ${credentials}` };
|
||||
}
|
||||
|
||||
case 'bearer': {
|
||||
if (!apiKey) return undefined;
|
||||
return { Authorization: `Bearer ${apiKey}` };
|
||||
}
|
||||
|
||||
case 'custom': {
|
||||
return customHeaders || undefined;
|
||||
}
|
||||
|
||||
case 'none': {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
180
packages/model-runtime/src/providers/comfyui/index.ts
Normal file
180
packages/model-runtime/src/providers/comfyui/index.ts
Normal file
|
|
@ -0,0 +1,180 @@
|
|||
import { createBasicAuthCredentials } from '@lobechat/utils';
|
||||
import debug from 'debug';
|
||||
|
||||
import type { ComfyUIKeyVault } from '@/types/index';
|
||||
|
||||
import { LobeRuntimeAI } from '../../core/BaseAI';
|
||||
import {
|
||||
AuthenticatedImageRuntime,
|
||||
CreateImagePayload,
|
||||
CreateImageResponse,
|
||||
} from '../../types/image';
|
||||
import { parseComfyUIErrorMessage } from '../../utils/comfyuiErrorParser';
|
||||
import { AgentRuntimeError } from '../../utils/createError';
|
||||
|
||||
const log = debug('lobe-image:comfyui');
|
||||
|
||||
/**
|
||||
* ComfyUI Runtime implementation
|
||||
* Supports text-to-image and image editing
|
||||
*/
|
||||
export class LobeComfyUI implements LobeRuntimeAI, AuthenticatedImageRuntime {
|
||||
private options: ComfyUIKeyVault;
|
||||
baseURL: string;
|
||||
|
||||
constructor(options: ComfyUIKeyVault = {}) {
|
||||
log('🏗️ ComfyUI Runtime initialized');
|
||||
|
||||
this.options = options;
|
||||
this.baseURL = options.baseURL || process.env.COMFYUI_DEFAULT_URL || 'http://localhost:8188';
|
||||
|
||||
log('✅ ComfyUI Runtime ready - baseURL: %s', this.baseURL);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get authentication headers for image download
|
||||
* Used by framework for authenticated image downloads
|
||||
*/
|
||||
getAuthHeaders(): Record<string, string> | undefined {
|
||||
log('🔐 Providing auth headers for image download');
|
||||
|
||||
const { authType = 'none', apiKey, username, password, customHeaders } = this.options;
|
||||
|
||||
switch (authType) {
|
||||
case 'basic': {
|
||||
if (username && password) {
|
||||
return { Authorization: `Basic ${createBasicAuthCredentials(username, password)}` };
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
case 'bearer': {
|
||||
if (apiKey) {
|
||||
return { Authorization: `Bearer ${apiKey}` };
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
case 'custom': {
|
||||
return customHeaders || undefined;
|
||||
}
|
||||
|
||||
case 'none': {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create image using internal API endpoint
|
||||
* Always uses full URL for consistency across environments
|
||||
*/
|
||||
async createImage(payload: CreateImagePayload): Promise<CreateImageResponse> {
|
||||
log('🎨 Creating image with model: %s', payload.model);
|
||||
|
||||
try {
|
||||
// Determine app URL with Vercel support
|
||||
const isInVercel = process.env.VERCEL === '1';
|
||||
const vercelUrl = `https://${process.env.VERCEL_URL}`;
|
||||
const appUrl =
|
||||
process.env.APP_URL ||
|
||||
(isInVercel ? vercelUrl : `http://localhost:${process.env.PORT || 3010}`);
|
||||
|
||||
// Build headers with authentication
|
||||
const headers: Record<string, string> = {
|
||||
'Content-Type': 'application/json',
|
||||
...this.getAuthHeaders(),
|
||||
};
|
||||
|
||||
// In development mode, use debug header to bypass auth
|
||||
if (process.env.NODE_ENV === 'development') {
|
||||
headers['lobe-auth-dev-backend-api'] = '1';
|
||||
}
|
||||
|
||||
// If KEY_VAULTS_SECRET is available (server-side), use it for internal service auth
|
||||
// But only if it's actually set (not empty string)
|
||||
const keyVaultSecret = process.env.KEY_VAULTS_SECRET;
|
||||
if (keyVaultSecret && keyVaultSecret.trim() !== '') {
|
||||
headers['Authorization'] = `Bearer ${keyVaultSecret}`;
|
||||
}
|
||||
|
||||
const response = await fetch(`${appUrl}/webapi/create-image/comfyui`, {
|
||||
body: JSON.stringify({
|
||||
model: payload.model,
|
||||
options: this.options,
|
||||
params: payload.params,
|
||||
}),
|
||||
headers,
|
||||
method: 'POST',
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorText = await response.text();
|
||||
let errorData: any;
|
||||
|
||||
try {
|
||||
errorData = JSON.parse(errorText);
|
||||
} catch {
|
||||
// If not JSON, use the text as error message
|
||||
errorData = { message: errorText, status: response.status };
|
||||
}
|
||||
|
||||
// Check if it's already an AgentRuntimeError from WebAPI
|
||||
if (
|
||||
errorData &&
|
||||
typeof errorData === 'object' &&
|
||||
'errorType' in errorData &&
|
||||
'error' in errorData &&
|
||||
'provider' in errorData
|
||||
) {
|
||||
// Already a properly formatted AgentRuntimeError from WebAPI
|
||||
// Reconstruct the error using the framework's method to ensure proper type
|
||||
throw AgentRuntimeError.createImage({
|
||||
error: errorData.error,
|
||||
errorType: errorData.errorType,
|
||||
provider: errorData.provider,
|
||||
});
|
||||
}
|
||||
|
||||
// Otherwise parse and create new error
|
||||
const { error: parsedError, errorType } = parseComfyUIErrorMessage(errorData);
|
||||
|
||||
throw AgentRuntimeError.createImage({
|
||||
error: parsedError,
|
||||
errorType,
|
||||
provider: 'comfyui',
|
||||
});
|
||||
}
|
||||
|
||||
const result = await response.json();
|
||||
log('✅ ComfyUI image created successfully');
|
||||
return result;
|
||||
} catch (error) {
|
||||
log('❌ ComfyUI createImage error: %O', error);
|
||||
|
||||
// If it looks like an AgentRuntimeError object structure (already processed), reconstruct it
|
||||
if (
|
||||
error &&
|
||||
typeof error === 'object' &&
|
||||
'errorType' in error &&
|
||||
'error' in error &&
|
||||
'provider' in error
|
||||
) {
|
||||
throw AgentRuntimeError.createImage({
|
||||
error: (error as any).error,
|
||||
errorType: (error as any).errorType,
|
||||
provider: (error as any).provider,
|
||||
});
|
||||
}
|
||||
|
||||
// Otherwise parse and format the error
|
||||
const { error: parsedError, errorType } = parseComfyUIErrorMessage(error);
|
||||
|
||||
throw AgentRuntimeError.createImage({
|
||||
error: parsedError,
|
||||
errorType,
|
||||
provider: 'comfyui',
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -13,6 +13,7 @@ import { LobeCerebrasAI } from './providers/cerebras';
|
|||
import { LobeCloudflareAI } from './providers/cloudflare';
|
||||
import { LobeCohereAI } from './providers/cohere';
|
||||
import { LobeCometAPIAI } from './providers/cometapi';
|
||||
import { LobeComfyUI } from './providers/comfyui';
|
||||
import { LobeDeepSeekAI } from './providers/deepseek';
|
||||
import { LobeFalAI } from './providers/fal';
|
||||
import { LobeFireworksAI } from './providers/fireworksai';
|
||||
|
|
@ -79,6 +80,7 @@ export const providerRuntimeMap = {
|
|||
cloudflare: LobeCloudflareAI,
|
||||
cohere: LobeCohereAI,
|
||||
cometapi: LobeCometAPIAI,
|
||||
comfyui: LobeComfyUI,
|
||||
deepseek: LobeDeepSeekAI,
|
||||
fal: LobeFalAI,
|
||||
fireworksai: LobeFireworksAI,
|
||||
|
|
|
|||
|
|
@ -19,6 +19,14 @@ export const AgentRuntimeErrorType = {
|
|||
OllamaBizError: 'OllamaBizError',
|
||||
OllamaServiceUnavailable: 'OllamaServiceUnavailable',
|
||||
|
||||
InvalidComfyUIArgs: 'InvalidComfyUIArgs',
|
||||
ComfyUIBizError: 'ComfyUIBizError',
|
||||
ComfyUIServiceUnavailable: 'ComfyUIServiceUnavailable',
|
||||
ComfyUIEmptyResult: 'ComfyUIEmptyResult',
|
||||
ComfyUIUploadFailed: 'ComfyUIUploadFailed',
|
||||
ComfyUIWorkflowError: 'ComfyUIWorkflowError',
|
||||
ComfyUIModelError: 'ComfyUIModelError',
|
||||
|
||||
InvalidBedrockCredentials: 'InvalidBedrockCredentials',
|
||||
InvalidVertexCredentials: 'InvalidVertexCredentials',
|
||||
StreamChunkError: 'StreamChunkError',
|
||||
|
|
|
|||
|
|
@ -34,3 +34,12 @@ export type CreateImageResponse = {
|
|||
*/
|
||||
modelUsage?: ModelUsage;
|
||||
};
|
||||
|
||||
// 新增：支持认证图片下载的运行时接口
|
||||
export interface AuthenticatedImageRuntime {
|
||||
/**
|
||||
* Get authentication headers for image download
|
||||
* Used when the image server requires authentication
|
||||
*/
|
||||
getAuthHeaders(): Record<string, string> | undefined;
|
||||
}
|
||||
|
|
|
|||
369
packages/model-runtime/src/utils/comfyuiErrorParser.test.ts
Normal file
369
packages/model-runtime/src/utils/comfyuiErrorParser.test.ts
Normal file
|
|
@ -0,0 +1,369 @@
|
|||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import { AgentRuntimeErrorType } from '../types/error';
|
||||
import { cleanComfyUIErrorMessage, parseComfyUIErrorMessage } from './comfyuiErrorParser';
|
||||
|
||||
describe('comfyuiErrorParser', () => {
|
||||
describe('cleanComfyUIErrorMessage', () => {
|
||||
it('should remove leading asterisks and spaces', () => {
|
||||
const message = '* Error message';
|
||||
expect(cleanComfyUIErrorMessage(message)).toBe('Error message');
|
||||
|
||||
// Test multiple asterisks
|
||||
const multiAsterisk = '* * * Error message';
|
||||
expect(cleanComfyUIErrorMessage(multiAsterisk)).toBe('* * Error message');
|
||||
});
|
||||
|
||||
it('should convert escaped newlines', () => {
|
||||
const message = 'Line 1\\nLine 2';
|
||||
expect(cleanComfyUIErrorMessage(message)).toBe('Line 1 Line 2');
|
||||
});
|
||||
|
||||
it('should replace multiple newlines with single space', () => {
|
||||
const message = 'Line 1\n\n\nLine 2';
|
||||
expect(cleanComfyUIErrorMessage(message)).toBe('Line 1 Line 2');
|
||||
});
|
||||
|
||||
it('should trim leading and trailing spaces', () => {
|
||||
const message = ' Error message ';
|
||||
expect(cleanComfyUIErrorMessage(message)).toBe('Error message');
|
||||
});
|
||||
});
|
||||
|
||||
describe('parseComfyUIErrorMessage', () => {
|
||||
describe('HTTP status code errors', () => {
|
||||
it('should identify 401 as InvalidProviderAPIKey', () => {
|
||||
const error = { message: 'Unauthorized', status: 401 };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.InvalidProviderAPIKey);
|
||||
});
|
||||
|
||||
it('should identify 403 as PermissionDenied', () => {
|
||||
const error = { message: 'Forbidden', status: 403 };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.PermissionDenied);
|
||||
});
|
||||
|
||||
it('should identify 404 as InvalidProviderAPIKey', () => {
|
||||
const error = { message: 'Not Found', status: 404 };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.InvalidProviderAPIKey);
|
||||
});
|
||||
|
||||
it('should identify 500+ as ComfyUIServiceUnavailable', () => {
|
||||
const error = { message: 'Internal Server Error', status: 500 };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIServiceUnavailable);
|
||||
});
|
||||
|
||||
it('should identify HTTP status in message when status field missing', () => {
|
||||
const error = { message: 'Request failed with HTTP 401' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.InvalidProviderAPIKey);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Network errors', () => {
|
||||
it('should return ComfyUIBizError for fetch failed (processed by server)', () => {
|
||||
const error = new Error('fetch failed');
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Network error detection moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('fetch failed');
|
||||
});
|
||||
|
||||
it('should return ComfyUIBizError for ECONNREFUSED (processed by server)', () => {
|
||||
const error = { message: 'Connection ECONNREFUSED', code: 'ECONNREFUSED' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Network error detection moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('Connection ECONNREFUSED');
|
||||
});
|
||||
|
||||
it('should return ComfyUIBizError for WebSocket errors (processed by server)', () => {
|
||||
const error = { message: 'WebSocket connection failed', code: 'WS_CONNECTION_FAILED' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Network error detection moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('WebSocket connection failed');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Model errors', () => {
|
||||
it('should return ComfyUIBizError for model not found (processed by server)', () => {
|
||||
const error = { message: 'Model not found: flux1-dev.safetensors' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Model error detection moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('Model not found: flux1-dev.safetensors');
|
||||
});
|
||||
|
||||
it('should return ComfyUIBizError for checkpoint not found (processed by server)', () => {
|
||||
const error = { message: 'Checkpoint not found' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Model error detection moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('Checkpoint not found');
|
||||
});
|
||||
|
||||
it('should return ComfyUIBizError for safetensors file errors (processed by server)', () => {
|
||||
const error = { message: 'Missing file: model.safetensors' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Model error detection moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('Missing file: model.safetensors');
|
||||
});
|
||||
|
||||
it('should preserve server-provided file info but return ComfyUIBizError', () => {
|
||||
const error = {
|
||||
message: 'Some error',
|
||||
missingFileName: 'flux1-dev.safetensors',
|
||||
missingFileType: 'model',
|
||||
};
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// File info is preserved but error type detection is server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.missingFileName).toBe('flux1-dev.safetensors');
|
||||
expect(result.error.missingFileType).toBe('model');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Workflow errors', () => {
|
||||
it('should return ComfyUIBizError for workflow validation errors (processed by server)', () => {
|
||||
const error = { message: 'Workflow validation failed' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Workflow error detection moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('Workflow validation failed');
|
||||
});
|
||||
|
||||
it('should return ComfyUIBizError for node execution errors (processed by server)', () => {
|
||||
const error = {
|
||||
message: 'Node execution failed',
|
||||
node_id: '5',
|
||||
node_type: 'KSampler',
|
||||
};
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Workflow error detection moved to server-side, but node info is preserved
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('Node execution failed');
|
||||
expect(result.error.details).toEqual({
|
||||
node_id: '5',
|
||||
node_type: 'KSampler',
|
||||
});
|
||||
});
|
||||
|
||||
it('should return ComfyUIBizError for queue errors (processed by server)', () => {
|
||||
const error = { message: 'Queue processing error' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Workflow error detection moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('Queue processing error');
|
||||
});
|
||||
});
|
||||
|
||||
describe('SDK custom errors', () => {
|
||||
it('should identify SDK error classes', () => {
|
||||
const error = {
|
||||
name: 'ExecutionFailedError',
|
||||
message: 'Execution failed',
|
||||
};
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
});
|
||||
|
||||
it('should identify SDK error messages', () => {
|
||||
const error = { message: 'SDK Error: Invalid configuration' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
});
|
||||
});
|
||||
|
||||
describe('JSON parsing errors', () => {
|
||||
it('should return ComfyUIBizError for SyntaxError (SyntaxError detection moved to server)', () => {
|
||||
const error = new SyntaxError('Unexpected token < in JSON at position 0');
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// SyntaxError detection and message enhancement moved to server-side
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('Unexpected token < in JSON at position 0');
|
||||
expect(result.error.type).toBe('SyntaxError');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Error information extraction', () => {
|
||||
it('should extract error info from string', () => {
|
||||
const error = 'Simple error message';
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.error.message).toBe('Simple error message');
|
||||
});
|
||||
|
||||
it('should extract error info from Error object', () => {
|
||||
const error = new Error('Error message');
|
||||
(error as any).code = 'ERROR_CODE';
|
||||
(error as any).status = 500;
|
||||
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.error.message).toBe('Error message');
|
||||
expect(result.error.code).toBe('ERROR_CODE');
|
||||
expect(result.error.status).toBe(500);
|
||||
expect(result.error.type).toBe('Error');
|
||||
});
|
||||
|
||||
it('should extract error info from structured object', () => {
|
||||
const error = {
|
||||
message: 'Error message',
|
||||
code: 'ERROR_CODE',
|
||||
status: 400,
|
||||
details: { foo: 'bar' },
|
||||
node_id: '5',
|
||||
node_type: 'KSampler',
|
||||
};
|
||||
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.error.message).toBe('Error message');
|
||||
expect(result.error.code).toBe('ERROR_CODE');
|
||||
expect(result.error.status).toBe(400);
|
||||
expect(result.error.details).toEqual({
|
||||
foo: 'bar',
|
||||
node_id: '5',
|
||||
node_type: 'KSampler',
|
||||
});
|
||||
});
|
||||
|
||||
it('should preserve server-generated file info and guidance', () => {
|
||||
const error = {
|
||||
message: 'Model file missing',
|
||||
missingFileName: 'flux1-dev.safetensors',
|
||||
missingFileType: 'model' as const,
|
||||
userGuidance: 'Please download the model from...',
|
||||
};
|
||||
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.error.missingFileName).toBe('flux1-dev.safetensors');
|
||||
expect(result.error.missingFileType).toBe('model');
|
||||
expect(result.error.userGuidance).toBe('Please download the model from...');
|
||||
});
|
||||
|
||||
it('should extract nested error info from various locations', () => {
|
||||
const error = {
|
||||
body: {
|
||||
error: {
|
||||
message: 'Nested error',
|
||||
missingFileName: 'ae.safetensors',
|
||||
userGuidance: 'Download VAE model',
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.error.missingFileName).toBe('ae.safetensors');
|
||||
expect(result.error.userGuidance).toBe('Download VAE model');
|
||||
});
|
||||
|
||||
it('should handle cause field (SDK pattern)', () => {
|
||||
const error = new Error('Wrapper error');
|
||||
(error as any).cause = {
|
||||
message: 'Actual error',
|
||||
code: 'ACTUAL_CODE',
|
||||
};
|
||||
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.error.message).toBe('Actual error');
|
||||
expect(result.error.code).toBe('ACTUAL_CODE');
|
||||
// Type comes from the cause object's constructor name (plain object = "Object")
|
||||
expect(result.error.type).toBe('Object');
|
||||
});
|
||||
|
||||
it('should extract message from various possible sources', () => {
|
||||
const error = {
|
||||
exception_message: 'ComfyUI exception',
|
||||
error: {
|
||||
message: 'Should not use this',
|
||||
},
|
||||
};
|
||||
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
// exception_message has highest priority
|
||||
expect(result.error.message).toBe('ComfyUI exception');
|
||||
});
|
||||
});
|
||||
|
||||
describe('AgentRuntimeError handling', () => {
|
||||
it('should detect and return AgentRuntimeError as-is', () => {
|
||||
const agentRuntimeError = {
|
||||
error: {
|
||||
message: 'Model not found',
|
||||
missingFileName: 'flux1-dev.safetensors',
|
||||
missingFileType: 'model',
|
||||
userGuidance: 'Please download the model',
|
||||
},
|
||||
errorType: AgentRuntimeErrorType.ModelNotFound,
|
||||
provider: 'comfyui',
|
||||
};
|
||||
|
||||
const result = parseComfyUIErrorMessage(agentRuntimeError);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ModelNotFound);
|
||||
expect(result.error).toEqual(agentRuntimeError.error);
|
||||
});
|
||||
|
||||
it('should handle AgentRuntimeError with InvalidProviderAPIKey', () => {
|
||||
const agentRuntimeError = {
|
||||
error: {
|
||||
message: 'Authentication failed',
|
||||
status: 401,
|
||||
},
|
||||
errorType: AgentRuntimeErrorType.InvalidProviderAPIKey,
|
||||
provider: 'comfyui',
|
||||
};
|
||||
|
||||
const result = parseComfyUIErrorMessage(agentRuntimeError);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.InvalidProviderAPIKey);
|
||||
expect(result.error.message).toBe('Authentication failed');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Default error handling', () => {
|
||||
it('should handle unknown error types', () => {
|
||||
const error = { random: 'data' };
|
||||
const result = parseComfyUIErrorMessage(error);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toContain('object');
|
||||
});
|
||||
|
||||
it('should handle null/undefined gracefully', () => {
|
||||
const result = parseComfyUIErrorMessage(null);
|
||||
|
||||
expect(result.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(result.error.message).toBe('null');
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
266
packages/model-runtime/src/utils/comfyuiErrorParser.ts
Normal file
266
packages/model-runtime/src/utils/comfyuiErrorParser.ts
Normal file
|
|
@ -0,0 +1,266 @@
|
|||
import { AgentRuntimeErrorType, ILobeAgentRuntimeErrorType } from '../types/error';
|
||||
|
||||
export interface ComfyUIError {
|
||||
code?: number | string;
|
||||
details?: any;
|
||||
message: string;
|
||||
missingFileName?: string;
|
||||
missingFileType?: 'model' | 'component';
|
||||
status?: number;
|
||||
type?: string;
|
||||
userGuidance?: string;
|
||||
}
|
||||
|
||||
export interface ParsedError {
|
||||
error: ComfyUIError;
|
||||
errorType: ILobeAgentRuntimeErrorType;
|
||||
}
|
||||
|
||||
/**
|
||||
* Clean ComfyUI error message by removing formatting characters and extra spaces
|
||||
* @param message - Original error message
|
||||
* @returns Cleaned error message
|
||||
*/
|
||||
export function cleanComfyUIErrorMessage(message: string): string {
|
||||
return message
|
||||
.replaceAll(/^\*\s*/gm, '') // Remove leading asterisks and spaces (multiline)
|
||||
.replaceAll('\\n', '\n') // Convert escaped newlines
|
||||
.replaceAll(/\n+/g, ' ') // Replace multiple newlines with single space
|
||||
.trim(); // Remove leading and trailing spaces
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract structured information from error object
|
||||
* Client-side version that preserves server-generated information
|
||||
* @param error - Original error object
|
||||
* @returns Structured ComfyUI error information
|
||||
*/
|
||||
function extractComfyUIErrorInfo(error: any): ComfyUIError {
|
||||
// Handle string errors
|
||||
if (typeof error === 'string') {
|
||||
const cleanedMessage = cleanComfyUIErrorMessage(error);
|
||||
|
||||
return {
|
||||
message: cleanedMessage,
|
||||
};
|
||||
}
|
||||
|
||||
// Handle Error objects - prioritize cause field (SDK pattern)
|
||||
if (error instanceof Error) {
|
||||
// Check if there's a cause field with actual error details (SDK pattern)
|
||||
if ((error as any).cause) {
|
||||
const cause = (error as any).cause;
|
||||
// Recursively extract error info from cause
|
||||
const causeInfo = extractComfyUIErrorInfo(cause);
|
||||
return {
|
||||
...causeInfo,
|
||||
// Preserve the original error type if cause doesn't have one
|
||||
type: causeInfo.type || error.name,
|
||||
};
|
||||
}
|
||||
|
||||
const cleanedMessage = cleanComfyUIErrorMessage(error.message);
|
||||
|
||||
return {
|
||||
code: (error as any).code,
|
||||
message: cleanedMessage,
|
||||
// Preserve server-generated file info and guidance
|
||||
missingFileName: (error as any).missingFileName,
|
||||
missingFileType: (error as any).missingFileType,
|
||||
status: (error as any).status || (error as any).statusCode,
|
||||
type: error.name,
|
||||
userGuidance: (error as any).userGuidance,
|
||||
};
|
||||
}
|
||||
|
||||
// Handle structured objects
|
||||
if (error && typeof error === 'object') {
|
||||
// Check for cause field first (SDK pattern)
|
||||
if (error.cause) {
|
||||
const causeInfo = extractComfyUIErrorInfo(error.cause);
|
||||
return {
|
||||
...causeInfo,
|
||||
type: causeInfo.type || error.type || error.name || error.constructor?.name,
|
||||
};
|
||||
}
|
||||
|
||||
// Extract message from various possible sources
|
||||
const possibleMessage = [
|
||||
error.exception_message, // ComfyUI specific field (highest priority)
|
||||
error.error?.exception_message, // Nested ComfyUI exception message
|
||||
error.error?.error, // Deeply nested error.error.error path
|
||||
error.message,
|
||||
error.error?.message,
|
||||
error.data?.message,
|
||||
error.body?.message,
|
||||
error.body?.error?.message,
|
||||
error.response?.data?.message,
|
||||
error.response?.data?.error?.message,
|
||||
error.response?.text,
|
||||
error.response?.body,
|
||||
error.statusText,
|
||||
].find(Boolean);
|
||||
|
||||
const message = possibleMessage || String(error);
|
||||
|
||||
// Extract status code from various possible locations
|
||||
const possibleStatus = [
|
||||
error.status,
|
||||
error.statusCode,
|
||||
error.details?.status, // ServicesError puts status in details
|
||||
error.response?.status,
|
||||
error.response?.statusCode,
|
||||
error.error?.status,
|
||||
error.error?.statusCode,
|
||||
].find(Number.isInteger);
|
||||
|
||||
const code = error.code || error.error?.code || error.response?.data?.code;
|
||||
|
||||
// Extract details including ComfyUI specific fields
|
||||
let details = error.response?.data || error.details || undefined;
|
||||
|
||||
// Include ComfyUI specific fields in details
|
||||
if (error.node_id || error.node_type || error.nodeId || error.nodeType || error.nodeName) {
|
||||
details = {
|
||||
...details,
|
||||
nodeName: error.nodeName,
|
||||
node_id: error.node_id || error.nodeId,
|
||||
node_type: error.node_type || error.nodeType,
|
||||
};
|
||||
}
|
||||
|
||||
const cleanedMessage = cleanComfyUIErrorMessage(message);
|
||||
|
||||
// Extract server-provided file info and guidance from various locations
|
||||
const missingFileName =
|
||||
error.missingFileName || error.body?.error?.missingFileName || error.error?.missingFileName;
|
||||
|
||||
const missingFileType =
|
||||
error.missingFileType || error.body?.error?.missingFileType || error.error?.missingFileType;
|
||||
|
||||
const userGuidance =
|
||||
error.userGuidance || error.body?.error?.userGuidance || error.error?.userGuidance;
|
||||
|
||||
return {
|
||||
code,
|
||||
details,
|
||||
message: cleanedMessage,
|
||||
missingFileName,
|
||||
missingFileType,
|
||||
status: possibleStatus,
|
||||
type: error.type || error.name || error.constructor?.name,
|
||||
userGuidance,
|
||||
};
|
||||
}
|
||||
|
||||
// Fallback handling
|
||||
const cleanedMessage = cleanComfyUIErrorMessage(String(error));
|
||||
|
||||
return {
|
||||
message: cleanedMessage,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse ComfyUI error message and return structured error information
|
||||
* Client-side version that focuses on error type categorization
|
||||
* File information and userGuidance are expected from server-side error handling
|
||||
* @param error - Original error object
|
||||
* @returns Parsed error object and error type
|
||||
*/
|
||||
export function parseComfyUIErrorMessage(error: any): ParsedError {
|
||||
// Check if it's already an AgentRuntimeError from WebAPI
|
||||
// AgentRuntimeError has structure: { error: object, errorType: string, provider: string }
|
||||
if (
|
||||
error &&
|
||||
typeof error === 'object' &&
|
||||
'errorType' in error &&
|
||||
'error' in error &&
|
||||
'provider' in error
|
||||
) {
|
||||
// Already parsed by server, return as-is
|
||||
return {
|
||||
error: error.error,
|
||||
errorType: error.errorType,
|
||||
};
|
||||
}
|
||||
|
||||
// Check if it's an error from checkAuth middleware
|
||||
// Format: { body: any, errorType: string }
|
||||
if (error && typeof error === 'object' && 'errorType' in error && 'body' in error) {
|
||||
// Extract error message from body
|
||||
let message = 'Authentication failed';
|
||||
if (error.body?.error?.message) {
|
||||
message = error.body.error.message;
|
||||
} else if (error.body?.error && typeof error.body.error === 'string') {
|
||||
message = error.body.error;
|
||||
} else if (error.body?.message) {
|
||||
message = error.body.message;
|
||||
}
|
||||
|
||||
return {
|
||||
error: {
|
||||
message,
|
||||
status: 401,
|
||||
},
|
||||
errorType: AgentRuntimeErrorType.InvalidProviderAPIKey,
|
||||
};
|
||||
}
|
||||
|
||||
const errorInfo = extractComfyUIErrorInfo(error);
|
||||
|
||||
// Default error type
|
||||
let errorType: ILobeAgentRuntimeErrorType = AgentRuntimeErrorType.ComfyUIBizError;
|
||||
|
||||
// Note: SyntaxError checking moved to server-side errorHandlerService
|
||||
// Client-side will never receive raw SyntaxError as it's already processed by server
|
||||
|
||||
// 1. HTTP status code errors (priority check)
|
||||
const status = errorInfo.status;
|
||||
const message = errorInfo.message;
|
||||
|
||||
switch (status) {
|
||||
case 400:
|
||||
case 401: {
|
||||
// These trigger ComfyUIAuth component
|
||||
errorType = AgentRuntimeErrorType.InvalidProviderAPIKey;
|
||||
break;
|
||||
}
|
||||
case 403: {
|
||||
// Permission denied
|
||||
errorType = AgentRuntimeErrorType.PermissionDenied;
|
||||
break;
|
||||
}
|
||||
case 404: {
|
||||
// 404 should trigger ComfyUIAuth for baseURL errors
|
||||
errorType = AgentRuntimeErrorType.InvalidProviderAPIKey;
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
if (status && status >= 500) {
|
||||
// Server errors
|
||||
errorType = AgentRuntimeErrorType.ComfyUIServiceUnavailable;
|
||||
}
|
||||
// 2. Check HTTP status code from error message (when status field doesn't exist)
|
||||
else if (!status && message) {
|
||||
if (message.includes('HTTP 401') || message.includes('401')) {
|
||||
errorType = AgentRuntimeErrorType.InvalidProviderAPIKey;
|
||||
} else if (message.includes('HTTP 403') || message.includes('403')) {
|
||||
errorType = AgentRuntimeErrorType.PermissionDenied;
|
||||
} else if (message.includes('HTTP 404') || message.includes('404')) {
|
||||
errorType = AgentRuntimeErrorType.InvalidProviderAPIKey;
|
||||
} else if (message.includes('HTTP 400') || message.includes('400')) {
|
||||
errorType = AgentRuntimeErrorType.InvalidProviderAPIKey;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Note: Error type determination is done server-side
|
||||
// Client receives pre-determined errorType from server
|
||||
|
||||
return {
|
||||
error: errorInfo,
|
||||
errorType,
|
||||
};
|
||||
}
|
||||
|
|
@ -145,6 +145,8 @@ describe('modelParse', () => {
|
|||
expect(detectModelProvider('deepseek-coder')).toBe('deepseek');
|
||||
expect(detectModelProvider('doubao-pro')).toBe('volcengine');
|
||||
expect(detectModelProvider('yi-large')).toBe('zeroone');
|
||||
expect(detectModelProvider('comfyui/flux-dev')).toBe('comfyui');
|
||||
expect(detectModelProvider('comfyui/sdxl-model')).toBe('comfyui');
|
||||
});
|
||||
|
||||
it('should default to OpenAI when no provider is detected', () => {
|
||||
|
|
@ -362,21 +364,28 @@ describe('modelParse', () => {
|
|||
{ id: 'claude-3-opus' }, // anthropic
|
||||
{ id: 'gemini-pro' }, // google
|
||||
{ id: 'qwen-turbo' }, // qwen
|
||||
{ id: 'comfyui/flux-dev', parameters: { width: 1024, height: 1024 } }, // comfyui
|
||||
];
|
||||
|
||||
const result = await processMultiProviderModelList(modelList);
|
||||
expect(result).toHaveLength(4);
|
||||
expect(result).toHaveLength(5);
|
||||
|
||||
const gpt4 = result.find((model) => model.id === 'gpt-4')!;
|
||||
const claude = result.find((model) => model.id === 'claude-3-opus')!;
|
||||
const gemini = result.find((model) => model.id === 'gemini-pro')!;
|
||||
const qwen = result.find((model) => model.id === 'qwen-turbo')!;
|
||||
const comfyui = result.find((model) => model.id === 'comfyui/flux-dev')!;
|
||||
|
||||
// Check abilities based on their respective provider configs and knownModels
|
||||
expect(gpt4.reasoning).toBe(false); // From knownModel (gpt-4)
|
||||
expect(claude.functionCall).toBe(true); // From knownModel (claude-3-opus)
|
||||
expect(gemini.functionCall).toBe(true); // From google keyword 'gemini'
|
||||
expect(qwen.functionCall).toBe(true); // From knownModel (qwen-turbo)
|
||||
|
||||
// ComfyUI models should have no chat capabilities (all false)
|
||||
expect(comfyui.functionCall).toBe(false); // ComfyUI config has empty arrays
|
||||
expect(comfyui.reasoning).toBe(false); // ComfyUI config has empty arrays
|
||||
expect(comfyui.vision).toBe(false); // ComfyUI config has empty arrays
|
||||
});
|
||||
|
||||
it('should recognize model capabilities based on keyword detection across providers', async () => {
|
||||
|
|
|
|||
|
|
@ -23,6 +23,12 @@ export const MODEL_LIST_CONFIGS = {
|
|||
reasoningKeywords: ['-3-7', '3.7', '-4'],
|
||||
visionKeywords: ['claude'],
|
||||
},
|
||||
comfyui: {
|
||||
// ComfyUI models are image generation models, no chat capabilities
|
||||
functionCallKeywords: [],
|
||||
reasoningKeywords: [],
|
||||
visionKeywords: [],
|
||||
},
|
||||
deepseek: {
|
||||
functionCallKeywords: ['v3', 'r1', 'deepseek-chat'],
|
||||
reasoningKeywords: ['r1', 'deepseek-reasoner', 'v3.1', 'v3.2'],
|
||||
|
|
@ -105,6 +111,7 @@ export const MODEL_LIST_CONFIGS = {
|
|||
// 模型所有者 (提供商) 关键词配置
|
||||
export const MODEL_OWNER_DETECTION_CONFIG = {
|
||||
anthropic: ['claude'],
|
||||
comfyui: ['comfyui/'], // ComfyUI models detection - all ComfyUI models have comfyui/ prefix
|
||||
deepseek: ['deepseek'],
|
||||
google: ['gemini', 'imagen'],
|
||||
inclusionai: ['ling-', 'ming-', 'ring-'],
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ export const AiProviderSDKEnum = {
|
|||
AzureAI: 'azureai',
|
||||
Bedrock: 'bedrock',
|
||||
Cloudflare: 'cloudflare',
|
||||
ComfyUI: 'comfyui',
|
||||
Google: 'google',
|
||||
Huggingface: 'huggingface',
|
||||
Ollama: 'ollama',
|
||||
|
|
@ -38,6 +39,7 @@ export type AiProviderSDKType = (typeof AiProviderSDKEnum)[keyof typeof AiProvid
|
|||
|
||||
const AiProviderSdkTypes = [
|
||||
'anthropic',
|
||||
'comfyui',
|
||||
'openai',
|
||||
'ollama',
|
||||
'azure',
|
||||
|
|
@ -240,7 +242,15 @@ export const UpdateAiProviderConfigSchema = z.object({
|
|||
})
|
||||
.optional(),
|
||||
fetchOnClient: z.boolean().nullable().optional(),
|
||||
keyVaults: z.record(z.string(), z.string().optional()).optional(),
|
||||
keyVaults: z
|
||||
.record(
|
||||
z.string(),
|
||||
z.union([
|
||||
z.string().optional(),
|
||||
z.record(z.string(), z.string()).optional(), // 支持嵌套对象，如 customHeaders
|
||||
]),
|
||||
)
|
||||
.optional(),
|
||||
});
|
||||
|
||||
export type UpdateAiProviderConfigParams = z.infer<typeof UpdateAiProviderConfigSchema>;
|
||||
|
|
|
|||
|
|
@ -14,6 +14,10 @@ export enum AsyncTaskStatus {
|
|||
export enum AsyncTaskErrorType {
|
||||
EmbeddingError = 'EmbeddingError',
|
||||
InvalidProviderAPIKey = 'InvalidProviderAPIKey',
|
||||
/**
|
||||
* Model not found on server
|
||||
*/
|
||||
ModelNotFound = 'ModelNotFound',
|
||||
/**
|
||||
* the chunk parse result it empty
|
||||
*/
|
||||
|
|
|
|||
|
|
@ -29,6 +29,14 @@ export interface ClientSecretPayload {
|
|||
|
||||
vertexAIRegion?: string;
|
||||
|
||||
/**
|
||||
* ComfyUI specific authentication fields
|
||||
*/
|
||||
authType?: string;
|
||||
username?: string;
|
||||
password?: string;
|
||||
customHeaders?: Record<string, string>;
|
||||
|
||||
/**
|
||||
* user id
|
||||
* in client db mode it's a uuid
|
||||
|
|
|
|||
|
|
@ -34,6 +34,15 @@ export interface CloudflareKeyVault {
|
|||
baseURLOrAccountID?: string;
|
||||
}
|
||||
|
||||
export interface ComfyUIKeyVault {
|
||||
apiKey?: string;
|
||||
authType?: 'none' | 'basic' | 'bearer' | 'custom';
|
||||
baseURL?: string;
|
||||
customHeaders?: Record<string, string>;
|
||||
password?: string;
|
||||
username?: string;
|
||||
}
|
||||
|
||||
export interface SearchEngineKeyVaults {
|
||||
searchxng?: {
|
||||
apiKey?: string;
|
||||
|
|
@ -57,6 +66,7 @@ export interface UserKeyVaults extends SearchEngineKeyVaults {
|
|||
cloudflare?: CloudflareKeyVault;
|
||||
cohere?: OpenAICompatibleKeyVault;
|
||||
cometapi?: OpenAICompatibleKeyVault;
|
||||
comfyui?: ComfyUIKeyVault;
|
||||
deepseek?: OpenAICompatibleKeyVault;
|
||||
fal?: FalKeyVault;
|
||||
fireworksai?: OpenAICompatibleKeyVault;
|
||||
|
|
|
|||
117
packages/utils/src/base64.test.ts
Normal file
117
packages/utils/src/base64.test.ts
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
import { describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { createBasicAuthCredentials, decodeFromBase64, encodeToBase64 } from './base64';
|
||||
|
||||
describe('base64 utilities', () => {
|
||||
describe('encodeToBase64', () => {
|
||||
it('should encode string to base64 in browser environment', () => {
|
||||
// Mock browser environment
|
||||
global.btoa = vi
|
||||
.fn()
|
||||
.mockImplementation((input) => Buffer.from(input, 'utf8').toString('base64'));
|
||||
|
||||
const result = encodeToBase64('test');
|
||||
|
||||
expect(global.btoa).toHaveBeenCalledWith('test');
|
||||
expect(result).toBe('dGVzdA==');
|
||||
});
|
||||
|
||||
it('should encode string to base64 in Node.js environment', () => {
|
||||
// Mock Node.js environment by removing btoa
|
||||
const originalBtoa = global.btoa;
|
||||
// @ts-ignore
|
||||
delete global.btoa;
|
||||
|
||||
const result = encodeToBase64('test');
|
||||
|
||||
expect(result).toBe('dGVzdA==');
|
||||
|
||||
// Restore btoa
|
||||
global.btoa = originalBtoa;
|
||||
});
|
||||
|
||||
it('should handle special characters', () => {
|
||||
const input = 'test@123:password';
|
||||
const result = encodeToBase64(input);
|
||||
|
||||
// Expected base64 for 'test@123:password' is 'dGVzdEAxMjM6cGFzc3dvcmQ='
|
||||
expect(result).toBe(Buffer.from(input, 'utf8').toString('base64'));
|
||||
});
|
||||
});
|
||||
|
||||
describe('decodeFromBase64', () => {
|
||||
it('should decode base64 string in browser environment', () => {
|
||||
// Mock browser environment
|
||||
global.atob = vi
|
||||
.fn()
|
||||
.mockImplementation((input) => Buffer.from(input, 'base64').toString('utf8'));
|
||||
|
||||
const result = decodeFromBase64('dGVzdA==');
|
||||
|
||||
expect(global.atob).toHaveBeenCalledWith('dGVzdA==');
|
||||
expect(result).toBe('test');
|
||||
});
|
||||
|
||||
it('should decode base64 string in Node.js environment', () => {
|
||||
// Mock Node.js environment by removing atob
|
||||
const originalAtob = global.atob;
|
||||
// @ts-ignore
|
||||
delete global.atob;
|
||||
|
||||
const result = decodeFromBase64('dGVzdA==');
|
||||
|
||||
expect(result).toBe('test');
|
||||
|
||||
// Restore atob
|
||||
global.atob = originalAtob;
|
||||
});
|
||||
});
|
||||
|
||||
describe('createBasicAuthCredentials', () => {
|
||||
it('should create basic auth credentials', () => {
|
||||
const username = 'testuser';
|
||||
const password = 'testpass';
|
||||
|
||||
const result = createBasicAuthCredentials(username, password);
|
||||
|
||||
// Expected base64 for 'testuser:testpass' is 'dGVzdHVzZXI6dGVzdHBhc3M='
|
||||
expect(result).toBe('dGVzdHVzZXI6dGVzdHBhc3M=');
|
||||
});
|
||||
|
||||
it('should handle special characters in credentials', () => {
|
||||
const username = 'user@domain.com';
|
||||
const password = 'p@ss:w0rd!';
|
||||
|
||||
const result = createBasicAuthCredentials(username, password);
|
||||
const decoded = decodeFromBase64(result);
|
||||
|
||||
expect(decoded).toBe('user@domain.com:p@ss:w0rd!');
|
||||
});
|
||||
|
||||
it('should handle empty credentials', () => {
|
||||
const result = createBasicAuthCredentials('', '');
|
||||
const decoded = decodeFromBase64(result);
|
||||
|
||||
expect(decoded).toBe(':');
|
||||
});
|
||||
});
|
||||
|
||||
describe('round-trip encoding/decoding', () => {
|
||||
it('should preserve data through encode/decode cycle', () => {
|
||||
const testStrings = [
|
||||
'simple text',
|
||||
'test@123:password',
|
||||
'中文测试',
|
||||
'user:pass',
|
||||
'special!@#$%^&*()chars',
|
||||
'',
|
||||
];
|
||||
|
||||
testStrings.forEach((input) => {
|
||||
const encoded = encodeToBase64(input);
|
||||
const decoded = decodeFromBase64(encoded);
|
||||
expect(decoded).toBe(input);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
44
packages/utils/src/base64.ts
Normal file
44
packages/utils/src/base64.ts
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
/**
|
||||
* Cross-platform base64 encoding utility
|
||||
* Works in both browser and Node.js environments
|
||||
*/
|
||||
|
||||
/**
|
||||
* Encode a string to base64
|
||||
* @param input - The string to encode
|
||||
* @returns Base64 encoded string
|
||||
*/
|
||||
export const encodeToBase64 = (input: string): string => {
|
||||
if (typeof btoa === 'function') {
|
||||
// Browser environment
|
||||
return btoa(input);
|
||||
} else {
|
||||
// Node.js environment
|
||||
return Buffer.from(input, 'utf8').toString('base64');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Decode a base64 string
|
||||
* @param input - The base64 string to decode
|
||||
* @returns Decoded string
|
||||
*/
|
||||
export const decodeFromBase64 = (input: string): string => {
|
||||
if (typeof atob === 'function') {
|
||||
// Browser environment
|
||||
return atob(input);
|
||||
} else {
|
||||
// Node.js environment
|
||||
return Buffer.from(input, 'base64').toString('utf8');
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Create Basic Authentication header value
|
||||
* @param username - Username for authentication
|
||||
* @param password - Password for authentication
|
||||
* @returns Base64 encoded credentials for Basic auth
|
||||
*/
|
||||
export const createBasicAuthCredentials = (username: string, password: string): string => {
|
||||
return encodeToBase64(`${username}:${password}`);
|
||||
};
|
||||
|
|
@ -1,8 +1,10 @@
|
|||
export * from './base64';
|
||||
export * from './client/cookie';
|
||||
export * from './detectChinese';
|
||||
export * from './format';
|
||||
export * from './imageToBase64';
|
||||
export * from './keyboard';
|
||||
export * from './number';
|
||||
export * from './object';
|
||||
export * from './parseModels';
|
||||
export * from './pricing';
|
||||
|
|
|
|||
98
src/app/(backend)/webapi/create-image/comfyui/route.ts
Normal file
98
src/app/(backend)/webapi/create-image/comfyui/route.ts
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
import { NextResponse } from 'next/server';
|
||||
|
||||
import { checkAuth } from '@/app/(backend)/middleware/auth';
|
||||
import { getServerDBConfig } from '@/config/db';
|
||||
import { createCallerFactory } from '@/libs/trpc/lambda';
|
||||
import { lambdaRouter } from '@/server/routers/lambda';
|
||||
|
||||
export const runtime = 'nodejs';
|
||||
export const maxDuration = 300;
|
||||
|
||||
const serverDBEnv = getServerDBConfig();
|
||||
|
||||
// Custom handler that supports both regular auth and internal service auth
|
||||
const handler = async (req: Request, { jwtPayload }: { jwtPayload?: any }) => {
|
||||
try {
|
||||
const body = await req.json();
|
||||
const { model, params, options } = body;
|
||||
|
||||
// Create tRPC caller with authentication context
|
||||
const createCaller = createCallerFactory(lambdaRouter);
|
||||
|
||||
const caller = createCaller({
|
||||
jwtPayload,
|
||||
nextAuth: undefined, // WebAPI routes don't have nextAuth session
|
||||
userId: jwtPayload?.userId, // Required for userAuth middleware
|
||||
});
|
||||
|
||||
// Call ComfyUI service through tRPC
|
||||
const result = await caller.comfyui.createImage({
|
||||
model,
|
||||
options,
|
||||
params,
|
||||
});
|
||||
|
||||
return NextResponse.json(result);
|
||||
} catch (error: any) {
|
||||
console.error('[ComfyUI WebAPI] Error:', error);
|
||||
|
||||
// Extract AgentRuntimeError from TRPCError's cause
|
||||
const agentError = error?.cause;
|
||||
|
||||
// If we have an AgentRuntimeError in the cause, return it
|
||||
if (agentError && typeof agentError === 'object' && 'errorType' in agentError) {
|
||||
// Convert errorType to HTTP status
|
||||
let status;
|
||||
switch (agentError.errorType) {
|
||||
case 'InvalidProviderAPIKey':
|
||||
case 401: {
|
||||
status = 401;
|
||||
break;
|
||||
}
|
||||
case 'PermissionDenied':
|
||||
case 403: {
|
||||
status = 403;
|
||||
break;
|
||||
}
|
||||
case 'ModelNotFound':
|
||||
case 404: {
|
||||
status = 404;
|
||||
break;
|
||||
}
|
||||
case 'ComfyUIServiceUnavailable':
|
||||
case 503: {
|
||||
status = 503;
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
status = 500;
|
||||
}
|
||||
}
|
||||
|
||||
// Return the AgentRuntimeError directly for the Provider to handle
|
||||
return NextResponse.json(agentError, { status });
|
||||
}
|
||||
|
||||
// Fallback for other errors
|
||||
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
|
||||
return NextResponse.json({ error: errorMessage }, { status: 500 });
|
||||
}
|
||||
};
|
||||
|
||||
export const POST = async (req: Request) => {
|
||||
// Check for internal service authentication (only if KEY_VAULTS_SECRET is set)
|
||||
if (serverDBEnv.KEY_VAULTS_SECRET) {
|
||||
const authorization = req.headers.get('Authorization');
|
||||
|
||||
// If request has internal service token, bypass regular auth
|
||||
if (authorization === `Bearer ${serverDBEnv.KEY_VAULTS_SECRET}`) {
|
||||
// Internal service call from ComfyUI provider
|
||||
// Pass a system user ID for internal service calls
|
||||
return handler(req, { jwtPayload: { userId: 'INTERNAL_SERVICE' } });
|
||||
}
|
||||
}
|
||||
|
||||
// Otherwise use regular checkAuth
|
||||
// ComfyUI doesn't have a provider param, but checkAuth requires it
|
||||
return checkAuth(handler)(req, { params: Promise.resolve({ provider: 'comfyui' }) });
|
||||
};
|
||||
|
|
@ -115,6 +115,7 @@ export const GenerationBatchItem = memo<GenerationBatchItemProps>(({ batch }) =>
|
|||
);
|
||||
|
||||
if (isInvalidApiKey) {
|
||||
// Use unified InvalidAPIKey component for all providers (including ComfyUI)
|
||||
return (
|
||||
<InvalidAPIKey
|
||||
bedrockDescription={t('bedrock.unlock.imageGenerationDescription', { ns: 'modelProvider' })}
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@
|
|||
|
||||
import { Block, Icon, Text } from '@lobehub/ui';
|
||||
import { ImageOffIcon } from 'lucide-react';
|
||||
import { memo } from 'react';
|
||||
import { memo, useMemo } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { Center } from 'react-layout-kit';
|
||||
|
||||
import { AgentRuntimeErrorType } from '@lobechat/model-runtime';
|
||||
|
||||
import { ActionButtons } from './ActionButtons';
|
||||
import { useStyles } from './styles';
|
||||
import { ErrorStateProps } from './types';
|
||||
|
|
@ -16,12 +18,46 @@ export const ErrorState = memo<ErrorStateProps>(
|
|||
({ generation, generationBatch, aspectRatio, onDelete, onCopyError }) => {
|
||||
const { styles, theme } = useStyles();
|
||||
const { t } = useTranslation('image');
|
||||
const { t: tError } = useTranslation('error');
|
||||
|
||||
const errorMessage = generation.task.error
|
||||
? typeof generation.task.error.body === 'string'
|
||||
? generation.task.error.body
|
||||
: generation.task.error.body?.detail || generation.task.error.name || 'Unknown error'
|
||||
: '';
|
||||
const errorMessage = useMemo(() => {
|
||||
if (!generation.task.error) return '';
|
||||
|
||||
const error = generation.task.error;
|
||||
const errorBody = typeof error.body === 'string' ? error.body : error.body?.detail;
|
||||
|
||||
// Try to translate based on error type if it matches known AgentRuntimeErrorType
|
||||
if (errorBody) {
|
||||
// Check if the error body is an AgentRuntimeErrorType that needs translation
|
||||
const knownErrorTypes = Object.values(AgentRuntimeErrorType);
|
||||
if (
|
||||
knownErrorTypes.includes(
|
||||
errorBody as (typeof AgentRuntimeErrorType)[keyof typeof AgentRuntimeErrorType],
|
||||
)
|
||||
) {
|
||||
// Use localized error message - ComfyUI errors are under 'response' namespace
|
||||
const translationKey = `response.${errorBody}`;
|
||||
const translated = tError(translationKey as any);
|
||||
|
||||
// If translation key is not found, it returns the key itself
|
||||
// Check if we got back the key (meaning translation failed)
|
||||
if (translated === translationKey || (translated as string).startsWith('response.')) {
|
||||
// Try without any prefix (for backwards compatibility)
|
||||
const directTranslated = tError(errorBody as any);
|
||||
if (directTranslated !== errorBody) {
|
||||
return directTranslated as string;
|
||||
}
|
||||
// Final fallback to the original error message
|
||||
return errorBody;
|
||||
}
|
||||
|
||||
return translated as string;
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to original error message
|
||||
return errorBody || error.name || 'Unknown error';
|
||||
}, [generation.task.error, generationBatch.provider, tError]);
|
||||
|
||||
return (
|
||||
<Block
|
||||
|
|
|
|||
|
|
@ -0,0 +1,138 @@
|
|||
'use client';
|
||||
|
||||
import { Select } from '@lobehub/ui';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
import { FormInput, FormPassword } from '@/components/FormInput';
|
||||
import KeyValueEditor from '@/components/KeyValueEditor';
|
||||
import { ComfyUIProviderCard } from '@/config/modelProviders';
|
||||
import { aiProviderSelectors, useAiInfraStore } from '@/store/aiInfra';
|
||||
import { GlobalLLMProviderKey } from '@/types/user/settings';
|
||||
|
||||
import { KeyVaultsConfigKey } from '../../const';
|
||||
import { SkeletonInput } from '../../features/ProviderConfig';
|
||||
import { ProviderItem } from '../../type';
|
||||
import ProviderDetail from '../default';
|
||||
|
||||
const providerKey: GlobalLLMProviderKey = 'comfyui';
|
||||
|
||||
const useComfyUICard = (): ProviderItem => {
|
||||
const { t } = useTranslation('modelProvider');
|
||||
|
||||
const isLoading = useAiInfraStore(aiProviderSelectors.isAiProviderConfigLoading(providerKey));
|
||||
|
||||
// Get current config and watch for auth type changes
|
||||
const config = useAiInfraStore((s) => s.aiProviderRuntimeConfig?.[providerKey]);
|
||||
const authType = config?.keyVaults?.authType || 'none';
|
||||
|
||||
const authTypeOptions = [
|
||||
{ label: t('comfyui.authType.options.none'), value: 'none' },
|
||||
{ label: t('comfyui.authType.options.basic'), value: 'basic' },
|
||||
{ label: t('comfyui.authType.options.bearer'), value: 'bearer' },
|
||||
{ label: t('comfyui.authType.options.custom'), value: 'custom' },
|
||||
];
|
||||
|
||||
const apiKeyItems = [
|
||||
// Base URL - Always shown
|
||||
{
|
||||
children: isLoading ? (
|
||||
<SkeletonInput />
|
||||
) : (
|
||||
<FormInput placeholder={t('comfyui.baseURL.placeholder')} />
|
||||
),
|
||||
desc: t('comfyui.baseURL.desc'),
|
||||
label: t('comfyui.baseURL.title'),
|
||||
name: [KeyVaultsConfigKey, 'baseURL'],
|
||||
},
|
||||
|
||||
// Authentication Type Selector - Always shown
|
||||
{
|
||||
children: isLoading ? (
|
||||
<SkeletonInput />
|
||||
) : (
|
||||
<Select
|
||||
allowClear={false}
|
||||
options={authTypeOptions}
|
||||
placeholder={t('comfyui.authType.placeholder')}
|
||||
/>
|
||||
),
|
||||
desc: t('comfyui.authType.desc'),
|
||||
label: t('comfyui.authType.title'),
|
||||
name: [KeyVaultsConfigKey, 'authType'],
|
||||
},
|
||||
];
|
||||
|
||||
// Conditionally add fields based on auth type
|
||||
if (authType === 'basic') {
|
||||
apiKeyItems.push(
|
||||
{
|
||||
children: isLoading ? (
|
||||
<SkeletonInput />
|
||||
) : (
|
||||
<FormInput autoComplete="username" placeholder={t('comfyui.username.placeholder')} />
|
||||
),
|
||||
desc: t('comfyui.username.desc'),
|
||||
label: t('comfyui.username.title'),
|
||||
name: [KeyVaultsConfigKey, 'username'],
|
||||
},
|
||||
{
|
||||
children: isLoading ? (
|
||||
<SkeletonInput />
|
||||
) : (
|
||||
<FormPassword
|
||||
autoComplete="new-password"
|
||||
placeholder={t('comfyui.password.placeholder')}
|
||||
/>
|
||||
),
|
||||
desc: t('comfyui.password.desc'),
|
||||
label: t('comfyui.password.title'),
|
||||
name: [KeyVaultsConfigKey, 'password'],
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
if (authType === 'bearer') {
|
||||
apiKeyItems.push({
|
||||
children: isLoading ? (
|
||||
<SkeletonInput />
|
||||
) : (
|
||||
<FormPassword autoComplete="new-password" placeholder={t('comfyui.apiKey.placeholder')} />
|
||||
),
|
||||
desc: t('comfyui.apiKey.desc'),
|
||||
label: t('comfyui.apiKey.title'),
|
||||
name: [KeyVaultsConfigKey, 'apiKey'],
|
||||
});
|
||||
}
|
||||
|
||||
if (authType === 'custom') {
|
||||
apiKeyItems.push({
|
||||
children: isLoading ? (
|
||||
<SkeletonInput />
|
||||
) : (
|
||||
<KeyValueEditor
|
||||
addButtonText={t('comfyui.customHeaders.addButton')}
|
||||
deleteTooltip={t('comfyui.customHeaders.deleteTooltip')}
|
||||
duplicateKeyErrorText={t('comfyui.customHeaders.duplicateKeyError')}
|
||||
keyPlaceholder={t('comfyui.customHeaders.keyPlaceholder')}
|
||||
valuePlaceholder={t('comfyui.customHeaders.valuePlaceholder')}
|
||||
/>
|
||||
),
|
||||
desc: t('comfyui.customHeaders.desc'),
|
||||
label: t('comfyui.customHeaders.title'),
|
||||
name: [KeyVaultsConfigKey, 'customHeaders'],
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
...ComfyUIProviderCard,
|
||||
apiKeyItems,
|
||||
};
|
||||
};
|
||||
|
||||
const Page = () => {
|
||||
const card = useComfyUICard();
|
||||
|
||||
return <ProviderDetail {...card} />;
|
||||
};
|
||||
|
||||
export default Page;
|
||||
|
|
@ -3,6 +3,7 @@ import Azure from './azure';
|
|||
import AzureAI from './azureai';
|
||||
import Bedrock from './bedrock';
|
||||
import Cloudflare from './cloudflare';
|
||||
import ComfyUI from './comfyui';
|
||||
import DefaultPage from './default/ProviderDetialPage';
|
||||
import GitHub from './github';
|
||||
import Ollama from './ollama';
|
||||
|
|
@ -28,6 +29,9 @@ const ProviderDetailPage = (props: { id?: string | null }) => {
|
|||
case 'cloudflare': {
|
||||
return <Cloudflare />;
|
||||
}
|
||||
case 'comfyui': {
|
||||
return <ComfyUI />;
|
||||
}
|
||||
case 'github': {
|
||||
return <GitHub />;
|
||||
}
|
||||
|
|
|
|||
251
src/components/InvalidAPIKey/APIKeyForm/ComfyUIForm.tsx
Normal file
251
src/components/InvalidAPIKey/APIKeyForm/ComfyUIForm.tsx
Normal file
|
|
@ -0,0 +1,251 @@
|
|||
'use client';
|
||||
|
||||
import { ComfyUI } from '@lobehub/icons';
|
||||
import { Button, Icon, Select } from '@lobehub/ui';
|
||||
import { createStyles, useTheme } from 'antd-style';
|
||||
import { Loader2Icon, Network } from 'lucide-react';
|
||||
import { memo, useContext, useEffect, useState } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { Center, Flexbox } from 'react-layout-kit';
|
||||
|
||||
import { FormInput, FormPassword } from '@/components/FormInput';
|
||||
import KeyValueEditor from '@/components/KeyValueEditor';
|
||||
import { FormAction } from '@/features/Conversation/Error/style';
|
||||
import { useAiInfraStore } from '@/store/aiInfra';
|
||||
import { ComfyUIKeyVault } from '@/types/user/settings';
|
||||
|
||||
import { LoadingContext } from './LoadingContext';
|
||||
|
||||
interface ComfyUIFormProps {
|
||||
description: string;
|
||||
}
|
||||
|
||||
const useStyles = createStyles(({ css, token }) => ({
|
||||
comfyuiFormWide: css`
|
||||
max-width: 900px !important;
|
||||
|
||||
/* Hide the avatar - target the first child which is the Avatar component */
|
||||
> *:first-child {
|
||||
display: none !important;
|
||||
}
|
||||
`,
|
||||
container: css`
|
||||
width: 100%;
|
||||
max-width: 900px;
|
||||
border: 1px solid ${token.colorSplit};
|
||||
border-radius: 8px;
|
||||
|
||||
color: ${token.colorText};
|
||||
|
||||
background: ${token.colorBgContainer};
|
||||
`,
|
||||
}));
|
||||
|
||||
const ComfyUIForm = memo<ComfyUIFormProps>(({ description }) => {
|
||||
const { t } = useTranslation('error');
|
||||
const { t: s } = useTranslation('modelProvider');
|
||||
const theme = useTheme();
|
||||
const { styles } = useStyles();
|
||||
|
||||
// Use aiInfraStore for updating config (same as settings page)
|
||||
const updateAiProviderConfig = useAiInfraStore((s) => s.updateAiProviderConfig);
|
||||
const useFetchAiProviderRuntimeState = useAiInfraStore((s) => s.useFetchAiProviderRuntimeState);
|
||||
|
||||
const { loading, setLoading } = useContext(LoadingContext);
|
||||
|
||||
// Fetch the runtime state to ensure config is loaded
|
||||
// Pass true since this is for auth dialog (not initialization)
|
||||
const fetchRuntimeState = useFetchAiProviderRuntimeState(true);
|
||||
|
||||
// Get ComfyUI config from aiInfraStore (same as settings page)
|
||||
const comfyUIConfig = useAiInfraStore(
|
||||
(s) => s.aiProviderRuntimeConfig?.['comfyui']?.keyVaults,
|
||||
) as ComfyUIKeyVault | undefined;
|
||||
|
||||
// State for showing base URL input - initially hidden
|
||||
const [showBaseURL, setShowBaseURL] = useState(false);
|
||||
|
||||
// State management for form values - initialize without config first
|
||||
const [formValues, setFormValues] = useState({
|
||||
apiKey: '',
|
||||
authType: 'none' as string,
|
||||
baseURL: 'http://127.0.0.1:8000',
|
||||
customHeaders: {} as Record<string, string>,
|
||||
password: '',
|
||||
username: '',
|
||||
});
|
||||
|
||||
// Update form values when comfyUIConfig changes (配置反读)
|
||||
// Use individual primitive values to avoid infinite re-renders
|
||||
useEffect(() => {
|
||||
if (comfyUIConfig) {
|
||||
const newValues = {
|
||||
apiKey: comfyUIConfig.apiKey || '',
|
||||
authType: comfyUIConfig.authType || 'none',
|
||||
baseURL: comfyUIConfig.baseURL || 'http://127.0.0.1:8000',
|
||||
customHeaders: comfyUIConfig.customHeaders || {},
|
||||
password: comfyUIConfig.password || '',
|
||||
username: comfyUIConfig.username || '',
|
||||
};
|
||||
setFormValues(newValues);
|
||||
}
|
||||
}, [
|
||||
comfyUIConfig?.apiKey,
|
||||
comfyUIConfig?.authType,
|
||||
comfyUIConfig?.baseURL,
|
||||
comfyUIConfig?.password,
|
||||
comfyUIConfig?.username,
|
||||
JSON.stringify(comfyUIConfig?.customHeaders),
|
||||
]);
|
||||
|
||||
const authTypeOptions = [
|
||||
{ label: s('comfyui.authType.options.none'), value: 'none' },
|
||||
{ label: s('comfyui.authType.options.basic'), value: 'basic' },
|
||||
{ label: s('comfyui.authType.options.bearer'), value: 'bearer' },
|
||||
{ label: s('comfyui.authType.options.custom'), value: 'custom' },
|
||||
];
|
||||
|
||||
const handleValueChange = async (field: string, value: any) => {
|
||||
const newValues = {
|
||||
...formValues,
|
||||
[field]: value,
|
||||
};
|
||||
setFormValues(newValues);
|
||||
|
||||
// Skip validation for certain fields that can be empty
|
||||
const skipValidation = ['customHeaders', 'apiKey', 'username', 'password'];
|
||||
|
||||
// Basic validation before saving
|
||||
if (!skipValidation.includes(field) && field === 'baseURL' && !value) {
|
||||
return; // Don't save if baseURL is empty
|
||||
}
|
||||
|
||||
// Real-time save like other providers
|
||||
setLoading(true);
|
||||
try {
|
||||
await updateAiProviderConfig('comfyui', {
|
||||
keyVaults: newValues,
|
||||
});
|
||||
// Refetch the runtime state to ensure config is synced
|
||||
await fetchRuntimeState.mutate();
|
||||
} catch (error) {
|
||||
console.error('Failed to update ComfyUI config:', error);
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<Center className={styles.container} gap={24} padding={24}>
|
||||
<Center gap={16} paddingBlock={32} style={{ width: '100%' }}>
|
||||
<ComfyUI.Combine size={64} type={'color'} />
|
||||
<FormAction
|
||||
avatar={<div />}
|
||||
className={styles.comfyuiFormWide}
|
||||
description={description}
|
||||
title={t('unlock.comfyui.title', { name: 'ComfyUI' })}
|
||||
>
|
||||
<Flexbox gap={16} width="100%">
|
||||
{/* Base URL */}
|
||||
{showBaseURL ? (
|
||||
<Flexbox gap={4}>
|
||||
<div style={{ fontSize: 14, fontWeight: 500 }}>{s('comfyui.baseURL.title')}</div>
|
||||
<FormInput
|
||||
onChange={(value) => handleValueChange('baseURL', value)}
|
||||
placeholder={s('comfyui.baseURL.placeholder')}
|
||||
suffix={<div>{loading && <Icon icon={Loader2Icon} spin />}</div>}
|
||||
value={formValues.baseURL}
|
||||
/>
|
||||
</Flexbox>
|
||||
) : (
|
||||
<Button
|
||||
icon={<Icon icon={Network} />}
|
||||
onClick={() => setShowBaseURL(true)}
|
||||
type={'text'}
|
||||
>
|
||||
{t('unlock.comfyui.modifyBaseUrl')}
|
||||
</Button>
|
||||
)}
|
||||
|
||||
{/* Auth Type */}
|
||||
<Flexbox gap={4}>
|
||||
<div style={{ fontSize: 14, fontWeight: 500 }}>{s('comfyui.authType.title')}</div>
|
||||
<Select
|
||||
allowClear={false}
|
||||
onChange={(value) => handleValueChange('authType', value)}
|
||||
options={authTypeOptions}
|
||||
placeholder={s('comfyui.authType.placeholder')}
|
||||
value={formValues.authType}
|
||||
/>
|
||||
</Flexbox>
|
||||
|
||||
{/* Basic Auth Fields */}
|
||||
{formValues.authType === 'basic' && (
|
||||
<>
|
||||
<Flexbox gap={4}>
|
||||
<div style={{ fontSize: 14, fontWeight: 500 }}>{s('comfyui.username.title')}</div>
|
||||
<FormInput
|
||||
autoComplete="username"
|
||||
onChange={(value) => handleValueChange('username', value)}
|
||||
placeholder={s('comfyui.username.placeholder')}
|
||||
suffix={<div>{loading && <Icon icon={Loader2Icon} spin />}</div>}
|
||||
value={formValues.username}
|
||||
/>
|
||||
</Flexbox>
|
||||
<Flexbox gap={4}>
|
||||
<div style={{ fontSize: 14, fontWeight: 500 }}>{s('comfyui.password.title')}</div>
|
||||
<FormPassword
|
||||
autoComplete="new-password"
|
||||
onChange={(value) => handleValueChange('password', value)}
|
||||
placeholder={s('comfyui.password.placeholder')}
|
||||
suffix={<div>{loading && <Icon icon={Loader2Icon} spin />}</div>}
|
||||
value={formValues.password}
|
||||
/>
|
||||
</Flexbox>
|
||||
</>
|
||||
)}
|
||||
|
||||
{/* Bearer Token Field */}
|
||||
{formValues.authType === 'bearer' && (
|
||||
<Flexbox gap={4}>
|
||||
<div style={{ fontSize: 14, fontWeight: 500 }}>{s('comfyui.apiKey.title')}</div>
|
||||
<FormPassword
|
||||
autoComplete="new-password"
|
||||
onChange={(value) => handleValueChange('apiKey', value)}
|
||||
placeholder={s('comfyui.apiKey.placeholder')}
|
||||
suffix={<div>{loading && <Icon icon={Loader2Icon} spin />}</div>}
|
||||
value={formValues.apiKey}
|
||||
/>
|
||||
</Flexbox>
|
||||
)}
|
||||
|
||||
{/* Custom Headers Field */}
|
||||
{formValues.authType === 'custom' && (
|
||||
<Flexbox gap={4}>
|
||||
<div style={{ fontSize: 14, fontWeight: 500 }}>
|
||||
{s('comfyui.customHeaders.title')}
|
||||
</div>
|
||||
<div style={{ color: theme.colorTextSecondary, fontSize: 12, marginBottom: 4 }}>
|
||||
{s('comfyui.customHeaders.desc')}
|
||||
</div>
|
||||
<KeyValueEditor
|
||||
addButtonText={s('comfyui.customHeaders.addButton')}
|
||||
deleteTooltip={s('comfyui.customHeaders.deleteTooltip')}
|
||||
duplicateKeyErrorText={s('comfyui.customHeaders.duplicateKeyError')}
|
||||
keyPlaceholder={s('comfyui.customHeaders.keyPlaceholder')}
|
||||
onChange={(value) => handleValueChange('customHeaders', value)}
|
||||
value={formValues.customHeaders}
|
||||
valuePlaceholder={s('comfyui.customHeaders.valuePlaceholder')}
|
||||
/>
|
||||
</Flexbox>
|
||||
)}
|
||||
</Flexbox>
|
||||
</FormAction>
|
||||
</Center>
|
||||
</Center>
|
||||
);
|
||||
});
|
||||
|
||||
ComfyUIForm.displayName = 'ComfyUIForm';
|
||||
|
||||
export default ComfyUIForm;
|
||||
|
|
@ -0,0 +1,137 @@
|
|||
import { ModelProvider } from 'model-bank';
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import APIKeyForm from '../index';
|
||||
|
||||
// Mock the dependencies
|
||||
vi.mock('@/store/aiInfra', () => ({
|
||||
useAiInfraStore: vi.fn(() => ({
|
||||
updateAiProviderConfig: vi.fn(),
|
||||
useFetchAiProviderRuntimeState: vi.fn(() => ({})),
|
||||
aiProviderRuntimeConfig: {},
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock('react-i18next', () => ({
|
||||
useTranslation: () => ({
|
||||
t: (key: string) => key,
|
||||
}),
|
||||
}));
|
||||
|
||||
vi.mock('antd-style', () => ({
|
||||
useTheme: () => ({
|
||||
colorTextSecondary: '#999',
|
||||
}),
|
||||
createStyles: vi.fn(() => () => ({ styles: {} })),
|
||||
}));
|
||||
|
||||
vi.mock('@/components/FormInput', () => ({
|
||||
FormInput: vi.fn(({ value, onChange, ...props }) => (
|
||||
<input
|
||||
data-testid="form-input"
|
||||
value={value}
|
||||
onChange={(e) => onChange?.(e.target.value)}
|
||||
{...props}
|
||||
/>
|
||||
)),
|
||||
FormPassword: vi.fn(({ value, onChange, ...props }) => (
|
||||
<input
|
||||
data-testid="form-password"
|
||||
type="password"
|
||||
value={value}
|
||||
onChange={(e) => onChange?.(e.target.value)}
|
||||
{...props}
|
||||
/>
|
||||
)),
|
||||
}));
|
||||
|
||||
vi.mock('@/components/KeyValueEditor', () => ({
|
||||
default: vi.fn(() => <div data-testid="key-value-editor">Key-Value Editor</div>),
|
||||
}));
|
||||
|
||||
vi.mock('@lobehub/ui', () => ({
|
||||
Icon: vi.fn(({ icon, ...props }) => (
|
||||
<div data-testid="icon" {...props}>
|
||||
{icon?.name}
|
||||
</div>
|
||||
)),
|
||||
Select: vi.fn(({ value, options, onChange, ...props }) => (
|
||||
<select
|
||||
data-testid="select"
|
||||
value={value}
|
||||
onChange={(e) => onChange?.(e.target.value)}
|
||||
{...props}
|
||||
>
|
||||
{options?.map((option: any) => (
|
||||
<option key={option.value} value={option.value}>
|
||||
{option.label}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
)),
|
||||
Button: vi.fn(({ children, onClick, ...props }) => (
|
||||
<button data-testid="button" onClick={onClick} {...props}>
|
||||
{children}
|
||||
</button>
|
||||
)),
|
||||
ProviderIcon: vi.fn(() => <div data-testid="provider-icon">Provider Icon</div>),
|
||||
}));
|
||||
|
||||
vi.mock('@lobehub/icons', () => ({
|
||||
ComfyUI: {
|
||||
Combine: vi.fn(() => <div data-testid="comfyui-icon">ComfyUI Icon</div>),
|
||||
},
|
||||
ProviderIcon: vi.fn(() => <div data-testid="provider-icon">Provider Icon</div>),
|
||||
}));
|
||||
|
||||
vi.mock('react-layout-kit', () => ({
|
||||
Center: vi.fn(({ children, ...props }) => (
|
||||
<div data-testid="center" {...props}>
|
||||
{children}
|
||||
</div>
|
||||
)),
|
||||
Flexbox: vi.fn(({ children, ...props }) => (
|
||||
<div data-testid="flexbox" {...props}>
|
||||
{children}
|
||||
</div>
|
||||
)),
|
||||
}));
|
||||
|
||||
vi.mock('@/features/Conversation/Error/style', () => ({
|
||||
FormAction: vi.fn(({ children, title, description, avatar, ...props }) => (
|
||||
<div data-testid="form-action" {...props}>
|
||||
<div data-testid="avatar">{avatar}</div>
|
||||
<div data-testid="title">{title}</div>
|
||||
<div data-testid="description">{description}</div>
|
||||
{children}
|
||||
</div>
|
||||
)),
|
||||
ErrorActionContainer: vi.fn(({ children, ...props }) => (
|
||||
<div data-testid="error-action-container" {...props}>
|
||||
{children}
|
||||
</div>
|
||||
)),
|
||||
}));
|
||||
|
||||
describe('ComfyUIForm Integration', () => {
|
||||
const mockProps = {
|
||||
bedrockDescription: 'bedrock.description',
|
||||
description: 'comfyui.description',
|
||||
id: 'test-batch-id',
|
||||
onClose: vi.fn(),
|
||||
onRecreate: vi.fn(),
|
||||
provider: ModelProvider.ComfyUI,
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('should use ComfyUI provider correctly', () => {
|
||||
expect(ModelProvider.ComfyUI).toBe('comfyui');
|
||||
});
|
||||
|
||||
it('should import APIKeyForm without errors', () => {
|
||||
expect(APIKeyForm).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
|
@ -8,6 +8,7 @@ import { Center, Flexbox } from 'react-layout-kit';
|
|||
import { GlobalLLMProviderKey } from '@/types/user/settings';
|
||||
|
||||
import BedrockForm from './Bedrock';
|
||||
import ComfyUIForm from './ComfyUIForm';
|
||||
import { LoadingContext } from './LoadingContext';
|
||||
import ProviderApiKeyForm from './ProviderApiKeyForm';
|
||||
|
||||
|
|
@ -67,9 +68,17 @@ const APIKeyForm = memo<APIKeyFormProps>(
|
|||
|
||||
return (
|
||||
<LoadingContext value={{ loading, setLoading }}>
|
||||
<Center gap={16} style={{ maxWidth: 300 }}>
|
||||
<Center
|
||||
gap={16}
|
||||
style={{
|
||||
maxWidth: provider === ModelProvider.ComfyUI ? 900 : 300,
|
||||
width: provider === ModelProvider.ComfyUI ? '80%' : 'auto',
|
||||
}}
|
||||
>
|
||||
{provider === ModelProvider.Bedrock ? (
|
||||
<BedrockForm description={bedrockDescription} />
|
||||
) : provider === ModelProvider.ComfyUI ? (
|
||||
<ComfyUIForm description={description} />
|
||||
) : (
|
||||
<ProviderApiKeyForm
|
||||
apiKeyPlaceholder={apiKeyPlaceholder}
|
||||
|
|
|
|||
40
src/config/modelProviders/comfyui.ts
Normal file
40
src/config/modelProviders/comfyui.ts
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
import { ModelProviderCard } from '@/types/llm';
|
||||
|
||||
/**
|
||||
* ComfyUI Provider Configuration
|
||||
*
|
||||
* 支持本地和远程 ComfyUI 服务器连接
|
||||
* 提供 FLUX 系列模型的图像生成能力
|
||||
*
|
||||
* @see https://www.comfy.org/
|
||||
*/
|
||||
const ComfyUI: ModelProviderCard = {
|
||||
chatModels: [],
|
||||
description:
|
||||
'强大的开源图像、视频、音频生成工作流引擎，支持 SD FLUX Qwen Hunyuan WAN 等先进模型，提供节点化工作流编辑和私有化部署能力',
|
||||
enabled: true,
|
||||
id: 'comfyui',
|
||||
name: 'ComfyUI',
|
||||
settings: {
|
||||
// 禁用浏览器直接请求，通过服务端代理
|
||||
disableBrowserRequest: true,
|
||||
|
||||
// SDK 类型标识
|
||||
sdkType: 'comfyui',
|
||||
|
||||
// 不显示添加新模型按钮（模型通过配置管理）
|
||||
showAddNewModel: false,
|
||||
|
||||
// 显示 API 密钥配置（用于认证配置）
|
||||
showApiKey: true,
|
||||
|
||||
// 不显示连通性检查（图像生成不支持聊天接口检查）
|
||||
showChecker: false,
|
||||
|
||||
// 不显示模型获取器（使用预定义模型）
|
||||
showModelFetcher: false,
|
||||
},
|
||||
url: 'https://www.comfy.org/',
|
||||
};
|
||||
|
||||
export default ComfyUI;
|
||||
|
|
@ -15,6 +15,7 @@ import CerebrasProvider from './cerebras';
|
|||
import CloudflareProvider from './cloudflare';
|
||||
import CohereProvider from './cohere';
|
||||
import CometAPIProvider from './cometapi';
|
||||
import ComfyUIProvider from './comfyui';
|
||||
import DeepSeekProvider from './deepseek';
|
||||
import FalProvider from './fal';
|
||||
import FireworksAIProvider from './fireworksai';
|
||||
|
|
@ -129,6 +130,7 @@ export const DEFAULT_MODEL_PROVIDER_LIST = [
|
|||
OllamaProvider,
|
||||
OllamaCloudProvider,
|
||||
VLLMProvider,
|
||||
ComfyUIProvider,
|
||||
XinferenceProvider,
|
||||
AnthropicProvider,
|
||||
BedrockProvider,
|
||||
|
|
@ -214,6 +216,7 @@ export { default as CerebrasProviderCard } from './cerebras';
|
|||
export { default as CloudflareProviderCard } from './cloudflare';
|
||||
export { default as CohereProviderCard } from './cohere';
|
||||
export { default as CometAPIProviderCard } from './cometapi';
|
||||
export { default as ComfyUIProviderCard } from './comfyui';
|
||||
export { default as DeepSeekProviderCard } from './deepseek';
|
||||
export { default as FalProviderCard } from './fal';
|
||||
export { default as FireworksAIProviderCard } from './fireworksai';
|
||||
|
|
|
|||
|
|
@ -174,6 +174,14 @@ export const getLLMConfig = () => {
|
|||
ENABLED_BFL: z.boolean(),
|
||||
BFL_API_KEY: z.string().optional(),
|
||||
|
||||
ENABLED_COMFYUI: z.boolean(),
|
||||
COMFYUI_BASE_URL: z.string().optional(),
|
||||
COMFYUI_AUTH_TYPE: z.string().optional(),
|
||||
COMFYUI_API_KEY: z.string().optional(),
|
||||
COMFYUI_USERNAME: z.string().optional(),
|
||||
COMFYUI_PASSWORD: z.string().optional(),
|
||||
COMFYUI_CUSTOM_HEADERS: z.string().optional(),
|
||||
|
||||
ENABLED_MODELSCOPE: z.boolean(),
|
||||
MODELSCOPE_API_KEY: z.string().optional(),
|
||||
|
||||
|
|
@ -370,6 +378,14 @@ export const getLLMConfig = () => {
|
|||
ENABLED_BFL: !!process.env.BFL_API_KEY,
|
||||
BFL_API_KEY: process.env.BFL_API_KEY,
|
||||
|
||||
ENABLED_COMFYUI: process.env.ENABLED_COMFYUI !== '0',
|
||||
COMFYUI_BASE_URL: process.env.COMFYUI_BASE_URL,
|
||||
COMFYUI_AUTH_TYPE: process.env.COMFYUI_AUTH_TYPE,
|
||||
COMFYUI_API_KEY: process.env.COMFYUI_API_KEY,
|
||||
COMFYUI_USERNAME: process.env.COMFYUI_USERNAME,
|
||||
COMFYUI_PASSWORD: process.env.COMFYUI_PASSWORD,
|
||||
COMFYUI_CUSTOM_HEADERS: process.env.COMFYUI_CUSTOM_HEADERS,
|
||||
|
||||
ENABLED_MODELSCOPE: !!process.env.MODELSCOPE_API_KEY,
|
||||
MODELSCOPE_API_KEY: process.env.MODELSCOPE_API_KEY,
|
||||
|
||||
|
|
|
|||
|
|
@ -55,7 +55,9 @@ const getErrorAlertConfig = (
|
|||
}
|
||||
|
||||
case AgentRuntimeErrorType.OllamaServiceUnavailable:
|
||||
case AgentRuntimeErrorType.NoOpenAIAPIKey: {
|
||||
case AgentRuntimeErrorType.NoOpenAIAPIKey:
|
||||
case AgentRuntimeErrorType.ComfyUIServiceUnavailable:
|
||||
case AgentRuntimeErrorType.InvalidComfyUIArgs: {
|
||||
return {
|
||||
extraDefaultExpand: true,
|
||||
extraIsolate: true,
|
||||
|
|
|
|||
|
|
@ -147,6 +147,15 @@ export default {
|
|||
OllamaServiceUnavailable:
|
||||
'Ollama 服务连接失败，请检查 Ollama 是否运行正常，或是否正确设置 Ollama 的跨域配置',
|
||||
|
||||
InvalidComfyUIArgs: 'ComfyUI 配置不正确，请检查 ComfyUI 配置后重试',
|
||||
ComfyUIBizError: '请求 ComfyUI 服务出错，请根据以下信息排查或重试',
|
||||
ComfyUIServiceUnavailable:
|
||||
'ComfyUI 服务连接失败，请检查 ComfyUI 是否运行正常，或检查服务地址配置是否正确',
|
||||
ComfyUIEmptyResult: 'ComfyUI 未生成任何图像，请检查模型配置或重试',
|
||||
ComfyUIUploadFailed: 'ComfyUI 图片上传失败，请检查服务器连接或重试',
|
||||
ComfyUIWorkflowError: 'ComfyUI 工作流执行失败，请检查工作流配置',
|
||||
ComfyUIModelError: 'ComfyUI 模型加载失败，请检查模型文件是否存在',
|
||||
|
||||
AgentRuntimeError: 'Lobe AI Runtime 执行出错，请根据以下信息排查或重试',
|
||||
|
||||
// cloud
|
||||
|
|
@ -179,6 +188,11 @@ export default {
|
|||
title: '使用自定义 {{name}} API Key',
|
||||
},
|
||||
closeMessage: '关闭提示',
|
||||
comfyui: {
|
||||
description: '请输入正确的 {{name}} 认证信息即可开始生图',
|
||||
modifyBaseUrl: '修改 Comfy UI 服务地址',
|
||||
title: '确认你的 {{name}} 认证信息',
|
||||
},
|
||||
confirm: '确认并重试',
|
||||
oauth: {
|
||||
description: '管理员已开启统一登录认证，点击下方按钮登录，即可解锁应用',
|
||||
|
|
|
|||
|
|
@ -85,6 +85,58 @@ export default {
|
|||
title: 'Cloudflare 账户 ID / API 地址',
|
||||
},
|
||||
},
|
||||
comfyui: {
|
||||
apiKey: {
|
||||
desc: 'Bearer Token 认证所需的 API 密钥',
|
||||
placeholder: '请输入 API 密钥',
|
||||
required: '请输入 API 密钥',
|
||||
title: 'API 密钥',
|
||||
},
|
||||
authType: {
|
||||
desc: '选择与 ComfyUI 服务器的认证方式',
|
||||
options: {
|
||||
basic: '账号/密码',
|
||||
bearer: 'Bearer (API 密钥)',
|
||||
custom: '自定义请求头',
|
||||
none: '无需认证',
|
||||
},
|
||||
placeholder: '请选择认证类型',
|
||||
title: '认证类型',
|
||||
},
|
||||
baseURL: {
|
||||
desc: 'ComfyUI 网页访问地址',
|
||||
placeholder: 'http://127.0.0.1:8000',
|
||||
required: '请输入 ComfyUI 服务地址',
|
||||
title: 'ComfyUI 服务地址',
|
||||
},
|
||||
checker: {
|
||||
desc: '测试连接是否正确配置',
|
||||
title: '连通性检查',
|
||||
},
|
||||
customHeaders: {
|
||||
addButton: '添加请求头',
|
||||
deleteTooltip: '删除此请求头',
|
||||
desc: '自定义认证方式所需的请求头，格式为键值对',
|
||||
duplicateKeyError: '请求头键名不能重复',
|
||||
keyPlaceholder: '键名',
|
||||
required: '请输入自定义请求头',
|
||||
title: '自定义请求头',
|
||||
valuePlaceholder: '值',
|
||||
},
|
||||
password: {
|
||||
desc: '基本认证所需的密码',
|
||||
placeholder: '请输入密码',
|
||||
required: '请输入密码',
|
||||
title: '密码',
|
||||
},
|
||||
title: 'ComfyUI',
|
||||
username: {
|
||||
desc: '基本认证所需的用户名',
|
||||
placeholder: '请输入用户名',
|
||||
required: '请输入用户名',
|
||||
title: '用户名',
|
||||
},
|
||||
},
|
||||
createNewAiProvider: {
|
||||
apiKey: {
|
||||
placeholder: '请填写你的 API Key',
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import { ProviderConfig } from '@lobechat/types';
|
||||
import { extractEnabledModels, transformToAiModelList } from '@lobechat/utils';
|
||||
import { ModelProvider , AiFullModelCard } from 'model-bank';
|
||||
import { AiFullModelCard, ModelProvider } from 'model-bank';
|
||||
import * as AiModels from 'model-bank';
|
||||
|
||||
import { getLLMConfig } from '@/envs/llm';
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import {
|
|||
LobeAnthropicAI,
|
||||
LobeAzureOpenAI,
|
||||
LobeBedrockAI,
|
||||
LobeComfyUI,
|
||||
LobeDeepSeekAI,
|
||||
LobeGoogleAI,
|
||||
LobeGroq,
|
||||
|
|
@ -249,6 +250,49 @@ describe('initModelRuntimeWithUserPayload method', () => {
|
|||
expect(runtime['_runtime']).toBeInstanceOf(LobeStepfunAI);
|
||||
});
|
||||
|
||||
it('ComfyUI provider: with multiple auth types', async () => {
|
||||
// Test basic auth
|
||||
const basicAuthPayload: ClientSecretPayload = {
|
||||
authType: 'basic',
|
||||
username: 'test-user',
|
||||
password: 'test-pass',
|
||||
baseURL: 'http://localhost:8188',
|
||||
};
|
||||
let runtime = await initModelRuntimeWithUserPayload(ModelProvider.ComfyUI, basicAuthPayload);
|
||||
expect(runtime).toBeInstanceOf(ModelRuntime);
|
||||
expect(runtime['_runtime']).toBeInstanceOf(LobeComfyUI);
|
||||
expect(runtime['_runtime'].baseURL).toBe(basicAuthPayload.baseURL);
|
||||
|
||||
// Test bearer auth
|
||||
const bearerAuthPayload: ClientSecretPayload = {
|
||||
authType: 'bearer',
|
||||
apiKey: 'test-token',
|
||||
baseURL: 'http://localhost:8188',
|
||||
};
|
||||
runtime = await initModelRuntimeWithUserPayload(ModelProvider.ComfyUI, bearerAuthPayload);
|
||||
expect(runtime).toBeInstanceOf(ModelRuntime);
|
||||
expect(runtime['_runtime']).toBeInstanceOf(LobeComfyUI);
|
||||
|
||||
// Test custom auth
|
||||
const customAuthPayload: ClientSecretPayload = {
|
||||
authType: 'custom',
|
||||
customHeaders: { 'X-API-Key': 'secret123' },
|
||||
baseURL: 'http://localhost:8188',
|
||||
};
|
||||
runtime = await initModelRuntimeWithUserPayload(ModelProvider.ComfyUI, customAuthPayload);
|
||||
expect(runtime).toBeInstanceOf(ModelRuntime);
|
||||
expect(runtime['_runtime']).toBeInstanceOf(LobeComfyUI);
|
||||
|
||||
// Test none auth
|
||||
const noAuthPayload: ClientSecretPayload = {
|
||||
authType: 'none',
|
||||
baseURL: 'http://localhost:8188',
|
||||
};
|
||||
runtime = await initModelRuntimeWithUserPayload(ModelProvider.ComfyUI, noAuthPayload);
|
||||
expect(runtime).toBeInstanceOf(ModelRuntime);
|
||||
expect(runtime['_runtime']).toBeInstanceOf(LobeComfyUI);
|
||||
});
|
||||
|
||||
it('Unknown Provider: with apikey and endpoint, should initialize to OpenAi', async () => {
|
||||
const jwtPayload: ClientSecretPayload = {
|
||||
apiKey: 'user-unknown-key',
|
||||
|
|
@ -420,6 +464,28 @@ describe('initModelRuntimeWithUserPayload method', () => {
|
|||
expect(runtime['_runtime'].baseURL).toBe('https://dashscope.aliyuncs.com/compatible-mode/v1');
|
||||
});
|
||||
|
||||
it('ComfyUI provider: without user payload (using environment variables)', async () => {
|
||||
const jwtPayload: ClientSecretPayload = {};
|
||||
const runtime = await initModelRuntimeWithUserPayload(ModelProvider.ComfyUI, jwtPayload);
|
||||
|
||||
expect(runtime).toBeInstanceOf(ModelRuntime);
|
||||
expect(runtime['_runtime']).toBeInstanceOf(LobeComfyUI);
|
||||
// Should use environment variable defaults
|
||||
expect(runtime['_runtime'].baseURL).toBe('http://127.0.0.1:8000');
|
||||
});
|
||||
|
||||
it('ComfyUI provider: partial payload (mixed with env vars)', async () => {
|
||||
const jwtPayload: ClientSecretPayload = {
|
||||
baseURL: 'http://custom-comfyui:8188',
|
||||
// authType, username, password will come from env vars
|
||||
};
|
||||
const runtime = await initModelRuntimeWithUserPayload(ModelProvider.ComfyUI, jwtPayload);
|
||||
|
||||
expect(runtime).toBeInstanceOf(ModelRuntime);
|
||||
expect(runtime['_runtime']).toBeInstanceOf(LobeComfyUI);
|
||||
expect(runtime['_runtime'].baseURL).toBe('http://custom-comfyui:8188');
|
||||
});
|
||||
|
||||
it('Unknown Provider', async () => {
|
||||
const jwtPayload = {};
|
||||
const runtime = await initModelRuntimeWithUserPayload('unknown', jwtPayload);
|
||||
|
|
|
|||
|
|
@ -89,6 +89,41 @@ const getParamsFromPayload = (provider: string, payload: ClientSecretPayload) =>
|
|||
return { apiKey, baseURLOrAccountID };
|
||||
}
|
||||
|
||||
case ModelProvider.ComfyUI: {
|
||||
const {
|
||||
COMFYUI_BASE_URL,
|
||||
COMFYUI_AUTH_TYPE,
|
||||
COMFYUI_API_KEY,
|
||||
COMFYUI_USERNAME,
|
||||
COMFYUI_PASSWORD,
|
||||
COMFYUI_CUSTOM_HEADERS,
|
||||
} = llmConfig;
|
||||
|
||||
// ComfyUI specific handling with environment variables fallback
|
||||
const baseURL = payload?.baseURL || COMFYUI_BASE_URL || 'http://127.0.0.1:8000';
|
||||
|
||||
// ComfyUI supports multiple auth types: none, basic, bearer, custom
|
||||
// Extract all relevant auth fields from the payload or environment
|
||||
const authType = payload?.authType || COMFYUI_AUTH_TYPE || 'none';
|
||||
const apiKey = payload?.apiKey || COMFYUI_API_KEY;
|
||||
const username = payload?.username || COMFYUI_USERNAME;
|
||||
const password = payload?.password || COMFYUI_PASSWORD;
|
||||
|
||||
// Parse customHeaders from JSON string (similar to Vertex AI credentials handling)
|
||||
// Support both payload object and environment variable JSON string
|
||||
const customHeaders = payload?.customHeaders || safeParseJSON(COMFYUI_CUSTOM_HEADERS);
|
||||
|
||||
// Return all authentication parameters
|
||||
return {
|
||||
apiKey,
|
||||
authType,
|
||||
baseURL,
|
||||
customHeaders,
|
||||
password,
|
||||
username,
|
||||
};
|
||||
}
|
||||
|
||||
case ModelProvider.GiteeAI: {
|
||||
const { GITEE_AI_API_KEY } = llmConfig;
|
||||
|
||||
|
|
|
|||
|
|
@ -59,15 +59,79 @@ const checkAbortSignal = (signal: AbortSignal) => {
|
|||
|
||||
/**
|
||||
* Categorizes errors into appropriate AsyncTaskErrorType
|
||||
* Returns the original error message if available, otherwise returns the error type as message
|
||||
* Client should handle localization based on errorType
|
||||
*/
|
||||
const categorizeError = (
|
||||
error: any,
|
||||
isAborted: boolean,
|
||||
): { errorMessage: string; errorType: AsyncTaskErrorType } => {
|
||||
log('🔥🔥🔥 [ASYNC] categorizeError called:', {
|
||||
errorMessage: error?.message,
|
||||
errorName: error?.name,
|
||||
errorStatus: error?.status,
|
||||
errorType: error?.errorType,
|
||||
fullError: JSON.stringify(error, null, 2),
|
||||
isAborted,
|
||||
});
|
||||
// Handle Comfy UI errors
|
||||
if (error.errorType === AgentRuntimeErrorType.ComfyUIServiceUnavailable) {
|
||||
return {
|
||||
errorMessage:
|
||||
error.error?.message || error.message || AgentRuntimeErrorType.ComfyUIServiceUnavailable,
|
||||
errorType: AsyncTaskErrorType.InvalidProviderAPIKey,
|
||||
};
|
||||
}
|
||||
|
||||
if (error.errorType === AgentRuntimeErrorType.ComfyUIBizError) {
|
||||
return {
|
||||
errorMessage: error.error?.message || error.message || AgentRuntimeErrorType.ComfyUIBizError,
|
||||
errorType: AsyncTaskErrorType.ServerError,
|
||||
};
|
||||
}
|
||||
|
||||
if (error.errorType === AgentRuntimeErrorType.ComfyUIWorkflowError) {
|
||||
return {
|
||||
errorMessage:
|
||||
error.error?.message || error.message || AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
errorType: AsyncTaskErrorType.ServerError,
|
||||
};
|
||||
}
|
||||
|
||||
if (error.errorType === AgentRuntimeErrorType.ComfyUIModelError) {
|
||||
return {
|
||||
errorMessage:
|
||||
error.error?.message || error.message || AgentRuntimeErrorType.ComfyUIModelError,
|
||||
errorType: AsyncTaskErrorType.ModelNotFound,
|
||||
};
|
||||
}
|
||||
|
||||
if (error.errorType === AgentRuntimeErrorType.ConnectionCheckFailed) {
|
||||
return {
|
||||
errorMessage: error.message || AgentRuntimeErrorType.ConnectionCheckFailed,
|
||||
errorType: AsyncTaskErrorType.ServerError,
|
||||
};
|
||||
}
|
||||
|
||||
if (error.errorType === AgentRuntimeErrorType.PermissionDenied) {
|
||||
return {
|
||||
errorMessage: error.error?.message || error.message || AgentRuntimeErrorType.PermissionDenied,
|
||||
errorType: AsyncTaskErrorType.InvalidProviderAPIKey,
|
||||
};
|
||||
}
|
||||
|
||||
if (error.errorType === AgentRuntimeErrorType.ModelNotFound) {
|
||||
return {
|
||||
errorMessage: error.error?.message || error.message || AgentRuntimeErrorType.ModelNotFound,
|
||||
errorType: AsyncTaskErrorType.ModelNotFound,
|
||||
};
|
||||
}
|
||||
|
||||
// FIXME: 401 的问题应该放到 agentRuntime 中处理会更好
|
||||
if (error.errorType === AgentRuntimeErrorType.InvalidProviderAPIKey || error?.status === 401) {
|
||||
return {
|
||||
errorMessage: 'Invalid provider API key, please check your API key',
|
||||
errorMessage:
|
||||
error.error?.message || error.message || AgentRuntimeErrorType.InvalidProviderAPIKey,
|
||||
errorType: AsyncTaskErrorType.InvalidProviderAPIKey,
|
||||
};
|
||||
}
|
||||
|
|
@ -81,27 +145,27 @@ const categorizeError = (
|
|||
|
||||
if (isAborted || error.message?.includes('aborted')) {
|
||||
return {
|
||||
errorMessage: 'Image generation task timed out, please try again',
|
||||
errorMessage: AsyncTaskErrorType.Timeout,
|
||||
errorType: AsyncTaskErrorType.Timeout,
|
||||
};
|
||||
}
|
||||
|
||||
if (error.message?.includes('timeout') || error.name === 'TimeoutError') {
|
||||
return {
|
||||
errorMessage: 'Image generation task timed out, please try again',
|
||||
errorMessage: AsyncTaskErrorType.Timeout,
|
||||
errorType: AsyncTaskErrorType.Timeout,
|
||||
};
|
||||
}
|
||||
|
||||
if (error.message?.includes('network') || error.name === 'NetworkError') {
|
||||
return {
|
||||
errorMessage: error.message || 'Network error occurred during image generation',
|
||||
errorMessage: error.message || AsyncTaskErrorType.ServerError,
|
||||
errorType: AsyncTaskErrorType.ServerError,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
errorMessage: error.message || 'Unknown error occurred during image generation',
|
||||
errorMessage: error.message || AsyncTaskErrorType.ServerError,
|
||||
errorType: AsyncTaskErrorType.ServerError,
|
||||
};
|
||||
};
|
||||
|
|
@ -139,7 +203,6 @@ export const imageRouter = router({
|
|||
|
||||
// Check if operation has been cancelled
|
||||
checkAbortSignal(signal);
|
||||
|
||||
log('Agent runtime initialized, calling createImage');
|
||||
const response = await agentRuntime.createImage!({
|
||||
model,
|
||||
|
|
@ -171,8 +234,24 @@ export const imageRouter = router({
|
|||
|
||||
log('Transforming image for generation');
|
||||
const { imageUrl, width, height } = response;
|
||||
const { image, thumbnailImage } =
|
||||
await ctx.generationService.transformImageForGeneration(imageUrl);
|
||||
|
||||
// Extract ComfyUI authentication headers if provider is ComfyUI
|
||||
let authHeaders: Record<string, string> | undefined;
|
||||
if (provider === 'comfyui') {
|
||||
// Use the public interface method to get auth headers
|
||||
// This avoids accessing private members and exposing credentials
|
||||
authHeaders = agentRuntime.getAuthHeaders();
|
||||
if (authHeaders) {
|
||||
log('Using authentication headers for ComfyUI image download');
|
||||
} else {
|
||||
log('No authentication configured for ComfyUI');
|
||||
}
|
||||
}
|
||||
|
||||
const { image, thumbnailImage } = await ctx.generationService.transformImageForGeneration(
|
||||
imageUrl,
|
||||
authHeaders,
|
||||
);
|
||||
|
||||
// Check if operation has been cancelled
|
||||
checkAbortSignal(signal);
|
||||
|
|
|
|||
96
src/server/routers/lambda/comfyui.ts
Normal file
96
src/server/routers/lambda/comfyui.ts
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
import type { ComfyUIKeyVault } from '@lobechat/types';
|
||||
import { z } from 'zod';
|
||||
|
||||
import { authedProcedure, router } from '@/libs/trpc/lambda';
|
||||
// Import Framework layer services
|
||||
import { ComfyUIClientService } from '@/server/services/comfyui/core/comfyUIClientService';
|
||||
import { ImageService } from '@/server/services/comfyui/core/imageService';
|
||||
import { ModelResolverService } from '@/server/services/comfyui/core/modelResolverService';
|
||||
import { WorkflowBuilderService } from '@/server/services/comfyui/core/workflowBuilderService';
|
||||
import type { WorkflowContext } from '@/server/services/comfyui/types';
|
||||
|
||||
// ComfyUI params validation - only validate required fields
|
||||
// Other RuntimeImageGenParams fields are passed through automatically
|
||||
const ComfyUIParamsSchema = z
|
||||
.object({
|
||||
prompt: z.string(), // 只验证必需字段
|
||||
})
|
||||
.passthrough();
|
||||
|
||||
/**
|
||||
* ComfyUI tRPC Router
|
||||
* Exposes Framework layer services to Runtime layer
|
||||
*/
|
||||
export const comfyuiRouter = router({
|
||||
/**
|
||||
* Create image with complete business logic
|
||||
*/
|
||||
createImage: authedProcedure
|
||||
.input(
|
||||
z.object({
|
||||
model: z.string(),
|
||||
options: z.custom<ComfyUIKeyVault>().optional(),
|
||||
params: ComfyUIParamsSchema,
|
||||
}),
|
||||
)
|
||||
.mutation(async ({ input }) => {
|
||||
const { model, params, options = {} } = input;
|
||||
|
||||
// Initialize Framework layer services
|
||||
const clientService = new ComfyUIClientService(options);
|
||||
const modelResolverService = new ModelResolverService(clientService);
|
||||
|
||||
// Create workflow context
|
||||
const context: WorkflowContext = {
|
||||
clientService,
|
||||
modelResolverService,
|
||||
};
|
||||
|
||||
const workflowBuilderService = new WorkflowBuilderService(context);
|
||||
|
||||
// Initialize image service with all dependencies
|
||||
const imageService = new ImageService(
|
||||
clientService,
|
||||
modelResolverService,
|
||||
workflowBuilderService,
|
||||
);
|
||||
|
||||
// Execute image creation
|
||||
return imageService.createImage({
|
||||
model,
|
||||
params,
|
||||
});
|
||||
}),
|
||||
|
||||
/**
|
||||
* Get authentication headers for image downloads
|
||||
*/
|
||||
getAuthHeaders: authedProcedure
|
||||
.input(
|
||||
z.object({
|
||||
options: z.custom<ComfyUIKeyVault>().optional(),
|
||||
}),
|
||||
)
|
||||
.query(async ({ input }) => {
|
||||
const clientService = new ComfyUIClientService(input.options || {});
|
||||
return clientService.getAuthHeaders();
|
||||
}),
|
||||
|
||||
/**
|
||||
* Get available models
|
||||
*/
|
||||
getModels: authedProcedure
|
||||
.input(
|
||||
z.object({
|
||||
options: z.custom<ComfyUIKeyVault>().optional(),
|
||||
}),
|
||||
)
|
||||
.query(async ({ input }) => {
|
||||
const clientService = new ComfyUIClientService(input.options || {});
|
||||
const modelResolverService = new ModelResolverService(clientService);
|
||||
|
||||
return modelResolverService.getAvailableModelFiles();
|
||||
}),
|
||||
});
|
||||
|
||||
export type ComfyUIRouter = typeof comfyuiRouter;
|
||||
|
|
@ -9,6 +9,7 @@ import { aiModelRouter } from './aiModel';
|
|||
import { aiProviderRouter } from './aiProvider';
|
||||
import { apiKeyRouter } from './apiKey';
|
||||
import { chunkRouter } from './chunk';
|
||||
import { comfyuiRouter } from './comfyui';
|
||||
import { configRouter } from './config';
|
||||
import { documentRouter } from './document';
|
||||
import { exporterRouter } from './exporter';
|
||||
|
|
@ -38,6 +39,7 @@ export const lambdaRouter = router({
|
|||
aiProvider: aiProviderRouter,
|
||||
apiKey: apiKeyRouter,
|
||||
chunk: chunkRouter,
|
||||
comfyui: comfyuiRouter,
|
||||
config: configRouter,
|
||||
document: documentRouter,
|
||||
exporter: exporterRouter,
|
||||
|
|
|
|||
146
src/server/services/comfyui/__tests__/config/constants.test.ts
Normal file
146
src/server/services/comfyui/__tests__/config/constants.test.ts
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import {
|
||||
COMFYUI_DEFAULTS,
|
||||
CUSTOM_SD_CONFIG,
|
||||
DEFAULT_NEGATIVE_PROMPT,
|
||||
FLUX_MODEL_CONFIG,
|
||||
SD_MODEL_CONFIG,
|
||||
WORKFLOW_DEFAULTS,
|
||||
} from '@/server/services/comfyui/config/constants';
|
||||
import { STYLE_KEYWORDS } from '@/server/services/comfyui/config/promptToolConst';
|
||||
|
||||
describe('ComfyUI Constants', () => {
|
||||
describe('COMFYUI_DEFAULTS', () => {
|
||||
it('should be a valid object', () => {
|
||||
expect(typeof COMFYUI_DEFAULTS).toBe('object');
|
||||
expect(COMFYUI_DEFAULTS).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('FLUX_MODEL_CONFIG', () => {
|
||||
it('should have correct filename prefixes', () => {
|
||||
expect(FLUX_MODEL_CONFIG.FILENAME_PREFIXES.SCHNELL).toContain('FLUX_Schnell');
|
||||
expect(FLUX_MODEL_CONFIG.FILENAME_PREFIXES.DEV).toContain('FLUX_Dev');
|
||||
expect(FLUX_MODEL_CONFIG.FILENAME_PREFIXES.KONTEXT).toContain('FLUX_Kontext');
|
||||
expect(FLUX_MODEL_CONFIG.FILENAME_PREFIXES.KREA).toContain('FLUX_Krea');
|
||||
});
|
||||
|
||||
it('should have all required prefixes', () => {
|
||||
const expectedKeys = ['SCHNELL', 'DEV', 'KONTEXT', 'KREA'];
|
||||
expect(Object.keys(FLUX_MODEL_CONFIG.FILENAME_PREFIXES)).toEqual(
|
||||
expect.arrayContaining(expectedKeys),
|
||||
);
|
||||
});
|
||||
|
||||
it('should be a readonly object (TypeScript as const)', () => {
|
||||
// `as const` provides readonly types in TypeScript, not runtime freezing
|
||||
expect(typeof FLUX_MODEL_CONFIG).toBe('object');
|
||||
});
|
||||
});
|
||||
|
||||
describe('WORKFLOW_DEFAULTS', () => {
|
||||
it('should have valid workflow parameters', () => {
|
||||
expect(WORKFLOW_DEFAULTS.IMAGE.BATCH_SIZE).toBeGreaterThan(0);
|
||||
expect(WORKFLOW_DEFAULTS.SAMPLING.DENOISE).toBeGreaterThanOrEqual(0);
|
||||
expect(WORKFLOW_DEFAULTS.SAMPLING.DENOISE).toBeLessThanOrEqual(1);
|
||||
expect(WORKFLOW_DEFAULTS.SAMPLING.MAX_SHIFT).toBeGreaterThan(0);
|
||||
expect(WORKFLOW_DEFAULTS.SD3.SHIFT).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('should be a readonly object (TypeScript as const)', () => {
|
||||
// `as const` provides readonly types in TypeScript, not runtime freezing
|
||||
expect(typeof WORKFLOW_DEFAULTS).toBe('object');
|
||||
});
|
||||
});
|
||||
|
||||
describe('STYLE_KEYWORDS', () => {
|
||||
it('should have all required categories', () => {
|
||||
const expectedCategories = [
|
||||
'ARTISTS',
|
||||
'ART_STYLES',
|
||||
'LIGHTING',
|
||||
'PHOTOGRAPHY',
|
||||
'QUALITY',
|
||||
'RENDERING',
|
||||
];
|
||||
expect(Object.keys(STYLE_KEYWORDS)).toEqual(expect.arrayContaining(expectedCategories));
|
||||
});
|
||||
|
||||
it('should have non-empty arrays for each category', () => {
|
||||
Object.values(STYLE_KEYWORDS).forEach((keywords) => {
|
||||
expect(Array.isArray(keywords)).toBe(true);
|
||||
expect(keywords.length).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
|
||||
it('should contain expected artist keywords', () => {
|
||||
expect(STYLE_KEYWORDS.ARTISTS).toEqual(
|
||||
expect.arrayContaining(['by greg rutkowski', 'by artgerm', 'trending on artstation']),
|
||||
);
|
||||
});
|
||||
|
||||
it('should contain expected art style keywords', () => {
|
||||
expect(STYLE_KEYWORDS.ART_STYLES).toEqual(
|
||||
expect.arrayContaining(['photorealistic', 'anime', 'digital art', '3d render']),
|
||||
);
|
||||
});
|
||||
|
||||
it('should contain expected lighting keywords', () => {
|
||||
expect(STYLE_KEYWORDS.LIGHTING).toEqual(
|
||||
expect.arrayContaining(['dramatic lighting', 'studio lighting', 'soft lighting']),
|
||||
);
|
||||
});
|
||||
|
||||
it('should contain expected photography keywords', () => {
|
||||
expect(STYLE_KEYWORDS.PHOTOGRAPHY).toEqual(
|
||||
expect.arrayContaining([
|
||||
'depth of field',
|
||||
'bokeh',
|
||||
'35mm photograph',
|
||||
'professional photograph',
|
||||
]),
|
||||
);
|
||||
});
|
||||
|
||||
it('should contain expected quality keywords', () => {
|
||||
expect(STYLE_KEYWORDS.QUALITY).toEqual(
|
||||
expect.arrayContaining([
|
||||
'masterpiece',
|
||||
'best quality',
|
||||
'high quality',
|
||||
'extremely detailed',
|
||||
]),
|
||||
);
|
||||
});
|
||||
|
||||
it('should contain expected rendering keywords', () => {
|
||||
expect(STYLE_KEYWORDS.RENDERING).toEqual(
|
||||
expect.arrayContaining(['octane render', 'unreal engine', 'ray tracing', 'cycles render']),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('DEFAULT_NEGATIVE_PROMPT', () => {
|
||||
it('should be defined and non-empty', () => {
|
||||
expect(DEFAULT_NEGATIVE_PROMPT).toBeDefined();
|
||||
expect(DEFAULT_NEGATIVE_PROMPT).not.toBe('');
|
||||
});
|
||||
});
|
||||
|
||||
describe('CUSTOM_SD_CONFIG', () => {
|
||||
it('should have model and VAE filenames', () => {
|
||||
expect(CUSTOM_SD_CONFIG.MODEL_FILENAME).toBeDefined();
|
||||
expect(CUSTOM_SD_CONFIG.VAE_FILENAME).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('SD_MODEL_CONFIG', () => {
|
||||
it('should have correct filename prefixes', () => {
|
||||
expect(SD_MODEL_CONFIG.FILENAME_PREFIXES.SD15).toContain('SD15');
|
||||
expect(SD_MODEL_CONFIG.FILENAME_PREFIXES.SD35).toContain('SD35');
|
||||
expect(SD_MODEL_CONFIG.FILENAME_PREFIXES.SDXL).toContain('SDXL');
|
||||
expect(SD_MODEL_CONFIG.FILENAME_PREFIXES.CUSTOM).toContain('CustomSD');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,277 @@
|
|||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import { MODEL_REGISTRY } from '@/server/services/comfyui/config/modelRegistry';
|
||||
import {
|
||||
getAllModelNames,
|
||||
getModelConfig,
|
||||
getModelsByVariant,
|
||||
} from '@/server/services/comfyui/utils/staticModelLookup';
|
||||
|
||||
describe('ModelRegistry', () => {
|
||||
describe('MODEL_REGISTRY', () => {
|
||||
it('should be a non-empty object with valid structure', () => {
|
||||
expect(typeof MODEL_REGISTRY).toBe('object');
|
||||
expect(Object.keys(MODEL_REGISTRY).length).toBeGreaterThan(0);
|
||||
|
||||
// Check that all models have required fields
|
||||
Object.entries(MODEL_REGISTRY).forEach(([, config]) => {
|
||||
expect(config).toBeDefined();
|
||||
expect(config.modelFamily).toBeDefined();
|
||||
expect(config.priority).toBeTypeOf('number');
|
||||
if (config.recommendedDtype) {
|
||||
expect(
|
||||
['default', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2'].includes(
|
||||
config.recommendedDtype,
|
||||
),
|
||||
).toBe(true);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
it('should contain essential model families', () => {
|
||||
const modelFamilies = Object.values(MODEL_REGISTRY).map((c) => c.modelFamily);
|
||||
const uniqueFamilies = [...new Set(modelFamilies)];
|
||||
|
||||
// Should have at least one model family and FLUX should be included
|
||||
expect(uniqueFamilies.length).toBeGreaterThan(0);
|
||||
expect(uniqueFamilies).toContain('FLUX');
|
||||
});
|
||||
|
||||
it('should have valid priority ranges', () => {
|
||||
Object.entries(MODEL_REGISTRY).forEach(([, config]) => {
|
||||
// Priorities should be positive numbers
|
||||
expect(config.priority).toBeGreaterThan(0);
|
||||
expect(config.priority).toBeLessThanOrEqual(10);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('getModelConfig', () => {
|
||||
it('should return model config for valid name', () => {
|
||||
// Get any available FLUX model instead of hardcoding
|
||||
const allModelNames = getAllModelNames();
|
||||
const fluxModels = allModelNames.filter((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return config?.modelFamily === 'FLUX';
|
||||
});
|
||||
|
||||
expect(fluxModels.length).toBeGreaterThan(0);
|
||||
|
||||
const config = getModelConfig(fluxModels[0]);
|
||||
expect(config).toBeDefined();
|
||||
expect(config?.modelFamily).toBe('FLUX');
|
||||
});
|
||||
|
||||
it('should return undefined for invalid name', () => {
|
||||
const config = getModelConfig('nonexistent.safetensors');
|
||||
expect(config).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAllModelNames', () => {
|
||||
it('should return all model names', () => {
|
||||
const names = getAllModelNames();
|
||||
expect(names.length).toBeGreaterThan(0);
|
||||
// Check if at least one FLUX model exists instead of hardcoding
|
||||
const hasFluxModel = names.some((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return config?.modelFamily === 'FLUX';
|
||||
});
|
||||
expect(hasFluxModel).toBe(true);
|
||||
});
|
||||
|
||||
it('should return unique names', () => {
|
||||
const names = getAllModelNames();
|
||||
const uniqueNames = [...new Set(names)];
|
||||
expect(uniqueNames.length).toBe(names.length);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getModelsByVariant', () => {
|
||||
it('should return model names for valid variant', () => {
|
||||
const modelNames = getModelsByVariant('dev');
|
||||
expect(modelNames.length).toBeGreaterThan(0);
|
||||
expect(Array.isArray(modelNames)).toBe(true);
|
||||
|
||||
// Verify all returned names are strings and correspond to dev variant models
|
||||
modelNames.forEach((name) => {
|
||||
expect(typeof name).toBe('string');
|
||||
const config = getModelConfig(name);
|
||||
expect(config).toBeDefined();
|
||||
expect(config?.variant).toBe('dev');
|
||||
});
|
||||
});
|
||||
|
||||
it('should return models sorted by priority', () => {
|
||||
const modelNames = getModelsByVariant('dev');
|
||||
expect(modelNames.length).toBeGreaterThan(1);
|
||||
|
||||
// Verify priority sorting (lower priority number = higher priority)
|
||||
for (let i = 0; i < modelNames.length - 1; i++) {
|
||||
const config1 = getModelConfig(modelNames[i]);
|
||||
const config2 = getModelConfig(modelNames[i + 1]);
|
||||
expect(config1?.priority).toBeLessThanOrEqual(config2?.priority || 0);
|
||||
}
|
||||
});
|
||||
|
||||
it('should return empty array for invalid variant', () => {
|
||||
const models = getModelsByVariant('nonexistent' as any);
|
||||
expect(models).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getModelConfig with options', () => {
|
||||
it('should support case-insensitive lookup', () => {
|
||||
// Get any FLUX dev model for testing case-insensitive lookup
|
||||
const allModels = getAllModelNames();
|
||||
const fluxDevModel = allModels.find((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return config?.modelFamily === 'FLUX' && config?.variant === 'dev';
|
||||
});
|
||||
|
||||
if (fluxDevModel) {
|
||||
const config = getModelConfig(fluxDevModel.toUpperCase(), { caseInsensitive: true });
|
||||
expect(config).toBeDefined();
|
||||
expect(config?.modelFamily).toBe('FLUX');
|
||||
expect(config?.variant).toBe('dev');
|
||||
} else {
|
||||
// If no dev variant exists, test with any FLUX model
|
||||
const fluxModel = allModels.find((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return config?.modelFamily === 'FLUX';
|
||||
});
|
||||
expect(fluxModel).toBeDefined();
|
||||
|
||||
const config = getModelConfig(fluxModel!.toUpperCase(), { caseInsensitive: true });
|
||||
expect(config).toBeDefined();
|
||||
expect(config?.modelFamily).toBe('FLUX');
|
||||
}
|
||||
});
|
||||
|
||||
it('should return undefined for non-matching case without caseInsensitive option', () => {
|
||||
// Find any FLUX model and test uppercase version without case-insensitive flag
|
||||
const allModels = getAllModelNames();
|
||||
const fluxModel = allModels.find((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return config?.modelFamily === 'FLUX';
|
||||
});
|
||||
|
||||
if (fluxModel) {
|
||||
const config = getModelConfig(fluxModel.toUpperCase());
|
||||
expect(config).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it('should filter by variant', () => {
|
||||
// Find models with different variants for testing
|
||||
const allModels = getAllModelNames();
|
||||
const devModel = allModels.find((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return config?.variant === 'dev';
|
||||
});
|
||||
|
||||
if (devModel) {
|
||||
// Test matching variant
|
||||
const config = getModelConfig(devModel, { variant: 'dev' });
|
||||
expect(config).toBeDefined();
|
||||
expect(config?.variant).toBe('dev');
|
||||
|
||||
// Test non-matching variant
|
||||
const nonMatchingConfig = getModelConfig(devModel, { variant: 'schnell' });
|
||||
expect(nonMatchingConfig).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it('should filter by modelFamily', () => {
|
||||
// 测试 SD3.5 模型家族
|
||||
const config = getModelConfig('sd3.5_large.safetensors', { modelFamily: 'SD3' });
|
||||
expect(config).toBeDefined();
|
||||
expect(config?.modelFamily).toBe('SD3');
|
||||
|
||||
// 测试不匹配的 modelFamily
|
||||
const nonMatchingConfig = getModelConfig('sd3.5_large.safetensors', { modelFamily: 'FLUX' });
|
||||
expect(nonMatchingConfig).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should filter by priority', () => {
|
||||
// Find a model with priority 1 for testing
|
||||
const allModels = getAllModelNames();
|
||||
const priority1Model = allModels.find((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return config?.priority === 1;
|
||||
});
|
||||
|
||||
if (priority1Model) {
|
||||
const config = getModelConfig(priority1Model, { priority: 1 });
|
||||
expect(config).toBeDefined();
|
||||
|
||||
// Test non-matching priority
|
||||
const nonMatchingConfig = getModelConfig(priority1Model, { priority: 999 });
|
||||
expect(nonMatchingConfig).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it('should filter by recommendedDtype', () => {
|
||||
// flux_shakker_labs_union_pro-fp8_e4m3fn 有 fp8_e4m3fn
|
||||
const config = getModelConfig('flux_shakker_labs_union_pro-fp8_e4m3fn.safetensors', {
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
});
|
||||
expect(config).toBeDefined();
|
||||
expect(config?.recommendedDtype).toBe('fp8_e4m3fn');
|
||||
|
||||
// 测试不匹配的 recommendedDtype
|
||||
const nonMatchingConfig = getModelConfig(
|
||||
'flux_shakker_labs_union_pro-fp8_e4m3fn.safetensors',
|
||||
{ recommendedDtype: 'default' },
|
||||
);
|
||||
expect(nonMatchingConfig).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should combine multiple filters', () => {
|
||||
// Find a FLUX dev model with priority 1 for testing
|
||||
const allModels = getAllModelNames();
|
||||
const testModel = allModels.find((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return (
|
||||
config?.modelFamily === 'FLUX' && config?.variant === 'dev' && config?.priority === 1
|
||||
);
|
||||
});
|
||||
|
||||
if (testModel) {
|
||||
// All filters match
|
||||
const config = getModelConfig(testModel, {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
variant: 'dev',
|
||||
});
|
||||
expect(config).toBeDefined();
|
||||
|
||||
// One filter doesn't match
|
||||
const nonMatchingConfig = getModelConfig(testModel, {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 999, // Wrong priority
|
||||
variant: 'dev',
|
||||
});
|
||||
expect(nonMatchingConfig).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle case-insensitive with other filters', () => {
|
||||
// Find a FLUX dev model for testing
|
||||
const allModels = getAllModelNames();
|
||||
const fluxDevModel = allModels.find((name) => {
|
||||
const config = getModelConfig(name);
|
||||
return config?.modelFamily === 'FLUX' && config?.variant === 'dev';
|
||||
});
|
||||
|
||||
if (fluxDevModel) {
|
||||
const config = getModelConfig(fluxDevModel.toUpperCase(), {
|
||||
caseInsensitive: true,
|
||||
modelFamily: 'FLUX',
|
||||
variant: 'dev',
|
||||
});
|
||||
expect(config).toBeDefined();
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,357 @@
|
|||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import {
|
||||
COMPOUND_STYLES,
|
||||
STYLE_ADJECTIVE_PATTERNS,
|
||||
STYLE_KEYWORDS,
|
||||
STYLE_SYNONYMS,
|
||||
extractStyleAdjectives,
|
||||
getAllStyleKeywords,
|
||||
getCompoundStyles,
|
||||
isStyleAdjective,
|
||||
normalizeStyleTerm,
|
||||
} from '@/server/services/comfyui/config/promptToolConst';
|
||||
|
||||
describe('promptToolConst', () => {
|
||||
describe('STYLE_KEYWORDS', () => {
|
||||
it('should have all expected categories', () => {
|
||||
const expectedCategories = [
|
||||
'ARTISTS',
|
||||
'ART_STYLES',
|
||||
'LIGHTING',
|
||||
'PHOTOGRAPHY',
|
||||
'QUALITY',
|
||||
'RENDERING',
|
||||
'COLOR_MOOD',
|
||||
'TEXTURE_MATERIAL',
|
||||
];
|
||||
expect(Object.keys(STYLE_KEYWORDS)).toEqual(expectedCategories);
|
||||
});
|
||||
|
||||
it('should have expanded keywords in each category', () => {
|
||||
// Minimum expectations to allow expansion
|
||||
expect(STYLE_KEYWORDS.ARTISTS.length).toBeGreaterThanOrEqual(20);
|
||||
expect(STYLE_KEYWORDS.ART_STYLES.length).toBeGreaterThanOrEqual(52);
|
||||
expect(STYLE_KEYWORDS.LIGHTING.length).toBeGreaterThanOrEqual(37);
|
||||
expect(STYLE_KEYWORDS.PHOTOGRAPHY.length).toBeGreaterThanOrEqual(49);
|
||||
expect(STYLE_KEYWORDS.QUALITY.length).toBeGreaterThanOrEqual(39);
|
||||
expect(STYLE_KEYWORDS.RENDERING.length).toBeGreaterThanOrEqual(39);
|
||||
expect(STYLE_KEYWORDS.COLOR_MOOD.length).toBeGreaterThanOrEqual(56);
|
||||
expect(STYLE_KEYWORDS.TEXTURE_MATERIAL.length).toBeGreaterThanOrEqual(60);
|
||||
});
|
||||
|
||||
it('should not have duplicate keywords within categories', () => {
|
||||
Object.entries(STYLE_KEYWORDS).forEach(([, keywords]: [string, readonly string[]]) => {
|
||||
const uniqueKeywords = [...new Set(keywords)];
|
||||
expect(keywords.length).toBe(uniqueKeywords.length);
|
||||
});
|
||||
});
|
||||
|
||||
it('should have lowercase keywords', () => {
|
||||
Object.values(STYLE_KEYWORDS).forEach((keywords: readonly string[]) => {
|
||||
keywords.forEach((keyword: string) => {
|
||||
expect(keyword).toBe(keyword.toLowerCase());
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('STYLE_SYNONYMS', () => {
|
||||
it('should have synonym mappings', () => {
|
||||
// Minimum number of synonym groups to allow expansion
|
||||
expect(Object.keys(STYLE_SYNONYMS).length).toBeGreaterThanOrEqual(15);
|
||||
expect(Object.keys(STYLE_SYNONYMS).length).toBeLessThanOrEqual(50); // Reasonable upper bound
|
||||
});
|
||||
|
||||
it('should map common variations', () => {
|
||||
expect(STYLE_SYNONYMS['photorealistic']).toContain('photo-realistic');
|
||||
expect(STYLE_SYNONYMS['photorealistic']).toContain('photo realistic');
|
||||
expect(STYLE_SYNONYMS['photorealistic']).toContain('lifelike');
|
||||
|
||||
expect(STYLE_SYNONYMS['4k']).toContain('4k resolution');
|
||||
expect(STYLE_SYNONYMS['4k']).toContain('ultra hd');
|
||||
|
||||
expect(STYLE_SYNONYMS['cinematic']).toContain('filmic');
|
||||
expect(STYLE_SYNONYMS['cinematic']).toContain('movie-like');
|
||||
});
|
||||
|
||||
it('should have unique synonyms for each key', () => {
|
||||
Object.entries(STYLE_SYNONYMS).forEach(([, synonyms]: [string, string[]]) => {
|
||||
const uniqueSynonyms = [...new Set(synonyms)];
|
||||
expect(synonyms.length).toBe(uniqueSynonyms.length);
|
||||
});
|
||||
});
|
||||
|
||||
it('should not have overlapping synonyms between different keys', () => {
|
||||
const allSynonyms: string[] = [];
|
||||
const duplicates: string[] = [];
|
||||
|
||||
Object.values(STYLE_SYNONYMS).forEach((synonyms: string[]) => {
|
||||
synonyms.forEach((synonym: string) => {
|
||||
if (allSynonyms.includes(synonym)) {
|
||||
duplicates.push(synonym);
|
||||
}
|
||||
allSynonyms.push(synonym);
|
||||
});
|
||||
});
|
||||
|
||||
expect(duplicates).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('COMPOUND_STYLES', () => {
|
||||
it('should have compound style definitions', () => {
|
||||
// Minimum range to allow expansion
|
||||
expect(COMPOUND_STYLES.length).toBeGreaterThanOrEqual(35);
|
||||
expect(COMPOUND_STYLES.length).toBeLessThanOrEqual(150); // Reasonable upper bound
|
||||
});
|
||||
|
||||
it('should include expected compound styles', () => {
|
||||
expect(COMPOUND_STYLES).toContain('studio ghibli style');
|
||||
expect(COMPOUND_STYLES).toContain('cinematic lighting');
|
||||
expect(COMPOUND_STYLES).toContain('dramatic lighting');
|
||||
expect(COMPOUND_STYLES).toContain('depth of field');
|
||||
expect(COMPOUND_STYLES).toContain('physically based rendering');
|
||||
expect(COMPOUND_STYLES).toContain('global illumination');
|
||||
});
|
||||
|
||||
it('should have unique compound styles', () => {
|
||||
const uniqueStyles = [...new Set(COMPOUND_STYLES)];
|
||||
expect(COMPOUND_STYLES.length).toBe(uniqueStyles.length);
|
||||
});
|
||||
|
||||
it('should have lowercase compound styles', () => {
|
||||
COMPOUND_STYLES.forEach((style: string) => {
|
||||
expect(style).toBe(style.toLowerCase());
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('STYLE_ADJECTIVE_PATTERNS', () => {
|
||||
it('should have all expected pattern categories', () => {
|
||||
const expectedPatterns = [
|
||||
'quality',
|
||||
'artistic',
|
||||
'visual',
|
||||
'mood',
|
||||
'texture',
|
||||
'scale',
|
||||
'detail',
|
||||
'professional',
|
||||
];
|
||||
expect(Object.keys(STYLE_ADJECTIVE_PATTERNS)).toEqual(expectedPatterns);
|
||||
});
|
||||
|
||||
it('should match expected adjectives', () => {
|
||||
// Quality patterns
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.quality.test('sharp')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.quality.test('blurry')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.quality.test('crisp')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.quality.test('walking')).toBe(false);
|
||||
|
||||
// Artistic patterns
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.artistic.test('abstract')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.artistic.test('surreal')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.artistic.test('minimal')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.artistic.test('minimalist')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.artistic.test('running')).toBe(false);
|
||||
|
||||
// Visual patterns
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.visual.test('bright')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.visual.test('dark')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.visual.test('vibrant')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.visual.test('opened')).toBe(false);
|
||||
|
||||
// Mood patterns
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.mood.test('dramatic')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.mood.test('peaceful')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.mood.test('mysterious')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.mood.test('walking')).toBe(false);
|
||||
});
|
||||
|
||||
it('should be case insensitive', () => {
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.quality.test('Sharp')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.quality.test('SHARP')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.artistic.test('Abstract')).toBe(true);
|
||||
expect(STYLE_ADJECTIVE_PATTERNS.artistic.test('ABSTRACT')).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAllStyleKeywords', () => {
|
||||
it('should return flattened array of all keywords', () => {
|
||||
const allKeywords = getAllStyleKeywords();
|
||||
|
||||
expect(Array.isArray(allKeywords)).toBe(true);
|
||||
// Minimum range for total keywords to allow expansion
|
||||
expect(allKeywords.length).toBeGreaterThanOrEqual(350);
|
||||
expect(allKeywords.length).toBeLessThanOrEqual(750);
|
||||
|
||||
// Check that it contains keywords from different categories
|
||||
expect(allKeywords).toContain('by greg rutkowski');
|
||||
expect(allKeywords).toContain('photorealistic');
|
||||
expect(allKeywords).toContain('dramatic lighting');
|
||||
expect(allKeywords).toContain('bokeh');
|
||||
expect(allKeywords).toContain('masterpiece');
|
||||
});
|
||||
|
||||
it('should return readonly array', () => {
|
||||
const keywords = getAllStyleKeywords();
|
||||
// TypeScript will enforce readonly at compile time
|
||||
expect(Object.isFrozen(keywords) || Array.isArray(keywords)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getCompoundStyles', () => {
|
||||
it('should return compound styles array', () => {
|
||||
const compounds = getCompoundStyles();
|
||||
|
||||
expect(Array.isArray(compounds)).toBe(true);
|
||||
// Minimum range to allow expansion
|
||||
expect(compounds.length).toBeGreaterThanOrEqual(35);
|
||||
expect(compounds.length).toBeLessThanOrEqual(150);
|
||||
expect(compounds).toContain('studio ghibli style');
|
||||
expect(compounds).toContain('cinematic lighting');
|
||||
});
|
||||
|
||||
it('should return the same array as COMPOUND_STYLES', () => {
|
||||
const compounds = getCompoundStyles();
|
||||
expect(compounds).toEqual(COMPOUND_STYLES);
|
||||
});
|
||||
});
|
||||
|
||||
describe('normalizeStyleTerm', () => {
|
||||
it('should normalize known synonyms', () => {
|
||||
expect(normalizeStyleTerm('photo-realistic')).toBe('photorealistic');
|
||||
expect(normalizeStyleTerm('photo realistic')).toBe('photorealistic');
|
||||
expect(normalizeStyleTerm('lifelike')).toBe('photorealistic');
|
||||
|
||||
expect(normalizeStyleTerm('4k resolution')).toBe('4k');
|
||||
expect(normalizeStyleTerm('ultra hd')).toBe('4k');
|
||||
|
||||
expect(normalizeStyleTerm('filmic')).toBe('cinematic');
|
||||
expect(normalizeStyleTerm('movie-like')).toBe('cinematic');
|
||||
});
|
||||
|
||||
it('should return original term if not a synonym', () => {
|
||||
expect(normalizeStyleTerm('unknown-term')).toBe('unknown-term');
|
||||
expect(normalizeStyleTerm('random')).toBe('random');
|
||||
expect(normalizeStyleTerm('test')).toBe('test');
|
||||
});
|
||||
|
||||
it('should handle case insensitive matching', () => {
|
||||
expect(normalizeStyleTerm('Photo-Realistic')).toBe('photorealistic');
|
||||
expect(normalizeStyleTerm('PHOTO REALISTIC')).toBe('photorealistic');
|
||||
expect(normalizeStyleTerm('Filmic')).toBe('cinematic');
|
||||
});
|
||||
|
||||
it('should handle empty or invalid input', () => {
|
||||
expect(normalizeStyleTerm('')).toBe('');
|
||||
expect(normalizeStyleTerm(' ')).toBe(' ');
|
||||
});
|
||||
});
|
||||
|
||||
describe('isStyleAdjective', () => {
|
||||
it('should identify style adjectives', () => {
|
||||
// Quality adjectives
|
||||
expect(isStyleAdjective('sharp')).toBe(true);
|
||||
expect(isStyleAdjective('blurry')).toBe(true);
|
||||
expect(isStyleAdjective('crisp')).toBe(true);
|
||||
|
||||
// Artistic adjectives
|
||||
expect(isStyleAdjective('abstract')).toBe(true);
|
||||
expect(isStyleAdjective('surreal')).toBe(true);
|
||||
expect(isStyleAdjective('minimal')).toBe(true);
|
||||
|
||||
// Visual adjectives
|
||||
expect(isStyleAdjective('bright')).toBe(true);
|
||||
expect(isStyleAdjective('dark')).toBe(true);
|
||||
expect(isStyleAdjective('vibrant')).toBe(true);
|
||||
|
||||
// Mood adjectives
|
||||
expect(isStyleAdjective('dramatic')).toBe(true);
|
||||
expect(isStyleAdjective('peaceful')).toBe(true);
|
||||
expect(isStyleAdjective('mysterious')).toBe(true);
|
||||
});
|
||||
|
||||
it('should reject non-style adjectives', () => {
|
||||
expect(isStyleAdjective('walking')).toBe(false);
|
||||
expect(isStyleAdjective('running')).toBe(false);
|
||||
expect(isStyleAdjective('opened')).toBe(false);
|
||||
expect(isStyleAdjective('closed')).toBe(false);
|
||||
expect(isStyleAdjective('basic')).toBe(false);
|
||||
expect(isStyleAdjective('normal')).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle case insensitive matching', () => {
|
||||
expect(isStyleAdjective('Sharp')).toBe(true);
|
||||
expect(isStyleAdjective('SHARP')).toBe(true);
|
||||
expect(isStyleAdjective('Abstract')).toBe(true);
|
||||
expect(isStyleAdjective('ABSTRACT')).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('extractStyleAdjectives', () => {
|
||||
it('should extract style adjectives from word array', () => {
|
||||
const words = ['a', 'sharp', 'walking', 'robot', 'with', 'dramatic', 'lighting'];
|
||||
const adjectives = extractStyleAdjectives(words);
|
||||
|
||||
expect(adjectives).toEqual(['sharp', 'dramatic']);
|
||||
});
|
||||
|
||||
it('should handle empty array', () => {
|
||||
expect(extractStyleAdjectives([])).toEqual([]);
|
||||
});
|
||||
|
||||
it('should handle array with no style adjectives', () => {
|
||||
const words = ['walking', 'running', 'jumping', 'swimming'];
|
||||
expect(extractStyleAdjectives(words)).toEqual([]);
|
||||
});
|
||||
|
||||
it('should handle array with all style adjectives', () => {
|
||||
const words = ['sharp', 'bright', 'dramatic', 'mysterious'];
|
||||
expect(extractStyleAdjectives(words)).toEqual(words);
|
||||
});
|
||||
|
||||
it('should preserve original case', () => {
|
||||
const words = ['Sharp', 'BRIGHT', 'Dramatic'];
|
||||
const adjectives = extractStyleAdjectives(words);
|
||||
|
||||
expect(adjectives).toEqual(['Sharp', 'BRIGHT', 'Dramatic']);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Integration tests', () => {
|
||||
it('should have consistent data across all exports', () => {
|
||||
const allKeywords = getAllStyleKeywords();
|
||||
const totalInCategories = Object.values(STYLE_KEYWORDS).reduce(
|
||||
(sum: number, keywords: readonly string[]) => sum + keywords.length,
|
||||
0,
|
||||
);
|
||||
|
||||
expect(allKeywords.length).toBe(totalInCategories);
|
||||
});
|
||||
|
||||
it('should not have keywords that are also synonyms', () => {
|
||||
const allKeywords = getAllStyleKeywords();
|
||||
const allSynonyms = new Set(Object.values(STYLE_SYNONYMS).flat());
|
||||
|
||||
const overlap = allKeywords.filter((keyword: string) => allSynonyms.has(keyword));
|
||||
|
||||
// Allow reasonable range of overlaps
|
||||
expect(overlap.length).toBeGreaterThanOrEqual(10);
|
||||
expect(overlap.length).toBeLessThanOrEqual(20); // Reasonable overlap range
|
||||
});
|
||||
|
||||
it('should have compound styles that contain style keywords', () => {
|
||||
const compounds = getCompoundStyles();
|
||||
const keywords = getAllStyleKeywords();
|
||||
|
||||
// At least some compound styles should contain individual keywords
|
||||
const compoundsWithKeywords = compounds.filter((compound: string) => {
|
||||
return keywords.some((keyword: string) => compound.includes(keyword));
|
||||
});
|
||||
|
||||
expect(compoundsWithKeywords.length).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,137 @@
|
|||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import {
|
||||
SYSTEM_COMPONENTS,
|
||||
getAllComponentsWithNames,
|
||||
getOptimalComponent,
|
||||
} from '@/server/services/comfyui/config/systemComponents';
|
||||
|
||||
describe('SystemComponents', () => {
|
||||
describe('SYSTEM_COMPONENTS', () => {
|
||||
it('should be a non-empty object with valid structure', () => {
|
||||
expect(typeof SYSTEM_COMPONENTS).toBe('object');
|
||||
expect(Object.keys(SYSTEM_COMPONENTS).length).toBeGreaterThan(0);
|
||||
|
||||
// Check that all components have required fields
|
||||
Object.entries(SYSTEM_COMPONENTS).forEach(([, config]) => {
|
||||
expect(config).toBeDefined();
|
||||
expect(config.type).toBeDefined();
|
||||
expect(config.priority).toBeTypeOf('number');
|
||||
expect(config.modelFamily).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
it('should contain essential component types', () => {
|
||||
const types = Object.values(SYSTEM_COMPONENTS).map((c) => c.type);
|
||||
const uniqueTypes = [...new Set(types)];
|
||||
|
||||
expect(uniqueTypes).toContain('vae');
|
||||
expect(uniqueTypes).toContain('clip');
|
||||
expect(uniqueTypes).toContain('t5');
|
||||
});
|
||||
|
||||
it('should allow direct access to component config by name', () => {
|
||||
const config = SYSTEM_COMPONENTS['ae.safetensors'];
|
||||
expect(config).toBeDefined();
|
||||
expect(config.type).toBe('vae');
|
||||
expect(config.modelFamily).toBe('FLUX');
|
||||
expect(config.priority).toBe(1);
|
||||
});
|
||||
|
||||
it('should return undefined for invalid component name', () => {
|
||||
const config = SYSTEM_COMPONENTS['nonexistent.safetensors'];
|
||||
expect(config).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAllComponentsWithNames', () => {
|
||||
it('should return components with names for valid type', () => {
|
||||
const result = getAllComponentsWithNames({ type: 'vae' });
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
result.forEach(({ name, config }) => {
|
||||
expect(name).toBeTypeOf('string');
|
||||
expect(config.type).toBe('vae');
|
||||
});
|
||||
});
|
||||
|
||||
it('should filter by modelFamily when specified', () => {
|
||||
const result = getAllComponentsWithNames({ modelFamily: 'FLUX', type: 'vae' });
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
result.forEach(({ config }) => {
|
||||
expect(config.modelFamily).toBe('FLUX');
|
||||
expect(config.type).toBe('vae');
|
||||
});
|
||||
});
|
||||
|
||||
it('should filter by priority when specified', () => {
|
||||
const result = getAllComponentsWithNames({ priority: 1 });
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
result.forEach(({ config }) => {
|
||||
expect(config.priority).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
it('should filter by multiple criteria', () => {
|
||||
const result = getAllComponentsWithNames({
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
});
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
result.forEach(({ config }) => {
|
||||
expect(config.type).toBe('lora');
|
||||
expect(config.modelFamily).toBe('FLUX');
|
||||
expect(config.priority).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
it('should filter by compatible variant', () => {
|
||||
const result = getAllComponentsWithNames({
|
||||
compatibleVariant: 'dev',
|
||||
type: 'lora',
|
||||
});
|
||||
expect(result.length).toBeGreaterThan(0);
|
||||
result.forEach(({ config }) => {
|
||||
expect(config.type).toBe('lora');
|
||||
expect(config.compatibleVariants).toContain('dev');
|
||||
});
|
||||
});
|
||||
|
||||
it('should return empty array for invalid filters', () => {
|
||||
const result = getAllComponentsWithNames({
|
||||
modelFamily: 'NONEXISTENT' as any,
|
||||
type: 'vae',
|
||||
});
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getOptimalComponent', () => {
|
||||
it('should return component with highest priority (lowest number) for FLUX VAE', () => {
|
||||
const component = getOptimalComponent('vae', 'FLUX');
|
||||
expect(component).toBeDefined();
|
||||
expect(typeof component).toBe('string');
|
||||
|
||||
// Should return ae.safetensors which has priority 1
|
||||
expect(component).toBe('ae.safetensors');
|
||||
});
|
||||
|
||||
it('should return component with highest priority for SD1 VAE', () => {
|
||||
const component = getOptimalComponent('vae', 'SD1');
|
||||
expect(component).toBeDefined();
|
||||
expect(typeof component).toBe('string');
|
||||
});
|
||||
|
||||
it('should return component with highest priority for FLUX clip', () => {
|
||||
const component = getOptimalComponent('clip', 'FLUX');
|
||||
expect(component).toBeDefined();
|
||||
expect(typeof component).toBe('string');
|
||||
});
|
||||
|
||||
it('should throw ConfigError when no components found', () => {
|
||||
expect(() => {
|
||||
getOptimalComponent('vae', 'NONEXISTENT' as any);
|
||||
}).toThrow('No vae components configured for model family NONEXISTENT');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,146 @@
|
|||
import { beforeEach, describe, expect, it } from 'vitest';
|
||||
|
||||
import { ComfyUIAuthService } from '@/server/services/comfyui/core/comfyUIAuthService';
|
||||
import { ServicesError } from '@/server/services/comfyui/errors';
|
||||
import type { ComfyUIKeyVault } from '@/types/user/settings/keyVaults';
|
||||
|
||||
describe('ComfyUIAuthService', () => {
|
||||
describe('Constructor and initialization', () => {
|
||||
it('should initialize with none auth type by default', () => {
|
||||
const service = new ComfyUIAuthService({});
|
||||
|
||||
expect(service.getCredentials()).toBeUndefined();
|
||||
expect(service.getAuthHeaders()).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should initialize with basic auth', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'basic',
|
||||
username: 'testuser',
|
||||
password: 'testpass',
|
||||
};
|
||||
|
||||
const service = new ComfyUIAuthService(options);
|
||||
|
||||
const credentials = service.getCredentials();
|
||||
expect(credentials).toEqual({
|
||||
type: 'basic',
|
||||
username: 'testuser',
|
||||
password: 'testpass',
|
||||
});
|
||||
|
||||
const headers = service.getAuthHeaders();
|
||||
expect(headers).toEqual({
|
||||
Authorization: `Basic ${btoa('testuser:testpass')}`,
|
||||
});
|
||||
});
|
||||
|
||||
it('should initialize with bearer auth', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'bearer',
|
||||
apiKey: 'test-api-key',
|
||||
};
|
||||
|
||||
const service = new ComfyUIAuthService(options);
|
||||
|
||||
const credentials = service.getCredentials();
|
||||
expect(credentials).toEqual({
|
||||
type: 'bearer_token',
|
||||
token: 'test-api-key',
|
||||
});
|
||||
|
||||
const headers = service.getAuthHeaders();
|
||||
expect(headers).toEqual({
|
||||
Authorization: 'Bearer test-api-key',
|
||||
});
|
||||
});
|
||||
|
||||
it('should initialize with custom auth', () => {
|
||||
const customHeaders = { 'X-API-Key': 'custom-key', 'X-Client': 'test' };
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'custom',
|
||||
customHeaders,
|
||||
};
|
||||
|
||||
const service = new ComfyUIAuthService(options);
|
||||
|
||||
const credentials = service.getCredentials();
|
||||
expect(credentials).toEqual({
|
||||
type: 'custom',
|
||||
headers: customHeaders,
|
||||
});
|
||||
|
||||
const headers = service.getAuthHeaders();
|
||||
expect(headers).toEqual(customHeaders);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Validation', () => {
|
||||
it('should throw error for basic auth without username', () => {
|
||||
expect(() => {
|
||||
new ComfyUIAuthService({
|
||||
authType: 'basic',
|
||||
password: 'testpass',
|
||||
});
|
||||
}).toThrow(ServicesError);
|
||||
});
|
||||
|
||||
it('should throw error for basic auth without password', () => {
|
||||
expect(() => {
|
||||
new ComfyUIAuthService({
|
||||
authType: 'basic',
|
||||
username: 'testuser',
|
||||
});
|
||||
}).toThrow(ServicesError);
|
||||
});
|
||||
|
||||
it('should throw error for bearer auth without apiKey', () => {
|
||||
expect(() => {
|
||||
new ComfyUIAuthService({
|
||||
authType: 'bearer',
|
||||
});
|
||||
}).toThrow(ServicesError);
|
||||
});
|
||||
|
||||
it('should throw error for custom auth without headers', () => {
|
||||
expect(() => {
|
||||
new ComfyUIAuthService({
|
||||
authType: 'custom',
|
||||
});
|
||||
}).toThrow(ServicesError);
|
||||
});
|
||||
|
||||
it('should throw error for custom auth with empty headers', () => {
|
||||
expect(() => {
|
||||
new ComfyUIAuthService({
|
||||
authType: 'custom',
|
||||
customHeaders: {},
|
||||
});
|
||||
}).toThrow(ServicesError);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Edge cases', () => {
|
||||
it('should handle partial basic auth gracefully in headers', () => {
|
||||
// This tests the createAuthHeaders method behavior
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'basic',
|
||||
username: 'testuser',
|
||||
password: 'testpass',
|
||||
};
|
||||
|
||||
const service = new ComfyUIAuthService(options);
|
||||
expect(service.getAuthHeaders()).toBeDefined();
|
||||
});
|
||||
|
||||
it('should handle partial bearer auth gracefully in headers', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'bearer',
|
||||
apiKey: 'test-key',
|
||||
};
|
||||
|
||||
const service = new ComfyUIAuthService(options);
|
||||
expect(service.getAuthHeaders()).toBeDefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,287 @@
|
|||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { ComfyUIConnectionService } from '@/server/services/comfyui/core/comfyUIConnectionService';
|
||||
import { ServicesError } from '@/server/services/comfyui/errors';
|
||||
|
||||
// Mock global fetch
|
||||
global.fetch = vi.fn();
|
||||
|
||||
describe('ComfyUIConnectionService', () => {
|
||||
let service: ComfyUIConnectionService;
|
||||
const mockFetch = vi.mocked(fetch);
|
||||
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIConnectionService();
|
||||
vi.clearAllMocks();
|
||||
vi.useFakeTimers();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
describe('Constructor and initialization', () => {
|
||||
it('should initialize with default state', () => {
|
||||
const newService = new ComfyUIConnectionService();
|
||||
|
||||
expect(newService.isValidated()).toBe(false);
|
||||
|
||||
const status = newService.getStatus();
|
||||
expect(status.isValidated).toBe(false);
|
||||
expect(status.lastValidationTime).toBe(null);
|
||||
expect(status.timeUntilExpiry).toBe(null);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Connection validation', () => {
|
||||
const baseURL = 'http://localhost:8188';
|
||||
const authHeaders = { Authorization: 'Bearer test-token' };
|
||||
|
||||
it('should validate connection successfully', async () => {
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ system: 'data' }),
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
const result = await service.validateConnection(baseURL, authHeaders);
|
||||
|
||||
expect(result).toBe(true);
|
||||
expect(service.isValidated()).toBe(true);
|
||||
expect(mockFetch).toHaveBeenCalledWith(`${baseURL}/system_stats`, {
|
||||
headers: {
|
||||
...authHeaders,
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
method: 'GET',
|
||||
mode: 'cors',
|
||||
});
|
||||
});
|
||||
|
||||
it('should validate connection without auth headers', async () => {
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ system: 'data' }),
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
const result = await service.validateConnection(baseURL);
|
||||
|
||||
expect(result).toBe(true);
|
||||
expect(mockFetch).toHaveBeenCalledWith(`${baseURL}/system_stats`, {
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
method: 'GET',
|
||||
mode: 'cors',
|
||||
});
|
||||
});
|
||||
|
||||
it('should return cached validation result within TTL', async () => {
|
||||
// First validation
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ system: 'data' }),
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
await service.validateConnection(baseURL, authHeaders);
|
||||
expect(mockFetch).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Second validation within TTL should use cached result
|
||||
const result = await service.validateConnection(baseURL, authHeaders);
|
||||
expect(result).toBe(true);
|
||||
expect(mockFetch).toHaveBeenCalledTimes(1); // No additional fetch call
|
||||
});
|
||||
|
||||
it('should re-validate after TTL expiry', async () => {
|
||||
// First validation
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ system: 'data' }),
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
await service.validateConnection(baseURL, authHeaders);
|
||||
expect(mockFetch).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Fast-forward time beyond TTL (5 minutes + 1 second)
|
||||
vi.advanceTimersByTime(5 * 60 * 1000 + 1000);
|
||||
|
||||
// Second validation after TTL should make new request
|
||||
const result = await service.validateConnection(baseURL, authHeaders);
|
||||
expect(result).toBe(true);
|
||||
expect(mockFetch).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should handle HTTP error responses', async () => {
|
||||
const mockResponse = {
|
||||
ok: false,
|
||||
status: 404,
|
||||
statusText: 'Not Found',
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
await expect(service.validateConnection(baseURL, authHeaders)).rejects.toThrow(ServicesError);
|
||||
expect(service.isValidated()).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle invalid JSON response', async () => {
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue(null),
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
await expect(service.validateConnection(baseURL, authHeaders)).rejects.toThrow(ServicesError);
|
||||
expect(service.isValidated()).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle network errors', async () => {
|
||||
mockFetch.mockRejectedValue(new Error('Network error'));
|
||||
|
||||
await expect(service.validateConnection(baseURL, authHeaders)).rejects.toThrow(Error);
|
||||
expect(service.isValidated()).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle JSON parse errors', async () => {
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: vi.fn().mockRejectedValue(new Error('JSON parse error')),
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
await expect(service.validateConnection(baseURL, authHeaders)).rejects.toThrow(Error);
|
||||
expect(service.isValidated()).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Connection state management', () => {
|
||||
it('should mark connection as validated', () => {
|
||||
service.markAsValidated();
|
||||
|
||||
expect(service.isValidated()).toBe(true);
|
||||
|
||||
const status = service.getStatus();
|
||||
expect(status.isValidated).toBe(true);
|
||||
expect(status.lastValidationTime).toBeGreaterThan(0);
|
||||
expect(status.timeUntilExpiry).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('should invalidate connection', () => {
|
||||
service.markAsValidated();
|
||||
expect(service.isValidated()).toBe(true);
|
||||
|
||||
service.invalidate();
|
||||
expect(service.isValidated()).toBe(false);
|
||||
|
||||
const status = service.getStatus();
|
||||
expect(status.isValidated).toBe(false);
|
||||
expect(status.lastValidationTime).toBe(null);
|
||||
expect(status.timeUntilExpiry).toBe(null);
|
||||
});
|
||||
|
||||
it('should expire validation after TTL', () => {
|
||||
service.markAsValidated();
|
||||
expect(service.isValidated()).toBe(true);
|
||||
|
||||
// Fast-forward time beyond TTL
|
||||
vi.advanceTimersByTime(5 * 60 * 1000 + 1000);
|
||||
|
||||
expect(service.isValidated()).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Connection status', () => {
|
||||
it('should return correct status for unvalidated connection', () => {
|
||||
const status = service.getStatus();
|
||||
|
||||
expect(status.isValidated).toBe(false);
|
||||
expect(status.lastValidationTime).toBe(null);
|
||||
expect(status.timeUntilExpiry).toBe(null);
|
||||
});
|
||||
|
||||
it('should return correct status for validated connection', () => {
|
||||
service.markAsValidated();
|
||||
|
||||
const status = service.getStatus();
|
||||
|
||||
expect(status.isValidated).toBe(true);
|
||||
expect(status.lastValidationTime).toBeGreaterThan(0);
|
||||
expect(status.timeUntilExpiry).toBeGreaterThan(0);
|
||||
expect(status.timeUntilExpiry).toBeLessThanOrEqual(5 * 60 * 1000); // Should be <= 5 minutes
|
||||
});
|
||||
|
||||
it('should calculate time until expiry correctly', () => {
|
||||
service.markAsValidated();
|
||||
|
||||
// Advance time by 2 minutes
|
||||
vi.advanceTimersByTime(2 * 60 * 1000);
|
||||
|
||||
const status = service.getStatus();
|
||||
expect(status.timeUntilExpiry).toBeCloseTo(3 * 60 * 1000, -2); // ~3 minutes remaining
|
||||
});
|
||||
|
||||
it('should return zero time until expiry when expired', () => {
|
||||
service.markAsValidated();
|
||||
|
||||
// Advance time beyond TTL
|
||||
vi.advanceTimersByTime(6 * 60 * 1000);
|
||||
|
||||
const status = service.getStatus();
|
||||
expect(status.timeUntilExpiry).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Edge cases', () => {
|
||||
const baseURL = 'http://localhost:8188';
|
||||
const authHeaders = { Authorization: 'Bearer test-token' };
|
||||
|
||||
it('should handle multiple rapid validation calls', async () => {
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ system: 'data' }),
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
// Make multiple concurrent validation calls
|
||||
const promises = [
|
||||
service.validateConnection(baseURL, authHeaders),
|
||||
service.validateConnection(baseURL, authHeaders),
|
||||
service.validateConnection(baseURL, authHeaders),
|
||||
];
|
||||
|
||||
const results = await Promise.all(promises);
|
||||
|
||||
// All should succeed
|
||||
expect(results.every((r) => r === true)).toBe(true);
|
||||
|
||||
// For concurrent calls, each call checks the cache independently
|
||||
// Since they start before any completes, they will all make HTTP requests
|
||||
expect(mockFetch).toHaveBeenCalledTimes(3);
|
||||
|
||||
// After all complete, subsequent calls should use cache
|
||||
await service.validateConnection(baseURL, authHeaders);
|
||||
expect(mockFetch).toHaveBeenCalledTimes(3); // No additional call
|
||||
});
|
||||
|
||||
it('should handle validation with empty auth headers object', async () => {
|
||||
const mockResponse = {
|
||||
ok: true,
|
||||
json: vi.fn().mockResolvedValue({ system: 'data' }),
|
||||
};
|
||||
mockFetch.mockResolvedValue(mockResponse as any);
|
||||
|
||||
const result = await service.validateConnection(baseURL, {});
|
||||
|
||||
expect(result).toBe(true);
|
||||
expect(mockFetch).toHaveBeenCalledWith(`${baseURL}/system_stats`, {
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
method: 'GET',
|
||||
mode: 'cors',
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
666
src/server/services/comfyui/__tests__/core/comfyuiClient.test.ts
Normal file
666
src/server/services/comfyui/__tests__/core/comfyuiClient.test.ts
Normal file
|
|
@ -0,0 +1,666 @@
|
|||
import { ComfyApi } from '@saintno/comfyui-sdk';
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { ComfyUIAuthService } from '@/server/services/comfyui/core/comfyUIAuthService';
|
||||
import { ComfyUIClientService } from '@/server/services/comfyui/core/comfyUIClientService';
|
||||
import { ComfyUIConnectionService } from '@/server/services/comfyui/core/comfyUIConnectionService';
|
||||
import { ServicesError } from '@/server/services/comfyui/errors';
|
||||
import { ModelResolverError } from '@/server/services/comfyui/errors/modelResolverError';
|
||||
import { ComfyUIKeyVault } from '@/types/user/settings/keyVaults';
|
||||
|
||||
// Mock the SDK
|
||||
vi.mock('@saintno/comfyui-sdk', () => ({
|
||||
CallWrapper: vi.fn(),
|
||||
ComfyApi: vi.fn(),
|
||||
}));
|
||||
|
||||
// Mock the modular services
|
||||
vi.mock('@/server/services/comfyui/core/comfyUIAuthService');
|
||||
vi.mock('@/server/services/comfyui/core/comfyUIConnectionService');
|
||||
|
||||
describe('ComfyUIClientService', () => {
|
||||
let service: ComfyUIClientService;
|
||||
let mockClient: any;
|
||||
let mockAuthService: any;
|
||||
let mockConnectionService: any;
|
||||
let originalDateNow: () => number;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
originalDateNow = Date.now;
|
||||
|
||||
// Create mock client
|
||||
mockClient = {
|
||||
fetchApi: vi.fn(),
|
||||
getCheckpoints: vi.fn(),
|
||||
getLoras: vi.fn(),
|
||||
getNodeDefs: vi.fn(),
|
||||
getPathImage: vi.fn(),
|
||||
getSamplerInfo: vi.fn(),
|
||||
init: vi.fn(),
|
||||
uploadImage: vi.fn(),
|
||||
};
|
||||
|
||||
// Create mock services
|
||||
mockAuthService = {
|
||||
getCredentials: vi.fn().mockReturnValue(undefined),
|
||||
getAuthHeaders: vi.fn().mockReturnValue(undefined),
|
||||
};
|
||||
|
||||
mockConnectionService = {
|
||||
validateConnection: vi.fn().mockResolvedValue(true),
|
||||
getStatus: vi.fn().mockReturnValue({
|
||||
isValidated: false,
|
||||
lastValidationTime: null,
|
||||
timeUntilExpiry: null,
|
||||
}),
|
||||
};
|
||||
|
||||
// Mock constructors
|
||||
vi.mocked(ComfyApi).mockImplementation(() => mockClient);
|
||||
vi.mocked(ComfyUIAuthService).mockImplementation(() => mockAuthService);
|
||||
vi.mocked(ComfyUIConnectionService).mockImplementation(() => mockConnectionService);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
Date.now = originalDateNow;
|
||||
});
|
||||
|
||||
describe('constructor', () => {
|
||||
it('should initialize with default settings', () => {
|
||||
service = new ComfyUIClientService();
|
||||
|
||||
expect(ComfyUIAuthService).toHaveBeenCalledWith({});
|
||||
expect(ComfyUIConnectionService).toHaveBeenCalled();
|
||||
expect(ComfyApi).toHaveBeenCalledWith(expect.stringContaining('http'), undefined, {
|
||||
credentials: undefined,
|
||||
});
|
||||
expect(mockClient.init).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should initialize with custom options', () => {
|
||||
const options: ComfyUIKeyVault = {
|
||||
authType: 'basic',
|
||||
baseURL: 'http://custom:8188',
|
||||
password: 'pass',
|
||||
username: 'user',
|
||||
};
|
||||
|
||||
service = new ComfyUIClientService(options);
|
||||
|
||||
expect(ComfyUIAuthService).toHaveBeenCalledWith(options);
|
||||
expect(ComfyUIConnectionService).toHaveBeenCalled();
|
||||
expect(ComfyApi).toHaveBeenCalledWith('http://custom:8188', undefined, {
|
||||
credentials: undefined,
|
||||
});
|
||||
expect(mockClient.init).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle auth service errors during initialization', () => {
|
||||
// Mock AuthService constructor to throw
|
||||
vi.mocked(ComfyUIAuthService).mockImplementation(() => {
|
||||
throw new ServicesError('Invalid auth config', ServicesError.Reasons.INVALID_ARGS);
|
||||
});
|
||||
|
||||
expect(() => new ComfyUIClientService({ authType: 'basic' })).toThrow();
|
||||
|
||||
// Verify it throws an error (ErrorHandlerService wraps it into TRPCError)
|
||||
try {
|
||||
new ComfyUIClientService({ authType: 'basic' });
|
||||
expect.fail('Should have thrown an error');
|
||||
} catch (error: any) {
|
||||
expect(error.cause).toHaveProperty('errorType', 'InvalidComfyUIArgs');
|
||||
expect(error.cause).toHaveProperty('provider', 'comfyui');
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('uploadImage', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
});
|
||||
|
||||
it('should successfully upload an image', async () => {
|
||||
// Setup mock
|
||||
const mockBuffer = Buffer.from('test image data');
|
||||
const mockFileName = 'test.png';
|
||||
const mockResult = {
|
||||
info: {
|
||||
filename: 'uploaded_test.png',
|
||||
},
|
||||
};
|
||||
|
||||
mockClient.uploadImage.mockResolvedValue(mockResult);
|
||||
|
||||
// Execute
|
||||
const result = await service.uploadImage(mockBuffer, mockFileName);
|
||||
|
||||
// Verify
|
||||
expect(result).toBe('uploaded_test.png');
|
||||
expect(mockClient.uploadImage).toHaveBeenCalledWith(mockBuffer, mockFileName);
|
||||
});
|
||||
|
||||
it('should handle upload failure when result is null', async () => {
|
||||
// Setup
|
||||
mockClient.uploadImage.mockResolvedValue(null);
|
||||
|
||||
// Execute and verify
|
||||
await expect(service.uploadImage(Buffer.from('data'), 'file.png')).rejects.toThrow(
|
||||
'Failed to upload image to ComfyUI server',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle network errors during upload', async () => {
|
||||
// Setup
|
||||
const networkError = new TypeError('Failed to fetch');
|
||||
mockClient.uploadImage.mockRejectedValue(networkError);
|
||||
|
||||
// Execute and verify - uploadImage just re-throws without transformation
|
||||
await expect(service.uploadImage(Buffer.from('data'), 'file.png')).rejects.toThrow(
|
||||
'Failed to fetch',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle 401 authentication error', async () => {
|
||||
// Setup
|
||||
const authError = new Error('Request failed with status: 401');
|
||||
mockClient.uploadImage.mockRejectedValue(authError);
|
||||
|
||||
// Execute and verify - uploadImage just re-throws without transformation
|
||||
await expect(service.uploadImage(Buffer.from('data'), 'file.png')).rejects.toThrow(
|
||||
'Request failed with status: 401',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle 403 forbidden error', async () => {
|
||||
// Setup
|
||||
const forbiddenError = new Error('Request failed with status: 403');
|
||||
mockClient.uploadImage.mockRejectedValue(forbiddenError);
|
||||
|
||||
// Execute and verify - uploadImage just re-throws without transformation
|
||||
await expect(service.uploadImage(Buffer.from('data'), 'file.png')).rejects.toThrow(
|
||||
'Request failed with status: 403',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle 500+ server errors', async () => {
|
||||
// Setup
|
||||
const serverError = new Error('Request failed with status: 503');
|
||||
mockClient.uploadImage.mockRejectedValue(serverError);
|
||||
|
||||
// Execute and verify - uploadImage just re-throws without transformation
|
||||
await expect(service.uploadImage(Buffer.from('data'), 'file.png')).rejects.toThrow(
|
||||
'Request failed with status: 503',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle unknown errors', async () => {
|
||||
// Setup
|
||||
const unknownError = 'Some unexpected error string';
|
||||
mockClient.uploadImage.mockRejectedValue(unknownError);
|
||||
|
||||
// Execute and verify - uploadImage just re-throws without transformation
|
||||
await expect(service.uploadImage(Buffer.from('data'), 'file.png')).rejects.toBe(
|
||||
'Some unexpected error string',
|
||||
);
|
||||
});
|
||||
|
||||
it('should support Blob upload', async () => {
|
||||
// Setup
|
||||
const mockBlob = new Blob(['test data']);
|
||||
const mockResult = {
|
||||
info: { filename: 'blob_upload.png' },
|
||||
};
|
||||
|
||||
mockClient.uploadImage.mockResolvedValue(mockResult);
|
||||
|
||||
// Execute
|
||||
const result = await service.uploadImage(mockBlob, 'blob.png');
|
||||
|
||||
// Verify
|
||||
expect(result).toBe('blob_upload.png');
|
||||
expect(mockClient.uploadImage).toHaveBeenCalledWith(mockBlob, 'blob.png');
|
||||
});
|
||||
});
|
||||
|
||||
describe('executeWorkflow', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
});
|
||||
|
||||
it('should execute workflow successfully', async () => {
|
||||
// Import CallWrapper mock
|
||||
const { CallWrapper } = await import('@saintno/comfyui-sdk');
|
||||
|
||||
// Setup mock workflow
|
||||
const mockWorkflow = { id: 'test-workflow' };
|
||||
const mockResult = {
|
||||
images: {
|
||||
images: [{ data: 'base64' }],
|
||||
},
|
||||
};
|
||||
|
||||
// Create CallWrapper mock instance
|
||||
const mockCallWrapper = {
|
||||
onFailed: vi.fn().mockReturnThis(),
|
||||
onFinished: vi.fn().mockReturnThis(),
|
||||
onProgress: vi.fn().mockReturnThis(),
|
||||
run: vi.fn(),
|
||||
};
|
||||
|
||||
// Setup CallWrapper mock
|
||||
vi.mocked(CallWrapper).mockImplementation(() => mockCallWrapper as any);
|
||||
|
||||
// Simulate successful execution
|
||||
mockCallWrapper.run.mockImplementation(() => {
|
||||
const finishCallback = mockCallWrapper.onFinished.mock.calls[0][0];
|
||||
finishCallback(mockResult);
|
||||
});
|
||||
|
||||
// Execute
|
||||
const result = await service.executeWorkflow(mockWorkflow as any);
|
||||
|
||||
// Verify
|
||||
expect(result).toEqual(mockResult);
|
||||
expect(CallWrapper).toHaveBeenCalledWith(mockClient, mockWorkflow);
|
||||
});
|
||||
|
||||
it('should handle workflow execution failure', async () => {
|
||||
const { CallWrapper } = await import('@saintno/comfyui-sdk');
|
||||
|
||||
const mockWorkflow = { id: 'test' };
|
||||
const mockError = new Error('Workflow failed');
|
||||
|
||||
const mockCallWrapper = {
|
||||
onFailed: vi.fn().mockReturnThis(),
|
||||
onFinished: vi.fn().mockReturnThis(),
|
||||
onProgress: vi.fn().mockReturnThis(),
|
||||
run: vi.fn(),
|
||||
};
|
||||
|
||||
vi.mocked(CallWrapper).mockImplementation(() => mockCallWrapper as any);
|
||||
|
||||
// Simulate failure
|
||||
mockCallWrapper.run.mockImplementation(() => {
|
||||
const failCallback = mockCallWrapper.onFailed.mock.calls[0][0];
|
||||
failCallback(mockError);
|
||||
});
|
||||
|
||||
// Execute and verify - executeWorkflow just passes through the error
|
||||
await expect(service.executeWorkflow(mockWorkflow as any)).rejects.toThrow('Workflow failed');
|
||||
});
|
||||
|
||||
it('should call progress callback', async () => {
|
||||
const { CallWrapper } = await import('@saintno/comfyui-sdk');
|
||||
|
||||
const mockWorkflow = { id: 'test' };
|
||||
const mockProgress = { step: 1, total: 10 };
|
||||
const progressCallback = vi.fn();
|
||||
|
||||
const mockCallWrapper = {
|
||||
onFailed: vi.fn().mockReturnThis(),
|
||||
onFinished: vi.fn().mockReturnThis(),
|
||||
onProgress: vi.fn().mockReturnThis(),
|
||||
run: vi.fn(),
|
||||
};
|
||||
|
||||
vi.mocked(CallWrapper).mockImplementation(() => mockCallWrapper as any);
|
||||
|
||||
// Simulate progress and completion
|
||||
mockCallWrapper.run.mockImplementation(() => {
|
||||
const progressCb = mockCallWrapper.onProgress.mock.calls[0][0];
|
||||
progressCb(mockProgress);
|
||||
|
||||
const finishCb = mockCallWrapper.onFinished.mock.calls[0][0];
|
||||
finishCb({ images: { images: [] } });
|
||||
});
|
||||
|
||||
// Execute
|
||||
await service.executeWorkflow(mockWorkflow as any, progressCallback);
|
||||
|
||||
// Verify
|
||||
expect(progressCallback).toHaveBeenCalledWith(mockProgress);
|
||||
});
|
||||
});
|
||||
|
||||
describe('validateConnection', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
});
|
||||
|
||||
it('should delegate to connection service', async () => {
|
||||
mockConnectionService.validateConnection.mockResolvedValue(true);
|
||||
|
||||
const result = await service.validateConnection();
|
||||
|
||||
expect(result).toBe(true);
|
||||
expect(mockConnectionService.validateConnection).toHaveBeenCalledWith(
|
||||
expect.stringContaining('http'), // baseURL
|
||||
undefined, // auth headers (undefined for no auth)
|
||||
);
|
||||
});
|
||||
|
||||
it('should pass auth headers to connection service', async () => {
|
||||
const authHeaders = { Authorization: 'Bearer test-token' };
|
||||
mockAuthService.getAuthHeaders.mockReturnValue(authHeaders);
|
||||
mockConnectionService.validateConnection.mockResolvedValue(true);
|
||||
|
||||
const result = await service.validateConnection();
|
||||
|
||||
expect(result).toBe(true);
|
||||
expect(mockConnectionService.validateConnection).toHaveBeenCalledWith(
|
||||
expect.stringContaining('http'), // baseURL
|
||||
authHeaders,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle connection service errors', async () => {
|
||||
const connectionError = new ServicesError(
|
||||
'Connection failed',
|
||||
ServicesError.Reasons.CONNECTION_FAILED,
|
||||
);
|
||||
mockConnectionService.validateConnection.mockRejectedValue(connectionError);
|
||||
|
||||
await expect(service.validateConnection()).rejects.toThrow(connectionError);
|
||||
});
|
||||
});
|
||||
|
||||
// fetchApi and getObjectInfo tests removed
|
||||
// These methods should not be used directly
|
||||
// Use SDK methods: getCheckpoints(), getNodeDefs(), getLoras(), getSamplerInfo()
|
||||
|
||||
describe('getPathImage', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
});
|
||||
|
||||
it('should delegate to client getPathImage', () => {
|
||||
// Setup
|
||||
const mockImageInfo = { filename: 'test.png' };
|
||||
const expectedPath = 'https://server/image/test.png';
|
||||
mockClient.getPathImage.mockReturnValue(expectedPath);
|
||||
|
||||
// Execute
|
||||
const result = service.getPathImage(mockImageInfo);
|
||||
|
||||
// Verify
|
||||
expect(result).toBe(expectedPath);
|
||||
expect(mockClient.getPathImage).toHaveBeenCalledWith(mockImageInfo);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAuthHeaders', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
});
|
||||
|
||||
it('should delegate to auth service', () => {
|
||||
const expectedHeaders = { Authorization: 'Bearer test-token' };
|
||||
mockAuthService.getAuthHeaders.mockReturnValue(expectedHeaders);
|
||||
|
||||
const headers = service.getAuthHeaders();
|
||||
|
||||
expect(headers).toEqual(expectedHeaders);
|
||||
expect(mockAuthService.getAuthHeaders).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return undefined when auth service returns undefined', () => {
|
||||
mockAuthService.getAuthHeaders.mockReturnValue(undefined);
|
||||
|
||||
const headers = service.getAuthHeaders();
|
||||
|
||||
expect(headers).toBeUndefined();
|
||||
expect(mockAuthService.getAuthHeaders).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('service access methods', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
});
|
||||
|
||||
it('should provide access to auth service', () => {
|
||||
const authService = service.getAuthService();
|
||||
expect(authService).toBe(mockAuthService);
|
||||
});
|
||||
|
||||
it('should provide access to connection service', () => {
|
||||
const connectionService = service.getConnectionService();
|
||||
expect(connectionService).toBe(mockConnectionService);
|
||||
});
|
||||
|
||||
it('should provide connection status', () => {
|
||||
const expectedStatus = {
|
||||
isValidated: true,
|
||||
lastValidationTime: Date.now(),
|
||||
timeUntilExpiry: 300000,
|
||||
};
|
||||
mockConnectionService.getStatus.mockReturnValue(expectedStatus);
|
||||
|
||||
const status = service.getConnectionStatus();
|
||||
|
||||
expect(status).toEqual(expectedStatus);
|
||||
expect(mockConnectionService.getStatus).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('getCheckpoints', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
mockClient.getCheckpoints = vi.fn();
|
||||
});
|
||||
|
||||
it('should get checkpoints successfully', async () => {
|
||||
const mockCheckpoints = ['flux1-dev.safetensors', 'sd3.5_large.safetensors'];
|
||||
mockClient.getCheckpoints.mockResolvedValue(mockCheckpoints);
|
||||
|
||||
const result = await service.getCheckpoints();
|
||||
|
||||
expect(result).toEqual(mockCheckpoints);
|
||||
expect(mockClient.getCheckpoints).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle error when getting checkpoints', async () => {
|
||||
mockClient.getCheckpoints.mockRejectedValue(new Error('Failed to fetch'));
|
||||
|
||||
await expect(service.getCheckpoints()).rejects.toThrow();
|
||||
});
|
||||
|
||||
it('should cache checkpoints for 1 minute TTL', async () => {
|
||||
const mockCheckpoints = ['flux1-dev.safetensors', 'sd3.5_large.safetensors'];
|
||||
mockClient.getCheckpoints.mockResolvedValue(mockCheckpoints);
|
||||
|
||||
// First call
|
||||
const result1 = await service.getCheckpoints();
|
||||
expect(result1).toEqual(mockCheckpoints);
|
||||
expect(mockClient.getCheckpoints).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Second call within TTL - should use cache
|
||||
const result2 = await service.getCheckpoints();
|
||||
expect(result2).toEqual(mockCheckpoints);
|
||||
expect(mockClient.getCheckpoints).toHaveBeenCalledTimes(1); // Still only 1 call
|
||||
|
||||
// Mock time passing (simulate cache expiry)
|
||||
vi.spyOn(Date, 'now').mockReturnValue(Date.now() + 61 * 1000); // 61 seconds later
|
||||
|
||||
// Third call after TTL expired - should make new SDK call
|
||||
const result3 = await service.getCheckpoints();
|
||||
expect(result3).toEqual(mockCheckpoints);
|
||||
expect(mockClient.getCheckpoints).toHaveBeenCalledTimes(2); // Now 2 calls
|
||||
});
|
||||
});
|
||||
|
||||
describe('getLoras', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
mockClient.getLoras = vi.fn();
|
||||
});
|
||||
|
||||
it('should get LoRAs successfully', async () => {
|
||||
const mockLoras = ['lora1.safetensors', 'lora2.safetensors'];
|
||||
mockClient.getLoras.mockResolvedValue(mockLoras);
|
||||
|
||||
const result = await service.getLoras();
|
||||
|
||||
expect(result).toEqual(mockLoras);
|
||||
expect(mockClient.getLoras).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle error when getting LoRAs', async () => {
|
||||
mockClient.getLoras.mockRejectedValue(new Error('Failed to fetch'));
|
||||
|
||||
await expect(service.getLoras()).rejects.toThrow();
|
||||
});
|
||||
|
||||
it('should cache LoRAs for 1 minute TTL', async () => {
|
||||
const mockLoras = ['lora1.safetensors', 'lora2.safetensors'];
|
||||
mockClient.getLoras.mockResolvedValue(mockLoras);
|
||||
|
||||
// First call
|
||||
const result1 = await service.getLoras();
|
||||
expect(result1).toEqual(mockLoras);
|
||||
expect(mockClient.getLoras).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Second call within TTL - should use cache
|
||||
const result2 = await service.getLoras();
|
||||
expect(result2).toEqual(mockLoras);
|
||||
expect(mockClient.getLoras).toHaveBeenCalledTimes(1); // Still only 1 call
|
||||
|
||||
// Mock time passing (simulate cache expiry)
|
||||
vi.spyOn(Date, 'now').mockReturnValue(Date.now() + 61 * 1000); // 61 seconds later
|
||||
|
||||
// Third call after TTL expired - should make new SDK call
|
||||
const result3 = await service.getLoras();
|
||||
expect(result3).toEqual(mockLoras);
|
||||
expect(mockClient.getLoras).toHaveBeenCalledTimes(2); // Now 2 calls
|
||||
});
|
||||
});
|
||||
|
||||
describe('getNodeDefs', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
mockClient.getNodeDefs = vi.fn();
|
||||
});
|
||||
|
||||
it('should get node definitions with caching', async () => {
|
||||
const mockNodeDefs = {
|
||||
CheckpointLoaderSimple: {
|
||||
input: {
|
||||
required: {
|
||||
ckpt_name: [['flux1-dev.safetensors']],
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
mockClient.getNodeDefs.mockResolvedValue(mockNodeDefs);
|
||||
|
||||
// First call - should fetch from API
|
||||
const result1 = await service.getNodeDefs();
|
||||
expect(result1).toEqual(mockNodeDefs);
|
||||
expect(mockClient.getNodeDefs).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Second call - should use cache
|
||||
const result2 = await service.getNodeDefs();
|
||||
expect(result2).toEqual(mockNodeDefs);
|
||||
expect(mockClient.getNodeDefs).toHaveBeenCalledTimes(1); // Still 1, used cache
|
||||
|
||||
// Get specific node - should return full cache since SDK doesn't support filtering
|
||||
const result3 = await service.getNodeDefs('CheckpointLoaderSimple');
|
||||
expect(result3).toEqual(mockNodeDefs);
|
||||
expect(mockClient.getNodeDefs).toHaveBeenCalledTimes(1); // Still 1, used cache
|
||||
});
|
||||
|
||||
it('should refresh cache after TTL expires', async () => {
|
||||
const mockNodeDefs1 = { node1: {} };
|
||||
const mockNodeDefs2 = { node2: {} };
|
||||
|
||||
mockClient.getNodeDefs
|
||||
.mockResolvedValueOnce(mockNodeDefs1)
|
||||
.mockResolvedValueOnce(mockNodeDefs2);
|
||||
|
||||
// First call
|
||||
const result1 = await service.getNodeDefs();
|
||||
expect(result1).toEqual(mockNodeDefs1);
|
||||
|
||||
// Simulate time passing (more than 1 minute)
|
||||
const originalNow = Date.now;
|
||||
Date.now = vi.fn(() => originalNow() + 61_000);
|
||||
|
||||
// Second call after TTL - should fetch again
|
||||
const result2 = await service.getNodeDefs();
|
||||
expect(result2).toEqual(mockNodeDefs2);
|
||||
expect(mockClient.getNodeDefs).toHaveBeenCalledTimes(2);
|
||||
|
||||
// Restore Date.now
|
||||
Date.now = originalNow;
|
||||
});
|
||||
|
||||
it('should handle error when getting node definitions', async () => {
|
||||
mockClient.getNodeDefs.mockRejectedValue(new Error('Failed to fetch'));
|
||||
|
||||
await expect(service.getNodeDefs()).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('getSamplerInfo', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
mockClient.getSamplerInfo = vi.fn();
|
||||
});
|
||||
|
||||
it('should get sampler info successfully', async () => {
|
||||
const mockSDKResponse = {
|
||||
sampler: ['euler', 'ddim'],
|
||||
scheduler: ['normal', 'karras'],
|
||||
};
|
||||
mockClient.getSamplerInfo.mockResolvedValue(mockSDKResponse);
|
||||
|
||||
const result = await service.getSamplerInfo();
|
||||
|
||||
// Service now returns samplerName instead of sampler for consistency
|
||||
expect(result).toEqual({
|
||||
samplerName: ['euler', 'ddim'],
|
||||
scheduler: ['normal', 'karras'],
|
||||
});
|
||||
expect(mockClient.getSamplerInfo).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle error when getting sampler info', async () => {
|
||||
mockClient.getSamplerInfo.mockRejectedValue(new Error('Failed to fetch'));
|
||||
|
||||
await expect(service.getSamplerInfo()).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('uploadImage', () => {
|
||||
beforeEach(() => {
|
||||
service = new ComfyUIClientService();
|
||||
mockClient.uploadImage = vi.fn();
|
||||
});
|
||||
|
||||
it('should upload image successfully', async () => {
|
||||
const mockFile = new File(['test'], 'test.png', { type: 'image/png' });
|
||||
const mockResponse = {
|
||||
info: {
|
||||
filename: 'uploaded.png',
|
||||
type: 'input',
|
||||
},
|
||||
};
|
||||
|
||||
mockClient.uploadImage.mockResolvedValue(mockResponse);
|
||||
|
||||
const result = await service.uploadImage(mockFile, 'test.png');
|
||||
|
||||
expect(result).toEqual('uploaded.png');
|
||||
expect(mockClient.uploadImage).toHaveBeenCalledWith(mockFile, 'test.png');
|
||||
});
|
||||
|
||||
it('should handle upload error', async () => {
|
||||
const mockFile = new File(['test'], 'test.png', { type: 'image/png' });
|
||||
|
||||
mockClient.uploadImage.mockRejectedValue(new Error('Upload failed'));
|
||||
|
||||
await expect(service.uploadImage(mockFile, 'test.png')).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
});
|
||||
230
src/server/services/comfyui/__tests__/core/errorHandler.test.ts
Normal file
230
src/server/services/comfyui/__tests__/core/errorHandler.test.ts
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
import { AgentRuntimeErrorType } from '@lobechat/model-runtime';
|
||||
import { beforeEach, describe, expect, it } from 'vitest';
|
||||
|
||||
import { ErrorHandlerService } from '@/server/services/comfyui/core/errorHandlerService';
|
||||
import {
|
||||
ConfigError,
|
||||
ServicesError,
|
||||
UtilsError,
|
||||
WorkflowError,
|
||||
} from '@/server/services/comfyui/errors';
|
||||
import { ModelResolverError } from '@/server/services/comfyui/errors/modelResolverError';
|
||||
|
||||
describe('ErrorHandlerService', () => {
|
||||
let service: ErrorHandlerService;
|
||||
|
||||
beforeEach(() => {
|
||||
service = new ErrorHandlerService();
|
||||
});
|
||||
|
||||
describe('handleError', () => {
|
||||
describe('ComfyUI internal errors', () => {
|
||||
it('should handle ConfigError correctly', () => {
|
||||
const error = new ConfigError('Config is invalid', ConfigError.Reasons.INVALID_CONFIG, {
|
||||
config: 'test',
|
||||
});
|
||||
|
||||
expect(() => service.handleError(error)).toThrow();
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
expect(e.cause.error.message).toBe('Config is invalid');
|
||||
expect(e.cause.error.details).toEqual({ config: 'test' });
|
||||
expect(e.cause.provider).toBe('comfyui');
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle WorkflowError with UNSUPPORTED_MODEL', () => {
|
||||
const error = new WorkflowError(
|
||||
'Model not supported',
|
||||
WorkflowError.Reasons.UNSUPPORTED_MODEL,
|
||||
{ model: 'flux1-dev.safetensors' },
|
||||
);
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ModelNotFound);
|
||||
expect(e.cause.error.message).toBe('Model not supported');
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle WorkflowError with MISSING_COMPONENT', () => {
|
||||
const error = new WorkflowError(
|
||||
'Component missing',
|
||||
WorkflowError.Reasons.MISSING_COMPONENT,
|
||||
{ component: 'vae' },
|
||||
);
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ComfyUIModelError);
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle UtilsError correctly', () => {
|
||||
const error = new UtilsError('Connection failed', UtilsError.Reasons.CONNECTION_ERROR);
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ComfyUIServiceUnavailable);
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle ModelResolverError correctly', () => {
|
||||
const error = new ModelResolverError('MODEL_NOT_FOUND', 'Model not found');
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ModelNotFound);
|
||||
}
|
||||
});
|
||||
|
||||
it('should use default error type for unknown reasons', () => {
|
||||
const error = new ConfigError('Unknown error', 'UNKNOWN_REASON' as any);
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle ServicesError with all mapped reasons', () => {
|
||||
// Test a mapped reason
|
||||
const error1 = new ServicesError('Model not found', ServicesError.Reasons.MODEL_NOT_FOUND, {
|
||||
model: 'test',
|
||||
});
|
||||
|
||||
try {
|
||||
service.handleError(error1);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ModelNotFound);
|
||||
}
|
||||
|
||||
// Test unmapped reason - should hit line 120 and return default
|
||||
const error2 = new ServicesError('Unknown error', 'UNMAPPED_REASON' as any, {});
|
||||
|
||||
try {
|
||||
service.handleError(error2);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Pre-formatted framework errors', () => {
|
||||
it('should pass through pre-formatted errors', () => {
|
||||
const error = {
|
||||
errorType: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
message: 'Already formatted',
|
||||
provider: 'comfyui',
|
||||
};
|
||||
|
||||
expect(() => service.handleError(error)).toThrowError();
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
// Verify the cause contains the same error properties
|
||||
expect(e.cause.errorType).toBe(error.errorType);
|
||||
expect(e.cause.message).toBe(error.message);
|
||||
expect(e.cause.provider).toBe(error.provider);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Other errors', () => {
|
||||
it('should parse string errors', () => {
|
||||
const error = 'Some error message';
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.provider).toBe('comfyui');
|
||||
expect(e.cause.error).toBeDefined();
|
||||
}
|
||||
});
|
||||
|
||||
it('should parse Error objects', () => {
|
||||
const error = new Error('Standard error');
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.provider).toBe('comfyui');
|
||||
expect(e.cause.error).toBeDefined();
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle null/undefined errors', () => {
|
||||
expect(() => service.handleError(null)).toThrow();
|
||||
expect(() => service.handleError(undefined)).toThrow();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Error mapping completeness', () => {
|
||||
it('should map all ConfigError reasons', () => {
|
||||
const reasons = Object.values(ConfigError.Reasons);
|
||||
|
||||
reasons.forEach((reason) => {
|
||||
const error = new ConfigError('Test', reason);
|
||||
|
||||
expect(() => service.handleError(error)).toThrow();
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBeDefined();
|
||||
expect(e.cause.errorType).toBe(AgentRuntimeErrorType.ComfyUIBizError);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
it('should map all WorkflowError reasons', () => {
|
||||
const reasons = Object.values(WorkflowError.Reasons);
|
||||
const expectedMapping: Record<string, string> = {
|
||||
[WorkflowError.Reasons.INVALID_CONFIG]: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
[WorkflowError.Reasons.MISSING_COMPONENT]: AgentRuntimeErrorType.ComfyUIModelError,
|
||||
[WorkflowError.Reasons.MISSING_ENCODER]: AgentRuntimeErrorType.ComfyUIModelError,
|
||||
[WorkflowError.Reasons.UNSUPPORTED_MODEL]: AgentRuntimeErrorType.ModelNotFound,
|
||||
[WorkflowError.Reasons.INVALID_PARAMS]: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
};
|
||||
|
||||
reasons.forEach((reason) => {
|
||||
const error = new WorkflowError('Test', reason);
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
const expected = expectedMapping[reason] || AgentRuntimeErrorType.ComfyUIWorkflowError;
|
||||
expect(e.cause.errorType).toBe(expected);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
it('should map all UtilsError reasons', () => {
|
||||
const reasons = Object.values(UtilsError.Reasons);
|
||||
|
||||
reasons.forEach((reason) => {
|
||||
const error = new UtilsError('Test', reason);
|
||||
|
||||
expect(() => service.handleError(error)).toThrow();
|
||||
|
||||
try {
|
||||
service.handleError(error);
|
||||
} catch (e: any) {
|
||||
expect(e.cause.errorType).toBeDefined();
|
||||
// Should not be undefined
|
||||
expect(e.cause.errorType).not.toBe(undefined);
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
134
src/server/services/comfyui/__tests__/core/errorHandling.test.ts
Normal file
134
src/server/services/comfyui/__tests__/core/errorHandling.test.ts
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
// @vitest-environment node
|
||||
import { CallWrapper } from '@saintno/comfyui-sdk';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { parametersFixture } from '@/server/services/comfyui/__tests__/fixtures/parameters.fixture';
|
||||
import { mockContext } from '@/server/services/comfyui/__tests__/helpers/mockContext';
|
||||
import { setupAllMocks } from '@/server/services/comfyui/__tests__/setup/unifiedMocks';
|
||||
import { ComfyUIClientService } from '@/server/services/comfyui/core/comfyUIClientService';
|
||||
// Import services for testing
|
||||
import { ImageService } from '@/server/services/comfyui/core/imageService';
|
||||
import { ModelResolverService } from '@/server/services/comfyui/core/modelResolverService';
|
||||
import { WorkflowBuilderService } from '@/server/services/comfyui/core/workflowBuilderService';
|
||||
|
||||
describe('Error Handling - SDK Integration', () => {
|
||||
let imageService: ImageService;
|
||||
let inputCalls: Map<string, any>;
|
||||
|
||||
beforeEach(() => {
|
||||
const mocks = setupAllMocks();
|
||||
inputCalls = mocks.inputCalls;
|
||||
|
||||
// Create service instances
|
||||
const clientService = new ComfyUIClientService();
|
||||
const modelResolverService = new ModelResolverService(clientService);
|
||||
const workflowBuilderService = new WorkflowBuilderService({
|
||||
clientService,
|
||||
modelResolverService,
|
||||
});
|
||||
|
||||
imageService = new ImageService(clientService, modelResolverService, workflowBuilderService);
|
||||
});
|
||||
|
||||
describe('SDK Error Handling', () => {
|
||||
it('should catch and transform SDK errors', async () => {
|
||||
// This test relies on the unified mock system for CallWrapper
|
||||
|
||||
const params = {
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
...parametersFixture.models['flux-dev'].defaults,
|
||||
},
|
||||
};
|
||||
|
||||
// Should catch and transform errors
|
||||
await expect(imageService.createImage(params)).rejects.toThrow();
|
||||
});
|
||||
|
||||
it('should handle workflow build errors gracefully', async () => {
|
||||
const incompleteParams = {
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: '', // Empty prompt may cause issues
|
||||
},
|
||||
};
|
||||
|
||||
// Should not crash, should return meaningful error
|
||||
await expect(imageService.createImage(incompleteParams)).rejects.toThrow();
|
||||
});
|
||||
|
||||
it('should handle invalid model errors', async () => {
|
||||
const invalidModelParams = {
|
||||
model: 'non-existent-model',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// 应该优雅地处理无效模型
|
||||
await expect(imageService.createImage(invalidModelParams)).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Parameter Validation Errors', () => {
|
||||
it('should validate required parameters', async () => {
|
||||
const missingParams = {
|
||||
model: 'flux-dev',
|
||||
params: { prompt: '' }, // 缺少必要参数，但至少包含必需的prompt
|
||||
};
|
||||
|
||||
// 应该验证并报错
|
||||
await expect(imageService.createImage(missingParams)).rejects.toThrow();
|
||||
});
|
||||
|
||||
it('should handle parameter boundary violations gracefully', async () => {
|
||||
const invalidParams = {
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
cfg: -1, // 无效值
|
||||
steps: 1000, // 超出范围
|
||||
},
|
||||
};
|
||||
|
||||
// 应该处理边界违规
|
||||
await expect(imageService.createImage(invalidParams)).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Service Error Propagation', () => {
|
||||
it('should propagate model resolution errors', async () => {
|
||||
// Mock 模型解析失败
|
||||
vi.spyOn(ModelResolverService.prototype, 'validateModel').mockRejectedValue(
|
||||
new Error('Model not found'),
|
||||
);
|
||||
|
||||
const params = {
|
||||
model: 'unknown-model',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
await expect(imageService.createImage(params)).rejects.toThrow();
|
||||
});
|
||||
|
||||
it('should handle connection validation failures', async () => {
|
||||
// Mock 连接验证失败
|
||||
vi.spyOn(ComfyUIClientService.prototype, 'validateConnection').mockRejectedValue(
|
||||
new Error('Connection failed'),
|
||||
);
|
||||
|
||||
const params = {
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
...parametersFixture.models['flux-dev'].defaults,
|
||||
},
|
||||
};
|
||||
|
||||
await expect(imageService.createImage(params)).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
});
|
||||
528
src/server/services/comfyui/__tests__/core/imageService.test.ts
Normal file
528
src/server/services/comfyui/__tests__/core/imageService.test.ts
Normal file
|
|
@ -0,0 +1,528 @@
|
|||
import { AgentRuntimeErrorType } from '@lobechat/model-runtime';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { CreateImagePayload } from '@lobechat/model-runtime';
|
||||
import { ComfyUIClientService } from '@/server/services/comfyui/core/comfyUIClientService';
|
||||
import { ErrorHandlerService } from '@/server/services/comfyui/core/errorHandlerService';
|
||||
import { ImageService } from '@/server/services/comfyui/core/imageService';
|
||||
import { ModelResolverService } from '@/server/services/comfyui/core/modelResolverService';
|
||||
import { WorkflowBuilderService } from '@/server/services/comfyui/core/workflowBuilderService';
|
||||
import { WorkflowDetector } from '@/server/services/comfyui/utils/workflowDetector';
|
||||
|
||||
// Mock dependencies
|
||||
vi.mock('@/server/services/comfyui/core/comfyUIClientService');
|
||||
vi.mock('@/server/services/comfyui/core/modelResolverService');
|
||||
vi.mock('@/server/services/comfyui/core/workflowBuilderService');
|
||||
vi.mock('@/server/services/comfyui/core/errorHandlerService');
|
||||
vi.mock('@/server/services/comfyui/utils/workflowDetector');
|
||||
|
||||
// Mock sharp module for image processing
|
||||
vi.mock('sharp', () => ({
|
||||
default: vi.fn((buffer) => ({
|
||||
metadata: vi.fn().mockResolvedValue({ height: 1024, width: 1024 }),
|
||||
resize: vi.fn().mockReturnThis(),
|
||||
toBuffer: vi.fn().mockResolvedValue(Buffer.from(buffer)),
|
||||
})),
|
||||
}));
|
||||
|
||||
describe('ImageService', () => {
|
||||
let imageService: ImageService;
|
||||
let mockClientService: any;
|
||||
let mockModelResolverService: any;
|
||||
let mockWorkflowBuilderService: any;
|
||||
let mockErrorHandler: any;
|
||||
let mockFetch: any;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
// Create mocks
|
||||
mockClientService = {
|
||||
executeWorkflow: vi.fn(),
|
||||
getPathImage: vi.fn(),
|
||||
uploadImage: vi.fn(),
|
||||
validateConnection: vi.fn().mockResolvedValue(true),
|
||||
};
|
||||
|
||||
mockModelResolverService = {
|
||||
validateModel: vi.fn(),
|
||||
};
|
||||
|
||||
mockWorkflowBuilderService = {
|
||||
buildWorkflow: vi.fn(),
|
||||
};
|
||||
|
||||
mockErrorHandler = {
|
||||
handleError: vi.fn(),
|
||||
};
|
||||
|
||||
// Mock global fetch
|
||||
mockFetch = vi.fn();
|
||||
global.fetch = mockFetch;
|
||||
|
||||
// Setup mocks for constructor
|
||||
vi.mocked(ComfyUIClientService, true).mockImplementation(() => mockClientService as any);
|
||||
vi.mocked(ModelResolverService, true).mockImplementation(() => mockModelResolverService as any);
|
||||
vi.mocked(WorkflowBuilderService, true).mockImplementation(
|
||||
() => mockWorkflowBuilderService as any,
|
||||
);
|
||||
vi.mocked(ErrorHandlerService, true).mockImplementation(() => mockErrorHandler as any);
|
||||
|
||||
// Create service instance
|
||||
imageService = new ImageService(
|
||||
mockClientService,
|
||||
mockModelResolverService,
|
||||
mockWorkflowBuilderService,
|
||||
);
|
||||
|
||||
// Mock workflow detector
|
||||
vi.mocked(WorkflowDetector, true).detectModelType = vi.fn().mockReturnValue({
|
||||
architecture: 'flux-schnell',
|
||||
isSupported: true,
|
||||
modelType: 'FLUX',
|
||||
});
|
||||
});
|
||||
|
||||
describe('createImage', () => {
|
||||
const mockPayload: CreateImagePayload = {
|
||||
model: 'flux-schnell',
|
||||
params: {
|
||||
height: 1024,
|
||||
prompt: 'test prompt',
|
||||
width: 1024,
|
||||
},
|
||||
};
|
||||
|
||||
it('should successfully create image with text2img workflow', async () => {
|
||||
// Setup mocks
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'flux1-schnell-fp8.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
const mockWorkflow = { id: 'test-workflow' };
|
||||
mockWorkflowBuilderService.buildWorkflow.mockResolvedValue(mockWorkflow);
|
||||
|
||||
const mockResult = {
|
||||
images: {
|
||||
images: [
|
||||
{
|
||||
data: 'base64data',
|
||||
height: 1024,
|
||||
width: 1024,
|
||||
},
|
||||
],
|
||||
},
|
||||
};
|
||||
mockClientService.executeWorkflow.mockResolvedValue(mockResult);
|
||||
mockClientService.getPathImage.mockReturnValue('https://comfyui.test/image.png');
|
||||
|
||||
// Execute
|
||||
const result = await imageService.createImage(mockPayload);
|
||||
|
||||
// Verify
|
||||
expect(result).toEqual({
|
||||
imageUrl: 'https://comfyui.test/image.png',
|
||||
});
|
||||
|
||||
expect(mockModelResolverService.validateModel).toHaveBeenCalledWith('flux-schnell');
|
||||
expect(mockWorkflowBuilderService.buildWorkflow).toHaveBeenCalled();
|
||||
expect(mockClientService.executeWorkflow).toHaveBeenCalledWith(
|
||||
mockWorkflow,
|
||||
expect.any(Function),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle model not found error', async () => {
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
exists: false,
|
||||
});
|
||||
|
||||
mockErrorHandler.handleError.mockImplementation((error: any) => {
|
||||
throw {
|
||||
error: { message: error.message },
|
||||
errorType: AgentRuntimeErrorType.ModelNotFound,
|
||||
provider: 'comfyui',
|
||||
};
|
||||
});
|
||||
|
||||
// Execute and verify
|
||||
await expect(imageService.createImage(mockPayload)).rejects.toMatchObject({
|
||||
errorType: AgentRuntimeErrorType.ModelNotFound,
|
||||
provider: 'comfyui',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle empty result from workflow', async () => {
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'flux1-schnell-fp8.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
mockWorkflowBuilderService.buildWorkflow.mockResolvedValue({});
|
||||
mockClientService.executeWorkflow.mockResolvedValue({
|
||||
images: { images: [] },
|
||||
});
|
||||
|
||||
mockErrorHandler.handleError.mockImplementation((error: any) => {
|
||||
throw {
|
||||
error: { message: error.message },
|
||||
errorType: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
provider: 'comfyui',
|
||||
};
|
||||
});
|
||||
|
||||
// Execute and verify
|
||||
await expect(imageService.createImage(mockPayload)).rejects.toMatchObject({
|
||||
errorType: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
provider: 'comfyui',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('processImageFetch', () => {
|
||||
const mockPayloadWithImage: CreateImagePayload = {
|
||||
model: 'flux-schnell',
|
||||
params: {
|
||||
height: 1024,
|
||||
imageUrl: 'https://s3.test/bucket/image.png',
|
||||
prompt: 'test prompt',
|
||||
width: 1024,
|
||||
},
|
||||
};
|
||||
|
||||
it('should fetch image from URL and upload to ComfyUI', async () => {
|
||||
// Setup mocks
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'flux1-schnell-fp8.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
// Fetch mocks
|
||||
const mockImageData = new Uint8Array([1, 2, 3, 4, 5]);
|
||||
mockFetch.mockResolvedValue({
|
||||
arrayBuffer: vi.fn().mockResolvedValue(mockImageData.buffer),
|
||||
ok: true,
|
||||
});
|
||||
|
||||
// Upload mock
|
||||
mockClientService.uploadImage.mockResolvedValue('img2img_123456.png');
|
||||
|
||||
// Workflow mocks
|
||||
mockWorkflowBuilderService.buildWorkflow.mockResolvedValue({});
|
||||
mockClientService.executeWorkflow.mockResolvedValue({
|
||||
images: { images: [{ height: 1024, width: 1024 }] },
|
||||
});
|
||||
mockClientService.getPathImage.mockReturnValue('https://comfyui.test/result.png');
|
||||
|
||||
// Execute
|
||||
await imageService.createImage(mockPayloadWithImage);
|
||||
|
||||
// Verify fetch was called with the image URL
|
||||
expect(mockFetch).toHaveBeenCalledWith('https://s3.test/bucket/image.png');
|
||||
// Note: uploadImage won't be called in test environment since window exists (jsdom)
|
||||
// and sharp code is skipped. The actual image processing is tested in integration tests.
|
||||
|
||||
// Verify the original params are NOT modified (we clone them now)
|
||||
expect(mockPayloadWithImage.params.imageUrl).toBe('https://s3.test/bucket/image.png');
|
||||
});
|
||||
|
||||
it('should skip processing if imageUrl is already a ComfyUI filename', async () => {
|
||||
const payloadWithFilename: CreateImagePayload = {
|
||||
model: 'flux-schnell',
|
||||
params: {
|
||||
imageUrl: 'existing_image.png',
|
||||
prompt: 'test prompt', // Not a URL
|
||||
},
|
||||
};
|
||||
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'flux1-schnell-fp8.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
mockWorkflowBuilderService.buildWorkflow.mockResolvedValue({});
|
||||
mockClientService.executeWorkflow.mockResolvedValue({
|
||||
images: { images: [{}] },
|
||||
});
|
||||
mockClientService.getPathImage.mockReturnValue('result.png');
|
||||
|
||||
// Execute
|
||||
await imageService.createImage(payloadWithFilename);
|
||||
|
||||
// Verify fetch was not called
|
||||
expect(mockFetch).not.toHaveBeenCalled();
|
||||
expect(mockClientService.uploadImage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle fetch error', async () => {
|
||||
const payload: CreateImagePayload = {
|
||||
model: 'flux-schnell',
|
||||
params: {
|
||||
imageUrl: 'https://s3.test/missing.png',
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'model.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
// Fetch error
|
||||
mockFetch.mockResolvedValue({
|
||||
ok: false,
|
||||
status: 404,
|
||||
statusText: 'Not Found',
|
||||
});
|
||||
|
||||
mockErrorHandler.handleError.mockImplementation((error: any) => {
|
||||
throw error;
|
||||
});
|
||||
|
||||
// Execute and verify
|
||||
await expect(imageService.createImage(payload)).rejects.toThrow(
|
||||
/Failed to fetch image: 404 Not Found/,
|
||||
);
|
||||
});
|
||||
|
||||
it('should not modify original params object', async () => {
|
||||
const originalImageUrl = 'https://s3.test/original.png';
|
||||
const payload: CreateImagePayload = {
|
||||
model: 'flux-schnell',
|
||||
params: {
|
||||
imageUrl: originalImageUrl,
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'flux1-schnell-fp8.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
// Mock WorkflowDetector to return proper architecture
|
||||
vi.mocked(WorkflowDetector, true).detectModelType = vi.fn().mockReturnValue({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
modelType: 'FLUX',
|
||||
});
|
||||
|
||||
// Mock fetch and upload
|
||||
mockFetch.mockResolvedValue({
|
||||
arrayBuffer: vi.fn().mockResolvedValue(new ArrayBuffer(100)),
|
||||
ok: true,
|
||||
});
|
||||
mockClientService.uploadImage.mockResolvedValue('uploaded.png');
|
||||
mockWorkflowBuilderService.buildWorkflow.mockResolvedValue({});
|
||||
mockClientService.executeWorkflow.mockResolvedValue({
|
||||
images: { images: [{}] },
|
||||
});
|
||||
mockClientService.getPathImage.mockReturnValue('result.png');
|
||||
|
||||
// Execute
|
||||
await imageService.createImage(payload);
|
||||
|
||||
// Verify original params are NOT modified
|
||||
expect(payload.params.imageUrl).toBe(originalImageUrl);
|
||||
});
|
||||
|
||||
it('should handle empty image data', async () => {
|
||||
const payload: CreateImagePayload = {
|
||||
model: 'flux-schnell',
|
||||
params: {
|
||||
imageUrl: 'https://s3.test/empty.png',
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'model.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
// Empty image data
|
||||
mockFetch.mockResolvedValue({
|
||||
arrayBuffer: vi.fn().mockResolvedValue(new ArrayBuffer(0)),
|
||||
ok: true,
|
||||
});
|
||||
|
||||
mockErrorHandler.handleError.mockImplementation((error: any) => {
|
||||
throw error;
|
||||
});
|
||||
|
||||
// Execute and verify
|
||||
await expect(imageService.createImage(payload)).rejects.toThrow(/Invalid image data/);
|
||||
});
|
||||
|
||||
it('should handle network fetch errors', async () => {
|
||||
const payload: CreateImagePayload = {
|
||||
model: 'flux-schnell',
|
||||
params: {
|
||||
imageUrl: 'https://s3.test/network-error.png',
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'model.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
// Network error
|
||||
mockFetch.mockRejectedValue(new TypeError('Failed to fetch'));
|
||||
|
||||
mockErrorHandler.handleError.mockImplementation((error: any) => {
|
||||
throw error;
|
||||
});
|
||||
|
||||
// Execute and verify
|
||||
await expect(imageService.createImage(payload)).rejects.toThrow(/Failed to fetch/);
|
||||
});
|
||||
|
||||
it('should handle imageUrls array format', async () => {
|
||||
const payloadWithArray: CreateImagePayload = {
|
||||
model: 'flux-schnell',
|
||||
params: {
|
||||
imageUrls: ['https://s3.test/image.png'],
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'model.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
// S3 mocks
|
||||
mockFetch.mockResolvedValue({
|
||||
arrayBuffer: vi.fn().mockResolvedValue(new Uint8Array([1, 2, 3]).buffer),
|
||||
ok: true,
|
||||
});
|
||||
mockClientService.uploadImage.mockResolvedValue('uploaded.png');
|
||||
|
||||
// Workflow mocks
|
||||
mockWorkflowBuilderService.buildWorkflow.mockResolvedValue({});
|
||||
mockClientService.executeWorkflow.mockResolvedValue({
|
||||
images: { images: [{}] },
|
||||
});
|
||||
mockClientService.getPathImage.mockReturnValue('result.png');
|
||||
|
||||
// Execute
|
||||
await imageService.createImage(payloadWithArray);
|
||||
|
||||
// Verify original params are NOT modified (we clone them now)
|
||||
expect(payloadWithArray.params.imageUrl).toBeUndefined();
|
||||
expect(payloadWithArray.params.imageUrls![0]).toBe('https://s3.test/image.png');
|
||||
});
|
||||
});
|
||||
|
||||
describe('buildWorkflow', () => {
|
||||
it('should detect unsupported models', async () => {
|
||||
const payload: CreateImagePayload = {
|
||||
model: 'unsupported-model',
|
||||
params: { prompt: 'test prompt' },
|
||||
};
|
||||
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'unsupported.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
// Mock unsupported detection
|
||||
vi.mocked(WorkflowDetector).detectModelType = vi.fn().mockReturnValue({
|
||||
isSupported: false,
|
||||
});
|
||||
|
||||
mockErrorHandler.handleError.mockImplementation((error: any) => {
|
||||
throw {
|
||||
error: { message: error.message },
|
||||
errorType: AgentRuntimeErrorType.ModelNotFound,
|
||||
provider: 'comfyui',
|
||||
};
|
||||
});
|
||||
|
||||
// Execute and verify
|
||||
await expect(imageService.createImage(payload)).rejects.toMatchObject({
|
||||
errorType: AgentRuntimeErrorType.ModelNotFound,
|
||||
});
|
||||
});
|
||||
|
||||
it('should pass correct parameters to workflow builder', async () => {
|
||||
const payload: CreateImagePayload = {
|
||||
model: 'sd3.5-large',
|
||||
params: {
|
||||
height: 768,
|
||||
prompt: 'test',
|
||||
width: 1024,
|
||||
},
|
||||
};
|
||||
|
||||
// Setup
|
||||
mockModelResolverService.validateModel.mockResolvedValue({
|
||||
actualFileName: 'sd3.5_large.safetensors',
|
||||
exists: true,
|
||||
});
|
||||
|
||||
const detectionResult = {
|
||||
architecture: 'sd35-large',
|
||||
isSupported: true,
|
||||
modelType: 'SD35',
|
||||
};
|
||||
|
||||
vi.mocked(WorkflowDetector).detectModelType = vi.fn().mockReturnValue(detectionResult);
|
||||
|
||||
mockWorkflowBuilderService.buildWorkflow.mockResolvedValue({});
|
||||
mockClientService.executeWorkflow.mockResolvedValue({
|
||||
images: { images: [{}] },
|
||||
});
|
||||
mockClientService.getPathImage.mockReturnValue('result.png');
|
||||
|
||||
// Execute
|
||||
await imageService.createImage(payload);
|
||||
|
||||
// Verify workflow builder was called correctly
|
||||
expect(mockWorkflowBuilderService.buildWorkflow).toHaveBeenCalledWith(
|
||||
'sd3.5-large',
|
||||
detectionResult,
|
||||
'sd3.5_large.safetensors',
|
||||
payload.params,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('error handling delegation', () => {
|
||||
it('should delegate all errors to ErrorHandlerService', async () => {
|
||||
const payload: CreateImagePayload = {
|
||||
model: 'test',
|
||||
params: { prompt: 'test prompt' },
|
||||
};
|
||||
|
||||
// Setup error
|
||||
const testError = new Error('Test error');
|
||||
mockModelResolverService.validateModel.mockRejectedValue(testError);
|
||||
|
||||
mockErrorHandler.handleError.mockImplementation(() => {
|
||||
throw { original: testError, transformed: true };
|
||||
});
|
||||
|
||||
// Execute
|
||||
await expect(imageService.createImage(payload)).rejects.toMatchObject({
|
||||
original: testError,
|
||||
transformed: true,
|
||||
});
|
||||
|
||||
// Verify error handler was called
|
||||
expect(mockErrorHandler.handleError).toHaveBeenCalledWith(testError);
|
||||
});
|
||||
});
|
||||
});
|
||||
454
src/server/services/comfyui/__tests__/core/modelResolver.test.ts
Normal file
454
src/server/services/comfyui/__tests__/core/modelResolver.test.ts
Normal file
|
|
@ -0,0 +1,454 @@
|
|||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import {
|
||||
TEST_CUSTOM_SD,
|
||||
TEST_FLUX_MODELS,
|
||||
TEST_MODEL_SETS,
|
||||
TEST_SD35_MODELS,
|
||||
} from '@/server/services/comfyui/__tests__/fixtures/testModels';
|
||||
import { ComfyUIClientService } from '@/server/services/comfyui/core/comfyUIClientService';
|
||||
import { ModelResolverService } from '@/server/services/comfyui/core/modelResolverService';
|
||||
import { ModelResolverError } from '@/server/services/comfyui/errors/modelResolverError';
|
||||
|
||||
// Mock ComfyUI Client Service
|
||||
vi.mock('@/server/services/comfyui/core/comfyUIClientService', () => ({
|
||||
ComfyUIClientService: vi.fn(),
|
||||
}));
|
||||
|
||||
// Mock the config module
|
||||
vi.mock('@/server/services/comfyui/config/modelRegistry', () => {
|
||||
const configs: Record<string, any> = {
|
||||
'flux1-dev.safetensors': {
|
||||
family: 'flux',
|
||||
modelFamily: 'FLUX',
|
||||
variant: 'dev',
|
||||
},
|
||||
'flux1-schnell.safetensors': {
|
||||
family: 'flux',
|
||||
modelFamily: 'FLUX',
|
||||
variant: 'schnell',
|
||||
},
|
||||
'sd3.5_large.safetensors': {
|
||||
family: 'sd35',
|
||||
features: { inclclip: false },
|
||||
modelFamily: 'SD3.5',
|
||||
variant: 'sd35',
|
||||
},
|
||||
'sd3.5_large_inclclip.safetensors': {
|
||||
family: 'sd35',
|
||||
features: { inclclip: true },
|
||||
modelFamily: 'SD3.5',
|
||||
variant: 'sd35-inclclip',
|
||||
},
|
||||
'sdxl_base.safetensors': {
|
||||
family: 'sdxl',
|
||||
modelFamily: 'SDXL',
|
||||
variant: 'sdxl-t2i',
|
||||
},
|
||||
};
|
||||
|
||||
return {
|
||||
MODEL_ID_VARIANT_MAP: {
|
||||
'flux-dev': 'dev',
|
||||
'flux-schnell': 'schnell', // Fixed to match actual mapping
|
||||
'stable-diffusion-35': 'sd35',
|
||||
},
|
||||
MODEL_REGISTRY: configs,
|
||||
};
|
||||
});
|
||||
|
||||
// Mock the staticModelLookup module
|
||||
vi.mock('../utils/staticModelLookup', () => {
|
||||
const configs: Record<string, any> = {
|
||||
'flux1-dev.safetensors': {
|
||||
family: 'flux',
|
||||
modelFamily: 'FLUX',
|
||||
variant: 'dev',
|
||||
},
|
||||
'flux1-schnell.safetensors': {
|
||||
family: 'flux',
|
||||
modelFamily: 'FLUX',
|
||||
variant: 'schnell',
|
||||
},
|
||||
'sd3.5_large.safetensors': {
|
||||
family: 'sd35',
|
||||
features: { inclclip: false },
|
||||
modelFamily: 'SD3.5',
|
||||
variant: 'sd35',
|
||||
},
|
||||
'sd3.5_large_inclclip.safetensors': {
|
||||
family: 'sd35',
|
||||
features: { inclclip: true },
|
||||
modelFamily: 'SD3.5',
|
||||
variant: 'sd35-inclclip',
|
||||
},
|
||||
'sdxl_base.safetensors': {
|
||||
family: 'sdxl',
|
||||
modelFamily: 'SDXL',
|
||||
variant: 'sdxl-t2i',
|
||||
},
|
||||
};
|
||||
|
||||
return {
|
||||
getModelConfig: vi.fn((filename: string) => {
|
||||
return configs[filename] || null;
|
||||
}),
|
||||
getModelsByVariant: vi.fn((variant: string) => {
|
||||
// Return models sorted by priority (mock implementation)
|
||||
const models = Object.entries(configs)
|
||||
.filter(([, config]) => config.variant === variant)
|
||||
.map(([filename]) => filename);
|
||||
return models;
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock('@/server/services/comfyui/config/systemComponents', () => ({
|
||||
SYSTEM_COMPONENTS: {
|
||||
'clip_g.safetensors': {
|
||||
modelFamily: 'SD3',
|
||||
priority: 1,
|
||||
type: 'clip',
|
||||
},
|
||||
'clip_l.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'clip',
|
||||
},
|
||||
't5-v1_1-xxl-encoder.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 't5',
|
||||
},
|
||||
't5xxl_fp16.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 't5',
|
||||
},
|
||||
},
|
||||
getSystemComponents: vi.fn(() => ({
|
||||
flux: {
|
||||
clip: ['t5xxl_fp16.safetensors', 'clip_l.safetensors'],
|
||||
t5: 't5-v1_1-xxl-encoder',
|
||||
},
|
||||
sd35: {
|
||||
clip: ['clip_g.safetensors', 'clip_l.safetensors', 't5xxl_fp16.safetensors'],
|
||||
},
|
||||
})),
|
||||
}));
|
||||
|
||||
describe('ModelResolverService', () => {
|
||||
let service: ModelResolverService;
|
||||
let mockClientService: any;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
mockClientService = {
|
||||
getCheckpoints: vi.fn(),
|
||||
getNodeDefs: vi.fn(),
|
||||
};
|
||||
|
||||
service = new ModelResolverService(mockClientService as ComfyUIClientService);
|
||||
});
|
||||
|
||||
describe('resolveModelFileName', () => {
|
||||
it('should return undefined for unregistered model ID', async () => {
|
||||
// Model not in registry and not on server should return undefined
|
||||
mockClientService.getCheckpoints.mockResolvedValue([TEST_SD35_MODELS.LARGE]);
|
||||
|
||||
const result = await service.resolveModelFileName('nonexistent-model');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return filename if already a file', async () => {
|
||||
// Mock getCheckpoints to include the file
|
||||
mockClientService.getCheckpoints.mockResolvedValue([
|
||||
TEST_FLUX_MODELS.DEV,
|
||||
TEST_FLUX_MODELS.SCHNELL,
|
||||
]);
|
||||
|
||||
const result = await service.resolveModelFileName(TEST_FLUX_MODELS.DEV);
|
||||
expect(result).toBe(TEST_FLUX_MODELS.DEV);
|
||||
});
|
||||
|
||||
it('should use cache on subsequent calls', async () => {
|
||||
// Use a non-registry model that requires server check
|
||||
const customModel = 'custom_test_model.safetensors';
|
||||
mockClientService.getCheckpoints.mockResolvedValue([customModel]);
|
||||
|
||||
// First call
|
||||
await service.resolveModelFileName(customModel);
|
||||
// Second call should use cache
|
||||
const result = await service.resolveModelFileName(customModel);
|
||||
|
||||
expect(result).toBe(customModel);
|
||||
// Should only call once due to caching
|
||||
expect(mockClientService.getCheckpoints).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should resolve custom SD model to fixed filename', async () => {
|
||||
mockClientService.getCheckpoints.mockResolvedValue([TEST_CUSTOM_SD, TEST_SD35_MODELS.LARGE]);
|
||||
|
||||
const result = await service.resolveModelFileName('stable-diffusion-custom');
|
||||
expect(result).toBe(TEST_CUSTOM_SD);
|
||||
});
|
||||
|
||||
it('should resolve custom SD refiner model to same fixed filename', async () => {
|
||||
mockClientService.getCheckpoints.mockResolvedValue([TEST_CUSTOM_SD, TEST_SD35_MODELS.LARGE]);
|
||||
|
||||
const result = await service.resolveModelFileName('stable-diffusion-custom-refiner');
|
||||
expect(result).toBe(TEST_CUSTOM_SD);
|
||||
});
|
||||
|
||||
it('should throw error if custom SD model file not found', async () => {
|
||||
mockClientService.getCheckpoints.mockResolvedValue([TEST_SD35_MODELS.LARGE]);
|
||||
|
||||
const result = await service.resolveModelFileName('stable-diffusion-custom');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('validateModel', () => {
|
||||
it('should validate existing model file on server', async () => {
|
||||
mockClientService.getCheckpoints.mockResolvedValue([
|
||||
TEST_FLUX_MODELS.DEV,
|
||||
TEST_FLUX_MODELS.SCHNELL,
|
||||
]);
|
||||
|
||||
const result = await service.validateModel(TEST_FLUX_MODELS.DEV);
|
||||
|
||||
expect(result.exists).toBe(true);
|
||||
expect(result.actualFileName).toBe(TEST_FLUX_MODELS.DEV);
|
||||
});
|
||||
|
||||
it('should throw error for non-existent model', async () => {
|
||||
mockClientService.getCheckpoints.mockResolvedValue([TEST_SD35_MODELS.LARGE]);
|
||||
|
||||
await expect(service.validateModel(TEST_MODEL_SETS.NON_EXISTENT[0])).rejects.toThrow(
|
||||
'Model not found: , please install one first.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should re-throw ModelResolverError from network errors', async () => {
|
||||
// Network error in getCheckpoints leads to CONNECTION_ERROR in handleApiError
|
||||
// But then resolveModelFileName catches it and throws MODEL_NOT_FOUND
|
||||
mockClientService.getCheckpoints.mockRejectedValue(new TypeError('Failed to fetch'));
|
||||
|
||||
try {
|
||||
await service.validateModel('test-model');
|
||||
expect(true).toBe(false); // Should not reach here
|
||||
} catch (error) {
|
||||
expect(error).toBeInstanceOf(ModelResolverError);
|
||||
// The error gets re-thrown as MODEL_NOT_FOUND by resolveModelFileName
|
||||
expect((error as any).reason).toBe('MODEL_NOT_FOUND');
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('cache management', () => {
|
||||
it('should use cached VAE data when available', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({
|
||||
VAELoader: {
|
||||
input: {
|
||||
required: {
|
||||
vae_name: [['vae1.safetensors', 'vae2.safetensors']],
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// First call - populates cache
|
||||
const result1 = await service.getAvailableVAEFiles();
|
||||
expect(result1).toEqual(['vae1.safetensors', 'vae2.safetensors']);
|
||||
expect(mockClientService.getNodeDefs).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Second call - ModelResolverService doesn't cache, but ClientService does
|
||||
const result2 = await service.getAvailableVAEFiles();
|
||||
expect(result2).toEqual(['vae1.safetensors', 'vae2.safetensors']);
|
||||
expect(mockClientService.getNodeDefs).toHaveBeenCalledTimes(2); // Called again, caching is in ClientService
|
||||
});
|
||||
|
||||
it('should use cached component data when available', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({
|
||||
CheckpointLoaderSimple: {
|
||||
input: {
|
||||
required: {
|
||||
ckpt_name: [['model1.safetensors', 'model2.safetensors']],
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// First call - populates cache
|
||||
const result1 = await service.getAvailableComponentFiles(
|
||||
'CheckpointLoaderSimple',
|
||||
'ckpt_name',
|
||||
);
|
||||
expect(result1).toEqual(['model1.safetensors', 'model2.safetensors']);
|
||||
expect(mockClientService.getNodeDefs).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Second call - ModelResolverService doesn't cache, but ClientService does
|
||||
const result2 = await service.getAvailableComponentFiles(
|
||||
'CheckpointLoaderSimple',
|
||||
'ckpt_name',
|
||||
);
|
||||
expect(result2).toEqual(['model1.safetensors', 'model2.safetensors']);
|
||||
expect(mockClientService.getNodeDefs).toHaveBeenCalledTimes(2); // Called again, caching is in ClientService
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAvailableVAEFiles edge cases', () => {
|
||||
it('should handle non-array VAE list from getNodeDefs', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({
|
||||
VAELoader: {
|
||||
input: {
|
||||
required: {
|
||||
vae_name: [{}], // Object instead of array
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const result = await service.getAvailableVAEFiles();
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should handle missing VAELoader node', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({});
|
||||
|
||||
const result = await service.getAvailableVAEFiles();
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should handle missing input in VAELoader', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({
|
||||
VAELoader: {},
|
||||
});
|
||||
|
||||
const result = await service.getAvailableVAEFiles();
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should handle missing required in VAELoader input', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({
|
||||
VAELoader: {
|
||||
input: {},
|
||||
},
|
||||
});
|
||||
|
||||
const result = await service.getAvailableVAEFiles();
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should handle missing vae_name in required', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({
|
||||
VAELoader: {
|
||||
input: {
|
||||
required: {},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const result = await service.getAvailableVAEFiles();
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getAvailableComponentFiles edge cases', () => {
|
||||
it('should handle non-array component list', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({
|
||||
VAELoader: {
|
||||
input: {
|
||||
required: {
|
||||
vae_name: [{}], // Object instead of array
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const result = await service.getAvailableComponentFiles('VAELoader', 'vae_name');
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
|
||||
it('should handle string component list', async () => {
|
||||
mockClientService.getNodeDefs.mockResolvedValue({
|
||||
VAELoader: {
|
||||
input: {
|
||||
required: {
|
||||
vae_name: ['not-an-array'], // String instead of array
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const result = await service.getAvailableComponentFiles('VAELoader', 'vae_name');
|
||||
expect(result).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('validateModel edge cases', () => {
|
||||
it('should re-throw non-ModelResolverError errors', async () => {
|
||||
// Mock to throw a regular error instead of ModelResolverError
|
||||
mockClientService.getCheckpoints.mockRejectedValue(new Error('Network error'));
|
||||
|
||||
await expect(service.validateModel('test-model.safetensors')).rejects.toThrow(
|
||||
'Network error',
|
||||
);
|
||||
});
|
||||
|
||||
it('should re-throw ModelResolverError', async () => {
|
||||
// Mock to throw ModelResolverError
|
||||
const modelError = new ModelResolverError('Test error', 'TEST_ERROR');
|
||||
mockClientService.getCheckpoints.mockRejectedValue(modelError);
|
||||
|
||||
await expect(service.validateModel('test-model.safetensors')).rejects.toThrow(
|
||||
ModelResolverError,
|
||||
);
|
||||
});
|
||||
|
||||
it('should include expected files in error message when model not found', async () => {
|
||||
// Mock getCheckpoints to return empty array (no models available)
|
||||
mockClientService.getCheckpoints.mockResolvedValue([]);
|
||||
|
||||
// Validate a known variant should throw with expected files
|
||||
await expect(service.validateModel('comfyui/flux-schnell')).rejects.toMatchObject({
|
||||
details: {
|
||||
// The actual variant from MODEL_ID_VARIANT_MAP
|
||||
expectedFiles: expect.arrayContaining(['flux1-schnell.safetensors']),
|
||||
|
||||
modelId: 'comfyui/flux-schnell',
|
||||
variant: 'schnell',
|
||||
},
|
||||
message: expect.stringContaining(
|
||||
'Model not found: flux1-schnell.safetensors, please install one first.',
|
||||
),
|
||||
});
|
||||
|
||||
// Also verify the message contains expected files
|
||||
await expect(service.validateModel('comfyui/flux-schnell')).rejects.toThrow(
|
||||
'Model not found: flux1-schnell.safetensors, please install one first.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should not include expected files for unknown models', async () => {
|
||||
// Mock getCheckpoints to return empty array
|
||||
mockClientService.getCheckpoints.mockResolvedValue([]);
|
||||
|
||||
// Validate an unknown model should throw without expected files
|
||||
await expect(service.validateModel('comfyui/unknown-model')).rejects.toMatchObject({
|
||||
details: {
|
||||
expectedFiles: [],
|
||||
modelId: 'comfyui/unknown-model',
|
||||
variant: undefined,
|
||||
},
|
||||
message: 'Model not found: , please install one first.',
|
||||
});
|
||||
|
||||
// Verify the message doesn't contain expected files
|
||||
await expect(service.validateModel('comfyui/unknown-model')).rejects.toThrow(
|
||||
'Model not found: , please install one first.',
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,294 @@
|
|||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
// Import real test data
|
||||
import {
|
||||
TEST_COMPONENTS,
|
||||
TEST_MODELS,
|
||||
} from '@/server/services/comfyui/__tests__/helpers/realConfigData';
|
||||
import { ComfyUIClientService } from '@/server/services/comfyui/core/comfyUIClientService';
|
||||
import { ModelResolverService } from '@/server/services/comfyui/core/modelResolverService';
|
||||
import {
|
||||
WorkflowBuilderService,
|
||||
WorkflowContext,
|
||||
} from '@/server/services/comfyui/core/workflowBuilderService';
|
||||
import { WorkflowError } from '@/server/services/comfyui/errors';
|
||||
|
||||
// Mock dependencies (must be before other imports)
|
||||
vi.mock('@/server/services/comfyui/core/comfyUIClientService');
|
||||
vi.mock('@/server/services/comfyui/core/modelResolverService');
|
||||
|
||||
// No need to mock modelRegistry - we want to use the real implementation
|
||||
|
||||
describe('WorkflowBuilderService', () => {
|
||||
let service: WorkflowBuilderService;
|
||||
let mockContext: WorkflowContext;
|
||||
let mockModelResolver: any;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
|
||||
mockModelResolver = {
|
||||
getAvailableVAEFiles: vi.fn().mockResolvedValue(['sdxl_vae.safetensors']),
|
||||
getOptimalComponent: vi.fn(),
|
||||
};
|
||||
|
||||
mockContext = {
|
||||
clientService: {} as ComfyUIClientService,
|
||||
modelResolverService: mockModelResolver as ModelResolverService,
|
||||
};
|
||||
|
||||
service = new WorkflowBuilderService(mockContext);
|
||||
});
|
||||
|
||||
describe('buildWorkflow', () => {
|
||||
it('should build FLUX workflow', async () => {
|
||||
// Mock component resolution with real component names
|
||||
mockModelResolver.getOptimalComponent
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.flux.t5) // First call for t5
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.flux.clip); // Second call for clip
|
||||
|
||||
const workflow = await service.buildWorkflow(
|
||||
'flux-1-dev',
|
||||
{ architecture: 'FLUX', isSupported: true },
|
||||
TEST_MODELS.flux,
|
||||
{
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: 'A beautiful landscape',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
},
|
||||
);
|
||||
|
||||
expect(workflow).toBeDefined();
|
||||
expect(workflow.workflow).toBeDefined();
|
||||
// Verify component resolution was called
|
||||
expect(mockModelResolver.getOptimalComponent).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should build SD/SDXL workflow with VAE', async () => {
|
||||
mockModelResolver.getOptimalComponent.mockResolvedValue(TEST_COMPONENTS.sd.vae);
|
||||
|
||||
const workflow = await service.buildWorkflow(
|
||||
'stable-diffusion-xl',
|
||||
{ architecture: 'SDXL', isSupported: true },
|
||||
TEST_MODELS.sdxl,
|
||||
{
|
||||
cfg: 7,
|
||||
height: 1024,
|
||||
negativePrompt: 'blurry, ugly',
|
||||
prompt: 'A beautiful landscape',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
},
|
||||
);
|
||||
|
||||
expect(workflow).toBeDefined();
|
||||
expect(workflow.workflow).toBeDefined();
|
||||
|
||||
// Check if VAE loader was added
|
||||
const nodes = workflow.workflow as any;
|
||||
const vaeNode = Object.values(nodes).find((node: any) => node.class_type === 'VAELoader');
|
||||
expect(vaeNode).toBeDefined();
|
||||
});
|
||||
|
||||
it('should build SD3.5 workflow', async () => {
|
||||
mockModelResolver.getOptimalComponent
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.sd.clip) // clip_g
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.flux.clip) // clip_l (reuse from FLUX)
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.flux.t5); // t5xxl (reuse from FLUX)
|
||||
|
||||
const workflow = await service.buildWorkflow(
|
||||
'stable-diffusion-35',
|
||||
{ architecture: 'SD3', isSupported: true, variant: 'sd35' },
|
||||
TEST_MODELS.sd35,
|
||||
{
|
||||
cfg: 4.5,
|
||||
height: 1024,
|
||||
prompt: 'A futuristic city',
|
||||
shift: 3,
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
},
|
||||
);
|
||||
|
||||
expect(workflow).toBeDefined();
|
||||
expect(workflow.workflow).toBeDefined();
|
||||
|
||||
// Check for SD3.5 specific nodes
|
||||
const nodes = workflow.workflow as any;
|
||||
const samplingNode = Object.values(nodes).find(
|
||||
(node: any) => node.class_type === 'ModelSamplingSD3',
|
||||
);
|
||||
expect(samplingNode).toBeDefined();
|
||||
});
|
||||
|
||||
it('should throw error for unsupported model type', async () => {
|
||||
await expect(
|
||||
service.buildWorkflow(
|
||||
'unknown-model',
|
||||
{ architecture: 'UNKNOWN' as any, isSupported: false },
|
||||
'unknown.safetensors',
|
||||
{},
|
||||
),
|
||||
).rejects.toThrow(WorkflowError);
|
||||
});
|
||||
});
|
||||
|
||||
describe('FLUX workflow specifics', () => {
|
||||
it('should throw error when required components not found', async () => {
|
||||
// Mock component resolution to fail - should throw error (not use defaults)
|
||||
mockModelResolver.getOptimalComponent.mockRejectedValue(
|
||||
new Error('Required CLIP component not found'),
|
||||
);
|
||||
|
||||
// Should throw error because required components are missing
|
||||
await expect(
|
||||
service.buildWorkflow(
|
||||
'flux-1-dev',
|
||||
{ architecture: 'FLUX', isSupported: true },
|
||||
TEST_MODELS.flux,
|
||||
{ prompt: 'test' },
|
||||
),
|
||||
).rejects.toThrow('Required CLIP component not found');
|
||||
});
|
||||
|
||||
it('should use default parameters when not provided', async () => {
|
||||
// Mock component resolution with real components
|
||||
mockModelResolver.getOptimalComponent
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.flux.t5)
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.flux.clip);
|
||||
|
||||
const workflow = await service.buildWorkflow(
|
||||
'flux-1-dev',
|
||||
{ architecture: 'FLUX', isSupported: true },
|
||||
TEST_MODELS.flux,
|
||||
{ prompt: 'test' },
|
||||
);
|
||||
|
||||
// The workflow should be built with default params
|
||||
expect(workflow).toBeDefined();
|
||||
// Verify component resolution was called
|
||||
expect(mockModelResolver.getOptimalComponent).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('SD/SDXL workflow specifics', () => {
|
||||
it('should build workflow with VAE loader when external VAE specified', async () => {
|
||||
// For SDXL, the workflow will use sdxl_vae.safetensors from getOptimalVAEForModel
|
||||
mockModelResolver.getOptimalComponent.mockResolvedValue('sdxl_vae.safetensors');
|
||||
|
||||
const workflow = await service.buildWorkflow(
|
||||
'stable-diffusion-xl',
|
||||
{ architecture: 'SDXL', isSupported: true },
|
||||
TEST_MODELS.sdxl,
|
||||
{ prompt: 'test' },
|
||||
);
|
||||
|
||||
const nodes = workflow.workflow as any;
|
||||
const vaeLoader = Object.values(nodes).find((node: any) => node.class_type === 'VAELoader');
|
||||
|
||||
// Should have VAE loader with the priority 1 SDXL VAE from config
|
||||
expect(vaeLoader).toBeDefined();
|
||||
expect((vaeLoader as any).inputs.vae_name).toBe('sdxl_vae.safetensors');
|
||||
});
|
||||
|
||||
it('should support custom sampler and scheduler', async () => {
|
||||
mockModelResolver.getOptimalComponent.mockResolvedValue(undefined);
|
||||
|
||||
const workflow = await service.buildWorkflow(
|
||||
'stable-diffusion-xl',
|
||||
{ architecture: 'SDXL', isSupported: true },
|
||||
TEST_MODELS.sdxl,
|
||||
{
|
||||
prompt: 'test',
|
||||
samplerName: 'dpmpp_2m',
|
||||
scheduler: 'karras',
|
||||
},
|
||||
);
|
||||
|
||||
const nodes = workflow.workflow as any;
|
||||
const samplerNode = Object.values(nodes).find(
|
||||
(node: any) => node.class_type === 'KSampler',
|
||||
) as any;
|
||||
|
||||
expect(samplerNode.inputs.sampler_name).toBe('dpmpp_2m');
|
||||
expect(samplerNode.inputs.scheduler).toBe('karras');
|
||||
});
|
||||
});
|
||||
|
||||
describe('SD3.5 workflow specifics', () => {
|
||||
it('should use Triple CLIP loader when components available', async () => {
|
||||
mockModelResolver.getOptimalComponent
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.sd.clip) // clip_g
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.flux.clip) // clip_l
|
||||
.mockResolvedValueOnce(TEST_COMPONENTS.flux.t5); // t5xxl
|
||||
|
||||
const workflow = await service.buildWorkflow(
|
||||
'stable-diffusion-35',
|
||||
{ architecture: 'SD3', isSupported: true, variant: 'sd35' },
|
||||
TEST_MODELS.sd35,
|
||||
{ prompt: 'test' },
|
||||
);
|
||||
|
||||
const nodes = workflow.workflow as any;
|
||||
const tripleClipNode = Object.values(nodes).find(
|
||||
(node: any) => node.class_type === 'TripleCLIPLoader',
|
||||
);
|
||||
|
||||
expect(tripleClipNode).toBeDefined();
|
||||
});
|
||||
|
||||
it('should throw error when components not available', async () => {
|
||||
// Mock no components available
|
||||
mockModelResolver.getOptimalComponent.mockResolvedValue(undefined);
|
||||
|
||||
await expect(
|
||||
service.buildWorkflow(
|
||||
'stable-diffusion-35',
|
||||
{ architecture: 'SD3', isSupported: true, variant: 'sd35' },
|
||||
TEST_MODELS.sd35,
|
||||
{ prompt: 'test' },
|
||||
),
|
||||
).rejects.toThrow(WorkflowError);
|
||||
|
||||
await expect(
|
||||
service.buildWorkflow(
|
||||
'stable-diffusion-35',
|
||||
{ architecture: 'SD3', isSupported: true, variant: 'sd35' },
|
||||
TEST_MODELS.sd35,
|
||||
{ prompt: 'test' },
|
||||
),
|
||||
).rejects.toThrow('SD3.5 models require external CLIP/T5 encoder files');
|
||||
});
|
||||
|
||||
it('should use default shift parameter', async () => {
|
||||
// Mock components available for this test
|
||||
let callCount = 0;
|
||||
mockModelResolver.getOptimalComponent.mockImplementation(() => {
|
||||
callCount++;
|
||||
if (callCount === 1) return Promise.resolve('clip_l.safetensors');
|
||||
if (callCount === 2) return Promise.resolve('clip_g.safetensors');
|
||||
return Promise.resolve('t5xxl_fp16.safetensors');
|
||||
});
|
||||
|
||||
const workflow = await service.buildWorkflow(
|
||||
'stable-diffusion-35',
|
||||
{ architecture: 'SD3', isSupported: true, variant: 'sd35' },
|
||||
TEST_MODELS.sd35,
|
||||
{
|
||||
cfg: 4.5,
|
||||
prompt: 'test',
|
||||
steps: 28,
|
||||
},
|
||||
);
|
||||
|
||||
const nodes = workflow.workflow as any;
|
||||
const samplingNode = Object.values(nodes).find(
|
||||
(node: any) => node.class_type === 'ModelSamplingSD3',
|
||||
) as any;
|
||||
|
||||
expect(samplingNode.inputs.shift).toBe(3); // Uses default from WORKFLOW_DEFAULTS.SD3.SHIFT
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,140 @@
|
|||
import {
|
||||
fluxDevParamsSchema,
|
||||
fluxKontextDevParamsSchema,
|
||||
fluxSchnellParamsSchema,
|
||||
sd15T2iParamsSchema,
|
||||
sd35ParamsSchema,
|
||||
sdxlT2iParamsSchema,
|
||||
} from 'model-bank/comfyui';
|
||||
|
||||
export const parametersFixture = {
|
||||
models: {
|
||||
'flux-dev': {
|
||||
boundaries: {
|
||||
max: {
|
||||
cfg: fluxDevParamsSchema.cfg!.max,
|
||||
steps: fluxDevParamsSchema.steps!.max,
|
||||
},
|
||||
min: {
|
||||
cfg: fluxDevParamsSchema.cfg!.min,
|
||||
steps: fluxDevParamsSchema.steps!.min,
|
||||
},
|
||||
},
|
||||
defaults: {
|
||||
cfg: fluxDevParamsSchema.cfg!.default,
|
||||
samplerName: fluxDevParamsSchema.samplerName!.default,
|
||||
scheduler: fluxDevParamsSchema.scheduler!.default,
|
||||
steps: fluxDevParamsSchema.steps!.default,
|
||||
},
|
||||
schema: fluxDevParamsSchema,
|
||||
},
|
||||
'flux-kontext': {
|
||||
boundaries: {
|
||||
max: {
|
||||
cfg: fluxKontextDevParamsSchema.cfg!.max,
|
||||
steps: fluxKontextDevParamsSchema.steps!.max,
|
||||
},
|
||||
min: {
|
||||
cfg: fluxKontextDevParamsSchema.cfg!.min,
|
||||
steps: fluxKontextDevParamsSchema.steps!.min,
|
||||
},
|
||||
},
|
||||
defaults: {
|
||||
cfg: fluxKontextDevParamsSchema.cfg!.default,
|
||||
steps: fluxKontextDevParamsSchema.steps!.default,
|
||||
strength: fluxKontextDevParamsSchema.strength!.default,
|
||||
},
|
||||
schema: fluxKontextDevParamsSchema,
|
||||
},
|
||||
'flux-schnell': {
|
||||
boundaries: {
|
||||
max: {
|
||||
cfg: 1,
|
||||
steps: fluxSchnellParamsSchema.steps!.max,
|
||||
},
|
||||
min: {
|
||||
cfg: 1,
|
||||
steps: fluxSchnellParamsSchema.steps!.min,
|
||||
},
|
||||
},
|
||||
defaults: {
|
||||
cfg: 1,
|
||||
samplerName: fluxSchnellParamsSchema.samplerName!.default,
|
||||
|
||||
scheduler: fluxSchnellParamsSchema.scheduler!.default,
|
||||
// Schnell fixed at 1
|
||||
steps: fluxSchnellParamsSchema.steps!.default,
|
||||
},
|
||||
schema: fluxSchnellParamsSchema,
|
||||
},
|
||||
'sd15': {
|
||||
boundaries: {
|
||||
max: {
|
||||
cfg: sd15T2iParamsSchema.cfg!.max,
|
||||
steps: sd15T2iParamsSchema.steps!.max,
|
||||
},
|
||||
min: {
|
||||
cfg: sd15T2iParamsSchema.cfg!.min,
|
||||
steps: sd15T2iParamsSchema.steps!.min,
|
||||
},
|
||||
},
|
||||
defaults: {
|
||||
cfg: sd15T2iParamsSchema.cfg!.default,
|
||||
samplerName: sd15T2iParamsSchema.samplerName!.default,
|
||||
scheduler: sd15T2iParamsSchema.scheduler!.default,
|
||||
steps: sd15T2iParamsSchema.steps!.default,
|
||||
},
|
||||
schema: sd15T2iParamsSchema,
|
||||
},
|
||||
'sd35': {
|
||||
boundaries: {
|
||||
max: {
|
||||
cfg: sd35ParamsSchema.cfg!.max,
|
||||
steps: sd35ParamsSchema.steps!.max,
|
||||
},
|
||||
min: {
|
||||
cfg: sd35ParamsSchema.cfg!.min,
|
||||
steps: sd35ParamsSchema.steps!.min,
|
||||
},
|
||||
},
|
||||
defaults: {
|
||||
cfg: sd35ParamsSchema.cfg!.default,
|
||||
samplerName: sd35ParamsSchema.samplerName!.default,
|
||||
scheduler: sd35ParamsSchema.scheduler!.default,
|
||||
steps: sd35ParamsSchema.steps!.default,
|
||||
},
|
||||
schema: sd35ParamsSchema,
|
||||
},
|
||||
'sdxl': {
|
||||
boundaries: {
|
||||
max: {
|
||||
cfg: sdxlT2iParamsSchema.cfg!.max,
|
||||
steps: sdxlT2iParamsSchema.steps!.max,
|
||||
},
|
||||
min: {
|
||||
cfg: sdxlT2iParamsSchema.cfg!.min,
|
||||
steps: sdxlT2iParamsSchema.steps!.min,
|
||||
},
|
||||
},
|
||||
defaults: {
|
||||
cfg: sdxlT2iParamsSchema.cfg!.default,
|
||||
samplerName: sdxlT2iParamsSchema.samplerName!.default,
|
||||
scheduler: sdxlT2iParamsSchema.scheduler!.default,
|
||||
steps: sdxlT2iParamsSchema.steps!.default,
|
||||
},
|
||||
schema: sdxlT2iParamsSchema,
|
||||
},
|
||||
},
|
||||
|
||||
transformations: {
|
||||
aspectRatio: [
|
||||
{ expected: { height: 576, width: 1024 }, input: '16:9' },
|
||||
{ expected: { height: 1024, width: 1024 }, input: '1:1' },
|
||||
{ expected: { height: 1024, width: 576 }, input: '9:16' },
|
||||
],
|
||||
imageUrl: [
|
||||
{ expectedParam: 'imageUrl', input: 'test.png', mode: 'img2img' },
|
||||
{ expectedParam: undefined, input: undefined, mode: 'txt2img' },
|
||||
],
|
||||
},
|
||||
};
|
||||
|
|
@ -0,0 +1,97 @@
|
|||
// 可扩展的支持配置接口
|
||||
export interface SupportedConfig {
|
||||
extensions?: Record<string, any>;
|
||||
models: Record<string, string[]>;
|
||||
workflows: string[];
|
||||
}
|
||||
|
||||
// 基础支持配置
|
||||
const baseConfig: SupportedConfig = {
|
||||
extensions: {},
|
||||
models: {
|
||||
flux: ['flux-dev', 'flux-schnell', 'flux-kontext', 'flux-krea'],
|
||||
sd: ['sd15', 'sdxl', 'sd35'],
|
||||
},
|
||||
workflows: ['flux-dev', 'flux-schnell', 'flux-kontext', 'flux-krea', 'simple-sd', 'sd35'],
|
||||
};
|
||||
|
||||
// 动态配置合并函数
|
||||
function mergeConfig(base: SupportedConfig, custom?: Partial<SupportedConfig>): SupportedConfig {
|
||||
if (!custom) return base;
|
||||
|
||||
return {
|
||||
// 去重
|
||||
extensions: {
|
||||
...base.extensions,
|
||||
...custom.extensions,
|
||||
},
|
||||
|
||||
models: {
|
||||
...base.models,
|
||||
...custom.models,
|
||||
},
|
||||
workflows: [...base.workflows, ...(custom.workflows || [])].filter(
|
||||
(workflow, index, array) => array.indexOf(workflow) === index,
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
// 可扩展的 fixture 对象
|
||||
export const supportedFixture = {
|
||||
|
||||
// 扩展工具函数
|
||||
addCustomModels: (modelType: string, models: string[]) => {
|
||||
baseConfig.models[modelType] = [...(baseConfig.models[modelType] || []), ...models].filter(
|
||||
(model, index, array) => array.indexOf(model) === index,
|
||||
);
|
||||
},
|
||||
|
||||
|
||||
|
||||
|
||||
addCustomWorkflows: (workflows: string[]) => {
|
||||
baseConfig.workflows = [...baseConfig.workflows, ...workflows].filter(
|
||||
(workflow, index, array) => array.indexOf(workflow) === index,
|
||||
);
|
||||
},
|
||||
|
||||
|
||||
|
||||
|
||||
// 获取当前配置（支持自定义扩展）
|
||||
getConfig: (customConfig?: Partial<SupportedConfig>): SupportedConfig => {
|
||||
return mergeConfig(baseConfig, customConfig);
|
||||
},
|
||||
|
||||
|
||||
// 验证帮助函数
|
||||
isSupported: (model: string, customConfig?: Partial<SupportedConfig>) => {
|
||||
const config = mergeConfig(baseConfig, customConfig);
|
||||
const allModels = Object.values(config.models).flat();
|
||||
return allModels.includes(model);
|
||||
},
|
||||
|
||||
|
||||
// 向后兼容的属性（保持现有测试不受影响）
|
||||
models: baseConfig.models,
|
||||
|
||||
// 重置为基础配置（用于测试隔离）
|
||||
reset: () => {
|
||||
baseConfig.models = {
|
||||
flux: ['flux-dev', 'flux-schnell', 'flux-kontext', 'flux-krea'],
|
||||
sd: ['sd15', 'sdxl', 'sd35'],
|
||||
};
|
||||
baseConfig.workflows = [
|
||||
'flux-dev',
|
||||
'flux-schnell',
|
||||
'flux-kontext',
|
||||
'flux-krea',
|
||||
'simple-sd',
|
||||
'sd35',
|
||||
];
|
||||
baseConfig.extensions = {};
|
||||
},
|
||||
|
||||
|
||||
workflows: baseConfig.workflows,
|
||||
};
|
||||
64
src/server/services/comfyui/__tests__/fixtures/testModels.ts
Normal file
64
src/server/services/comfyui/__tests__/fixtures/testModels.ts
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
/**
|
||||
* Real model names from registry for testing
|
||||
* Using actual registered models instead of fake names
|
||||
*/
|
||||
|
||||
// Real FLUX models from registry
|
||||
export const TEST_FLUX_MODELS = {
|
||||
DEV: 'flux1-dev.safetensors',
|
||||
KONTEXT: 'flux1-kontext-dev.safetensors',
|
||||
KREA: 'flux1-krea-dev.safetensors',
|
||||
SCHNELL: 'flux1-schnell.safetensors',
|
||||
} as const;
|
||||
|
||||
// Real SD3.5 models from registry
|
||||
export const TEST_SD35_MODELS = {
|
||||
LARGE: 'sd3.5_large.safetensors',
|
||||
LARGE_FP8: 'sd3.5_large_fp8_scaled.safetensors',
|
||||
LARGE_TURBO: 'sd3.5_large_turbo.safetensors',
|
||||
MEDIUM: 'sd3.5_medium.safetensors',
|
||||
} as const;
|
||||
|
||||
// Real SDXL models from registry
|
||||
export const TEST_SDXL_MODELS = {
|
||||
BASE: 'sd_xl_base_1.0.safetensors',
|
||||
TURBO: 'sd_xl_turbo_1.0_fp16.safetensors',
|
||||
} as const;
|
||||
|
||||
// Custom SD model
|
||||
export const TEST_CUSTOM_SD = 'custom_sd_lobe.safetensors';
|
||||
|
||||
// Real component names from system components
|
||||
export const TEST_COMPONENTS = {
|
||||
FLUX: {
|
||||
CLIP_L: 'clip_l.safetensors',
|
||||
T5: 't5xxl_fp16.safetensors',
|
||||
VAE: 'ae.safetensors',
|
||||
},
|
||||
SD: {
|
||||
CLIP_G: 'clip_g.safetensors',
|
||||
CLIP_L: 'clip_l.safetensors',
|
||||
VAE: 'sdxl_vae_fp16fix.safetensors',
|
||||
},
|
||||
} as const;
|
||||
|
||||
// Common test model sets for different scenarios
|
||||
export const TEST_MODEL_SETS = {
|
||||
// Models that don't exist (for error testing)
|
||||
NON_EXISTENT: [
|
||||
'nonexistent-model.safetensors',
|
||||
'unknown-model.safetensors',
|
||||
'fake-model.safetensors',
|
||||
],
|
||||
|
||||
// Models that should exist in registry
|
||||
REGISTERED: [
|
||||
TEST_FLUX_MODELS.DEV,
|
||||
TEST_FLUX_MODELS.SCHNELL,
|
||||
TEST_SD35_MODELS.LARGE,
|
||||
TEST_SDXL_MODELS.BASE,
|
||||
],
|
||||
} as const;
|
||||
|
||||
// Default test model for general use
|
||||
export const DEFAULT_TEST_MODEL = TEST_FLUX_MODELS.DEV;
|
||||
98
src/server/services/comfyui/__tests__/helpers/mockContext.ts
Normal file
98
src/server/services/comfyui/__tests__/helpers/mockContext.ts
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
import { vi } from 'vitest';
|
||||
|
||||
import {
|
||||
TEST_COMPONENTS,
|
||||
TEST_FLUX_MODELS,
|
||||
} from '@/server/services/comfyui/__tests__/fixtures/testModels';
|
||||
import type { WorkflowContext } from '@/server/services/comfyui/core/workflowBuilderService';
|
||||
|
||||
/**
|
||||
* Create a mock WorkflowContext for testing
|
||||
* 创建测试用的模拟 WorkflowContext
|
||||
*/
|
||||
export function createMockContext(): WorkflowContext {
|
||||
return {
|
||||
clientService: {
|
||||
executeWorkflow: vi.fn().mockResolvedValue({ images: { images: [] } }),
|
||||
// New SDK wrapper methods
|
||||
getCheckpoints: vi.fn().mockResolvedValue([TEST_FLUX_MODELS.DEV, TEST_FLUX_MODELS.SCHNELL]),
|
||||
getLoras: vi.fn().mockResolvedValue(['lora1.safetensors', 'lora2.safetensors']),
|
||||
getNodeDefs: vi.fn().mockResolvedValue({
|
||||
CLIPLoader: {
|
||||
input: {
|
||||
required: {
|
||||
clip_name: [['clip_l.safetensors', 'clip_g.safetensors']],
|
||||
},
|
||||
},
|
||||
},
|
||||
DualCLIPLoader: {
|
||||
input: {
|
||||
required: {
|
||||
clip_name1: [['t5-v1_1-xxl-encoder.safetensors', 't5xxl_fp16.safetensors']],
|
||||
},
|
||||
},
|
||||
},
|
||||
VAELoader: {
|
||||
input: {
|
||||
required: {
|
||||
vae_name: [
|
||||
[
|
||||
'ae.safetensors',
|
||||
'sdxl_vae_fp16fix.safetensors',
|
||||
'vae-ft-mse-840000-ema-pruned.safetensors',
|
||||
],
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
getObjectInfo: vi.fn().mockResolvedValue({}),
|
||||
getPathImage: vi.fn().mockReturnValue('http://example.com/image.png'),
|
||||
getSamplerInfo: vi.fn().mockResolvedValue({
|
||||
sampler: ['euler', 'ddim', 'dpm_2'],
|
||||
scheduler: ['normal', 'karras', 'exponential'],
|
||||
}),
|
||||
validateConnection: vi.fn().mockResolvedValue(undefined),
|
||||
},
|
||||
modelResolverService: {
|
||||
// 新的服务层方法
|
||||
getOptimalComponent: vi.fn().mockImplementation((type: string, modelFamily: string) => {
|
||||
// 根据不同的组件类型和模型家族返回相应的默认值
|
||||
if (type === 't5') {
|
||||
return Promise.resolve(TEST_COMPONENTS.FLUX.T5);
|
||||
}
|
||||
if (type === 'vae') {
|
||||
if (modelFamily === 'FLUX') {
|
||||
return Promise.resolve(TEST_COMPONENTS.FLUX.VAE);
|
||||
}
|
||||
return Promise.resolve(TEST_COMPONENTS.SD.VAE);
|
||||
}
|
||||
if (type === 'clip') {
|
||||
if (modelFamily === 'FLUX') {
|
||||
return Promise.resolve(TEST_COMPONENTS.FLUX.CLIP_L);
|
||||
}
|
||||
return Promise.resolve(TEST_COMPONENTS.SD.CLIP_G);
|
||||
}
|
||||
return Promise.resolve(null);
|
||||
}),
|
||||
|
||||
// 保留旧方法以兼容
|
||||
selectComponents: vi.fn().mockResolvedValue({
|
||||
clip: [TEST_COMPONENTS.FLUX.T5, TEST_COMPONENTS.FLUX.CLIP_L],
|
||||
t5: TEST_COMPONENTS.FLUX.T5,
|
||||
vae: TEST_COMPONENTS.FLUX.VAE,
|
||||
}),
|
||||
|
||||
validateModel: vi.fn().mockResolvedValue({
|
||||
actualFileName: TEST_FLUX_MODELS.DEV,
|
||||
exists: true,
|
||||
}),
|
||||
} as any,
|
||||
} as unknown as WorkflowContext;
|
||||
}
|
||||
|
||||
/**
|
||||
* Default mock context instance
|
||||
* 默认的模拟上下文实例
|
||||
*/
|
||||
export const mockContext = createMockContext();
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
/**
|
||||
* Real configuration data helper for tests
|
||||
* Uses actual data from configuration files instead of mock data
|
||||
*/
|
||||
import { MODEL_REGISTRY } from '@/server/services/comfyui/config/modelRegistry';
|
||||
import { SYSTEM_COMPONENTS } from '@/server/services/comfyui/config/systemComponents';
|
||||
import { getModelConfig } from '@/server/services/comfyui/utils/staticModelLookup';
|
||||
|
||||
// Export real model entries for tests
|
||||
export const REAL_MODEL_ENTRIES = Object.entries(MODEL_REGISTRY);
|
||||
|
||||
// Get real FLUX models
|
||||
export const REAL_FLUX_MODELS = REAL_MODEL_ENTRIES.filter(
|
||||
([, config]) => config.modelFamily === 'FLUX',
|
||||
).map(([fileName]) => fileName);
|
||||
|
||||
// Get real SD models
|
||||
export const REAL_SD_MODELS = REAL_MODEL_ENTRIES.filter(([, config]) =>
|
||||
['SD1', 'SDXL', 'SD3'].includes(config.modelFamily),
|
||||
).map(([fileName]) => fileName);
|
||||
|
||||
// Get real system components
|
||||
export const REAL_COMPONENT_ENTRIES = Object.entries(SYSTEM_COMPONENTS);
|
||||
|
||||
// Get real FLUX components
|
||||
export const REAL_FLUX_COMPONENTS = {
|
||||
clip: REAL_COMPONENT_ENTRIES.filter(
|
||||
([, config]) => config.type === 'clip' && config.modelFamily === 'FLUX',
|
||||
).map(([name]) => name),
|
||||
t5: REAL_COMPONENT_ENTRIES.filter(
|
||||
([, config]) => config.type === 't5' && config.modelFamily === 'FLUX',
|
||||
).map(([name]) => name),
|
||||
vae: REAL_COMPONENT_ENTRIES.filter(
|
||||
([, config]) => config.type === 'vae' && config.modelFamily === 'FLUX',
|
||||
).map(([name]) => name),
|
||||
};
|
||||
|
||||
// Get real SD components
|
||||
export const REAL_SD_COMPONENTS = {
|
||||
clip: REAL_COMPONENT_ENTRIES.filter(
|
||||
([, config]) => config.type === 'clip' && ['SD1', 'SDXL', 'SD3'].includes(config.modelFamily),
|
||||
).map(([name]) => name),
|
||||
vae: REAL_COMPONENT_ENTRIES.filter(
|
||||
([, config]) => config.type === 'vae' && ['SD1', 'SDXL', 'SD3'].includes(config.modelFamily),
|
||||
).map(([name]) => name),
|
||||
};
|
||||
|
||||
// Export real workflow defaults
|
||||
|
||||
// Export real component node mappings
|
||||
|
||||
// Helper to get real model config
|
||||
export const getRealModelConfig = getModelConfig;
|
||||
|
||||
// Test data selections (using real data)
|
||||
export const TEST_MODELS = {
|
||||
flux: REAL_FLUX_MODELS[0] || 'flux1-dev.safetensors', // Use first real FLUX model
|
||||
sd35:
|
||||
REAL_SD_MODELS.find((m) => getRealModelConfig(m)?.modelFamily === 'SD3') ||
|
||||
'sd3.5_large.safetensors',
|
||||
sdxl:
|
||||
REAL_SD_MODELS.find((m) => getRealModelConfig(m)?.modelFamily === 'SDXL') ||
|
||||
'sdxl_base.safetensors',
|
||||
};
|
||||
|
||||
export const TEST_COMPONENTS = {
|
||||
flux: {
|
||||
clip: REAL_FLUX_COMPONENTS.clip[0] || 'clip_l.safetensors',
|
||||
t5: REAL_FLUX_COMPONENTS.t5[0] || 't5xxl_fp16.safetensors',
|
||||
vae: REAL_FLUX_COMPONENTS.vae[0] || 'ae.safetensors',
|
||||
},
|
||||
sd: {
|
||||
clip: REAL_SD_COMPONENTS.clip[0] || 'clip_g.safetensors',
|
||||
vae: REAL_SD_COMPONENTS.vae[0] || 'sdxl_vae_fp16fix.safetensors',
|
||||
},
|
||||
};
|
||||
export {
|
||||
COMPONENT_NODE_MAPPINGS as REAL_COMPONENT_MAPPINGS,
|
||||
WORKFLOW_DEFAULTS as REAL_WORKFLOW_DEFAULTS,
|
||||
} from '@/server/services/comfyui/config/constants';
|
||||
219
src/server/services/comfyui/__tests__/helpers/testSetup.ts
Normal file
219
src/server/services/comfyui/__tests__/helpers/testSetup.ts
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
// @vitest-environment node
|
||||
import { vi } from 'vitest';
|
||||
|
||||
// Common mock setup for ComfyUI tests
|
||||
export function setupComfyUIMocks() {
|
||||
// Mock the ComfyUI SDK - keep it simple, tests will override
|
||||
vi.mock('@saintno/comfyui-sdk', () => ({
|
||||
CallWrapper: vi.fn(),
|
||||
ComfyApi: vi.fn(),
|
||||
PromptBuilder: vi.fn(),
|
||||
}));
|
||||
|
||||
// Mock the ModelResolver
|
||||
vi.mock('../utils/modelResolver', () => ({
|
||||
ModelResolver: vi.fn(),
|
||||
getAllModels: vi.fn().mockReturnValue(['flux-schnell.safetensors', 'flux-dev.safetensors']),
|
||||
isValidModel: vi.fn().mockReturnValue(true),
|
||||
resolveModel: vi.fn().mockImplementation(() => {
|
||||
return {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default' as const,
|
||||
variant: 'dev' as const,
|
||||
};
|
||||
}),
|
||||
resolveModelStrict: vi.fn().mockImplementation(() => {
|
||||
return {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default' as const,
|
||||
variant: 'dev' as const,
|
||||
};
|
||||
}),
|
||||
}));
|
||||
|
||||
// Mock fetch globally
|
||||
global.fetch = vi.fn();
|
||||
|
||||
// Mock console.error to avoid polluting test output
|
||||
vi.spyOn(console, 'error').mockImplementation(() => {});
|
||||
|
||||
// Mock WorkflowDetector
|
||||
vi.mock('../utils/workflowDetector', () => ({
|
||||
WorkflowDetector: {
|
||||
detectModelType: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock processModels utility
|
||||
vi.mock('../utils/modelParse', () => ({
|
||||
MODEL_LIST_CONFIGS: {
|
||||
comfyui: {
|
||||
id: 'comfyui',
|
||||
modelList: [],
|
||||
},
|
||||
},
|
||||
detectModelProvider: vi.fn().mockImplementation((modelId: string) => {
|
||||
if (modelId.includes('claude')) return 'anthropic';
|
||||
if (modelId.includes('gpt')) return 'openai';
|
||||
if (modelId.includes('gemini')) return 'google';
|
||||
return 'unknown';
|
||||
}),
|
||||
processModelList: vi.fn(),
|
||||
}));
|
||||
}
|
||||
|
||||
export function createMockComfyApi() {
|
||||
return {
|
||||
fetchApi: vi.fn().mockResolvedValue({
|
||||
CheckpointLoaderSimple: {
|
||||
input: {
|
||||
required: {
|
||||
ckpt_name: [['flux-schnell.safetensors', 'flux-dev.safetensors', 'sd15-base.ckpt']],
|
||||
},
|
||||
},
|
||||
},
|
||||
}),
|
||||
getPathImage: vi.fn().mockReturnValue('http://localhost:8000/view?filename=test.png'),
|
||||
init: vi.fn(),
|
||||
waitForReady: vi.fn().mockResolvedValue(undefined),
|
||||
};
|
||||
}
|
||||
|
||||
export function createMockCallWrapper() {
|
||||
return {
|
||||
onFailed: vi.fn().mockReturnThis(),
|
||||
onFinished: vi.fn().mockReturnThis(),
|
||||
onProgress: vi.fn().mockReturnThis(),
|
||||
run: vi.fn().mockReturnThis(),
|
||||
};
|
||||
}
|
||||
|
||||
export function createMockPromptBuilder() {
|
||||
return {
|
||||
input: vi.fn().mockReturnThis(),
|
||||
prompt: {},
|
||||
setInputNode: vi.fn().mockReturnThis(),
|
||||
setOutputNode: vi.fn().mockReturnThis(),
|
||||
} as any;
|
||||
}
|
||||
|
||||
export function createMockModelResolver() {
|
||||
return {
|
||||
getAvailableModelFiles: vi
|
||||
.fn()
|
||||
.mockResolvedValue(['flux-schnell.safetensors', 'flux-dev.safetensors', 'sd15-base.ckpt']),
|
||||
resolveModelFileName: vi.fn().mockImplementation((modelId: string) => {
|
||||
if (
|
||||
modelId.includes('non-existent') ||
|
||||
modelId.includes('unknown') ||
|
||||
modelId.includes('non-verified')
|
||||
) {
|
||||
return Promise.reject(new Error(`Model not found: ${modelId}`));
|
||||
}
|
||||
const fileName = modelId.split('/').pop() || modelId;
|
||||
return Promise.resolve(fileName + '.safetensors');
|
||||
}),
|
||||
transformModelFilesToList: vi.fn().mockReturnValue([]),
|
||||
validateModel: vi.fn().mockImplementation((modelId: string) => {
|
||||
if (
|
||||
modelId.includes('non-existent') ||
|
||||
modelId.includes('unknown') ||
|
||||
modelId.includes('non-verified')
|
||||
) {
|
||||
return Promise.resolve({ exists: false });
|
||||
}
|
||||
const fileName = modelId.split('/').pop() || modelId;
|
||||
return Promise.resolve({ actualFileName: fileName + '.safetensors', exists: true });
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
// Mock workflow builders
|
||||
export function setupWorkflowMocks() {
|
||||
const createMockBuilder = () => ({
|
||||
input: vi.fn().mockReturnThis(),
|
||||
prompt: {
|
||||
'1': {
|
||||
_meta: { title: 'Checkpoint Loader' },
|
||||
class_type: 'CheckpointLoaderSimple',
|
||||
inputs: { ckpt_name: 'test.safetensors' },
|
||||
},
|
||||
},
|
||||
setInputNode: vi.fn().mockReturnThis(),
|
||||
setOutputNode: vi.fn().mockReturnThis(),
|
||||
});
|
||||
|
||||
// Mock the workflows index
|
||||
vi.mock('../../workflows', () => ({
|
||||
buildFluxDevWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
buildFluxKontextWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
buildFluxKreaWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
buildFluxSchnellWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
buildSD35NoClipWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
buildSD35Workflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
}));
|
||||
|
||||
// Mock individual workflow builders
|
||||
vi.mock('../../workflows/flux-schnell', () => ({
|
||||
buildFluxSchnellWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
}));
|
||||
|
||||
vi.mock('../../workflows/flux-dev', () => ({
|
||||
buildFluxDevWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
}));
|
||||
|
||||
vi.mock('../../workflows/flux-kontext', () => ({
|
||||
buildFluxKontextWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
}));
|
||||
|
||||
vi.mock('../../workflows/sd35', () => ({
|
||||
buildSD35Workflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
}));
|
||||
|
||||
vi.mock('../../workflows/simple-sd', () => ({
|
||||
buildSimpleSDWorkflow: vi.fn().mockImplementation(() => createMockBuilder()),
|
||||
}));
|
||||
|
||||
// Mock WorkflowRouter
|
||||
vi.mock('../utils/workflowRouter', () => {
|
||||
class WorkflowRoutingError extends Error {
|
||||
constructor(message?: string) {
|
||||
super(message);
|
||||
this.name = 'WorkflowRoutingError';
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
WorkflowRouter: {
|
||||
getExactlySupportedModels: () => ['comfyui/flux-dev', 'comfyui/flux-schnell'],
|
||||
getSupportedFluxVariants: () => ['dev', 'schnell', 'kontext', 'krea'],
|
||||
routeWorkflow: () => createMockBuilder(),
|
||||
},
|
||||
WorkflowRoutingError,
|
||||
};
|
||||
});
|
||||
|
||||
// Mock systemComponents
|
||||
vi.mock('../../config/systemComponents', () => ({
|
||||
getAllComponentsWithNames: vi.fn().mockImplementation((options: any) => {
|
||||
if (options?.type === 'clip') {
|
||||
return [
|
||||
{ config: { priority: 1 }, name: 'clip_l.safetensors' },
|
||||
{ config: { priority: 2 }, name: 'clip_g.safetensors' },
|
||||
];
|
||||
}
|
||||
if (options?.type === 't5') {
|
||||
return [{ config: { priority: 1 }, name: 't5xxl_fp16.safetensors' }];
|
||||
}
|
||||
return [];
|
||||
}),
|
||||
getOptimalComponent: vi.fn().mockImplementation((type: string) => {
|
||||
if (type === 't5') return 't5xxl_fp16.safetensors';
|
||||
if (type === 'vae') return 'ae.safetensors';
|
||||
if (type === 'clip') return 'clip_l.safetensors';
|
||||
return 'default.safetensors';
|
||||
}),
|
||||
}));
|
||||
}
|
||||
|
|
@ -0,0 +1,138 @@
|
|||
// @vitest-environment node
|
||||
import { PromptBuilder } from '@saintno/comfyui-sdk';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { parametersFixture } from '@/server/services/comfyui/__tests__/fixtures/parameters.fixture';
|
||||
import { supportedFixture } from '@/server/services/comfyui/__tests__/fixtures/supported.fixture';
|
||||
import { mockContext } from '@/server/services/comfyui/__tests__/helpers/mockContext';
|
||||
import { setupAllMocks } from '@/server/services/comfyui/__tests__/setup/unifiedMocks';
|
||||
// Import workflow builders
|
||||
import { buildFluxDevWorkflow } from '@/server/services/comfyui/workflows/flux-dev';
|
||||
import { buildFluxKontextWorkflow } from '@/server/services/comfyui/workflows/flux-kontext';
|
||||
import { buildFluxSchnellWorkflow } from '@/server/services/comfyui/workflows/flux-schnell';
|
||||
import { buildSD35Workflow } from '@/server/services/comfyui/workflows/sd35';
|
||||
import { buildSimpleSDWorkflow } from '@/server/services/comfyui/workflows/simple-sd';
|
||||
|
||||
describe('Parameter Mapping - Core Business Logic', () => {
|
||||
const { models } = parametersFixture;
|
||||
let inputCalls: Map<string, any>;
|
||||
|
||||
beforeEach(() => {
|
||||
const mocks = setupAllMocks();
|
||||
inputCalls = mocks.inputCalls;
|
||||
});
|
||||
|
||||
// Workflow builder mapping
|
||||
const workflowBuilders = {
|
||||
'flux-dev': buildFluxDevWorkflow,
|
||||
'flux-schnell': buildFluxSchnellWorkflow,
|
||||
'flux-kontext': buildFluxKontextWorkflow,
|
||||
'sd15': buildSimpleSDWorkflow,
|
||||
'sdxl': buildSimpleSDWorkflow,
|
||||
'sd35': buildSD35Workflow,
|
||||
};
|
||||
|
||||
// Parameterized tests for all supported models
|
||||
describe.each(
|
||||
Object.entries(models).filter(
|
||||
([name]) => workflowBuilders[name as keyof typeof workflowBuilders],
|
||||
),
|
||||
)('%s parameter mapping', (modelName, modelConfig) => {
|
||||
const builder = workflowBuilders[modelName as keyof typeof workflowBuilders];
|
||||
|
||||
it('should map schema parameters to workflow', async () => {
|
||||
const params: any = {
|
||||
prompt: 'test prompt',
|
||||
...modelConfig.defaults,
|
||||
};
|
||||
|
||||
// Special parameter handling
|
||||
if (modelName === 'flux-kontext') {
|
||||
params.imageUrl = 'test.png';
|
||||
}
|
||||
if (modelName.startsWith('sd')) {
|
||||
params.width = 512;
|
||||
params.height = 512;
|
||||
} else if (modelName.startsWith('flux') && modelName !== 'flux-kontext') {
|
||||
params.width = 1024;
|
||||
params.height = 1024;
|
||||
}
|
||||
|
||||
// Build workflow
|
||||
const workflow = await builder(`${modelName}.safetensors`, params, mockContext);
|
||||
|
||||
// Verify workflow is built successfully
|
||||
expect(workflow).toBeDefined();
|
||||
|
||||
// Verify PromptBuilder was used for workflow construction
|
||||
expect(workflow).toBeDefined();
|
||||
expect(typeof workflow).toBe('object');
|
||||
});
|
||||
|
||||
it('should handle boundary values', async () => {
|
||||
const { min, max } = modelConfig.boundaries!;
|
||||
|
||||
const baseParams = {
|
||||
prompt: 'test prompt',
|
||||
width: 512,
|
||||
height: 512,
|
||||
};
|
||||
|
||||
// Minimum values should not error
|
||||
const minResult = await builder(
|
||||
`${modelName}.safetensors`,
|
||||
{ ...baseParams, ...modelConfig.defaults, ...min },
|
||||
mockContext,
|
||||
);
|
||||
expect(minResult).toBeDefined();
|
||||
|
||||
// Maximum values should not error
|
||||
const maxResult = await builder(
|
||||
`${modelName}.safetensors`,
|
||||
{ ...baseParams, ...modelConfig.defaults, ...max },
|
||||
mockContext,
|
||||
);
|
||||
expect(maxResult).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
// Parameter transformation tests
|
||||
describe('Parameter Transformations', () => {
|
||||
it.each(parametersFixture.transformations.aspectRatio)(
|
||||
'should transform aspectRatio $input to width/height',
|
||||
async ({ input, expected }) => {
|
||||
const params = {
|
||||
prompt: 'test prompt',
|
||||
...models['flux-dev'].defaults,
|
||||
aspectRatio: input,
|
||||
};
|
||||
|
||||
const workflow = await buildFluxDevWorkflow('flux-dev.safetensors', params, mockContext);
|
||||
|
||||
// Verify workflow builds successfully
|
||||
expect(workflow).toBeDefined();
|
||||
|
||||
// aspectRatio should be processed (verified through successful workflow build)
|
||||
const workflowStr = JSON.stringify(workflow.workflow || workflow);
|
||||
expect(workflowStr).not.toContain('aspectRatio');
|
||||
},
|
||||
);
|
||||
|
||||
it('should handle imageUrl for img2img mode', async () => {
|
||||
const params = {
|
||||
prompt: 'test prompt',
|
||||
imageUrl: 'test-image.png',
|
||||
strength: 0.8,
|
||||
};
|
||||
|
||||
const workflow = await buildFluxKontextWorkflow(
|
||||
'flux-kontext.safetensors',
|
||||
params,
|
||||
mockContext,
|
||||
);
|
||||
|
||||
// Verify workflow builds successfully (img2img parameters processed)
|
||||
expect(workflow).toBeDefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
// @vitest-environment node
|
||||
import { beforeEach, describe, expect, it } from 'vitest';
|
||||
|
||||
import { parametersFixture } from '@/server/services/comfyui/__tests__/fixtures/parameters.fixture';
|
||||
import { mockContext } from '@/server/services/comfyui/__tests__/helpers/mockContext';
|
||||
import { setupAllMocks } from '@/server/services/comfyui/__tests__/setup/unifiedMocks';
|
||||
// Import transformation utilities
|
||||
import { buildFluxDevWorkflow } from '@/server/services/comfyui/workflows/flux-dev';
|
||||
import { buildFluxKontextWorkflow } from '@/server/services/comfyui/workflows/flux-kontext';
|
||||
|
||||
describe('Parameter Transformation Tests', () => {
|
||||
let inputCalls: Map<string, any>;
|
||||
|
||||
beforeEach(() => {
|
||||
const mocks = setupAllMocks();
|
||||
inputCalls = mocks.inputCalls;
|
||||
});
|
||||
|
||||
describe('AspectRatio Transformation', () => {
|
||||
it.each(parametersFixture.transformations.aspectRatio)(
|
||||
'should handle aspectRatio $input correctly',
|
||||
async ({ input }) => {
|
||||
const params = {
|
||||
prompt: 'test prompt',
|
||||
...parametersFixture.models['flux-dev'].defaults,
|
||||
aspectRatio: input,
|
||||
};
|
||||
|
||||
// Should successfully build workflow
|
||||
await expect(
|
||||
buildFluxDevWorkflow('flux-dev.safetensors', params, mockContext),
|
||||
).resolves.toBeDefined();
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe('Image URL Processing', () => {
|
||||
it('should process imageUrl for img2img workflows', async () => {
|
||||
const params = {
|
||||
prompt: 'test prompt',
|
||||
imageUrl: 'test-image.png',
|
||||
strength: 0.8,
|
||||
};
|
||||
|
||||
// Kontext supports img2img
|
||||
await expect(
|
||||
buildFluxKontextWorkflow('flux-kontext.safetensors', params, mockContext),
|
||||
).resolves.toBeDefined();
|
||||
});
|
||||
|
||||
it('should handle missing imageUrl gracefully', async () => {
|
||||
const params = {
|
||||
prompt: 'test prompt',
|
||||
// No imageUrl provided
|
||||
};
|
||||
|
||||
// Should build normally (may fallback to txt2img mode)
|
||||
await expect(
|
||||
buildFluxKontextWorkflow('flux-kontext.safetensors', params, mockContext),
|
||||
).resolves.toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Parameter Validation', () => {
|
||||
it('should handle valid parameter ranges', () => {
|
||||
Object.entries(parametersFixture.models).forEach(([modelName, config]) => {
|
||||
const { min, max } = (config as any).boundaries!;
|
||||
|
||||
// Minimum values within range
|
||||
Object.entries(min).forEach(([key, value]) => {
|
||||
expect(typeof value).toBe('number');
|
||||
expect(value).toBeGreaterThanOrEqual(0);
|
||||
});
|
||||
|
||||
// Maximum values within reasonable range
|
||||
Object.entries(max).forEach(([key, value]) => {
|
||||
expect(typeof value).toBe('number');
|
||||
if (key === 'cfg') {
|
||||
expect(value).toBeLessThanOrEqual(20);
|
||||
}
|
||||
if (key === 'steps') {
|
||||
expect(value).toBeLessThanOrEqual(150);
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,160 @@
|
|||
// @vitest-environment node
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { parametersFixture } from '@/server/services/comfyui/__tests__/fixtures/parameters.fixture';
|
||||
import { setupAllMocks } from '@/server/services/comfyui/__tests__/setup/unifiedMocks';
|
||||
import { ComfyUIClientService } from '@/server/services/comfyui/core/comfyUIClientService';
|
||||
// Import services for testing
|
||||
import { ImageService } from '@/server/services/comfyui/core/imageService';
|
||||
import { ModelResolverService } from '@/server/services/comfyui/core/modelResolverService';
|
||||
import { WorkflowBuilderService } from '@/server/services/comfyui/core/workflowBuilderService';
|
||||
|
||||
describe('Service Integration - Module Level', () => {
|
||||
let imageService: ImageService;
|
||||
let clientService: ComfyUIClientService;
|
||||
let modelResolverService: ModelResolverService;
|
||||
let workflowBuilderService: WorkflowBuilderService;
|
||||
let inputCalls: Map<string, any>;
|
||||
|
||||
beforeEach(() => {
|
||||
const mocks = setupAllMocks();
|
||||
inputCalls = mocks.inputCalls;
|
||||
|
||||
// 创建服务实例
|
||||
clientService = new ComfyUIClientService();
|
||||
modelResolverService = new ModelResolverService(clientService);
|
||||
workflowBuilderService = new WorkflowBuilderService({
|
||||
clientService,
|
||||
modelResolverService,
|
||||
});
|
||||
|
||||
imageService = new ImageService(clientService, modelResolverService, workflowBuilderService);
|
||||
});
|
||||
|
||||
describe('Service Coordination', () => {
|
||||
it('should coordinate model resolution and workflow building', async () => {
|
||||
const modelResolverSpy = vi.spyOn(modelResolverService, 'validateModel');
|
||||
const validateConnectionSpy = vi.spyOn(clientService, 'validateConnection');
|
||||
|
||||
// Mock successful connection validation
|
||||
validateConnectionSpy.mockResolvedValue(true);
|
||||
|
||||
// Mock successful model validation
|
||||
modelResolverSpy.mockResolvedValue({
|
||||
exists: true,
|
||||
actualFileName: 'flux-dev.safetensors',
|
||||
});
|
||||
|
||||
const params = {
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
...parametersFixture.models['flux-dev'].defaults,
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
},
|
||||
};
|
||||
|
||||
try {
|
||||
await imageService.createImage(params);
|
||||
} catch (error) {
|
||||
// 预期在 mock 环境中可能有错误
|
||||
console.log('Expected error in mock environment:', error);
|
||||
}
|
||||
|
||||
// 验证服务调用顺序
|
||||
expect(validateConnectionSpy).toHaveBeenCalled();
|
||||
expect(modelResolverSpy).toHaveBeenCalledWith('flux-dev');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Context Passing', () => {
|
||||
it('should pass context between services correctly', async () => {
|
||||
const context = {
|
||||
clientService,
|
||||
modelResolverService,
|
||||
};
|
||||
|
||||
// 验证 WorkflowBuilderService 接收正确的 context
|
||||
expect(workflowBuilderService).toBeDefined();
|
||||
|
||||
// 测试 context 中的服务是否可用
|
||||
expect(clientService).toBeDefined();
|
||||
expect(modelResolverService).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Error Propagation Between Services', () => {
|
||||
it('should propagate errors from model resolver to image service', async () => {
|
||||
const modelResolverSpy = vi.spyOn(modelResolverService, 'validateModel');
|
||||
modelResolverSpy.mockRejectedValue(new Error('Model validation failed'));
|
||||
|
||||
const params = {
|
||||
model: 'invalid-model',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
},
|
||||
};
|
||||
|
||||
await expect(imageService.createImage(params)).rejects.toBeDefined();
|
||||
});
|
||||
|
||||
it('should handle workflow builder errors', async () => {
|
||||
const workflowBuilderSpy = vi.spyOn(workflowBuilderService, 'buildWorkflow');
|
||||
workflowBuilderSpy.mockRejectedValue(new Error('Workflow build failed'));
|
||||
|
||||
const params = {
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
...parametersFixture.models['flux-dev'].defaults,
|
||||
},
|
||||
};
|
||||
|
||||
await expect(imageService.createImage(params)).rejects.toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Service Dependencies', () => {
|
||||
it('should maintain proper service dependencies', () => {
|
||||
// ImageService 依赖其他三个服务
|
||||
expect(imageService).toBeDefined();
|
||||
|
||||
// ModelResolverService 依赖 ClientService
|
||||
expect(modelResolverService).toBeDefined();
|
||||
|
||||
// WorkflowBuilderService 依赖 context
|
||||
expect(workflowBuilderService).toBeDefined();
|
||||
|
||||
// ClientService 是基础服务
|
||||
expect(clientService).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Mock Integration', () => {
|
||||
it('should work with unified mocks', async () => {
|
||||
// 验证统一 mock 正常工作
|
||||
expect(inputCalls).toBeDefined();
|
||||
expect(inputCalls).toBeInstanceOf(Map);
|
||||
|
||||
// 测试 mock 是否被正确设置
|
||||
const params = {
|
||||
model: 'flux-dev',
|
||||
params: {
|
||||
prompt: 'test prompt',
|
||||
...parametersFixture.models['flux-dev'].defaults,
|
||||
},
|
||||
};
|
||||
|
||||
// 这应该使用统一的 mocks
|
||||
try {
|
||||
await imageService.createImage(params);
|
||||
} catch (error) {
|
||||
// 预期在 mock 环境中可能有错误
|
||||
}
|
||||
|
||||
// 验证基本功能正常
|
||||
expect(true).toBe(true);
|
||||
});
|
||||
});
|
||||
});
|
||||
48
src/server/services/comfyui/__tests__/setup/unifiedMocks.ts
Normal file
48
src/server/services/comfyui/__tests__/setup/unifiedMocks.ts
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
// @vitest-environment node
|
||||
import { vi } from 'vitest';
|
||||
|
||||
// Create mock PromptBuilder class first
|
||||
const MockPromptBuilder = vi.fn().mockImplementation((workflow: any) => ({
|
||||
input: vi.fn().mockReturnThis(),
|
||||
setInputNode: vi.fn().mockReturnThis(),
|
||||
setOutputNode: vi.fn().mockReturnThis(),
|
||||
workflow, // Expose the workflow for testing
|
||||
}));
|
||||
|
||||
// Module-level mock for @saintno/comfyui-sdk
|
||||
vi.mock('@saintno/comfyui-sdk', () => ({
|
||||
CallWrapper: vi.fn().mockImplementation(() => ({
|
||||
call: vi.fn(),
|
||||
execute: vi.fn(),
|
||||
})),
|
||||
ComfyApi: vi.fn().mockImplementation((baseURL: string, clientId?: string, options?: any) => ({
|
||||
baseURL,
|
||||
clientId,
|
||||
connect: vi.fn(),
|
||||
disconnect: vi.fn(),
|
||||
getObjectInfo: vi.fn().mockResolvedValue({}),
|
||||
init: vi.fn(),
|
||||
options,
|
||||
})),
|
||||
PromptBuilder: MockPromptBuilder,
|
||||
seed: vi.fn(() => 42),
|
||||
}));
|
||||
|
||||
export const setupAllMocks = () => {
|
||||
// Mock other utility functions
|
||||
vi.mock('../utils/promptSplitter', () => ({
|
||||
splitPromptForDualCLIP: vi.fn((prompt: string) => ({
|
||||
clipLPrompt: prompt,
|
||||
t5xxlPrompt: prompt,
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock('../utils/weightDType', () => ({
|
||||
selectOptimalWeightDtype: vi.fn(() => 'default'),
|
||||
}));
|
||||
|
||||
// Enhanced PromptBuilder mock to record parameters
|
||||
const inputCalls = new Map<string, any>();
|
||||
|
||||
return { inputCalls };
|
||||
};
|
||||
571
src/server/services/comfyui/__tests__/utils/cacheManager.test.ts
Normal file
571
src/server/services/comfyui/__tests__/utils/cacheManager.test.ts
Normal file
|
|
@ -0,0 +1,571 @@
|
|||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { TTLCacheManager } from '@/server/services/comfyui/utils/cacheManager';
|
||||
|
||||
// Mock debug module
|
||||
vi.mock('debug', () => ({
|
||||
default: vi.fn(() => vi.fn()),
|
||||
}));
|
||||
|
||||
describe('cacheManager.ts', () => {
|
||||
let cacheManager: TTLCacheManager;
|
||||
let mockFetcher: any;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
vi.useFakeTimers();
|
||||
cacheManager = new TTLCacheManager(60000); // 60 second TTL
|
||||
mockFetcher = vi.fn();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
describe('TTLCacheManager constructor', () => {
|
||||
it('should create instance with default TTL', () => {
|
||||
const cache = new TTLCacheManager();
|
||||
expect(cache).toBeInstanceOf(TTLCacheManager);
|
||||
});
|
||||
|
||||
it('should create instance with custom TTL', () => {
|
||||
const cache = new TTLCacheManager(30000);
|
||||
expect(cache).toBeInstanceOf(TTLCacheManager);
|
||||
});
|
||||
|
||||
it('should handle zero TTL', () => {
|
||||
const cache = new TTLCacheManager(0);
|
||||
expect(cache).toBeInstanceOf(TTLCacheManager);
|
||||
});
|
||||
|
||||
it('should handle negative TTL', () => {
|
||||
const cache = new TTLCacheManager(-1000);
|
||||
expect(cache).toBeInstanceOf(TTLCacheManager);
|
||||
});
|
||||
});
|
||||
|
||||
describe('get method', () => {
|
||||
it('should fetch and cache value on first call', async () => {
|
||||
const testValue = 'test-value';
|
||||
mockFetcher.mockResolvedValue(testValue);
|
||||
|
||||
const result = await cacheManager.get('test-key', mockFetcher);
|
||||
|
||||
expect(result).toBe(testValue);
|
||||
expect(mockFetcher).toHaveBeenCalledTimes(1);
|
||||
expect(cacheManager.size()).toBe(1);
|
||||
});
|
||||
|
||||
it('should return cached value on subsequent calls within TTL', async () => {
|
||||
const testValue = 'cached-value';
|
||||
mockFetcher.mockResolvedValue(testValue);
|
||||
|
||||
// First call
|
||||
const result1 = await cacheManager.get('test-key', mockFetcher);
|
||||
|
||||
// Advance time by 30 seconds (within TTL)
|
||||
vi.advanceTimersByTime(30000);
|
||||
|
||||
// Second call
|
||||
const result2 = await cacheManager.get('test-key', mockFetcher);
|
||||
|
||||
expect(result1).toBe(testValue);
|
||||
expect(result2).toBe(testValue);
|
||||
expect(mockFetcher).toHaveBeenCalledTimes(1); // Fetcher called only once
|
||||
});
|
||||
|
||||
it('should re-fetch value after TTL expires', async () => {
|
||||
const firstValue = 'first-value';
|
||||
const secondValue = 'second-value';
|
||||
mockFetcher.mockResolvedValueOnce(firstValue).mockResolvedValueOnce(secondValue);
|
||||
|
||||
// First call
|
||||
const result1 = await cacheManager.get('test-key', mockFetcher);
|
||||
|
||||
// Advance time beyond TTL
|
||||
vi.advanceTimersByTime(70000);
|
||||
|
||||
// Second call after TTL expiration
|
||||
const result2 = await cacheManager.get('test-key', mockFetcher);
|
||||
|
||||
expect(result1).toBe(firstValue);
|
||||
expect(result2).toBe(secondValue);
|
||||
expect(mockFetcher).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should handle multiple different keys', async () => {
|
||||
const value1 = 'value-1';
|
||||
const value2 = 'value-2';
|
||||
const fetcher1 = vi.fn().mockResolvedValue(value1);
|
||||
const fetcher2 = vi.fn().mockResolvedValue(value2);
|
||||
|
||||
const result1 = await cacheManager.get('key-1', fetcher1);
|
||||
const result2 = await cacheManager.get('key-2', fetcher2);
|
||||
|
||||
expect(result1).toBe(value1);
|
||||
expect(result2).toBe(value2);
|
||||
expect(cacheManager.size()).toBe(2);
|
||||
expect(fetcher1).toHaveBeenCalledTimes(1);
|
||||
expect(fetcher2).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should handle fetcher that throws error', async () => {
|
||||
const error = new Error('Fetcher failed');
|
||||
mockFetcher.mockRejectedValue(error);
|
||||
|
||||
await expect(cacheManager.get('test-key', mockFetcher)).rejects.toThrow('Fetcher failed');
|
||||
expect(cacheManager.size()).toBe(0); // Should not cache failed results
|
||||
});
|
||||
|
||||
it('should handle async fetcher correctly', async () => {
|
||||
const asyncFetcher = vi
|
||||
.fn()
|
||||
.mockImplementation(
|
||||
() => new Promise((resolve) => setTimeout(() => resolve('async-value'), 100)),
|
||||
);
|
||||
|
||||
const resultPromise = cacheManager.get('async-key', asyncFetcher);
|
||||
|
||||
// Advance time to resolve the async fetcher
|
||||
vi.advanceTimersByTime(100);
|
||||
|
||||
const result = await resultPromise;
|
||||
expect(result).toBe('async-value');
|
||||
expect(cacheManager.size()).toBe(1);
|
||||
});
|
||||
|
||||
it('should handle concurrent requests for same key', async () => {
|
||||
let callCount = 0;
|
||||
const concurrentFetcher = vi.fn().mockImplementation(() => {
|
||||
callCount++;
|
||||
return Promise.resolve(`value-${callCount}`);
|
||||
});
|
||||
|
||||
// Start multiple concurrent requests for the same key
|
||||
const promises = [
|
||||
cacheManager.get('concurrent-key', concurrentFetcher),
|
||||
cacheManager.get('concurrent-key', concurrentFetcher),
|
||||
cacheManager.get('concurrent-key', concurrentFetcher),
|
||||
];
|
||||
|
||||
const results = await Promise.all(promises);
|
||||
|
||||
// Note: In the current implementation, concurrent calls may each trigger the fetcher
|
||||
// since there's no deduplication. This test verifies that the cache works correctly.
|
||||
expect(results.length).toBe(3);
|
||||
expect(cacheManager.size()).toBe(1);
|
||||
|
||||
// At least some of the results should be the same if caching worked
|
||||
const uniqueResults = [...new Set(results)];
|
||||
expect(uniqueResults.length).toBeGreaterThanOrEqual(1);
|
||||
});
|
||||
|
||||
it('should handle different data types', async () => {
|
||||
const objectValue = { foo: 'bar', num: 42 };
|
||||
const arrayValue = [1, 2, 3, 'test'];
|
||||
const numberValue = 123.45;
|
||||
const booleanValue = true;
|
||||
const nullValue = null;
|
||||
const undefinedValue = undefined;
|
||||
|
||||
const results = await Promise.all([
|
||||
cacheManager.get('object', () => Promise.resolve(objectValue)),
|
||||
cacheManager.get('array', () => Promise.resolve(arrayValue)),
|
||||
cacheManager.get('number', () => Promise.resolve(numberValue)),
|
||||
cacheManager.get('boolean', () => Promise.resolve(booleanValue)),
|
||||
cacheManager.get('null', () => Promise.resolve(nullValue)),
|
||||
cacheManager.get('undefined', () => Promise.resolve(undefinedValue)),
|
||||
]);
|
||||
|
||||
expect(results[0]).toEqual(objectValue);
|
||||
expect(results[1]).toEqual(arrayValue);
|
||||
expect(results[2]).toBe(numberValue);
|
||||
expect(results[3]).toBe(booleanValue);
|
||||
expect(results[4]).toBe(nullValue);
|
||||
expect(results[5]).toBe(undefinedValue);
|
||||
expect(cacheManager.size()).toBe(6);
|
||||
});
|
||||
|
||||
it('should handle empty string key', async () => {
|
||||
const testValue = 'empty-key-value';
|
||||
mockFetcher.mockResolvedValue(testValue);
|
||||
|
||||
const result = await cacheManager.get('', mockFetcher);
|
||||
expect(result).toBe(testValue);
|
||||
expect(cacheManager.has('')).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle special characters in key', async () => {
|
||||
const specialKey = 'key-with-special-chars!@#$%^&*()[]{}|;:,.<>?';
|
||||
const testValue = 'special-value';
|
||||
mockFetcher.mockResolvedValue(testValue);
|
||||
|
||||
const result = await cacheManager.get(specialKey, mockFetcher);
|
||||
expect(result).toBe(testValue);
|
||||
expect(cacheManager.has(specialKey)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('invalidate method', () => {
|
||||
it('should remove specific cache entry', async () => {
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await cacheManager.get('test-key', mockFetcher);
|
||||
expect(cacheManager.size()).toBe(1);
|
||||
expect(cacheManager.has('test-key')).toBe(true);
|
||||
|
||||
cacheManager.invalidate('test-key');
|
||||
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
expect(cacheManager.has('test-key')).toBe(false);
|
||||
});
|
||||
|
||||
it('should not affect other cache entries', async () => {
|
||||
const fetcher1 = vi.fn().mockResolvedValue('value-1');
|
||||
const fetcher2 = vi.fn().mockResolvedValue('value-2');
|
||||
|
||||
await cacheManager.get('key-1', fetcher1);
|
||||
await cacheManager.get('key-2', fetcher2);
|
||||
expect(cacheManager.size()).toBe(2);
|
||||
|
||||
cacheManager.invalidate('key-1');
|
||||
|
||||
expect(cacheManager.size()).toBe(1);
|
||||
expect(cacheManager.has('key-1')).toBe(false);
|
||||
expect(cacheManager.has('key-2')).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle invalidating non-existent key gracefully', () => {
|
||||
expect(() => cacheManager.invalidate('non-existent')).not.toThrow();
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
});
|
||||
|
||||
it('should cause re-fetch after invalidation', async () => {
|
||||
const firstValue = 'first-value';
|
||||
const secondValue = 'second-value';
|
||||
mockFetcher.mockResolvedValueOnce(firstValue).mockResolvedValueOnce(secondValue);
|
||||
|
||||
// First call
|
||||
const result1 = await cacheManager.get('test-key', mockFetcher);
|
||||
expect(result1).toBe(firstValue);
|
||||
expect(mockFetcher).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Invalidate
|
||||
cacheManager.invalidate('test-key');
|
||||
|
||||
// Second call should re-fetch
|
||||
const result2 = await cacheManager.get('test-key', mockFetcher);
|
||||
expect(result2).toBe(secondValue);
|
||||
expect(mockFetcher).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe('invalidateAll method', () => {
|
||||
it('should clear all cache entries', async () => {
|
||||
const fetcher1 = vi.fn().mockResolvedValue('value-1');
|
||||
const fetcher2 = vi.fn().mockResolvedValue('value-2');
|
||||
const fetcher3 = vi.fn().mockResolvedValue('value-3');
|
||||
|
||||
await cacheManager.get('key-1', fetcher1);
|
||||
await cacheManager.get('key-2', fetcher2);
|
||||
await cacheManager.get('key-3', fetcher3);
|
||||
expect(cacheManager.size()).toBe(3);
|
||||
|
||||
cacheManager.invalidateAll();
|
||||
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
expect(cacheManager.has('key-1')).toBe(false);
|
||||
expect(cacheManager.has('key-2')).toBe(false);
|
||||
expect(cacheManager.has('key-3')).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle clearing empty cache', () => {
|
||||
expect(() => cacheManager.invalidateAll()).not.toThrow();
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
});
|
||||
|
||||
it('should cause re-fetch for all keys after clear', async () => {
|
||||
const fetcher1 = vi.fn().mockResolvedValue('value-1').mockResolvedValue('new-value-1');
|
||||
const fetcher2 = vi.fn().mockResolvedValue('value-2').mockResolvedValue('new-value-2');
|
||||
|
||||
// Initial calls
|
||||
await cacheManager.get('key-1', fetcher1);
|
||||
await cacheManager.get('key-2', fetcher2);
|
||||
expect(fetcher1).toHaveBeenCalledTimes(1);
|
||||
expect(fetcher2).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Clear all
|
||||
cacheManager.invalidateAll();
|
||||
|
||||
// Subsequent calls should re-fetch
|
||||
await cacheManager.get('key-1', fetcher1);
|
||||
await cacheManager.get('key-2', fetcher2);
|
||||
expect(fetcher1).toHaveBeenCalledTimes(2);
|
||||
expect(fetcher2).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe('size method', () => {
|
||||
it('should return zero for empty cache', () => {
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
});
|
||||
|
||||
it('should return correct count after adding entries', async () => {
|
||||
const fetcher1 = vi.fn().mockResolvedValue('value-1');
|
||||
const fetcher2 = vi.fn().mockResolvedValue('value-2');
|
||||
const fetcher3 = vi.fn().mockResolvedValue('value-3');
|
||||
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
|
||||
await cacheManager.get('key-1', fetcher1);
|
||||
expect(cacheManager.size()).toBe(1);
|
||||
|
||||
await cacheManager.get('key-2', fetcher2);
|
||||
expect(cacheManager.size()).toBe(2);
|
||||
|
||||
await cacheManager.get('key-3', fetcher3);
|
||||
expect(cacheManager.size()).toBe(3);
|
||||
});
|
||||
|
||||
it('should return correct count after removing entries', async () => {
|
||||
const fetcher1 = vi.fn().mockResolvedValue('value-1');
|
||||
const fetcher2 = vi.fn().mockResolvedValue('value-2');
|
||||
|
||||
await cacheManager.get('key-1', fetcher1);
|
||||
await cacheManager.get('key-2', fetcher2);
|
||||
expect(cacheManager.size()).toBe(2);
|
||||
|
||||
cacheManager.invalidate('key-1');
|
||||
expect(cacheManager.size()).toBe(1);
|
||||
|
||||
cacheManager.invalidateAll();
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('has method', () => {
|
||||
it('should return false for non-existent key', () => {
|
||||
expect(cacheManager.has('non-existent')).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true for existing key', async () => {
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
expect(cacheManager.has('test-key')).toBe(false);
|
||||
|
||||
await cacheManager.get('test-key', mockFetcher);
|
||||
|
||||
expect(cacheManager.has('test-key')).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false after invalidation', async () => {
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await cacheManager.get('test-key', mockFetcher);
|
||||
expect(cacheManager.has('test-key')).toBe(true);
|
||||
|
||||
cacheManager.invalidate('test-key');
|
||||
expect(cacheManager.has('test-key')).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true for expired entries (key exists but may be expired)', async () => {
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await cacheManager.get('test-key', mockFetcher);
|
||||
expect(cacheManager.has('test-key')).toBe(true);
|
||||
|
||||
// Advance time beyond TTL
|
||||
vi.advanceTimersByTime(70000);
|
||||
|
||||
// has() checks existence regardless of expiration
|
||||
expect(cacheManager.has('test-key')).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('isValid method', () => {
|
||||
it('should return false for non-existent key', () => {
|
||||
expect(cacheManager.isValid('non-existent')).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true for valid (not expired) entry', async () => {
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await cacheManager.get('test-key', mockFetcher);
|
||||
expect(cacheManager.isValid('test-key')).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false for expired entry', async () => {
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await cacheManager.get('test-key', mockFetcher);
|
||||
expect(cacheManager.isValid('test-key')).toBe(true);
|
||||
|
||||
// Advance time beyond TTL
|
||||
vi.advanceTimersByTime(70000);
|
||||
|
||||
expect(cacheManager.isValid('test-key')).toBe(false);
|
||||
});
|
||||
|
||||
it('should return true just before expiration', async () => {
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await cacheManager.get('test-key', mockFetcher);
|
||||
expect(cacheManager.isValid('test-key')).toBe(true);
|
||||
|
||||
// Advance time to just before TTL expiration
|
||||
vi.advanceTimersByTime(59999);
|
||||
|
||||
expect(cacheManager.isValid('test-key')).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false at exact expiration time', async () => {
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await cacheManager.get('test-key', mockFetcher);
|
||||
expect(cacheManager.isValid('test-key')).toBe(true);
|
||||
|
||||
// Advance time to exact TTL expiration
|
||||
vi.advanceTimersByTime(60000);
|
||||
|
||||
expect(cacheManager.isValid('test-key')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('TTL behavior', () => {
|
||||
it('should handle zero TTL correctly', async () => {
|
||||
const zeroTTLCache = new TTLCacheManager(0);
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await zeroTTLCache.get('test-key', mockFetcher);
|
||||
expect(zeroTTLCache.isValid('test-key')).toBe(false); // Should be immediately invalid
|
||||
});
|
||||
|
||||
it('should handle negative TTL correctly', async () => {
|
||||
const negativeTTLCache = new TTLCacheManager(-1000);
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await negativeTTLCache.get('test-key', mockFetcher);
|
||||
expect(negativeTTLCache.isValid('test-key')).toBe(false); // Should be immediately invalid
|
||||
});
|
||||
|
||||
it('should handle very short TTL', async () => {
|
||||
const shortTTLCache = new TTLCacheManager(100); // 100ms TTL
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await shortTTLCache.get('test-key', mockFetcher);
|
||||
expect(shortTTLCache.isValid('test-key')).toBe(true);
|
||||
|
||||
vi.advanceTimersByTime(150);
|
||||
expect(shortTTLCache.isValid('test-key')).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle very long TTL', async () => {
|
||||
const longTTLCache = new TTLCacheManager(1000000000); // Very long TTL
|
||||
mockFetcher.mockResolvedValue('test-value');
|
||||
|
||||
await longTTLCache.get('test-key', mockFetcher);
|
||||
expect(longTTLCache.isValid('test-key')).toBe(true);
|
||||
|
||||
vi.advanceTimersByTime(999999999);
|
||||
expect(longTTLCache.isValid('test-key')).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('edge cases and error scenarios', () => {
|
||||
it('should handle fetcher returning Promise.reject', async () => {
|
||||
const rejectedFetcher = vi.fn().mockRejectedValue(new Error('Async rejection'));
|
||||
|
||||
await expect(cacheManager.get('reject-key', rejectedFetcher)).rejects.toThrow(
|
||||
'Async rejection',
|
||||
);
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
});
|
||||
|
||||
it('should handle fetcher throwing synchronous error', async () => {
|
||||
const throwingFetcher = vi.fn().mockImplementation(() => {
|
||||
throw new Error('Synchronous error');
|
||||
});
|
||||
|
||||
await expect(cacheManager.get('throw-key', throwingFetcher)).rejects.toThrow(
|
||||
'Synchronous error',
|
||||
);
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
});
|
||||
|
||||
it('should handle very long keys', async () => {
|
||||
const longKey = 'a'.repeat(10000);
|
||||
mockFetcher.mockResolvedValue('long-key-value');
|
||||
|
||||
const result = await cacheManager.get(longKey, mockFetcher);
|
||||
expect(result).toBe('long-key-value');
|
||||
expect(cacheManager.has(longKey)).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle unicode keys', async () => {
|
||||
const unicodeKey = '测试-键-🔑-نام';
|
||||
mockFetcher.mockResolvedValue('unicode-value');
|
||||
|
||||
const result = await cacheManager.get(unicodeKey, mockFetcher);
|
||||
expect(result).toBe('unicode-value');
|
||||
expect(cacheManager.has(unicodeKey)).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle null and undefined values from fetcher', async () => {
|
||||
const nullFetcher = vi.fn().mockResolvedValue(null);
|
||||
const undefinedFetcher = vi.fn().mockResolvedValue(undefined);
|
||||
|
||||
const nullResult = await cacheManager.get('null-key', nullFetcher);
|
||||
const undefinedResult = await cacheManager.get('undefined-key', undefinedFetcher);
|
||||
|
||||
expect(nullResult).toBe(null);
|
||||
expect(undefinedResult).toBe(undefined);
|
||||
expect(cacheManager.size()).toBe(2);
|
||||
});
|
||||
|
||||
it('should handle rapid sequential operations', async () => {
|
||||
const operations = [];
|
||||
|
||||
for (let i = 0; i < 100; i++) {
|
||||
const fetcher = vi.fn().mockResolvedValue(`value-${i}`);
|
||||
operations.push(cacheManager.get(`key-${i}`, fetcher));
|
||||
}
|
||||
|
||||
const results = await Promise.all(operations);
|
||||
expect(results).toHaveLength(100);
|
||||
expect(cacheManager.size()).toBe(100);
|
||||
});
|
||||
});
|
||||
|
||||
describe('memory management', () => {
|
||||
it('should not grow indefinitely with different keys', async () => {
|
||||
// Add many entries
|
||||
for (let i = 0; i < 1000; i++) {
|
||||
const fetcher = vi.fn().mockResolvedValue(`value-${i}`);
|
||||
await cacheManager.get(`key-${i}`, fetcher);
|
||||
}
|
||||
|
||||
expect(cacheManager.size()).toBe(1000);
|
||||
|
||||
// Clear all
|
||||
cacheManager.invalidateAll();
|
||||
expect(cacheManager.size()).toBe(0);
|
||||
});
|
||||
|
||||
it('should handle invalidation of many entries efficiently', async () => {
|
||||
// Add many entries
|
||||
for (let i = 0; i < 100; i++) {
|
||||
const fetcher = vi.fn().mockResolvedValue(`value-${i}`);
|
||||
await cacheManager.get(`key-${i}`, fetcher);
|
||||
}
|
||||
|
||||
expect(cacheManager.size()).toBe(100);
|
||||
|
||||
// Invalidate half
|
||||
for (let i = 0; i < 50; i++) {
|
||||
cacheManager.invalidate(`key-${i}`);
|
||||
}
|
||||
|
||||
expect(cacheManager.size()).toBe(50);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,329 @@
|
|||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { COMPONENT_NODE_MAPPINGS } from '@/server/services/comfyui/config/constants';
|
||||
import { SYSTEM_COMPONENTS } from '@/server/services/comfyui/config/systemComponents';
|
||||
import {
|
||||
type ComponentInfo,
|
||||
getComponentDisplayName,
|
||||
getComponentFolderPath,
|
||||
getComponentInfo,
|
||||
isSystemComponent,
|
||||
} from '@/server/services/comfyui/utils/componentInfo';
|
||||
|
||||
// Mock the config modules to have full control over test data
|
||||
vi.mock('@/server/services/comfyui/config/constants', () => ({
|
||||
COMPONENT_NODE_MAPPINGS: {
|
||||
clip: { node: 'CLIPTextEncode' },
|
||||
controlnet: { node: 'ControlNetApply' },
|
||||
lora: { node: 'LoraLoader' },
|
||||
t5: { node: 'T5TextEncode' },
|
||||
vae: { node: 'VAEDecode' },
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('@/server/services/comfyui/config/systemComponents', () => ({
|
||||
SYSTEM_COMPONENTS: {
|
||||
'clip-l.safetensors': { type: 'clip', modelFamily: 'FLUX', priority: 1 },
|
||||
'flux-dev.safetensors': { type: 'unet', modelFamily: 'FLUX', priority: 1 },
|
||||
'invalid-component.bin': { type: 'unknown', modelFamily: 'TEST', priority: 3 },
|
||||
't5-xxl.safetensors': { type: 't5', modelFamily: 'FLUX', priority: 1 },
|
||||
'vae.safetensors': { type: 'vae', modelFamily: 'FLUX', priority: 2 },
|
||||
},
|
||||
}));
|
||||
|
||||
describe('componentInfo.ts', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('getComponentDisplayName', () => {
|
||||
it('should return correct display name for t5 type', () => {
|
||||
const result = getComponentDisplayName('t5');
|
||||
expect(result).toBe('T5 text encoder');
|
||||
});
|
||||
|
||||
it('should return correct display name for clip type', () => {
|
||||
const result = getComponentDisplayName('clip');
|
||||
expect(result).toBe('CLIP text encoder');
|
||||
});
|
||||
|
||||
it('should return correct display name for vae type', () => {
|
||||
const result = getComponentDisplayName('vae');
|
||||
expect(result).toBe('VAE model');
|
||||
});
|
||||
|
||||
it('should return generic display name for unknown type', () => {
|
||||
const result = getComponentDisplayName('unknown');
|
||||
expect(result).toBe('UNKNOWN component');
|
||||
});
|
||||
|
||||
it('should return uppercase display name for custom type', () => {
|
||||
const result = getComponentDisplayName('lora');
|
||||
expect(result).toBe('LORA component');
|
||||
});
|
||||
|
||||
it('should handle empty string', () => {
|
||||
const result = getComponentDisplayName('');
|
||||
expect(result).toBe(' component');
|
||||
});
|
||||
|
||||
it('should handle null and undefined gracefully', () => {
|
||||
// The function expects string input and will handle null/undefined by converting to string
|
||||
expect(() => getComponentDisplayName(null as any)).toThrow();
|
||||
expect(() => getComponentDisplayName(undefined as any)).toThrow();
|
||||
});
|
||||
|
||||
it('should handle numeric input', () => {
|
||||
const result = getComponentDisplayName('123' as any);
|
||||
expect(result).toBe('123 component');
|
||||
});
|
||||
});
|
||||
|
||||
describe('getComponentFolderPath', () => {
|
||||
it('should return models/clip for clip type', () => {
|
||||
const result = getComponentFolderPath('clip');
|
||||
expect(result).toBe('models/clip');
|
||||
});
|
||||
|
||||
it('should return models/clip for t5 type', () => {
|
||||
const result = getComponentFolderPath('t5');
|
||||
expect(result).toBe('models/clip');
|
||||
});
|
||||
|
||||
it('should return models/vae for vae type', () => {
|
||||
const result = getComponentFolderPath('vae');
|
||||
expect(result).toBe('models/vae');
|
||||
});
|
||||
|
||||
it('should return generic models path for unknown type', () => {
|
||||
const result = getComponentFolderPath('unknown');
|
||||
expect(result).toBe('models/unknown');
|
||||
});
|
||||
|
||||
it('should return generic models path for custom type', () => {
|
||||
const result = getComponentFolderPath('lora');
|
||||
expect(result).toBe('models/lora');
|
||||
});
|
||||
|
||||
it('should handle empty string', () => {
|
||||
const result = getComponentFolderPath('');
|
||||
expect(result).toBe('models/');
|
||||
});
|
||||
|
||||
it('should handle null and undefined', () => {
|
||||
const resultNull = getComponentFolderPath(null as any);
|
||||
expect(resultNull).toBe('models/null');
|
||||
|
||||
const resultUndefined = getComponentFolderPath(undefined as any);
|
||||
expect(resultUndefined).toBe('models/undefined');
|
||||
});
|
||||
|
||||
it('should handle special characters in type', () => {
|
||||
const result = getComponentFolderPath('my-custom_type.v2');
|
||||
expect(result).toBe('models/my-custom_type.v2');
|
||||
});
|
||||
});
|
||||
|
||||
describe('getComponentInfo', () => {
|
||||
it('should return complete component info for valid clip component', () => {
|
||||
const result = getComponentInfo('clip-l.safetensors');
|
||||
|
||||
expect(result).toEqual({
|
||||
displayName: 'CLIP text encoder',
|
||||
folderPath: 'models/clip',
|
||||
nodeType: 'CLIPTextEncode',
|
||||
type: 'clip',
|
||||
});
|
||||
});
|
||||
|
||||
it('should return complete component info for valid t5 component', () => {
|
||||
const result = getComponentInfo('t5-xxl.safetensors');
|
||||
|
||||
expect(result).toEqual({
|
||||
displayName: 'T5 text encoder',
|
||||
folderPath: 'models/clip',
|
||||
nodeType: 'T5TextEncode',
|
||||
type: 't5',
|
||||
});
|
||||
});
|
||||
|
||||
it('should return complete component info for valid vae component', () => {
|
||||
const result = getComponentInfo('vae.safetensors');
|
||||
|
||||
expect(result).toEqual({
|
||||
displayName: 'VAE model',
|
||||
folderPath: 'models/vae',
|
||||
nodeType: 'VAEDecode',
|
||||
type: 'vae',
|
||||
});
|
||||
});
|
||||
|
||||
it('should return undefined for non-existent component', () => {
|
||||
const result = getComponentInfo('non-existent.safetensors');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should return undefined for component with unknown type', () => {
|
||||
const result = getComponentInfo('invalid-component.bin');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should handle empty string filename', () => {
|
||||
const result = getComponentInfo('');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should handle null and undefined filename', () => {
|
||||
const resultNull = getComponentInfo(null as any);
|
||||
expect(resultNull).toBeUndefined();
|
||||
|
||||
const resultUndefined = getComponentInfo(undefined as any);
|
||||
expect(resultUndefined).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should work with different file extensions', () => {
|
||||
// Since we're mocking the entire module, we need to test the existing mock data
|
||||
// This test validates that the function works with the current mock setup
|
||||
const result = getComponentInfo('clip-l.safetensors');
|
||||
expect(result).toEqual({
|
||||
displayName: 'CLIP text encoder',
|
||||
folderPath: 'models/clip',
|
||||
nodeType: 'CLIPTextEncode',
|
||||
type: 'clip',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('isSystemComponent', () => {
|
||||
it('should return true for known system components', () => {
|
||||
expect(isSystemComponent('clip-l.safetensors')).toBe(true);
|
||||
expect(isSystemComponent('t5-xxl.safetensors')).toBe(true);
|
||||
expect(isSystemComponent('vae.safetensors')).toBe(true);
|
||||
expect(isSystemComponent('flux-dev.safetensors')).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false for unknown components', () => {
|
||||
expect(isSystemComponent('unknown.safetensors')).toBe(false);
|
||||
expect(isSystemComponent('random-file.txt')).toBe(false);
|
||||
});
|
||||
|
||||
it('should handle edge cases', () => {
|
||||
expect(isSystemComponent('')).toBe(false);
|
||||
expect(isSystemComponent(null as any)).toBe(false);
|
||||
expect(isSystemComponent(undefined as any)).toBe(false);
|
||||
});
|
||||
|
||||
it('should be case sensitive', () => {
|
||||
expect(isSystemComponent('CLIP-L.SAFETENSORS')).toBe(false);
|
||||
expect(isSystemComponent('clip-l.safetensors')).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle special characters', () => {
|
||||
// Test with existing mock data that contains standard characters
|
||||
expect(isSystemComponent('clip-l.safetensors')).toBe(true);
|
||||
expect(isSystemComponent('t5-xxl.safetensors')).toBe(true);
|
||||
|
||||
// Test with non-existent special character filename
|
||||
expect(isSystemComponent('special-file@2.0_beta.safetensors')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('integration tests', () => {
|
||||
it('should provide consistent results across all functions for same component', () => {
|
||||
const fileName = 'clip-l.safetensors';
|
||||
const componentInfo = getComponentInfo(fileName);
|
||||
const isSystem = isSystemComponent(fileName);
|
||||
|
||||
expect(isSystem).toBe(true);
|
||||
expect(componentInfo).toBeDefined();
|
||||
expect(componentInfo!.type).toBe('clip');
|
||||
|
||||
const displayName = getComponentDisplayName(componentInfo!.type);
|
||||
const folderPath = getComponentFolderPath(componentInfo!.type);
|
||||
|
||||
expect(displayName).toBe(componentInfo!.displayName);
|
||||
expect(folderPath).toBe(componentInfo!.folderPath);
|
||||
});
|
||||
|
||||
it('should handle workflow where component exists but node mapping missing', () => {
|
||||
// Test with existing mock data - flux-dev.safetensors has type 'unet' which is not in node mappings
|
||||
const result = getComponentInfo('flux-dev.safetensors');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should handle concurrent access safely', async () => {
|
||||
// Test concurrent access to functions
|
||||
const promises = Array.from({ length: 100 }, (_, i) =>
|
||||
Promise.all([
|
||||
Promise.resolve(getComponentInfo('clip-l.safetensors')),
|
||||
Promise.resolve(isSystemComponent('t5-xxl.safetensors')),
|
||||
Promise.resolve(getComponentDisplayName('vae')),
|
||||
Promise.resolve(getComponentFolderPath('clip')),
|
||||
]),
|
||||
);
|
||||
|
||||
const results = await Promise.all(promises);
|
||||
|
||||
// All results should be consistent
|
||||
results.forEach(([info, isSystem, displayName, folderPath]) => {
|
||||
expect(info).toBeDefined();
|
||||
expect(isSystem).toBe(true);
|
||||
expect(displayName).toBe('VAE model');
|
||||
expect(folderPath).toBe('models/clip');
|
||||
});
|
||||
});
|
||||
|
||||
it('should maintain type safety with ComponentInfo interface', () => {
|
||||
const info = getComponentInfo('clip-l.safetensors');
|
||||
|
||||
if (info) {
|
||||
// These should not cause TypeScript errors
|
||||
const displayName: string = info.displayName;
|
||||
const folderPath: string = info.folderPath;
|
||||
const nodeType: string = info.nodeType;
|
||||
const type: string = info.type;
|
||||
|
||||
expect(typeof displayName).toBe('string');
|
||||
expect(typeof folderPath).toBe('string');
|
||||
expect(typeof nodeType).toBe('string');
|
||||
expect(typeof type).toBe('string');
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('error handling and robustness', () => {
|
||||
it('should handle corrupted SYSTEM_COMPONENTS gracefully', () => {
|
||||
// Since we can't easily mock return values, test with invalid input
|
||||
expect(() => getComponentInfo('any-file.safetensors')).not.toThrow();
|
||||
expect(() => isSystemComponent('any-file.safetensors')).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle corrupted COMPONENT_NODE_MAPPINGS gracefully', () => {
|
||||
// Test with a component that exists in SYSTEM_COMPONENTS but has invalid type
|
||||
const result = getComponentInfo('invalid-component.bin');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should handle missing properties in config', () => {
|
||||
// Test with a component that has invalid type in our mock
|
||||
const result = getComponentInfo('invalid-component.bin');
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should handle very long filenames', () => {
|
||||
const longFilename = 'a'.repeat(1000) + '.safetensors';
|
||||
expect(() => {
|
||||
getComponentInfo(longFilename);
|
||||
isSystemComponent(longFilename);
|
||||
}).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle unicode characters in filenames', () => {
|
||||
const unicodeFilename = '测试-文件-🤖.safetensors';
|
||||
expect(() => {
|
||||
getComponentInfo(unicodeFilename);
|
||||
isSystemComponent(unicodeFilename);
|
||||
}).not.toThrow();
|
||||
});
|
||||
});
|
||||
});
|
||||
424
src/server/services/comfyui/__tests__/utils/imageResizer.test.ts
Normal file
424
src/server/services/comfyui/__tests__/utils/imageResizer.test.ts
Normal file
|
|
@ -0,0 +1,424 @@
|
|||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import {
|
||||
type Architecture,
|
||||
ImageResizer,
|
||||
imageResizer,
|
||||
} from '@/server/services/comfyui/utils/imageResizer';
|
||||
|
||||
// Mock debug module
|
||||
vi.mock('debug', () => ({
|
||||
default: vi.fn(() => vi.fn()),
|
||||
}));
|
||||
|
||||
describe('imageResizer.ts', () => {
|
||||
let resizer: ImageResizer;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
resizer = new ImageResizer();
|
||||
});
|
||||
|
||||
describe('ImageResizer class', () => {
|
||||
describe('calculateTargetDimensions', () => {
|
||||
describe('FLUX architecture', () => {
|
||||
it('should not resize when dimensions are within limits', () => {
|
||||
const result = resizer.calculateTargetDimensions(1024, 768, 'FLUX');
|
||||
|
||||
expect(result).toEqual({
|
||||
width: 1024,
|
||||
height: 768,
|
||||
needsResize: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should resize when width exceeds maximum', () => {
|
||||
const result = resizer.calculateTargetDimensions(2000, 800, 'FLUX');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeLessThanOrEqual(1440);
|
||||
expect(result.height).toBeLessThanOrEqual(1440);
|
||||
expect(result.width % 32).toBe(0); // Should be rounded to step
|
||||
expect(result.height % 32).toBe(0);
|
||||
});
|
||||
|
||||
it('should resize when height exceeds maximum', () => {
|
||||
const result = resizer.calculateTargetDimensions(800, 2000, 'FLUX');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeLessThanOrEqual(1440);
|
||||
expect(result.height).toBeLessThanOrEqual(1440);
|
||||
expect(result.width % 32).toBe(0);
|
||||
expect(result.height % 32).toBe(0);
|
||||
});
|
||||
|
||||
it('should resize when dimensions are too small', () => {
|
||||
const result = resizer.calculateTargetDimensions(100, 200, 'FLUX');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeGreaterThanOrEqual(256);
|
||||
expect(result.height).toBeGreaterThanOrEqual(256);
|
||||
expect(result.width % 32).toBe(0);
|
||||
expect(result.height % 32).toBe(0);
|
||||
});
|
||||
|
||||
it('should handle exact limit dimensions', () => {
|
||||
const result = resizer.calculateTargetDimensions(1440, 256, 'FLUX');
|
||||
|
||||
expect(result).toEqual({
|
||||
width: 1440,
|
||||
height: 256,
|
||||
needsResize: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should round to step size (32) correctly', () => {
|
||||
const result = resizer.calculateTargetDimensions(1000, 600, 'FLUX');
|
||||
|
||||
if (result.needsResize) {
|
||||
expect(result.width % 32).toBe(0);
|
||||
expect(result.height % 32).toBe(0);
|
||||
} else {
|
||||
// Original dimensions should be accepted
|
||||
expect(result.width).toBe(1000);
|
||||
expect(result.height).toBe(600);
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle aspect ratio constraints', () => {
|
||||
// Test extreme aspect ratios
|
||||
const wideResult = resizer.calculateTargetDimensions(2100, 900, 'FLUX'); // 2.33 ratio
|
||||
const tallResult = resizer.calculateTargetDimensions(900, 2100, 'FLUX'); // 0.43 ratio
|
||||
|
||||
expect(wideResult.needsResize).toBe(true);
|
||||
expect(tallResult.needsResize).toBe(true);
|
||||
|
||||
// Should still produce valid dimensions within model limits
|
||||
expect(wideResult.width).toBeLessThanOrEqual(1440);
|
||||
expect(wideResult.height).toBeLessThanOrEqual(1440);
|
||||
expect(tallResult.width).toBeLessThanOrEqual(1440);
|
||||
expect(tallResult.height).toBeLessThanOrEqual(1440);
|
||||
});
|
||||
});
|
||||
|
||||
describe('SD1 architecture', () => {
|
||||
it('should not resize for optimal dimensions', () => {
|
||||
const result = resizer.calculateTargetDimensions(512, 512, 'SD1');
|
||||
|
||||
expect(result).toEqual({
|
||||
width: 512,
|
||||
height: 512,
|
||||
needsResize: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should resize when exceeding maximum', () => {
|
||||
const result = resizer.calculateTargetDimensions(1000, 800, 'SD1');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeLessThanOrEqual(768);
|
||||
expect(result.height).toBeLessThanOrEqual(768);
|
||||
});
|
||||
|
||||
it('should resize when below minimum', () => {
|
||||
const result = resizer.calculateTargetDimensions(100, 200, 'SD1');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeGreaterThanOrEqual(256);
|
||||
expect(result.height).toBeGreaterThanOrEqual(256);
|
||||
});
|
||||
|
||||
it('should handle edge case at exact limits', () => {
|
||||
const maxResult = resizer.calculateTargetDimensions(768, 768, 'SD1');
|
||||
const minResult = resizer.calculateTargetDimensions(256, 256, 'SD1');
|
||||
|
||||
expect(maxResult.needsResize).toBe(false);
|
||||
expect(minResult.needsResize).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('SD3 architecture', () => {
|
||||
it('should not resize for optimal dimensions', () => {
|
||||
const result = resizer.calculateTargetDimensions(1024, 1024, 'SD3');
|
||||
|
||||
expect(result).toEqual({
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
needsResize: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should resize when exceeding maximum', () => {
|
||||
const result = resizer.calculateTargetDimensions(3000, 2000, 'SD3');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeLessThanOrEqual(2048);
|
||||
expect(result.height).toBeLessThanOrEqual(2048);
|
||||
});
|
||||
|
||||
it('should resize when below minimum', () => {
|
||||
const result = resizer.calculateTargetDimensions(300, 400, 'SD3');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeGreaterThanOrEqual(512);
|
||||
expect(result.height).toBeGreaterThanOrEqual(512);
|
||||
});
|
||||
});
|
||||
|
||||
describe('SDXL architecture', () => {
|
||||
it('should not resize for optimal dimensions', () => {
|
||||
const result = resizer.calculateTargetDimensions(1024, 1024, 'SDXL');
|
||||
|
||||
expect(result).toEqual({
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
needsResize: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should resize when exceeding maximum', () => {
|
||||
const result = resizer.calculateTargetDimensions(2500, 2500, 'SDXL');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeLessThanOrEqual(2048);
|
||||
expect(result.height).toBeLessThanOrEqual(2048);
|
||||
});
|
||||
|
||||
it('should resize when below minimum', () => {
|
||||
const result = resizer.calculateTargetDimensions(400, 300, 'SDXL');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeGreaterThanOrEqual(512);
|
||||
expect(result.height).toBeGreaterThanOrEqual(512);
|
||||
});
|
||||
});
|
||||
|
||||
describe('edge cases and error handling', () => {
|
||||
it('should handle zero dimensions', () => {
|
||||
expect(() => resizer.calculateTargetDimensions(0, 100, 'FLUX')).not.toThrow();
|
||||
expect(() => resizer.calculateTargetDimensions(100, 0, 'FLUX')).not.toThrow();
|
||||
expect(() => resizer.calculateTargetDimensions(0, 0, 'FLUX')).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle negative dimensions', () => {
|
||||
expect(() => resizer.calculateTargetDimensions(-100, 100, 'FLUX')).not.toThrow();
|
||||
expect(() => resizer.calculateTargetDimensions(100, -100, 'FLUX')).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle very large dimensions', () => {
|
||||
const result = resizer.calculateTargetDimensions(10000, 10000, 'FLUX');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeLessThanOrEqual(1440);
|
||||
expect(result.height).toBeLessThanOrEqual(1440);
|
||||
});
|
||||
|
||||
it('should handle very small dimensions', () => {
|
||||
const result = resizer.calculateTargetDimensions(1, 1, 'FLUX');
|
||||
|
||||
expect(result.needsResize).toBe(true);
|
||||
expect(result.width).toBeGreaterThanOrEqual(256);
|
||||
expect(result.height).toBeGreaterThanOrEqual(256);
|
||||
});
|
||||
|
||||
it('should handle unknown architecture gracefully', () => {
|
||||
expect(() => {
|
||||
resizer.calculateTargetDimensions(1024, 768, 'UNKNOWN' as Architecture);
|
||||
}).toThrow();
|
||||
});
|
||||
|
||||
it('should handle floating point dimensions', () => {
|
||||
const result = resizer.calculateTargetDimensions(1024.5, 768.7, 'FLUX');
|
||||
|
||||
// Results should be numbers (may be integers or floats depending on calculation)
|
||||
expect(typeof result.width).toBe('number');
|
||||
expect(typeof result.height).toBe('number');
|
||||
expect(Number.isFinite(result.width)).toBe(true);
|
||||
expect(Number.isFinite(result.height)).toBe(true);
|
||||
});
|
||||
|
||||
it('should maintain aspect ratio during scaling', () => {
|
||||
const originalAspectRatio = 16 / 9;
|
||||
const originalWidth = 1920;
|
||||
const originalHeight = 1080;
|
||||
|
||||
const result = resizer.calculateTargetDimensions(originalWidth, originalHeight, 'FLUX');
|
||||
|
||||
if (result.needsResize) {
|
||||
const newAspectRatio = result.width / result.height;
|
||||
expect(Math.abs(newAspectRatio - originalAspectRatio)).toBeLessThan(0.1);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('performance and consistency', () => {
|
||||
it('should be deterministic for same inputs', () => {
|
||||
const inputs = [
|
||||
[1024, 768, 'FLUX'],
|
||||
[512, 512, 'SD1'],
|
||||
[2048, 1024, 'SDXL'],
|
||||
] as const;
|
||||
|
||||
inputs.forEach(([width, height, arch]) => {
|
||||
const result1 = resizer.calculateTargetDimensions(width, height, arch);
|
||||
const result2 = resizer.calculateTargetDimensions(width, height, arch);
|
||||
|
||||
expect(result1).toEqual(result2);
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle rapid successive calls', () => {
|
||||
const results = [];
|
||||
|
||||
for (let i = 0; i < 1000; i++) {
|
||||
const result = resizer.calculateTargetDimensions(1024 + i, 768 + i, 'FLUX');
|
||||
results.push(result);
|
||||
}
|
||||
|
||||
expect(results).toHaveLength(1000);
|
||||
results.forEach((result) => {
|
||||
expect(typeof result.width).toBe('number');
|
||||
expect(typeof result.height).toBe('number');
|
||||
expect(typeof result.needsResize).toBe('boolean');
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('private method testing through public interface', () => {
|
||||
describe('calculateRatio method (through public interface)', () => {
|
||||
it('should calculate ratios correctly in various scenarios', () => {
|
||||
// Test through calculateTargetDimensions which uses calculateRatioFromDimensions
|
||||
const wideResult = resizer.calculateTargetDimensions(1600, 900, 'FLUX'); // 16:9
|
||||
const squareResult = resizer.calculateTargetDimensions(1000, 1000, 'FLUX'); // 1:1
|
||||
const tallResult = resizer.calculateTargetDimensions(900, 1600, 'FLUX'); // 9:16
|
||||
|
||||
expect(typeof wideResult).toBe('object');
|
||||
expect(typeof squareResult).toBe('object');
|
||||
expect(typeof tallResult).toBe('object');
|
||||
});
|
||||
});
|
||||
|
||||
describe('isWithinRatioRange method (through public interface)', () => {
|
||||
it('should handle ratio range validation through FLUX constraints', () => {
|
||||
// Test extreme ratios that should trigger ratio range checks
|
||||
const extremeWide = resizer.calculateTargetDimensions(3000, 1000, 'FLUX'); // 3:1 (exceeds 21:9)
|
||||
const extremeTall = resizer.calculateTargetDimensions(1000, 3000, 'FLUX'); // 1:3 (exceeds 9:21)
|
||||
|
||||
expect(extremeWide.needsResize).toBe(true);
|
||||
expect(extremeTall.needsResize).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getModelLimits method (through public interface)', () => {
|
||||
it('should use correct limits for each architecture', () => {
|
||||
// Test by checking if results respect known limits
|
||||
const fluxResult = resizer.calculateTargetDimensions(2000, 2000, 'FLUX');
|
||||
const sd1Result = resizer.calculateTargetDimensions(1000, 1000, 'SD1');
|
||||
const sd3Result = resizer.calculateTargetDimensions(3000, 3000, 'SD3');
|
||||
const sdxlResult = resizer.calculateTargetDimensions(3000, 3000, 'SDXL');
|
||||
|
||||
expect(fluxResult.width).toBeLessThanOrEqual(1440); // FLUX max
|
||||
expect(sd1Result.width).toBeLessThanOrEqual(768); // SD1 max
|
||||
expect(sd3Result.width).toBeLessThanOrEqual(2048); // SD3 max
|
||||
expect(sdxlResult.width).toBeLessThanOrEqual(2048); // SDXL max
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('singleton instance', () => {
|
||||
it('should export a singleton instance', () => {
|
||||
expect(imageResizer).toBeInstanceOf(ImageResizer);
|
||||
});
|
||||
|
||||
it('should maintain state across calls', () => {
|
||||
const result1 = imageResizer.calculateTargetDimensions(1024, 768, 'FLUX');
|
||||
const result2 = imageResizer.calculateTargetDimensions(1024, 768, 'FLUX');
|
||||
|
||||
expect(result1).toEqual(result2);
|
||||
});
|
||||
|
||||
it('should be the same instance when imported multiple times', () => {
|
||||
// This tests that the singleton pattern is working correctly
|
||||
expect(imageResizer).toBeDefined();
|
||||
expect(typeof imageResizer.calculateTargetDimensions).toBe('function');
|
||||
});
|
||||
});
|
||||
|
||||
describe('integration scenarios', () => {
|
||||
it('should handle common use cases correctly', () => {
|
||||
const commonSizes = [
|
||||
{ width: 512, height: 512, arch: 'SD1' as Architecture },
|
||||
{ width: 1024, height: 1024, arch: 'SDXL' as Architecture },
|
||||
{ width: 1440, height: 1024, arch: 'FLUX' as Architecture },
|
||||
{ width: 768, height: 512, arch: 'SD1' as Architecture },
|
||||
];
|
||||
|
||||
commonSizes.forEach(({ width, height, arch }) => {
|
||||
const result = resizer.calculateTargetDimensions(width, height, arch);
|
||||
|
||||
expect(result).toHaveProperty('width');
|
||||
expect(result).toHaveProperty('height');
|
||||
expect(result).toHaveProperty('needsResize');
|
||||
expect(typeof result.needsResize).toBe('boolean');
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle img2img workflow dimensions', () => {
|
||||
// Common img2img input sizes
|
||||
const img2imgSizes = [
|
||||
[1920, 1080], // Full HD
|
||||
[1280, 720], // HD
|
||||
[800, 600], // SVGA
|
||||
[640, 480], // VGA
|
||||
];
|
||||
|
||||
img2imgSizes.forEach(([width, height]) => {
|
||||
['FLUX', 'SDXL', 'SD3', 'SD1'].forEach((arch) => {
|
||||
const result = resizer.calculateTargetDimensions(width, height, arch as Architecture);
|
||||
|
||||
// All results should be valid
|
||||
expect(result.width).toBeGreaterThan(0);
|
||||
expect(result.height).toBeGreaterThan(0);
|
||||
expect(Number.isInteger(result.width)).toBe(true);
|
||||
expect(Number.isInteger(result.height)).toBe(true);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('boundary testing', () => {
|
||||
it('should handle boundary conditions for each architecture', () => {
|
||||
const architectures: Architecture[] = ['FLUX', 'SD1', 'SD3', 'SDXL'];
|
||||
|
||||
architectures.forEach((arch) => {
|
||||
// Test minimum boundary
|
||||
const minResult = resizer.calculateTargetDimensions(1, 1, arch);
|
||||
expect(minResult.needsResize).toBe(true);
|
||||
|
||||
// Test maximum boundary
|
||||
const maxResult = resizer.calculateTargetDimensions(10000, 10000, arch);
|
||||
expect(maxResult.needsResize).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle step size boundaries for FLUX', () => {
|
||||
// Test dimensions that are close to but not exactly divisible by 32
|
||||
const testCases = [
|
||||
[1023, 767], // Just below step boundary
|
||||
[1025, 769], // Just above step boundary
|
||||
[1056, 800], // Exactly on step boundary
|
||||
];
|
||||
|
||||
testCases.forEach(([width, height]) => {
|
||||
const result = resizer.calculateTargetDimensions(width, height, 'FLUX');
|
||||
|
||||
if (result.needsResize) {
|
||||
expect(result.width % 32).toBe(0);
|
||||
expect(result.height % 32).toBe(0);
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,191 @@
|
|||
// @vitest-environment node
|
||||
import { describe, expect, it } from 'vitest';
|
||||
|
||||
import { splitPromptForDualCLIP } from '@/server/services/comfyui/utils/promptSplitter';
|
||||
|
||||
describe('splitPromptForDualCLIP', () => {
|
||||
it('should handle empty or null prompt', () => {
|
||||
expect(splitPromptForDualCLIP('')).toEqual({
|
||||
clipLPrompt: '',
|
||||
t5xxlPrompt: '',
|
||||
});
|
||||
|
||||
expect(splitPromptForDualCLIP(null as any)).toEqual({
|
||||
clipLPrompt: '',
|
||||
t5xxlPrompt: '',
|
||||
});
|
||||
|
||||
expect(splitPromptForDualCLIP(undefined as any)).toEqual({
|
||||
clipLPrompt: '',
|
||||
t5xxlPrompt: '',
|
||||
});
|
||||
});
|
||||
|
||||
it('should split prompt with style keywords', () => {
|
||||
const prompt = 'a beautiful landscape, photorealistic, high quality, cinematic lighting';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('photorealistic');
|
||||
expect(result.clipLPrompt).toContain('high quality');
|
||||
expect(result.clipLPrompt).toContain('cinematic');
|
||||
});
|
||||
|
||||
it('should extract single-word style keywords', () => {
|
||||
const prompt = 'a cat sitting, realistic, detailed, masterpiece';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('realistic');
|
||||
expect(result.clipLPrompt).toContain('detailed');
|
||||
expect(result.clipLPrompt).toContain('masterpiece');
|
||||
expect(result.clipLPrompt).not.toContain('cat');
|
||||
expect(result.clipLPrompt).not.toContain('sitting');
|
||||
});
|
||||
|
||||
it('should extract multi-word style keywords', () => {
|
||||
const prompt = 'beautiful girl portrait, digital art, depth of field, trending on artstation';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('digital art');
|
||||
expect(result.clipLPrompt).toContain('depth of field');
|
||||
expect(result.clipLPrompt).toContain('trending on artstation');
|
||||
});
|
||||
|
||||
it('should handle lighting keywords', () => {
|
||||
const prompt = 'sunset over ocean, dramatic lighting, golden hour, soft lighting';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('dramatic lighting');
|
||||
expect(result.clipLPrompt).toContain('golden hour');
|
||||
expect(result.clipLPrompt).toContain('soft lighting');
|
||||
});
|
||||
|
||||
it('should handle quality keywords', () => {
|
||||
const prompt = 'mountain view, 4k, ultra detailed, best quality, highly detailed';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('4k');
|
||||
expect(result.clipLPrompt).toContain('ultra detailed');
|
||||
expect(result.clipLPrompt).toContain('best quality');
|
||||
expect(result.clipLPrompt).toContain('highly detailed');
|
||||
});
|
||||
|
||||
it('should handle photography terms', () => {
|
||||
const prompt = 'city street, bokeh, motion blur, wide angle, macro shot';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('bokeh');
|
||||
expect(result.clipLPrompt).toContain('motion blur');
|
||||
expect(result.clipLPrompt).toContain('wide angle');
|
||||
expect(result.clipLPrompt).toContain('macro');
|
||||
});
|
||||
|
||||
it('should handle artist and platform keywords', () => {
|
||||
const prompt = 'fantasy landscape, by greg rutkowski, concept art, octane render';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('by greg rutkowski');
|
||||
expect(result.clipLPrompt).toContain('concept art');
|
||||
expect(result.clipLPrompt).toContain('octane render');
|
||||
});
|
||||
|
||||
it('should fallback to adjectives when no style keywords found', () => {
|
||||
const prompt = 'a beautiful sunny day with colorful flowers blooming magnificently';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
// Should contain adjective-like words that match the regex pattern
|
||||
expect(result.clipLPrompt).toMatch(/blooming|magnificently|colorful|beautiful|sunny/);
|
||||
expect(result.clipLPrompt.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('should use same prompt for both when no style words or adjectives', () => {
|
||||
const prompt = 'cat dog house tree';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toBe(prompt);
|
||||
});
|
||||
|
||||
it('should preserve original case in style words', () => {
|
||||
const prompt = 'Portrait of girl, Photorealistic, High Quality, Digital Art';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('Photorealistic');
|
||||
expect(result.clipLPrompt).toContain('High Quality');
|
||||
expect(result.clipLPrompt).toContain('Digital Art');
|
||||
});
|
||||
|
||||
it('should handle comma-separated prompts', () => {
|
||||
const prompt = 'forest path, cinematic, dramatic lighting, 8k, masterpiece';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('cinematic');
|
||||
expect(result.clipLPrompt).toContain('dramatic lighting');
|
||||
expect(result.clipLPrompt).toContain('8k');
|
||||
expect(result.clipLPrompt).toContain('masterpiece');
|
||||
});
|
||||
|
||||
it('should handle partial keyword matches correctly', () => {
|
||||
const prompt = 'realistic portrait, photo-realistic style, realism art';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
// Should match "realistic" but exact behavior depends on implementation
|
||||
expect(result.clipLPrompt.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('should handle overlapping multi-word keywords', () => {
|
||||
const prompt = 'art gallery, digital art work, concept art design';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('digital art');
|
||||
expect(result.clipLPrompt).toContain('concept art');
|
||||
});
|
||||
|
||||
it('should work with very long prompts', () => {
|
||||
const prompt =
|
||||
'An incredibly detailed and photorealistic portrait of a young woman with flowing hair, sitting in a beautiful garden during golden hour, with soft lighting and dramatic shadows, rendered in 8k ultra high quality with perfect focus and depth of field, trending on artstation, masterpiece';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('photorealistic');
|
||||
expect(result.clipLPrompt).toContain('golden hour');
|
||||
expect(result.clipLPrompt).toContain('soft lighting');
|
||||
expect(result.clipLPrompt).toContain('8k');
|
||||
expect(result.clipLPrompt).toContain('depth of field');
|
||||
expect(result.clipLPrompt).toContain('trending on artstation');
|
||||
expect(result.clipLPrompt).toContain('masterpiece');
|
||||
});
|
||||
|
||||
it('should handle mixed content with various separators', () => {
|
||||
const prompt = 'sunset landscape; cinematic mood, soft lighting. 4k resolution!';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
expect(result.clipLPrompt).toContain('cinematic');
|
||||
// soft lighting might be treated as two separate words due to separator handling
|
||||
expect(result.clipLPrompt).toMatch(/soft|lighting|cinematic|4k/);
|
||||
});
|
||||
|
||||
it('should prioritize style keywords over content words', () => {
|
||||
const prompt = 'beautiful mountain landscape, photorealistic, detailed';
|
||||
const result = splitPromptForDualCLIP(prompt);
|
||||
|
||||
expect(result.t5xxlPrompt).toBe(prompt);
|
||||
// Should contain style keywords
|
||||
expect(result.clipLPrompt).toContain('photorealistic');
|
||||
expect(result.clipLPrompt).toContain('detailed');
|
||||
// The algorithm extracts style keywords first, so may not contain content words
|
||||
expect(result.clipLPrompt.length).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
192
src/server/services/comfyui/__tests__/utils/weightDType.test.ts
Normal file
192
src/server/services/comfyui/__tests__/utils/weightDType.test.ts
Normal file
|
|
@ -0,0 +1,192 @@
|
|||
// @vitest-environment node
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { selectOptimalWeightDtype } from '@/server/services/comfyui/utils/weightDType';
|
||||
|
||||
// Mock the modelRegistry module
|
||||
vi.mock('@/server/services/comfyui/config/modelRegistry', () => {
|
||||
const models = {
|
||||
'flux1-dev-fp8-e4m3fn.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
variant: 'flux1-dev-fp8-e4m3fn',
|
||||
},
|
||||
'flux1-dev.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'default',
|
||||
variant: 'flux1-dev',
|
||||
},
|
||||
'flux1-kontext-dev.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'default',
|
||||
variant: 'flux1-kontext-dev',
|
||||
},
|
||||
'flux1-schnell-fp8-e4m3fn.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
variant: 'flux1-schnell-fp8-e4m3fn',
|
||||
},
|
||||
'flux1-schnell.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'default',
|
||||
variant: 'flux1-schnell',
|
||||
},
|
||||
'vision_realistic_flux_dev_fp8_no_clip_v2.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
variant: 'vision_realistic_flux_dev_fp8_no_clip_v2',
|
||||
},
|
||||
};
|
||||
|
||||
return {
|
||||
MODEL_ID_VARIANT_MAP: {
|
||||
'flux-dev': 'flux1-dev',
|
||||
'flux-schnell': 'flux1-schnell',
|
||||
},
|
||||
MODEL_REGISTRY: models,
|
||||
};
|
||||
});
|
||||
|
||||
// Mock the staticModelLookup module
|
||||
vi.mock('../utils/staticModelLookup', () => {
|
||||
const models = {
|
||||
'flux1-dev-fp8-e4m3fn.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
variant: 'flux1-dev-fp8-e4m3fn',
|
||||
},
|
||||
'flux1-dev.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'default',
|
||||
variant: 'flux1-dev',
|
||||
},
|
||||
'flux1-kontext-dev.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'default',
|
||||
variant: 'flux1-kontext-dev',
|
||||
},
|
||||
'flux1-schnell-fp8-e4m3fn.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
variant: 'flux1-schnell-fp8-e4m3fn',
|
||||
},
|
||||
'flux1-schnell.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'default',
|
||||
variant: 'flux1-schnell',
|
||||
},
|
||||
'vision_realistic_flux_dev_fp8_no_clip_v2.safetensors': {
|
||||
family: 'flux',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
variant: 'vision_realistic_flux_dev_fp8_no_clip_v2',
|
||||
},
|
||||
};
|
||||
|
||||
return {
|
||||
resolveModel: vi.fn((modelName: string) => {
|
||||
const cleanName = modelName.replace(/^comfyui\//, '');
|
||||
|
||||
// Case-insensitive lookup
|
||||
const lowerModelName = cleanName.toLowerCase();
|
||||
for (const [key, config] of Object.entries(models)) {
|
||||
if (key.toLowerCase() === lowerModelName) {
|
||||
return config;
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}),
|
||||
};
|
||||
});
|
||||
|
||||
describe('selectOptimalWeightDtype', () => {
|
||||
it('should return model recommendedDtype for known FLUX models', () => {
|
||||
// FLUX Dev models should use default for quality
|
||||
expect(selectOptimalWeightDtype('flux1-dev.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_dev.safetensors')).toBe('default');
|
||||
|
||||
// FLUX Schnell models use default in current registry (fps8 variants have separate entries)
|
||||
expect(selectOptimalWeightDtype('flux1-schnell.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_schnell.safetensors')).toBe('default'); // Not in registry
|
||||
|
||||
// FLUX Kontext models should use default
|
||||
expect(selectOptimalWeightDtype('flux1-kontext-dev.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_kontext.safetensors')).toBe('default');
|
||||
|
||||
// FLUX Krea models should use default
|
||||
expect(selectOptimalWeightDtype('flux_krea.safetensors')).toBe('default');
|
||||
});
|
||||
|
||||
it('should return default for GGUF models', () => {
|
||||
expect(selectOptimalWeightDtype('flux1-dev-Q4_K_S.gguf')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux1-schnell-Q6_K.gguf')).toBe('default');
|
||||
});
|
||||
|
||||
it('should return correct dtype for quantized models that exist in registry', () => {
|
||||
// FP8 quantized models that exist in the registry with exact names
|
||||
expect(selectOptimalWeightDtype('flux1-dev-fp8-e4m3fn.safetensors')).toBe('fp8_e4m3fn');
|
||||
expect(selectOptimalWeightDtype('flux1-schnell-fp8-e4m3fn.safetensors')).toBe('fp8_e4m3fn');
|
||||
|
||||
// Models with approximate names that don't exactly match registry return default
|
||||
expect(selectOptimalWeightDtype('flux1-dev-fp8.safetensors')).toBe('default'); // Not exact match
|
||||
});
|
||||
|
||||
it('should return default for enterprise lite models', () => {
|
||||
expect(selectOptimalWeightDtype('flux.1-lite-8B.safetensors')).toBe('default');
|
||||
});
|
||||
|
||||
it('should return default fallback for unknown models', () => {
|
||||
expect(selectOptimalWeightDtype('unknown_model.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('custom_flux.bin')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('not_a_flux_model.ckpt')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('model.pt')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('weird@model&name.safetensors')).toBe('default');
|
||||
|
||||
// Models with precision in filename but not in registry fall back to default
|
||||
expect(selectOptimalWeightDtype('flux_model_fp32.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_model_fp16.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_model_int8.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_model_int4.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_model_nf4.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_model_bnb.safetensors')).toBe('default');
|
||||
});
|
||||
|
||||
it('should be case-insensitive for model detection', () => {
|
||||
expect(selectOptimalWeightDtype('FLUX1-DEV.SAFETENSORS')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('FLUX1-SCHNELL.SAFETENSORS')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('FLUX1-DEV-FP8-E4M3FN.SAFETENSORS')).toBe('fp8_e4m3fn');
|
||||
});
|
||||
|
||||
it('should handle community models correctly', () => {
|
||||
// Most community models will fall back to default unless specifically in registry
|
||||
expect(selectOptimalWeightDtype('Jib_mix_Flux_V11_Krea_b_00001_.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('RealFlux_1.0b_Dev_Transformer.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('RealFlux_1.0b_Schnell_Transformer.safetensors')).toBe(
|
||||
'default',
|
||||
);
|
||||
expect(selectOptimalWeightDtype('Jib_Mix_Flux_Krea_b_fp8_00001_.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('vision_realistic_flux_dev_fp8_no_clip_v2.safetensors')).toBe(
|
||||
'fp8_e4m3fn', // This model is actually in the registry
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle edge cases', () => {
|
||||
expect(selectOptimalWeightDtype('flux_model')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux_model.')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('.gguf')).toBe('default');
|
||||
});
|
||||
|
||||
describe('simplified logic without user parameters', () => {
|
||||
it('should only accept modelName parameter', () => {
|
||||
// Function now only takes modelName parameter
|
||||
// JavaScript allows extra parameters, so this won't throw, but TypeScript will catch it
|
||||
expect(selectOptimalWeightDtype('flux1-dev.safetensors')).toBe('default');
|
||||
});
|
||||
|
||||
it('should always use model-based selection', () => {
|
||||
// No user choice - always use model configuration or default fallback
|
||||
expect(selectOptimalWeightDtype('flux1-dev.safetensors')).toBe('default');
|
||||
expect(selectOptimalWeightDtype('flux1-schnell.safetensors')).toBe('default'); // Base model uses default
|
||||
expect(selectOptimalWeightDtype('unknown_model.safetensors')).toBe('default');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,507 @@
|
|||
import { type Mock, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import type { ModelConfig } from '@/server/services/comfyui/config/modelRegistry';
|
||||
import { resolveModel } from '@/server/services/comfyui/utils/staticModelLookup';
|
||||
import {
|
||||
type SD3Variant,
|
||||
WorkflowDetector,
|
||||
} from '@/server/services/comfyui/utils/workflowDetector';
|
||||
|
||||
// Mock static model lookup functions
|
||||
vi.mock('../../utils/staticModelLookup', () => ({
|
||||
resolveModel: vi.fn(),
|
||||
getModelConfig: vi.fn(),
|
||||
}));
|
||||
|
||||
describe('WorkflowDetector', () => {
|
||||
const mockedResolveModel = resolveModel as Mock;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('detectModelType', () => {
|
||||
describe('Input Processing', () => {
|
||||
it('should remove "comfyui/" prefix from modelId', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'dev',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('comfyui/flux-dev');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('flux-dev');
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'dev',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle modelId without comfyui prefix', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'schnell',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('flux-schnell');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('flux-schnell');
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'schnell',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle multiple comfyui prefixes correctly', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'SD3',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'sd35',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
// Only the first "comfyui/" should be removed
|
||||
const result = WorkflowDetector.detectModelType('comfyui/comfyui/model');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('comfyui/model');
|
||||
expect(result).toEqual({
|
||||
architecture: 'SD3',
|
||||
isSupported: true,
|
||||
variant: 'sd35',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('FLUX Model Detection', () => {
|
||||
it('should detect FLUX dev variant', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'dev',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('flux-dev');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'dev',
|
||||
});
|
||||
});
|
||||
|
||||
it('should detect FLUX schnell variant', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
variant: 'schnell',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('flux-schnell-fp8');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'schnell',
|
||||
});
|
||||
});
|
||||
|
||||
it('should detect FLUX kontext variant', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'kontext',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('flux-kontext-dev');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'kontext',
|
||||
});
|
||||
});
|
||||
|
||||
it('should detect FLUX krea model with dev variant', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'dev',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('flux-krea-dev');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'dev',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle FLUX model with comfyui prefix', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
recommendedDtype: 'fp8_e5m2',
|
||||
variant: 'dev',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('comfyui/custom-flux-model');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('custom-flux-model');
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'dev',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Custom SD Model Detection', () => {
|
||||
it('should detect custom SD model', () => {
|
||||
const result = WorkflowDetector.detectModelType('stable-diffusion-custom');
|
||||
|
||||
// Custom SD models are hardcoded and don't use resolveModel
|
||||
expect(mockedResolveModel).not.toHaveBeenCalled();
|
||||
expect(result).toEqual({
|
||||
architecture: 'SDXL', // Uses SDXL for img2img support
|
||||
isSupported: true,
|
||||
variant: 'custom-sd',
|
||||
});
|
||||
});
|
||||
|
||||
it('should detect custom SD refiner model', () => {
|
||||
const result = WorkflowDetector.detectModelType('stable-diffusion-custom-refiner');
|
||||
|
||||
// Custom SD models are hardcoded and don't use resolveModel
|
||||
expect(mockedResolveModel).not.toHaveBeenCalled();
|
||||
expect(result).toEqual({
|
||||
architecture: 'SDXL', // Uses SDXL for img2img support
|
||||
isSupported: true,
|
||||
variant: 'custom-sd',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle custom SD with comfyui prefix', () => {
|
||||
const result = WorkflowDetector.detectModelType('comfyui/stable-diffusion-custom');
|
||||
|
||||
// Custom SD models are hardcoded and don't use resolveModel
|
||||
expect(mockedResolveModel).not.toHaveBeenCalled();
|
||||
expect(result).toEqual({
|
||||
architecture: 'SDXL', // Uses SDXL for img2img support
|
||||
isSupported: true,
|
||||
variant: 'custom-sd',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('SD3 Model Detection', () => {
|
||||
it('should detect SD3 sd35 variant', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'SD3',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'sd35',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('sd3.5_large');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'SD3',
|
||||
isSupported: true,
|
||||
variant: 'sd35',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle SD3 model with comfyui prefix', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'SD3',
|
||||
priority: 2,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'sd35',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('comfyui/sd3.5_medium');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('sd3.5_medium');
|
||||
expect(result).toEqual({
|
||||
architecture: 'SD3',
|
||||
isSupported: true,
|
||||
variant: 'sd35',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Unknown/Unsupported Model Detection', () => {
|
||||
it('should return unknown architecture when model is not found', () => {
|
||||
mockedResolveModel.mockReturnValue(null);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('unknown-model');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'unknown',
|
||||
isSupported: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should return SDXL architecture for SDXL model family', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'SDXL' as any,
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'sdxl-t2i',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('sdxl-base');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'SDXL',
|
||||
isSupported: true,
|
||||
variant: 'sdxl-t2i',
|
||||
});
|
||||
});
|
||||
|
||||
it('should return SD1 architecture for SD1 model family', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'SD1' as any,
|
||||
priority: 3,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'sd15-t2i',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('stable-diffusion-v1-5');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'SD1',
|
||||
isSupported: true,
|
||||
variant: 'sd15-t2i',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle null modelId by causing runtime error (expected behavior)', () => {
|
||||
// According to the function signature, modelId is expected to be a string
|
||||
// Passing null/undefined would cause a runtime error, which is expected behavior
|
||||
expect(() => {
|
||||
WorkflowDetector.detectModelType(null as any);
|
||||
}).toThrow('Cannot read properties of null');
|
||||
});
|
||||
|
||||
it('should handle undefined modelId by causing runtime error (expected behavior)', () => {
|
||||
// According to the function signature, modelId is expected to be a string
|
||||
// Passing null/undefined would cause a runtime error, which is expected behavior
|
||||
expect(() => {
|
||||
WorkflowDetector.detectModelType(undefined as any);
|
||||
}).toThrow('Cannot read properties of undefined');
|
||||
});
|
||||
|
||||
it('should handle empty string modelId', () => {
|
||||
mockedResolveModel.mockReturnValue(null);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('');
|
||||
expect(result).toEqual({
|
||||
architecture: 'unknown',
|
||||
isSupported: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle whitespace-only modelId', () => {
|
||||
mockedResolveModel.mockReturnValue(null);
|
||||
|
||||
const result = WorkflowDetector.detectModelType(' ');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith(' ');
|
||||
expect(result).toEqual({
|
||||
architecture: 'unknown',
|
||||
isSupported: false,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Type Casting', () => {
|
||||
it('should properly cast FLUX variant to FluxVariant type', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'dev',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('flux-model');
|
||||
|
||||
expect(result.variant).toBe('dev');
|
||||
expect(typeof result.variant).toBe('string');
|
||||
|
||||
// Test with dev variant (krea uses dev workflow)
|
||||
const mockKreaConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'dev',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockKreaConfig);
|
||||
|
||||
const kreaResult = WorkflowDetector.detectModelType('flux-krea-model');
|
||||
expect(kreaResult.variant).toBe('dev');
|
||||
});
|
||||
|
||||
it('should properly cast SD3 variant to SD3Variant type', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'SD3',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'sd35',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('sd3-model');
|
||||
|
||||
expect(result.variant).toBe('sd35');
|
||||
expect(typeof result.variant).toBe('string');
|
||||
|
||||
// Verify it matches SD3Variant type expectations
|
||||
const sd3Variants: SD3Variant[] = ['sd35'];
|
||||
expect(sd3Variants).toContain(result.variant as SD3Variant);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Edge Cases', () => {
|
||||
it('should handle special characters in modelId', () => {
|
||||
mockedResolveModel.mockReturnValue(null);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('model-with-special!@#$%^&*()_+');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('model-with-special!@#$%^&*()_+');
|
||||
expect(result).toEqual({
|
||||
architecture: 'unknown',
|
||||
isSupported: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle modelId with path separators', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'dev',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('path/to/model.safetensors');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('path/to/model.safetensors');
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'dev',
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle very long modelId', () => {
|
||||
const longModelId = 'a'.repeat(1000);
|
||||
mockedResolveModel.mockReturnValue(null);
|
||||
|
||||
const result = WorkflowDetector.detectModelType(longModelId);
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith(longModelId);
|
||||
expect(result).toEqual({
|
||||
architecture: 'unknown',
|
||||
isSupported: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle modelId that is only "comfyui/"', () => {
|
||||
mockedResolveModel.mockReturnValue(null);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('comfyui/');
|
||||
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('');
|
||||
expect(result).toEqual({
|
||||
architecture: 'unknown',
|
||||
isSupported: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle case sensitivity in modelId', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: 'dev',
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('COMFYUI/FLUX-DEV');
|
||||
|
||||
// Should not match the prefix replacement since it's case sensitive
|
||||
expect(mockedResolveModel).toHaveBeenCalledWith('COMFYUI/FLUX-DEV');
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: 'dev',
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('Configuration Edge Cases', () => {
|
||||
it('should handle config with missing variant property', () => {
|
||||
const mockConfig: Partial<ModelConfig> = {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
// variant is missing
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig as ModelConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('flux-model');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'FLUX',
|
||||
isSupported: true,
|
||||
variant: undefined, // Will be cast to FluxVariant but is undefined
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle config with null variant', () => {
|
||||
const mockConfig: ModelConfig = {
|
||||
modelFamily: 'SD3',
|
||||
priority: 1,
|
||||
recommendedDtype: 'default',
|
||||
variant: null as any,
|
||||
};
|
||||
mockedResolveModel.mockReturnValue(mockConfig);
|
||||
|
||||
const result = WorkflowDetector.detectModelType('sd3-model');
|
||||
|
||||
expect(result).toEqual({
|
||||
architecture: 'SD3',
|
||||
isSupported: true,
|
||||
variant: null, // Will be cast to SD3Variant but is null
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,381 @@
|
|||
// @vitest-environment node
|
||||
import { PromptBuilder } from '@saintno/comfyui-sdk';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import { TEST_FLUX_MODELS } from '@/server/services/comfyui/__tests__/fixtures/testModels';
|
||||
import { mockContext } from '@/server/services/comfyui/__tests__/helpers/mockContext';
|
||||
import { setupAllMocks } from '@/server/services/comfyui/__tests__/setup/unifiedMocks';
|
||||
import { buildFluxKontextWorkflow } from '@/server/services/comfyui/workflows/flux-kontext';
|
||||
|
||||
// Setup basic mocks
|
||||
vi.mock('../utils/promptSplitter', () => ({
|
||||
splitPromptForDualCLIP: vi.fn((prompt) => ({
|
||||
clipLPrompt: prompt,
|
||||
t5xxlPrompt: prompt,
|
||||
})),
|
||||
}));
|
||||
vi.mock('../utils/weightDType', () => ({
|
||||
selectOptimalWeightDtype: vi.fn(() => 'default'),
|
||||
}));
|
||||
vi.mock('@lobechat/utils', () => ({
|
||||
generateUniqueSeeds: vi.fn(() => ({ seed: 123456, noiseSeed: 654321 })),
|
||||
}));
|
||||
vi.mock('../utils/workflowUtils', () => ({
|
||||
getWorkflowFilenamePrefix: vi.fn(() => 'kontext'),
|
||||
}));
|
||||
|
||||
const { inputCalls } = setupAllMocks();
|
||||
|
||||
describe('buildFluxKontextWorkflow - Complex Dual-Mode Architecture', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('Dual-Mode Architecture Tests', () => {
|
||||
it('should build text-to-image workflow when no input image provided', async () => {
|
||||
const modelName = TEST_FLUX_MODELS.KONTEXT;
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: 'A beautiful landscape',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(modelName, params, mockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Verify text-to-image mode: no image loader nodes
|
||||
expect(mockContext.modelResolverService.getOptimalComponent).toHaveBeenCalledWith(
|
||||
't5',
|
||||
'FLUX',
|
||||
);
|
||||
expect(mockContext.modelResolverService.getOptimalComponent).toHaveBeenCalledWith(
|
||||
'vae',
|
||||
'FLUX',
|
||||
);
|
||||
expect(mockContext.modelResolverService.getOptimalComponent).toHaveBeenCalledWith(
|
||||
'clip',
|
||||
'FLUX',
|
||||
);
|
||||
});
|
||||
|
||||
it('should build image-to-image workflow when input image provided', async () => {
|
||||
const modelName = TEST_FLUX_MODELS.KONTEXT;
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
imageUrl: 'https://example.com/input.jpg',
|
||||
prompt: 'Transform this image',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(modelName, params, mockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Verify image-to-image mode configuration
|
||||
expect(mockContext.modelResolverService.getOptimalComponent).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
|
||||
it('should handle imageUrls array parameter', async () => {
|
||||
const modelName = TEST_FLUX_MODELS.KONTEXT;
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
imageUrls: ['https://example.com/input1.jpg', 'https://example.com/input2.jpg'],
|
||||
prompt: 'Process first image',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(modelName, params, mockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Should use first image from array for i2i mode
|
||||
});
|
||||
});
|
||||
|
||||
describe('Dynamic Node Management Tests', () => {
|
||||
it('should create workflow with all required nodes for t2i mode', async () => {
|
||||
const modelName = TEST_FLUX_MODELS.KONTEXT;
|
||||
const params = {
|
||||
cfg: 4.0,
|
||||
height: 768,
|
||||
prompt: 'Dynamic node test',
|
||||
steps: 28,
|
||||
width: 768,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(modelName, params, mockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Verify essential nodes are properly connected
|
||||
expect(mockContext.modelResolverService.getOptimalComponent).toHaveBeenCalledWith(
|
||||
't5',
|
||||
'FLUX',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle different CFG values for guidance', async () => {
|
||||
const testCases = [
|
||||
{ cfg: 1.0, expected: 'minimal guidance' },
|
||||
{ cfg: 3.5, expected: 'default guidance' },
|
||||
{ cfg: 7.0, expected: 'high guidance' },
|
||||
];
|
||||
|
||||
for (const testCase of testCases) {
|
||||
const params = {
|
||||
cfg: testCase.cfg,
|
||||
height: 1024,
|
||||
prompt: `Test with ${testCase.expected}`,
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(
|
||||
TEST_FLUX_MODELS.KONTEXT,
|
||||
params,
|
||||
mockContext,
|
||||
);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Component Integration Tests', () => {
|
||||
it('should integrate with GetImageSize for dynamic dimensions', async () => {
|
||||
const modelName = TEST_FLUX_MODELS.KONTEXT;
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
imageUrl: 'https://example.com/variable-size.jpg',
|
||||
prompt: 'Resize based on input',
|
||||
steps: 28,
|
||||
// Note: height/width should be dynamically determined by GetImageSize
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(modelName, params, mockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Verify GetImageSize integration would be handled
|
||||
});
|
||||
|
||||
it('should handle component resolution failures gracefully', async () => {
|
||||
// Mock component resolution failure
|
||||
const failingContext = {
|
||||
...mockContext,
|
||||
modelResolverService: {
|
||||
...mockContext.modelResolverService,
|
||||
getOptimalComponent: vi.fn().mockRejectedValue(new Error('Component not found')),
|
||||
} as any,
|
||||
};
|
||||
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: 'Test component failure',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
await expect(
|
||||
buildFluxKontextWorkflow(TEST_FLUX_MODELS.KONTEXT, params, failingContext),
|
||||
).rejects.toThrow('Component not found');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Parameter Validation and Edge Cases', () => {
|
||||
it('should handle empty prompt', async () => {
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: '',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(TEST_FLUX_MODELS.KONTEXT, params, mockContext);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
});
|
||||
|
||||
it('should handle custom dimensions', async () => {
|
||||
const customDimensions = [
|
||||
{ width: 512, height: 768 }, // Portrait
|
||||
{ width: 1152, height: 896 }, // Landscape
|
||||
{ width: 896, height: 896 }, // Square
|
||||
];
|
||||
|
||||
for (const dims of customDimensions) {
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: dims.height,
|
||||
prompt: `Test ${dims.width}x${dims.height}`,
|
||||
steps: 28,
|
||||
width: dims.width,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(
|
||||
TEST_FLUX_MODELS.KONTEXT,
|
||||
params,
|
||||
mockContext,
|
||||
);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle different step counts', async () => {
|
||||
const stepCounts = [20, 28, 35, 50];
|
||||
|
||||
for (const steps of stepCounts) {
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: `Test with ${steps} steps`,
|
||||
steps,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(
|
||||
TEST_FLUX_MODELS.KONTEXT,
|
||||
params,
|
||||
mockContext,
|
||||
);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Advanced Feature Tests', () => {
|
||||
it('should support prompt splitting for dual CLIP', async () => {
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: 'Complex prompt requiring dual CLIP processing',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(TEST_FLUX_MODELS.KONTEXT, params, mockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Mock function should be called (tested via workflow execution)
|
||||
});
|
||||
|
||||
it('should handle weight dtype optimization', async () => {
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: 'Weight dtype test',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(TEST_FLUX_MODELS.KONTEXT, params, mockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Mock function should be called (tested via workflow execution)
|
||||
});
|
||||
|
||||
it('should generate unique seeds for workflow', async () => {
|
||||
const { generateUniqueSeeds } = await import('@lobechat/utils');
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: 'Seed generation test',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
await buildFluxKontextWorkflow(TEST_FLUX_MODELS.KONTEXT, params, mockContext);
|
||||
|
||||
expect(generateUniqueSeeds).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Complex Workflow Architecture', () => {
|
||||
it('should handle 28-step workflow complexity', async () => {
|
||||
const params = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: 'Complex 28-step workflow test',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildFluxKontextWorkflow(TEST_FLUX_MODELS.KONTEXT, params, mockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Verify proper step configuration
|
||||
});
|
||||
|
||||
it('should maintain node connection integrity across modes', async () => {
|
||||
// Test both modes to ensure consistent node connections
|
||||
const baseParams = {
|
||||
cfg: 3.5,
|
||||
height: 1024,
|
||||
prompt: 'Connection integrity test',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
// Test t2i mode
|
||||
const t2iResult = await buildFluxKontextWorkflow(
|
||||
TEST_FLUX_MODELS.KONTEXT,
|
||||
baseParams,
|
||||
mockContext,
|
||||
);
|
||||
expect(t2iResult).toHaveProperty('input');
|
||||
expect(t2iResult).toHaveProperty('setInputNode');
|
||||
expect(t2iResult).toHaveProperty('setOutputNode');
|
||||
expect(t2iResult).toHaveProperty('workflow');
|
||||
|
||||
// Test i2i mode
|
||||
const i2iParams = { ...baseParams, imageUrl: 'https://example.com/test.jpg' };
|
||||
const i2iResult = await buildFluxKontextWorkflow(
|
||||
TEST_FLUX_MODELS.KONTEXT,
|
||||
i2iParams,
|
||||
mockContext,
|
||||
);
|
||||
expect(i2iResult).toHaveProperty('input');
|
||||
expect(i2iResult).toHaveProperty('setInputNode');
|
||||
expect(i2iResult).toHaveProperty('setOutputNode');
|
||||
expect(i2iResult).toHaveProperty('workflow');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,558 @@
|
|||
// @vitest-environment node
|
||||
import { PromptBuilder } from '@saintno/comfyui-sdk';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import {
|
||||
TEST_CUSTOM_SD,
|
||||
TEST_SD35_MODELS,
|
||||
TEST_SDXL_MODELS,
|
||||
} from '@/server/services/comfyui/__tests__/fixtures/testModels';
|
||||
import { mockContext } from '@/server/services/comfyui/__tests__/helpers/mockContext';
|
||||
import { setupAllMocks } from '@/server/services/comfyui/__tests__/setup/unifiedMocks';
|
||||
import { buildSimpleSDWorkflow } from '@/server/services/comfyui/workflows/simple-sd';
|
||||
|
||||
// Setup basic mocks
|
||||
vi.mock('@lobechat/utils', () => ({
|
||||
generateUniqueSeeds: vi.fn(() => ({ seed: 123456, noiseSeed: 654321 })),
|
||||
}));
|
||||
vi.mock('../utils/workflowUtils', () => ({
|
||||
getWorkflowFilenamePrefix: vi.fn(() => 'simple-sd'),
|
||||
}));
|
||||
vi.mock('../utils/staticModelLookup', () => ({
|
||||
getModelConfig: vi.fn((modelName: string) => {
|
||||
// Mock model configuration mapping
|
||||
if (modelName.includes('sd3.5') || modelName.includes('sd35')) {
|
||||
return {
|
||||
modelFamily: 'SD3',
|
||||
variant: 'medium',
|
||||
};
|
||||
}
|
||||
if (modelName.includes('sdxl') || modelName.includes('xl')) {
|
||||
return {
|
||||
modelFamily: 'SDXL',
|
||||
variant: 'base',
|
||||
};
|
||||
}
|
||||
if (modelName.includes('sd1') || modelName.includes('v1')) {
|
||||
return {
|
||||
modelFamily: 'SD1',
|
||||
variant: '5',
|
||||
};
|
||||
}
|
||||
if (modelName === TEST_CUSTOM_SD) {
|
||||
return {
|
||||
modelFamily: 'SD1',
|
||||
variant: 'custom',
|
||||
};
|
||||
}
|
||||
return null;
|
||||
}),
|
||||
}));
|
||||
|
||||
const { inputCalls } = setupAllMocks();
|
||||
|
||||
// Extended mock context for SD testing
|
||||
const createSDMockContext = () => ({
|
||||
...mockContext,
|
||||
modelResolverService: {
|
||||
...mockContext.modelResolverService,
|
||||
getAvailableVAEFiles: vi
|
||||
.fn()
|
||||
.mockResolvedValue([
|
||||
'vae-ft-mse-840000-ema-pruned.safetensors',
|
||||
'sdxl_vae_fp16fix.safetensors',
|
||||
'custom_sd_lobe_vae.safetensors',
|
||||
]),
|
||||
getOptimalComponent: vi.fn().mockImplementation((type: string, modelFamily: string) => {
|
||||
if (type === 'vae') {
|
||||
if (modelFamily === 'SDXL') {
|
||||
return Promise.resolve('sdxl_vae_fp16fix.safetensors');
|
||||
}
|
||||
if (modelFamily === 'SD1') {
|
||||
return Promise.resolve('vae-ft-mse-840000-ema-pruned.safetensors');
|
||||
}
|
||||
// SD3 models have built-in VAE
|
||||
return Promise.resolve(undefined);
|
||||
}
|
||||
return Promise.resolve(null);
|
||||
}),
|
||||
},
|
||||
});
|
||||
|
||||
describe('buildSimpleSDWorkflow - Universal SD Support', () => {
|
||||
let sdMockContext: any;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
sdMockContext = createSDMockContext();
|
||||
});
|
||||
|
||||
describe('Model Family Detection Tests', () => {
|
||||
it('should detect SD3.5 model family', async () => {
|
||||
const modelName = TEST_SD35_MODELS.MEDIUM;
|
||||
const params = {
|
||||
cfg: 4.0,
|
||||
height: 1024,
|
||||
prompt: 'SD3.5 family test',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// SD3 models shouldn't need external VAE
|
||||
});
|
||||
|
||||
it('should detect SDXL model family', async () => {
|
||||
const modelName = TEST_SDXL_MODELS.BASE;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: 'SDXL family test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// SDXL models should use external VAE
|
||||
expect(sdMockContext.modelResolverService.getOptimalComponent).toHaveBeenCalledWith(
|
||||
'vae',
|
||||
'SDXL',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle custom SD model', async () => {
|
||||
const modelName = TEST_CUSTOM_SD;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 512,
|
||||
prompt: 'Custom SD model test',
|
||||
steps: 20,
|
||||
width: 512,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Custom models should check for custom VAE
|
||||
expect(sdMockContext.modelResolverService.getAvailableVAEFiles).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Smart VAE Selection Tests', () => {
|
||||
it('should not attach VAE for SD3 models (built-in VAE)', async () => {
|
||||
const modelName = TEST_SD35_MODELS.LARGE;
|
||||
const params = {
|
||||
cfg: 4.0,
|
||||
height: 1024,
|
||||
prompt: 'SD3 built-in VAE test',
|
||||
steps: 28,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// SD3 should not request external VAE
|
||||
});
|
||||
|
||||
it('should attach optimal VAE for SDXL models', async () => {
|
||||
const modelName = TEST_SDXL_MODELS.BASE;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: 'SDXL VAE selection test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
expect(sdMockContext.modelResolverService.getOptimalComponent).toHaveBeenCalledWith(
|
||||
'vae',
|
||||
'SDXL',
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle custom VAE for custom models', async () => {
|
||||
const modelName = TEST_CUSTOM_SD;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 512,
|
||||
prompt: 'Custom VAE test',
|
||||
steps: 20,
|
||||
width: 512,
|
||||
};
|
||||
|
||||
// Mock custom VAE availability
|
||||
sdMockContext.modelResolverService.getAvailableVAEFiles.mockResolvedValue([
|
||||
'custom_sd_lobe_vae.safetensors',
|
||||
'vae-ft-mse-840000-ema-pruned.safetensors',
|
||||
]);
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
expect(sdMockContext.modelResolverService.getAvailableVAEFiles).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should fallback to built-in VAE when custom VAE unavailable', async () => {
|
||||
const modelName = TEST_CUSTOM_SD;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 512,
|
||||
prompt: 'VAE fallback test',
|
||||
steps: 20,
|
||||
width: 512,
|
||||
};
|
||||
|
||||
// Mock custom VAE not available
|
||||
sdMockContext.modelResolverService.getAvailableVAEFiles.mockResolvedValue([
|
||||
'vae-ft-mse-840000-ema-pruned.safetensors',
|
||||
]);
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Should fall back to built-in VAE
|
||||
});
|
||||
});
|
||||
|
||||
describe('Dual Mode Support Tests', () => {
|
||||
it('should build text-to-image workflow', async () => {
|
||||
const modelName = TEST_SDXL_MODELS.BASE;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: 'Text to image test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Should be in t2i mode (no input image)
|
||||
});
|
||||
|
||||
it('should build image-to-image workflow', async () => {
|
||||
const modelName = TEST_SDXL_MODELS.BASE;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
denoise: 0.75,
|
||||
height: 1024,
|
||||
imageUrl: 'https://example.com/input.jpg',
|
||||
prompt: 'Image to image test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Should be in i2i mode with input image
|
||||
});
|
||||
|
||||
it('should handle strength parameter mapping to denoise', async () => {
|
||||
const modelName = TEST_SDXL_MODELS.BASE;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
imageUrl: 'https://example.com/input.jpg',
|
||||
prompt: 'Strength mapping test',
|
||||
steps: 20,
|
||||
strength: 0.8, // Frontend parameter
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Strength should be mapped to denoise internally
|
||||
});
|
||||
|
||||
it('should handle imageUrls array parameter', async () => {
|
||||
const modelName = TEST_SDXL_MODELS.BASE;
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
imageUrls: ['https://example.com/input1.jpg', 'https://example.com/input2.jpg'],
|
||||
prompt: 'Multiple images test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Should use first image from array
|
||||
});
|
||||
});
|
||||
|
||||
describe('Parameter Validation Tests', () => {
|
||||
it('should handle different CFG values by model family', async () => {
|
||||
const testCases = [
|
||||
{ model: TEST_SD35_MODELS.MEDIUM, cfg: 4.0, family: 'SD3' },
|
||||
{ model: TEST_SDXL_MODELS.BASE, cfg: 7.5, family: 'SDXL' },
|
||||
{ model: TEST_CUSTOM_SD, cfg: 7.5, family: 'SD1' },
|
||||
];
|
||||
|
||||
for (const testCase of testCases) {
|
||||
const params = {
|
||||
cfg: testCase.cfg,
|
||||
height: 1024,
|
||||
prompt: `CFG test for ${testCase.family}`,
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(testCase.model, params, sdMockContext);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle different schedulers by model family', async () => {
|
||||
const testCases = [
|
||||
{ model: TEST_SD35_MODELS.MEDIUM, scheduler: 'sgm_uniform' },
|
||||
{ model: TEST_SDXL_MODELS.BASE, scheduler: 'normal' },
|
||||
{ model: TEST_CUSTOM_SD, scheduler: 'normal' },
|
||||
];
|
||||
|
||||
for (const testCase of testCases) {
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: `Scheduler test: ${testCase.scheduler}`,
|
||||
scheduler: testCase.scheduler,
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(testCase.model, params, sdMockContext);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle various image dimensions', async () => {
|
||||
const dimensions = [
|
||||
{ width: 512, height: 512 }, // SD1.5 default
|
||||
{ width: 1024, height: 1024 }, // SDXL default
|
||||
{ width: 768, height: 1024 }, // Portrait
|
||||
{ width: 1344, height: 768 }, // Landscape
|
||||
];
|
||||
|
||||
for (const dim of dimensions) {
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: dim.height,
|
||||
prompt: `Dimension test ${dim.width}x${dim.height}`,
|
||||
steps: 20,
|
||||
width: dim.width,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(TEST_SDXL_MODELS.BASE, params, sdMockContext);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Error Handling Tests', () => {
|
||||
it('should handle unknown model gracefully', async () => {
|
||||
const modelName = 'unknown-model.safetensors';
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: 'Unknown model test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(modelName, params, sdMockContext);
|
||||
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
// Should work with default configuration
|
||||
});
|
||||
|
||||
it('should handle VAE resolution failure', async () => {
|
||||
const failingContext = {
|
||||
...sdMockContext,
|
||||
modelResolverService: {
|
||||
...sdMockContext.modelResolverService,
|
||||
getOptimalComponent: vi.fn().mockRejectedValue(new Error('VAE not found')),
|
||||
},
|
||||
};
|
||||
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: 'VAE failure test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
// Should not throw for SD3 models (built-in VAE)
|
||||
const result = await buildSimpleSDWorkflow(TEST_SD35_MODELS.MEDIUM, params, failingContext);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
});
|
||||
|
||||
it('should handle empty prompt', async () => {
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: '',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(TEST_SDXL_MODELS.BASE, params, sdMockContext);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Advanced Features Tests', () => {
|
||||
it('should support negative prompts', async () => {
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
negativePrompt: 'low quality, blurry',
|
||||
prompt: 'High quality image',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(TEST_SDXL_MODELS.BASE, params, sdMockContext);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
});
|
||||
|
||||
it('should handle seed generation', async () => {
|
||||
const { generateUniqueSeeds } = await import('@lobechat/utils');
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: 'Seed test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
await buildSimpleSDWorkflow(TEST_SDXL_MODELS.BASE, params, sdMockContext);
|
||||
|
||||
expect(generateUniqueSeeds).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should support custom sampler settings', async () => {
|
||||
const params = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
prompt: 'Custom sampler test',
|
||||
samplerName: 'dpmpp_2m_sde',
|
||||
scheduler: 'karras',
|
||||
steps: 25,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(TEST_SDXL_MODELS.BASE, params, sdMockContext);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Backward Compatibility Tests', () => {
|
||||
it('should maintain API compatibility with existing calls', async () => {
|
||||
// Test with minimal parameters (existing API pattern)
|
||||
const minimalParams = {
|
||||
prompt: 'Backward compatibility test',
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(
|
||||
TEST_SDXL_MODELS.BASE,
|
||||
minimalParams,
|
||||
sdMockContext,
|
||||
);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
});
|
||||
|
||||
it('should handle legacy parameter names', async () => {
|
||||
const legacyParams = {
|
||||
cfg: 7.5,
|
||||
height: 1024,
|
||||
inputImage: 'https://example.com/legacy.jpg', // Legacy parameter
|
||||
prompt: 'Legacy parameter test',
|
||||
steps: 20,
|
||||
width: 1024,
|
||||
};
|
||||
|
||||
const result = await buildSimpleSDWorkflow(
|
||||
TEST_SDXL_MODELS.BASE,
|
||||
legacyParams,
|
||||
sdMockContext,
|
||||
);
|
||||
expect(result).toHaveProperty('input');
|
||||
expect(result).toHaveProperty('setInputNode');
|
||||
expect(result).toHaveProperty('setOutputNode');
|
||||
expect(result).toHaveProperty('workflow');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,392 @@
|
|||
// @vitest-environment node
|
||||
import { PromptBuilder } from '@saintno/comfyui-sdk';
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
import {
|
||||
TEST_FLUX_MODELS,
|
||||
TEST_SD35_MODELS,
|
||||
} from '@/server/services/comfyui/__tests__/fixtures/testModels';
|
||||
import { mockContext } from '@/server/services/comfyui/__tests__/helpers/mockContext';
|
||||
import { setupAllMocks } from '@/server/services/comfyui/__tests__/setup/unifiedMocks';
|
||||
import { WorkflowError } from '@/server/services/comfyui/errors';
|
||||
import { buildFluxDevWorkflow } from '@/server/services/comfyui/workflows/flux-dev';
|
||||
import { buildFluxKontextWorkflow } from '@/server/services/comfyui/workflows/flux-kontext';
|
||||
import { buildFluxSchnellWorkflow } from '@/server/services/comfyui/workflows/flux-schnell';
|
||||
import { buildSD35Workflow } from '@/server/services/comfyui/workflows/sd35';
|
||||
import { buildSimpleSDWorkflow } from '@/server/services/comfyui/workflows/simple-sd';
|
||||
|
||||
// Create inline test parameters to avoid external dependencies
|
||||
const TEST_PARAMETERS = {
|
||||
'flux-dev': {
|
||||
defaults: { cfg: 3.5, steps: 20, samplerName: 'euler', scheduler: 'simple' },
|
||||
boundaries: { min: { cfg: 1, steps: 1 }, max: { cfg: 30, steps: 50 } },
|
||||
},
|
||||
'flux-schnell': {
|
||||
defaults: { cfg: 1, steps: 4, samplerName: 'euler', scheduler: 'simple' },
|
||||
boundaries: { min: { cfg: 1, steps: 1 }, max: { cfg: 1, steps: 8 } },
|
||||
},
|
||||
'flux-kontext': {
|
||||
defaults: { strength: 0.8 },
|
||||
},
|
||||
'sd35': {
|
||||
defaults: { cfg: 4, steps: 20, samplerName: 'euler', scheduler: 'sgm_uniform' },
|
||||
boundaries: { min: { cfg: 1, steps: 1 }, max: { cfg: 20, steps: 100 } },
|
||||
},
|
||||
'sdxl': {
|
||||
defaults: { cfg: 7.5, steps: 20, samplerName: 'euler', scheduler: 'normal' },
|
||||
boundaries: { min: { cfg: 1, steps: 1 }, max: { cfg: 20, steps: 100 } },
|
||||
},
|
||||
} as const;
|
||||
|
||||
// Mock the utility functions globally
|
||||
vi.mock('../utils/promptSplitter', () => ({
|
||||
splitPromptForDualCLIP: vi.fn((prompt: string) => ({
|
||||
clipLPrompt: prompt,
|
||||
t5xxlPrompt: prompt,
|
||||
})),
|
||||
}));
|
||||
|
||||
vi.mock('../utils/weightDType', () => ({
|
||||
selectOptimalWeightDtype: vi.fn(() => 'default'),
|
||||
}));
|
||||
|
||||
vi.mock('../utils/modelResolver', () => ({
|
||||
resolveModel: vi.fn((modelName: string) => {
|
||||
const cleanName = modelName.replace(/^comfyui\//, '');
|
||||
|
||||
// Return mock configuration based on model name patterns
|
||||
if (cleanName.includes('flux_dev') || cleanName.includes('flux-dev')) {
|
||||
return { family: 'flux', modelFamily: 'FLUX', variant: 'dev' };
|
||||
}
|
||||
if (cleanName.includes('flux_schnell') || cleanName.includes('flux-schnell')) {
|
||||
return { family: 'flux', modelFamily: 'FLUX', variant: 'schnell' };
|
||||
}
|
||||
if (cleanName.includes('flux_kontext') || cleanName.includes('kontext')) {
|
||||
return { family: 'flux', modelFamily: 'FLUX', variant: 'kontext' };
|
||||
}
|
||||
if (cleanName.includes('sd3.5') || cleanName.includes('sd35')) {
|
||||
return { family: 'sd35', modelFamily: 'SD3', variant: 'sd35' };
|
||||
}
|
||||
if (cleanName.includes('sdxl') || cleanName.includes('xl')) {
|
||||
return { family: 'sdxl', modelFamily: 'SDXL', variant: 'sdxl' };
|
||||
}
|
||||
if (cleanName.includes('v1-5') || cleanName.includes('sd15')) {
|
||||
return { family: 'sd15', modelFamily: 'SD1', variant: 'sd15' };
|
||||
}
|
||||
return null;
|
||||
}),
|
||||
}));
|
||||
|
||||
// Workflow builders configuration
|
||||
type WorkflowBuilderFunction = (modelFileName: string, params: any, context: any) => Promise<any>;
|
||||
|
||||
interface WorkflowTestConfig {
|
||||
name: string;
|
||||
builder: WorkflowBuilderFunction;
|
||||
modelName: string;
|
||||
parameterKey: keyof typeof TEST_PARAMETERS;
|
||||
specialFeatures?: string[];
|
||||
errorTests?: boolean;
|
||||
}
|
||||
|
||||
const WORKFLOW_CONFIGS: WorkflowTestConfig[] = [
|
||||
{
|
||||
name: 'FLUX Dev',
|
||||
builder: buildFluxDevWorkflow,
|
||||
modelName: TEST_FLUX_MODELS.DEV,
|
||||
parameterKey: 'flux-dev',
|
||||
specialFeatures: ['variable CFG', 'advanced sampler'],
|
||||
},
|
||||
{
|
||||
name: 'FLUX Schnell',
|
||||
builder: buildFluxSchnellWorkflow,
|
||||
modelName: TEST_FLUX_MODELS.SCHNELL,
|
||||
parameterKey: 'flux-schnell',
|
||||
specialFeatures: ['fixed CFG', 'fast generation'],
|
||||
},
|
||||
{
|
||||
name: 'FLUX Kontext',
|
||||
builder: buildFluxKontextWorkflow,
|
||||
modelName: TEST_FLUX_MODELS.KONTEXT,
|
||||
parameterKey: 'flux-kontext',
|
||||
specialFeatures: ['img2img support', 'vision capabilities'],
|
||||
},
|
||||
{
|
||||
name: 'SD3.5',
|
||||
builder: buildSD35Workflow,
|
||||
modelName: TEST_SD35_MODELS.LARGE,
|
||||
parameterKey: 'sd35',
|
||||
specialFeatures: ['external encoders', 'SGM scheduler'],
|
||||
errorTests: true,
|
||||
},
|
||||
{
|
||||
name: 'Simple SD',
|
||||
builder: buildSimpleSDWorkflow,
|
||||
modelName: 'sd_xl_base_1.0.safetensors',
|
||||
parameterKey: 'sdxl',
|
||||
specialFeatures: ['VAE handling', 'legacy support'],
|
||||
},
|
||||
];
|
||||
|
||||
describe('Unified Workflow Tests', () => {
|
||||
const { inputCalls } = setupAllMocks();
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe.each(WORKFLOW_CONFIGS)('$name Workflow', (config) => {
|
||||
it('should create workflow with default parameters', async () => {
|
||||
const fixture = TEST_PARAMETERS[config.parameterKey];
|
||||
const params = {
|
||||
prompt: 'A beautiful landscape',
|
||||
...fixture!.defaults,
|
||||
// Add standard dimensions for text-to-image models
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
|
||||
const result = await config.builder(config.modelName, params, mockContext);
|
||||
|
||||
// Verify workflow result is returned
|
||||
expect(result).toBeDefined();
|
||||
expect(result).toHaveProperty('input'); // PromptBuilder mock returns object with input method
|
||||
});
|
||||
|
||||
it('should create workflow with custom parameters', async () => {
|
||||
const fixture = TEST_PARAMETERS[config.parameterKey];
|
||||
const customParams = {
|
||||
prompt: 'Custom prompt for testing',
|
||||
width: 768,
|
||||
height: 512,
|
||||
steps: (fixture as any).boundaries?.max?.steps || 30,
|
||||
cfg: (fixture as any).boundaries?.max?.cfg || 7.5,
|
||||
};
|
||||
|
||||
const result = await config.builder(config.modelName, customParams, mockContext);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result).toHaveProperty('input');
|
||||
});
|
||||
|
||||
it('should handle empty prompt gracefully', async () => {
|
||||
const fixture = TEST_PARAMETERS[config.parameterKey];
|
||||
const params = {
|
||||
prompt: '',
|
||||
...fixture!.defaults,
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
|
||||
const result = await config.builder(config.modelName, params, mockContext);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result).toHaveProperty('input');
|
||||
});
|
||||
|
||||
it('should handle boundary values correctly', async () => {
|
||||
const fixture = TEST_PARAMETERS[config.parameterKey];
|
||||
|
||||
// Only test boundaries if they exist - Linus principle: don't test what doesn't exist
|
||||
if ((fixture as any).boundaries) {
|
||||
const minParams = {
|
||||
prompt: 'Minimum value test',
|
||||
width: 512,
|
||||
height: 512,
|
||||
steps: (fixture as any).boundaries.min.steps,
|
||||
cfg: (fixture as any).boundaries.min.cfg,
|
||||
};
|
||||
const minResult = await config.builder(config.modelName, minParams, mockContext);
|
||||
expect(minResult).toBeDefined();
|
||||
|
||||
const maxParams = {
|
||||
prompt: 'Maximum value test',
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
steps: (fixture as any).boundaries.max.steps,
|
||||
cfg: (fixture as any).boundaries.max.cfg,
|
||||
};
|
||||
const maxResult = await config.builder(config.modelName, maxParams, mockContext);
|
||||
expect(maxResult).toBeDefined();
|
||||
}
|
||||
});
|
||||
|
||||
// Special feature tests
|
||||
if (config.specialFeatures?.includes('img2img support')) {
|
||||
it('should handle image-to-image parameters', async () => {
|
||||
const params = {
|
||||
prompt: 'Transform this image',
|
||||
imageUrl: 'https://example.com/test.jpg',
|
||||
strength: 0.8,
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
|
||||
const result = await config.builder(config.modelName, params, mockContext);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result).toHaveProperty('input');
|
||||
});
|
||||
|
||||
it('should handle multiple image URLs', async () => {
|
||||
const params = {
|
||||
prompt: 'Process multiple images',
|
||||
imageUrls: ['https://example.com/img1.jpg', 'https://example.com/img2.jpg'],
|
||||
strength: 0.75,
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
|
||||
const result = await config.builder(config.modelName, params, mockContext);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result).toHaveProperty('input');
|
||||
});
|
||||
}
|
||||
|
||||
if (config.specialFeatures?.includes('variable CFG')) {
|
||||
it('should support variable CFG values', async () => {
|
||||
const params = {
|
||||
prompt: 'Variable CFG test',
|
||||
cfg: 5.0, // Different from default
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
|
||||
const result = await config.builder(config.modelName, params, mockContext);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result).toHaveProperty('input');
|
||||
});
|
||||
}
|
||||
|
||||
if (config.specialFeatures?.includes('fixed CFG')) {
|
||||
it('should use fixed CFG regardless of input', async () => {
|
||||
const params = {
|
||||
prompt: 'Fixed CFG test',
|
||||
cfg: 7.0, // Should be ignored for Schnell
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
|
||||
const result = await config.builder(config.modelName, params, mockContext);
|
||||
|
||||
expect(result).toBeDefined();
|
||||
expect(result).toHaveProperty('input');
|
||||
});
|
||||
}
|
||||
|
||||
// Error handling tests for models that support them
|
||||
if (config.errorTests) {
|
||||
it('should throw WorkflowError when required components are missing', async () => {
|
||||
// Create a context that simulates missing encoders
|
||||
const mockContextNoEncoders = {
|
||||
...mockContext,
|
||||
modelResolverService: {
|
||||
...mockContext.modelResolverService,
|
||||
getOptimalComponent: vi.fn().mockResolvedValue(undefined),
|
||||
},
|
||||
};
|
||||
|
||||
const params = {
|
||||
prompt: 'Test with missing encoders',
|
||||
};
|
||||
|
||||
await expect(
|
||||
config.builder(config.modelName, params, mockContextNoEncoders),
|
||||
).rejects.toThrow(WorkflowError);
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// Cross-workflow comparison tests
|
||||
describe('Cross-Workflow Validation', () => {
|
||||
it('should handle aspect ratio transformations consistently', async () => {
|
||||
const aspectRatioTests = [
|
||||
{ input: '16:9', expected: { width: 1024, height: 576 } },
|
||||
{ input: '1:1', expected: { width: 1024, height: 1024 } },
|
||||
{ input: '9:16', expected: { width: 576, height: 1024 } },
|
||||
];
|
||||
|
||||
for (const ratioTest of aspectRatioTests) {
|
||||
const params = {
|
||||
prompt: 'Aspect ratio test',
|
||||
width: ratioTest.expected.width,
|
||||
height: ratioTest.expected.height,
|
||||
};
|
||||
|
||||
// Test with multiple workflows
|
||||
for (const config of WORKFLOW_CONFIGS.slice(0, 3)) {
|
||||
// Test first 3 workflows
|
||||
const result = await config.builder(config.modelName, params, mockContext);
|
||||
expect(result).toBeDefined();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle seed parameter consistently', async () => {
|
||||
const testSeeds = [undefined, 0, 12345, 999999];
|
||||
|
||||
for (const seed of testSeeds) {
|
||||
const params = {
|
||||
prompt: 'Seed consistency test',
|
||||
seed,
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
};
|
||||
|
||||
// Test with workflows that support seed
|
||||
for (const config of WORKFLOW_CONFIGS.filter((c) => c.name !== 'FLUX Kontext')) {
|
||||
const result = await config.builder(config.modelName, params, mockContext);
|
||||
expect(result).toBeDefined();
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// Performance and validation tests
|
||||
describe('Performance and Validation', () => {
|
||||
it('should create workflows efficiently', async () => {
|
||||
const startTime = Date.now();
|
||||
|
||||
// Create multiple workflows in parallel
|
||||
const promises = WORKFLOW_CONFIGS.map((config) =>
|
||||
config.builder(config.modelName, { prompt: 'Performance test' }, mockContext),
|
||||
);
|
||||
|
||||
const results = await Promise.all(promises);
|
||||
const endTime = Date.now();
|
||||
|
||||
// Verify all workflows were created
|
||||
results.forEach((result) => {
|
||||
expect(result).toBeDefined();
|
||||
});
|
||||
|
||||
// Simple performance check - should complete within reasonable time
|
||||
expect(endTime - startTime).toBeLessThan(1000); // Less than 1 second
|
||||
});
|
||||
|
||||
it('should handle malformed parameters gracefully', async () => {
|
||||
const malformedParams = [
|
||||
{ prompt: null },
|
||||
{ prompt: 'test', width: -100 },
|
||||
{ prompt: 'test', height: 0 },
|
||||
{ prompt: 'test', steps: -5 },
|
||||
];
|
||||
|
||||
for (const params of malformedParams) {
|
||||
for (const config of WORKFLOW_CONFIGS.slice(0, 2)) {
|
||||
// Test with 2 workflows
|
||||
// Should not throw - workflows should handle invalid params gracefully
|
||||
try {
|
||||
const result = await config.builder(config.modelName, params as any, mockContext);
|
||||
expect(result).toBeDefined();
|
||||
} catch (error) {
|
||||
// If it throws, it should be a specific workflow error, not a generic JS error
|
||||
expect(error).toBeInstanceOf(Error);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
110
src/server/services/comfyui/config/constants.ts
Normal file
110
src/server/services/comfyui/config/constants.ts
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
/**
|
||||
* ComfyUI framework constants configuration
|
||||
* Unified management of hardcoded values with environment variable overrides / 统一管理硬编码值，支持环境变量覆盖
|
||||
*/
|
||||
|
||||
/**
|
||||
* Default configuration / 默认配置
|
||||
* 注意：BASE_URL不再处理环境变量，由构造函数统一处理优先级
|
||||
*/
|
||||
export const COMFYUI_DEFAULTS = {
|
||||
BASE_URL: 'http://localhost:8000',
|
||||
CONNECTION_TIMEOUT: 30_000,
|
||||
MAX_RETRIES: 3,
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* FLUX model configuration / FLUX 模型配置
|
||||
* Removed over-engineered dynamic T5 selection, maintain simple fixed configuration / 移除过度工程化的动态T5选择，保持简单固定配置
|
||||
*/
|
||||
export const FLUX_MODEL_CONFIG = {
|
||||
FILENAME_PREFIXES: {
|
||||
DEV: 'LobeChat/%year%-%month%-%day%/FLUX_Dev',
|
||||
KONTEXT: 'LobeChat/%year%-%month%-%day%/FLUX_Kontext',
|
||||
KREA: 'LobeChat/%year%-%month%-%day%/FLUX_Krea',
|
||||
SCHNELL: 'LobeChat/%year%-%month%-%day%/FLUX_Schnell',
|
||||
},
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* SD model configuration
|
||||
* Fixed model and filename prefixes for SD models
|
||||
*/
|
||||
export const SD_MODEL_CONFIG = {
|
||||
FILENAME_PREFIXES: {
|
||||
CUSTOM: 'LobeChat/%year%-%month%-%day%/CustomSD',
|
||||
SD15: 'LobeChat/%year%-%month%-%day%/SD15',
|
||||
SD35: 'LobeChat/%year%-%month%-%day%/SD35',
|
||||
SDXL: 'LobeChat/%year%-%month%-%day%/SDXL',
|
||||
},
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Default workflow node parameters / 工作流节点默认参数
|
||||
* Based on 2024 community best practices configuration / 基于 2024 年社区最佳实践配置
|
||||
*/
|
||||
|
||||
/**
|
||||
* Essential workflow defaults for internal use only
|
||||
* These are hardcoded values used by workflow internals, not user-configurable parameters
|
||||
*/
|
||||
export const WORKFLOW_DEFAULTS = {
|
||||
// FLUX specific settings
|
||||
FLUX: {
|
||||
BASE_SHIFT: 0.5,
|
||||
CLIP_GUIDANCE: 1,
|
||||
SAMPLER: 'euler',
|
||||
SCHEDULER: 'simple', // Higher denoise for Kontext img2img
|
||||
},
|
||||
// Image dimensions and batch settings
|
||||
IMAGE: {
|
||||
BATCH_SIZE: 1, // workflow internal use
|
||||
},
|
||||
// Internal noise and sampling settings
|
||||
SAMPLING: {
|
||||
DENOISE: 1, // t2i mode internal use
|
||||
MAX_SHIFT: 1.15, // FLUX internal parameter
|
||||
},
|
||||
// SD3.5 specific internal settings
|
||||
SD3: {
|
||||
SHIFT: 3, // SD3.5 ModelSamplingSD3 internal parameter
|
||||
},
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Default negative prompt for all SD models
|
||||
*/
|
||||
export const DEFAULT_NEGATIVE_PROMPT = `worst quality, normal quality, low quality, low res, blurry, distortion, text, watermark, logo, banner, extra digits, cropped, jpeg artifacts, signature, username, error, sketch, duplicate, ugly, monochrome, horror, geometry, mutation, disgusting, bad anatomy, bad proportions, bad quality, deformed, disconnected limbs, out of frame, out of focus, dehydrated, disfigured, extra arms, extra limbs, extra hands, fused fingers, gross proportions, long neck, jpeg, malformed limbs, mutated, mutated hands, mutated limbs, missing arms, missing fingers, picture frame, poorly drawn hands, poorly drawn face, collage, pixel, pixelated, grainy, color aberration, amputee, autograph, bad illustration, beyond the borders, blank background, body out of frame, boring background, branding, cut off, dismembered, disproportioned, distorted, draft, duplicated features, extra fingers, extra legs, fault, flaw, grains, hazy, identifying mark, improper scale, incorrect physiology, incorrect ratio, indistinct, kitsch, low resolution, macabre, malformed, mark, misshapen, missing hands, missing legs, mistake, morbid, mutilated, off-screen, outside the picture, poorly drawn feet, printed words, render, repellent, replicate, reproduce, revolting dimensions, script, shortened, sign, split image, squint, storyboard, tiling, trimmed, unfocused, unattractive, unnatural pose, unreal engine, unsightly, written language`;
|
||||
|
||||
/**
|
||||
* Supported model file formats
|
||||
* Used for model file validation and detection
|
||||
*/
|
||||
export const SUPPORTED_MODEL_FORMATS = [
|
||||
'.safetensors',
|
||||
'.ckpt',
|
||||
'.pt',
|
||||
'.pth',
|
||||
'.bin',
|
||||
'.gguf', // GGUF format for quantized models
|
||||
] as const;
|
||||
|
||||
/**
|
||||
* Custom SD model configuration
|
||||
* Fixed model and VAE filenames for custom SD models
|
||||
*/
|
||||
export const CUSTOM_SD_CONFIG = {
|
||||
MODEL_FILENAME: 'custom_sd_lobe.safetensors', // Both custom models use same file
|
||||
VAE_FILENAME: 'custom_sd_vae_lobe.safetensors', // Optional VAE file
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Component to ComfyUI node mappings
|
||||
* Maps component types to their corresponding ComfyUI loader nodes and input fields
|
||||
*/
|
||||
export const COMPONENT_NODE_MAPPINGS: Record<string, { field: string; node: string }> = {
|
||||
clip: { field: 'clip_name', node: 'CLIPLoader' },
|
||||
t5: { field: 'clip_name', node: 'CLIPLoader' }, // T5 is also CLIP type
|
||||
vae: { field: 'vae_name', node: 'VAELoader' },
|
||||
// Main models (UNET) are fetched via getCheckpoints(), not here
|
||||
} as const;
|
||||
843
src/server/services/comfyui/config/fluxModelRegistry.ts
Normal file
843
src/server/services/comfyui/config/fluxModelRegistry.ts
Normal file
|
|
@ -0,0 +1,843 @@
|
|||
/**
|
||||
* FLUX Model Registry - Separated for maintainability
|
||||
* Contains all FLUX model family registrations
|
||||
*/
|
||||
import type { ModelConfig } from './modelRegistry';
|
||||
|
||||
/* eslint-disable sort-keys-fix/sort-keys-fix */
|
||||
export const FLUX_MODEL_REGISTRY: Record<string, ModelConfig> = {
|
||||
// === Priority 1: Official Models (4 models) ===
|
||||
'flux1-dev.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// === Priority 2: Enterprise Optimized Models (106 models) ===
|
||||
|
||||
// 2.1 Enterprise Lightweight Models
|
||||
'flux.1-lite-8B.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux.1-lite-8B-alpha.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux-mini.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'FLUX_Mini_3_2B.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux_shakker_labs_union_pro-fp8_e4m3fn.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
|
||||
// 2.2 GGUF Series - FLUX.1-dev (11 models)
|
||||
'flux1-dev-F16.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q8_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q6_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q5_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q5_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q4_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q4_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q4_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q3_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q3_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-Q2_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 2.3 GGUF Series - FLUX.1-schnell (11 models)
|
||||
'flux1-schnell-F16.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q8_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q6_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q5_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q5_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q4_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q4_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q4_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q3_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q3_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-Q2_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 2.4 GGUF Series - FLUX.1-kontext (11 models)
|
||||
'flux1-kontext-dev-F16.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q8_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q6_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q5_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q5_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q4_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q4_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q4_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q3_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q3_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-Q2_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 2.5 GGUF Series - FLUX.1-krea (11 models)
|
||||
'flux1-krea-dev-F16.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q8_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q6_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q5_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q5_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q4_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q4_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q4_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q3_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q3_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-Q2_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 2.6 FP8 Series - dev variant (10 models)
|
||||
'flux1-dev-fp8.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux1-dev-fp8-e4m3fn.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux1-dev-fp8-e5m2.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e5m2',
|
||||
},
|
||||
|
||||
// 2.7 FP8 Series - schnell variant (4 models)
|
||||
'flux1-schnell-fp8.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux1-schnell-fp8-e4m3fn.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux1-schnell-fp8-e5m2.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e5m2',
|
||||
},
|
||||
|
||||
// 2.8 FP8 Series - kontext variant (3 models)
|
||||
'flux1-dev-kontext_fp8_scaled.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux1-kontext-dev-fp8-e4m3fn.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux1-kontext-dev-fp8-e5m2.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e5m2',
|
||||
},
|
||||
|
||||
// 2.9 FP8 Series - krea variant (3 models)
|
||||
'flux1-krea-dev_fp8_scaled.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux1-krea-dev-fp8-e4m3fn.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux1-krea-dev-fp8-e5m2.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e5m2',
|
||||
},
|
||||
|
||||
// 2.10 NF4 Quantization Series (7 models)
|
||||
'flux1-dev-bnb-nf4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-bnb-nf4-v2.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-bnb-nf4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-kontext-dev-bnb-nf4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-krea-dev-bnb-nf4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 2.11 Technical Experimental Models - SVDQuant Series
|
||||
'flux1-dev-svdquant-w4a4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-svdquant-w4a4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 2.12 Technical Experimental Models - TorchAO Series
|
||||
'flux1-dev-torchao-int8.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-dev-torchao-int4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-torchao-int8.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 2.13 Technical Experimental Models - optimum-quanto Series
|
||||
'flux1-dev-quanto-qfloat8.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-quanto-qfloat8.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 2.14 Technical Experimental Models - MFLUX Series (Apple Silicon Optimized)
|
||||
'flux1-dev-mflux-q4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux1-schnell-mflux-q4.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// === Priority 3: Community Fine-tuned Models (48 models) ===
|
||||
|
||||
// 3.1 Jib Mix Flux系列
|
||||
'Jib_Mix_Flux_v8_schnell.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'Jib_mix_Flux_V11_Krea_b_00001_.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'jibMixFlux_v8.q4_0.gguf': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.2 Real Dream FLUX系列
|
||||
'real_dream_flux_v1.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'real_dream_flux_beta.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'real_dream_flux_release.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'realDream_flux1V1.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'realDream_flux1V1_schnell.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.3 Vision Realistic FLUX系列
|
||||
'vision_realistic_flux_dev_v2.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'vision_realistic_flux_dev_fp8_no_clip_v2.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'vision_realistic_flux_v2_fp8.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'vision_realistic_flux_v2_dev.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'vision_realistic_flux_shakker.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.4 Flux Fusion系列
|
||||
'flux_fusion_v2_4steps.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux_fusion_ds_merge.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux_fusion_v2_tensorart.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.5 PixelWave FLUX系列
|
||||
'PixelWave_FLUX.1-dev_03.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'PixelWave_FLUX.1-schnell_04.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.6 Fux Capacity系列
|
||||
'Fux_Capacity_NSFW_v3.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'FuxCapacity2.1-Q8_0.gguf': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'FuxCapacity3.0_FP8.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'FuxCapacity3.1_FP16.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.7 Fluxmania系列
|
||||
'FluxMania_Kreamania_v1.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'Fluxmania_IV_fp8.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'Fluxmania_V6I.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'Fluxmania_V6I_fp16.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.8 Fluxed Up系列
|
||||
'Fluxed_Up_NSFW_v2.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.9 企业级LiblibAI模型
|
||||
'flux.1-ultra-realphoto-v2.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'f.1-dev-schnell-8steps-fp8.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'flux-muchen-asian.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'moyou-film-flux.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'firefly-fantasy-flux.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux-yanling-anime.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
|
||||
// 3.10 Legacy Models for Compatibility
|
||||
'Acorn_Spinning_FLUX_photorealism.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'CreArt_Hyper_Flux_Dev_8steps.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'Flux_Unchained_SCG_mixed.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'RealFlux_1.0b_Dev_Transformer.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'RealFlux_1.0b_Schnell.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'UltraReal_FineTune_v4.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'UltraRealistic_FineTune_Project_v4.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'XPlus_2(GGUF_Q4).gguf': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'XPlus_2(GGUF_Q6).gguf': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'XPlus_2(GGUF_Q8).gguf': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'educational-flux-simplified.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux-depth-fp16.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'commercial-flux-toolkit.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'krea',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux-fill-object-removal.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux-medical-environment-lora.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'kontext',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'flux-schnell-dev-merged-fp8.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'fp8_e4m3fn',
|
||||
},
|
||||
'schnellMODE_FLUX_S_v5_1.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'schnell',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
'NF4_BnB_FLUX_dev_optimized.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'dev',
|
||||
modelFamily: 'FLUX',
|
||||
recommendedDtype: 'default',
|
||||
},
|
||||
};
|
||||
/* eslint-enable sort-keys-fix/sort-keys-fix */
|
||||
48
src/server/services/comfyui/config/modelRegistry.ts
Normal file
48
src/server/services/comfyui/config/modelRegistry.ts
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
/**
|
||||
* ComfyUI Model Registry - Linus-style simple design
|
||||
* Interface shared, registries split for maintainability
|
||||
*/
|
||||
import { FLUX_MODEL_REGISTRY } from './fluxModelRegistry';
|
||||
import { SD_MODEL_REGISTRY } from './sdModelRegistry';
|
||||
|
||||
export interface ModelConfig {
|
||||
modelFamily: string;
|
||||
priority: number;
|
||||
recommendedDtype?: 'default' | 'fp8_e4m3fn' | 'fp8_e4m3fn_fast' | 'fp8_e5m2';
|
||||
variant: string;
|
||||
}
|
||||
|
||||
// ===================================================================
|
||||
// Combined Model Registry - FLUX + SD families
|
||||
// ===================================================================
|
||||
|
||||
export const MODEL_REGISTRY: Record<string, ModelConfig> = {
|
||||
...FLUX_MODEL_REGISTRY,
|
||||
...SD_MODEL_REGISTRY,
|
||||
};
|
||||
|
||||
/**
|
||||
* Model ID to Variant mapping
|
||||
* Maps actual frontend model IDs to their corresponding variants in registry
|
||||
* Based on src/config/aiModels/comfyui.ts definitions
|
||||
*/
|
||||
/* eslint-disable sort-keys-fix/sort-keys-fix */
|
||||
export const MODEL_ID_VARIANT_MAP: Record<string, string> = {
|
||||
// FLUX models
|
||||
'flux-schnell': 'schnell', // comfyui/flux-schnell
|
||||
'flux-dev': 'dev', // comfyui/flux-dev
|
||||
'flux-krea-dev': 'krea', // comfyui/flux-krea-dev
|
||||
'flux-kontext-dev': 'kontext', // comfyui/flux-kontext-dev
|
||||
|
||||
// SD3 models
|
||||
'stable-diffusion-35': 'sd35', // comfyui/stable-diffusion-35
|
||||
'stable-diffusion-35-inclclip': 'sd35-inclclip', // comfyui/stable-diffusion-35-inclclip
|
||||
|
||||
// SD1/SDXL models
|
||||
'stable-diffusion-15': 'sd15-t2i', // comfyui/stable-diffusion-15
|
||||
'stable-diffusion-xl': 'sdxl-t2i', // comfyui/stable-diffusion-xl
|
||||
'stable-diffusion-refiner': 'sdxl-i2i', // comfyui/stable-diffusion-refiner
|
||||
'stable-diffusion-custom': 'custom-sd', // comfyui/stable-diffusion-custom
|
||||
'stable-diffusion-custom-refiner': 'custom-sd', // comfyui/stable-diffusion-custom-refiner
|
||||
};
|
||||
/* eslint-enable sort-keys-fix/sort-keys-fix */
|
||||
624
src/server/services/comfyui/config/promptToolConst.ts
Normal file
624
src/server/services/comfyui/config/promptToolConst.ts
Normal file
|
|
@ -0,0 +1,624 @@
|
|||
/**
|
||||
* Prompt Optimizer Configuration
|
||||
* 提示词优化器配置 - 用于智能分割和优化提示词
|
||||
*/
|
||||
|
||||
/**
|
||||
* Style keywords configuration - organized by category
|
||||
* 风格关键词配置 - 按类别组织便于维护和扩展
|
||||
*/
|
||||
|
||||
/* eslint-disable sort-keys-fix/sort-keys-fix */
|
||||
export const STYLE_KEYWORDS = {
|
||||
// Artists and platforms / 艺术家和平台
|
||||
ARTISTS: [
|
||||
'by greg rutkowski',
|
||||
'by artgerm',
|
||||
'by wlop',
|
||||
'by alphonse mucha',
|
||||
'by james gurney',
|
||||
'by makoto shinkai',
|
||||
'by ghibli',
|
||||
'by hayao miyazaki',
|
||||
'by tim burton',
|
||||
'by banksy',
|
||||
'trending on artstation',
|
||||
'artstation',
|
||||
'deviantart',
|
||||
'pixiv',
|
||||
'concept art',
|
||||
'illustration',
|
||||
'artwork',
|
||||
'painting',
|
||||
'drawing',
|
||||
'digital painting',
|
||||
],
|
||||
|
||||
// Art styles / 艺术风格
|
||||
ART_STYLES: [
|
||||
'photorealistic',
|
||||
'photo realistic',
|
||||
'realistic',
|
||||
'hyperrealistic',
|
||||
'hyper realistic',
|
||||
'anime',
|
||||
'anime style',
|
||||
'manga',
|
||||
'manga style',
|
||||
'cartoon',
|
||||
'cartoon style',
|
||||
'oil painting',
|
||||
'watercolor',
|
||||
'watercolor painting',
|
||||
'acrylic painting',
|
||||
'sketch',
|
||||
'pencil sketch',
|
||||
'charcoal drawing',
|
||||
'digital art',
|
||||
'3d render',
|
||||
'3d rendering',
|
||||
'cgi',
|
||||
'pixel art',
|
||||
'8bit',
|
||||
'16bit',
|
||||
'retro pixel',
|
||||
'cinematic',
|
||||
'film still',
|
||||
'movie scene',
|
||||
'abstract',
|
||||
'abstract art',
|
||||
'surreal',
|
||||
'surrealism',
|
||||
'impressionist',
|
||||
'impressionism',
|
||||
'expressionist',
|
||||
'expressionism',
|
||||
'minimalist',
|
||||
'minimalism',
|
||||
'pop art',
|
||||
'art nouveau',
|
||||
'art deco',
|
||||
'baroque',
|
||||
'renaissance',
|
||||
'gothic',
|
||||
'cyberpunk',
|
||||
'steampunk',
|
||||
'dieselpunk',
|
||||
'solarpunk',
|
||||
'vaporwave',
|
||||
'synthwave',
|
||||
'retrowave',
|
||||
],
|
||||
|
||||
// Lighting effects / 光照效果
|
||||
LIGHTING: [
|
||||
'dramatic lighting',
|
||||
'soft lighting',
|
||||
'hard lighting',
|
||||
'studio lighting',
|
||||
'golden hour',
|
||||
'golden hour lighting',
|
||||
'blue hour',
|
||||
'magic hour',
|
||||
'sunset lighting',
|
||||
'sunrise lighting',
|
||||
'neon lights',
|
||||
'neon lighting',
|
||||
'rim lighting',
|
||||
'backlit',
|
||||
'backlighting',
|
||||
'volumetric lighting',
|
||||
'god rays',
|
||||
'crepuscular rays',
|
||||
'natural lighting',
|
||||
'ambient lighting',
|
||||
'warm lighting',
|
||||
'cold lighting',
|
||||
'cool lighting',
|
||||
'moody lighting',
|
||||
'atmospheric lighting',
|
||||
'cinematic lighting',
|
||||
'chiaroscuro',
|
||||
'low key lighting',
|
||||
'high key lighting',
|
||||
'diffused lighting',
|
||||
'harsh lighting',
|
||||
'candlelight',
|
||||
'firelight',
|
||||
'moonlight',
|
||||
'sunlight',
|
||||
'fluorescent',
|
||||
'incandescent',
|
||||
],
|
||||
|
||||
// Photography terms / 摄影术语
|
||||
PHOTOGRAPHY: [
|
||||
'depth of field',
|
||||
'shallow depth of field',
|
||||
'deep depth of field',
|
||||
'bokeh',
|
||||
'bokeh effect',
|
||||
'motion blur',
|
||||
'film grain',
|
||||
'lens grain',
|
||||
'chromatic aberration',
|
||||
'lens distortion',
|
||||
'fisheye',
|
||||
'macro',
|
||||
'macro photography',
|
||||
'wide angle',
|
||||
'ultra wide angle',
|
||||
'telephoto',
|
||||
'telephoto lens',
|
||||
'portrait',
|
||||
'portrait photography',
|
||||
'landscape',
|
||||
'landscape photography',
|
||||
'street photography',
|
||||
'aerial photography',
|
||||
'drone photography',
|
||||
'long exposure',
|
||||
'time-lapse',
|
||||
'close-up',
|
||||
'extreme close-up',
|
||||
'medium shot',
|
||||
'wide shot',
|
||||
'establishing shot',
|
||||
'dof',
|
||||
'35mm',
|
||||
'35mm photograph',
|
||||
'50mm',
|
||||
'85mm',
|
||||
'professional photograph',
|
||||
'professional photography',
|
||||
'dslr',
|
||||
'mirrorless',
|
||||
'medium format',
|
||||
'hasselblad',
|
||||
'canon',
|
||||
'nikon',
|
||||
'sony alpha',
|
||||
'film photography',
|
||||
'analog photography',
|
||||
'polaroid',
|
||||
'instant photo',
|
||||
],
|
||||
|
||||
// Quality descriptions / 质量描述
|
||||
QUALITY: [
|
||||
'high quality',
|
||||
'best quality',
|
||||
'highest quality',
|
||||
'top quality',
|
||||
'masterpiece',
|
||||
'award winning',
|
||||
'professional',
|
||||
'professional quality',
|
||||
'4k',
|
||||
'4k resolution',
|
||||
'8k',
|
||||
'8k resolution',
|
||||
'uhd',
|
||||
'ultra hd',
|
||||
'full hd',
|
||||
'hd',
|
||||
'high resolution',
|
||||
'high res',
|
||||
'ultra detailed',
|
||||
'highly detailed',
|
||||
'super detailed',
|
||||
'extremely detailed',
|
||||
'insanely detailed',
|
||||
'intricate',
|
||||
'intricate details',
|
||||
'fine details',
|
||||
'sharp',
|
||||
'sharp focus',
|
||||
'crisp',
|
||||
'crystal clear',
|
||||
'pristine',
|
||||
'flawless',
|
||||
'perfect',
|
||||
'stunning',
|
||||
'beautiful',
|
||||
'gorgeous',
|
||||
'breathtaking',
|
||||
'magnificent',
|
||||
'exquisite',
|
||||
],
|
||||
|
||||
// Rendering and effects / 渲染和效果
|
||||
RENDERING: [
|
||||
'octane render',
|
||||
'octane',
|
||||
'unreal engine',
|
||||
'unreal engine 5',
|
||||
'ue5',
|
||||
'unity',
|
||||
'blender',
|
||||
'maya',
|
||||
'cinema 4d',
|
||||
'c4d',
|
||||
'houdini',
|
||||
'zbrush',
|
||||
'substance painter',
|
||||
'marmoset',
|
||||
'keyshot',
|
||||
'vray',
|
||||
'v-ray',
|
||||
'arnold render',
|
||||
'redshift',
|
||||
'cycles',
|
||||
'cycles render',
|
||||
'ray tracing',
|
||||
'path tracing',
|
||||
'global illumination',
|
||||
'gi',
|
||||
'ambient occlusion',
|
||||
'ao',
|
||||
'subsurface scattering',
|
||||
'sss',
|
||||
'pbr',
|
||||
'physically based rendering',
|
||||
'bloom',
|
||||
'bloom effect',
|
||||
'lens flare',
|
||||
'post processing',
|
||||
'color grading',
|
||||
'tone mapping',
|
||||
'hdr',
|
||||
'high dynamic range',
|
||||
],
|
||||
|
||||
// Color and mood / 颜色和氛围
|
||||
COLOR_MOOD: [
|
||||
'vibrant',
|
||||
'vibrant colors',
|
||||
'vivid',
|
||||
'vivid colors',
|
||||
'muted',
|
||||
'muted colors',
|
||||
'pastel',
|
||||
'pastel colors',
|
||||
'monochrome',
|
||||
'black and white',
|
||||
'grayscale',
|
||||
'sepia',
|
||||
'warm colors',
|
||||
'cool colors',
|
||||
'cold colors',
|
||||
'neon colors',
|
||||
'psychedelic',
|
||||
'psychedelic colors',
|
||||
'rainbow',
|
||||
'iridescent',
|
||||
'holographic',
|
||||
'metallic',
|
||||
'chrome',
|
||||
'golden',
|
||||
'silver',
|
||||
'dark',
|
||||
'dark mood',
|
||||
'moody',
|
||||
'atmospheric',
|
||||
'ethereal',
|
||||
'dreamy',
|
||||
'dreamlike',
|
||||
'surreal atmosphere',
|
||||
'mysterious',
|
||||
'mystical',
|
||||
'magical',
|
||||
'fantasy',
|
||||
'epic',
|
||||
'dramatic',
|
||||
'intense',
|
||||
'peaceful',
|
||||
'serene',
|
||||
'calm',
|
||||
'tranquil',
|
||||
'melancholic',
|
||||
'nostalgic',
|
||||
'romantic',
|
||||
'whimsical',
|
||||
'playful',
|
||||
'cheerful',
|
||||
'gloomy',
|
||||
'ominous',
|
||||
'eerie',
|
||||
'creepy',
|
||||
'horror',
|
||||
'gothic atmosphere',
|
||||
],
|
||||
|
||||
// Texture and materials / 纹理和材质
|
||||
TEXTURE_MATERIAL: [
|
||||
'glossy',
|
||||
'matte',
|
||||
'satin',
|
||||
'rough',
|
||||
'smooth',
|
||||
'polished',
|
||||
'brushed',
|
||||
'textured',
|
||||
'glass',
|
||||
'crystal',
|
||||
'transparent',
|
||||
'translucent',
|
||||
'reflective',
|
||||
'refractive',
|
||||
'metallic texture',
|
||||
'chrome finish',
|
||||
'copper',
|
||||
'brass',
|
||||
'bronze',
|
||||
'steel',
|
||||
'aluminum',
|
||||
'titanium',
|
||||
'wood',
|
||||
'wooden',
|
||||
'oak',
|
||||
'mahogany',
|
||||
'bamboo',
|
||||
'marble',
|
||||
'granite',
|
||||
'stone',
|
||||
'concrete',
|
||||
'brick',
|
||||
'fabric',
|
||||
'cloth',
|
||||
'silk',
|
||||
'velvet',
|
||||
'cotton',
|
||||
'denim',
|
||||
'leather',
|
||||
'suede',
|
||||
'fur',
|
||||
'plastic',
|
||||
'rubber',
|
||||
'latex',
|
||||
'organic',
|
||||
'bio',
|
||||
'liquid',
|
||||
'fluid',
|
||||
'gel',
|
||||
'ice',
|
||||
'frost',
|
||||
'frozen',
|
||||
'wet',
|
||||
'dry',
|
||||
'dusty',
|
||||
'rusty',
|
||||
'weathered',
|
||||
'aged',
|
||||
'vintage texture',
|
||||
'retro texture',
|
||||
],
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Style synonyms mapping for better recognition
|
||||
* 同义词映射，提高识别准确率
|
||||
*/
|
||||
export const STYLE_SYNONYMS: Record<string, string[]> = {
|
||||
// Photography variations
|
||||
'photorealistic': [
|
||||
'photo-realistic',
|
||||
'photo realistic',
|
||||
'lifelike',
|
||||
'true-to-life',
|
||||
'true to life',
|
||||
],
|
||||
'hyperrealistic': ['hyper-realistic', 'hyper realistic', 'ultra realistic', 'ultrarealistic'],
|
||||
'depth of field': ['dof', 'depth-of-field', 'focal depth', 'focus depth'],
|
||||
'bokeh': ['bokeh effect', 'background blur', 'out of focus background'],
|
||||
|
||||
// Art style variations
|
||||
'cinematic': ['filmic', 'movie-like', 'film-style', 'theatrical', 'cinema style'],
|
||||
'anime': ['anime-style', 'japanese animation', 'animestyle'],
|
||||
'manga': ['manga-style', 'japanese comic', 'mangastyle'],
|
||||
'3d render': ['3d-render', '3d rendering', '3d-rendering', 'three dimensional', 'cgi render'],
|
||||
'digital art': ['digital-art', 'digital artwork', 'digital painting', 'digitalart'],
|
||||
|
||||
// Quality variations
|
||||
'4k': ['4k resolution', '4k quality', 'ultra hd', 'uhd', '3840x2160', '4096x2160'],
|
||||
'8k': ['8k resolution', '8k quality', 'ultra hd+', '7680x4320', '8192x4320'],
|
||||
'high quality': ['high-quality', 'hq', 'hi quality', 'hi-quality', 'highquality'],
|
||||
'masterpiece': ['master piece', 'master-piece', 'opus', 'magnum opus'],
|
||||
|
||||
// Lighting variations
|
||||
'golden hour': ['golden-hour', 'magic hour', 'sunset light', 'sunrise light'],
|
||||
'rim lighting': ['rim-lighting', 'rimlight', 'rim light', 'edge lighting'],
|
||||
'volumetric lighting': ['volumetric-lighting', 'god rays', 'light rays', 'sun rays'],
|
||||
|
||||
// Rendering variations
|
||||
'octane render': ['octane-render', 'octanerender', 'otoy octane'],
|
||||
'unreal engine': ['unreal-engine', 'ue4', 'ue5', 'unrealengine'],
|
||||
'ray tracing': ['ray-tracing', 'raytracing', 'rt', 'rtx'],
|
||||
|
||||
// Artist variations
|
||||
'by greg rutkowski': ['greg rutkowski', 'rutkowski', 'greg-rutkowski'],
|
||||
'by artgerm': ['artgerm', 'stanley lau', 'artgerm lau'],
|
||||
'trending on artstation': ['artstation trending', 'artstation hq', 'artstation-hq'],
|
||||
};
|
||||
|
||||
/**
|
||||
* Compound styles that should be recognized as a whole
|
||||
* 组合风格，应该作为整体识别
|
||||
*/
|
||||
export const COMPOUND_STYLES = [
|
||||
// Studio and brand styles
|
||||
'studio ghibli style',
|
||||
'pixar style',
|
||||
'disney style',
|
||||
'dreamworks style',
|
||||
'marvel style',
|
||||
'dc comics style',
|
||||
|
||||
// Specific art movements
|
||||
'art nouveau style',
|
||||
'art deco style',
|
||||
'pop art style',
|
||||
'street art style',
|
||||
'graffiti style',
|
||||
|
||||
// Game and media styles
|
||||
'league of legends style',
|
||||
'overwatch style',
|
||||
'world of warcraft style',
|
||||
'final fantasy style',
|
||||
'zelda style',
|
||||
'pokemon style',
|
||||
|
||||
// Photography styles
|
||||
'national geographic style',
|
||||
'vogue style',
|
||||
'fashion photography',
|
||||
'portrait photography',
|
||||
'landscape photography',
|
||||
'street photography',
|
||||
'wildlife photography',
|
||||
'macro photography',
|
||||
|
||||
// Specific artist styles
|
||||
'van gogh style',
|
||||
'picasso style',
|
||||
'monet style',
|
||||
'rembrandt style',
|
||||
'da vinci style',
|
||||
'warhol style',
|
||||
'banksy style',
|
||||
'tim burton style',
|
||||
'wes anderson style',
|
||||
'christopher nolan style',
|
||||
|
||||
// Technical compound terms
|
||||
'physically based rendering',
|
||||
'global illumination',
|
||||
'subsurface scattering',
|
||||
'ambient occlusion',
|
||||
'chromatic aberration',
|
||||
'depth of field',
|
||||
'motion blur',
|
||||
'lens flare',
|
||||
|
||||
// Atmosphere combinations
|
||||
'cinematic lighting',
|
||||
'dramatic lighting',
|
||||
'studio lighting',
|
||||
'natural lighting',
|
||||
'volumetric fog',
|
||||
'atmospheric perspective',
|
||||
'aerial perspective',
|
||||
|
||||
// Quality combinations
|
||||
'ultra high definition',
|
||||
'ultra high quality',
|
||||
'super high resolution',
|
||||
'professional quality',
|
||||
'production quality',
|
||||
'broadcast quality',
|
||||
'print quality',
|
||||
|
||||
// Complex styles
|
||||
'cyberpunk aesthetic',
|
||||
'steampunk aesthetic',
|
||||
'vaporwave aesthetic',
|
||||
'synthwave aesthetic',
|
||||
'cottagecore aesthetic',
|
||||
'dark academia aesthetic',
|
||||
'y2k aesthetic',
|
||||
'minimalist design',
|
||||
'maximalist design',
|
||||
'brutalist architecture',
|
||||
'gothic architecture',
|
||||
'baroque architecture',
|
||||
] as const;
|
||||
|
||||
/**
|
||||
* Precise adjective patterns for style extraction
|
||||
* 精确的形容词模式，用于风格提取
|
||||
*/
|
||||
export const STYLE_ADJECTIVE_PATTERNS = {
|
||||
// Visual quality related / 视觉质量相关
|
||||
quality:
|
||||
/^(sharp|blur(ry)?|clear|crisp|clean|smooth|rough|grainy|noisy|pristine|flawless|perfect|polished)$/i,
|
||||
|
||||
// Artistic style related / 艺术风格相关
|
||||
artistic:
|
||||
/^(abstract|surreal|minimal(ist)?|ornate|baroque|gothic|modern|contemporary|traditional|classical|vintage|retro|antique|futuristic|avant-garde)$/i,
|
||||
|
||||
// Color and lighting / 颜色和光照
|
||||
visual:
|
||||
/^(bright|dark|dim|vibrant|vivid|muted|saturated|desaturated|warm|cool|cold|hot|soft|hard|harsh|gentle|subtle|bold|pale|rich|deep)$/i,
|
||||
|
||||
// Mood and atmosphere / 情绪和氛围
|
||||
mood: /^(dramatic|peaceful|chaotic|serene|calm|mysterious|mystical|magical|epic|legendary|heroic|romantic|melancholic|nostalgic|whimsical|playful|serious|solemn|cheerful|gloomy|ominous|eerie|creepy|scary|dreamy|ethereal|fantastical|moody|atmospheric)$/i,
|
||||
|
||||
// Texture and material / 纹理和材质
|
||||
texture:
|
||||
/^(metallic|wooden|glass(y)?|crystalline|fabric|leather|plastic|rubber|organic|synthetic|liquid|solid|transparent|translucent|opaque|reflective|matte|glossy|satin|rough|smooth|wet|dry|dusty|rusty|weathered|aged|new|fresh|worn)$/i,
|
||||
|
||||
// Size and scale / 尺寸和规模
|
||||
scale:
|
||||
/^(tiny|small|medium|large|huge|massive|gigantic|colossal|enormous|microscopic|miniature|oversized|epic-scale|human-scale|intimate|vast|infinite)$/i,
|
||||
|
||||
// Complexity and detail / 复杂度和细节
|
||||
detail:
|
||||
/^(simple|complex|intricate|elaborate|detailed|minimal|advanced|sophisticated|primitive|refined|crude|delicate|robust)$/i,
|
||||
|
||||
// Professional quality / 专业质量
|
||||
professional:
|
||||
/^(professional|amateur|masterful|skilled|expert|novice|polished|raw|finished|unfinished|complete|incomplete|refined|rough)$/i,
|
||||
} as const;
|
||||
|
||||
/* eslint-enable sort-keys-fix/sort-keys-fix */
|
||||
|
||||
/**
|
||||
* Get all style keywords as a flattened array
|
||||
* 获取所有风格关键词的扁平数组
|
||||
*/
|
||||
export function getAllStyleKeywords(): readonly string[] {
|
||||
return Object.values(STYLE_KEYWORDS).flat();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all compound styles
|
||||
* 获取所有组合风格
|
||||
*/
|
||||
export function getCompoundStyles(): readonly string[] {
|
||||
return COMPOUND_STYLES;
|
||||
}
|
||||
|
||||
/**
|
||||
* Normalize a style term using synonyms
|
||||
* 使用同义词标准化风格术语
|
||||
*/
|
||||
export function normalizeStyleTerm(term: string): string {
|
||||
const lowerTerm = term.toLowerCase();
|
||||
|
||||
// Check if this term is a synonym
|
||||
for (const [canonical, synonyms] of Object.entries(STYLE_SYNONYMS)) {
|
||||
if (synonyms.includes(lowerTerm)) {
|
||||
return canonical;
|
||||
}
|
||||
}
|
||||
|
||||
return term;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a word matches any style adjective pattern
|
||||
* 检查词语是否匹配任何风格形容词模式
|
||||
*/
|
||||
export function isStyleAdjective(word: string): boolean {
|
||||
const lowerWord = word.toLowerCase();
|
||||
return Object.values(STYLE_ADJECTIVE_PATTERNS).some((pattern) => pattern.test(lowerWord));
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract style adjectives from words based on precise patterns
|
||||
* 基于精确模式从词语中提取风格形容词
|
||||
*/
|
||||
export function extractStyleAdjectives(words: string[]): string[] {
|
||||
return words.filter((word) => isStyleAdjective(word));
|
||||
}
|
||||
508
src/server/services/comfyui/config/sdModelRegistry.ts
Normal file
508
src/server/services/comfyui/config/sdModelRegistry.ts
Normal file
|
|
@ -0,0 +1,508 @@
|
|||
/**
|
||||
* Stable Diffusion Model Registry - Separated for maintainability
|
||||
* Contains all SD1.5, SDXL, and SD3.5 model family registrations
|
||||
*/
|
||||
import type { ModelConfig } from './modelRegistry';
|
||||
|
||||
/* eslint-disable sort-keys-fix/sort-keys-fix */
|
||||
export const SD_MODEL_REGISTRY: Record<string, ModelConfig> = {
|
||||
// ===================================================================
|
||||
// SD3.5 Model Family Registry
|
||||
// ===================================================================
|
||||
|
||||
// SD3.5 Models (requires clip_g.safetensors)
|
||||
'sd3.5_large.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sd35',
|
||||
modelFamily: 'SD3',
|
||||
},
|
||||
'sd3.5_large_turbo.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sd35',
|
||||
modelFamily: 'SD3',
|
||||
},
|
||||
'sd3.5_medium.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd35',
|
||||
modelFamily: 'SD3',
|
||||
},
|
||||
'sd3.5_large_fp8_scaled.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sd35',
|
||||
modelFamily: 'SD3',
|
||||
},
|
||||
|
||||
// SD3.5 Models (With CLIP - includes CLIP/T5 internally)
|
||||
'sd3.5_medium_incl_clips_t5xxlfp8scaled.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sd35-inclclip',
|
||||
modelFamily: 'SD3',
|
||||
},
|
||||
|
||||
// === Custom SD Models (for user-uploaded models) ===
|
||||
// These entries serve as examples for custom model support
|
||||
'custom-sd-model.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'custom-sd',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'custom-sd-refiner.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'custom-sd',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
|
||||
// ===================================================================
|
||||
// SD1.5 Model Family Registry (Built-in CLIP/VAE)
|
||||
// ===================================================================
|
||||
|
||||
// === SD1.5 Official Models (Priority 1) ===
|
||||
'v1-5-pruned-emaonly.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-fp16.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly.ckpt': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
|
||||
// === SD1.5 Quantized Models (Priority 2) ===
|
||||
'v1-5-pruned-emaonly-F16.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q8_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q6_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q5_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q5_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q4_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q4_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q4_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q3_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q3_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'v1-5-pruned-emaonly-Q2_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
|
||||
// === SD1.5 Community Models (Priority 3) ===
|
||||
'dreamshaper_8.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'DreamShaper_8_pruned.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'Deliberate_v2.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'Deliberate_v6.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'Realistic_Vision_V5.1.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'Realistic_Vision_V5.1_fp16-no-ema.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'realisticVisionV60B1_v60B1VAE.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'Chilloutmix.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'chilloutmix-Ni.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'chilloutmix_NiPrunedFp16Fix.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'chilloutmix_NiPrunedFp32Fix.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'braV7.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'guofeng3_v34.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'koreanDollLikeness_v20.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'AnythingV5Ink_ink.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'neverendingDream_v122.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'majestixMix_v70.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'kissMix2_v20.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'xxmix9realistic_v40.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'tangYuan_v50.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'flat2DAnimerge_v45Sharp.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'cyberrealistic_v33.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
'analog-diffusion-1.0.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sd15-t2i',
|
||||
modelFamily: 'SD1',
|
||||
},
|
||||
|
||||
// ===================================================================
|
||||
// SDXL Model Family Registry (Built-in CLIP/VAE)
|
||||
// ===================================================================
|
||||
|
||||
// === SDXL Text-to-Image Models (Priority 1) ===
|
||||
'sd_xl_base_1.0.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_turbo_1.0_fp16.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0_0.9vae.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
|
||||
// === SDXL Image-to-Image Models (Refiner) ===
|
||||
'sd_xl_refiner_1.0.safetensors': {
|
||||
priority: 1,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
|
||||
// === SDXL Quantized Models (Priority 2) ===
|
||||
'sd_xl_base_1.0-F16.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q8_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q6_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q5_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q5_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q4_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q4_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q4_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q3_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q3_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_base_1.0-Q2_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
|
||||
// === SDXL Refiner Quantized Models ===
|
||||
'sd_xl_refiner_1.0-F16.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q8_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q6_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q5_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q5_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q4_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q4_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q4_0.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q3_K_M.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q3_K_S.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sd_xl_refiner_1.0-Q2_K.gguf': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-i2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
|
||||
// === SDXL Enterprise Models (Priority 2) ===
|
||||
'SSD-1B.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'SSD-1B-modelspec.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sdxl_lightning_1step.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sdxl_lightning_4step.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sdxl_lightning_8step.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'diffusion_pytorch_model.fp16.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'lcm_lora_sdxl.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
|
||||
// === SDXL Community Models (Priority 3) ===
|
||||
'juggernautXL_v9Rdphoto2Lightning.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'realvisxlV50_v50Bakedvae.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'dreamshaperXL_v21TurboDPMSDE.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'ponyDiffusionV6XL_v6StartWithThisOne.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'novaAnimeXL_il.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'nebulaeAnimeStyleSDXL_v20.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'counterfeitxl_v25.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'animagineXLV31_v31.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'bluepencilXL_v100.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'sudachi_v10.safetensors': {
|
||||
priority: 3,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
|
||||
// === Playground Models (Based on SDXL Architecture) ===
|
||||
'playground-v2.5-1024px-aesthetic.fp16.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'playground-v2.5-1024px-aesthetic.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'playground-v2-1024px-aesthetic.fp16.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
'playground-v2-1024px-aesthetic.safetensors': {
|
||||
priority: 2,
|
||||
variant: 'sdxl-t2i',
|
||||
modelFamily: 'SDXL',
|
||||
},
|
||||
};
|
||||
/* eslint-enable sort-keys-fix/sort-keys-fix */
|
||||
385
src/server/services/comfyui/config/systemComponents.ts
Normal file
385
src/server/services/comfyui/config/systemComponents.ts
Normal file
|
|
@ -0,0 +1,385 @@
|
|||
/**
|
||||
* System Components Registry Configuration
|
||||
*/
|
||||
import { ConfigError } from '@/server/services/comfyui/errors';
|
||||
|
||||
export interface ComponentConfig {
|
||||
/** Compatible model variants (for LoRA and ControlNet) */
|
||||
compatibleVariants?: string[];
|
||||
/** ControlNet type (for ControlNet components only) */
|
||||
controlnetType?: string;
|
||||
/** Model family this component is designed for */
|
||||
modelFamily: string;
|
||||
/** Priority level: 1=Essential/Official, 2=Standard/Professional, 3=Optional/Community */
|
||||
priority: number;
|
||||
/** Component type */
|
||||
type: string;
|
||||
}
|
||||
|
||||
// Model family constants (for business logic reference)
|
||||
export const MODEL_FAMILIES = {
|
||||
FLUX: 'FLUX',
|
||||
SD1: 'SD1',
|
||||
SD3: 'SD3',
|
||||
SDXL: 'SDXL',
|
||||
} as const;
|
||||
|
||||
// Component type constants (for business logic reference)
|
||||
export const COMPONENT_TYPES = {
|
||||
CLIP: 'clip',
|
||||
CONTROLNET: 'controlnet',
|
||||
LORA: 'lora',
|
||||
T5: 't5',
|
||||
VAE: 'vae',
|
||||
} as const;
|
||||
|
||||
// ControlNet type constants
|
||||
export const CONTROLNET_TYPES = {
|
||||
CANNY: 'canny',
|
||||
DEPTH: 'depth',
|
||||
HED: 'hed',
|
||||
NORMAL: 'normal',
|
||||
POSE: 'pose',
|
||||
SCRIBBLE: 'scribble',
|
||||
SEMANTIC: 'semantic',
|
||||
} as const;
|
||||
|
||||
/* eslint-disable sort-keys-fix/sort-keys-fix */
|
||||
export const SYSTEM_COMPONENTS: Record<string, ComponentConfig> = {
|
||||
// ===================================================================
|
||||
// === ESSENTIAL COMPONENTS (Priority 1) ===
|
||||
// ===================================================================
|
||||
|
||||
'ae.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'vae',
|
||||
},
|
||||
|
||||
'clip_l.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'clip',
|
||||
},
|
||||
|
||||
'clip_g.safetensors': {
|
||||
modelFamily: 'SD3',
|
||||
priority: 1,
|
||||
type: 'clip',
|
||||
},
|
||||
|
||||
't5xxl_fp16.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 't5',
|
||||
},
|
||||
|
||||
// ===================================================================
|
||||
// === OPTIONAL COMPONENTS (Priority 2-3) ===
|
||||
// ===================================================================
|
||||
't5xxl_fp8_e4m3fn.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 't5',
|
||||
},
|
||||
|
||||
't5xxl_fp8_e4m3fn_scaled.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 't5',
|
||||
},
|
||||
|
||||
't5xxl_fp8_e5m2.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 't5',
|
||||
},
|
||||
|
||||
'google_t5-v1_1-xxl_encoderonly-fp16.safetensors': {
|
||||
modelFamily: 'FLUX',
|
||||
priority: 3,
|
||||
type: 't5',
|
||||
},
|
||||
|
||||
// ===================================================================
|
||||
// === VAE MODELS ===
|
||||
// ===================================================================
|
||||
|
||||
// SD1 VAE Models
|
||||
'vae-ft-mse-840000-ema-pruned.safetensors': {
|
||||
modelFamily: 'SD1',
|
||||
priority: 1,
|
||||
type: 'vae',
|
||||
},
|
||||
|
||||
'sd-vae-ft-ema.safetensors': {
|
||||
modelFamily: 'SD1',
|
||||
priority: 1,
|
||||
type: 'vae',
|
||||
},
|
||||
|
||||
// SDXL VAE Models
|
||||
'sdxl_vae.safetensors': {
|
||||
modelFamily: 'SDXL',
|
||||
priority: 1,
|
||||
type: 'vae',
|
||||
},
|
||||
|
||||
'sdxl.vae.safetensors': {
|
||||
modelFamily: 'SDXL',
|
||||
priority: 1,
|
||||
type: 'vae',
|
||||
},
|
||||
|
||||
'sd_xl_base_1.0_0.9vae.safetensors': {
|
||||
modelFamily: 'SDXL',
|
||||
priority: 2,
|
||||
type: 'vae',
|
||||
},
|
||||
|
||||
// ===================================================================
|
||||
// === LORA ADAPTERS ===
|
||||
// ===================================================================
|
||||
|
||||
// XLabs-AI Official FLUX LoRA Adapters (Priority 1 - Official)
|
||||
'realism_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'anime_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'disney_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'scenery_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'art_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'mjv6_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'flux-realism-lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'flux-lora-collection-8-styles.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'disney-anime-art-lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
// LiblibAI Professional LoRA (Priority 2 - Professional)
|
||||
'flux-kodak-grain-lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'flux-first-person-selfie-lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'flux-anime-rainbow-light-lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'flux-detailer-enhancement-lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
// CivitAI Special Effects LoRA (Priority 2 - Professional)
|
||||
'Envy_Flux_Reanimated_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'Photon_Construct_Flux_V1_0_lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
// ModelScope LoRA Collection (Priority 2 - Professional)
|
||||
'flux-ultimate-lora-collection.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'artaug-flux-enhancement-lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'flux-canny-dev-lora.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 2,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
// Community and Experimental LoRA (Priority 3 - Community)
|
||||
'watercolor_painting_schnell_lora.safetensors': {
|
||||
compatibleVariants: ['schnell'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 3,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'juggernaut_lora_flux.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 3,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'chinese-style-flux-lora-collection.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 3,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'flux-medical-environment-lora.safetensors': {
|
||||
compatibleVariants: ['kontext'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 3,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
'flux-fill-object-removal.safetensors': {
|
||||
compatibleVariants: ['kontext'],
|
||||
modelFamily: 'FLUX',
|
||||
priority: 3,
|
||||
type: 'lora',
|
||||
},
|
||||
|
||||
// ===================================================================
|
||||
// === CONTROLNET MODELS ===
|
||||
// ===================================================================
|
||||
|
||||
// XLabs-AI Official FLUX ControlNet Models (Priority 1 - Official)
|
||||
'flux-controlnet-canny-v3.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
controlnetType: 'canny',
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'controlnet',
|
||||
},
|
||||
|
||||
'flux-controlnet-depth-v3.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
controlnetType: 'depth',
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'controlnet',
|
||||
},
|
||||
|
||||
'flux-controlnet-hed-v3.safetensors': {
|
||||
compatibleVariants: ['dev'],
|
||||
controlnetType: 'hed',
|
||||
modelFamily: 'FLUX',
|
||||
priority: 1,
|
||||
type: 'controlnet',
|
||||
},
|
||||
} as const;
|
||||
/* eslint-enable sort-keys-fix/sort-keys-fix */
|
||||
|
||||
/**
|
||||
* Get all components with names matching filters
|
||||
*/
|
||||
export function getAllComponentsWithNames(options?: {
|
||||
compatibleVariant?: string;
|
||||
controlnetType?: ComponentConfig['controlnetType'];
|
||||
modelFamily?: ComponentConfig['modelFamily'];
|
||||
priority?: number;
|
||||
type?: ComponentConfig['type'];
|
||||
}): Array<{ config: ComponentConfig; name: string }> {
|
||||
return Object.entries(SYSTEM_COMPONENTS)
|
||||
.filter(
|
||||
([, config]) =>
|
||||
(!options?.type || config.type === options.type) &&
|
||||
(!options?.priority || config.priority === options.priority) &&
|
||||
(!options?.modelFamily || config.modelFamily === options.modelFamily) &&
|
||||
(!options?.compatibleVariant ||
|
||||
(config.compatibleVariants &&
|
||||
config.compatibleVariants.includes(options.compatibleVariant))) &&
|
||||
(!options?.controlnetType || config.controlnetType === options.controlnetType),
|
||||
)
|
||||
.map(([name, config]) => ({ config, name }));
|
||||
}
|
||||
|
||||
/**
|
||||
* Get optimal component of specified type
|
||||
*/
|
||||
export function getOptimalComponent(
|
||||
type: ComponentConfig['type'],
|
||||
modelFamily: ComponentConfig['modelFamily'],
|
||||
): string {
|
||||
const components = getAllComponentsWithNames({ modelFamily, type }).sort(
|
||||
(a, b) => a.config.priority - b.config.priority,
|
||||
);
|
||||
|
||||
if (components.length === 0) {
|
||||
throw new ConfigError(
|
||||
`No ${type} components configured for model family ${modelFamily}`,
|
||||
ConfigError.Reasons.MISSING_CONFIG,
|
||||
{ modelFamily, type },
|
||||
);
|
||||
}
|
||||
|
||||
return components[0].name;
|
||||
}
|
||||
70
src/server/services/comfyui/config/workflowRegistry.ts
Normal file
70
src/server/services/comfyui/config/workflowRegistry.ts
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
import type { PromptBuilder } from '@saintno/comfyui-sdk';
|
||||
|
||||
import type { WorkflowContext } from '@/server/services/comfyui/core/workflowBuilderService';
|
||||
// Import all workflow builders
|
||||
import { buildFluxDevWorkflow } from '@/server/services/comfyui/workflows/flux-dev';
|
||||
import { buildFluxKontextWorkflow } from '@/server/services/comfyui/workflows/flux-kontext';
|
||||
import { buildFluxSchnellWorkflow } from '@/server/services/comfyui/workflows/flux-schnell';
|
||||
import { buildSD35Workflow } from '@/server/services/comfyui/workflows/sd35';
|
||||
import { buildSimpleSDWorkflow } from '@/server/services/comfyui/workflows/simple-sd';
|
||||
|
||||
// Workflow builder type
|
||||
type WorkflowBuilder = (
|
||||
modelFileName: string,
|
||||
params: Record<string, any>,
|
||||
context: WorkflowContext,
|
||||
) => Promise<PromptBuilder<any, any, any>>;
|
||||
|
||||
/**
|
||||
* Variant to Workflow mapping
|
||||
* Based on actual model registry variant values
|
||||
*/
|
||||
/* eslint-disable sort-keys-fix/sort-keys-fix */
|
||||
export const VARIANT_WORKFLOW_MAP: Record<string, WorkflowBuilder> = {
|
||||
// FLUX variants
|
||||
'dev': buildFluxDevWorkflow,
|
||||
'schnell': buildFluxSchnellWorkflow,
|
||||
'kontext': buildFluxKontextWorkflow,
|
||||
'krea': buildFluxDevWorkflow,
|
||||
|
||||
// SD3 variants
|
||||
'sd35': buildSD35Workflow, // needs external encoders
|
||||
'sd35-inclclip': buildSimpleSDWorkflow, // built-in encoders
|
||||
|
||||
// SD1/SDXL variants
|
||||
'sd15-t2i': buildSimpleSDWorkflow,
|
||||
'sdxl-t2i': buildSimpleSDWorkflow,
|
||||
'sdxl-i2i': buildSimpleSDWorkflow,
|
||||
'custom-sd': buildSimpleSDWorkflow,
|
||||
};
|
||||
|
||||
/**
|
||||
* Architecture default workflows (when variant not matched)
|
||||
*/
|
||||
export const ARCHITECTURE_DEFAULT_MAP: Record<string, WorkflowBuilder> = {
|
||||
FLUX: buildFluxDevWorkflow,
|
||||
SD3: buildSD35Workflow,
|
||||
SD1: buildSimpleSDWorkflow,
|
||||
SDXL: buildSimpleSDWorkflow,
|
||||
};
|
||||
/* eslint-enable sort-keys-fix/sort-keys-fix */
|
||||
|
||||
/**
|
||||
* Get the appropriate workflow builder for a given architecture and variant
|
||||
*
|
||||
* @param architecture - Model architecture (FLUX, SD3, SD1, SDXL)
|
||||
* @param variant - Model variant (dev, schnell, kontext, sd35, etc.)
|
||||
* @returns Workflow builder function or undefined if not found
|
||||
*/
|
||||
export function getWorkflowBuilder(
|
||||
architecture: string,
|
||||
variant?: string,
|
||||
): WorkflowBuilder | undefined {
|
||||
// Prefer variant mapping
|
||||
if (variant && VARIANT_WORKFLOW_MAP[variant]) {
|
||||
return VARIANT_WORKFLOW_MAP[variant];
|
||||
}
|
||||
|
||||
// Fallback to architecture default
|
||||
return ARCHITECTURE_DEFAULT_MAP[architecture];
|
||||
}
|
||||
145
src/server/services/comfyui/core/comfyUIAuthService.ts
Normal file
145
src/server/services/comfyui/core/comfyUIAuthService.ts
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
/**
|
||||
* ComfyUI Authentication Service
|
||||
*
|
||||
* Handles all authentication-related logic for ComfyUI connections
|
||||
* Supports 4 authentication modes: none, basic, bearer, custom
|
||||
*/
|
||||
import type { ComfyUIKeyVault } from '@lobechat/types';
|
||||
import { createBasicAuthCredentials } from '@lobechat/utils';
|
||||
import type {
|
||||
BasicCredentials,
|
||||
BearerTokenCredentials,
|
||||
CustomCredentials,
|
||||
} from '@saintno/comfyui-sdk';
|
||||
import debug from 'debug';
|
||||
|
||||
import { ServicesError } from '@/server/services/comfyui/errors';
|
||||
|
||||
const log = debug('lobe-image:comfyui:auth');
|
||||
|
||||
export class ComfyUIAuthService {
|
||||
private credentials: BasicCredentials | BearerTokenCredentials | CustomCredentials | undefined;
|
||||
private authHeaders: Record<string, string> | undefined;
|
||||
|
||||
constructor(options: ComfyUIKeyVault) {
|
||||
log('🔐 Initializing authentication service');
|
||||
|
||||
this.validateOptions(options);
|
||||
this.credentials = this.createCredentials(options);
|
||||
this.authHeaders = this.createAuthHeaders(options);
|
||||
|
||||
log('✅ Authentication service initialized with type:', options.authType || 'none');
|
||||
}
|
||||
|
||||
/**
|
||||
* Get credentials for ComfyUI SDK
|
||||
*/
|
||||
getCredentials(): BasicCredentials | BearerTokenCredentials | CustomCredentials | undefined {
|
||||
return this.credentials;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get authentication headers for HTTP requests
|
||||
*/
|
||||
getAuthHeaders(): Record<string, string> | undefined {
|
||||
return this.authHeaders;
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate authentication options
|
||||
*/
|
||||
private validateOptions(options: ComfyUIKeyVault): void {
|
||||
const { authType = 'none', apiKey, username, password, customHeaders } = options;
|
||||
|
||||
if (authType === 'basic' && (!username || !password)) {
|
||||
throw new ServicesError(
|
||||
'Basic authentication requires username and password',
|
||||
ServicesError.Reasons.INVALID_ARGS,
|
||||
{ authType },
|
||||
);
|
||||
}
|
||||
|
||||
if (authType === 'bearer' && !apiKey) {
|
||||
throw new ServicesError(
|
||||
'Bearer token authentication requires API key',
|
||||
ServicesError.Reasons.INVALID_AUTH,
|
||||
{ authType },
|
||||
);
|
||||
}
|
||||
|
||||
if (authType === 'custom' && (!customHeaders || Object.keys(customHeaders).length === 0)) {
|
||||
throw new ServicesError(
|
||||
'Custom authentication requires custom headers',
|
||||
ServicesError.Reasons.INVALID_ARGS,
|
||||
{ authType },
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create credentials object for ComfyUI SDK
|
||||
*/
|
||||
private createCredentials(
|
||||
options: ComfyUIKeyVault,
|
||||
): BasicCredentials | BearerTokenCredentials | CustomCredentials | undefined {
|
||||
const { authType = 'none', apiKey, username, password, customHeaders } = options;
|
||||
|
||||
switch (authType) {
|
||||
case 'basic': {
|
||||
return {
|
||||
password: password!,
|
||||
type: 'basic',
|
||||
username: username!,
|
||||
} as BasicCredentials;
|
||||
}
|
||||
|
||||
case 'bearer': {
|
||||
return {
|
||||
token: apiKey!,
|
||||
type: 'bearer_token',
|
||||
} as BearerTokenCredentials;
|
||||
}
|
||||
|
||||
case 'custom': {
|
||||
return {
|
||||
headers: customHeaders!,
|
||||
type: 'custom',
|
||||
} as CustomCredentials;
|
||||
}
|
||||
|
||||
default: {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create authentication headers for direct HTTP requests
|
||||
*/
|
||||
private createAuthHeaders(options: ComfyUIKeyVault): Record<string, string> | undefined {
|
||||
const { authType = 'none', apiKey, username, password, customHeaders } = options;
|
||||
|
||||
switch (authType) {
|
||||
case 'basic': {
|
||||
if (username && password) {
|
||||
const basicAuth = createBasicAuthCredentials(username, password);
|
||||
return { Authorization: `Basic ${basicAuth}` };
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
case 'bearer': {
|
||||
if (apiKey) {
|
||||
return { Authorization: `Bearer ${apiKey}` };
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
case 'custom': {
|
||||
return customHeaders;
|
||||
}
|
||||
}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
249
src/server/services/comfyui/core/comfyUIClientService.ts
Normal file
249
src/server/services/comfyui/core/comfyUIClientService.ts
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
/**
|
||||
* ComfyUI Client Service
|
||||
*
|
||||
* Central service layer for all ComfyUI SDK interactions
|
||||
* Provides unified error handling and abstraction over SDK
|
||||
* Uses modular services for authentication, connection, and caching
|
||||
*/
|
||||
import type { ComfyUIKeyVault } from '@lobechat/types';
|
||||
import { CallWrapper, ComfyApi, PromptBuilder } from '@saintno/comfyui-sdk';
|
||||
import debug from 'debug';
|
||||
|
||||
import { COMFYUI_DEFAULTS } from '@/server/services/comfyui/config/constants';
|
||||
import { ComfyUIAuthService } from '@/server/services/comfyui/core/comfyUIAuthService';
|
||||
import { ComfyUIConnectionService } from '@/server/services/comfyui/core/comfyUIConnectionService';
|
||||
import { ErrorHandlerService } from '@/server/services/comfyui/core/errorHandlerService';
|
||||
import { ServicesError } from '@/server/services/comfyui/errors';
|
||||
import { TTLCacheManager } from '@/server/services/comfyui/utils/cacheManager';
|
||||
|
||||
const log = debug('lobe-image:comfyui:client');
|
||||
|
||||
/**
|
||||
* Workflow execution result
|
||||
*/
|
||||
export interface WorkflowResult {
|
||||
// Raw output data from workflow execution, keyed by node ID
|
||||
_raw?: Record<string, any>;
|
||||
images?: {
|
||||
images?: Array<{
|
||||
data: string;
|
||||
height?: number;
|
||||
mimeType: string;
|
||||
width?: number;
|
||||
}>;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Progress callback type
|
||||
*/
|
||||
export type ProgressCallback = (info: any) => void;
|
||||
|
||||
/**
|
||||
* ComfyUI Client Service
|
||||
* Encapsulates all SDK interactions using modular services
|
||||
*/
|
||||
export class ComfyUIClientService {
|
||||
private client: ComfyApi;
|
||||
private baseURL: string;
|
||||
|
||||
// Modular services for separation of concerns
|
||||
private cacheManager: TTLCacheManager;
|
||||
private authService: ComfyUIAuthService;
|
||||
private connectionService: ComfyUIConnectionService;
|
||||
private errorHandler: ErrorHandlerService;
|
||||
|
||||
constructor(options: ComfyUIKeyVault = {}) {
|
||||
log('🏗️ Initializing ComfyUI Client Service');
|
||||
|
||||
this.errorHandler = new ErrorHandlerService();
|
||||
|
||||
try {
|
||||
// Initialize modular services
|
||||
this.authService = new ComfyUIAuthService(options);
|
||||
this.cacheManager = new TTLCacheManager(60_000); // 1 minute TTL
|
||||
this.connectionService = new ComfyUIConnectionService();
|
||||
|
||||
// Setup base URL
|
||||
this.baseURL =
|
||||
options.baseURL || process.env.COMFYUI_DEFAULT_URL || COMFYUI_DEFAULTS.BASE_URL;
|
||||
|
||||
// Initialize client with credentials from AuthService
|
||||
this.client = new ComfyApi(this.baseURL, undefined, {
|
||||
credentials: this.authService.getCredentials(),
|
||||
});
|
||||
this.client.init();
|
||||
|
||||
log('✅ ComfyUI Client Service initialized with baseURL:', this.baseURL);
|
||||
} catch (error) {
|
||||
// Use ErrorHandlerService to transform internal errors to framework errors
|
||||
this.errorHandler.handleError(error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get authentication headers for image download
|
||||
* This method provides auth headers to framework layer without exposing credentials
|
||||
* @returns Authentication headers object, or undefined if no auth is configured
|
||||
*/
|
||||
getAuthHeaders(): Record<string, string> | undefined {
|
||||
// Delegate to AuthService
|
||||
return this.authService.getAuthHeaders();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the path for an image result
|
||||
*/
|
||||
getPathImage(imageInfo: any): string {
|
||||
return this.client.getPathImage(imageInfo);
|
||||
}
|
||||
|
||||
/**
|
||||
* Upload an image to ComfyUI server
|
||||
* @param file - The image data as Buffer or Blob
|
||||
* @param fileName - The name for the uploaded file
|
||||
* @returns The filename on ComfyUI server
|
||||
*/
|
||||
async uploadImage(file: Buffer | Blob, fileName: string): Promise<string> {
|
||||
log('📤 Uploading image to ComfyUI:', fileName);
|
||||
|
||||
const result = await this.client.uploadImage(file, fileName);
|
||||
|
||||
if (!result) {
|
||||
throw new ServicesError(
|
||||
'Failed to upload image to ComfyUI server',
|
||||
ServicesError.Reasons.UPLOAD_FAILED,
|
||||
{ fileName, response: result },
|
||||
);
|
||||
}
|
||||
|
||||
log('✅ Image uploaded successfully:', result.info.filename);
|
||||
return result.info.filename;
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a workflow
|
||||
*/
|
||||
async executeWorkflow(
|
||||
workflow: PromptBuilder<any, any, any>,
|
||||
onProgress?: ProgressCallback,
|
||||
): Promise<WorkflowResult> {
|
||||
log('🚀 Executing workflow...');
|
||||
|
||||
return new Promise<WorkflowResult>((resolve, reject) => {
|
||||
new CallWrapper(this.client, workflow)
|
||||
.onFinished((result: any) => {
|
||||
log('✅ Workflow execution finished successfully');
|
||||
log('🔍 Raw workflow result structure:', {
|
||||
hasImages: 'images' in result,
|
||||
hasRaw: '_raw' in result,
|
||||
keys: Object.keys(result),
|
||||
rawKeys: result._raw ? Object.keys(result._raw) : null,
|
||||
});
|
||||
resolve(result);
|
||||
})
|
||||
.onFailed((error: any) => {
|
||||
log('❌ Workflow execution failed:', error?.message || error);
|
||||
reject(error);
|
||||
})
|
||||
.onProgress((info: any) => {
|
||||
log('⏳ Progress:', info);
|
||||
onProgress?.(info);
|
||||
})
|
||||
.run();
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Get available checkpoints from ComfyUI
|
||||
* Wraps SDK method to avoid Law of Demeter violation
|
||||
* Uses unified TTL cache for performance optimization
|
||||
*/
|
||||
async getCheckpoints(): Promise<string[]> {
|
||||
return await this.cacheManager.get('checkpoints', async () => {
|
||||
return await this.client.getCheckpoints();
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Get available LoRAs from ComfyUI
|
||||
* Wraps SDK method to avoid Law of Demeter violation
|
||||
* Uses unified TTL cache for performance optimization
|
||||
*/
|
||||
async getLoras(): Promise<string[]> {
|
||||
return await this.cacheManager.get('loras', async () => {
|
||||
return await this.client.getLoras();
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Get node definitions from ComfyUI
|
||||
* Wraps SDK method to avoid Law of Demeter violation
|
||||
* Uses unified TTL cache for performance optimization
|
||||
* @param nodeName - Optional specific node name to query
|
||||
*/
|
||||
async getNodeDefs(nodeName?: string): Promise<any> {
|
||||
const allNodeDefs = await this.cacheManager.get('nodeDefs', async () => {
|
||||
return await this.client.getNodeDefs();
|
||||
});
|
||||
|
||||
// Return specific node or all nodes
|
||||
return nodeName && allNodeDefs ? { [nodeName]: allNodeDefs[nodeName] } : allNodeDefs;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get sampler info from ComfyUI
|
||||
* Wraps SDK method to avoid Law of Demeter violation
|
||||
*/
|
||||
async getSamplerInfo(): Promise<{ samplerName: string[]; scheduler: string[] }> {
|
||||
const info = await this.client.getSamplerInfo();
|
||||
|
||||
return {
|
||||
samplerName: this.extractStrings(info.sampler),
|
||||
scheduler: this.extractStrings(info.scheduler),
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract string values from sampler info arrays
|
||||
* Handle both string arrays and tuple arrays like ['euler', { tooltip: 'info' }]
|
||||
*/
|
||||
private extractStrings(arr: any): string[] {
|
||||
if (!Array.isArray(arr)) return [];
|
||||
return arr
|
||||
.map((item) => (Array.isArray(item) ? item[0] : item))
|
||||
.filter((item) => typeof item === 'string');
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate connection to ComfyUI server
|
||||
* Delegates to ConnectionService for connection management
|
||||
*/
|
||||
async validateConnection(): Promise<boolean> {
|
||||
return await this.connectionService.validateConnection(
|
||||
this.baseURL,
|
||||
this.authService.getAuthHeaders(),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get connection status information
|
||||
*/
|
||||
getConnectionStatus() {
|
||||
return this.connectionService.getStatus();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get authentication service instance (for advanced usage)
|
||||
*/
|
||||
getAuthService(): ComfyUIAuthService {
|
||||
return this.authService;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get connection service instance (for advanced usage)
|
||||
*/
|
||||
getConnectionService(): ComfyUIConnectionService {
|
||||
return this.connectionService;
|
||||
}
|
||||
}
|
||||
136
src/server/services/comfyui/core/comfyUIConnectionService.ts
Normal file
136
src/server/services/comfyui/core/comfyUIConnectionService.ts
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
/**
|
||||
* ComfyUI Connection Service
|
||||
*
|
||||
* Handles connection validation and state management for ComfyUI server
|
||||
* Provides TTL-based connection validation caching
|
||||
*/
|
||||
import debug from 'debug';
|
||||
|
||||
import { ServicesError } from '@/server/services/comfyui/errors';
|
||||
|
||||
const log = debug('lobe-image:comfyui:connection');
|
||||
|
||||
export class ComfyUIConnectionService {
|
||||
private validated: boolean = false;
|
||||
private lastValidationTime: number = 0;
|
||||
private readonly validationTTL = 5 * 60 * 1000; // 5 minutes
|
||||
|
||||
constructor() {
|
||||
log('🔗 Initializing connection service');
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if connection is validated and not expired
|
||||
*/
|
||||
isValidated(): boolean {
|
||||
if (!this.validated) return false;
|
||||
|
||||
const now = Date.now();
|
||||
if (now - this.lastValidationTime > this.validationTTL) {
|
||||
this.validated = false;
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark connection as validated
|
||||
*/
|
||||
markAsValidated(): void {
|
||||
this.validated = true;
|
||||
this.lastValidationTime = Date.now();
|
||||
log('✅ Connection marked as validated');
|
||||
}
|
||||
|
||||
/**
|
||||
* Invalidate connection (force re-validation)
|
||||
*/
|
||||
invalidate(): void {
|
||||
this.validated = false;
|
||||
log('❌ Connection invalidated, will require re-validation');
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate connection to ComfyUI server
|
||||
* Uses system_stats endpoint for health check
|
||||
*/
|
||||
async validateConnection(
|
||||
baseURL: string,
|
||||
authHeaders?: Record<string, string>,
|
||||
): Promise<boolean> {
|
||||
// Check if already validated and not expired
|
||||
if (this.isValidated()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
try {
|
||||
// Use system_stats endpoint for health check
|
||||
// This is a lightweight endpoint that returns system information
|
||||
const url = `${baseURL}/system_stats`;
|
||||
const headers = authHeaders || {};
|
||||
|
||||
log('🔍 Validating connection to:', url);
|
||||
|
||||
const response = await fetch(url, {
|
||||
headers: {
|
||||
...headers,
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
method: 'GET',
|
||||
mode: 'cors',
|
||||
});
|
||||
|
||||
// Just check if we got a successful response
|
||||
if (!response.ok) {
|
||||
this.invalidate();
|
||||
// Throw ServicesError with status for error parser to handle
|
||||
throw new ServicesError(
|
||||
`HTTP ${response.status}: ${response.statusText}`,
|
||||
ServicesError.Reasons.CONNECTION_FAILED,
|
||||
{ endpoint: '/system_stats', status: response.status, statusText: response.statusText },
|
||||
);
|
||||
}
|
||||
|
||||
// Verify response is valid JSON
|
||||
const data = await response.json();
|
||||
if (!data || typeof data !== 'object') {
|
||||
throw new ServicesError(
|
||||
'Invalid response from ComfyUI server',
|
||||
ServicesError.Reasons.CONNECTION_FAILED,
|
||||
{ endpoint: '/system_stats' },
|
||||
);
|
||||
}
|
||||
|
||||
this.markAsValidated();
|
||||
log('✅ Connection validated successfully');
|
||||
return true;
|
||||
} catch (error) {
|
||||
// Reset connection state on any error
|
||||
this.invalidate();
|
||||
|
||||
// Re-throw all errors - let the service layer handle error classification
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get connection status information
|
||||
*/
|
||||
getStatus(): {
|
||||
isValidated: boolean;
|
||||
lastValidationTime: number | null;
|
||||
timeUntilExpiry: number | null;
|
||||
} {
|
||||
const now = Date.now();
|
||||
const timeUntilExpiry = this.validated
|
||||
? Math.max(0, this.validationTTL - (now - this.lastValidationTime))
|
||||
: null;
|
||||
|
||||
return {
|
||||
isValidated: this.validated,
|
||||
lastValidationTime: this.validated ? this.lastValidationTime : null,
|
||||
timeUntilExpiry,
|
||||
};
|
||||
}
|
||||
}
|
||||
538
src/server/services/comfyui/core/errorHandlerService.ts
Normal file
538
src/server/services/comfyui/core/errorHandlerService.ts
Normal file
|
|
@ -0,0 +1,538 @@
|
|||
/**
|
||||
* Error Handler Service
|
||||
*
|
||||
* Centralized error handling for ComfyUI runtime
|
||||
* Maps internal errors to framework errors
|
||||
*/
|
||||
import {
|
||||
AgentRuntimeError,
|
||||
AgentRuntimeErrorType,
|
||||
ILobeAgentRuntimeErrorType,
|
||||
} from '@lobechat/model-runtime';
|
||||
import { TRPCError } from '@trpc/server';
|
||||
|
||||
import { SYSTEM_COMPONENTS } from '@/server/services/comfyui/config/systemComponents';
|
||||
import {
|
||||
ConfigError,
|
||||
ServicesError,
|
||||
UtilsError,
|
||||
WorkflowError,
|
||||
isComfyUIInternalError,
|
||||
} from '@/server/services/comfyui/errors';
|
||||
import { ModelResolverError } from '@/server/services/comfyui/errors/modelResolverError';
|
||||
import { getComponentInfo } from '@/server/services/comfyui/utils/componentInfo';
|
||||
|
||||
interface ComfyUIError {
|
||||
code?: number | string;
|
||||
details?: any;
|
||||
message: string;
|
||||
missingFileName?: string;
|
||||
missingFileType?: 'model' | 'component';
|
||||
status?: number;
|
||||
type?: string;
|
||||
userGuidance?: string;
|
||||
}
|
||||
|
||||
interface ParsedError {
|
||||
error: ComfyUIError;
|
||||
errorType: ILobeAgentRuntimeErrorType;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate user guidance message based on missing file info
|
||||
* Server-side version with access to full component information
|
||||
* @param fileName - The missing file name
|
||||
* @param fileType - The type of missing file
|
||||
* @returns User-friendly guidance message
|
||||
*/
|
||||
function generateUserGuidance(fileName: string, fileType: 'model' | 'component'): string {
|
||||
if (fileType === 'component') {
|
||||
const componentInfo = getComponentInfo(fileName);
|
||||
|
||||
if (componentInfo) {
|
||||
return `Missing ${componentInfo.displayName}: ${fileName}. Please download and place it in the ${componentInfo.folderPath} folder.`;
|
||||
}
|
||||
|
||||
// Fallback for unknown components
|
||||
return `Missing component file: ${fileName}. Please download and place it in the appropriate ComfyUI models folder.`;
|
||||
}
|
||||
|
||||
// Main model files
|
||||
return `Missing model file: ${fileName}. Please download and place it in the models/checkpoints folder.`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract missing file information from error message
|
||||
* Server-side version with access to SYSTEM_COMPONENTS
|
||||
* @param message - Error message that may contain file names
|
||||
* @returns Object with extracted file name and type, or null if no file found
|
||||
*/
|
||||
function extractMissingFileInfo(message: string): {
|
||||
fileName: string;
|
||||
fileType: 'model' | 'component';
|
||||
} | null {
|
||||
if (!message) return null;
|
||||
|
||||
// Check for "Expected one of:" pattern from enhanced model errors
|
||||
const expectedPattern = /expected one of:\s*([^.]+\.(?:safetensors|ckpt|pt|pth))/i;
|
||||
const expectedMatch = message.match(expectedPattern);
|
||||
|
||||
if (expectedMatch) {
|
||||
// Extract the first file from the match
|
||||
const fileName = expectedMatch[1].trim().split(',')[0].trim();
|
||||
if (fileName) {
|
||||
return {
|
||||
fileName,
|
||||
fileType: 'model',
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Common model file extensions - allow dots in filename
|
||||
const modelFilePattern = /([\w.-]+\.(?:safetensors|ckpt|pt|pth))\b/gi;
|
||||
const fileMatch = message.match(modelFilePattern);
|
||||
|
||||
if (fileMatch) {
|
||||
const fileName = fileMatch[0];
|
||||
|
||||
// Use server-side SYSTEM_COMPONENTS to check if it's a system component
|
||||
if (fileName in SYSTEM_COMPONENTS) {
|
||||
return {
|
||||
fileName,
|
||||
fileType: 'component',
|
||||
};
|
||||
}
|
||||
|
||||
// If not found in SYSTEM_COMPONENTS, treat as main model
|
||||
return {
|
||||
fileName,
|
||||
fileType: 'model',
|
||||
};
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the error is a model-related error
|
||||
* @param error - Error object
|
||||
* @param message - Pre-extracted message
|
||||
* @returns Whether it's a model error
|
||||
*/
|
||||
function isModelError(error: any, message?: string): boolean {
|
||||
const errorMessage = message || error?.message || String(error);
|
||||
const lowerMessage = errorMessage.toLowerCase();
|
||||
|
||||
// Check for explicit model error patterns
|
||||
const hasModelErrorPattern =
|
||||
lowerMessage.includes('model not found') ||
|
||||
lowerMessage.includes('checkpoint not found') ||
|
||||
lowerMessage.includes('model file not found') ||
|
||||
lowerMessage.includes('ckpt_name') ||
|
||||
lowerMessage.includes('no models available') ||
|
||||
lowerMessage.includes('safetensors') ||
|
||||
lowerMessage.includes('.ckpt') ||
|
||||
lowerMessage.includes('.pt') ||
|
||||
lowerMessage.includes('.pth') ||
|
||||
error?.code === 'MODEL_NOT_FOUND';
|
||||
|
||||
// Also check if the error contains a model file that's missing
|
||||
if (!hasModelErrorPattern) {
|
||||
const fileInfo = extractMissingFileInfo(errorMessage);
|
||||
return fileInfo !== null; // Any missing model file is considered a model error
|
||||
}
|
||||
|
||||
return hasModelErrorPattern;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the error is a ComfyUI SDK custom error
|
||||
* @param error - Error object
|
||||
* @returns Whether it's a SDK custom error
|
||||
*/
|
||||
function isSDKCustomError(error: any): boolean {
|
||||
if (!error) return false;
|
||||
|
||||
// Check for SDK error class names
|
||||
const errorName = error?.name || error?.constructor?.name || '';
|
||||
const sdkErrorTypes = [
|
||||
// Base error class
|
||||
'CallWrapperError',
|
||||
// Actual SDK error classes from comfyui-sdk
|
||||
'WentMissingError',
|
||||
'FailedCacheError',
|
||||
'EnqueueFailedError',
|
||||
'DisconnectedError',
|
||||
'ExecutionFailedError',
|
||||
'CustomEventError',
|
||||
'ExecutionInterruptedError',
|
||||
'MissingNodeError',
|
||||
];
|
||||
|
||||
if (sdkErrorTypes.includes(errorName)) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Check for SDK error messages patterns
|
||||
const message = error?.message || String(error);
|
||||
const lowerMessage = message.toLowerCase();
|
||||
|
||||
return (
|
||||
lowerMessage.includes('sdk error:') ||
|
||||
lowerMessage.includes('call wrapper') ||
|
||||
lowerMessage.includes('execution interrupted') ||
|
||||
lowerMessage.includes('missing node type') ||
|
||||
lowerMessage.includes('invalid model configuration') ||
|
||||
lowerMessage.includes('workflow validation failed') ||
|
||||
lowerMessage.includes('sdk timeout') ||
|
||||
lowerMessage.includes('sdk configuration error')
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the error is a network connection error (including WebSocket)
|
||||
* @param error - Error object
|
||||
* @param message - Pre-extracted message
|
||||
* @param code - Pre-extracted code
|
||||
* @returns Whether it's a network connection error
|
||||
*/
|
||||
function isNetworkError(error: any, message?: string, code?: string | number): boolean {
|
||||
const errorMessage = message || error?.message || String(error);
|
||||
const lowerMessage = errorMessage.toLowerCase();
|
||||
const errorCode = code || error?.code;
|
||||
|
||||
return (
|
||||
// Basic network errors
|
||||
errorMessage === 'fetch failed' ||
|
||||
lowerMessage.includes('econnrefused') ||
|
||||
lowerMessage.includes('enotfound') ||
|
||||
lowerMessage.includes('etimedout') ||
|
||||
lowerMessage.includes('network error') ||
|
||||
lowerMessage.includes('connection refused') ||
|
||||
lowerMessage.includes('connection timeout') ||
|
||||
errorCode === 'ECONNREFUSED' ||
|
||||
errorCode === 'ENOTFOUND' ||
|
||||
errorCode === 'ETIMEDOUT' ||
|
||||
// WebSocket specific errors
|
||||
lowerMessage.includes('websocket') ||
|
||||
lowerMessage.includes('ws connection') ||
|
||||
lowerMessage.includes('connection lost to comfyui server') ||
|
||||
errorCode === 'WS_CONNECTION_FAILED' ||
|
||||
errorCode === 'WS_TIMEOUT' ||
|
||||
errorCode === 'WS_HANDSHAKE_FAILED'
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the error is a ComfyUI workflow error
|
||||
* @param error - Error object
|
||||
* @param message - Pre-extracted message
|
||||
* @returns Whether it's a workflow error
|
||||
*/
|
||||
function isWorkflowError(error: any, message?: string): boolean {
|
||||
const errorMessage = message || error?.message || String(error);
|
||||
const lowerMessage = errorMessage.toLowerCase();
|
||||
|
||||
// Check for structured workflow error fields
|
||||
if (
|
||||
error &&
|
||||
typeof error === 'object' &&
|
||||
(error.node_id || error.nodeId || error.node_type || error.nodeType)
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
||||
return (
|
||||
lowerMessage.includes('node') ||
|
||||
lowerMessage.includes('workflow') ||
|
||||
lowerMessage.includes('execution') ||
|
||||
lowerMessage.includes('prompt') ||
|
||||
lowerMessage.includes('queue') ||
|
||||
lowerMessage.includes('invalid input') ||
|
||||
lowerMessage.includes('missing required') ||
|
||||
lowerMessage.includes('node execution failed') ||
|
||||
lowerMessage.includes('workflow validation') ||
|
||||
error?.type === 'workflow_error'
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Simple ComfyUI error parser
|
||||
* Extracts error information and determines error type
|
||||
*/
|
||||
function parseComfyUIErrorMessage(error: any): ParsedError {
|
||||
// Default error info
|
||||
let message = 'Unknown error';
|
||||
let status: number | undefined;
|
||||
let code: string | undefined;
|
||||
let missingFileName: string | undefined;
|
||||
let missingFileType: 'model' | 'component' | undefined;
|
||||
let userGuidance: string | undefined;
|
||||
let errorType: ILobeAgentRuntimeErrorType = AgentRuntimeErrorType.ComfyUIBizError;
|
||||
|
||||
// Check for JSON parsing errors (indicates non-ComfyUI service)
|
||||
if (
|
||||
error instanceof SyntaxError ||
|
||||
(error && typeof error === 'object' && error.name === 'SyntaxError')
|
||||
) {
|
||||
const syntaxMessage = error?.message || String(error);
|
||||
if (syntaxMessage.includes('JSON') || syntaxMessage.includes('Unexpected token')) {
|
||||
// JSON parsing failed - service is not ComfyUI
|
||||
return {
|
||||
error: {
|
||||
message: 'Service is not ComfyUI - received non-JSON response',
|
||||
type: 'SyntaxError',
|
||||
userGuidance:
|
||||
'The service at this URL is not a ComfyUI server. Please check your baseURL configuration.',
|
||||
},
|
||||
errorType: AgentRuntimeErrorType.InvalidProviderAPIKey, // Trigger auth dialog
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Extract message
|
||||
if (typeof error === 'string') {
|
||||
message = error;
|
||||
} else if (error instanceof Error) {
|
||||
message = error.message;
|
||||
code = (error as any).code;
|
||||
} else if (error && typeof error === 'object') {
|
||||
// Extract message from various possible sources (matching original logic)
|
||||
const possibleMessage = [
|
||||
error.exception_message, // ComfyUI specific field (highest priority)
|
||||
error.error?.exception_message, // Nested ComfyUI exception message
|
||||
error.error?.error, // Deeply nested error.error.error path
|
||||
error.message,
|
||||
error.error?.message,
|
||||
error.data?.message,
|
||||
error.body?.message,
|
||||
error.response?.data?.message,
|
||||
error.response?.data?.error?.message,
|
||||
error.response?.text,
|
||||
error.response?.body,
|
||||
error.statusText,
|
||||
].find(Boolean);
|
||||
|
||||
// Use the message or fallback to a generic error
|
||||
if (!possibleMessage) {
|
||||
message = 'Unknown error occurred';
|
||||
} else {
|
||||
message = possibleMessage;
|
||||
}
|
||||
|
||||
// Extract status code from various possible locations
|
||||
const possibleStatus = [
|
||||
error.status,
|
||||
error.statusCode,
|
||||
error.details?.status, // ServicesError puts status in details
|
||||
error.response?.status,
|
||||
error.response?.statusCode,
|
||||
error.error?.status,
|
||||
error.error?.statusCode,
|
||||
].find(Number.isInteger);
|
||||
|
||||
status = possibleStatus;
|
||||
code = error.code || error.error?.code || error.response?.data?.code;
|
||||
}
|
||||
|
||||
// Extract missing file information and generate guidance
|
||||
const fileInfo = extractMissingFileInfo(message);
|
||||
if (fileInfo) {
|
||||
missingFileName = fileInfo.fileName;
|
||||
missingFileType = fileInfo.fileType;
|
||||
userGuidance = generateUserGuidance(fileInfo.fileName, fileInfo.fileType);
|
||||
}
|
||||
|
||||
// Determine error type based on status code
|
||||
if (status) {
|
||||
switch (status) {
|
||||
case 400:
|
||||
case 401:
|
||||
case 404: {
|
||||
errorType = AgentRuntimeErrorType.InvalidProviderAPIKey;
|
||||
break;
|
||||
}
|
||||
case 403: {
|
||||
errorType = AgentRuntimeErrorType.PermissionDenied;
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
if (status >= 500) {
|
||||
errorType = AgentRuntimeErrorType.ComfyUIServiceUnavailable;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check for more specific error types only if it's still a generic ComfyUIBizError
|
||||
if (errorType === AgentRuntimeErrorType.ComfyUIBizError) {
|
||||
if (isSDKCustomError(error)) {
|
||||
// SDK errors remain as ComfyUIBizError
|
||||
errorType = AgentRuntimeErrorType.ComfyUIBizError;
|
||||
} else if (isNetworkError(error, message, code)) {
|
||||
errorType = AgentRuntimeErrorType.ComfyUIServiceUnavailable;
|
||||
} else if (isWorkflowError(error, message)) {
|
||||
errorType = AgentRuntimeErrorType.ComfyUIWorkflowError;
|
||||
} else if (isModelError(error, message)) {
|
||||
errorType = AgentRuntimeErrorType.ModelNotFound;
|
||||
}
|
||||
}
|
||||
|
||||
const result = {
|
||||
error: {
|
||||
code,
|
||||
message,
|
||||
missingFileName,
|
||||
missingFileType,
|
||||
status,
|
||||
type: error?.name || error?.type,
|
||||
userGuidance,
|
||||
},
|
||||
errorType,
|
||||
};
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Error Handler Service
|
||||
* Provides unified error handling and transformation
|
||||
*/
|
||||
export class ErrorHandlerService {
|
||||
/**
|
||||
* Handle and transform any error into framework error
|
||||
* Enhanced to preserve more debugging information while maintaining compatibility
|
||||
* @param error - The error to handle
|
||||
* @throws {TRPCError} Always throws a properly formatted error with cause
|
||||
*/
|
||||
handleError(error: unknown): never {
|
||||
// 1. If already a framework error, wrap in TRPCError
|
||||
if (error && typeof error === 'object' && 'errorType' in error) {
|
||||
throw new TRPCError({
|
||||
cause: error,
|
||||
code: 'INTERNAL_SERVER_ERROR',
|
||||
message: 'ComfyUI service error',
|
||||
});
|
||||
}
|
||||
|
||||
// 2. Handle ComfyUI internal errors - enhance information preservation
|
||||
if (isComfyUIInternalError(error)) {
|
||||
const errorType = this.mapInternalErrorToRuntimeError(error);
|
||||
|
||||
// Enhanced: preserve more context information
|
||||
const enhancedError = {
|
||||
details: error.details || {},
|
||||
message: error.message,
|
||||
// Preserve original error type and reason
|
||||
originalErrorType: error.constructor.name,
|
||||
originalReason: error.reason,
|
||||
// Note: Removed originalError to avoid serialization issues
|
||||
};
|
||||
|
||||
const agentError = AgentRuntimeError.createImage({
|
||||
error: enhancedError,
|
||||
errorType: errorType as ILobeAgentRuntimeErrorType,
|
||||
provider: 'comfyui',
|
||||
});
|
||||
|
||||
throw new TRPCError({
|
||||
cause: agentError,
|
||||
code: 'INTERNAL_SERVER_ERROR',
|
||||
message: error.message,
|
||||
});
|
||||
}
|
||||
|
||||
// 3. Parse other errors - use enhanced parser with more information
|
||||
const { error: parsedError, errorType } = parseComfyUIErrorMessage(error);
|
||||
|
||||
// Enhanced: add more context
|
||||
const enhancedParsedError = {
|
||||
...parsedError,
|
||||
// Add timestamp for debugging
|
||||
timestamp: new Date().toISOString(),
|
||||
// Note: Removed originalError to avoid serialization issues
|
||||
};
|
||||
|
||||
const agentError = AgentRuntimeError.createImage({
|
||||
error: enhancedParsedError,
|
||||
errorType,
|
||||
provider: 'comfyui',
|
||||
});
|
||||
|
||||
throw new TRPCError({
|
||||
cause: agentError,
|
||||
code: 'INTERNAL_SERVER_ERROR',
|
||||
message: parsedError.message || 'ComfyUI service error',
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Map internal ComfyUI errors to runtime error types
|
||||
*/
|
||||
private mapInternalErrorToRuntimeError(
|
||||
error: ConfigError | WorkflowError | UtilsError | ServicesError | ModelResolverError,
|
||||
): string {
|
||||
if (error instanceof ConfigError) {
|
||||
const mapping: Record<string, string> = {
|
||||
[ConfigError.Reasons.INVALID_CONFIG]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
[ConfigError.Reasons.MISSING_CONFIG]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
[ConfigError.Reasons.CONFIG_PARSE_ERROR]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
[ConfigError.Reasons.REGISTRY_ERROR]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
};
|
||||
return mapping[error.reason] || AgentRuntimeErrorType.ComfyUIBizError;
|
||||
}
|
||||
|
||||
if (error instanceof WorkflowError) {
|
||||
const mapping: Record<string, string> = {
|
||||
[WorkflowError.Reasons.INVALID_CONFIG]: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
[WorkflowError.Reasons.MISSING_COMPONENT]: AgentRuntimeErrorType.ComfyUIModelError,
|
||||
[WorkflowError.Reasons.MISSING_ENCODER]: AgentRuntimeErrorType.ComfyUIModelError,
|
||||
[WorkflowError.Reasons.UNSUPPORTED_MODEL]: AgentRuntimeErrorType.ModelNotFound,
|
||||
[WorkflowError.Reasons.INVALID_PARAMS]: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
};
|
||||
return mapping[error.reason] || AgentRuntimeErrorType.ComfyUIWorkflowError;
|
||||
}
|
||||
|
||||
if (error instanceof ServicesError) {
|
||||
// If error already has parsed errorType in details, use it directly
|
||||
if (error.details?.errorType) {
|
||||
return error.details.errorType;
|
||||
}
|
||||
|
||||
// Otherwise use mapping table
|
||||
const mapping: Record<string, string> = {
|
||||
[ServicesError.Reasons.INVALID_ARGS]: AgentRuntimeErrorType.InvalidComfyUIArgs,
|
||||
[ServicesError.Reasons.INVALID_AUTH]: AgentRuntimeErrorType.InvalidProviderAPIKey,
|
||||
[ServicesError.Reasons.INVALID_CONFIG]: AgentRuntimeErrorType.InvalidComfyUIArgs,
|
||||
[ServicesError.Reasons.CONNECTION_FAILED]: AgentRuntimeErrorType.InvalidProviderAPIKey, // Trigger auth dialog for connection issues
|
||||
[ServicesError.Reasons.UPLOAD_FAILED]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
[ServicesError.Reasons.EXECUTION_FAILED]: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
[ServicesError.Reasons.MODEL_NOT_FOUND]: AgentRuntimeErrorType.ModelNotFound,
|
||||
[ServicesError.Reasons.EMPTY_RESULT]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
[ServicesError.Reasons.IMAGE_FETCH_FAILED]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
[ServicesError.Reasons.IMAGE_TOO_LARGE]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
[ServicesError.Reasons.UNSUPPORTED_PROTOCOL]: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
[ServicesError.Reasons.MODEL_VALIDATION_FAILED]: AgentRuntimeErrorType.ModelNotFound,
|
||||
[ServicesError.Reasons.WORKFLOW_BUILD_FAILED]: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
};
|
||||
return mapping[error.reason] || AgentRuntimeErrorType.ComfyUIBizError;
|
||||
}
|
||||
|
||||
if (error instanceof UtilsError || error instanceof ModelResolverError) {
|
||||
const mapping: Record<string, string> = {
|
||||
CONNECTION_ERROR: AgentRuntimeErrorType.ComfyUIServiceUnavailable,
|
||||
DETECTION_FAILED: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
INVALID_API_KEY: AgentRuntimeErrorType.InvalidProviderAPIKey,
|
||||
INVALID_MODEL_FORMAT: AgentRuntimeErrorType.ComfyUIBizError,
|
||||
MODEL_NOT_FOUND: AgentRuntimeErrorType.ModelNotFound,
|
||||
NO_BUILDER_FOUND: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
PERMISSION_DENIED: AgentRuntimeErrorType.PermissionDenied,
|
||||
ROUTING_FAILED: AgentRuntimeErrorType.ComfyUIWorkflowError,
|
||||
SERVICE_UNAVAILABLE: AgentRuntimeErrorType.ComfyUIServiceUnavailable,
|
||||
};
|
||||
return mapping[error.reason] || AgentRuntimeErrorType.ComfyUIBizError;
|
||||
}
|
||||
|
||||
return AgentRuntimeErrorType.ComfyUIBizError;
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue