84 lines
2.3 KiB
Go
84 lines
2.3 KiB
Go
package zhipu
|
|
|
|
import (
|
|
"apigo.cc/ai/zhipu/zhipu"
|
|
"context"
|
|
"errors"
|
|
"time"
|
|
)
|
|
|
|
func (ag *Agent) FastMakeImage(prompt, size, refImage string) ([]string, error) {
|
|
return ag.MakeImage(ModelCogView3Plus, prompt, size, refImage)
|
|
}
|
|
|
|
func (ag *Agent) BestMakeImage(prompt, size, refImage string) ([]string, error) {
|
|
return ag.MakeImage(ModelCogView3, prompt, size, refImage)
|
|
}
|
|
|
|
func (ag *Agent) MakeImage(model, prompt, size, refImage string) ([]string, error) {
|
|
c, err := zhipu.NewClient(zhipu.WithAPIKey(ag.config.ApiKey), zhipu.WithBaseURL(ag.config.Endpoint))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
cc := c.ImageGeneration(model).SetPrompt(prompt)
|
|
if size != "" {
|
|
cc.SetSize(size)
|
|
}
|
|
|
|
if r, err := cc.Do(context.Background()); err == nil {
|
|
results := make([]string, 0)
|
|
for _, item := range r.Data {
|
|
results = append(results, item.URL)
|
|
}
|
|
return results, nil
|
|
} else {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
func (ag *Agent) FastMakeVideo(prompt, size, refImage string) ([]string, []string, error) {
|
|
return ag.MakeVideo(ModelCogVideoX, prompt, size, refImage)
|
|
}
|
|
|
|
func (ag *Agent) BestMakeVideo(prompt, size, refImage string) ([]string, []string, error) {
|
|
return ag.MakeVideo(ModelCogVideoX, prompt, size, refImage)
|
|
}
|
|
|
|
func (ag *Agent) MakeVideo(model, prompt, size, refImage string) ([]string, []string, error) {
|
|
c, err := zhipu.NewClient(zhipu.WithAPIKey(ag.config.ApiKey), zhipu.WithBaseURL(ag.config.Endpoint))
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
|
|
cc := c.VideoGeneration(model).SetPrompt(prompt)
|
|
if refImage != "" {
|
|
cc.SetImageURL(refImage)
|
|
}
|
|
|
|
if resp, err := cc.Do(context.Background()); err == nil {
|
|
for i := 0; i < 1200; i++ {
|
|
r, err := c.AsyncResult(resp.ID).Do(context.Background())
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
if r.TaskStatus == zhipu.VideoGenerationTaskStatusSuccess {
|
|
covers := make([]string, 0)
|
|
results := make([]string, 0)
|
|
for _, item := range r.VideoResult {
|
|
results = append(results, item.URL)
|
|
covers = append(covers, item.CoverImageURL)
|
|
}
|
|
return results, covers, nil
|
|
}
|
|
if r.TaskStatus == zhipu.VideoGenerationTaskStatusFail {
|
|
return nil, nil, errors.New("fail on task " + resp.ID)
|
|
}
|
|
time.Sleep(3 * time.Second)
|
|
}
|
|
return nil, nil, errors.New("timeout on task " + resp.ID)
|
|
} else {
|
|
return nil, nil, err
|
|
}
|
|
}
|