95 lines
2.7 KiB
Go
95 lines
2.7 KiB
Go
|
package zhipu
|
||
|
|
|
||
|
|
import (
|
||
|
|
"apigo.cc/ai/llm/llm"
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"github.com/yankeguo/zhipu"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func (lm *LLM) FastMakeImage(prompt string, config llm.GCConfig) ([]string, llm.Usage, error) {
|
||
|
|
config.Model = ModelCogView3Plus
|
||
|
|
return lm.MakeImage(prompt, config)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (lm *LLM) BestMakeImage(prompt string, config llm.GCConfig) ([]string, llm.Usage, error) {
|
||
|
|
config.Model = ModelCogView3
|
||
|
|
return lm.MakeImage(prompt, config)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (lm *LLM) MakeImage(prompt string, config llm.GCConfig) ([]string, llm.Usage, error) {
|
||
|
|
c, err := zhipu.NewClient(zhipu.WithAPIKey(lm.config.ApiKey), zhipu.WithBaseURL(lm.config.Endpoint))
|
||
|
|
if err != nil {
|
||
|
|
return nil, llm.Usage{}, err
|
||
|
|
}
|
||
|
|
|
||
|
|
config.SetDefault(&lm.config.GCConfig)
|
||
|
|
cc := c.ImageGeneration(config.Model).SetPrompt(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)
|
||
|
|
}
|
||
|
|
return results, llm.Usage{
|
||
|
|
UsedTime: t2,
|
||
|
|
}, nil
|
||
|
|
} else {
|
||
|
|
return nil, llm.Usage{}, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (lm *LLM) FastMakeVideo(prompt string, config llm.GCConfig) ([]string, []string, llm.Usage, error) {
|
||
|
|
config.Model = ModelCogVideoX
|
||
|
|
return lm.MakeVideo(prompt, config)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (lm *LLM) BestMakeVideo(prompt string, config llm.GCConfig) ([]string, []string, llm.Usage, error) {
|
||
|
|
config.Model = ModelCogVideoX
|
||
|
|
return lm.MakeVideo(prompt, config)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (lm *LLM) MakeVideo(prompt string, config llm.GCConfig) ([]string, []string, llm.Usage, error) {
|
||
|
|
c, err := zhipu.NewClient(zhipu.WithAPIKey(lm.config.ApiKey), zhipu.WithBaseURL(lm.config.Endpoint))
|
||
|
|
if err != nil {
|
||
|
|
return nil, nil, llm.Usage{}, err
|
||
|
|
}
|
||
|
|
|
||
|
|
config.SetDefault(&lm.config.GCConfig)
|
||
|
|
cc := c.VideoGeneration(config.Model).SetPrompt(prompt)
|
||
|
|
cc.SetImageURL(config.GetRef())
|
||
|
|
|
||
|
|
t1 := time.Now().UnixMilli()
|
||
|
|
if resp, err := cc.Do(context.Background()); err == nil {
|
||
|
|
t2 := time.Now().UnixMilli() - t1
|
||
|
|
for i := 0; i < 1200; i++ {
|
||
|
|
r, err := c.AsyncResult(resp.ID).Do(context.Background())
|
||
|
|
if err != nil {
|
||
|
|
return nil, nil, llm.Usage{}, 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, llm.Usage{
|
||
|
|
UsedTime: t2,
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
if r.TaskStatus == zhipu.VideoGenerationTaskStatusFail {
|
||
|
|
return nil, nil, llm.Usage{}, errors.New("fail on task " + resp.ID)
|
||
|
|
}
|
||
|
|
time.Sleep(3 * time.Second)
|
||
|
|
}
|
||
|
|
return nil, nil, llm.Usage{}, errors.New("timeout on task " + resp.ID)
|
||
|
|
} else {
|
||
|
|
return nil, nil, llm.Usage{}, err
|
||
|
|
}
|
||
|
|
}
|