Move the prompts onto the model's device and cut the timbre prompt to 2 s, as the authors' demo does.

app.py

@@ -117,12 +117,15 @@ def generate_audio(timbre_path, rhythm_path, cfg_scale, top_p, mask_temperature,
     rhythm_in = load_audio(rhythm_path)
     original_sr = rhythm_in.sample_rate
 
-    timbre_sig = load_audio(timbre_path).resample(sample_rate)
-    rhythm_sig = rhythm_in.resample(sample_rate)
+    timbre_sig = load_audio(timbre_path).resample(sample_rate).to(DEVICE)
+    rhythm_sig = rhythm_in.clone().resample(sample_rate).to(DEVICE)
     timbre_sig.ensure_max_of_audio()
     rhythm_sig.ensure_max_of_audio()
 
+    # The model was trained with at most a third of its buffer as the timbre
+    # prefix, and the authors' demo cuts the timbre prompt to it as well
     prefix_dur = int(LOADED["max_duration"] / 3)
+    timbre_sig.truncate_samples(int(prefix_dur * sample_rate))
     timbre_tokens = tokenizer.encode(timbre_sig)
     rhythm_tokens = tokenizer.encode(rhythm_sig)
     tokens = torch.cat([timbre_tokens.tokens, rhythm_tokens.tokens], dim=-1)
@@ -191,7 +194,7 @@ if __name__ == "__main__":
     with gr.Blocks(title="TRIA") as demo:
         timbre_in = gr.Audio(type="filepath", label="Timbre Prompt").harp_required(True).set_info((
             "A recording of the drum sounds you want, such as a kit or a loop. "
-            "This decides what the result sounds like."
+            "This decides what the result sounds like. Only its first 2 seconds are used."
         ))
         rhythm_in = gr.Audio(type="filepath", label="Rhythm Prompt").harp_required(True).set_info((
             "A recording of the groove you want, which can be tapped or "