286 lines
6.9 KiB
Go
286 lines
6.9 KiB
Go
package utils
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"errors"
|
||
"image"
|
||
"image/png"
|
||
"io"
|
||
)
|
||
|
||
// CharacterCardV2 SillyTavern 角色卡 V2 格式
|
||
type CharacterCardV2 struct {
|
||
Spec string `json:"spec"`
|
||
SpecVersion string `json:"spec_version"`
|
||
Data CharacterCardV2Data `json:"data"`
|
||
}
|
||
|
||
type CharacterCardV2Data struct {
|
||
Name string `json:"name"`
|
||
Description string `json:"description"`
|
||
Personality string `json:"personality"`
|
||
Scenario string `json:"scenario"`
|
||
FirstMes string `json:"first_mes"`
|
||
MesExample string `json:"mes_example"`
|
||
CreatorNotes string `json:"creator_notes"`
|
||
SystemPrompt string `json:"system_prompt"`
|
||
PostHistoryInstructions string `json:"post_history_instructions"`
|
||
Tags []string `json:"tags"`
|
||
Creator string `json:"creator"`
|
||
CharacterVersion string `json:"character_version"`
|
||
AlternateGreetings []string `json:"alternate_greetings"`
|
||
CharacterBook map[string]interface{} `json:"character_book,omitempty"`
|
||
Extensions map[string]interface{} `json:"extensions"`
|
||
}
|
||
|
||
// ExtractCharacterFromPNG 从 PNG 图片中提取角色卡数据
|
||
func ExtractCharacterFromPNG(pngData []byte) (*CharacterCardV2, error) {
|
||
reader := bytes.NewReader(pngData)
|
||
|
||
// 验证 PNG 格式（解码但不保存图片）
|
||
_, err := png.Decode(reader)
|
||
if err != nil {
|
||
return nil, errors.New("无效的 PNG 文件")
|
||
}
|
||
|
||
// 重新读取以获取 tEXt chunks
|
||
reader.Seek(0, 0)
|
||
|
||
// 查找 tEXt chunk 中的 "chara" 字段
|
||
charaJSON, err := extractTextChunk(reader, "chara")
|
||
if err != nil {
|
||
return nil, errors.New("PNG 中没有找到角色卡数据")
|
||
}
|
||
|
||
// 尝试 Base64 解码
|
||
decodedJSON, err := base64.StdEncoding.DecodeString(charaJSON)
|
||
if err != nil {
|
||
// 如果不是 Base64，直接使用原始 JSON
|
||
decodedJSON = []byte(charaJSON)
|
||
}
|
||
|
||
// 解析 JSON
|
||
var card CharacterCardV2
|
||
err = json.Unmarshal(decodedJSON, &card)
|
||
if err != nil {
|
||
return nil, errors.New("解析角色卡数据失败: " + err.Error())
|
||
}
|
||
|
||
return &card, nil
|
||
}
|
||
|
||
// extractTextChunk 从 PNG 中提取指定 key 的 tEXt chunk
|
||
func extractTextChunk(r io.Reader, key string) (string, error) {
|
||
// 跳过 PNG signature (8 bytes)
|
||
signature := make([]byte, 8)
|
||
if _, err := io.ReadFull(r, signature); err != nil {
|
||
return "", err
|
||
}
|
||
|
||
// 验证 PNG signature
|
||
expectedSig := []byte{137, 80, 78, 71, 13, 10, 26, 10}
|
||
if !bytes.Equal(signature, expectedSig) {
|
||
return "", errors.New("invalid PNG signature")
|
||
}
|
||
|
||
// 读取所有 chunks
|
||
for {
|
||
// 读取 chunk length (4 bytes)
|
||
lengthBytes := make([]byte, 4)
|
||
if _, err := io.ReadFull(r, lengthBytes); err != nil {
|
||
if err == io.EOF {
|
||
break
|
||
}
|
||
return "", err
|
||
}
|
||
length := uint32(lengthBytes[0])<<24 | uint32(lengthBytes[1])<<16 |
|
||
uint32(lengthBytes[2])<<8 | uint32(lengthBytes[3])
|
||
|
||
// 读取 chunk type (4 bytes)
|
||
chunkType := make([]byte, 4)
|
||
if _, err := io.ReadFull(r, chunkType); err != nil {
|
||
return "", err
|
||
}
|
||
|
||
// 读取 chunk data
|
||
data := make([]byte, length)
|
||
if _, err := io.ReadFull(r, data); err != nil {
|
||
return "", err
|
||
}
|
||
|
||
// 读取 CRC (4 bytes)
|
||
crc := make([]byte, 4)
|
||
if _, err := io.ReadFull(r, crc); err != nil {
|
||
return "", err
|
||
}
|
||
|
||
// 检查是否是 tEXt chunk
|
||
if string(chunkType) == "tEXt" {
|
||
// tEXt chunk 格式: keyword\0text
|
||
nullIndex := bytes.IndexByte(data, 0)
|
||
if nullIndex == -1 {
|
||
continue
|
||
}
|
||
|
||
keyword := string(data[:nullIndex])
|
||
text := string(data[nullIndex+1:])
|
||
|
||
if keyword == key {
|
||
return text, nil
|
||
}
|
||
}
|
||
|
||
// IEND chunk 表示结束
|
||
if string(chunkType) == "IEND" {
|
||
break
|
||
}
|
||
}
|
||
|
||
return "", errors.New("text chunk not found")
|
||
}
|
||
|
||
// EmbedCharacterToPNG 将角色卡数据嵌入到 PNG 图片中
|
||
func EmbedCharacterToPNG(img image.Image, card *CharacterCardV2) ([]byte, error) {
|
||
// 序列化角色卡数据
|
||
cardJSON, err := json.Marshal(card)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// Base64 编码
|
||
encodedJSON := base64.StdEncoding.EncodeToString(cardJSON)
|
||
|
||
// 创建一个 buffer 来写入 PNG
|
||
var buf bytes.Buffer
|
||
|
||
// 写入 PNG signature
|
||
buf.Write([]byte{137, 80, 78, 71, 13, 10, 26, 10})
|
||
|
||
// 编码原始图片到临时 buffer
|
||
var imgBuf bytes.Buffer
|
||
if err := png.Encode(&imgBuf, img); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 跳过原始 PNG 的 signature
|
||
imgData := imgBuf.Bytes()[8:]
|
||
|
||
// 将原始图片的 chunks 复制到输出，在 IEND 之前插入 tEXt chunk
|
||
r := bytes.NewReader(imgData)
|
||
|
||
for {
|
||
// 读取 chunk length
|
||
lengthBytes := make([]byte, 4)
|
||
if _, err := io.ReadFull(r, lengthBytes); err != nil {
|
||
if err == io.EOF {
|
||
break
|
||
}
|
||
return nil, err
|
||
}
|
||
length := uint32(lengthBytes[0])<<24 | uint32(lengthBytes[1])<<16 |
|
||
uint32(lengthBytes[2])<<8 | uint32(lengthBytes[3])
|
||
|
||
// 读取 chunk type
|
||
chunkType := make([]byte, 4)
|
||
if _, err := io.ReadFull(r, chunkType); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 读取 chunk data
|
||
data := make([]byte, length)
|
||
if _, err := io.ReadFull(r, data); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 读取 CRC
|
||
crc := make([]byte, 4)
|
||
if _, err := io.ReadFull(r, crc); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 如果是 IEND chunk，先写入 tEXt chunk
|
||
if string(chunkType) == "IEND" {
|
||
// 写入 tEXt chunk
|
||
writeTextChunk(&buf, "chara", encodedJSON)
|
||
}
|
||
|
||
// 写入原始 chunk
|
||
buf.Write(lengthBytes)
|
||
buf.Write(chunkType)
|
||
buf.Write(data)
|
||
buf.Write(crc)
|
||
|
||
if string(chunkType) == "IEND" {
|
||
break
|
||
}
|
||
}
|
||
|
||
return buf.Bytes(), nil
|
||
}
|
||
|
||
// writeTextChunk 写入 tEXt chunk
|
||
func writeTextChunk(w io.Writer, keyword, text string) error {
|
||
data := append([]byte(keyword), 0)
|
||
data = append(data, []byte(text)...)
|
||
|
||
// 写入 length
|
||
length := uint32(len(data))
|
||
lengthBytes := []byte{
|
||
byte(length >> 24),
|
||
byte(length >> 16),
|
||
byte(length >> 8),
|
||
byte(length),
|
||
}
|
||
w.Write(lengthBytes)
|
||
|
||
// 写入 type
|
||
w.Write([]byte("tEXt"))
|
||
|
||
// 写入 data
|
||
w.Write(data)
|
||
|
||
// 计算并写入 CRC
|
||
crcData := append([]byte("tEXt"), data...)
|
||
crc := calculateCRC(crcData)
|
||
crcBytes := []byte{
|
||
byte(crc >> 24),
|
||
byte(crc >> 16),
|
||
byte(crc >> 8),
|
||
byte(crc),
|
||
}
|
||
w.Write(crcBytes)
|
||
|
||
return nil
|
||
}
|
||
|
||
// calculateCRC 计算 CRC32
|
||
func calculateCRC(data []byte) uint32 {
|
||
crc := uint32(0xFFFFFFFF)
|
||
|
||
for _, b := range data {
|
||
crc ^= uint32(b)
|
||
for i := 0; i < 8; i++ {
|
||
if crc&1 != 0 {
|
||
crc = (crc >> 1) ^ 0xEDB88320
|
||
} else {
|
||
crc >>= 1
|
||
}
|
||
}
|
||
}
|
||
|
||
return crc ^ 0xFFFFFFFF
|
||
}
|
||
|
||
// ParseCharacterCardJSON 解析 JSON 格式的角色卡
|
||
func ParseCharacterCardJSON(jsonData []byte) (*CharacterCardV2, error) {
|
||
var card CharacterCardV2
|
||
err := json.Unmarshal(jsonData, &card)
|
||
if err != nil {
|
||
return nil, errors.New("解析角色卡 JSON 失败: " + err.Error())
|
||
}
|
||
|
||
return &card, nil
|
||
}
|