Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,392 Bytes
96dc011 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import os
import argparse
import silentcipher
import torch
import torchaudio
CSM_1B_HF_WATERMARK = list(map(int, os.getenv("WATERMARK_KEY").split(" ")))
def cli_check_audio() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--audio_path", type=str, required=True)
args = parser.parse_args()
check_audio_from_file(args.audio_path)
def load_watermarker(device: str = "cuda") -> silentcipher.server.Model:
model = silentcipher.get_model(
model_type="44.1k",
device=device,
)
return model
@torch.inference_mode()
def watermark(
watermarker: silentcipher.server.Model,
audio_array: torch.Tensor,
sample_rate: int,
watermark_key: list[int],
) -> tuple[torch.Tensor, int]:
audio_array_44khz = torchaudio.functional.resample(audio_array, orig_freq=sample_rate, new_freq=44100)
encoded, _ = watermarker.encode_wav(audio_array_44khz, 44100, watermark_key, calc_sdr=False, message_sdr=36)
output_sample_rate = min(44100, sample_rate)
encoded = torchaudio.functional.resample(encoded, orig_freq=44100, new_freq=output_sample_rate)
return encoded, output_sample_rate
@torch.inference_mode()
def verify(
watermarker: silentcipher.server.Model,
watermarked_audio: torch.Tensor,
sample_rate: int,
watermark_key: list[int],
) -> bool:
watermarked_audio_44khz = torchaudio.functional.resample(watermarked_audio, orig_freq=sample_rate, new_freq=44100)
result = watermarker.decode_wav(watermarked_audio_44khz, 44100, phase_shift_decoding=True)
is_watermarked = result["status"]
if is_watermarked:
is_csm_watermarked = result["messages"][0] == watermark_key
else:
is_csm_watermarked = False
return is_watermarked and is_csm_watermarked
def check_audio_from_file(audio_path: str) -> None:
watermarker = load_watermarker(device="cuda")
audio_array, sample_rate = load_audio(audio_path)
is_watermarked = verify(watermarker, audio_array, sample_rate, CSM_1B_HF_WATERMARK)
outcome = "Watermarked" if is_watermarked else "Not watermarked"
print(f"{outcome}: {audio_path}")
def load_audio(audio_path: str) -> tuple[torch.Tensor, int]:
audio_array, sample_rate = torchaudio.load(audio_path)
audio_array = audio_array.mean(dim=0)
return audio_array, int(sample_rate)
if __name__ == "__main__":
cli_check_audio()
|