back

service/serviceLogic/diffusionAiService/diffusionAiPublicLogicBack.go

+package diffusionAiService
+
+import (
+	"aiapi/models/req"
+	"aiapi/repositories"
+	"bufio"
+	"encoding/json"
+	"github.com/sirupsen/logrus"
+	"os"
+	"strconv"
+	"time"
+)
+
+var DiffusionAiCha chan PubStruct
+
+type PubStruct struct {
+	ApiName   string `json:"apiName"`
+	InStruct  interface{}
+	requestId int64 `json:"requestId"`
+}
+
+type DiffusionRes struct {
+	RequestId int64 `structs:"request_id" json:"requestId"`
+}
+
+// 获取空闲的AI算力机器，如果无空闲返回""
+func queryComputePower() map[string]string {
+	dataReq := make(map[string]interface{})
+	computePower := make(map[string]string)
+	dataReq["status"] = int64(0)
+
+	var queryAiComputePower = new(repositories.TbAiComputePower)
+
+	dataRes := queryAiComputePower.Query(dataReq)
+	if len(dataRes) > 0 {
+		computePower["baseUrl"] = dataRes[0].BaseUrl
+		computePower["id"] = strconv.FormatInt(dataRes[0].Id, 10)
+	}
+	return computePower
+}
+
+// 更新算力状态
+func updateComputePower(updateMapReq map[string]int64) (count int64) {
+	var updateComputerPowerReq = new(repositories.TbAiComputePower)
+	updateComputerPowerReq.Id, _ = updateMapReq["id"]
+	updateComputerPowerReq.Status = updateMapReq["status"]
+	updateComputerPowerRes := updateComputerPowerReq.Update()
+
+	return updateComputerPowerRes
+}
+
+// 获取userAiRequest请求数据
+func queryUserAiRequest(requestId int64) (data repositories.TbUserAIRequest) {
+	var queryUserAIRequestReq = new(repositories.TbUserAIRequest)
+	queryUserAIRequestRes := queryUserAIRequestReq.GetById(strconv.FormatInt(requestId, 10))
+
+	return queryUserAIRequestRes
+}
+
+// 生成userAiRequest请求数据
+func addUserAiRequest(addUserAiReq req.UserAIRequestReq) (id int64) {
+
+	var addUserAIRequestReq = new(repositories.TbUserAIRequest)
+
+	addUserAIRequestReq.UserId = addUserAiReq.UserId
+	addUserAIRequestReq.Status = addUserAiReq.Status
+	addUserAIRequestRes, err := addUserAIRequestReq.Add()
+
+	logrus.Info("request data = %v\n", addUserAIRequestReq)
+	if err != nil {
+		logrus.Info("AddUserAIRequest error = %v\n", err)
+		return
+	}
+	id = addUserAIRequestRes.Id
+	return
+}
+
+// 更新userAiRequest
+func updateUserAiRequest(updateUserAiReq map[string]int64) (count int64) {
+	var updateUserAiRequestReq = new(repositories.TbAiComputePower)
+	updateUserAiRequestReq.Id, _ = updateUserAiReq["id"]
+	updateUserAiRequestReq.Status = updateUserAiReq["status"]
+	updateComputerPowerRes := updateUserAiRequestReq.Update()
+
+	return updateComputerPowerRes
+}
+
+// 管道
+func DiffusionAiChanel() {
+	for {
+		select {
+		case chanReq := <-DiffusionAiCha:
+			go func(gi PubStruct) {
+				if chanReq.ApiName == "genImage" {
+					byteData, err := json.Marshal(gi.InStruct)
+					if err != nil {
+						logrus.Info(err.Error())
+					}
+					var inStrc GenImageReq
+					err = json.Unmarshal(byteData, &inStrc)
+					if err != nil {
+						logrus.Info(err.Error())
+					}
+					GoGenImageLogic(inStrc, gi.requestId)
+				}
+
+			}(chanReq)
+		}
+	}
+}
+
+// base64上传至服务器
+func Base64Upload(imgStr string) (imgInfo map[string]string) {
+	timeNow := time.Now().Unix()
+	imgName := strconv.FormatInt(timeNow, 10) + ".png"
+	file, err := os.OpenFile("http://static.rejo9.com/stable-diffusion-webui/txt2img-images/"+imgName, os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		logrus.Info(err.Error())
+	}
+	w := bufio.NewWriter(file)
+	_, err3 := w.WriteString(string(imgStr))
+	if err3 != nil {
+		logrus.Info(err3.Error())
+	}
+	w.Flush()
+	defer file.Close()
+
+	imgInfo["Name"] = imgName
+	imgInfo["Path"] = "/stable-diffusion-webui/txt2img-images/"
+	imgInfo["url"] = "http://static.rejo9.com/stable-diffusion-webui/txt2img-images/" + imgName
+
+	return
+}

service/serviceLogic/diffusionAiService/genImageServiceLogicback.go

+package diffusionAiService
+
+import (
+	"aiapi/models/req"
+	"aiapi/repositories"
+	"aiapi/utils"
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"github.com/kataras/iris/v12"
+	"github.com/kataras/iris/v12/mvc"
+	"github.com/sirupsen/logrus"
+	"io"
+	"net/http"
+	"strconv"
+	"time"
+)
+
+type GenImageReq struct {
+	//用户Id
+	UserId int64 `structs:"user_id" json:"userId"`
+	// 提示词（必要）
+	Prompt string `structs:"prompt" json:"prompt"`
+	//状态
+	Status int64 `structs:"status" json:"status"`
+	// 种子，控制随机性 默认-1
+	seed int64 `structs:"seed" default:"-1"`
+	//一次生成几张图片 默认 1
+	batchSize int64 `structs:"batch_size"`
+	//提词器相关度 默认7
+	cfgScale int64 `structs:"cfg_scale"`
+	//负向提词器 默认 “”
+	negativePrompt string `structs:"negative_prompt"`
+	//覆盖默认参数的配置 修改使用的模型{"sd_model_checkpoint":  "768-v-ema.ckpt [a2a802b2]"}
+	overrideSettings struct {
+		sdModelCheckpoint string `structs:"sd_model_checkpoint"`
+	}
+	//图片宽度 512
+	width int64 `structs:"width"`
+	//图片高度 512
+	height int64 `structs:"height"`
+}
+
+// GenImageLogic
+// @Summary  根据文字生成图片，参数与python一致，返回为一个requestId
+// @Description
+// @Tags　diffusionAiServiceService
+// @Accept application/json
+// @Produce application/json
+// @Param prompt body GenImageReq true "json入参"
+// @Success 200
+// @Failure 500 "{"message":"fail","code":500,"data":null}"
+// @Failure 520 "{"message":"station exists","code":520,"data":null}"
+// @Router /diffusionAiServiceService/GenImage [post]
+func GenImageLogic(ctx iris.Context) (mr mvc.Response) {
+	logrus.Printf(ctx.String())
+	var genImage GenImageReq
+
+	body, _ := ctx.GetBody()
+	err := json.Unmarshal(body, &genImage)
+	if err != nil {
+		utils.FailWithMsg(&mr, err)
+		logrus.Error(err)
+		return
+	}
+	var addUserAiReq req.UserAIRequestReq
+	addUserAiReq.UserId = genImage.UserId
+	addUserAiReq.Status = genImage.Status
+	// add userAiRequest
+	var res DiffusionRes
+	res.RequestId = addUserAiRequest(addUserAiReq)
+	utils.SuccessWithMsg(&mr, "success", res)
+	var diffusionAiCha PubStruct
+	diffusionAiCha.ApiName = "genImage"
+	diffusionAiCha.InStruct = genImage
+	diffusionAiCha.requestId = res.RequestId
+	//go to chanel
+	DiffusionAiCha <- diffusionAiCha
+
+	return
+}
+
+func GoGenImageLogic(genImage GenImageReq, requestId int64) {
+
+	computePower := make(map[string]string)
+	var userAiReq = new(repositories.TbUserAIRequest)
+
+	userAiReq.Id = requestId
+	// 如果没有空闲的算力，进程一致处于等待中 获取等待时间
+	//estimateSec等待时间计时
+	estimateSecStart := time.Now()
+	for {
+		computePower = queryComputePower()
+		if len(computePower) > 0 {
+			break
+		}
+	}
+	//获取用户请求最新状态，如果状态为已取消，则不调用python服务
+
+	//更新算力状态 update
+
+	estimateSec := time.Since(estimateSecStart)
+	userAiReq.ProcessPowerId, _ = strconv.ParseInt(computePower["id"], 10, 64)
+	userAiReq.EstimateSec = int64(estimateSec)
+	count := userAiReq.Update()
+	fmt.Println(count)
+	fmt.Println(estimateSec)
+	//根据客户是否取消，来判定是否是要调用python服务
+	//更新算力机器以及执行状态
+	//updateUserAiRequest
+
+	var url = "http://" + computePower["baseUrl"] + "/sdapi/v1/txt2img"
+	genImageReq := make(map[string]string)
+	genImageReq["prompt"] = genImage.Prompt
+	bytesData, _ := json.Marshal(genImageReq)
+	//获取执行时间
+	//执行时间计时processSec
+	processSecStart := time.Now()
+	res, err := http.Post(url, "application/json", bytes.NewBuffer(bytesData))
+	processSec := time.Since(processSecStart)
+	updateUserAiReq.ProcessSec = int64(processSec)
+	count = updateUserAiReq.Update()
+	fmt.Println(count)
+	fmt.Println(processSec)
+	if err != nil {
+		fmt.Println("Fatal error ", err.Error())
+	}
+	//更新执行接口
+	//updateUserAiRequest
+	//更新算力状态 update
+
+	defer res.Body.Close()
+
+	bodyRes, _ := io.ReadAll(res.Body)
+	fmt.Println(bodyRes)
+
+	return
+
+}