Merge branch 'master' into feat/aggressive-offload

Author: uxtechie

comfy_extras/nodes_upscale_model.py

@@ -6,6 +6,7 @@
 import folder_paths
 from typing_extensions import override
 from comfy_api.latest import ComfyExtension, io
+import comfy.model_management
 
 try:
     from spandrel_extra_arches import EXTRA_REGISTRY
@@ -78,13 +79,15 @@ def execute(cls, upscale_model, image) -> io.NodeOutput:
         tile = 512
         overlap = 32
 
+        output_device = comfy.model_management.intermediate_device()
+
         oom = True
         try:
             while oom:
                 try:
                     steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap)
                     pbar = comfy.utils.ProgressBar(steps)
-                    s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
+                    s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a.float()), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar, output_device=output_device)
                     oom = False
                 except Exception as e:
                     model_management.raise_non_oom(e)
@@ -94,7 +97,7 @@ def execute(cls, upscale_model, image) -> io.NodeOutput:
         finally:
             upscale_model.to("cpu")
 
-        s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
+        s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0).to(comfy.model_management.intermediate_dtype())
         return io.NodeOutput(s)
 
     upscale = execute  # TODO: remove