forked from xiangwang25/whale-town-end-v2
400 lines
14 KiB
TypeScript
400 lines
14 KiB
TypeScript
import { BadRequestException, ForbiddenException, Injectable, Logger, NotFoundException } from '@nestjs/common';
|
||
import { ConfigService } from '@nestjs/config';
|
||
import { spawn } from 'child_process';
|
||
import { randomUUID } from 'crypto';
|
||
import { existsSync } from 'fs';
|
||
import { mkdir, readFile, readdir, writeFile } from 'fs/promises';
|
||
import { join, resolve } from 'path';
|
||
import { AccountProfileService } from '../auth/account_profile.service';
|
||
import { CreateSkinGenerationJobDto } from './dto/create_skin_generation_job.dto';
|
||
import { SkinGenerationJob, SkinGenerationJobResponse } from './skin_generation.types';
|
||
|
||
@Injectable()
|
||
export class SkinGenerationService {
|
||
private readonly logger = new Logger(SkinGenerationService.name);
|
||
private readonly jobs = new Map<string, SkinGenerationJob>();
|
||
|
||
constructor(
|
||
private readonly configService: ConfigService,
|
||
private readonly accountProfileService: AccountProfileService,
|
||
) {}
|
||
|
||
async createJob(userId: bigint, dto: CreateSkinGenerationJobDto): Promise<SkinGenerationJobResponse> {
|
||
this.accountProfileService.assertCustomSkinCreationAvailable();
|
||
const apiKey = this.configService.get<string>('NOVAMAILIO_API_KEY') || process.env.NOVAMAILIO_API_KEY;
|
||
if (!apiKey || apiKey.trim().length === 0) {
|
||
throw new BadRequestException('服务端尚未配置 NOVAMAILIO_API_KEY，无法生成角色皮肤');
|
||
}
|
||
await this.ensureWorkerRuntime();
|
||
if (!(await this.accountProfileService.canUseRegistrationSkinGeneration(userId))) {
|
||
throw new BadRequestException('该账号没有可用的注册角色生成机会');
|
||
}
|
||
if (await this.accountProfileService.hasRegistrationGeneratedSkin(userId)) {
|
||
throw new BadRequestException('该账号已经使用过注册角色生成机会');
|
||
}
|
||
const activeJob = Array.from(this.jobs.values()).find(
|
||
(job) => job.userId === userId.toString() && (job.status === 'queued' || job.status === 'running'),
|
||
);
|
||
if (activeJob) {
|
||
return this.toResponse(activeJob);
|
||
}
|
||
|
||
const jobId = randomUUID();
|
||
const ownerId = userId.toString();
|
||
const now = Date.now();
|
||
const safeName = this.sanitizeName(dto.name || '我的角色');
|
||
const outDir = join(this.getOutputRoot(), ownerId, jobId);
|
||
const sourceImagePath = join(outDir, 'source_character.png');
|
||
await mkdir(outDir, { recursive: true });
|
||
await this.saveSourceImage(dto.source_image_base64, sourceImagePath);
|
||
const job: SkinGenerationJob = {
|
||
jobId,
|
||
userId: ownerId,
|
||
status: 'queued',
|
||
stage: 'queued',
|
||
message: '生成任务已创建，等待开始',
|
||
name: safeName,
|
||
sourceImagePath,
|
||
createdAt: now,
|
||
updatedAt: now,
|
||
outDir,
|
||
statusJsonPath: join(outDir, 'status.json'),
|
||
resultJsonPath: join(outDir, 'result.json'),
|
||
};
|
||
|
||
this.jobs.set(jobId, job);
|
||
void this.runJob(job);
|
||
return this.toResponse(job);
|
||
}
|
||
|
||
async getJob(userId: bigint, jobId: string): Promise<SkinGenerationJobResponse> {
|
||
const job = this.jobs.get(jobId) || await this.restoreJobFromDisk(jobId);
|
||
if (!job) {
|
||
throw new NotFoundException('皮肤生成任务不存在或已过期');
|
||
}
|
||
if (job.userId !== userId.toString()) {
|
||
throw new ForbiddenException('无权查看该皮肤生成任务');
|
||
}
|
||
|
||
await this.refreshJobFromFiles(job);
|
||
return this.toResponse(job);
|
||
}
|
||
|
||
private async runJob(job: SkinGenerationJob): Promise<void> {
|
||
job.status = 'running';
|
||
job.stage = 'start';
|
||
job.message = '正在启动服务器角色生成流程';
|
||
job.updatedAt = Date.now();
|
||
|
||
try {
|
||
const scriptPath = this.getScriptPath();
|
||
const pythonPath = this.getPythonPath();
|
||
if (!existsSync(scriptPath)) {
|
||
throw new Error(`皮肤生成脚本不存在: ${scriptPath}`);
|
||
}
|
||
|
||
const child = spawn(
|
||
pythonPath,
|
||
[
|
||
scriptPath,
|
||
'--source-image',
|
||
job.sourceImagePath,
|
||
'--out-dir',
|
||
job.outDir,
|
||
'--name',
|
||
job.name,
|
||
'--result-json',
|
||
job.resultJsonPath,
|
||
'--status-json',
|
||
job.statusJsonPath,
|
||
],
|
||
{
|
||
cwd: this.getBackendRoot(),
|
||
env: process.env,
|
||
stdio: ['ignore', 'pipe', 'pipe'],
|
||
},
|
||
);
|
||
|
||
child.stdout.on('data', (chunk: Buffer) => this.consumeWorkerOutput(job, chunk));
|
||
child.stderr.on('data', (chunk: Buffer) => this.consumeWorkerOutput(job, chunk));
|
||
|
||
const workerExit = await new Promise<{ exitCode: number | null; signal: NodeJS.Signals | null }>(
|
||
(resolveExit, rejectExit) => {
|
||
child.on('error', rejectExit);
|
||
child.on('close', (exitCode, signal) => resolveExit({ exitCode, signal }));
|
||
},
|
||
);
|
||
|
||
await this.refreshJobFromFiles(job);
|
||
if (workerExit.signal) {
|
||
throw new Error(`生成进程被中断：${workerExit.signal}`);
|
||
}
|
||
if (workerExit.exitCode !== 0 || job.stage === 'failed' || job.error) {
|
||
throw new Error(job.error || job.message || `生成进程退出码: ${workerExit.exitCode}`);
|
||
}
|
||
|
||
job.status = 'completed';
|
||
job.stage = 'done';
|
||
job.message = '新角色皮肤已生成';
|
||
job.updatedAt = Date.now();
|
||
if (job.spritesheetPath) {
|
||
const skinAsset = await this.accountProfileService.saveGeneratedSkinForUser(
|
||
BigInt(job.userId),
|
||
job.spritesheetPath,
|
||
job.name,
|
||
'generated_registration',
|
||
);
|
||
await this.accountProfileService.consumeRegistrationSkinGeneration(BigInt(job.userId));
|
||
job.skinId = skinAsset.skinId;
|
||
job.textureUrl = skinAsset.texture_url;
|
||
}
|
||
} catch (error) {
|
||
const message = error instanceof Error ? error.message : String(error);
|
||
this.logger.error(`皮肤生成失败: ${job.jobId} ${message}`);
|
||
job.status = 'failed';
|
||
job.stage = 'failed';
|
||
job.message = `生成失败：${message}`;
|
||
job.error = message;
|
||
job.updatedAt = Date.now();
|
||
}
|
||
}
|
||
|
||
private consumeWorkerOutput(job: SkinGenerationJob, chunk: Buffer): void {
|
||
const line = chunk.toString('utf8').trim();
|
||
if (!line) {
|
||
return;
|
||
}
|
||
job.message = line.split('\n').pop() || job.message;
|
||
job.updatedAt = Date.now();
|
||
}
|
||
|
||
private async refreshJobFromFiles(job: SkinGenerationJob): Promise<void> {
|
||
await this.refreshStatusJson(job);
|
||
await this.refreshResultJson(job);
|
||
}
|
||
|
||
private async restoreJobFromDisk(jobId: string): Promise<SkinGenerationJob | null> {
|
||
const outputRoot = this.getOutputRoot();
|
||
try {
|
||
const ownerDirs = await readdir(outputRoot, { withFileTypes: true });
|
||
for (const ownerDir of ownerDirs) {
|
||
if (!ownerDir.isDirectory()) {
|
||
continue;
|
||
}
|
||
const outDir = join(outputRoot, ownerDir.name, jobId);
|
||
if (!existsSync(outDir)) {
|
||
continue;
|
||
}
|
||
const now = Date.now();
|
||
const job: SkinGenerationJob = {
|
||
jobId,
|
||
userId: ownerDir.name,
|
||
status: 'running',
|
||
stage: 'restored',
|
||
message: '正在恢复生成任务状态',
|
||
name: 'custom_whale_human',
|
||
sourceImagePath: join(outDir, 'source_character.png'),
|
||
createdAt: now,
|
||
updatedAt: now,
|
||
outDir,
|
||
statusJsonPath: join(outDir, 'status.json'),
|
||
resultJsonPath: join(outDir, 'result.json'),
|
||
};
|
||
await this.refreshJobFromFiles(job);
|
||
this.jobs.set(jobId, job);
|
||
return job;
|
||
}
|
||
} catch {
|
||
return null;
|
||
}
|
||
return null;
|
||
}
|
||
|
||
private async refreshStatusJson(job: SkinGenerationJob): Promise<void> {
|
||
const payload = await this.readJson(job.statusJsonPath);
|
||
if (!payload) {
|
||
return;
|
||
}
|
||
job.stage = String(payload.stage || job.stage);
|
||
job.message = String(payload.message || job.message);
|
||
job.updatedAt = Number(payload.updated_at ? Math.floor(payload.updated_at * 1000) : Date.now());
|
||
if (payload.ok === false && job.stage === 'failed') {
|
||
job.status = 'failed';
|
||
job.error = job.message.replace(/^生成失败：/, '');
|
||
}
|
||
}
|
||
|
||
private async refreshResultJson(job: SkinGenerationJob): Promise<void> {
|
||
const payload = await this.readJson(job.resultJsonPath);
|
||
if (!payload) {
|
||
return;
|
||
}
|
||
job.logPath = String(payload.log_path || job.logPath || '');
|
||
if (payload.ok === true) {
|
||
job.status = 'completed';
|
||
job.stage = 'done';
|
||
job.message = '新角色皮肤已生成';
|
||
job.spritesheetPath = String(payload.spritesheet_path || '');
|
||
job.reviewPath = String(payload.review_path || '');
|
||
job.feetZoomPath = String(payload.feet_zoom_path || '');
|
||
job.updatedAt = Date.now();
|
||
return;
|
||
}
|
||
if (payload.ok === false && payload.error) {
|
||
job.status = 'failed';
|
||
job.stage = 'failed';
|
||
job.error = String(payload.error);
|
||
job.message = `生成失败：${job.error}`;
|
||
job.updatedAt = Date.now();
|
||
}
|
||
}
|
||
|
||
private async readJson(path: string): Promise<Record<string, any> | null> {
|
||
try {
|
||
if (!existsSync(path)) {
|
||
return null;
|
||
}
|
||
const text = await readFile(path, 'utf8');
|
||
if (!text.trim()) {
|
||
return null;
|
||
}
|
||
return JSON.parse(text);
|
||
} catch {
|
||
return null;
|
||
}
|
||
}
|
||
|
||
private async toResponse(job: SkinGenerationJob): Promise<SkinGenerationJobResponse> {
|
||
await this.ensureCompletedJobSkinAsset(job);
|
||
const response: SkinGenerationJobResponse = {
|
||
job_id: job.jobId,
|
||
status: job.status,
|
||
stage: job.stage,
|
||
message: job.message,
|
||
name: job.name,
|
||
created_at: job.createdAt,
|
||
updated_at: job.updatedAt,
|
||
hframes: 8,
|
||
vframes: 4,
|
||
};
|
||
|
||
if (job.status === 'failed' && job.error) {
|
||
response.error = job.error;
|
||
}
|
||
|
||
if (job.status === 'completed' && job.spritesheetPath && existsSync(job.spritesheetPath)) {
|
||
const image = await readFile(job.spritesheetPath);
|
||
response.spritesheet_base64 = image.toString('base64');
|
||
response.mime_type = 'image/png';
|
||
response.skin_id = job.skinId;
|
||
response.texture_url = job.textureUrl;
|
||
}
|
||
|
||
return response;
|
||
}
|
||
|
||
private async ensureCompletedJobSkinAsset(job: SkinGenerationJob): Promise<void> {
|
||
if (job.status !== 'completed' || !job.spritesheetPath || job.skinId) {
|
||
return;
|
||
}
|
||
if (!existsSync(job.spritesheetPath)) {
|
||
return;
|
||
}
|
||
const skinAsset = await this.accountProfileService.saveGeneratedSkinForUser(
|
||
BigInt(job.userId),
|
||
job.spritesheetPath,
|
||
job.name,
|
||
'generated_registration',
|
||
);
|
||
await this.accountProfileService.consumeRegistrationSkinGeneration(BigInt(job.userId));
|
||
job.skinId = skinAsset.skinId;
|
||
job.textureUrl = skinAsset.texture_url;
|
||
}
|
||
|
||
private getOutputRoot(): string {
|
||
return resolve(this.getBackendRoot(), this.configService.get<string>('SKIN_GENERATION_OUTPUT_DIR') || 'generated/skins');
|
||
}
|
||
|
||
private getScriptPath(): string {
|
||
return resolve(
|
||
this.getBackendRoot(),
|
||
this.configService.get<string>('SKIN_GENERATION_SCRIPT_PATH') ||
|
||
'scripts/skin_generation/generate_skin_from_prompt.py',
|
||
);
|
||
}
|
||
|
||
private getPythonPath(): string {
|
||
return this.configService.get<string>('SKIN_GENERATION_PYTHON') || process.env.PYTHON || 'python3';
|
||
}
|
||
|
||
private getBackendRoot(): string {
|
||
return resolve(__dirname, '../../..');
|
||
}
|
||
|
||
private async ensureWorkerRuntime(): Promise<void> {
|
||
const scriptPath = this.getScriptPath();
|
||
if (!existsSync(scriptPath)) {
|
||
throw new BadRequestException(`服务端角色生成脚本不存在: ${scriptPath}`);
|
||
}
|
||
|
||
const pythonPath = this.getPythonPath();
|
||
const result = await new Promise<{ exitCode: number | null; stderr: string }>((resolveResult) => {
|
||
const child = spawn(
|
||
pythonPath,
|
||
[
|
||
'-c',
|
||
'import einops, kornia, numpy, scipy, timm, torch, torchvision, transformers; from PIL import Image',
|
||
],
|
||
{
|
||
cwd: this.getBackendRoot(),
|
||
env: process.env,
|
||
stdio: ['ignore', 'ignore', 'pipe'],
|
||
},
|
||
);
|
||
let stderr = '';
|
||
child.stderr.on('data', (chunk: Buffer) => {
|
||
if (stderr.length < 2000) stderr += chunk.toString('utf8');
|
||
});
|
||
child.on('error', (error) => resolveResult({ exitCode: null, stderr: error.message }));
|
||
child.on('close', (exitCode) => resolveResult({ exitCode, stderr }));
|
||
});
|
||
if (result.exitCode !== 0) {
|
||
this.logger.error(`角色生成运行时不可用: python=${pythonPath} ${result.stderr.trim()}`);
|
||
throw new BadRequestException('服务端角色生成运行时未就绪，请联系管理员');
|
||
}
|
||
}
|
||
|
||
private async saveSourceImage(base64: string, destinationPath: string): Promise<void> {
|
||
const normalized = base64.trim().replace(/^data:image\/[a-zA-Z0-9.+-]+;base64,/, '');
|
||
let buffer: Buffer;
|
||
try {
|
||
buffer = Buffer.from(normalized, 'base64');
|
||
} catch {
|
||
throw new BadRequestException('角色参考图解析失败');
|
||
}
|
||
if (buffer.length < 512) {
|
||
throw new BadRequestException('角色参考图内容为空或过小');
|
||
}
|
||
if (buffer.length > 8 * 1024 * 1024) {
|
||
throw new BadRequestException('角色参考图不能超过8MB');
|
||
}
|
||
const isPng = buffer.length >= 8 && buffer.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]));
|
||
const isJpeg = buffer.length >= 3 && buffer[0] === 0xff && buffer[1] === 0xd8 && buffer[2] === 0xff;
|
||
const isWebp = buffer.length >= 12 && buffer.subarray(0, 4).toString('ascii') === 'RIFF' && buffer.subarray(8, 12).toString('ascii') === 'WEBP';
|
||
if (!isPng && !isJpeg && !isWebp) {
|
||
throw new BadRequestException('角色参考图必须是 PNG、JPG 或 WebP');
|
||
}
|
||
await writeFile(destinationPath, buffer);
|
||
}
|
||
|
||
private sanitizeName(value: string): string {
|
||
const normalized = value
|
||
.trim()
|
||
.toLowerCase()
|
||
.replace(/[^a-zA-Z0-9_\u4e00-\u9fa5]+/g, '_')
|
||
.replace(/_+/g, '_')
|
||
.replace(/^_+|_+$/g, '');
|
||
return (normalized || 'custom_whale_human').slice(0, 40);
|
||
}
|
||
}
|