LaiTool/src/main/Service/watermark.ts

import path from 'path'
import fs from 'fs'
import util from 'util'
import { exec } from 'child_process'
const execAsync = util.promisify(exec);
import { CheckFileOrDirExist, CheckFolderExistsOrCreate } from '../../define/Tools/file'
import { errorMessage, successMessage } from '../Public/generalTools'
import { SoftwareService } from '../../define/db/service/SoftWare/softwareService'
import { isEmpty } from 'lodash'
import { RemoveWatermarkType } from '../../define/enum/waterMarkAndSubtitle'
import { ValidateJson } from '../../define/Tools/validate'
import { define } from '../../define/define'
import { LOGGER_DEFINE } from '../../define/logger_define'
import axios from 'axios'
import { Base64ToFile } from '../../define/Tools/image'
import { TaskScheduler } from './taskScheduler';
import { LoggerStatus, OtherData } from '../../define/enum/softwareEnum';
import { basicApi } from '../../api/apiBasic';

export class Watermark {
  softwareService: SoftwareService
  taskScheduler: TaskScheduler;
  constructor() { }

  async InitService() {
    if (!this.softwareService) {
      this.softwareService = await SoftwareService.getInstance()
    }
    if (!this.taskScheduler) {
      this.taskScheduler = new TaskScheduler()
    }
  }

  //#region  设置

  /**
   * 初始化水印设置
   */
  InitWatermarkSetting() {
    return {
      selectModel: RemoveWatermarkType.LOCAL_LAMA,
      iopaint: {
        url: undefined
      }
    } as ImageModel.RemoveWatermarkSetting
  }

  /**
   * 获取去除水印的设置
   */
  async GetWatermarkSetting() {
    try {
      await this.InitService();
      let watermarkSetting = this.InitWatermarkSetting();
      let setting = this.softwareService.GetSoftWarePropertyData("watermarkSetting");
      if (!isEmpty(setting)) {
        // 直接初始化
        let tryP = ValidateJson(setting);
        if (!tryP) {
          throw new Error('解析去除水印的设置失败，数据格式不正确')
        }
        watermarkSetting = JSON.parse(setting);
      }
      return successMessage(watermarkSetting, '获取去除水印的设置成功', 'Image_GetWatermarkSetting')
    } catch (error) {
      return errorMessage("获取去除水印的设置失败，失败信息如下：" + error.toString(), 'Image_GetWatermarkSetting')
    }
  }

  /**
   * 保存去除水印的设置
   * @param value
   * @returns 
   */
  async SaveWatermarkSetting(value: ImageModel.RemoveWatermarkSetting) {
    try {
      await this.InitService();
      if (value.selectModel == RemoveWatermarkType.IOPAINT && !value.iopaint.url) {
        throw new Error('iopaint模式，iopaint的地址不能为空')
      }

      if (value.selectModel == RemoveWatermarkType.LOCAL_LAMA) {
        // 简单判断环境是不是存在
        let lamaPath = path.join(define.scripts_path, 'lama/lama_inpaint.exe');
        if (!(await CheckFileOrDirExist(lamaPath))) {
          throw new Error('本地LAMA环墋不存在，请先安装，详细信息看下方的 ‘lama安装教程’')
        }
      }

      if (value.iopaint.url) {
        // 验证iopaint的地址是不是一个合法的url
        let reg = new RegExp('^(http|https)://')
        if (!reg.test(value.iopaint.url)) {
          throw new Error('iopaint的地址不是一个合法的网络地址')
        }
        if (!value.iopaint.url.endsWith('/')) {
          value.iopaint.url = value.iopaint.url + '/'
        }
      }


      // 保存数据
      let res = this.softwareService.SaveSoftwarePropertyData('watermarkSetting', JSON.stringify(value));
      return successMessage(null, "保存去除水印的设置成功", 'Image_SaveWatermarkSetting')
    } catch (error) {
      return errorMessage("保存去除水印的设置失败，失败信息如下：" + error.toString(), 'Image_SaveWatermarkSetting')
    }
  }

  //#endregion

  //#region 去除水印

  /**
   * 去除水印，试用本地lama
   * @param value 
   */
  async ProcessImageByLama(value: ImageModel.ProcessImageParams): Promise<Buffer | string> {
    try {
      let lama_script = path.resolve(define.scripts_path, `lama/lama_inpaint.exe`)
      // 就是判断指定的文件和文件夹是不是存在
      let has_exe = await CheckFileOrDirExist(lama_script)
      if (!has_exe) {
        throw new Error('图片水印处理组件不存在，请看教程自行下载')
      }
      let has_model = await CheckFileOrDirExist(
        path.resolve(define.scripts_path, 'lama/model/big-lama.pt')
      )
      if (!has_model) {
        throw new Error('图片水印处理的模型不存在，请看教程自行下载')
      }
      global.logger.info(
        LOGGER_DEFINE.REMOVE_WATERMARK,
        `开始使用lama去除水印，开始调用lama程序`
      )
      let script = `cd "${path.dirname(lama_script)}" && "${lama_script}" "-l" "${value.inputFilePath}" "${value.maskPath}" "${value.outFilePath}"`
      let scriptRes = await execAsync(script, { maxBuffer: 1024 * 1024 * 10, encoding: 'utf-8' })
      // 判断是不是有输出文件
      let has_out = await CheckFileOrDirExist(value.outFilePath)
      if (!has_out) {
        throw new Error('lama去除水印失败，没有输出文件')
      }

      if (value.type == 'arrayBuffer') {
        // 读取导出的文件
        let res_data = await fs.promises.readFile(value.outFilePath)
        return res_data
      } else if (value.type == 'file') {
        return value.outFilePath
      }

    } catch (error) {
      throw new Error('使用Lama去除水印失败，失败信息如下：' + error.toString())
    }
  }

  /**
   * 去除s水印，使用iopaint 
   * @param value 传入的参数
   * @param iopaintUrl 
   */
  async ProcessImageByIopaint(value: ImageModel.ProcessImageParams, iopaintUrl: string): Promise<Buffer | string> {
    try {
      iopaintUrl = iopaintUrl + 'api/v1/inpaint'
      let headers = {
        accept: '*/*',
        'accept-language': 'zh-CN,zh;q=0.9',
        'content-type': 'application/json'
      }
      let data = {
        image: value.imageBase64,
        mask: value.maskBase64,
        ldm_steps: 30,
        ldm_sampler: 'ddim',
        zits_wireframe: true,
        cv2_flag: 'INPAINT_NS',
        cv2_radius: 5,
        hd_strategy: 'Crop',
        hd_strategy_crop_triger_size: 640,
        hd_strategy_crop_margin: 128,
        hd_trategy_resize_imit: 2048 * 5,
        prompt: '',
        negative_prompt:
          'out of frame, lowres, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, disfigured, gross proportions, malformed limbs, watermark, signature',
        use_croper: false,
        croper_x: 284,
        croper_y: 284,
        croper_height: 512,
        croper_width: 512,
        use_extender: false,
        extender_x: 0,
        extender_y: 0,
        extender_height: 1080,
        extender_width: 1080,
        sd_mask_blur: 12,
        sd_strength: 1,
        sd_steps: 50,
        sd_guidance_scale: 7.5,
        sd_sampler: 'DPM++ 2M',
        sd_seed: -1,
        sd_match_histograms: false,
        sd_lcm_lora: false,
        paint_by_example_example_image: null,
        p2p_image_guidance_scale: 1.5,
        enable_controlnet: false,
        controlnet_conditioning_scale: 0.4,
        controlnet_method: '',
        enable_brushnet: false,
        brushnet_method: 'random_mask',
        brushnet_conditioning_scale: 1,
        enable_powerpaint_v2: false,
        powerpaint_task: 'text-guided'
      }
      let res = await basicApi.post(iopaintUrl, data, headers)

      if (value.type == 'arrayBuffer') {
        return res.data
      } else if (value.type == 'file') {
        let buffer = Buffer.from(res.data)
        await fs.promises.writeFile(value.outFilePath, buffer)
        return value.outFilePath
      }
    } catch (error) {
      throw new Error('使用Iopaint去除水印失败，失败信息如下：' + error.toString())
    }
  }

  /**
   * 去除图片水印，当前只支持本地，后续会支持 iopaint 
   * @param value 
   */
  async ProcessImage(value: ImageModel.ProcessImageParams) {
    let outDir = path.dirname(value.outFilePath)
    await CheckFolderExistsOrCreate(outDir);

    let watermarkSettingRes = await this.GetWatermarkSetting();
    if (watermarkSettingRes.code == 0) {
      throw new Error(watermarkSettingRes.message)
    }
    let watermarkSetting = watermarkSettingRes.data as ImageModel.RemoveWatermarkSetting;

    switch (watermarkSetting.selectModel) {
      case RemoveWatermarkType.LOCAL_LAMA:
        return await this.ProcessImageByLama(value)
      case RemoveWatermarkType.IOPAINT:
        return await this.ProcessImageByIopaint(value, watermarkSetting.iopaint.url)
      default:
        throw new Error('未知的去除水印模式')
    }
  }

  //#endregion
}