Spaces:
Running on Zero
Running on Zero
| """Demucsの実際の区間完了通知を、単調な作業進捗に変換する。音声演算は変更しない。""" | |
| import math | |
| import time | |
| from contextlib import contextmanager | |
| class InferenceProgress: | |
| def __init__(self, callback, models, shifts, length, samplerate, overlap): | |
| self.callback = callback | |
| self.shifts = max(1, shifts) | |
| self.models = models | |
| self.length = length + int(samplerate * 0.5) | |
| self.samplerate = samplerate | |
| self.overlap = overlap | |
| self.last = -1.0 | |
| self.last_at = 0.0 | |
| self.completed = [0] * len(models) | |
| def advance(self, index): | |
| if not self.callback: | |
| return | |
| if not 0 <= index < len(self.models): | |
| return | |
| stride = max(1, int((1 - self.overlap) * self.samplerate * self.models[index].segment)) | |
| segments = max(1, math.ceil(self.length / stride)) | |
| self.completed[index] += 1 | |
| value = (index + min(1.0, self.completed[index] / (segments * self.shifts))) / len(self.models) | |
| value = min(0.999, max(self.last, value)) | |
| now = time.monotonic() | |
| if value - self.last >= 0.005 or now - self.last_at >= 1: | |
| self.last, self.last_at = value, now | |
| self.callback({'progress': value}) | |
| def track(self): | |
| # Demucs 4.0.1にはcallback引数がないため、公開PyTorch hookで区間完了だけを観測。 | |
| handles = [] | |
| try: | |
| if self.callback: | |
| for index, model in enumerate(self.models): | |
| def observed(_module, _inputs, _output, index=index): | |
| self.advance(index) | |
| handles.append(model.register_forward_hook(observed)) | |
| yield | |
| finally: | |
| for handle in handles: | |
| handle.remove() | |
| def finish(self): | |
| if self.callback: | |
| self.callback({'progress': 1.0}) | |