Make --gpu-only put intermediate values in GPU memory instead of cpu.

This commit is contained in:
comfyanonymous
2023-12-08 02:35:45 -05:00
parent cdff081023
commit 9ac0b487ac
9 changed files with 36 additions and 29 deletions

View File

@@ -291,7 +291,7 @@ class Canny:
def detect_edge(self, image, low_threshold, high_threshold):
output = canny(image.to(comfy.model_management.get_torch_device()).movedim(-1, 1), low_threshold, high_threshold)
img_out = output[1].cpu().repeat(1, 3, 1, 1).movedim(1, -1)
img_out = output[1].to(comfy.model_management.intermediate_device()).repeat(1, 3, 1, 1).movedim(1, -1)
return (img_out,)
NODE_CLASS_MAPPINGS = {

View File

@@ -226,7 +226,7 @@ class Sharpen:
batch_size, height, width, channels = image.shape
kernel_size = sharpen_radius * 2 + 1
kernel = gaussian_kernel(kernel_size, sigma) * -(alpha*10)
kernel = gaussian_kernel(kernel_size, sigma, device=image.device) * -(alpha*10)
center = kernel_size // 2
kernel[center, center] = kernel[center, center] - kernel.sum() + 1.0
kernel = kernel.repeat(channels, 1, 1).unsqueeze(1)