feat: dalle

This commit is contained in:
Jack Merrill 2023-08-18 23:01:43 -05:00
parent ab852f001d
commit 62706e8309
No known key found for this signature in database
GPG Key ID: B8E3CDF57DD80CA5
2 changed files with 82 additions and 1 deletions

View File

@ -0,0 +1,77 @@
package fun
import (
"context"
"fmt"
"os"
"strings"
"github.com/jackmerrill/hampbot/internal/utils/config"
"github.com/jackmerrill/hampbot/internal/utils/embed"
"github.com/sashabaranov/go-openai"
"github.com/zekroTJA/shireikan"
)
type Dalle struct {
}
func (c *Dalle) GetInvokes() []string {
return []string{"dalle", "ai-art"}
}
func (c *Dalle) GetDescription() string {
return "Generate AI art with DALL-E"
}
func (c *Dalle) GetHelp() string {
return "`dalle [prompt]` - Generate AI art with DALL-E"
}
func (c *Dalle) GetGroup() string {
return config.GroupFun
}
func (c *Dalle) GetDomainName() string {
return "hamp.fun.dalle"
}
func (c *Dalle) GetSubPermissionRules() []shireikan.SubPermission {
return nil
}
func (c *Dalle) IsExecutableInDMChannels() bool {
return true
}
func (c *Dalle) Exec(ctx shireikan.Context) error {
openaiToken := os.Getenv("OPENAI_TOKEN")
if openaiToken == "" {
ctx.ReplyEmbed(embed.NewErrorEmbed(ctx).SetDescription("No OpenAI token set.").MessageEmbed)
return fmt.Errorf("no openai token set")
}
prompt := strings.TrimPrefix(ctx.GetMessage().Content, fmt.Sprintf("%sdalle ", config.BotPrefix))
client := openai.NewClient(openaiToken)
ctx.GetSession().ChannelTyping(ctx.GetChannel().ID)
resp, err := client.CreateImage(context.Background(), openai.ImageRequest{
Prompt: prompt,
Size: openai.CreateImageSize1024x1024,
ResponseFormat: openai.CreateImageResponseFormatURL,
N: 1,
})
if err != nil {
ctx.ReplyEmbed(embed.NewErrorEmbed(ctx).SetDescription("Failed to generate image.").MessageEmbed)
return err
}
ctx.ReplyEmbed(embed.NewSuccessEmbed(ctx).
SetTitle("DALL-E").
SetDescription(prompt).
SetImage(resp.Data[0].URL).
MessageEmbed)
return nil
}

View File

@ -11,6 +11,7 @@ import (
"github.com/jackmerrill/hampbot/internal/commands/fun"
studentlife "github.com/jackmerrill/hampbot/internal/commands/studentlife"
util "github.com/jackmerrill/hampbot/internal/commands/util"
"github.com/jackmerrill/hampbot/internal/utils/config"
"github.com/zekroTJA/shireikan"
"github.com/charmbracelet/log"
@ -38,7 +39,7 @@ func main() {
}()
handler := shireikan.New(&shireikan.Config{
GeneralPrefix: ">",
GeneralPrefix: config.BotPrefix,
AllowBots: false,
AllowDM: true,
ExecuteOnEdit: true,
@ -69,6 +70,9 @@ func main() {
handler.Register(&fun.XKCD{})
log.Debug("Registered xkcd command")
handler.Register(&fun.Dalle{})
log.Debug("Registered dalle command")
log.Info("Registered all commands")
handler.Setup(session)