363 lines
12 KiB
Python
363 lines
12 KiB
Python
import gc
|
|
import os
|
|
from typing import Dict, Tuple
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
try:
|
|
import folder_paths
|
|
except Exception:
|
|
folder_paths = None
|
|
|
|
try:
|
|
import comfy.model_management as model_management
|
|
except Exception:
|
|
model_management = None
|
|
|
|
from .pbasr_core.archs import build_network
|
|
|
|
|
|
if folder_paths is not None:
|
|
MODEL_DIR = os.path.join(folder_paths.models_dir, 'pbasr')
|
|
else:
|
|
MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models')
|
|
os.makedirs(MODEL_DIR, exist_ok=True)
|
|
|
|
_MODEL_CACHE: Dict[Tuple[str, str, str, str], torch.nn.Module] = {}
|
|
|
|
|
|
def _list_models():
|
|
if not os.path.isdir(MODEL_DIR):
|
|
return []
|
|
out = []
|
|
for name in os.listdir(MODEL_DIR):
|
|
if name.lower().endswith(('.pth', '.pt', '.safetensors')):
|
|
out.append(name)
|
|
return sorted(out)
|
|
|
|
|
|
def _get_device():
|
|
if model_management is not None:
|
|
return model_management.get_torch_device()
|
|
return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
|
|
def _get_offload_device():
|
|
if model_management is not None:
|
|
return model_management.unet_offload_device()
|
|
return torch.device('cpu')
|
|
|
|
|
|
def _srformer_config():
|
|
return {
|
|
'type': 'SRFormer',
|
|
'upscale': 4,
|
|
'in_chans': 3,
|
|
'img_size': 64,
|
|
'window_size': 22,
|
|
'img_range': 1.0,
|
|
'depths': [6, 6, 6, 6, 6, 6],
|
|
'embed_dim': 180,
|
|
'num_heads': [6, 6, 6, 6, 6, 6],
|
|
'mlp_ratio': 2,
|
|
'upsampler': 'nearest+conv',
|
|
'resi_connection': '1conv',
|
|
}
|
|
|
|
|
|
def _esrgan_config():
|
|
return {
|
|
'type': 'RRDBNet',
|
|
'num_in_ch': 3,
|
|
'num_out_ch': 3,
|
|
'num_feat': 64,
|
|
'num_block': 23,
|
|
'num_grow_ch': 32,
|
|
'scale': 4,
|
|
}
|
|
|
|
|
|
def _select_config(arch: str):
|
|
# Inference only needs the final averaged PBaSR combinet branch.
|
|
# Loading a single branch avoids allocating general_branch + defocus_branch + combinet.
|
|
if arch == 'PBaSR_SRFormer_x4':
|
|
return _srformer_config()
|
|
if arch == 'PBaSR_ESRGAN_x4':
|
|
return _esrgan_config()
|
|
raise ValueError(f'Unsupported arch: {arch}')
|
|
|
|
|
|
def _strip_prefix_if_present(state, prefix):
|
|
keys = list(state.keys())
|
|
if not keys:
|
|
return None
|
|
if all(k.startswith(prefix) for k in keys):
|
|
return {k[len(prefix):]: v for k, v in state.items()}
|
|
selected = {k[len(prefix):]: v for k, v in state.items() if k.startswith(prefix)}
|
|
return selected if selected else None
|
|
|
|
|
|
def _branch_state_for_single_model(state):
|
|
"""Convert a PBaSR ParallelNet checkpoint to a single-branch checkpoint.
|
|
|
|
Preferred path: use combinet.* keys because PBaSR test-time uses the averaged combinet branch.
|
|
Fallback path: if only general_branch.* and defocus_branch.* exist, average matching tensors.
|
|
"""
|
|
# Remove DataParallel prefix first.
|
|
cleaned = {}
|
|
for k, v in state.items():
|
|
if k.startswith('module.'):
|
|
k = k[7:]
|
|
cleaned[k] = v
|
|
|
|
combinet = _strip_prefix_if_present(cleaned, 'combinet.')
|
|
if combinet:
|
|
return combinet
|
|
|
|
general = _strip_prefix_if_present(cleaned, 'general_branch.')
|
|
defocus = _strip_prefix_if_present(cleaned, 'defocus_branch.')
|
|
if general and defocus:
|
|
merged = {}
|
|
for k, gv in general.items():
|
|
dv = defocus.get(k)
|
|
if dv is not None and torch.is_tensor(gv) and torch.is_tensor(dv) and gv.shape == dv.shape and gv.is_floating_point():
|
|
merged[k] = gv * 0.5 + dv * 0.5
|
|
else:
|
|
merged[k] = gv
|
|
return merged
|
|
|
|
return cleaned
|
|
|
|
|
|
def _extract_state_dict(ckpt, param_key='params'):
|
|
"""Return a plain state_dict from common BasicSR/PBaSR checkpoint wrappers.
|
|
|
|
Supported formats:
|
|
- {'params': state_dict}
|
|
- {'params_ema': state_dict}
|
|
- {'state_dict': state_dict}
|
|
- raw state_dict
|
|
"""
|
|
if not isinstance(ckpt, dict):
|
|
raise TypeError(f'Unsupported checkpoint type: {type(ckpt)}')
|
|
|
|
candidate_keys = [
|
|
param_key,
|
|
'params_ema',
|
|
'params',
|
|
'state_dict',
|
|
'model',
|
|
'net',
|
|
'network_g',
|
|
]
|
|
|
|
seen = set()
|
|
for key in candidate_keys:
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
|
|
value = ckpt.get(key)
|
|
if isinstance(value, dict) and value:
|
|
return value
|
|
|
|
# If checkpoint itself is already a state_dict
|
|
if ckpt and all(isinstance(k, str) for k in ckpt.keys()):
|
|
tensor_like = [v for v in ckpt.values() if torch.is_tensor(v)]
|
|
if tensor_like:
|
|
return ckpt
|
|
|
|
raise KeyError(
|
|
'Could not find a state_dict in checkpoint. Tried keys: '
|
|
+ ', '.join(candidate_keys)
|
|
)
|
|
|
|
|
|
def _load_torch_file(path, device):
|
|
if path.lower().endswith('.safetensors'):
|
|
try:
|
|
from safetensors.torch import load_file
|
|
except Exception as e:
|
|
raise RuntimeError('safetensors is not installed, cannot load this checkpoint') from e
|
|
return load_file(path, device=str(device))
|
|
try:
|
|
# PyTorch 2.6+ defaults weights_only=True in some builds. Explicit False is more compatible
|
|
# with older BasicSR .pth wrappers, but only use trusted model files.
|
|
return torch.load(path, map_location=device, weights_only=False)
|
|
except TypeError:
|
|
return torch.load(path, map_location=device)
|
|
|
|
|
|
def _load_model(arch: str, model_name: str, precision: str, device):
|
|
model_path = os.path.join(MODEL_DIR, model_name)
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f'Model not found: {model_path}')
|
|
key = (arch, model_path, precision, str(device))
|
|
if key in _MODEL_CACHE:
|
|
model = _MODEL_CACHE[key]
|
|
model.to(device)
|
|
model.eval()
|
|
return model
|
|
|
|
model = build_network(_select_config(arch))
|
|
ckpt = _load_torch_file(model_path, device='cpu')
|
|
state = _extract_state_dict(ckpt, param_key='params')
|
|
state = _branch_state_for_single_model(state)
|
|
|
|
missing, unexpected = model.load_state_dict(state, strict=True)
|
|
if missing or unexpected:
|
|
print(f'[ComfyUI-PBaSR] missing={len(missing)}, unexpected={len(unexpected)}')
|
|
|
|
model.eval()
|
|
model.to(device)
|
|
if precision == 'fp16':
|
|
model.half()
|
|
elif precision == 'bf16':
|
|
model.bfloat16()
|
|
else:
|
|
model.float()
|
|
|
|
_MODEL_CACHE[key] = model
|
|
return model
|
|
|
|
|
|
def _mod_pad(x, scale=4, window_size=22):
|
|
# ParallelNet_ncb.test pads to multiples of 16 for x4, while SRFormer itself pads to window_size.
|
|
# For tiled inference, explicit padding keeps border cropping deterministic.
|
|
_, _, h, w = x.shape
|
|
base = 16
|
|
pad_h = (base - h % base) % base
|
|
pad_w = (base - w % base) % base
|
|
if pad_h or pad_w:
|
|
x = F.pad(x, (0, pad_w, 0, pad_h), mode='reflect')
|
|
return x, h, w
|
|
|
|
|
|
def _forward_full(model, x):
|
|
_, _, h_old, w_old = x.shape
|
|
base = 16
|
|
pad_h = (base - h_old % base) % base
|
|
pad_w = (base - w_old % base) % base
|
|
if pad_h or pad_w:
|
|
# Reflect padding matches SRFormer/BasicSR style better than zero padding,
|
|
# but reflect requires pad < input dimension. Use replicate for tiny tiles/tests.
|
|
mode = 'reflect' if pad_h < h_old and pad_w < w_old else 'replicate'
|
|
x = F.pad(x, (0, pad_w, 0, pad_h), mode=mode)
|
|
y = model(x)
|
|
return y[..., :h_old * 4, :w_old * 4]
|
|
|
|
|
|
def _forward_tiled(model, x, tile_size=512, tile_pad=64, scale=4):
|
|
b, c, h, w = x.shape
|
|
assert b == 1, 'Tiled inference currently expects batch size 1.'
|
|
if tile_size <= 0 or (h <= tile_size and w <= tile_size):
|
|
return _forward_full(model, x)
|
|
|
|
tile_size = int(tile_size)
|
|
tile_pad = int(tile_pad)
|
|
out = torch.zeros((b, c, h * scale, w * scale), device=x.device, dtype=x.dtype)
|
|
weight = torch.zeros((b, 1, h * scale, w * scale), device=x.device, dtype=x.dtype)
|
|
|
|
ys = list(range(0, h, tile_size))
|
|
xs = list(range(0, w, tile_size))
|
|
for y in ys:
|
|
for x0 in xs:
|
|
y1 = min(y + tile_size, h)
|
|
x1 = min(x0 + tile_size, w)
|
|
py0 = max(y - tile_pad, 0)
|
|
px0 = max(x0 - tile_pad, 0)
|
|
py1 = min(y1 + tile_pad, h)
|
|
px1 = min(x1 + tile_pad, w)
|
|
|
|
tile = x[:, :, py0:py1, px0:px1]
|
|
sr_tile = _forward_full(model, tile).clamp(0, 1)
|
|
|
|
inner_y0 = (y - py0) * scale
|
|
inner_x0 = (x0 - px0) * scale
|
|
inner_y1 = inner_y0 + (y1 - y) * scale
|
|
inner_x1 = inner_x0 + (x1 - x0) * scale
|
|
|
|
oy0, ox0 = y * scale, x0 * scale
|
|
oy1, ox1 = y1 * scale, x1 * scale
|
|
out[:, :, oy0:oy1, ox0:ox1] += sr_tile[:, :, inner_y0:inner_y1, inner_x0:inner_x1]
|
|
weight[:, :, oy0:oy1, ox0:ox1] += 1
|
|
|
|
return out / weight.clamp_min(1e-8)
|
|
|
|
|
|
class PBaSRUpscale:
|
|
CATEGORY = 'image/upscaling'
|
|
FUNCTION = 'upscale'
|
|
RETURN_TYPES = ('IMAGE',)
|
|
RETURN_NAMES = ('image',)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
models = _list_models()
|
|
if not models:
|
|
models = ['put_PBaSR_SRFormer_best.pth_or_PBaSR_ESRGAN_best.pth_in_ComfyUI_models_pbasr']
|
|
return {
|
|
'required': {
|
|
'image': ('IMAGE',),
|
|
'model_name': (models,),
|
|
'arch': (['PBaSR_SRFormer_x4', 'PBaSR_ESRGAN_x4'], {'default': 'PBaSR_SRFormer_x4'}),
|
|
'tile_size': ('INT', {'default': 512, 'min': 0, 'max': 2048, 'step': 32}),
|
|
'tile_pad': ('INT', {'default': 64, 'min': 0, 'max': 256, 'step': 8}),
|
|
'precision': (['fp16', 'bf16', 'fp32'], {'default': 'fp16'}),
|
|
'keep_model_loaded': ('BOOLEAN', {'default': True}),
|
|
}
|
|
}
|
|
|
|
def upscale(self, image, model_name, arch, tile_size, tile_pad, precision, keep_model_loaded):
|
|
if model_name.startswith('put_'):
|
|
raise FileNotFoundError(f'Put PBaSR weights into: {MODEL_DIR}')
|
|
device = _get_device()
|
|
if getattr(device, 'type', str(device)) == 'cpu' and precision != 'fp32':
|
|
print('[ComfyUI-PBaSR] CPU device detected; forcing fp32.')
|
|
precision = 'fp32'
|
|
model = _load_model(arch, model_name, precision, device)
|
|
model.eval()
|
|
|
|
x = image.to(device=device, dtype=torch.float32).permute(0, 3, 1, 2).contiguous().clamp(0, 1)
|
|
if precision == 'fp16':
|
|
x = x.half()
|
|
elif precision == 'bf16':
|
|
x = x.bfloat16()
|
|
|
|
outs = []
|
|
with torch.inference_mode():
|
|
for i in range(x.shape[0]):
|
|
xi = x[i:i + 1]
|
|
h0, w0 = xi.shape[-2:]
|
|
# Tiled inference naturally handles non-multiple input sizes, but per-tile model.test pads internally.
|
|
yi = _forward_tiled(model, xi, tile_size=tile_size, tile_pad=tile_pad, scale=4)
|
|
yi = yi[:, :, :h0 * 4, :w0 * 4]
|
|
outs.append(yi.float().clamp(0, 1))
|
|
y = torch.cat(outs, dim=0)
|
|
|
|
out = y.permute(0, 2, 3, 1).contiguous().detach().cpu()
|
|
|
|
if not keep_model_loaded:
|
|
try:
|
|
# remove only current model from cache
|
|
model_path = os.path.join(MODEL_DIR, model_name)
|
|
_MODEL_CACHE.pop((arch, model_path, precision, str(device)), None)
|
|
model.to(_get_offload_device())
|
|
except Exception:
|
|
pass
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
return (out,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
'PBaSRUpscale': PBaSRUpscale,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
'PBaSRUpscale': 'PBaSR Upscale',
|
|
}
|