117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
import torch
|
|
import numpy as np
|
|
import comfy.model_management as model_management
|
|
import comfy.samplers
|
|
import comfy.utils
|
|
import comfy.sd
|
|
from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel
|
|
|
|
class YSC_HighResFix:
|
|
|
|
def __init__(self):
|
|
self.device = model_management.get_torch_device()
|
|
self.image_scaler = ImageUpscaleWithModel()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"vae": ("VAE",),
|
|
"image": ("IMAGE",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"upscale_factor": ("FLOAT", {
|
|
"default": 2.0,
|
|
"min": 1.0,
|
|
"max": 4.0,
|
|
"step": 0.01,
|
|
"display": "number"
|
|
}),
|
|
"upscale_method": (["nearest-exact", "bilinear", "bicubic", "lanczos"],),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff
|
|
}),
|
|
"steps": ("INT", {
|
|
"default": 20,
|
|
"min": 1,
|
|
"max": 1000
|
|
}),
|
|
"cfg": ("FLOAT", {
|
|
"default": 8.0,
|
|
"min": 1.0,
|
|
"max": 100.0,
|
|
"step": 0.1,
|
|
"round": 0.01
|
|
}),
|
|
"sampler_name": (comfy.samplers.SAMPLER_NAMES,),
|
|
"scheduler": (comfy.samplers.SCHEDULER_NAMES,),
|
|
"denoise": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01
|
|
}),
|
|
"upscale_model": ("UPSCALE_MODEL",),
|
|
},
|
|
"optional": {}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("LATENT",)
|
|
FUNCTION = "refine_upscale"
|
|
CATEGORY = "image/upscaling"
|
|
|
|
|
|
def refine_upscale(self, model, vae, image: torch.Tensor, positive, negative, upscale_factor,
|
|
upscale_method, seed, steps, cfg, sampler_name,
|
|
scheduler, denoise, upscale_model):
|
|
|
|
samples_upscaled_bhwc = self.image_scaler.upscale(upscale_model, image)[0]
|
|
|
|
samples = samples_upscaled_bhwc.movedim(-1,1)
|
|
|
|
target_width = round(image.shape[2] * upscale_factor)
|
|
target_height = round(image.shape[1] * upscale_factor)
|
|
|
|
upscaled_width = samples.shape[3]
|
|
upscaled_height = samples.shape[2]
|
|
|
|
if upscaled_width > target_width or upscaled_height > target_height:
|
|
print(f"[YSC HighResFix]: Resizing from {upscaled_width}x{upscaled_height} to {target_width}x{target_height}")
|
|
samples = comfy.utils.common_upscale(samples, target_width, target_height, upscale_method, "disabled")
|
|
|
|
pixels = samples.movedim(1, -1)
|
|
latent_image = vae.encode(pixels[:,:,:,:3])
|
|
|
|
|
|
latent = {"samples": latent_image}
|
|
print("[YSC HighResFix]: Upscaled image is now latent")
|
|
|
|
torch.manual_seed(seed)
|
|
noise = torch.randn_like(latent_image)
|
|
|
|
latent_sample_out = comfy.sample.sample(
|
|
model=model,
|
|
noise=noise,
|
|
steps=steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent_image=latent_image,
|
|
denoise=denoise
|
|
)
|
|
|
|
return ({"samples": latent_sample_out}, )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"YSC_HighresFix": YSC_HighResFix
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"YSC_HighresFix": "Highres Fix (YSC)"
|
|
} |