blepping
aad095ab14
Other cleanups Allow setting a downscale factor and mode for CA downsampling Add more type annotations Make it so MSW-MSA attention failing is a warning rather than hard error Adjust a few Ruff lint settings
95 lines
2.7 KiB
Python
95 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib
|
|
|
|
import torch.nn.functional as torchf
|
|
from comfy.utils import bislerp
|
|
|
|
UPSCALE_METHODS = ("bicubic", "bislerp", "bilinear", "nearest-exact", "nearest", "area")
|
|
|
|
|
|
def parse_blocks(name: str, s: str) -> set:
|
|
vals = (rawval.strip() for rawval in s.split(","))
|
|
return {(name, int(val.strip())) for val in vals if val}
|
|
|
|
|
|
def convert_time(
|
|
ms: object,
|
|
time_mode: str,
|
|
start_time: float,
|
|
end_time: float,
|
|
) -> tuple:
|
|
if time_mode == "sigma":
|
|
return (start_time, end_time)
|
|
if time_mode in {"percent", "timestep"}:
|
|
if time_mode == "timestep":
|
|
start_time = 1.0 - (start_time / 999.0)
|
|
end_time = 1.0 - (end_time / 999.0)
|
|
else:
|
|
if start_time > 1.0 or start_time < 0.0:
|
|
raise ValueError(
|
|
"invalid value for start percent",
|
|
)
|
|
if end_time > 1.0 or end_time < 0.0:
|
|
raise ValueError(
|
|
"invalid value for end percent",
|
|
)
|
|
return (ms.percent_to_sigma(start_time), ms.percent_to_sigma(end_time))
|
|
raise ValueError("invalid time mode")
|
|
|
|
|
|
def get_sigma(options: dict, key: str = "sigmas") -> None | float:
|
|
if not isinstance(options, dict):
|
|
return None
|
|
sigmas = options.get(key)
|
|
if sigmas is None:
|
|
return None
|
|
if isinstance(sigmas, float):
|
|
return sigmas
|
|
return sigmas.detach().cpu().max().item()
|
|
|
|
|
|
def check_time(time_arg: dict | float, start_sigma: float, end_sigma: float) -> bool:
|
|
sigma = get_sigma(time_arg) if not isinstance(time_arg, float) else time_arg
|
|
if sigma is None:
|
|
return False
|
|
return sigma <= start_sigma and sigma >= end_sigma
|
|
|
|
|
|
try:
|
|
bleh = importlib.import_module("custom_nodes.ComfyUI-bleh")
|
|
bleh_latentutils = getattr(bleh.py, "latent_utils", None)
|
|
if bleh_latentutils is None:
|
|
raise ImportError # noqa: TRY301
|
|
bleh_version = getattr(bleh, "BLEH_VERSION", -1)
|
|
if bleh_version < 0:
|
|
|
|
def scale_samples(*args: list, sigma=None, **kwargs: dict): # noqa: ARG001
|
|
return bleh_latentutils.scale_samples(*args, **kwargs)
|
|
|
|
else:
|
|
scale_samples = bleh_latentutils.scale_samples
|
|
UPSCALE_METHODS = bleh_latentutils.UPSCALE_METHODS
|
|
except (ImportError, NotImplementedError):
|
|
|
|
def scale_samples(
|
|
samples,
|
|
width,
|
|
height,
|
|
mode="bicubic",
|
|
sigma=None, # noqa: ARG001
|
|
):
|
|
if mode == "bislerp":
|
|
return bislerp(samples, width, height)
|
|
return torchf.interpolate(samples, size=(height, width), mode=mode)
|
|
|
|
|
|
__all__ = (
|
|
"UPSCALE_METHODS",
|
|
"check_time",
|
|
"convert_time",
|
|
"get_sigma",
|
|
"parse_blocks",
|
|
"scale_samples",
|
|
)
|