diff --git a/internal/telegram/generate.go b/internal/telegram/generate.go index 6fce3f5..2bf3aa5 100644 --- a/internal/telegram/generate.go +++ b/internal/telegram/generate.go @@ -13,7 +13,6 @@ import ( "github.com/LightningTipBot/LightningTipBot/internal/dalle" "github.com/LightningTipBot/LightningTipBot/internal/lnbits" "github.com/LightningTipBot/LightningTipBot/internal/runtime" - "github.com/LightningTipBot/LightningTipBot/internal/runtime/mutex" "github.com/LightningTipBot/LightningTipBot/internal/telegram/intercept" log "github.com/sirupsen/logrus" "github.com/skip2/go-qrcode" @@ -98,6 +97,21 @@ func (bot *TipBot) confirmGenerateImages(ctx intercept.Context) (intercept.Conte return ctx, nil } +var jobChan chan func(workerId int) +var workers = 2 + +func init() { + jobChan = make(chan func(workerId int), workers) + for i := 0; i < workers; i++ { + go worker(jobChan, i) + } +} +func worker(linkChan chan func(workerId int), workerId int) { + for generatePrompt := range linkChan { + generatePrompt(workerId) + } +} + // generateDalleImages is called by the invoice event when the user has paid func (bot *TipBot) generateDalleImages(event Event) { invoiceEvent := event.(*InvoiceEvent) @@ -107,74 +121,71 @@ func (bot *TipBot) generateDalleImages(event Event) { return } bot.trySendMessage(user.Telegram, "🔄 Your images are being generated. Please wait a few moments.") - - // we can have only one user using dalle - mutex.Lock("dalle-image-task") - defer mutex.Unlock("dalle-image-task") - time.Sleep(time.Second * 1) - - // create the client with the bearer token api key - dalleClient, err := dalle.NewHTTPClient(internal.Configuration.Generate.DalleKey) - // handle err - if err != nil { - log.Errorf("[NewHTTPClient] %v", err.Error()) - bot.dalleRefundUser(user, "") - return - } - - ctx, cancel := context.WithTimeout(context.Background(), time.Minute*10) - defer cancel() - // generate a task to create an image with a prompt - task, err := dalleClient.Generate(ctx, invoiceEvent.CallbackData) - if err != nil { - log.Errorf("[Generate] %v", err.Error()) - bot.dalleRefundUser(user, "") - return - } - // poll the task.ID until status is succeeded - var t *dalle.Task - timeout := time.After(5 * time.Minute) - ticker := time.Tick(5 * time.Second) - // Keep trying until we're timed out or get a result/error - for { - select { - case <-ctx.Done(): - bot.dalleRefundUser(user, "") - log.Errorf("[DALLE] ctx done") + var job = func(workerId int) { + // create the client with the bearer token api key + dalleClient, err := dalle.NewHTTPClient(internal.Configuration.Generate.DalleKey) + // handle err + if err != nil { + log.Errorf("[NewHTTPClient-%d] %v", workerId, err.Error()) + bot.dalleRefundUser(user,"") return - // Got a timeout! fail with a timeout error - case <-timeout: - bot.dalleRefundUser(user, "Timeout. Please try again later.") - log.Errorf("[DALLE] timeout") + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute*10) + defer cancel() + // generate a task to create an image with a prompt + task, err := dalleClient.Generate(ctx, invoiceEvent.CallbackData) + if err != nil { + log.Errorf("[Generate-%d] %v", workerId, err.Error()) + bot.dalleRefundUser(user,"") return - // Got a tick, we should check on checkSomething() - case <-ticker: - t, err = dalleClient.GetTask(ctx, task.ID) - // handle err - if err != nil { - log.Errorf("[GetTask] %v", err.Error()) - bot.dalleRefundUser(user, "") + } + // poll the task.ID until status is succeeded + var t *dalle.Task + timeout := time.After(5 * time.Minute) + ticker := time.Tick(5 * time.Second) + // Keep trying until we're timed out or get a result/error + for { + select { + case <-ctx.Done(): + bot.dalleRefundUser(user,"") + log.Errorf("[DALLE-%d] ctx done", workerId) return - } - if t.Status == dalle.StatusSucceeded { - log.Infof("[DALLE] task succeeded for user %s", GetUserStr(user.Telegram)) - // download the first generated image - for _, data := range t.Generations.Data { - err = bot.downloadAndSendImages(ctx, dalleClient, data, invoiceEvent) - if err != nil { - log.Errorf("[downloadAndSendImages] %v", err.Error()) - } + // Got a timeout! fail with a timeout error + case <-timeout: + bot.dalleRefundUser(user,"Timeout. Please try again later.") + log.Errorf("[DALLE-%d] timeout", workerId) + return + // Got a tick, we should check on checkSomething() + case <-ticker: + t, err = dalleClient.GetTask(ctx, task.ID) + // handle err + if err != nil { + log.Errorf("[GetTask] %v", err.Error()) + bot.dalleRefundUser(user,"") + return } - return + if t.Status == dalle.StatusSucceeded { + fmt.Printf("[DALLE-%d] task succeeded for user %s", worker, GetUserStr(user.Telegram)) + // download the first generated image + for _, data := range t.Generations.Data { + err = bot.downloadAndSendImages(ctx, dalleClient, data, invoiceEvent) + if err != nil { + log.Errorf("[downloadAndSendImages-%d] %v", worker, err.Error()) + } + } + return - } else if t.Status == dalle.StatusRejected { - log.Errorf("[DALLE] rejected: %s", t.ID) - bot.dalleRefundUser(user, "Your prompt has been rejected by OpenAI. Do not use celebrity names, sexual expressions, or any other harmful content as prompt.") - return + } else if t.Status == dalle.StatusRejected { + log.Errorf("[DALLE-%d] rejected: %s", workerId, t.ID) + bot.dalleRefundUser(user,"Your prompt has been rejected by OpenAI. Do not use celebrity names, sexual expressions, or any other harmful content as prompt.") + return + } + log.Debugf("[DALLE-%d] pending for user %s", worker, GetUserStr(user.Telegram)) } - log.Debugf("[DALLE] pending for user %s", GetUserStr(user.Telegram)) } } + jobChan <- job } // downloadAndSendImages will download dalle images and send them to the payer.