Spaces:
Running on Zero
Running on Zero
| """CPU専用子プロセス。Gradio・spaces・GPU予約を読み込まない。""" | |
| import json,os,re,sys,uuid,time | |
| from pathlib import Path | |
| def atomic_json(path,data): | |
| temporary=path.with_name(path.name+'.'+uuid.uuid4().hex+'.tmp') | |
| temporary.write_text(json.dumps(data,ensure_ascii=False),encoding='utf-8');os.replace(temporary,path) | |
| def run(root,jid): | |
| started=time.perf_counter() | |
| if not re.fullmatch(r'[a-f0-9]{32}',jid):raise ValueError('invalid_job') | |
| root=Path(root).resolve();directory=root/jid | |
| spec=json.loads((directory/'cpu-request.json').read_text(encoding='utf-8')) | |
| import torch | |
| torch.set_num_threads(1) | |
| from ultra.lunaron_model import LunaronBaseSeparator | |
| from audio_export import save_stem_wav | |
| import job_control | |
| def notify(progress,message): | |
| job_control.check_cancelled(root,jid) | |
| atomic_json(directory/'cpu-progress.json',{'progress':progress,'status':message}) | |
| quality=spec['quality'];stems=spec['stems'] | |
| if stems not in (2,4,6) or quality not in ('standard','hq','ultra'):raise ValueError('invalid_settings') | |
| path=Path(spec['input_path']) | |
| if spec.get('is_yt'): | |
| import yt_dlp | |
| with yt_dlp.YoutubeDL({'format':'bestaudio/best','outtmpl':str(directory/'input.%(ext)s'),'postprocessors':[{'key':'FFmpegExtractAudio','preferredcodec':'wav','preferredquality':'0'}],'quiet':True,'no_warnings':True,'noplaylist':True}) as downloader:downloader.download([spec['input_path']]) | |
| path=directory/'input.wav' | |
| if not path.is_file():raise FileNotFoundError('input_missing') | |
| if quality=='ultra': | |
| from ultra.inference_policy import ultra_settings | |
| shifts,overlap=ultra_settings(stems);model_stems=stems | |
| else: | |
| shifts=(4 if stems==6 else 10) if quality=='hq' else 2 | |
| overlap=.50;model_stems=6 if stems==6 else 4 | |
| notify(0,'CPUで分離モデルを準備中です。GPU処理より時間がかかります。') | |
| separator=LunaronBaseSeparator(model_name='htdemucs_6s' if stems==6 else 'htdemucs_ft',stems_count=model_stems,device='cpu',shifts=shifts,overlap=overlap,callback=lambda d:notify(round(max(0,min(1,d.get('progress',0)))*80,1),'CPUで音声を分離しています。'),preserve_model_output=quality=='ultra') | |
| prepared=time.perf_counter() | |
| with job_control.inference_scope(root,jid):origin,parts=separator.separate_audio_file(str(path)) | |
| inferred=time.perf_counter() | |
| if quality=='ultra' and origin is not None and parts and all(v is not None for v in parts.values()): | |
| length=min(origin.shape[-1],*(v.shape[-1] for v in parts.values())) | |
| parts={name:value[..., :length].cpu() for name,value in parts.items()} | |
| if stems==2 and quality!='ultra': | |
| accompaniment=None | |
| for name,tensor in parts.items(): | |
| if name!='vocals':accompaniment=tensor.clone() if accompaniment is None else accompaniment+tensor | |
| parts={'vocals':parts.get('vocals'),'accompaniment':accompaniment} | |
| if len(parts)!=stems or any(v is None for v in parts.values()):raise RuntimeError('incomplete_stems') | |
| for index,(name,tensor) in enumerate(parts.items(),1): | |
| if name not in ('vocals','accompaniment','drums','bass','other','guitar','piano'):raise ValueError('invalid_stem') | |
| job_control.check_cancelled(root,jid) | |
| save_stem_wav(tensor,str(directory/(name+'.wav')),samplerate=separator.samplerate) | |
| notify(round(80+index/len(parts)*19,1),'CPU分離のトラックを保存しています。') | |
| atomic_json(directory/'cpu-result.json',{'stems':list(parts),'input_path':str(path),'timings_ms':{'prepare':round((prepared-started)*1000),'inference':round((inferred-prepared)*1000),'export':round((time.perf_counter()-inferred)*1000)}}) | |
| if __name__=='__main__': | |
| try:run(sys.argv[1],sys.argv[2]) | |
| except Exception as error: | |
| # 自由文・音源名・パスを公開エラーに含めない。 | |
| try:atomic_json(Path(sys.argv[1])/sys.argv[2]/'cpu-error.json',{'error_class':type(error).__name__}) | |
| except Exception:pass | |
| sys.exit(1) | |