summaryrefslogtreecommitdiff
path: root/pngmeta
diff options
context:
space:
mode:
Diffstat (limited to 'pngmeta')
-rw-r--r--pngmeta/altwriter.go133
-rw-r--r--pngmeta/metareader.go147
-rw-r--r--pngmeta/metareader_test.go194
-rw-r--r--pngmeta/partsreader.go75
-rw-r--r--pngmeta/partswriter.go112
5 files changed, 661 insertions, 0 deletions
diff --git a/pngmeta/altwriter.go b/pngmeta/altwriter.go
new file mode 100644
index 0000000..206b563
--- /dev/null
+++ b/pngmeta/altwriter.go
@@ -0,0 +1,133 @@
+package pngmeta
+
+import (
+ "bytes"
+ "gf-lt/models"
+ "encoding/base64"
+ "encoding/binary"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "hash/crc32"
+ "io"
+ "os"
+)
+
+const (
+ pngHeader = "\x89PNG\r\n\x1a\n"
+ textChunkType = "tEXt"
+)
+
+// WriteToPng embeds the metadata into the specified PNG file and writes the result to outfile.
+func WriteToPng(metadata *models.CharCardSpec, sourcePath, outfile string) error {
+ pngData, err := os.ReadFile(sourcePath)
+ if err != nil {
+ return err
+ }
+ jsonData, err := json.Marshal(metadata)
+ if err != nil {
+ return err
+ }
+ base64Data := base64.StdEncoding.EncodeToString(jsonData)
+ embedData := PngEmbed{
+ Key: "gf-lt", // Replace with appropriate key constant
+ Value: base64Data,
+ }
+ var outputBuffer bytes.Buffer
+ if _, err := outputBuffer.Write([]byte(pngHeader)); err != nil {
+ return err
+ }
+ chunks, iend, err := processChunks(pngData[8:])
+ if err != nil {
+ return err
+ }
+ for _, chunk := range chunks {
+ outputBuffer.Write(chunk)
+ }
+ newChunk, err := createTextChunk(embedData)
+ if err != nil {
+ return err
+ }
+ outputBuffer.Write(newChunk)
+ outputBuffer.Write(iend)
+ return os.WriteFile(outfile, outputBuffer.Bytes(), 0666)
+}
+
+// processChunks extracts non-tEXt chunks and locates the IEND chunk
+func processChunks(data []byte) ([][]byte, []byte, error) {
+ var (
+ chunks [][]byte
+ iendChunk []byte
+ reader = bytes.NewReader(data)
+ )
+ for {
+ var chunkLength uint32
+ if err := binary.Read(reader, binary.BigEndian, &chunkLength); err != nil {
+ if errors.Is(err, io.EOF) {
+ break
+ }
+ return nil, nil, fmt.Errorf("error reading chunk length: %w", err)
+ }
+ chunkType := make([]byte, 4)
+ if _, err := reader.Read(chunkType); err != nil {
+ return nil, nil, fmt.Errorf("error reading chunk type: %w", err)
+ }
+ chunkData := make([]byte, chunkLength)
+ if _, err := reader.Read(chunkData); err != nil {
+ return nil, nil, fmt.Errorf("error reading chunk data: %w", err)
+ }
+ crc := make([]byte, 4)
+ if _, err := reader.Read(crc); err != nil {
+ return nil, nil, fmt.Errorf("error reading CRC: %w", err)
+ }
+ fullChunk := bytes.NewBuffer(nil)
+ if err := binary.Write(fullChunk, binary.BigEndian, chunkLength); err != nil {
+ return nil, nil, fmt.Errorf("error writing chunk length: %w", err)
+ }
+ if _, err := fullChunk.Write(chunkType); err != nil {
+ return nil, nil, fmt.Errorf("error writing chunk type: %w", err)
+ }
+ if _, err := fullChunk.Write(chunkData); err != nil {
+ return nil, nil, fmt.Errorf("error writing chunk data: %w", err)
+ }
+ if _, err := fullChunk.Write(crc); err != nil {
+ return nil, nil, fmt.Errorf("error writing CRC: %w", err)
+ }
+ switch string(chunkType) {
+ case "IEND":
+ iendChunk = fullChunk.Bytes()
+ return chunks, iendChunk, nil
+ case textChunkType:
+ continue // Skip existing tEXt chunks
+ default:
+ chunks = append(chunks, fullChunk.Bytes())
+ }
+ }
+ return nil, nil, errors.New("IEND chunk not found")
+}
+
+// createTextChunk generates a valid tEXt chunk with proper CRC
+func createTextChunk(embed PngEmbed) ([]byte, error) {
+ content := bytes.NewBuffer(nil)
+ content.WriteString(embed.Key)
+ content.WriteByte(0) // Null separator
+ content.WriteString(embed.Value)
+ data := content.Bytes()
+ crc := crc32.NewIEEE()
+ crc.Write([]byte(textChunkType))
+ crc.Write(data)
+ chunk := bytes.NewBuffer(nil)
+ if err := binary.Write(chunk, binary.BigEndian, uint32(len(data))); err != nil {
+ return nil, fmt.Errorf("error writing chunk length: %w", err)
+ }
+ if _, err := chunk.Write([]byte(textChunkType)); err != nil {
+ return nil, fmt.Errorf("error writing chunk type: %w", err)
+ }
+ if _, err := chunk.Write(data); err != nil {
+ return nil, fmt.Errorf("error writing chunk data: %w", err)
+ }
+ if err := binary.Write(chunk, binary.BigEndian, crc.Sum32()); err != nil {
+ return nil, fmt.Errorf("error writing CRC: %w", err)
+ }
+ return chunk.Bytes(), nil
+}
diff --git a/pngmeta/metareader.go b/pngmeta/metareader.go
new file mode 100644
index 0000000..369345a
--- /dev/null
+++ b/pngmeta/metareader.go
@@ -0,0 +1,147 @@
+package pngmeta
+
+import (
+ "bytes"
+ "gf-lt/models"
+ "encoding/base64"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "log/slog"
+ "os"
+ "path"
+ "strings"
+)
+
+const (
+ embType = "tEXt"
+ cKey = "chara"
+ IEND = "IEND"
+ header = "\x89PNG\r\n\x1a\n"
+ writeHeader = "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A"
+)
+
+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
+ }
+ specWrap := &models.Spec2Wrapper{}
+ if card.Name == "" {
+ if err := json.Unmarshal(data, &specWrap); err != nil {
+ return nil, err
+ }
+ return &specWrap.Data, nil
+ }
+ 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
+ }
+ if charSpec.Name == "" {
+ return nil, fmt.Errorf("failed to find role; fname %s", fname)
+ }
+ return charSpec.Simplify(uname, fname), nil
+}
+
+func readCardJson(fname string) (*models.CharCard, error) {
+ data, err := os.ReadFile(fname)
+ if err != nil {
+ return nil, err
+ }
+ card := models.CharCard{}
+ if err := json.Unmarshal(data, &card); err != nil {
+ return nil, err
+ }
+ return &card, nil
+}
+
+func ReadDirCards(dirname, uname string, log *slog.Logger) ([]*models.CharCard, error) {
+ files, err := os.ReadDir(dirname)
+ if err != nil {
+ return nil, err
+ }
+ resp := []*models.CharCard{}
+ for _, f := range files {
+ if f.IsDir() {
+ continue
+ }
+ if strings.HasSuffix(f.Name(), ".png") {
+ fpath := path.Join(dirname, f.Name())
+ cc, err := ReadCard(fpath, uname)
+ if err != nil {
+ log.Warn("failed to load card", "error", err)
+ continue
+ // return nil, err // better to log and continue
+ }
+ resp = append(resp, cc)
+ }
+ if strings.HasSuffix(f.Name(), ".json") {
+ fpath := path.Join(dirname, f.Name())
+ cc, err := readCardJson(fpath)
+ if err != nil {
+ return nil, err // better to log and continue
+ }
+ cc.FirstMsg = strings.ReplaceAll(strings.ReplaceAll(cc.FirstMsg, "{{char}}", cc.Role), "{{user}}", uname)
+ cc.SysPrompt = strings.ReplaceAll(strings.ReplaceAll(cc.SysPrompt, "{{char}}", cc.Role), "{{user}}", uname)
+ 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..f88de06
--- /dev/null
+++ b/pngmeta/metareader_test.go
@@ -0,0 +1,194 @@
+package pngmeta
+
+import (
+ "bytes"
+ "gf-lt/models"
+ "encoding/base64"
+ "encoding/binary"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "image"
+ "image/color"
+ "image/png"
+ "io"
+ "os"
+ "path/filepath"
+ "testing"
+)
+
+func TestReadMeta(t *testing.T) {
+ cases := []struct {
+ Filename string
+ }{
+ {
+ 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", tc.Filename))
+ })
+ }
+}
+
+// Test helper: Create a simple PNG image with test shapes
+func createTestImage(t *testing.T) string {
+ img := image.NewRGBA(image.Rect(0, 0, 200, 200))
+ // Fill background with white
+ for y := 0; y < 200; y++ {
+ for x := 0; x < 200; x++ {
+ img.Set(x, y, color.White)
+ }
+ }
+ // Draw a red square
+ for y := 50; y < 150; y++ {
+ for x := 50; x < 150; x++ {
+ img.Set(x, y, color.RGBA{R: 255, A: 255})
+ }
+ }
+ // Draw a blue circle
+ center := image.Point{100, 100}
+ radius := 40
+ for y := center.Y - radius; y <= center.Y+radius; y++ {
+ for x := center.X - radius; x <= center.X+radius; x++ {
+ dx := x - center.X
+ dy := y - center.Y
+ if dx*dx+dy*dy <= radius*radius {
+ img.Set(x, y, color.RGBA{B: 255, A: 255})
+ }
+ }
+ }
+ // Create temp file
+ tmpDir := t.TempDir()
+ fpath := filepath.Join(tmpDir, "test-image.png")
+ f, err := os.Create(fpath)
+ if err != nil {
+ t.Fatalf("Error creating temp file: %v", err)
+ }
+ defer f.Close()
+ if err := png.Encode(f, img); err != nil {
+ t.Fatalf("Error encoding PNG: %v", err)
+ }
+ return fpath
+}
+
+func TestWriteToPng(t *testing.T) {
+ // Create test image
+ srcPath := createTestImage(t)
+ dstPath := filepath.Join(filepath.Dir(srcPath), "output.png")
+ // dstPath := "test.png"
+ // Create test metadata
+ metadata := &models.CharCardSpec{
+ Description: "Test image containing a red square and blue circle on white background",
+ }
+ // Embed metadata
+ if err := WriteToPng(metadata, srcPath, dstPath); err != nil {
+ t.Fatalf("WriteToPng failed: %v", err)
+ }
+ // Verify output file exists
+ if _, err := os.Stat(dstPath); os.IsNotExist(err) {
+ t.Fatalf("Output file not created: %v", err)
+ }
+ // Read and verify metadata
+ t.Run("VerifyMetadata", func(t *testing.T) {
+ data, err := os.ReadFile(dstPath)
+ if err != nil {
+ t.Fatalf("Error reading output file: %v", err)
+ }
+ // Verify PNG header
+ if string(data[:8]) != pngHeader {
+ t.Errorf("Invalid PNG header")
+ }
+ // Extract metadata
+ embedded := extractMetadata(t, data)
+ if embedded.Description != metadata.Description {
+ t.Errorf("Metadata mismatch\nWant: %q\nGot: %q",
+ metadata.Description, embedded.Description)
+ }
+ })
+ // Optional: Add cleanup if needed
+ // t.Cleanup(func() {
+ // os.Remove(dstPath)
+ // })
+}
+
+// Helper to extract embedded metadata from PNG bytes
+func extractMetadata(t *testing.T, data []byte) *models.CharCardSpec {
+ r := bytes.NewReader(data[8:]) // Skip PNG header
+ for {
+ var length uint32
+ if err := binary.Read(r, binary.BigEndian, &length); err != nil {
+ if errors.Is(err, io.EOF) {
+ break
+ }
+ t.Fatalf("Error reading chunk length: %v", err)
+ }
+ chunkType := make([]byte, 4)
+ if _, err := r.Read(chunkType); err != nil {
+ t.Fatalf("Error reading chunk type: %v", err)
+ }
+ // Read chunk data
+ chunkData := make([]byte, length)
+ if _, err := r.Read(chunkData); err != nil {
+ t.Fatalf("Error reading chunk data: %v", err)
+ }
+ // Read and discard CRC
+ if _, err := r.Read(make([]byte, 4)); err != nil {
+ t.Fatalf("Error reading CRC: %v", err)
+ }
+ if string(chunkType) == embType {
+ parts := bytes.SplitN(chunkData, []byte{0}, 2)
+ if len(parts) != 2 {
+ t.Fatalf("Invalid tEXt chunk format")
+ }
+ decoded, err := base64.StdEncoding.DecodeString(string(parts[1]))
+ if err != nil {
+ t.Fatalf("Base64 decode error: %v", err)
+ }
+ var result models.CharCardSpec
+ if err := json.Unmarshal(decoded, &result); err != nil {
+ t.Fatalf("JSON unmarshal error: %v", err)
+ }
+ return &result
+ }
+ }
+ t.Fatal("Metadata not found in PNG")
+ return nil
+}
+
+func readTextChunk(t *testing.T, r io.ReadSeeker) *models.CharCardSpec {
+ var length uint32
+ binary.Read(r, binary.BigEndian, &length)
+ chunkType := make([]byte, 4)
+ r.Read(chunkType)
+ data := make([]byte, length)
+ r.Read(data)
+ // Read CRC (but skip validation for test purposes)
+ crc := make([]byte, 4)
+ r.Read(crc)
+ parts := bytes.SplitN(data, []byte{0}, 2) // Split key-value pair
+ if len(parts) != 2 {
+ t.Fatalf("Invalid tEXt chunk format")
+ }
+ // key := string(parts[0])
+ value := parts[1]
+ decoded, err := base64.StdEncoding.DecodeString(string(value))
+ if err != nil {
+ t.Fatalf("Base64 decode error: %v; value: %s", err, string(value))
+ }
+ var result models.CharCardSpec
+ if err := json.Unmarshal(decoded, &result); err != nil {
+ t.Fatalf("JSON unmarshal error: %v", err)
+ }
+ return &result
+}
diff --git a/pngmeta/partsreader.go b/pngmeta/partsreader.go
new file mode 100644
index 0000000..d345a16
--- /dev/null
+++ b/pngmeta/partsreader.go
@@ -0,0 +1,75 @@
+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")
+)
+
+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
+}
diff --git a/pngmeta/partswriter.go b/pngmeta/partswriter.go
new file mode 100644
index 0000000..7282df6
--- /dev/null
+++ b/pngmeta/partswriter.go
@@ -0,0 +1,112 @@
+package pngmeta
+
+// import (
+// "bytes"
+// "encoding/binary"
+// "errors"
+// "fmt"
+// "hash/crc32"
+// "io"
+// )
+
+// type Writer struct {
+// w io.Writer
+// }
+
+// func NewPNGWriter(w io.Writer) (*Writer, error) {
+// if _, err := io.WriteString(w, writeHeader); err != nil {
+// return nil, err
+// }
+// return &Writer{w}, nil
+// }
+
+// func (w *Writer) WriteChunk(length int32, typ string, r io.Reader) error {
+// if err := binary.Write(w.w, binary.BigEndian, length); err != nil {
+// return err
+// }
+// if _, err := w.w.Write([]byte(typ)); err != nil {
+// return err
+// }
+// checksummer := crc32.NewIEEE()
+// checksummer.Write([]byte(typ))
+// if _, err := io.CopyN(io.MultiWriter(w.w, checksummer), r, int64(length)); err != nil {
+// return err
+// }
+// if err := binary.Write(w.w, binary.BigEndian, checksummer.Sum32()); err != nil {
+// return err
+// }
+// return nil
+// }
+
+// func WWriteToPngriteToPng(c *models.CharCardSpec, fpath, outfile string) error {
+// data, err := os.ReadFile(fpath)
+// if err != nil {
+// return err
+// }
+// jsonData, err := json.Marshal(c)
+// if err != nil {
+// return err
+// }
+// // Base64 encode the JSON data
+// base64Data := base64.StdEncoding.EncodeToString(jsonData)
+// pe := PngEmbed{
+// Key: cKey,
+// Value: base64Data,
+// }
+// w, err := WritetEXtToPngBytes(data, pe)
+// if err != nil {
+// return err
+// }
+// return os.WriteFile(outfile, w.Bytes(), 0666)
+// }
+
+// func WritetEXtToPngBytes(inputBytes []byte, pe PngEmbed) (outputBytes bytes.Buffer, err error) {
+// if !(string(inputBytes[:8]) == header) {
+// return outputBytes, errors.New("wrong file format")
+// }
+// reader := bytes.NewReader(inputBytes)
+// pngr, err := NewPNGStepReader(reader)
+// if err != nil {
+// return outputBytes, fmt.Errorf("NewReader(): %s", err)
+// }
+// pngw, err := NewPNGWriter(&outputBytes)
+// if err != nil {
+// return outputBytes, fmt.Errorf("NewWriter(): %s", err)
+// }
+// for {
+// chunk, err := pngr.Next()
+// if err != nil {
+// if errors.Is(err, io.EOF) {
+// break
+// }
+// return outputBytes, fmt.Errorf("NextChunk(): %s", err)
+// }
+// if chunk.Type() != embType {
+// // IENDChunkType will only appear on the final iteration of a valid PNG
+// if chunk.Type() == IEND {
+// // This is where we inject tEXtChunkType as the penultimate chunk with the new value
+// newtEXtChunk := []byte(fmt.Sprintf(tEXtChunkDataSpecification, pe.Key, pe.Value))
+// if err := pngw.WriteChunk(int32(len(newtEXtChunk)), embType, bytes.NewBuffer(newtEXtChunk)); err != nil {
+// return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
+// }
+// // Now we end the buffer with IENDChunkType chunk
+// if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
+// return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
+// }
+// } else {
+// // writes back original chunk to buffer
+// if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
+// return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
+// }
+// }
+// } else {
+// if _, err := io.Copy(io.Discard, chunk); err != nil {
+// return outputBytes, fmt.Errorf("io.Copy(io.Discard, chunk): %s", err)
+// }
+// }
+// if err := chunk.Close(); err != nil {
+// return outputBytes, fmt.Errorf("chunk.Close(): %s", err)
+// }
+// }
+// return outputBytes, nil
+// }
generated by cgit v1.2.3 (git 2.25.1) at 2025-07-03 02:54:44 +0000