DrawAI.go

package main

import (
	"WorldEpcho/src/utils"
	"bytes"
	"encoding/json"
	"fmt"
	"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 Response 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"`
}

func AiDraw(OBJname string, id int64) {
	// 指定图片保存路径和文件名
	saveDir := "./static/Worlds/items/"
	filePath := saveDir + OBJname + ".png"

	// 检查文件是否已经存在
	if _, err := os.Stat(filePath); err == nil {
		// 如果文件存在，直接返回文件路径
		fmt.Println("文件已存在, 路径为: ", filePath)
		return
	}

	api := "http://52.83.116.11:13679/Flow"
	// 这里的Rcd5和Rc的值未给出，需要根据实际情况定义或获取

	rc := utils.Strval(id) + "_" + utils.Strval(time.Now().Unix())
	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.Fatalf("Error occurred during unmarshaling. Error: %s", err.Error())
		return
	}

	// 使用 json.Marshal 将 map 转换回 JSON 字符串，以发送 HTTP 请求
	requestBytes, err := json.Marshal(requestData)
	if err != nil {
		log.Fatalf("Error occurred during marshaling. Error: %s", 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.Fatalf("Error occurred during POST request. Error: %s", err.Error())
		return
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatalf("Error occurred during reading response body. Error: %s", err.Error())
		return
	}

	var response Response
	err = json.Unmarshal(body, &response)
	if err != nil {
		log.Fatalf("Error occurred during unmarshaling response body. Error: %s", err.Error())
		return
	}

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

	if resp.StatusCode != 200 || response.Errors != nil {
		log.Fatalf("Server returned error. Status: %d, Errors: %v", resp.StatusCode, response.Errors)
		return
	}

	// 处理WebSocket连接和消息
	HandleWebSocket(api, Rcd5, OBJname)
}

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

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

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

	// 确保我们至少得到了主机名
	if len(hostParts) < 1 {
		log.Fatal("API 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)
	}
	defer c.Close()

	done := make(chan bool)

	go func() {
		defer close(done)
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				log.Printf("read error: %v", err)
				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)
				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)
				fmt.Println("imageURL ==> : ", imageURL)
				if err := DownloadAndSaveImage(imageURL, OBJname); err != nil {
					log.Printf("error downloading or saving image: %v", err)
				}
				done <- true
				return
			}
		}
	}()

	// 等待消息处理完成
	<-done
}

// 下载并保存图片到指定路径
func DownloadAndSaveImage(imageURL, OBJname 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 + OBJname + ".png"
	fmt.Println("filePath: ==> ", filePath)
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

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

func main() {
	AiDraw("notes", 1)
}