robocadey2

Signed-off-by: Xe <me@christine.website>

6 files changed, 225 insertions, 3 deletions

diff --git a/go.mod b/go.mod
index 9ac7fbc..ef131db 100644
--- a/go.mod
+++ b/go.mod
@@ -31,6 +31,7 @@ require (
 	github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2
 	github.com/pkg/errors v0.9.1
 	github.com/posener/complete v1.2.3
+	github.com/rs/cors v1.8.2
 	github.com/sebest/xff v0.0.0-20210106013422-671bd2870b3a
 	github.com/streamrail/concurrent-map v0.0.0-20160823150647-8bf1e9bacbf6
 	github.com/sycamoreone/orc v0.0.0-20150204213357-1627eaec2699
@@ -110,7 +111,6 @@ require (
 	github.com/olekukonko/tablewriter v0.0.5 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/rogpeppe/go-internal v1.8.1-0.20211023094830-115ce09fd6b4 // indirect
-	github.com/rs/cors v1.8.2 // indirect
 	github.com/sendgrid/rest v2.4.1+incompatible // indirect
 	github.com/ssor/bom v0.0.0-20170718123548-6386211fdfcf // indirect
 	github.com/tailscale/certstore v0.1.1-0.20220316223106-78d6e1c49d8d // indirect
@@ -148,3 +148,5 @@ require (
 )
 
 go 1.18
+
+replace github.com/McKael/madon/v2 => /home/cadey/gohack/github.com/McKael/madon/v2
diff --git a/go.sum b/go.sum
index ecd6e40..6162c96 100644
--- a/go.sum
+++ b/go.sum
@@ -39,8 +39,6 @@ filippo.io/mkcert v1.4.3 h1:axpnmtrZMM8u5Hf4N3UXxboGemMOV+Tn+e+pkHM6E3o=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/toml v1.1.0 h1:ksErzDEI1khOiGPgpwuI7x2ebx/uXQNw7xJpn9Eq1+I=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
-github.com/McKael/madon/v2 v2.0.0-20180929094633-c679abc985d6 h1:9cJcTOeILzInNo+DCYmXKME1QfAP07FYdo3M9/9jyc4=
-github.com/McKael/madon/v2 v2.0.0-20180929094633-c679abc985d6/go.mod h1:mvlJhxZCchfiasx3XvN3hBu5RekGwTDm09dKlSM/dQQ=
 github.com/Yawning/bulb v0.0.0-20170405033506-85d80d893c3d h1:ASfbc4vPGK69qbMO5qqRtGxfPvUmfbaK1ryEBv9WnQQ=
 github.com/Yawning/bulb v0.0.0-20170405033506-85d80d893c3d/go.mod h1:v1YQIDSfaeRRehor57unRn66SBtmgAzxfsCpKIV998I=
 github.com/aclements/go-moremath v0.0.0-20190506201756-286cc0be6f75 h1:OMm3VB/wrXTgQaFh/2z7xTH6fqUy3Q77ybHjmAJ/Ca4=
diff --git a/mastodon/robocadey/gpt2/main.py b/mastodon/robocadey/gpt2/main.py
index 750f338..4bb0ef2 100755
--- a/mastodon/robocadey/gpt2/main.py
+++ b/mastodon/robocadey/gpt2/main.py
@@ -27,6 +27,7 @@ while True:
                 top_p=0.9,
                 )[0].split("\n")[1:][:-1]
     print("shitpost generated")
+    print(json.dumps(result))
     connection.send(json.dumps(result).encode())
     connection.close()
 
diff --git a/mastodon/robocadey2/do_image_gen.sh b/mastodon/robocadey2/do_image_gen.sh
new file mode 100755
index 0000000..18978e1
--- /dev/null
+++ b/mastodon/robocadey2/do_image_gen.sh
@@ -0,0 +1,10 @@
+#!/usr/bin/env nix-shell
+#! nix-shell -p conda -i conda-shell
+
+set -euo pipefail
+
+cd $HOME/code/stable-diffusion
+source ~/.bashrc
+conda activate ldm
+
+python optimizedSD/robocadey.py --H 512 --W 512 --turbo --seed $SEED --from-file $OUTDIR/prompt.txt --outdir $OUTDIR
diff --git a/mastodon/robocadey2/main.go b/mastodon/robocadey2/main.go
new file mode 100644
index 0000000..1b49399
--- /dev/null
+++ b/mastodon/robocadey2/main.go
@@ -0,0 +1,205 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"math/rand"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"time"
+
+	madon "github.com/McKael/madon/v2"
+	"github.com/jaytaylor/html2text"
+	"within.website/ln"
+	"within.website/ln/opname"
+	"within.website/x/internal"
+)
+
+var (
+	instance  = flag.String("instance", "", "mastodon instance")
+	appID     = flag.String("app-id", "", "oauth2 app id")
+	appSecret = flag.String("app-secret", "", "oauth2 app secret")
+	token     = flag.String("token", "", "oauth2 token")
+)
+
+func main() {
+	internal.HandleStartup()
+
+	ctx := opname.With(context.Background(), "main")
+	rand.Seed(time.Now().Unix())
+
+	c, err := madon.RestoreApp("Robocadey2", *instance, *appID, *appSecret, &madon.UserToken{AccessToken: *token})
+	if err != nil {
+		ln.FatalErr(opname.With(ctx, "restore-app"), err)
+	}
+
+	ln.Log(ctx, ln.Info("waiting for messages"))
+
+	for {
+		evChan := make(chan madon.StreamEvent, 10)
+		stop := make(chan bool)
+		done := make(chan bool)
+		ctx = opname.With(context.Background(), "notifications-stream")
+
+		err = c.StreamListener("user", "", evChan, stop, done)
+		if err != nil {
+			ln.FatalErr(ctx, err)
+		}
+
+		for {
+			select {
+			case _, _ = <-done:
+				goto redo
+			case ev := <-evChan:
+				ln.Log(ctx, ln.F{"event": ev.Event})
+				switch ev.Event {
+				case "error":
+					ln.Error(opname.With(ctx, "event-parse"), err)
+					break
+				case "notification":
+					n := ev.Data.(madon.Notification)
+
+					if n.Type != "mention" {
+						continue
+					}
+
+					if err := handleNotification(c, n); err != nil {
+						ln.Error(ctx, err, ln.F{"content": n.Status.Content})
+						continue
+					}
+				}
+			}
+		}
+	}
+
+redo:
+}
+
+type DiffusionMetadata struct {
+	Prompt     string  `json:"prompt"`
+	Outdir     string  `json:"outdir"`
+	SkipGrid   bool    `json:"skip_grid"`
+	SkipSave   bool    `json:"skip_save"`
+	DDIMSteps  int     `json:"ddim_steps"`
+	FixedCode  bool    `json:"fixed_code"`
+	DDIMEta    float64 `json:"ddim_eta"`
+	NIter      int     `json:"n_iter"`
+	H          int     `json:"H"`
+	W          int     `json:"W"`
+	C          int     `json:"C"`
+	F          int     `json:"f"`
+	NSamples   int     `json:"n_samples"`
+	NRows      int     `json:"n_rows"`
+	Scale      float64 `json:"scale"`
+	Device     string  `json:"device"`
+	Seed       int     `json:"seed"`
+	UNETBs     int     `json:"unet_bs"`
+	Turbo      bool    `json:"turbo"`
+	Precision  string  `json:"precision"`
+	Format     string  `json:"format"`
+	Sampler    string  `json:"sampler"`
+	Checkpoint string  `json:"ckpt"`
+}
+
+func handleNotification(c *madon.Client, n madon.Notification) error {
+	text, err := html2text.FromString(n.Status.Content, html2text.Options{OmitLinks: true})
+	if err != nil {
+		return nil
+	}
+	text = strings.ReplaceAll(text, "@ ", "")
+
+	for _, m := range n.Status.Mentions {
+		text = strings.ReplaceAll(text, m.Username, "")
+	}
+
+	text = strings.TrimSpace(text)
+
+	fmt.Printf("text: %q\n", text)
+
+	dir, err := os.MkdirTemp("", "robocadey2")
+	if err != nil {
+		return err
+	}
+	defer os.RemoveAll(dir)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+	defer cancel()
+
+	metadata, err := makeImage(ctx, text, dir)
+	if err != nil {
+		return err
+	}
+
+	response := &strings.Builder{}
+
+	response.WriteString("@")
+	response.WriteString(n.Status.Account.Acct)
+	response.WriteString(" ")
+
+	for _, m := range n.Status.Mentions {
+		if m.Acct == "robocadey" {
+			continue
+		}
+
+		response.WriteString("@")
+		response.WriteString(m.Acct)
+		response.WriteString(" ")
+	}
+
+	att, err := c.UploadMedia(filepath.Join(dir, "00002.png"), "Waifu Diffusion v1.3 with the prompt: "+text, "")
+	if err != nil {
+		c.PostStatus(madon.PostStatusParams{
+			Text:       response.String() + " @cadey please help: " + err.Error(),
+			Visibility: n.Status.Visibility,
+			InReplyTo:  n.Status.ID,
+		})
+		return err
+	}
+
+	response.WriteString("here is your image:\n\n")
+	fmt.Fprintf(response, "prompt: %s\n", text)
+	fmt.Fprintf(response, "seed: %d\n", metadata.Seed)
+	response.WriteString("Generated with Waifu Diffusion v1.3 (float16)")
+
+	c.PostStatus(madon.PostStatusParams{
+		Text:        response.String(),
+		MediaIDs:    []int64{att.ID},
+		Sensitive:   true,
+		SpoilerText: "AI generated image (can be NSFW)",
+		Visibility:  n.Status.Visibility,
+		InReplyTo:   n.Status.ID,
+	})
+
+	return nil
+}
+
+func makeImage(ctx context.Context, prompt, dir string) (*DiffusionMetadata, error) {
+	err := os.WriteFile(filepath.Join(dir, "prompt.txt"), []byte(prompt), 0666)
+	if err != nil {
+		return nil, err
+	}
+
+	cmd := exec.CommandContext(ctx, "./do_image_gen.sh")
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	cmd.Env = append(os.Environ(), fmt.Sprintf("OUTDIR=%s", dir), fmt.Sprintf("SEED=%d", rand.Int31()))
+	if err := cmd.Run(); err != nil {
+		return nil, err
+	}
+
+	data, err := os.ReadFile(filepath.Join(dir, "metadata.json"))
+	if err != nil {
+		return nil, err
+	}
+
+	var result DiffusionMetadata
+	if err := json.Unmarshal(data, &result); err != nil {
+		return nil, err
+	}
+
+	return &result, nil
+}
diff --git a/mastodon/robocadey2/run.sh b/mastodon/robocadey2/run.sh
new file mode 100644
index 0000000..91de2bb
--- /dev/null
+++ b/mastodon/robocadey2/run.sh
@@ -0,0 +1,6 @@
+#!/bin/sh
+
+while true; do
+	go run .
+	sleep 2
+done