100 lines
2.5 KiB
Go
100 lines
2.5 KiB
Go
package zhipu
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"time"
|
|
|
|
"apigo.cc/ai"
|
|
"github.com/yankeguo/zhipu"
|
|
)
|
|
|
|
func MakeImage(aiConf *ai.AIConfig, conf ai.ImageConfig) (ai.ImageResult, error) {
|
|
c, err := getClient(aiConf)
|
|
if err != nil {
|
|
return ai.ImageResult{}, err
|
|
}
|
|
|
|
cc := c.ImageGeneration(conf.Model).SetPrompt(conf.SystemPrompt + conf.Prompt)
|
|
//cc.SetSize(config.GetSize())
|
|
|
|
t1 := time.Now().UnixMilli()
|
|
if r, err := cc.Do(context.Background()); err == nil {
|
|
t2 := time.Now().UnixMilli() - t1
|
|
results := make([]string, 0)
|
|
for _, item := range r.Data {
|
|
results = append(results, item.URL)
|
|
}
|
|
if len(results) == 0 {
|
|
results = append(results, "")
|
|
}
|
|
return ai.ImageResult{
|
|
Results: results,
|
|
UsedTime: t2,
|
|
}, nil
|
|
} else {
|
|
return ai.ImageResult{}, err
|
|
}
|
|
}
|
|
|
|
func MakeVideo(aiConf *ai.AIConfig, conf ai.VideoConfig) (string, error) {
|
|
c, err := getClient(aiConf)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
cc := c.VideoGeneration(conf.Model).SetPrompt(conf.SystemPrompt + conf.Prompt)
|
|
if len(conf.Ref) > 0 {
|
|
cc.SetImageURL(conf.Ref[0])
|
|
}
|
|
|
|
if resp, err := cc.Do(context.Background()); err == nil {
|
|
return resp.ID, nil
|
|
} else {
|
|
return "", err
|
|
}
|
|
}
|
|
|
|
func GetVideoResult(aiConf *ai.AIConfig, taskId string, waitSeconds int) (ai.VideoResult, error) {
|
|
c, err := getClient(aiConf)
|
|
if err != nil {
|
|
return ai.VideoResult{}, err
|
|
}
|
|
|
|
var r zhipu.AsyncResultResponse
|
|
for i := 0; i < waitSeconds; i += 3 {
|
|
r, err = c.AsyncResult(taskId).Do(context.Background())
|
|
if err != nil {
|
|
return ai.VideoResult{}, err
|
|
}
|
|
if r.TaskStatus == zhipu.VideoGenerationTaskStatusSuccess {
|
|
results := make([]string, 0)
|
|
previews := make([]string, 0)
|
|
for _, item := range r.VideoResult {
|
|
results = append(results, item.URL)
|
|
previews = append(previews, item.CoverImageURL)
|
|
}
|
|
if len(results) == 0 {
|
|
results = append(results, "")
|
|
previews = append(previews, "")
|
|
}
|
|
return ai.VideoResult{
|
|
Results: results,
|
|
Previews: previews,
|
|
UsedTime: 0,
|
|
}, nil
|
|
}
|
|
if r.TaskStatus == zhipu.VideoGenerationTaskStatusFail {
|
|
return ai.VideoResult{}, errors.New("fail on video task " + taskId)
|
|
} else if r.TaskStatus == zhipu.VideoGenerationTaskStatusProcessing {
|
|
if waitSeconds == 0 {
|
|
return ai.VideoResult{IsProcessing: true}, nil
|
|
}
|
|
time.Sleep(3 * time.Second)
|
|
} else {
|
|
return ai.VideoResult{}, errors.New("unknow status " + r.TaskStatus + " on video task " + taskId)
|
|
}
|
|
}
|
|
return ai.VideoResult{IsProcessing: true}, errors.New("timeout on video task " + taskId)
|
|
}
|