From 69a64928149c28855bea68927a6bccc47b424abb Mon Sep 17 00:00:00 2001 From: Yuanpei-Wu Date: Tue, 20 Aug 2024 20:31:27 +0800 Subject: [PATCH] =?UTF-8?q?=E5=8F=96=E6=B6=88=E8=BF=9B=E5=BA=A6=E6=9D=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- comfy/utils.py | 4 ++-- comfy_extras/nodes_upscale_model.py | 24 ++++++++++++++++-------- latent_preview.py | 1 + server.py | 24 ++++++++++++------------ 4 files changed, 31 insertions(+), 22 deletions(-) diff --git a/comfy/utils.py b/comfy/utils.py index 0db9fbb62..2a5f191d2 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -627,12 +627,12 @@ def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_am PROGRESS_BAR_ENABLED = True def set_progress_bar_enabled(enabled): global PROGRESS_BAR_ENABLED - PROGRESS_BAR_ENABLED = enabled + PROGRESS_BAR_ENABLED = False PROGRESS_BAR_HOOK = None def set_progress_bar_global_hook(function): global PROGRESS_BAR_HOOK - PROGRESS_BAR_HOOK = function + PROGRESS_BAR_HOOK = None class ProgressBar: def __init__(self, total): diff --git a/comfy_extras/nodes_upscale_model.py b/comfy_extras/nodes_upscale_model.py index bca79ef2e..341c533e8 100644 --- a/comfy_extras/nodes_upscale_model.py +++ b/comfy_extras/nodes_upscale_model.py @@ -1,5 +1,7 @@ import os import logging +import time + from spandrel import ModelLoader, ImageModelDescriptor from comfy import model_management import torch @@ -49,12 +51,14 @@ class ImageUpscaleWithModel: CATEGORY = "image/upscaling" def upscale(self, upscale_model, image): + begin = time.time() device = model_management.get_torch_device() + print(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) + # 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) @@ -65,17 +69,21 @@ class ImageUpscaleWithModel: 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) + #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) + pbar = None + 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,output_device=device, pbar=pbar) + print("output.device: ", s.device) oom = False except model_management.OOM_EXCEPTION as e: tile //= 2 if tile < 128: raise e - upscale_model.to("cpu") + #upscale_model.to("cpu") s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0) + end = time.time()-begin + logging.info("upscale time: {}".format(end)) return (s,) NODE_CLASS_MAPPINGS = { diff --git a/latent_preview.py b/latent_preview.py index ae6c106e4..341c1dc77 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -46,6 +46,7 @@ class Latent2RGBPreviewer(LatentPreviewer): def get_previewer(device, latent_format): + return None previewer = None method = args.preview_method if method != LatentPreviewMethod.NoPreviews: diff --git a/server.py b/server.py index 0f98a5184..411e8355a 100644 --- a/server.py +++ b/server.py @@ -655,18 +655,18 @@ class PromptServer(): await send_socket_catch_exception(self.sockets[sid].send_json, message) def send_sync(self, event, data, sid=None): - # if hasattr(self, "last_prompt_id") is not True: - # logging.warning("server has not attr last_prompt_id") - # return - # prompt_id = self.last_prompt_id - # key = "comfy_response_" + prompt_id - # message = { - # "type": event, - # "data": data, - # } - # res = self.redis_op.zadd_with_expire(key, json.dumps(message), int(time.time()), 120) - # if not res: - # logging.warning(f"Failed to send {event}") + if hasattr(self, "last_prompt_id") is not True: + logging.warning("server has not attr last_prompt_id") + return + prompt_id = self.last_prompt_id + key = "comfy_response_" + prompt_id + message = { + "type": event, + "data": data, + } + res = self.redis_op.zadd_with_expire(key, json.dumps(message), int(time.time()), 120) + if not res: + logging.warning(f"Failed to send {event}") self.loop.call_soon_threadsafe( self.messages.put_nowait, (event, data, sid))