fix: missing support for OGG cover and atom corruption

Author: acolombier

requirements.txt

@@ -1,5 +1,5 @@
 demucs @ git+https://github.com/facebookresearch/demucs.git@583db9df0213ba5f5b3491eca5c993e7629f1949#egg=demucs
-tagpy @ git+https://github.com/acolombier/tagpy.git@c5de51fe9636b312bfe95dee8b051ab640976d43#egg=tagpy
+tagpy @ git+https://github.com/acolombier/tagpy.git@7be8801d1079cfb3c3d84839959e954781722eb0#egg=tagpy
 setuptools>=61.0
 ffmpeg-python==0.2.0
 torch>=2.1.2

setup.py

@@ -13,7 +13,7 @@
 URL = "https://github.com/acolombier/stemgen"
 EMAIL = "stemgen@acolombier.dev"
 AUTHOR = "Antoine Colombier"
-REQUIRES_PYTHON = ">=3.11.0"
+REQUIRES_PYTHON = ">=3.10.0"
 
 # Get version without explicitly loading the module.
 for line in open("stemgen/__init__.py"):

stemgen/__main__.py

@@ -197,6 +197,7 @@ def main(
         overlap=overlap,
         jobs=jobs,
     )
+    has_failure = False
     for file in files:
         file = str(Path(file).resolve())
         filename = ".".join(os.path.basename(file).split(".")[:-1])
@@ -209,7 +210,8 @@ def main(
                 fg="red",
                 err=True,
             )
-            exit(1)
+            has_failure |= True
+            continue
         click.echo(f"Processing {filename}...")
 
         src = Track(file)
@@ -230,6 +232,8 @@ def main(
             stem_4_color=vocal_stem_color,
         )
         click.secho(f"Stem generated in {os.path.basename(dst)}", bold=True, fg="green")
+    if has_failure:
+        exit(1)
 
 
 if __name__ == "__main__":

stemgen/nistemfile.py

@@ -2,6 +2,7 @@
 
 import tagpy
 import tagpy.id3v2
+import tagpy.ogg.flac
 from torchaudio.io import StreamWriter, CodecConfig
 import stembox
 import torch
@@ -27,29 +28,34 @@ def _extract_cover(f):
     covers = []
     if hasattr(tag, "covers"):
         covers = tag.covers
+    elif hasattr(tag, "pictureList"):
+        covers = tag.pictureList()
     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
+        data = None
+        if isinstance(cover, tagpy.ogg.flac.Picture):
+            data = cover.data()
         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())
+            data = cover.picture()
+        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, data)
 
 
 class NIStemFile:
@@ -118,6 +124,20 @@ def write(self, original, stems):
                 progress.finish()
 
     def update_metadata(self, src, **stem_metadata):
+        # FIXME generating metadata atom after the file tags
+        with stembox.Stem(self.__path) as f:
+            f.stems = [
+                dict(
+                    color=stem_metadata.get(
+                        f"stem_{i+1}_color",
+                    )
+                    or self.STEM_DEFAULT_COLOR[i],
+                    name=stem_metadata.get(f"stem_{i+1}_label")
+                    or self.STEM_DEFAULT_LABEL[i].title(),
+                )
+                for i in range(4)
+            ]
+
         src = tagpy.FileRef(src)
         dst = tagpy.FileRef(self.__path)
 
@@ -132,16 +152,3 @@ def update_metadata(self, src, **stem_metadata):
             c.append(cover)
             dst_tag.covers = c
         dst.save()
-
-        with stembox.Stem(self.__path) as f:
-            f.stems = [
-                dict(
-                    color=stem_metadata.get(
-                        f"stem_{i+1}_color",
-                    )
-                    or self.STEM_DEFAULT_COLOR[i],
-                    name=stem_metadata.get(f"stem_{i+1}_label")
-                    or self.STEM_DEFAULT_LABEL[i].title(),
-                )
-                for i in range(4)
-            ]