diff --git a/SDMeta/Metadata/PngMetadataExtractor.cs b/SDMeta/Metadata/PngMetadataExtractor.cs index 77d5de4..87bd680 100644 --- a/SDMeta/Metadata/PngMetadataExtractor.cs +++ b/SDMeta/Metadata/PngMetadataExtractor.cs @@ -1,6 +1,7 @@ using System; using System.Collections.Generic; using System.IO; +using System.IO.Compression; using System.Linq; using System.Text; using System.Threading.Tasks; @@ -34,6 +35,11 @@ public class PngMetadataExtractor SkipBytes(fs, 4); // Move past the current chunk CRC yield return keywordValuePair; break; + case "iTXt": + var itxtKeywordValuePair = await ReadInternationalTextualData(fs, chunkLength); + SkipBytes(fs, 4); // Move past the current chunk CRC + yield return itxtKeywordValuePair; + break; default: SkipBytes(fs, chunkLength + 4); // Move past the current chunk and its CRC break; @@ -80,10 +86,71 @@ private async static Task ReadChunkType(Stream stream) var nullIndex = dataString.IndexOf(NullTerminator); return new ( nullIndex > -1 ? dataString.Substring(0, nullIndex) : string.Empty, - nullIndex > -1 && nullIndex + 1 < length ? dataString.Substring(nullIndex + 1).TrimEnd(NullTerminator) : string.Empty + nullIndex > -1 && nullIndex + 1 < length ? dataString.Substring(nullIndex + 1).Trim(NullTerminator) : string.Empty ); } + private async static Task<(string Key, string Value)> ReadInternationalTextualData(Stream stream, int length) + { + var buffer = new byte[length]; + if (await stream.ReadAsync(buffer) != length) + { + throw new EndOfStreamException("Unexpected end of file while reading iTXt data."); + } + + // iTXt layout (all indices into buffer): + // [keyword]\0 [compression_flag:1] [compression_method:1] + // [language_tag]\0 [translated_keyword]\0 [text...] + var keywordEnd = Array.IndexOf(buffer, (byte)0); + if (keywordEnd < 0 || keywordEnd + 4 > buffer.Length) + { + return (string.Empty, string.Empty); + } + + var keyword = Encoding.UTF8.GetString(buffer, 0, keywordEnd); + var compressionFlag = buffer[keywordEnd + 1]; + var compressionMethod = buffer[keywordEnd + 2]; + // compression_method is at keywordEnd + 2; only method 0 (zlib) is defined + var afterFlags = keywordEnd + 3; + + // Skip language tag (null-terminated) + var languageEnd = Array.IndexOf(buffer, (byte)0, afterFlags); + if (languageEnd < 0 || languageEnd + 1 > buffer.Length) + { + return (keyword, string.Empty); + } + + // Skip translated keyword (null-terminated) + var translatedKeywordEnd = Array.IndexOf(buffer, (byte)0, languageEnd + 1); + if (translatedKeywordEnd < 0 || translatedKeywordEnd + 1 > buffer.Length) + { + return (keyword, string.Empty); + } + + var textStart = translatedKeywordEnd + 1; + var textBytes = buffer.AsSpan(textStart).ToArray(); + + string text; + if (compressionFlag == 1) + { + if (compressionMethod != 0) + { + throw new InvalidDataException($"Unsupported iTXt compression method: {compressionMethod}. Only zlib deflate (method 0) is supported."); + } + using var compressed = new MemoryStream(textBytes); + using var zlib = new ZLibStream(compressed, CompressionMode.Decompress); + using var decompressed = new MemoryStream(); + await zlib.CopyToAsync(decompressed); + text = Encoding.UTF8.GetString(decompressed.ToArray()); + } + else + { + text = Encoding.UTF8.GetString(textBytes); + } + + return (keyword, text); + } + private static void SkipBytes(Stream stream, int bytesToSkip) { var originalPosition = stream.Position; diff --git a/SDMetaTest/Metadata/PngMetadataExtractorTests.cs b/SDMetaTest/Metadata/PngMetadataExtractorTests.cs index 26eed3b..844e684 100644 --- a/SDMetaTest/Metadata/PngMetadataExtractorTests.cs +++ b/SDMetaTest/Metadata/PngMetadataExtractorTests.cs @@ -25,5 +25,39 @@ public async Task PngMetadataExtractorTest() Assert.Contains("Sánchez", prompt); Assert.EndsWith("v1.8.0", prompt); } + + [TestMethod] + public async Task PngMetadataExtractorTest_iTXt_Uncompressed() + { + using var fs = new FileSystem().FileStream.New("./Metadata/itxt-uncompressed.png", FileMode.Open); + + var items = PngMetadataExtractor.ExtractTextualInformation(fs); + var metadata = await items.ToDictionaryAsync(p => p.Key, p => p.Value); + + Assert.IsNotNull(metadata); + Assert.HasCount(1, metadata); + Assert.IsTrue(metadata.ContainsKey("parameters")); + + var prompt = metadata["parameters"]; + Assert.IsNotNull(prompt); + Assert.AreEqual("steps: 20, sampler: Euler a, cfg scale: 7, seed: 12345, size: 512x512, model: v1-5-pruned", prompt); + } + + [TestMethod] + public async Task PngMetadataExtractorTest_iTXt_Compressed() + { + using var fs = new FileSystem().FileStream.New("./Metadata/itxt-compressed.png", FileMode.Open); + + var items = PngMetadataExtractor.ExtractTextualInformation(fs); + var metadata = await items.ToDictionaryAsync(p => p.Key, p => p.Value); + + Assert.IsNotNull(metadata); + Assert.HasCount(1, metadata); + Assert.IsTrue(metadata.ContainsKey("parameters")); + + var prompt = metadata["parameters"]; + Assert.IsNotNull(prompt); + Assert.AreEqual("steps: 30, sampler: DPM++ 2M Karras, cfg scale: 8, seed: 67890, size: 768x768, model: v2-1", prompt); + } } } diff --git a/SDMetaTest/Metadata/itxt-compressed.png b/SDMetaTest/Metadata/itxt-compressed.png new file mode 100644 index 0000000..3389da4 Binary files /dev/null and b/SDMetaTest/Metadata/itxt-compressed.png differ diff --git a/SDMetaTest/Metadata/itxt-uncompressed.png b/SDMetaTest/Metadata/itxt-uncompressed.png new file mode 100644 index 0000000..f3c730d Binary files /dev/null and b/SDMetaTest/Metadata/itxt-uncompressed.png differ diff --git a/SDMetaTest/SDMetaTest.csproj b/SDMetaTest/SDMetaTest.csproj index 680d459..b6ab0bc 100644 --- a/SDMetaTest/SDMetaTest.csproj +++ b/SDMetaTest/SDMetaTest.csproj @@ -29,6 +29,12 @@ Always + + Always + + + Always +