fix(api): mark all convert methods as no_grad

ssube · Mar 1, 2023 · 21fc7c5 · 21fc7c5
1 parent b44e644
commit 21fc7c5
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 0 deletions.
diff --git a/api/onnx_web/convert/diffusion/original.py b/api/onnx_web/convert/diffusion/original.py
@@ -1650,6 +1650,7 @@ def extract_checkpoint(
     logger.info(result_status)
 
 
+@torch.no_grad()
 def convert_diffusion_original(
     ctx: ConversionContext,
     model: ModelDict,

diff --git a/api/onnx_web/convert/diffusion/textual_inversion.py b/api/onnx_web/convert/diffusion/textual_inversion.py
@@ -11,6 +11,7 @@
 logger = getLogger(__name__)
 
 
+@torch.no_grad()
 def convert_diffusion_textual_inversion(
     context: ConversionContext, name: str, base_model: str, inversion: str
 ):

diff --git a/api/onnx_web/worker/context.py b/api/onnx_web/worker/context.py
@@ -74,14 +74,22 @@ def clear_flags(self) -> None:
 
 
 class JobStatus:
+    name: str
+    device: str
+    progress: int
+    cancelled: bool
+    finished: bool
+
     def __init__(
         self,
         name: str,
+        device: DeviceParams,
         progress: int = 0,
         cancelled: bool = False,
         finished: bool = False,
     ) -> None:
         self.name = name
+        self.device = device.device
         self.progress = progress
         self.cancelled = cancelled
         self.finished = finished