소스 검색

Fix memory usage

Lengyue 1 년 전
부모
커밋
5b999e9197
2개의 변경된 파일2개의 추가작업 그리고 0개의 파일을 삭제
  1. 1 0
      tools/api.py
  2. 1 0
      tools/webui.py

+ 1 - 0
tools/api.py

@@ -71,6 +71,7 @@ class InvokeRequest(BaseModel):
     format: Literal["wav", "mp3", "flac"] = "wav"
 
 
+@torch.inference_mode()
 def inference(req: InvokeRequest):
     # Parse reference audio aka prompt
     prompt_tokens = None

+ 1 - 0
tools/webui.py

@@ -59,6 +59,7 @@ def build_html_error_message(error):
 
 
 @GPU_DECORATOR
+@torch.inference_mode()
 def inference(
     text,
     enable_reference_audio,