From 09548204aa147e8c5fee222ea04fc8b2f3ad7c6a Mon Sep 17 00:00:00 2001
From: Antoine C <hi@acolombier.dev>
Date: Tue, 21 May 2024 20:41:28 +0100
Subject: [PATCH] feat: add support for MP4 and ID3v2 cover

---
 pyproject.toml        |  4 +++
 requirements.txt      |  2 +-
 stemgen/nistemfile.py | 64 +++++++++++++++++++++++++++++++++++++++----
 3 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 0266342..723df76 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,3 +1,7 @@
 [build-system]
 requires = ["setuptools>=61.0", "Cython>=3.0.10"]
 build-backend = "setuptools.build_meta"
+
+
+[tool.licensecheck]
+using = "requirements:requirements.txt"
diff --git a/requirements.txt b/requirements.txt
index 1d2e8e7..a446ade 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
 demucs @ git+https://github.com/facebookresearch/demucs.git@583db9df0213ba5f5b3491eca5c993e7629f1949#egg=demucs
+tagpy @ git+https://github.com/acolombier/tagpy.git@c5de51fe9636b312bfe95dee8b051ab640976d43#egg=tagpy
 setuptools>=61.0
-pytaglib==3.0.0
 ffmpeg-python==0.2.0
 torch>=2.1.2
 torchaudio>=2.1.2
diff --git a/stemgen/nistemfile.py b/stemgen/nistemfile.py
index 574f42d..cf89f43 100644
--- a/stemgen/nistemfile.py
+++ b/stemgen/nistemfile.py
@@ -1,12 +1,56 @@
 import click
 
-import taglib
+import tagpy
+import tagpy.id3v2
 from torchaudio.io import StreamWriter, CodecConfig
 import stembox
 import torch
 
 from .constant import SAMPLE_RATE, CHUNK_SIZE
 
+SUPPORTED_TAGS = [
+    "title",
+    "artist",
+    "album",
+    "comment",
+    "genre",
+    "year",
+    "track",
+]
+
+
+def _extract_cover(f):
+    tag = None
+    if isinstance(f, tagpy.FileRef):
+        tag = f.tag()
+        f = f.file()
+    covers = []
+    if hasattr(tag, "covers"):
+        covers = tag.covers
+    elif hasattr(f, "ID3v2Tag"):
+        covers = [
+            a
+            for a in f.ID3v2Tag().frameList()
+            if isinstance(a, tagpy.id3v2.AttachedPictureFrame)
+        ]
+
+    if covers:
+        cover = covers[0]
+        fmt = tagpy.mp4.CoverArtFormats.Unknown
+        if isinstance(cover, tagpy.mp4.CoverArt):
+            return cover
+        else:
+            mime = cover.mimeType().lower().strip()
+            if "image/jpeg":
+                fmt = tagpy.mp4.CoverArtFormats.JPEG
+            elif "image/png":
+                fmt = tagpy.mp4.CoverArtFormats.PNG
+            elif "image/bmp":
+                fmt = tagpy.mp4.CoverArtFormats.BMP
+            elif "image/gif":
+                fmt = tagpy.mp4.CoverArtFormats.GIF
+            return tagpy.mp4.CoverArt(fmt, cover.picture())
+
 
 class NIStemFile:
     STEM_DEFAULT_LABEL = [
@@ -74,10 +118,20 @@ def write(self, original, stems):
                 progress.finish()
 
     def update_metadata(self, src, **stem_metadata):
-        with taglib.File(src) as src, taglib.File(
-            self.__path, save_on_exit=True
-        ) as dst:
-            dst.tags = src.tags
+        src = tagpy.FileRef(src)
+        dst = tagpy.FileRef(self.__path)
+
+        src_tag = src.tag()
+        dst_tag = dst.tag()
+        for tag in SUPPORTED_TAGS:
+            setattr(dst_tag, tag, getattr(src_tag, tag))
+
+        cover = _extract_cover(src)
+        if cover:
+            c = tagpy.mp4.CoverArtList()
+            c.append(cover)
+            dst_tag.covers = c
+        dst.save()
 
         with stembox.Stem(self.__path) as f:
             f.stems = [