🪛 Refine the code snippet x2

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";