From da6041dd937169c2b1fd7ebfa37848e81b9fab36 Mon Sep 17 00:00:00 2001 From: lllyasviel Date: Fri, 2 Feb 2024 00:34:18 -0800 Subject: [PATCH] elegant patcher --- .../scripts/preprocessor_marigold.py | 5 ++--- modules_forge/diffusers_patcher.py | 3 +++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/extensions-builtin/forge_preprocessor_marigold/scripts/preprocessor_marigold.py b/extensions-builtin/forge_preprocessor_marigold/scripts/preprocessor_marigold.py index 6add9b5c..62c38f5a 100644 --- a/extensions-builtin/forge_preprocessor_marigold/scripts/preprocessor_marigold.py +++ b/extensions-builtin/forge_preprocessor_marigold/scripts/preprocessor_marigold.py @@ -62,9 +62,8 @@ class PreprocessorMarigold(Preprocessor): ) with torch.no_grad(): - img = numpy_to_pytorch(input_image).movedim(-1, 1).to( - device=self.diffusers_patcher.patcher.current_device, - dtype=self.diffusers_patcher.dtype) + img = numpy_to_pytorch(input_image).movedim(-1, 1) + img = self.diffusers_patcher.move_tensor_to_current_device(img) img = img * 2.0 - 1.0 depth = self.diffusers_patcher.pipeline(img, num_inference_steps=20, show_pbar=False) diff --git a/modules_forge/diffusers_patcher.py b/modules_forge/diffusers_patcher.py index 1cd2ff49..8aa25386 100644 --- a/modules_forge/diffusers_patcher.py +++ b/modules_forge/diffusers_patcher.py @@ -46,3 +46,6 @@ class DiffusersModelPatcher: models=[self.patcher], memory_required=inference_memory ) + + def move_tensor_to_current_device(self, x): + return x.to(device=self.patcher.current_device, dtype=self.dtype)