Skip to content

Commit

Permalink
fix: missing support for OGG cover and atom corruption
Browse files Browse the repository at this point in the history
  • Loading branch information
acolombier committed Jun 23, 2024
1 parent 4b28236 commit 6431668
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 27 deletions.
2 changes: 1 addition & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
URL = "https://github.com/acolombier/stemgen"
EMAIL = "[email protected]"
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"):
Expand Down
6 changes: 5 additions & 1 deletion stemgen/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand All @@ -209,7 +210,8 @@ def main(
fg="red",
err=True,
)
exit(1)
has_failure |= True
continue
click.echo(f"Processing {filename}...")

src = Track(file)
Expand All @@ -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__":
Expand Down
55 changes: 31 additions & 24 deletions stemgen/nistemfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import tagpy
import tagpy.id3v2
import tagpy.ogg.flac
from torchaudio.io import StreamWriter, CodecConfig
import stembox
import torch
Expand All @@ -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:
Expand Down Expand Up @@ -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)

Expand All @@ -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)
]

0 comments on commit 6431668

Please sign in to comment.