XPrintServer/XPrintServer.Api/Controllers/Ai/ComfyController.cs

using CrazyStudio.Core.Common;
using CrazyStudio.Core.Common.AliyunExt.Oss.Tools;
using CrazyStudio.Core.Common.Eitities;
using CrazyStudio.Core.Common.Tools;
using Microsoft.AspNetCore.Mvc;
using Microsoft.Identity.Client;
using Newtonsoft.Json;
using System;
using System.Collections.Concurrent;
using System.ComponentModel.DataAnnotations;
using System.Text.Json;
using System.Threading.Tasks;
using TouchSocket.Core;
using XPrintServer.Api.Dto;
using XPrintServer.Api.Entities;
using XPrintServer.Api.Tools;
using XPrintServer.Business.Services;
using XPrintServer.Business.Workflows;
using XPrintServer.DataModel;
using XPrintServer.DataModel.Models;
using File = System.IO.File;

namespace XPrintServer.Api.Controllers.Ai
{

    [Route("api/[controller]/[action]")]
    [ApiController]
    public class ComfyController : ControllerBase
    {

        private readonly IHttpContextAccessor _httpContextAccessor;

        private readonly AliyunInfos _aliyunInfos;
        private readonly OssTool _ossTool;

        private readonly ComfyUIService _compyUIService;

        private static ConcurrentDictionary<string, string> _tmpInputPath = new ConcurrentDictionary<string, string>();
        public ComfyController(ComfyUIService compyUIService, IHttpContextAccessor httpContextAccessor, AliyunInfos aliyunInfos, OssTool ossTool)
        {
            _aliyunInfos = aliyunInfos;
            _ossTool = ossTool;
            _httpContextAccessor = httpContextAccessor;
            _compyUIService = compyUIService;
        }


        [HttpPost]
        public async Task<string> GenerationImage([FromBody] PromptRequestInfo requestInfo)
        {
            return string.Empty;
            string promptId = await _compyUIService.QwenWorkflow.GenerateImageAsync(requestInfo.Prompt, "", long.Parse(requestInfo.Seed), requestInfo.Width, requestInfo.Height);
            return promptId;
        }


        [HttpPost]
        public async Task<ActionResult<string?>> SplitImage([FromBody] RequestImage requestInfo)
        {
            byte[]? bytes = await HttpTool.GetBytes(requestInfo.Url);
            try
            {
                _ossTool.DeleteObject(_aliyunInfos.Oss.Bucket, requestInfo.ObjectName);
            }
            catch (Exception)
            {

            }
            if (bytes == null)
            {
                return StatusCode(404);
            }
            string tempFileName = Guid.NewGuid().ToString() + Path.GetExtension(requestInfo.Url);
            string tempFilePath = Path.Combine(Workflow.InputTempPath, tempFileName);
            System.IO.File.WriteAllBytes(tempFilePath, bytes);

            string promptId = await _compyUIService.SplitImageWorkflow.GenerateImageAsync(tempFileName);

            if (string.IsNullOrWhiteSpace(promptId))
            {
                if (System.IO.File.Exists(tempFilePath))
                {
                    System.IO.File.Delete(tempFilePath);
                }
            }
            else
            {
                _tmpInputPath[promptId] = tempFilePath;
            }

            return promptId;
        }


        private void DeleteTempInputFile(string promptId)
        {
            if (_tmpInputPath.TryRemove(promptId, out string? tempInputPath))
            {
                if (!string.IsNullOrWhiteSpace(tempInputPath))
                {
                    if (System.IO.File.Exists(tempInputPath))
                    {
                        System.IO.File.Delete(tempInputPath);
                    }
                }
            }
        }

        [HttpPost]
        public async Task<ActionResult<ImageRequestResult>?> GetGenerationResult([FromBody] RequestTaskPrompt requestInfo)
        {
            string result = string.Empty;
            try
            {
                result = await _compyUIService.GetGenerationResult(requestInfo.WorkflowType, requestInfo.PromptId);
            }
            catch (Exception)
            {
                DeleteTempInputFile(requestInfo.PromptId);
            }

            if (!string.IsNullOrWhiteSpace(result) && System.IO.File.Exists(result))
            {
                DeleteTempInputFile(requestInfo.PromptId);
                byte[] bytes = System.IO.File.ReadAllBytes(result);
                if (System.IO.File.Exists(result))
                {
                    System.IO.File.Delete(result);
                }
                string ext = Path.GetExtension(result);

                var objectName = $"cache/image/{Guid.NewGuid().ToString("N")}{ext}";
                using (var mem = new MemoryStream(bytes!))
                {
                    _ossTool.PutObject(_aliyunInfos.Oss.Bucket, objectName, mem, DateTime.UtcNow.AddDays(7));
                }

                string url = $"https://{_aliyunInfos.Oss.Bucket}.{_aliyunInfos.Oss.Region.Replace("-internal", "")}.aliyuncs.com/{objectName}";

                return new ImageRequestResult
                {
                    ImageUrl = url,
                    ObjectName = objectName,
                };
            }

            return null;
        }
    }
}