huoshan/gc.go

package huoshan

import (
	"errors"
	"strings"
	"time"

	"apigo.cc/ai"
	"github.com/ssgo/u"
	"github.com/volcengine/volc-sdk-golang/service/visual"
	"github.com/volcengine/volc-sdk-golang/service/visual/model"
)

func getVisualClient(aiConf *ai.AIConfig) *visual.Visual {
	ak, sk := getAKSK(aiConf)
	vis := visual.NewInstance()
	vis.Client.SetAccessKey(ak)
	vis.Client.SetSecretKey(sk)
	if aiConf.Extra["region"] != nil {
		vis.SetRegion(u.String(aiConf.Extra["region"]))
	}
	if aiConf.Extra["host"] != nil {
		vis.SetHost(u.String(aiConf.Extra["host"]))
	}
	return vis
}

func MakeImage(aiConf *ai.AIConfig, conf ai.ImageConfig) (ai.ImageResult, error) {
	modelA := strings.SplitN(conf.Model, ":", 2)
	data := map[string]any{
		"req_key":    modelA[0],
		"prompt":     conf.SystemPrompt + conf.Prompt,
		"return_url": true,
	}
	if len(modelA) > 1 {
		data["model_version"] = modelA[1]
	}
	if conf.NegativePrompt != "" {
		data["negative_prompt"] = conf.NegativePrompt
	}
	if conf.Width > 0 {
		data["width"] = conf.Width
	}
	if conf.Height > 0 {
		data["height"] = conf.Height
	}
	if conf.Scale > 0 {
		if conf.Scale < 1 {
			// 取值 0-1 放大到 1-30
			data["scale"] = conf.Scale * 30
		} else {
			data["scale"] = conf.Scale
		}
	}
	if conf.Steps > 0 {
		data["ddim_steps"] = conf.Steps
	}
	if len(conf.Ref) > 0 {
		// 如果有参考图，自动切换到图生图模型
		if strings.Contains(modelA[0], "t2i") {
			data["req_key"] = strings.ReplaceAll(modelA[0], "t2i", "i2i")
		}
		// 根据参考图类型设置(url和base64只能2选1)
		image_url := make([]string, 0)
		binary_data_base64 := make([]string, 0)
		for _, ref := range conf.Ref {
			if strings.Contains(ref, "://") {
				image_url = append(image_url, ref)
			} else {
				binary_data_base64 = append(binary_data_base64, ref)
			}
		}
		if len(image_url) > 0 {
			data["image_urls"] = conf.Ref
		} else {
			data["binary_data_base64"] = binary_data_base64
		}

		// 参考图权重设置
		if conf.Cref > 0 || conf.Sref > 0 {
			if strings.Contains(conf.Model, "ip_keep") {
				// 人脸保持模型
				if conf.Cref > 0 {
					data["ref_id_weight"] = conf.Cref
				}
				if conf.Sref > 0 {
					data["ref_ip_weight"] = conf.Sref
				}
			} else {
				// 图生图模型
				style_reference_args := map[string]any{"binary_data_index": 0}
				if conf.Cref > 0 {
					style_reference_args["id_weight"] = conf.Cref
				}
				if conf.Sref > 0 {
					style_reference_args["style_weight"] = conf.Sref
				}
				data["style_reference_args"] = style_reference_args
			}
		}
	}

	// 其他参数
	for k, v := range conf.Extra {
		data[k] = v
	}

	// fmt.Println(u.BMagenta(u.JsonP(data)), 111)
	t1 := time.Now().UnixMilli()
	c := getVisualClient(aiConf)
	respMap, status, err := c.CVProcess(data)
	// fmt.Println(u.BCyan(u.JsonP(respMap)), 222)
	resp := &model.VisualPubResult{}
	u.Convert(respMap, resp)
	t2 := time.Now().UnixMilli() - t1

	if err != nil {
		return ai.ImageResult{}, err
	}
	if status != 200 {
		return ai.ImageResult{}, errors.New(u.String(resp.Message))
	}
	return ai.ImageResult{
		Results:  resp.Data.ImageUrls,
		UsedTime: t2,
	}, nil
}