download-img.js 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. /**
  2. * fun 结点:根据描述词(prompt)用 python/imagedl 搜索并下载一张图片到 savePath
  3. * 入参:prompt(描述词), savePath(相对路径时基于 folderPath)
  4. */
  5. const path = require('path')
  6. const fs = require('fs')
  7. const { spawnSync } = require('child_process')
  8. const { getPythonExeFromConfig } = require('../../../python-exe-from-config.js')
  9. const configPath = process.env.STATIC_ROOT
  10. ? path.join(path.dirname(path.resolve(process.env.STATIC_ROOT)), 'config.js')
  11. : path.join(__dirname, '..', '..', '..', '..', 'config.js')
  12. const config = fs.existsSync(configPath) ? require(configPath) : {}
  13. const projectRoot = (config.projectRoot && fs.existsSync(config.projectRoot))
  14. ? config.projectRoot
  15. : path.dirname(path.resolve(configPath))
  16. const scriptPath = path.join(projectRoot, 'python', 'scripts', 'download-img-by-prompt.py')
  17. const imagedlParent = path.join(projectRoot, 'python')
  18. const imagedlRequirements = path.join(projectRoot, 'python', 'imagedl', 'requirements.txt')
  19. function buildSavePath(savePath, folderPath) {
  20. if (!savePath || typeof savePath !== 'string') return null
  21. const trimmed = savePath.trim()
  22. if (path.isAbsolute(trimmed) || trimmed.match(/^[A-Za-z]:/)) return trimmed
  23. return folderPath ? path.join(folderPath, trimmed) : path.resolve(projectRoot, trimmed)
  24. }
  25. /** spawnSync 非 0 退出时拼详细日志(含启动失败、信号、截断后的 stdout/stderr) */
  26. function formatSpawnFailure (r, pythonPath, scriptPath) {
  27. const maxLen = 1500
  28. const trunc = (s) => {
  29. const t = String(s || '').trim()
  30. if (!t) return ''
  31. return t.length > maxLen ? `${t.slice(0, maxLen)}…(truncated)` : t
  32. }
  33. const parts = []
  34. if (pythonPath) {
  35. const abs = path.isAbsolute(pythonPath) || /^[A-Za-z]:/.test(pythonPath)
  36. const missing = abs && !fs.existsSync(pythonPath)
  37. parts.push(`python=${pythonPath}${missing ? ' (file missing)' : ''}`)
  38. }
  39. if (scriptPath) parts.push(`script=${scriptPath}`)
  40. if (r.error) {
  41. const e = r.error
  42. const code = e.code != null ? String(e.code) : ''
  43. parts.push(`spawnError${code ? `[${code}]` : ''}: ${e.message || e}`)
  44. }
  45. if (r.signal) parts.push(`signal=${r.signal}`)
  46. if (r.status !== null && r.status !== undefined) parts.push(`exitCode=${r.status}`)
  47. else parts.push('exitCode=(null)')
  48. const stderr = trunc(r.stderr)
  49. const stdout = trunc(r.stdout)
  50. if (stderr) parts.push(`stderr: ${stderr}`)
  51. if (stdout) parts.push(`stdout: ${stdout}`)
  52. return parts.length ? parts.join(' | ') : 'download-img 执行失败(无子进程输出)'
  53. }
  54. async function executeDownloadImg({ prompt, savePath, folderPath }) {
  55. if (!prompt || typeof prompt !== 'string' || !prompt.trim()) return { success: false, error: 'download-img 缺少 prompt 参数' }
  56. if (savePath == null) return { success: false, error: 'download-img 缺少 savePath 参数' }
  57. const absolutePath = buildSavePath(String(savePath).trim(), folderPath)
  58. if (!absolutePath) return { success: false, error: 'download-img savePath 无效' }
  59. if (!fs.existsSync(scriptPath)) return { success: false, error: `脚本不存在: ${scriptPath}` }
  60. const pythonPath = getPythonExeFromConfig(config)
  61. const args = [scriptPath, '--prompt', prompt.trim(), '--save-path', absolutePath.replace(/\\/g, '/')]
  62. const runScript = () => spawnSync(pythonPath, args, {
  63. encoding: 'utf-8',
  64. timeout: 120000,
  65. env: { ...process.env, PYTHONIOENCODING: 'utf-8', IMAGEDL_PARENT: imagedlParent },
  66. cwd: projectRoot
  67. })
  68. let r = runScript()
  69. let out = (r.stdout || '').trim()
  70. const err = (r.stderr || '').trim()
  71. const fullOut = out + (err ? '\n' + err : '')
  72. if (r.status !== 0 && fs.existsSync(imagedlRequirements) && (fullOut.includes('No module named') || fullOut.includes('ModuleNotFoundError'))) {
  73. spawnSync(pythonPath, ['-m', 'pip', 'install', '-r', imagedlRequirements, '-q'], {
  74. encoding: 'utf-8',
  75. timeout: 180000,
  76. cwd: projectRoot
  77. })
  78. r = runScript()
  79. out = (r.stdout || '').trim()
  80. }
  81. if (r.status !== 0) {
  82. return { success: false, error: formatSpawnFailure(r, pythonPath, scriptPath) }
  83. }
  84. // stdout 可能包含进度条等,取最后一行以 { 开头的行作为 JSON
  85. const lines = out.split(/\r?\n/).map(s => s.trim()).filter(Boolean)
  86. let jsonStr = lines[lines.length - 1]
  87. if (!jsonStr || !jsonStr.startsWith('{')) {
  88. const lastJson = lines.filter(l => l.startsWith('{')).pop()
  89. jsonStr = lastJson || out
  90. }
  91. let result
  92. try {
  93. result = JSON.parse(jsonStr)
  94. } catch (_) {
  95. const raw = String(out || '').trim()
  96. const tail = raw.length > 800 ? raw.slice(-800) : raw
  97. return { success: false, error: `无法解析 JSON 输出 | stdout(tail): ${tail || '(empty)'}` }
  98. }
  99. if (!result.success) return { success: false, error: result.error || '未下载到图片' }
  100. return { success: true, path: result.path }
  101. }
  102. module.exports = { executeDownloadImg }