node_upscale_by.py: v1 nodes copied

node_upscale_by.py

@@ -9,8 +9,8 @@ from inspect import cleandoc as _cleandoc
 from frozendict import deepfreeze as _deepfreeze
 
 from comfy.comfy_types.node_typing import IO as _IO
-from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel as _ImageUpscaleWithModel
-from nodes import ImageScaleBy as _ImageScaleBy
+# from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel as _ImageUpscaleWithModel
+# from nodes import ImageScaleBy as _ImageScaleBy
 
 from . import _meta
 from .docstring_formatter import format_docstring as _format_docstring
@@ -18,6 +18,85 @@ from .node_scale import _scale_type_dict as __scale_type_dict_base
 from ._funcs import _show_text_on_node
 
 
+# ==========================================================
+# Copy of built-in ComfyUI nodes with v1 schema,
+# as a temporary workaround
+
+
+import torch
+
+from comfy import model_management
+import comfy.utils
+
+
+class _ImageUpscaleWithModel:
+	@classmethod
+	def INPUT_TYPES(s):
+		return {"required": {
+			"upscale_model": ("UPSCALE_MODEL",), "image": ("IMAGE",),
+		}}
+	RETURN_TYPES = ("IMAGE",)
+	FUNCTION = "upscale"
+
+	CATEGORY = "image/upscaling"
+
+	def upscale(self, upscale_model, image):
+		device = model_management.get_torch_device()
+
+		memory_required = model_management.module_size(upscale_model.model)
+		memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
+		memory_required += image.nelement() * image.element_size()
+		model_management.free_memory(memory_required, device)
+
+		upscale_model.to(device)
+		in_img = image.movedim(-1,-3).to(device)
+
+		tile = 512
+		overlap = 32
+
+		oom = True
+		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)
+				oom = False
+			except model_management.OOM_EXCEPTION as e:
+				tile //= 2
+				if tile < 128:
+					raise e
+
+		upscale_model.to("cpu")
+		s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
+		return (s,)
+
+
+class _ImageScaleBy:
+	upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
+
+	@classmethod
+	def INPUT_TYPES(s):
+		return {"required": {
+			"image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
+			"scale_by": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 8.0, "step": 0.01}),
+		}}
+	RETURN_TYPES = ("IMAGE",)
+	FUNCTION = "upscale"
+
+	CATEGORY = "image/upscaling"
+
+	def upscale(self, image, upscale_method, scale_by):
+		samples = image.movedim(-1,1)
+		width = round(samples.shape[3] * scale_by)
+		height = round(samples.shape[2] * scale_by)
+		s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled")
+		s = s.movedim(1,-1)
+		return (s,)
+
+
+# ==========================================================
+
+
 _ImageUpscaleWithModel_instance = _ImageUpscaleWithModel()
 __ImageUpscaleWithModel_input_types: _t.Dict[str, dict] = _ImageUpscaleWithModel_instance.INPUT_TYPES()
 __ImageUpscaleWithModel_input_types_required = __ImageUpscaleWithModel_input_types.get('required', dict())