summaryrefslogtreecommitdiff
path: root/pngmeta
diff options
context:
space:
mode:
Diffstat (limited to 'pngmeta')
-rw-r--r--pngmeta/metareader.go107
-rw-r--r--pngmeta/metareader_test.go33
-rw-r--r--pngmeta/partsreader.go77
3 files changed, 217 insertions, 0 deletions
diff --git a/pngmeta/metareader.go b/pngmeta/metareader.go
new file mode 100644
index 0000000..ea726c9
--- /dev/null
+++ b/pngmeta/metareader.go
@@ -0,0 +1,107 @@
+package pngmeta
+
+import (
+ "bytes"
+ "elefant/models"
+ "encoding/base64"
+ "encoding/json"
+ "errors"
+ "io"
+ "os"
+ "path"
+ "strings"
+)
+
+const (
+ embType = "tEXt"
+)
+
+type PngEmbed struct {
+ Key string
+ Value string
+}
+
+func (c PngEmbed) GetDecodedValue() (*models.CharCardSpec, error) {
+ data, err := base64.StdEncoding.DecodeString(c.Value)
+ if err != nil {
+ return nil, err
+ }
+ card := &models.CharCardSpec{}
+ if err := json.Unmarshal(data, &card); err != nil {
+ return nil, err
+ }
+ return card, nil
+}
+
+func extractChar(fname string) (*PngEmbed, error) {
+ data, err := os.ReadFile(fname)
+ if err != nil {
+ return nil, err
+ }
+ reader := bytes.NewReader(data)
+ pr, err := NewPNGStepReader(reader)
+ if err != nil {
+ return nil, err
+ }
+ for {
+ step, err := pr.Next()
+ if err != nil {
+ if errors.Is(err, io.EOF) {
+ break
+ }
+ }
+ if step.Type() != embType {
+ if _, err := io.Copy(io.Discard, step); err != nil {
+ return nil, err
+ }
+ } else {
+ buf, err := io.ReadAll(step)
+ if err != nil {
+ return nil, err
+ }
+ dataInstep := string(buf)
+ values := strings.Split(dataInstep, "\x00")
+ if len(values) == 2 {
+ return &PngEmbed{Key: values[0], Value: values[1]}, nil
+ }
+ }
+ if err := step.Close(); err != nil {
+ return nil, err
+ }
+ }
+ return nil, errors.New("failed to find embedded char in png: " + fname)
+}
+
+func ReadCard(fname, uname string) (*models.CharCard, error) {
+ pe, err := extractChar(fname)
+ if err != nil {
+ return nil, err
+ }
+ charSpec, err := pe.GetDecodedValue()
+ if err != nil {
+ return nil, err
+ }
+ return charSpec.Simplify(uname, fname), nil
+}
+
+func ReadDirCards(dirname, uname string) ([]*models.CharCard, error) {
+ files, err := os.ReadDir(dirname)
+ if err != nil {
+ return nil, err
+ }
+ resp := []*models.CharCard{}
+ for _, f := range files {
+ if !strings.HasSuffix(f.Name(), ".png") {
+ continue
+ }
+ fpath := path.Join(dirname, f.Name())
+ cc, err := ReadCard(fpath, uname)
+ if err != nil {
+ // log err
+ return nil, err
+ // continue
+ }
+ resp = append(resp, cc)
+ }
+ return resp, nil
+}
diff --git a/pngmeta/metareader_test.go b/pngmeta/metareader_test.go
new file mode 100644
index 0000000..51dadc6
--- /dev/null
+++ b/pngmeta/metareader_test.go
@@ -0,0 +1,33 @@
+package pngmeta
+
+import (
+ "fmt"
+ "testing"
+)
+
+func TestReadMeta(t *testing.T) {
+ cases := []struct {
+ Filename string
+ }{
+ {
+ Filename: "../sysprompts/default_Seraphina.png",
+ },
+ {
+ Filename: "../sysprompts/llama.png",
+ },
+ }
+ for i, tc := range cases {
+ t.Run(fmt.Sprintf("test_%d", i), func(t *testing.T) {
+ // Call the readMeta function
+ pembed, err := extractChar(tc.Filename)
+ if err != nil {
+ t.Errorf("Expected no error, but got %v", err)
+ }
+ v, err := pembed.GetDecodedValue()
+ if err != nil {
+ t.Errorf("Expected no error, but got %v\n", err)
+ }
+ fmt.Printf("%+v\n", v.Simplify("Adam"))
+ })
+ }
+}
diff --git a/pngmeta/partsreader.go b/pngmeta/partsreader.go
new file mode 100644
index 0000000..b69e4c3
--- /dev/null
+++ b/pngmeta/partsreader.go
@@ -0,0 +1,77 @@
+package pngmeta
+
+import (
+ "encoding/binary"
+ "errors"
+ "hash"
+ "hash/crc32"
+ "io"
+)
+
+var (
+ ErrCRC32Mismatch = errors.New("crc32 mismatch")
+ ErrNotPNG = errors.New("not png")
+ ErrBadLength = errors.New("bad length")
+)
+
+const header = "\x89PNG\r\n\x1a\n"
+
+type PngChunk struct {
+ typ string
+ length int32
+ r io.Reader
+ realR io.Reader
+ checksummer hash.Hash32
+}
+
+func (c *PngChunk) Read(p []byte) (int, error) {
+ return io.TeeReader(c.r, c.checksummer).Read(p)
+}
+
+func (c *PngChunk) Close() error {
+ var crc32 uint32
+ if err := binary.Read(c.realR, binary.BigEndian, &crc32); err != nil {
+ return err
+ }
+ if crc32 != c.checksummer.Sum32() {
+ return ErrCRC32Mismatch
+ }
+ return nil
+}
+
+func (c *PngChunk) Type() string {
+ return c.typ
+}
+
+type Reader struct {
+ r io.Reader
+}
+
+func NewPNGStepReader(r io.Reader) (*Reader, error) {
+ expectedHeader := make([]byte, len(header))
+ if _, err := io.ReadFull(r, expectedHeader); err != nil {
+ return nil, err
+ }
+ if string(expectedHeader) != header {
+ return nil, ErrNotPNG
+ }
+ return &Reader{r}, nil
+}
+
+func (r *Reader) Next() (*PngChunk, error) {
+ var length int32
+ if err := binary.Read(r.r, binary.BigEndian, &length); err != nil {
+ return nil, err
+ }
+ if length < 0 {
+ return nil, ErrBadLength
+ }
+ var rawTyp [4]byte
+ if _, err := io.ReadFull(r.r, rawTyp[:]); err != nil {
+ return nil, err
+ }
+ typ := string(rawTyp[:])
+ checksummer := crc32.NewIEEE()
+ checksummer.Write([]byte(typ))
+ return &PngChunk{typ, length, io.LimitReader(r.r, int64(length)), r.r, checksummer}, nil
+}