audiobook_gen / src /output.py
Matthew Kutarna
Streamlit app development (#5)
74f2c64
raw
history blame
1.85 kB
"""
Notes
-----
This module contains the functions for audiobook_gen that take the generated audio tensors and output to audio files,
as well as assembling the final zip archive for user download.
"""
import logging
from src import config
def write_audio(audio_list, sample_path):
"""
Invokes torchaudio to save generated audio tensors to a file.
Parameters
----------
audio_list : torch.tensor
pytorch tensor containing generated audio
sample_path : str
file name and path for outputting tensor to audio file
Returns
-------
None
"""
import torch
import torchaudio
from src import config as cf
if not config.output_path.exists():
config.output_path.mkdir()
if len(audio_list) > 0:
audio_file = torch.cat(audio_list).reshape(1, -1)
torchaudio.save(sample_path, audio_file, cf.SAMPLE_RATE)
logging.info(f'Audio generated at: {sample_path}')
else:
logging.info(f'Audio at: {sample_path} is empty.')
def assemble_zip(title):
"""
Creates a zip file and inserts all .wav files in the output directory,
and returns the name / path of the zip file.
Parameters
----------
title : str
title of document, used to name zip directory
Returns
-------
zip_name : str
name and path of zip directory generated
"""
import zipfile
from stqdm import stqdm
if not config.output_path.exists():
config.output_path.mkdir()
zip_name = config.output_path / f'{title}.zip'
with zipfile.ZipFile(zip_name, mode="w") as archive:
for file_path in stqdm(config.output_path.iterdir()):
if file_path.suffix == '.wav':
archive.write(file_path, arcname=file_path.name)
file_path.unlink()
return zip_name