wrice commited on
Commit
3b84cc8
·
1 Parent(s): 85c50e2

Handle audio_path being None

Browse files
Files changed (1) hide show
  1. app.py +30 -27
app.py CHANGED
@@ -23,42 +23,45 @@ def denoise(model_name: str, audio_path: str) -> str:
23
  if torch.cuda.is_available():
24
  model = model.cuda()
25
 
26
- stream_reader = torchaudio.io.StreamReader(audio_path)
27
- stream_reader.add_basic_audio_stream(
28
- frames_per_chunk=model.config.max_length,
29
- sample_rate=model.config.sample_rate,
30
- num_channels=1,
31
- )
 
32
 
33
- stream_writer = torchaudio.io.StreamWriter("denoised.wav")
34
- stream_writer.add_audio_stream(sample_rate=model.config.sample_rate, num_channels=1)
 
 
35
 
36
- chunk_size = model.config.max_length
37
 
38
- with stream_writer.open():
39
- for (audio_chunk,) in tqdm(stream_reader.stream()):
40
- if audio_chunk is None:
41
- break
42
 
43
- audio_chunk = audio_chunk.permute(1, 0)
44
- original_chunk_size = audio_chunk.size(-1)
45
 
46
- if audio_chunk.size(-1) < chunk_size:
47
- padding = chunk_size - audio_chunk.size(-1)
48
- audio_chunk = torch.nn.functional.pad(audio_chunk, (0, padding))
49
 
50
- if torch.cuda.is_available():
51
- audio_chunk = audio_chunk.cuda()
52
 
53
- with torch.no_grad():
54
- denoised_chunk = model(audio_chunk[None]).audio
55
- denoised_chunk = denoised_chunk[:, :, :original_chunk_size]
56
 
57
- stream_writer.write_audio_chunk(
58
- 0, denoised_chunk.squeeze(0).permute(1, 0).cpu()
59
- )
60
 
61
- return "denoised.wav"
62
 
63
 
64
  iface = gr.Interface(
 
23
  if torch.cuda.is_available():
24
  model = model.cuda()
25
 
26
+ if audio_path:
27
+ stream_reader = torchaudio.io.StreamReader(audio_path)
28
+ stream_reader.add_basic_audio_stream(
29
+ frames_per_chunk=model.config.max_length,
30
+ sample_rate=model.config.sample_rate,
31
+ num_channels=1,
32
+ )
33
 
34
+ stream_writer = torchaudio.io.StreamWriter("denoised.wav")
35
+ stream_writer.add_audio_stream(
36
+ sample_rate=model.config.sample_rate, num_channels=1
37
+ )
38
 
39
+ chunk_size = model.config.max_length
40
 
41
+ with stream_writer.open():
42
+ for (audio_chunk,) in tqdm(stream_reader.stream()):
43
+ if audio_chunk is None:
44
+ break
45
 
46
+ audio_chunk = audio_chunk.permute(1, 0)
47
+ original_chunk_size = audio_chunk.size(-1)
48
 
49
+ if audio_chunk.size(-1) < chunk_size:
50
+ padding = chunk_size - audio_chunk.size(-1)
51
+ audio_chunk = torch.nn.functional.pad(audio_chunk, (0, padding))
52
 
53
+ if torch.cuda.is_available():
54
+ audio_chunk = audio_chunk.cuda()
55
 
56
+ with torch.no_grad():
57
+ denoised_chunk = model(audio_chunk[None]).audio
58
+ denoised_chunk = denoised_chunk[..., :original_chunk_size]
59
 
60
+ stream_writer.write_audio_chunk(
61
+ 0, denoised_chunk.squeeze(0).permute(1, 0).cpu()
62
+ )
63
 
64
+ return "denoised.wav"
65
 
66
 
67
  iface = gr.Interface(