Resample an edit that runs on, up to three times, and explain the edit interval

app.py

@@ -94,6 +94,14 @@ def process_fn(audio_path: str, target_text: str, start_time: float, end_time: f
     return regenerate(audio_path, target_text, start_time, end_time, duration)
 
 
+# The replacement is generated token by token until the model decides it has ended, and it
+# sometimes runs on, stretching a two-second edit into ten seconds of speech. A result is
+# accepted when the recording has grown by no more than this many seconds, and otherwise the
+# edit is sampled again, up to the limit below, keeping the shortest result.
+MAX_GROWTH_SECONDS = 2.0
+MAX_ATTEMPTS = 3
+
+
 @spaces.GPU(duration=60)
 def regenerate(audio_path: str, target_text: str, start_time: float, end_time: float,
                duration: float) -> str:
@@ -109,22 +117,35 @@ def regenerate(audio_path: str, target_text: str, start_time: float, end_time: f
         "codec_sr": CODEC_RATE,
         "silence_tokens": [1388, 1898, 131],
     }
-    generated = inference_one_sample(
-        model,
-        model.args,
-        model.args.phn2num,
-        text_tokenizer,
-        audio_tokenizer,
-        audio_path,
-        target_text,
-        mask,
-        DEVICE,
-        decode_config,
-    )
+    best = None
+    for attempt in range(MAX_ATTEMPTS):
+        generated = inference_one_sample(
+            model,
+            model.args,
+            model.args.phn2num,
+            text_tokenizer,
+            audio_tokenizer,
+            audio_path,
+            target_text,
+            mask,
+            DEVICE,
+            decode_config,
+        )
+        candidate = generated[0].detach().cpu().float()
+        seconds = candidate.shape[-1] / 16000
+        print(f"VoiceCraft: attempt {attempt + 1} gives {seconds:.1f} s for a {duration:.1f} s "
+              "recording", flush=True)
+        if best is None or candidate.shape[-1] < best.shape[-1]:
+            best = candidate
+        if seconds <= duration + MAX_GROWTH_SECONDS:
+            break
+        # A sample that ran on is more likely to stop with silence penalized harder
+        decode_config["stop_repetition"] = max(1, decode_config["stop_repetition"] - 1)
+
     # The model is mono at 16 kHz, so the result is resampled back to the input's rate and copied
     # across its channels
     input_info = sf.info(audio_path)
-    audio = generated[0].detach().cpu().float()
+    audio = best
     if input_info.samplerate != 16000:
         audio = torchaudio.functional.resample(audio, 16000, input_info.samplerate)
     audio = audio[:1].T.repeat(1, input_info.channels)
@@ -145,12 +166,15 @@ if __name__ == "__main__":
                 .set_info("English speech to edit, up to 60 seconds."),
                 gr.Textbox(label="Target Transcript", lines=4)
                 .harp_required(True)
-                .set_info("Full transcript of the recording, with the replacement words in place, "
-                          "up to 500 characters."),
+                .set_info("The whole transcript of the recording, as it should read after the "
+                          "edit, up to 500 characters. Only the words inside the edit interval "
+                          "are regenerated, so change those and leave the rest as spoken."),
                 gr.Number(label="Edit Start (seconds)", value=0.5, precision=2)
-                .set_info("Start of the words to replace."),
+                .set_info("Where the words to replace begin, in seconds from the start of the "
+                          "recording. Read the time off the ruler of the input track."),
                 gr.Number(label="Edit End (seconds)", value=1.5, precision=2)
-                .set_info("End of the words to replace."),
+                .set_info("Where the words to replace end, in seconds. The interval should cover "
+                          "whole words, and a little silence on either side helps."),
             ],
             output_components=[
                 gr.Audio(type="filepath", label="Edited Speech")