diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts index 73b4346c3..87cd4c23a 100644 --- a/packages/tasks/src/model-libraries-snippets.ts +++ b/packages/tasks/src/model-libraries-snippets.ts @@ -46,24 +46,25 @@ export const asteroid = (model: ModelData): string[] => [ model = BaseModel.from_pretrained("${model.id}")`, ]; -export const audioseal = (model: ModelData): string[] => [ - `# Watermark Generator -from audioseal import AudioSeal - -model = AudioSeal.load_generator("${model.id}") -# pass a tensor (tensor_wav) of shape (batch, channels, samples) and a sample rate -wav, sr = tensor_wav, 16000 - -watermark = model.get_watermark(wav, sr) -watermarked_audio = wav + watermark - -# Watermark Detector +export const audioseal = (model: ModelData): string[] => { + const watermarkSnippet = `# Watermark Generator + from audioseal import AudioSeal + + model = AudioSeal.load_generator("${model.id}") + # pass a tensor (tensor_wav) of shape (batch, channels, samples) and a sample rate + wav, sr = tensor_wav, 16000 + + watermark = model.get_watermark(wav, sr) + watermarked_audio = wav + watermark`; -detector = AudioSeal.load_detector("${model.id}") + const detectorSnippet = `# Watermark Detector + from audioseal import AudioSeal -result, message = detector.detect_watermark(watermarked_audio, sr) -`, -]; + detector = AudioSeal.load_detector("${model.id}") + + result, message = detector.detect_watermark(watermarked_audio, sr)`; + return [watermarkSnippet, detectorSnippet]; +}; function get_base_diffusers_model(model: ModelData): string { return model.cardData?.base_model?.toString() ?? "fill-in-base-model";