ai_old/llm/zhipu/gc.go

87 lines
2.4 KiB
Go
Raw Normal View History

2024-09-17 18:44:21 +08:00
package zhipu
import (
"apigo.cc/ai/ai/llm"
"apigo.cc/ai/ai/llm/zhipu/zhipu"
"context"
"errors"
"time"
)
func (lm *LLM) FastMakeImage(prompt string, config llm.GCConfig) ([]string, error) {
config.Model = ModelCogView3Plus
return lm.MakeImage(prompt, config)
}
func (lm *LLM) BestMakeImage(prompt string, config llm.GCConfig) ([]string, error) {
config.Model = ModelCogView3
return lm.MakeImage(prompt, config)
}
func (lm *LLM) MakeImage(prompt string, config llm.GCConfig) ([]string, error) {
c, err := zhipu.NewClient(zhipu.WithAPIKey(lm.config.ApiKey), zhipu.WithBaseURL(lm.config.Endpoint))
if err != nil {
return nil, err
}
2024-09-18 18:29:21 +08:00
config.SetDefault(&lm.config.GCConfig)
2024-09-17 18:44:21 +08:00
cc := c.ImageGeneration(config.Model).SetPrompt(prompt)
2024-09-18 18:29:21 +08:00
cc.SetSize(config.GetSize())
2024-09-17 18:44:21 +08:00
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 (lm *LLM) FastMakeVideo(prompt string, config llm.GCConfig) ([]string, []string, error) {
config.Model = ModelCogVideoX
return lm.MakeVideo(prompt, config)
}
func (lm *LLM) BestMakeVideo(prompt string, config llm.GCConfig) ([]string, []string, error) {
config.Model = ModelCogVideoX
return lm.MakeVideo(prompt, config)
}
func (lm *LLM) MakeVideo(prompt string, config llm.GCConfig) ([]string, []string, error) {
c, err := zhipu.NewClient(zhipu.WithAPIKey(lm.config.ApiKey), zhipu.WithBaseURL(lm.config.Endpoint))
if err != nil {
return nil, nil, err
}
2024-09-18 18:29:21 +08:00
config.SetDefault(&lm.config.GCConfig)
2024-09-17 18:44:21 +08:00
cc := c.VideoGeneration(config.Model).SetPrompt(prompt)
2024-09-18 18:29:21 +08:00
cc.SetImageURL(config.GetRef())
2024-09-17 18:44:21 +08:00
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
}
}