取消进度条

This commit is contained in:
Yuanpei-Wu 2024-08-20 20:31:27 +08:00
parent da7b4fd11a
commit 69a6492814
4 changed files with 31 additions and 22 deletions

View File

@ -627,12 +627,12 @@ def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_am
PROGRESS_BAR_ENABLED = True PROGRESS_BAR_ENABLED = True
def set_progress_bar_enabled(enabled): def set_progress_bar_enabled(enabled):
global PROGRESS_BAR_ENABLED global PROGRESS_BAR_ENABLED
PROGRESS_BAR_ENABLED = enabled PROGRESS_BAR_ENABLED = False
PROGRESS_BAR_HOOK = None PROGRESS_BAR_HOOK = None
def set_progress_bar_global_hook(function): def set_progress_bar_global_hook(function):
global PROGRESS_BAR_HOOK global PROGRESS_BAR_HOOK
PROGRESS_BAR_HOOK = function PROGRESS_BAR_HOOK = None
class ProgressBar: class ProgressBar:
def __init__(self, total): def __init__(self, total):

View File

@ -1,5 +1,7 @@
import os import os
import logging import logging
import time
from spandrel import ModelLoader, ImageModelDescriptor from spandrel import ModelLoader, ImageModelDescriptor
from comfy import model_management from comfy import model_management
import torch import torch
@ -49,12 +51,14 @@ class ImageUpscaleWithModel:
CATEGORY = "image/upscaling" CATEGORY = "image/upscaling"
def upscale(self, upscale_model, image): def upscale(self, upscale_model, image):
begin = time.time()
device = model_management.get_torch_device() device = model_management.get_torch_device()
print(device)
memory_required = model_management.module_size(upscale_model.model) # 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 += (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() # memory_required += image.nelement() * image.element_size()
model_management.free_memory(memory_required, device) # model_management.free_memory(memory_required, device)
upscale_model.to(device) upscale_model.to(device)
in_img = image.movedim(-1,-3).to(device) in_img = image.movedim(-1,-3).to(device)
@ -65,17 +69,21 @@ class ImageUpscaleWithModel:
oom = True oom = True
while oom: while oom:
try: 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) #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 = 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) 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 oom = False
except model_management.OOM_EXCEPTION as e: except model_management.OOM_EXCEPTION as e:
tile //= 2 tile //= 2
if tile < 128: if tile < 128:
raise e raise e
upscale_model.to("cpu") #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)
end = time.time()-begin
logging.info("upscale time: {}".format(end))
return (s,) return (s,)
NODE_CLASS_MAPPINGS = { NODE_CLASS_MAPPINGS = {

View File

@ -46,6 +46,7 @@ class Latent2RGBPreviewer(LatentPreviewer):
def get_previewer(device, latent_format): def get_previewer(device, latent_format):
return None
previewer = None previewer = None
method = args.preview_method method = args.preview_method
if method != LatentPreviewMethod.NoPreviews: if method != LatentPreviewMethod.NoPreviews:

View File

@ -655,18 +655,18 @@ class PromptServer():
await send_socket_catch_exception(self.sockets[sid].send_json, message) await send_socket_catch_exception(self.sockets[sid].send_json, message)
def send_sync(self, event, data, sid=None): def send_sync(self, event, data, sid=None):
# if hasattr(self, "last_prompt_id") is not True: if hasattr(self, "last_prompt_id") is not True:
# logging.warning("server has not attr last_prompt_id") logging.warning("server has not attr last_prompt_id")
# return return
# prompt_id = self.last_prompt_id prompt_id = self.last_prompt_id
# key = "comfy_response_" + prompt_id key = "comfy_response_" + prompt_id
# message = { message = {
# "type": event, "type": event,
# "data": data, "data": data,
# } }
# res = self.redis_op.zadd_with_expire(key, json.dumps(message), int(time.time()), 120) res = self.redis_op.zadd_with_expire(key, json.dumps(message), int(time.time()), 120)
# if not res: if not res:
# logging.warning(f"Failed to send {event}") logging.warning(f"Failed to send {event}")
self.loop.call_soon_threadsafe( self.loop.call_soon_threadsafe(
self.messages.put_nowait, (event, data, sid)) self.messages.put_nowait, (event, data, sid))