diff options
Diffstat (limited to 'pngmeta')
-rw-r--r-- | pngmeta/metareader.go | 107 | ||||
-rw-r--r-- | pngmeta/metareader_test.go | 33 | ||||
-rw-r--r-- | pngmeta/partsreader.go | 77 |
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 +} |