Update app.py

app.py

@@ -28,13 +28,13 @@ description = """
 [[Project Page]](https://audio-agi.github.io/Separate-Anything-You-Describe) [[Paper]](https://audio-agi.github.io/Separate-Anything-You-Describe/AudioSep_arXiv.pdf) [[Code]](https://github.com/Audio-AGI/AudioSep)
 
 AudioSep is a foundation model for open-domain sound separation with natural language queries.
-AudioSep demonstrates strong separation performance and impressivezero-shot generalization ability on
+AudioSep demonstrates strong separation performance and impressive zero-shot generalization ability on
 numerous tasks such as audio event separation, musical instrument separation, and speech enhancement.
 """
 
 
-def inference(audio_file_path: str, text: str):
-    print(f"Separate audio from [{audio_file_path}] with textual query [{text}]")
+def inference(audio_file_path: str, text: str, use_advanced: bool):
+    print(f"Separate audio from [{audio_file_path}] with textual query [{text}] and advanced mode [{use_advanced}]")
     mixture, _ = librosa.load(audio_file_path, sr=32000, mono=True)
 
     with torch.no_grad():
@@ -49,9 +49,12 @@ def inference(audio_file_path: str, text: str):
             "condition": conditions,
         }
 
-        sep_segment = MODEL.ss_model(input_dict)["waveform"]
-
-        sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy()
+        if use_chunk:
+            sep_segment = model.ss_model.chunk_inference(input_dict)
+            sep_segment = np.squeeze(sep_segment)
+        else:
+            sep_segment = model.ss_model(input_dict)["waveform"]
+            sep_segment = sep_segment.squeeze(0).squeeze(0).data.cpu().numpy()
 
         return 32000, np.round(sep_segment * 32767).astype(np.int16)
 
@@ -62,21 +65,21 @@ with gr.Blocks(title="AudioSep") as demo:
         with gr.Column():
             input_audio = gr.Audio(label="Mixture", type="filepath")
             text = gr.Textbox(label="Text Query")
+            use_chunk = gr.Checkbox(label="Use Chunk Mode (Less Quality but Less Memory Used)", value=False)
         with gr.Column():
-            with gr.Column():
-                output_audio = gr.Audio(label="Separation Result", scale=10)
-                button = gr.Button(
-                    "Separate",
-                    variant="primary",
-                    scale=2,
-                    size="lg",
-                    interactive=True,
-                )
-                button.click(
-                    fn=inference, inputs=[input_audio, text], outputs=[output_audio]
-                )
+            output_audio = gr.Audio(label="Separation Result", scale=10)
+            button = gr.Button(
+                "Separate",
+                variant="primary",
+                scale=2,
+                size="lg",
+                interactive=True,
+            )
+            button.click(
+                fn=inference, inputs=[input_audio, text, use_chunk], outputs=[output_audio]
+            )
 
     gr.Markdown("## Examples")
     gr.Examples(examples=EXAMPLES, inputs=[input_audio, text])
 
-demo.queue().launch(share=True)
+demo.queue().launch(share=True)