Spaces:
Running
Running
File size: 2,730 Bytes
96dc011 31c247f 96dc011 13d6287 96dc011 31c247f |
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 |
import os
import argparse
import silentcipher
import torch
import torchaudio
# Set a default watermark key if environment variable is not set
watermark_key_str = os.getenv("WATERMARK_KEY")
if watermark_key_str is None:
CSM_1B_HF_WATERMARK = [0, 0, 0, 0] # Default placeholder
else:
CSM_1B_HF_WATERMARK = list(map(int, watermark_key_str.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)
# Fix the syntax error by properly unpacking the return values
result = watermarker.encode_wav(audio_array_44khz, 44100, watermark_key, calc_sdr=False, message_sdr=36)
encoded = result[0] # Assuming the first element is the encoded audio
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() |