AiDraw.go

package service

import (
	"WorldEpcho/src/config"
	"WorldEpcho/src/utils"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

type PromptRequest struct {
	ClientID string                 `json:"client_id"`
	Prompt   map[string]interface{} `json:"prompt"`
}

type DrawAIResponse struct {
	Status     int         `json:"status"`
	PromptID   string      `json:"prompt_id"`
	NodeErrors interface{} `json:"node_errors"`
	Errors     interface{} `json:"errors"`
	Output     struct {
		Images []struct {
			Filename string `json:"filename"`
		} `json:"images"`
	} `json:"data"`
}

// PropTranslateRequest 定义了发送给翻译接口的请求体结构
type PropTranslateRequest struct {
	ID    string `json:"id"`
	Name  string `json:"name"`
	Upset string `json:"upset"`
}

// PropTranslateResponse 定义了翻译接口，返回数据的结构
type PropTranslateResponse struct {
	Status  string `json:"status"`
	Content string `json:"content"`
}

// TranslateProp 发送翻译请求到服务器并处理返回的数据
func TranslateProp(name string) (string, error) {

	// 设置请求体中的固定值
	var requestBody PropTranslateRequest
	requestBody.ID = "96"
	requestBody.Upset = "mipm18vvi"
	requestBody.Name = name + "φ"

	// 转换请求体为JSON
	requestBodyBytes, err := json.Marshal(requestBody)
	if err != nil {
		//c.JSON(http.StatusBadRequest, gin.H{"code": 0, "message": "转换请求体为JSON出错"})
		return "", errors.New("转换请求体为JSON出错")
	}

	// 从配置文件获取请求路径
	url := config.Conf.TranslatePropUrl

	// 发送POST请求
	response, err := http.Post(url, "application/json", bytes.NewBuffer(requestBodyBytes))
	if err != nil {
		//c.JSON(http.StatusInternalServerError, gin.H{"code": 0, "message": "发送请求出错"})
		return "", errors.New("道具名称翻译发送请求出错")
	}
	defer response.Body.Close()

	// 读取返回的body内容
	responseBodyBytes, err := ioutil.ReadAll(response.Body)
	if err != nil {
		//c.JSON(http.StatusInternalServerError, gin.H{"code": 0, "message": "读取返回的body内容出错"})
		return "", errors.New("道具名称翻译发送请求出错")
	}

	// 解析返回的JSON数据
	var translateResponse PropTranslateResponse
	if err := json.Unmarshal(responseBodyBytes, &translateResponse); err != nil {
		//c.JSON(http.StatusInternalServerError, gin.H{"code": 0, "message": "解析返回的JSON数据出错"})
		return "", errors.New("解析返回的JSON数据出错")
	}

	// 记录翻译后的结果
	log.Println(config.ColorPurple, "翻译结果: ", translateResponse.Content, config.ColorReset)

	// 根据翻译结果的状态，向客户端发送不同的响应
	if translateResponse.Status == "OK" {
		//c.JSON(http.StatusOK, gin.H{"code": 1, "message": "翻译成功", "translateResult": translateResponse})
		return translateResponse.Content, nil
	} else {
		//c.JSON(http.StatusInternalServerError, gin.H{"code": 0, "message": "翻译服务返回错误"})
		return "", errors.New("翻译服务返回错误")
	}
}

// AI 绘画接口
func AiDraw(c *gin.Context) {
	/*
	   判断用户是否登录
	*/
	id, isLogin := IsUserLoggedIn(c)
	if !isLogin {
		log.Println("用户未登录")
		c.JSON(http.StatusOK, gin.H{"code": 0, "message": "用户未登录"})
		return
	}

	imagName := c.PostForm("name")
	// 确保目标文件夹存在
	//saveDir := "./static/Worlds/items/"
	saveDir := "./static/Worlds/items/"
	filePath := saveDir + imagName + ".png"
	aiUrl := config.Conf.DrawAiUrl

	imagUrl := "/worlds/items/" + imagName + ".png"

	// 检查文件是否已经存在
	if _, err := os.Stat(filePath); err == nil {
		// 如果文件存在，直接返回文件路径
		fmt.Println("文件已存在, 路径为: ", filePath)
		c.JSON(http.StatusOK, gin.H{"code": 1, "message": "道具图片已经存在", "propImagePath": imagUrl})
		return
	}

	propName, err := TranslateProp(imagName)
	if err != nil {
		fmt.Println(config.ColorYellow, "翻译道具名称出错：", err.Error(), config.ColorReset)
		c.JSON(http.StatusOK, gin.H{"code": 0, "message": "翻译道具名称出错"})
		return
	}
	OBJname := propName

	//api := "http://52.83.116.11:13679/Flow"
	api := aiUrl + "/Flow"

	// 这里的Rcd5和Rc的值未给出，需要根据实际情况定义或获取

	rc := utils.Strval(id) + "_" + utils.Strval(time.Now().UnixNano())
	fmt.Println("rc: =====> ", rc)
	Rcd5 := utils.MD5(rc)

	// 使用 fmt.Sprintf 构建 JSON 字符串，并插入 ClientID (Rcd5) 和 OBJname 变量
	jsonString := fmt.Sprintf(`{
		"client_id": "%s",
		"prompt": {
			"3": {
				"inputs": {
					"seed": 224721986114697,
					"steps": 20,
					"cfg": 7,
					"sampler_name": "euler_ancestral",
					"scheduler": "normal",
					"denoise": 1,
					"model": ["4", 0],
					"positive": ["6", 0],
					"negative": ["7", 0],
					"latent_image": ["5", 0]
				},
				"class_type": "KSampler"
			},
			"4": {
				"inputs": {
					"ckpt_name": "models/Stable-diffusion/RealPhoto.safetensors"
				},
				"class_type": "CheckpointLoaderSimple"
			},
			"5": {
				"inputs": {
					"width": 512,
					"height": 512,
					"batch_size": 1
				},
				"class_type": "EmptyLatentImage"
			},
			"6": {
				"inputs": {
					"text": "8k,realistic,game icon,(3d), black background,(%s:1.2)",
					"clip": ["4",1]
				},
				"class_type": "CLIPTextEncode"
			},
			"7": {
				"inputs": {
					"text": "paintings, sketches,nsfw, (worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), skin spots, acnes, skin blemishes, age spot, manboobs, backlight, glasses, panty, anime, cartoon, drawing, illustration, boring,long neck, out of frame, extra fingers, mutated hands, monochrome, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, glitchy, bokeh, (((long neck))), ((flat chested)), red eyes, extra heads, close up, text ,watermarks, logo,Six fingers,fingers.",
					"clip": ["4",1]
				},
				"class_type": "CLIPTextEncode"
			},
			"8": {
				"inputs": {
					"samples": ["3",0],
					"vae": ["4",2]
				},
				"class_type": "VAEDecode"
			},
			"9": {
				"inputs": {
					"filename_prefix": "ComfyUI",
					"images": ["8",0]
				},
				"class_type": "SaveImage"
			}
		}
	}`, Rcd5, OBJname)

	// 解析 JSON 字符串为 map[string]interface{}
	var requestData map[string]interface{}
	err = json.Unmarshal([]byte(jsonString), &requestData)
	if err != nil {
		log.Printf("Error occurred during unmarshaling. Error: %s", err.Error())
		c.JSON(http.StatusOK, gin.H{
			"code":    0,
			"message": "解析请求体为JSON字符串出错",
		})
		return
	}

	// 使用 json.Marshal 将 map 转换回 JSON 字符串，以发送 HTTP 请求
	requestBytes, err := json.Marshal(requestData)
	if err != nil {
		log.Printf("Error occurred during marshaling. Error: %s", err.Error())
		c.JSON(http.StatusOK, gin.H{"code": 0, "message": err.Error()})
		return
	}
	// 发送 HTTP 请求
	url := api + "/prompt"
	resp, err := http.Post(url, "application/json", bytes.NewBuffer(requestBytes))
	fmt.Println("发送请求url: ", url, "clientID: ", Rcd5)
	if err != nil {
		log.Printf("Error occurred during POST request. Error: %s", err.Error())
		c.JSON(http.StatusOK, gin.H{"code": 0, "message": err.Error()})
		return
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Error occurred during reading response body. Error: %s", err.Error())
		c.JSON(http.StatusOK, gin.H{"code": 0, "message": err.Error()})
		return
	}

	var response DrawAIResponse
	err = json.Unmarshal(body, &response)
	if err != nil {
		log.Printf("Error occurred during unmarshaling response body. Error: %s", err.Error())
		c.JSON(http.StatusOK, gin.H{"code": 0, "message": err.Error()})
		return
	}

	fmt.Println("response: ==> ", string(body))

	if resp.StatusCode != 200 || response.Errors != nil {
		log.Printf("Server returned error. Status: %d, Errors: %v", resp.StatusCode, response.Errors)
		c.JSON(http.StatusOK, gin.H{"code": 0, "message": err.Error()})
		return
	}

	// 处理WebSocket连接和消息
	url, err = HandleWebSocket(api, Rcd5, imagName, saveDir)
	if err != nil {
		fmt.Println("websocket 连接生成图片出错")
		c.JSON(http.StatusOK, gin.H{"code": 0, "message": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"code": 1, "message": "道具图片绘制完成", "propImagePath": imagUrl})
}

func HandleWebSocket(api, clientID, imagName, savePath string) (string, error) {
	// 使用 "//" 分割api字符串
	parts := strings.Split(api, "//")

	// 检查分割结果的长度，确保它至少有2个元素
	if len(parts) < 2 {
		log.Fatal("API URL格式不正确")
		return "", errors.New("URL格式不正确")
	}

	// 提取主机和路径部分
	hostPath := parts[1]
	hostParts := strings.SplitN(hostPath, "/", 2)

	// 确保我们至少得到了主机名
	if len(hostParts) < 1 {
		log.Fatal("API URL缺少主机名")
		return "", errors.New("URL缺少主机名")
	}

	// 提取主机名（可能包含端口）
	host := hostParts[0]

	// 确定完整的WebSocket路径
	wsPath := "/ws"
	if len(hostParts) > 1 {
		// 如果原始URL包含路径，则附加到WebSocket路径
		wsPath = "/" + hostParts[1] + wsPath
	}

	u := url.URL{Scheme: "ws", Host: host, Path: wsPath, RawQuery: "clientId=" + clientID}
	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Fatal("dial:", err)
		return "", errors.New("连接websocket出错")
	}
	defer c.Close()

	done := make(chan string)

	go func() {
		defer close(done)
		var PropimagPath string
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				log.Printf("read error: %v", err)
				PropimagPath = "error"
				done <- PropimagPath
				return
			}
			log.Printf("recv: %s", message)

			var response struct {
				Type   string `json:"type"`
				Status string `json:"status"`
				Data   struct {
					Output struct {
						Images []struct {
							Filename string `json:"filename"`
						} `json:"images"`
					} `json:"output"`
				} `json:"data"`
			}

			if err := json.Unmarshal(message, &response); err != nil {
				log.Printf("error unmarshaling message: %v", err)
				PropimagPath = "error"
				done <- PropimagPath
				continue
			}

			// 检查条件是否满足以保存图片
			if response.Type == "executed" && len(response.Data.Output.Images) > 0 {
				filename := response.Data.Output.Images[0].Filename
				log.Printf("Image filename: %s", filename)

				// 下载并保存图片
				imageURL := fmt.Sprintf("%s/view?filename=%s&type=output", api, filename)

				PropimagPath, err := DownloadAndSaveImage(savePath, imageURL, imagName)
				if err != nil {
					log.Printf("error downloading or saving image: %v", err)
					PropimagPath = "error"

				}
				done <- PropimagPath

				return
			}
		}
	}()

	// 等待消息处理完成
	for {
		PropUrl := <-done
		if PropUrl == "error" {
			return "", errors.New("error downloading or saving image")
		} else if PropUrl != "" {
			return PropUrl, nil
		}

	}

}

// 下载并保存图片到指定路径
func DownloadAndSaveImage(saveDir, imageURL, imagName string) (string, error) {
	resp, err := http.Get(imageURL)
	fmt.Println("imageURL get resp.StatusCode :", resp.StatusCode)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("failed to download image: status code %d", resp.StatusCode)
	}

	// 确保目标文件夹存在
	//saveDir := "./static/Worlds/items/"
	if err := os.MkdirAll(saveDir, 0755); err != nil {
		return "", err
	}

	// 创建文件
	filePath := saveDir + imagName + ".png"
	fmt.Println("filePath: ==> ", filePath)
	file, err := os.Create(filePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	// 保存图片内容到文件
	_, err = io.Copy(file, resp.Body)
	if err != nil {
		return "", err
	}
	return filePath, nil
}