diff --git a/modules/caption/joytag.py b/modules/caption/joytag.py
index 97a05b8b5..ac28e55a1 100644
--- a/modules/caption/joytag.py
+++ b/modules/caption/joytag.py
@@ -14,7 +14,7 @@ import torch.nn as nn
import torch.nn.functional as F
from transformers.activations import QuickGELUActivation
import torchvision
-import torchvision.transforms.functional as TVF
+from modules import images_sharpfin
import einops
from einops.layers.torch import Rearrange
import huggingface_hub
@@ -1035,8 +1035,8 @@ def prepare_image(image: Image.Image, target_size: int) -> torch.Tensor:
padded_image.paste(image, (pad_left, pad_top))
if max_dim != target_size:
padded_image = padded_image.resize((target_size, target_size), Image.Resampling.LANCZOS)
- image_tensor = TVF.pil_to_tensor(padded_image) / 255.0
- image_tensor = TVF.normalize(image_tensor, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])
+ image_tensor = images_sharpfin.to_tensor(padded_image)
+ image_tensor = images_sharpfin.normalize(image_tensor, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])
return image_tensor
diff --git a/modules/images_resize.py b/modules/images_resize.py
index bbe25ba68..9fd8757bf 100644
--- a/modules/images_resize.py
+++ b/modules/images_resize.py
@@ -4,7 +4,7 @@ import time
import numpy as np
import torch
from PIL import Image
-from modules import shared, upscaler
+from modules import shared, upscaler, images_sharpfin
def resize_image(resize_mode: int, im: Union[Image.Image, torch.Tensor], width: int, height: int, upscaler_name: str=None, output_type: str='image', context: str=None):
@@ -36,7 +36,7 @@ def resize_image(resize_mode: int, im: Union[Image.Image, torch.Tensor], width:
def resize(im: Union[Image.Image, torch.Tensor], w, h):
w, h = int(w), int(h)
if upscaler_name is None or upscaler_name == "None" or (hasattr(im, 'mode') and im.mode == 'L'):
- return im.resize((w, h), resample=Image.Resampling.LANCZOS) # force for mask
+ return images_sharpfin.resize(im, (w, h), linearize=False) # force for mask
if isinstance(im, torch.Tensor):
scale = max(w // 8 / im.shape[-1] , h // 8 / im.shape[-2])
else:
@@ -53,7 +53,7 @@ def resize_image(resize_mode: int, im: Union[Image.Image, torch.Tensor], width:
shared.log.warning(f"Resize upscaler: invalid={upscaler_name} fallback={selected_upscaler.name}")
shared.log.debug(f"Resize upscaler: available={[u.name for u in shared.sd_upscalers]}")
if isinstance(im, Image.Image) and (im.width != w or im.height != h): # probably downsample after upscaler created larger image
- im = im.resize((w, h), resample=Image.Resampling.LANCZOS)
+ im = images_sharpfin.resize(im, (w, h))
return im
def crop(im: Image.Image):
diff --git a/modules/images_sharpfin.py b/modules/images_sharpfin.py
new file mode 100644
index 000000000..3a303a083
--- /dev/null
+++ b/modules/images_sharpfin.py
@@ -0,0 +1,298 @@
+"""Sharpfin wrapper for high-quality image resize and tensor conversion.
+
+Provides drop-in replacements for torchvision.transforms.functional operations
+with higher quality resampling (Magic Kernel Sharp 2021), sRGB linearization,
+and Triton GPU acceleration when available.
+
+All public functions include try/except fallback to PIL/torchvision.
+"""
+
+import torch
+import numpy as np
+from PIL import Image
+
+_sharpfin_checked = False
+_sharpfin_ok = False
+_triton_ok = False
+_log = None
+
+
+def _get_log():
+ global _log
+ if _log is None:
+ try:
+ from modules.shared import log
+ _log = log
+ except Exception:
+ import logging
+ _log = logging.getLogger(__name__)
+ return _log
+
+
+def _check():
+ global _sharpfin_checked, _sharpfin_ok, _triton_ok
+ if not _sharpfin_checked:
+ # DEBUG: no try/except — let import errors propagate
+ from modules.sharpfin.functional import scale # pylint: disable=unused-import
+ _sharpfin_ok = True
+ try:
+ from modules.sharpfin import TRITON_AVAILABLE
+ _triton_ok = TRITON_AVAILABLE
+ except Exception:
+ _triton_ok = False
+ _sharpfin_checked = True
+
+
+def is_available():
+ """Check if sharpfin functional module loaded."""
+ _check()
+ return _sharpfin_ok
+
+
+KERNEL_MAP = {
+ "Sharpfin MKS2021": "MAGIC_KERNEL_SHARP_2021",
+ "Sharpfin Lanczos3": "LANCZOS3",
+ "Sharpfin Mitchell": "MITCHELL",
+ "Sharpfin Catmull-Rom": "CATMULL_ROM",
+}
+
+
+def _resolve_kernel(kernel=None):
+ """Resolve kernel name to ResizeKernel enum. Returns None for PIL fallback."""
+ if kernel is not None:
+ name = kernel
+ else:
+ try:
+ from modules import shared
+ name = getattr(shared.opts, 'resize_quality', 'Sharpfin MKS2021')
+ except Exception:
+ name = 'Sharpfin MKS2021'
+ if name == "PIL Lanczos" or name not in KERNEL_MAP:
+ return None
+ from modules.sharpfin.util import ResizeKernel
+ return getattr(ResizeKernel, KERNEL_MAP[name])
+
+
+def _resolve_linearize(linearize=None, is_mask=False):
+ """Determine sRGB linearization setting."""
+ if is_mask:
+ return False
+ if linearize is not None:
+ return linearize
+ try:
+ from modules import shared
+ return getattr(shared.opts, 'resize_linearize_srgb', True)
+ except Exception:
+ return True
+
+
+def _get_device_dtype(device=None, dtype=None):
+ """Get optimal device/dtype for sharpfin operations."""
+ if device is not None and dtype is not None:
+ return device, dtype
+ try:
+ from modules import devices
+ dev = device or devices.device
+ if dev.type == 'cuda':
+ return dev, dtype or torch.float16
+ return dev, dtype or torch.float32
+ except Exception:
+ return device or torch.device('cpu'), dtype or torch.float32
+
+
+def resize(image, target_size, *, kernel=None, linearize=None, device=None, dtype=None):
+ """Resize PIL.Image or torch.Tensor, returning same type.
+
+ Args:
+ image: PIL.Image or torch.Tensor [B,C,H,W] / [C,H,W]
+ target_size: (width, height) for PIL, (H, W) for tensor
+ kernel: Override kernel name, or None for settings
+ linearize: Override sRGB linearization, or None for settings
+ device: Override compute device
+ dtype: Override compute dtype
+ """
+ _check()
+ if isinstance(image, Image.Image):
+ return _resize_pil(image, target_size, kernel=kernel, linearize=linearize, device=device, dtype=dtype)
+ elif isinstance(image, torch.Tensor):
+ return resize_tensor(image, target_size, kernel=kernel, linearize=linearize if linearize is not None else False)
+ return image
+
+
+def _want_sparse(dev, rk, both_down):
+ """Check if Triton sparse acceleration should be attempted."""
+ return _triton_ok and dev.type == 'cuda' and rk.value == 'magic_kernel_sharp_2021' and both_down
+
+
+def _scale_pil(scale_fn, tensor, out_res, rk, dev, dt, do_linear, src_h, src_w, h, w, both_down, both_up):
+ """Run sharpfin scale with sparse fallback. Returns result tensor."""
+ global _triton_ok # pylint: disable=global-statement
+ if both_down or both_up:
+ use_sparse = _want_sparse(dev, rk, both_down)
+ if use_sparse:
+ try:
+ return scale_fn(tensor, out_res, resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=True)
+ except Exception:
+ _triton_ok = False
+ _get_log().info("Sharpfin: Triton sparse disabled, using dense path")
+ return scale_fn(tensor, out_res, resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
+ # Mixed axis: split into two single-axis resizes
+ if h > src_h: # H up, W down
+ intermediate = scale_fn(tensor, (h, src_w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
+ use_sparse = _want_sparse(dev, rk, True)
+ if use_sparse:
+ try:
+ return scale_fn(intermediate, (h, w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=True)
+ except Exception:
+ _triton_ok = False
+ _get_log().info("Sharpfin: Triton sparse disabled, using dense path")
+ return scale_fn(intermediate, (h, w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
+ # H down, W up
+ use_sparse = _want_sparse(dev, rk, True)
+ if use_sparse:
+ try:
+ intermediate = scale_fn(tensor, (h, src_w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=True)
+ return scale_fn(intermediate, (h, w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
+ except Exception:
+ _triton_ok = False
+ _get_log().info("Sharpfin: Triton sparse disabled, using dense path")
+ intermediate = scale_fn(tensor, (h, src_w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
+ return scale_fn(intermediate, (h, w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
+
+
+def _resize_pil(image, target_size, *, kernel=None, linearize=None, device=None, dtype=None):
+ """Resize a PIL Image via sharpfin, falling back to PIL on error."""
+ w, h = target_size
+ if image.width == w and image.height == h:
+ return image
+ is_mask = image.mode == 'L'
+ rk = _resolve_kernel(kernel)
+ if rk is None:
+ # DEBUG: only "PIL Lanczos" setting should reach here
+ assert _resolve_kernel.__doc__, "unreachable" # keeps linter happy
+ return image.resize((w, h), resample=Image.Resampling.LANCZOS)
+ try:
+ from modules.sharpfin.functional import scale
+ do_linear = _resolve_linearize(linearize, is_mask=is_mask)
+ dev, dt = _get_device_dtype(device, dtype)
+ tensor = to_tensor(image)
+ if tensor.dim() == 3:
+ tensor = tensor.unsqueeze(0)
+ tensor = tensor.to(device=dev, dtype=dt)
+ out_res = (h, w) # sharpfin uses (H, W)
+ src_h, src_w = tensor.shape[-2], tensor.shape[-1]
+ both_down = (h <= src_h and w <= src_w)
+ both_up = (h >= src_h and w >= src_w)
+ result = _scale_pil(scale, tensor, out_res, rk, dev, dt, do_linear, src_h, src_w, h, w, both_down, both_up)
+ return to_pil(result)
+ # except Exception as e: # DEBUG: PIL fallback disabled for testing
+ # _get_log().warning(f"Sharpfin resize failed, falling back to PIL: {e}")
+ # return image.resize((w, h), resample=Image.Resampling.LANCZOS)
+ finally:
+ pass
+
+
+def resize_tensor(tensor, target_size, *, kernel=None, linearize=False):
+ """Resize tensor [B,C,H,W] or [C,H,W] -> Tensor. For in-pipeline tensor resizes.
+
+ Args:
+ tensor: Input tensor
+ target_size: (H, W) tuple
+ kernel: Override kernel name
+ linearize: sRGB linearization (default False for latent/mask data)
+ """
+ _check()
+ rk = _resolve_kernel(kernel)
+ if rk is None:
+ # DEBUG: only "PIL Lanczos" setting should reach here
+ mode = 'bilinear' if target_size[0] * target_size[1] > tensor.shape[-2] * tensor.shape[-1] else 'area'
+ return torch.nn.functional.interpolate(tensor if tensor.dim() == 4 else tensor.unsqueeze(0), size=target_size, mode=mode, antialias=True).squeeze(0) if tensor.dim() == 3 else torch.nn.functional.interpolate(tensor, size=target_size, mode=mode, antialias=True)
+ try:
+ from modules.sharpfin.functional import scale
+ dev, dt = _get_device_dtype()
+ squeezed = False
+ if tensor.dim() == 3:
+ tensor = tensor.unsqueeze(0)
+ squeezed = True
+ src_h, src_w = tensor.shape[-2], tensor.shape[-1]
+ th, tw = target_size
+ both_down = (th <= src_h and tw <= src_w)
+ both_up = (th >= src_h and tw >= src_w)
+ if both_down or both_up:
+ use_sparse = _triton_ok and dev.type == 'cuda' and rk.value == 'magic_kernel_sharp_2021' and both_down
+ result = scale(tensor, target_size, resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=linearize, use_sparse=use_sparse)
+ else:
+ if th > src_h:
+ intermediate = scale(tensor, (th, src_w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=linearize, use_sparse=False)
+ result = scale(intermediate, (th, tw), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=linearize, use_sparse=False)
+ else:
+ intermediate = scale(tensor, (th, src_w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=linearize, use_sparse=False)
+ result = scale(intermediate, (th, tw), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=linearize, use_sparse=False)
+ if squeezed:
+ result = result.squeeze(0)
+ return result
+ # except Exception as e: # DEBUG: F.interpolate fallback disabled for testing
+ # _get_log().warning(f"Sharpfin resize_tensor failed, falling back to F.interpolate: {e}")
+ # mode = 'bilinear' if target_size[0] * target_size[1] > tensor.shape[-2] * tensor.shape[-1] else 'area'
+ # inp = tensor if tensor.dim() == 4 else tensor.unsqueeze(0)
+ # result = torch.nn.functional.interpolate(inp, size=target_size, mode=mode, antialias=True)
+ # return result.squeeze(0) if tensor.dim() == 3 else result
+ finally:
+ pass
+
+
+def to_tensor(image):
+ """PIL Image -> float32 CHW tensor [0,1]. Pure torch, no torchvision."""
+ if not isinstance(image, Image.Image):
+ raise TypeError(f"Expected PIL Image, got {type(image)}")
+ pic = np.array(image, copy=True)
+ if pic.ndim == 2:
+ pic = pic[:, :, np.newaxis]
+ tensor = torch.from_numpy(pic.transpose((2, 0, 1))).contiguous()
+ if tensor.dtype == torch.uint8:
+ return tensor.to(torch.float32).div_(255.0)
+ return tensor.to(torch.float32)
+
+
+def to_pil(tensor):
+ """Float CHW/HWC or BCHW/BHWC tensor [0,1] -> PIL Image. Pure torch, no torchvision."""
+ if not isinstance(tensor, torch.Tensor):
+ raise TypeError(f"Expected torch.Tensor, got {type(tensor)}")
+ tensor = tensor.detach().cpu()
+ if tensor.dim() == 4:
+ if tensor.shape[-1] in (1, 3, 4) and tensor.shape[-1] < tensor.shape[-2]: # BHWC
+ tensor = tensor.permute(0, 3, 1, 2)
+ tensor = tensor[0]
+ elif tensor.dim() == 3:
+ if tensor.shape[-1] in (1, 3, 4) and tensor.shape[-1] < tensor.shape[-2] and tensor.shape[-1] < tensor.shape[-3]: # HWC
+ tensor = tensor.permute(2, 0, 1)
+ if tensor.dtype != torch.uint8:
+ tensor = (tensor.clamp(0, 1) * 255).round().to(torch.uint8)
+ ndarr = tensor.permute(1, 2, 0).numpy()
+ if ndarr.shape[2] == 1:
+ return Image.fromarray(ndarr[:, :, 0], mode='L')
+ return Image.fromarray(ndarr)
+
+
+def pil_to_tensor(image):
+ """PIL Image -> uint8 CHW tensor (no float scaling). Replaces TF.pil_to_tensor."""
+ if not isinstance(image, Image.Image):
+ raise TypeError(f"Expected PIL Image, got {type(image)}")
+ pic = np.array(image, copy=True)
+ if pic.ndim == 2:
+ pic = pic[:, :, np.newaxis]
+ return torch.from_numpy(pic.transpose((2, 0, 1))).contiguous()
+
+
+def normalize(tensor, mean, std, inplace=False):
+ """Tensor normalization. Replaces TF.normalize."""
+ if not inplace:
+ tensor = tensor.clone()
+ mean_t = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device)
+ std_t = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device)
+ if mean_t.ndim == 1:
+ mean_t = mean_t[:, None, None]
+ if std_t.ndim == 1:
+ std_t = std_t[:, None, None]
+ tensor.sub_(mean_t).div_(std_t)
+ return tensor
diff --git a/modules/lama.py b/modules/lama.py
index 24d1f8fa3..32029da85 100644
--- a/modules/lama.py
+++ b/modules/lama.py
@@ -5,7 +5,7 @@ import torch
import numpy as np
from torch.hub import download_url_to_file, get_dir
from PIL import Image
-from modules import devices
+from modules import devices, images_sharpfin
from installer import log
@@ -96,7 +96,5 @@ class SimpleLama:
image, mask = prepare_img_and_mask(image, mask, self.device)
with devices.inference_context():
inpainted = self.model(image, mask)
- cur_res = inpainted[0].permute(1, 2, 0).detach().float().cpu().numpy()
- cur_res = np.clip(cur_res * 255, 0, 255).astype(np.uint8)
- cur_res = Image.fromarray(cur_res)
+ cur_res = images_sharpfin.to_pil(inpainted[0])
return cur_res
diff --git a/modules/postprocess/codeformer_model.py b/modules/postprocess/codeformer_model.py
index ad83e1d0c..60e7ef3b6 100644
--- a/modules/postprocess/codeformer_model.py
+++ b/modules/postprocess/codeformer_model.py
@@ -70,7 +70,7 @@ def setup_model(dirname):
self.face_helper.face_parse.to(device)
def restore(self, np_image, p=None, w=None): # pylint: disable=unused-argument
- from torchvision.transforms.functional import normalize
+ from modules import images_sharpfin
from basicsr.utils import img2tensor, tensor2img
np_image = np_image[:, :, ::-1]
original_resolution = np_image.shape[0:2]
@@ -84,7 +84,7 @@ def setup_model(dirname):
self.face_helper.align_warp_face()
for cropped_face in self.face_helper.cropped_faces:
cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
- normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
+ images_sharpfin.normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
cropped_face_t = cropped_face_t.unsqueeze(0).to(devices.device)
try:
with devices.inference_context():
diff --git a/modules/postprocess/pixelart.py b/modules/postprocess/pixelart.py
index 6b027011e..8b1ef1254 100644
--- a/modules/postprocess/pixelart.py
+++ b/modules/postprocess/pixelart.py
@@ -2,8 +2,8 @@ from typing import List
import math
import torch
-import torchvision
import numpy as np
+from modules import images_sharpfin
from PIL import Image
from diffusers.utils import CONFIG_NAME
@@ -65,11 +65,9 @@ def edge_detect_for_pixelart(image: PipelineImageInput, image_weight: float = 1.
greyscale_reshaped = greyscale_reshaped.reshape(batch_size, block_size_sq, block_height, block_width)
greyscale_range = greyscale_reshaped.amax(dim=1, keepdim=True).sub_(greyscale_reshaped.amin(dim=1, keepdim=True))
- upsample = torchvision.transforms.Resize((height, width), interpolation=torchvision.transforms.InterpolationMode.BICUBIC)
-
- range_weight = upsample(greyscale_range)
+ range_weight = images_sharpfin.resize_tensor(greyscale_range, (height, width), linearize=False)
range_weight = range_weight.div_(range_weight.max())
- weight_map = upsample((greyscale > greyscale.median()).to(dtype=torch.float32))
+ weight_map = images_sharpfin.resize_tensor((greyscale > greyscale.median()).to(dtype=torch.float32), (height, width), linearize=False)
weight_map = weight_map.unsqueeze(0).add_(range_weight).mul_(image_weight / 2)
new_image = new_image.mul_(weight_map).addcmul_(min_pool, (1-weight_map))
@@ -161,8 +159,7 @@ def encode_jpeg_tensor(img: torch.FloatTensor, block_size: int=16, cbcr_downscal
img = img[:, :, :(img.shape[-2]//block_size)*block_size, :(img.shape[-1]//block_size)*block_size] # crop to a multiply of block_size
cbcr_block_size = block_size//cbcr_downscale
_, _, height, width = img.shape
- downsample = torchvision.transforms.Resize((height//cbcr_downscale, width//cbcr_downscale), interpolation=torchvision.transforms.InterpolationMode.BICUBIC)
- down_img = downsample(img[:, 1:,:,:])
+ down_img = images_sharpfin.resize_tensor(img[:, 1:,:,:], (height//cbcr_downscale, width//cbcr_downscale), linearize=False)
y = encode_single_channel_dct_2d(img[:, 0, :,:], block_size=block_size, norm=norm)
cb = encode_single_channel_dct_2d(down_img[:, 0, :,:], block_size=cbcr_block_size, norm=norm)
cr = encode_single_channel_dct_2d(down_img[:, 1, :,:], block_size=cbcr_block_size, norm=norm)
@@ -180,9 +177,8 @@ def decode_jpeg_tensor(jpeg_img: torch.FloatTensor, block_size: int=16, cbcr_dow
y = decode_single_channel_dct_2d(y, norm=norm)
cb = decode_single_channel_dct_2d(cb, norm=norm)
cr = decode_single_channel_dct_2d(cr, norm=norm)
- upsample = torchvision.transforms.Resize((h_blocks*block_size, w_blocks*block_size), interpolation=torchvision.transforms.InterpolationMode.BICUBIC)
- cb = upsample(cb)
- cr = upsample(cr)
+ cb = images_sharpfin.resize_tensor(cb, (h_blocks*block_size, w_blocks*block_size), linearize=False)
+ cr = images_sharpfin.resize_tensor(cr, (h_blocks*block_size, w_blocks*block_size), linearize=False)
return torch.stack([y,cb,cr], dim=1)
diff --git a/modules/postprocess/seedvr_model.py b/modules/postprocess/seedvr_model.py
index 577d8609f..b78c174bc 100644
--- a/modules/postprocess/seedvr_model.py
+++ b/modules/postprocess/seedvr_model.py
@@ -3,8 +3,7 @@ import random
import numpy as np
import torch
from PIL import Image
-from torchvision.transforms import ToPILImage
-from modules import devices
+from modules import devices, images_sharpfin
from modules.shared import opts, log
from modules.upscaler import Upscaler, UpscalerData
@@ -14,7 +13,7 @@ MODELS_MAP = {
"SeedVR2 7B": "seedvr2_ema_7b_fp16.safetensors",
"SeedVR2 7B Sharp": "seedvr2_ema_7b_sharp_fp16.safetensors",
}
-to_pil = ToPILImage()
+to_pil = images_sharpfin.to_pil
class UpscalerSeedVR(Upscaler):
@@ -159,7 +158,7 @@ class UpscalerSeedVR(Upscaler):
)
t1 = time.time()
log.info(f'Upscaler: type="{self.name}" model="{selected_file}" scale={self.scale} cfg={opts.seedvt_cfg_scale} seed={seed} time={t1 - t0:.2f}')
- img = to_pil(result_tensor.squeeze().permute((2, 0, 1)))
+ img = to_pil(result_tensor.squeeze())
if opts.upscaler_unload:
self.model.dit = None
diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py
index 04ae68ea2..e559f3ce0 100644
--- a/modules/processing_diffusers.py
+++ b/modules/processing_diffusers.py
@@ -3,9 +3,8 @@ import os
import time
import numpy as np
import torch
-import torchvision.transforms.functional as TF
from PIL import Image
-from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, timer, modelstats, extra_networks, attention
+from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, timer, modelstats, extra_networks, attention, images_sharpfin
from modules.processing_helpers import resize_hires, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, save_intermediate, update_sampler, is_txt2img, is_refiner_enabled, get_job_name
from modules.processing_args import set_pipeline_args
from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed
@@ -270,9 +269,9 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
sd_hijack_hypertile.hypertile_set(p, hr=True)
elif torch.is_tensor(output.images) and output.images.shape[-1] == 3: # nhwc
if output.images.dim() == 3:
- output.images = TF.to_pil_image(output.images.permute(2,0,1))
+ output.images = images_sharpfin.to_pil(output.images)
elif output.images.dim() == 4:
- output.images = [TF.to_pil_image(output.images[i].permute(2,0,1)) for i in range(output.images.shape[0])]
+ output.images = [images_sharpfin.to_pil(output.images[i]) for i in range(output.images.shape[0])]
strength = p.hr_denoising_strength if p.hr_denoising_strength > 0 else p.denoising_strength
if (p.hr_upscaler is not None) and (p.hr_upscaler.lower().startswith('latent') or p.hr_force) and strength > 0:
@@ -572,7 +571,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
if hasattr(shared.sd_model, 'unet') and hasattr(shared.sd_model.unet, 'config') and hasattr(shared.sd_model.unet.config, 'in_channels') and shared.sd_model.unet.config.in_channels == 9 and not is_control:
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.INPAINTING) # force pipeline
if len(getattr(p, 'init_images', [])) == 0:
- p.init_images = [TF.to_pil_image(torch.rand((3, getattr(p, 'height', 512), getattr(p, 'width', 512))))]
+ p.init_images = [images_sharpfin.to_pil(torch.rand((3, getattr(p, 'height', 512), getattr(p, 'width', 512))))]
if not p.prompts:
p.prompts = p.all_prompts[p.iteration * p.batch_size:(p.iteration+1) * p.batch_size]
if not p.negative_prompts:
diff --git a/modules/processing_vae.py b/modules/processing_vae.py
index 72c385ac5..77da7ed57 100644
--- a/modules/processing_vae.py
+++ b/modules/processing_vae.py
@@ -334,7 +334,7 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he
def vae_encode(image, model, vae_type='Full'): # pylint: disable=unused-variable
jobid = shared.state.begin('VAE Encode')
- import torchvision.transforms.functional as f
+ from modules import images_sharpfin
if shared.state.interrupted or shared.state.skipped:
return []
if not hasattr(model, 'vae') and hasattr(model, 'pipe'):
@@ -342,7 +342,7 @@ def vae_encode(image, model, vae_type='Full'): # pylint: disable=unused-variable
if not hasattr(model, 'vae'):
shared.log.error('VAE not found in model')
return []
- tensor = f.to_tensor(image.convert("RGB")).unsqueeze(0).to(devices.device, devices.dtype_vae)
+ tensor = images_sharpfin.to_tensor(image.convert("RGB")).unsqueeze(0).to(devices.device, devices.dtype_vae)
if vae_type == 'Tiny':
latents = taesd_vae_encode(image=tensor)
elif vae_type == 'Full' and hasattr(model, 'vae'):
diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index c49dcf11d..6e7358be8 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -2,9 +2,8 @@ import time
import threading
from collections import namedtuple
import torch
-import torchvision.transforms.functional as TF
from PIL import Image
-from modules import shared, devices, processing, images, sd_samplers, timer
+from modules import shared, devices, processing, images, sd_samplers, timer, images_sharpfin
from modules.vae import sd_vae_approx, sd_vae_taesd, sd_vae_stablecascade
@@ -84,7 +83,7 @@ def single_sample_to_image(sample, approximation=None):
x_sample = (255.0 * x_sample).to(torch.uint8)
if len(x_sample.shape) == 4:
x_sample = x_sample[0]
- image = TF.to_pil_image(x_sample)
+ image = images_sharpfin.to_pil(x_sample)
except Exception as e:
warn_once(f'Preview: {e}')
image = Image.new(mode="RGB", size=(512, 512))
diff --git a/modules/shared.py b/modules/shared.py
index fb3b919fa..5893ca980 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -671,6 +671,10 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), {
"upscaler_latent_steps": OptionInfo(20, "Upscaler latent steps", gr.Slider, {"minimum": 4, "maximum": 100, "step": 1}),
"upscaler_tile_size": OptionInfo(192, "Upscaler tile size", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}),
"upscaler_tile_overlap": OptionInfo(8, "Upscaler tile overlap", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}),
+
+ "postprocessing_sep_resize": OptionInfo("
Resize
", "", gr.HTML),
+ "resize_quality": OptionInfo("Sharpfin MKS2021", "Image resize algorithm", gr.Dropdown, {"choices": ["PIL Lanczos", "Sharpfin MKS2021", "Sharpfin Lanczos3", "Sharpfin Mitchell", "Sharpfin Catmull-Rom"]}),
+ "resize_linearize_srgb": OptionInfo(True, "Apply sRGB linearization during image resize"),
}))
diff --git a/modules/sharpfin/LICENSE b/modules/sharpfin/LICENSE
new file mode 100644
index 000000000..c845fba18
--- /dev/null
+++ b/modules/sharpfin/LICENSE
@@ -0,0 +1,190 @@
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ Copyright 2024 drhead
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/modules/sharpfin/__init__.py b/modules/sharpfin/__init__.py
new file mode 100644
index 000000000..535430d0f
--- /dev/null
+++ b/modules/sharpfin/__init__.py
@@ -0,0 +1,20 @@
+"""Sharpfin - High quality image resizing with GPU acceleration.
+
+Vendored from https://github.com/drhead/sharpfin (Apache 2.0)
+Provides Magic Kernel Sharp 2021 resampling, sRGB linearization,
+and Triton sparse GPU acceleration.
+"""
+
+from .util import ResizeKernel, SharpenKernel, QuantHandling, srgb_to_linear, linear_to_srgb
+
+try:
+ from .functional import scale, _upscale, _downscale, _get_resize_kernel
+ FUNCTIONAL_AVAILABLE = True
+except Exception:
+ FUNCTIONAL_AVAILABLE = False
+
+try:
+ from .triton_functional import downscale_sparse
+ TRITON_AVAILABLE = True
+except Exception:
+ TRITON_AVAILABLE = False
diff --git a/modules/sharpfin/cms.py b/modules/sharpfin/cms.py
new file mode 100644
index 000000000..18f88bfe3
--- /dev/null
+++ b/modules/sharpfin/cms.py
@@ -0,0 +1,174 @@
+"""Sharpfin color management (ICC profile handling).
+
+Vendored from https://github.com/drhead/sharpfin (Apache 2.0)
+"""
+
+from io import BytesIO
+from typing import Any, cast
+from warnings import warn
+
+import numpy as np
+from torch import Tensor
+
+import PIL.Image as image
+import PIL.ImageCms as image_cms
+
+from PIL.Image import Image
+from PIL.ImageCms import (
+ Direction, Intent, ImageCmsProfile, PyCMSError,
+ createProfile, getDefaultIntent, isIntentSupported, profileToProfile
+)
+from PIL.ImageOps import exif_transpose
+
+image.MAX_IMAGE_PIXELS = None
+
+_SRGB = createProfile(colorSpace='sRGB')
+
+_INTENT_FLAGS = {
+ Intent.PERCEPTUAL: image_cms.FLAGS["HIGHRESPRECALC"],
+ Intent.RELATIVE_COLORIMETRIC: (
+ image_cms.FLAGS["HIGHRESPRECALC"] |
+ image_cms.FLAGS["BLACKPOINTCOMPENSATION"]
+ ),
+ Intent.ABSOLUTE_COLORIMETRIC: image_cms.FLAGS["HIGHRESPRECALC"]
+}
+
+class CMSWarning(UserWarning):
+ def __init__(
+ self,
+ message: str,
+ *,
+ path: str | None = None,
+ cms_info: dict[str, Any] | None = None,
+ cause: Exception | None = None,
+ ):
+ super().__init__(message)
+ self.__cause__ = cause
+
+ self.path = path
+ self.cms_info = cms_info
+
+def _coalesce_intent(intent: Intent | int) -> Intent:
+ if isinstance(intent, Intent):
+ return intent
+
+ match intent:
+ case 0:
+ return Intent.PERCEPTUAL
+ case 1:
+ return Intent.RELATIVE_COLORIMETRIC
+ case 2:
+ return Intent.SATURATION
+ case 3:
+ return Intent.ABSOLUTE_COLORIMETRIC
+ case _:
+ raise ValueError("invalid intent")
+
+def _add_info(info: dict[str, Any], source: object, key: str) -> None:
+ try:
+ if (value := getattr(source, key, None)) is not None:
+ info[key] = value
+ except Exception:
+ pass
+
+def apply_srgb(
+ img: Image
+) -> Image:
+ if hasattr(img, 'filename'):
+ path = img.filename
+ else:
+ path = ""
+
+ try:
+ img.load()
+
+ try:
+ exif_transpose(img, in_place=True)
+ except Exception:
+ pass # corrupt EXIF metadata is fine
+
+ if (icc_raw := img.info.get("icc_profile")) is not None:
+ cms_info: dict[str, Any] = {
+ "native_mode": img.mode,
+ "transparency": img.has_transparency_data,
+ }
+
+ try:
+ profile = ImageCmsProfile(BytesIO(icc_raw))
+ _add_info(cms_info, profile.profile, "profile_description")
+ _add_info(cms_info, profile.profile, "target")
+ _add_info(cms_info, profile.profile, "xcolor_space")
+ _add_info(cms_info, profile.profile, "connection_space")
+ _add_info(cms_info, profile.profile, "colorimetric_intent")
+ _add_info(cms_info, profile.profile, "rendering_intent")
+
+ working_mode = img.mode
+ if img.mode.startswith(("RGB", "BGR", "P")):
+ working_mode = "RGBA" if img.has_transparency_data else "RGB"
+ elif img.mode.startswith(("L", "I", "F")) or img.mode == "1":
+ working_mode = "LA" if img.has_transparency_data else "L"
+
+ if img.mode != working_mode:
+ cms_info["working_mode"] = working_mode
+ img = img.convert(working_mode)
+
+ mode = "RGBA" if img.has_transparency_data else "RGB"
+
+ intent = Intent.RELATIVE_COLORIMETRIC
+ if isIntentSupported(profile, intent, Direction.INPUT) != 1:
+ intent = _coalesce_intent(getDefaultIntent(profile))
+
+ cms_info["conversion_intent"] = intent
+
+ if (flags := _INTENT_FLAGS.get(intent)) is not None:
+ if img.mode == mode:
+ profileToProfile(
+ img,
+ profile,
+ _SRGB,
+ renderingIntent=intent,
+ inPlace=True,
+ flags=flags
+ )
+ else:
+ img = cast(Image, profileToProfile(
+ img,
+ profile,
+ _SRGB,
+ renderingIntent=intent,
+ outputMode=mode,
+ flags=flags
+ ))
+ else:
+ warn(CMSWarning(
+ f"unsupported intent on {path} assuming sRGB: {cms_info}",
+ path=path,
+ cms_info=cms_info
+ ))
+ except PyCMSError as ex:
+ warn(CMSWarning(
+ f"{ex} on {path}, assuming sRGB: {cms_info}",
+ path=path,
+ cms_info=cms_info,
+ cause=ex,
+ ))
+
+ except Exception as ex:
+ print(f"{ex} on {path}")
+
+ if img.has_transparency_data:
+ if img.mode != "RGBA":
+ try:
+ img = img.convert("RGBA")
+ except ValueError:
+ img = img.convert("RGBa").convert("RGBA")
+ elif img.mode != "RGB":
+ img = img.convert("RGB")
+
+ return img
+
+def put_srgb(img: Image, tensor: Tensor) -> None:
+ if img.mode not in ("RGB", "RGBA", "RGBa"):
+ raise ValueError(f"Image has non-RGB mode {img.mode}.")
+
+ np.copyto(tensor.numpy(), np.asarray(img)[:, :, :3], casting="no")
diff --git a/modules/sharpfin/functional.py b/modules/sharpfin/functional.py
new file mode 100644
index 000000000..c70c87719
--- /dev/null
+++ b/modules/sharpfin/functional.py
@@ -0,0 +1,285 @@
+"""Sharpfin functional image scaling operations.
+
+Vendored from https://github.com/drhead/sharpfin (Apache 2.0)
+Imports patched: absolute sharpfin.X → relative .X, triton import guarded.
+"""
+
+import torch
+import numpy as np
+import torch.nn.functional as F
+from typing import Callable, Tuple
+import math
+from contextlib import nullcontext
+
+from .util import ResizeKernel, linear_to_srgb, srgb_to_linear
+
+# from Pytorch >= 2.6
+set_stance = getattr(torch.compiler, "set_stance", None)
+
+
+def _get_resize_kernel(k: ResizeKernel):
+ match k:
+ case ResizeKernel.NEAREST:
+ resize_kernel = nearest
+ kernel_window = 0.5
+ case ResizeKernel.BILINEAR:
+ resize_kernel = bilinear
+ kernel_window = 1.
+ case ResizeKernel.MITCHELL:
+ resize_kernel = mitchell # B = 1/3, C = 1/3
+ kernel_window = 2.
+ case ResizeKernel.CATMULL_ROM:
+ resize_kernel = lambda x: mitchell(x, 0.0, 0.5)
+ kernel_window = 2.
+ case ResizeKernel.B_SPLINE:
+ resize_kernel = lambda x: mitchell(x, 1.0, 0.0)
+ kernel_window = 2.
+ case ResizeKernel.LANCZOS2:
+ resize_kernel = lambda x: lanczos(x, 2)
+ kernel_window = 2.
+ case ResizeKernel.LANCZOS3:
+ resize_kernel = lambda x: lanczos(x, 3)
+ kernel_window = 3.
+ case ResizeKernel.MAGIC_KERNEL:
+ resize_kernel = magic_kernel
+ kernel_window = 1.5
+ case ResizeKernel.MAGIC_KERNEL_SHARP_2013:
+ resize_kernel = magic_kernel_sharp_2013
+ kernel_window = 2.5
+ case ResizeKernel.MAGIC_KERNEL_SHARP_2021:
+ resize_kernel = magic_kernel_sharp_2021
+ kernel_window = 4.5
+ case _:
+ raise ValueError(f"Unknown resize kernel {k}")
+ return resize_kernel, kernel_window
+
+
+### Resampling kernels
+def nearest(x: torch.Tensor) -> torch.Tensor:
+ x = torch.abs(x)
+
+ weights = torch.where(x <= 0.5, 1., 0.)
+
+ return weights
+
+def bilinear(x: torch.Tensor) -> torch.Tensor:
+ x = torch.abs(x)
+
+ weights = torch.where(x <= 1.0, 1 - x, 0.)
+
+ return weights
+
+def mitchell(x: torch.Tensor, B: float = 1 / 3, C: float = 1 / 3) -> torch.Tensor:
+ x = torch.abs(x)
+
+ weights = torch.where(x <= 2, (-B - 6 * C) * x**3 + (6 * B + 30 * C) * x**2 + (-12 * B - 48 * C) * x + (8 * B + 24 * C), 0)
+ weights = torch.where(x <= 1, (12 - 9 * B - 6 * C) * x**3 + (-18 + 12 * B + 6 * C) * x**2 + (6 - 2 * B), weights)
+
+ return weights
+
+def magic_kernel(x: torch.Tensor) -> torch.Tensor:
+ x = torch.abs(x)
+
+ weights = torch.where(x <= 1.5, (1/2) * (x - 3/2) ** 2, 0)
+ weights = torch.where(x <= 0.5, (3/4) - x ** 2, weights)
+
+ return weights
+
+def magic_kernel_sharp_2013(x: torch.Tensor):
+ x = torch.abs(x)
+
+ weights = torch.where(x <= 2.5, (-1/8) * (x - 5/2) ** 2, 0)
+ weights = torch.where(x <= 1.5, (1 - x) * (7/4 - x), weights)
+ weights = torch.where(x <= 0.5, (17/16) - (7/4) * x ** 2, weights)
+
+ return weights
+
+def magic_kernel_sharp_2021(x: torch.Tensor):
+ x = torch.abs(x)
+
+ weights = torch.where(x <= 4.5, (-1/288) * (x - 9/2) ** 2, 0)
+ weights = torch.where(x <= 3.5, (1/36) * (x - 3) * (x - 15/4), weights)
+ weights = torch.where(x <= 2.5, (1/6) * (x - 2) * (65/24 - x), weights)
+ weights = torch.where(x <= 1.5, (35/36) * (x - 1) * (x - 239/140), weights)
+ weights = torch.where(x <= 0.5, (577/576) - (239/144) * x ** 2, weights)
+
+ return weights
+
+def lanczos(x: torch.Tensor, n: int):
+ return torch.where(torch.abs(x) < n, torch.sinc(x) * torch.sinc(x/n), 0)
+
+def sharpen_conv2d(image: torch.Tensor, kernel: torch.Tensor, pad: int) -> torch.Tensor:
+ image = F.pad(image, (pad,pad,pad,pad), mode='replicate')
+ return F.conv2d(image, kernel, groups=image.shape[-3])
+
+### Dithering and related functions.
+def stochastic_round(
+ x: torch.Tensor,
+ out_dtype: torch.dtype,
+ generator: torch.Generator = torch.Generator(),
+ ):
+ image = x * torch.iinfo(out_dtype).max
+ image_quant = image.to(out_dtype)
+ quant_error = image - image_quant.to(image.dtype)
+ dither = torch.empty_like(image_quant, dtype=torch.bool)
+ torch.bernoulli(quant_error, generator=generator, out=dither)
+ return image_quant + dither
+
+def generate_bayer_matrix(n):
+ """Generate an n x n Bayer matrix where n is a power of 2."""
+ assert (n & (n - 1)) == 0 and n > 0, "n must be a power of 2"
+
+ if n == 1:
+ return np.array([[0]]) # Base case
+
+ smaller_matrix = generate_bayer_matrix(n // 2)
+
+ return np.block([
+ [4 * smaller_matrix + 0, 4 * smaller_matrix + 2],
+ [4 * smaller_matrix + 3, 4 * smaller_matrix + 1]
+ ])
+
+### Scaling transforms
+
+def _downscale_axis(
+ image: torch.Tensor,
+ size: int,
+ resize_kernel: ResizeKernel,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ kernel, window = _get_resize_kernel(resize_kernel)
+ k = size / image.shape[-1]
+ PAD = math.ceil((window - 0.5) / k)
+
+ # Optimization note: doing torch.arange like this will compile to doing a int64 arange. Float arange
+ # is much slower. So don't try to get clever and "optimize" by adding the +0.5 and *k to this.
+ # Source grid is padded to allow "out of range" sampling from the source image.
+ coords_source = (torch.arange(-PAD, image.shape[-1]+PAD, 1, dtype=torch.float32, device=device) + 0.5) * k
+ coords_dest = (torch.arange(0, size, 1, dtype=torch.float32, device=device) + 0.5)
+
+ # Create a grid of relative distances between each point on this axis.
+ coord_grid = torch.empty((coords_source.shape[0], coords_dest.shape[0]), dtype=dtype, device=device)
+ # Coord grid always constructed in torch.float32 because float16 precision breaks down for this
+ # after 1024.0. This subtraction is the first opportunity we have to safely cast to float16.
+ torch.sub(coords_source.unsqueeze(-1), other=coords_dest, out=coord_grid)
+
+ weights = kernel(coord_grid)
+
+ # Normalizing weights to sum to 1 along axis we are resizing on
+ weights /= weights.sum(dim=0, keepdim=True)
+ # weights /= (1/k)
+
+ # Padded dimension is reduced by the matmul here.
+ return F.pad(image, (PAD,PAD,0,0), mode='replicate') @ weights
+
+def _upscale_axis(
+ image: torch.Tensor,
+ size: int,
+ resize_kernel: ResizeKernel,
+ device: torch.device,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ kernel, window = _get_resize_kernel(resize_kernel)
+ k = size / image.shape[-1]
+ PAD = math.ceil((window - 0.5) * k)
+
+ # For upsizing, we expect out of range sampling from the destination image.
+ coords_source = (torch.arange(0, image.shape[-1], 1, dtype=torch.float32, device=device) + 0.5)
+ coords_dest = (torch.arange(-PAD, size+PAD, 1, dtype=torch.float32, device=device) + 0.5) / k
+
+ coord_grid = torch.empty((coords_source.shape[0], coords_dest.shape[0]), dtype=dtype, device=device)
+ torch.sub(coords_source.unsqueeze(-1), other=coords_dest, out=coord_grid)
+
+ weights = kernel(coord_grid)
+
+ # We need to explicitly trim padding by summing it into the real area of the destination grid.
+ weights[:, PAD] += weights[:, :PAD].sum(dim=1)
+ weights[:, -PAD-1] += weights[:, -PAD:].sum(dim=1)
+ weights = weights[:, PAD:-PAD]
+
+ weights /= weights.sum(dim=0, keepdim=True)
+
+ return image @ weights
+
+@torch.compile
+def _downscale(
+ image: torch.Tensor,
+ out_res: tuple[int, int],
+ resize_kernel: ResizeKernel,
+ device: torch.device,
+ dtype: torch.dtype,
+ do_srgb_conversion: bool,
+ ):
+ H, W = out_res
+ image = image.to(device=device, dtype=dtype)
+ if do_srgb_conversion:
+ image = srgb_to_linear(image)
+
+ image = _downscale_axis(image, W, resize_kernel, device, dtype)
+ image = _downscale_axis(image.mT, H, resize_kernel, device, dtype).mT
+
+ if do_srgb_conversion:
+ image = linear_to_srgb(image)
+ image = image.clamp(0,1)
+ return image
+
+@torch.compile
+def _upscale(
+ image: torch.Tensor,
+ out_res: tuple[int, int],
+ resize_kernel: ResizeKernel,
+ device: torch.device,
+ dtype: torch.dtype,
+ do_srgb_conversion: bool,
+ ):
+ H, W = out_res
+ image = image.to(device=device, dtype=dtype)
+ if do_srgb_conversion:
+ image = srgb_to_linear(image)
+
+ image = _upscale_axis(image, W, resize_kernel, device, dtype)
+ image = _upscale_axis(image.mT, H, resize_kernel, device, dtype).mT
+
+ if do_srgb_conversion:
+ image = linear_to_srgb(image)
+ image = image.clamp(0,1)
+ return image
+
+# Triton sparse downscale - only available with Triton (CUDA)
+try:
+ from .triton_functional import downscale_sparse
+except ImportError:
+ downscale_sparse = None
+
+def scale(
+ image: torch.Tensor,
+ out_res: Tuple[int, int],
+ resize_kernel: ResizeKernel = ResizeKernel.MAGIC_KERNEL_SHARP_2021,
+ device: torch.device = torch.device('cpu'),
+ dtype: torch.dtype = torch.float32,
+ do_srgb_conversion: bool = True,
+ use_sparse: bool = False,
+ ) -> torch.Tensor:
+ if isinstance(device, str):
+ device = torch.device(device)
+ if use_sparse:
+ assert device.type != "cpu", "sparse implementation is only for GPU!"
+ if resize_kernel != ResizeKernel.MAGIC_KERNEL_SHARP_2021:
+ raise NotImplementedError
+ if downscale_sparse is None:
+ raise ImportError("Triton is required for sparse GPU acceleration")
+
+ context_manager = (
+ set_stance("force_eager") if set_stance and device.type == "cpu" else nullcontext()
+ )
+ with context_manager:
+ if image.shape[-1] <= out_res[-1] and image.shape[-2] <= out_res[-2]:
+ assert not use_sparse
+ return _upscale(image, out_res, resize_kernel, device, dtype, do_srgb_conversion)
+ elif image.shape[-1] >= out_res[-1] and image.shape[-2] >= out_res[-2]:
+ if use_sparse:
+ return downscale_sparse(image, out_res, resize_kernel)
+ return _downscale(image, out_res, resize_kernel, device, dtype, do_srgb_conversion)
+ else:
+ raise ValueError("Mixed axis resizing (e.g. scaling one axis up and the other down) is not supported. File a bug report with your use case if needed.")
diff --git a/modules/sharpfin/sparse_backend.py b/modules/sharpfin/sparse_backend.py
new file mode 100644
index 000000000..2ea9084ea
--- /dev/null
+++ b/modules/sharpfin/sparse_backend.py
@@ -0,0 +1,845 @@
+"""Sharpfin sparse matrix backend for Triton DDS matmul.
+
+Vendored from https://github.com/drhead/sharpfin (Apache 2.0)
+Adapted from https://github.com/stanford-futuredata/stk (Apache 2.0)
+"""
+
+import numpy as np
+import torch
+import triton
+import triton.language as tl
+from typing import Tuple
+from dataclasses import dataclass
+from .triton_functional import linear_to_srgb_triton, srgb_to_linear_triton, magic_kernel_sharp_2021_triton, lanczos_triton
+
+# Code is all adapted from https://github.com/stanford-futuredata/stk, licensed under Apache-2.0
+# Very reduced set of functions for handling DDS (Dense = Dense @ Sparse) matmul only, with the
+# DDS kernel modified to be more flexible on input shapes.
+
+def _validate_matrix(shape, data, row_indices, column_indices, offsets):
+ if data.dim() == 1:
+ data = torch.reshape(data, [data.numel(), 1, 1])
+
+ if data.shape[-2] != data.shape[-1]:
+ raise ValueError(
+ "Expected square blocking in data. "
+ f"Got block shape {[data.shape[-2], data.shape[-1]]}")
+
+ block_size = data.shape[-1]
+ data = data.view([-1, block_size, block_size])
+
+ if data.dim() != 3:
+ raise ValueError(
+ "Expected 3D shape for data (nnz, block, block). "
+ f"Got shape {data.dim()}D shape.")
+
+ block_size = data.shape[1]
+ if shape[-2] % block_size != 0 or shape[-1] % block_size != 0:
+ raise ValueError(
+ "Matrix shape must be dividible by blocking. "
+ f"Got shape {shape} with "
+ f"{[block_size, block_size]} blocking.")
+
+ if np.prod(shape) < data.numel():
+ raise ValueError(
+ "Invalid matrix. Number of nonzeros exceeds matrix capacity "
+ f"({data.numel()} v. {np.prod(shape)})")
+
+ if row_indices.dim() != 1:
+ raise ValueError(
+ f"Expected 1D row_indices. Got {row_indices.dim()}D row_indices.")
+
+ if column_indices.dim() != 1:
+ raise ValueError(
+ f"Expected 1D column_indices. Got {column_indices.dim()}D column_indices.")
+
+ if offsets.dim() != 1:
+ raise ValueError(
+ f"Expected 1D offsets. Got {offsets.dim()}D offsets.")
+
+ if row_indices.numel() != data.shape[0]:
+ raise ValueError(
+ "Expected 1 index per nonzero block. "
+ f"Got {row_indices.numel()} row_indices for {data.shape[0]} blocks")
+
+ if column_indices.numel() != data.shape[0]:
+ raise ValueError(
+ "Expected 1 index per nonzero block. "
+ f"Got {column_indices.numel()} column_indices for {data.shape[0]} blocks")
+
+ block_rows = np.prod(shape[:-1]) / block_size
+ if offsets.numel() != block_rows + 1:
+ raise ValueError(
+ "Expected one offset per block row plus one. "
+ f"Got {offsets.numel()} offsets with {block_rows} block rows.")
+
+ is_cuda = (data.is_cuda and
+ row_indices.is_cuda and
+ column_indices.is_cuda and
+ offsets.is_cuda)
+ is_cpu = (not data.is_cuda and
+ not row_indices.is_cuda and
+ not column_indices.is_cuda and
+ not offsets.is_cuda)
+ if not (is_cuda or is_cpu):
+ raise ValueError(
+ "Expected data & meta-data on common device. "
+ f"Got data on {data.device}, row_indices on {row_indices.device} "
+ f"column_indices on {column_indices.device} and "
+ f"offsets on {offsets.device}.")
+
+ if data.dtype != torch.float16:
+ raise ValueError(
+ f"Expected float16 data. Got {data.dtype} data.")
+ if row_indices.dtype != torch.int16:
+ raise ValueError(
+ f"Expected int16 row_indices. Got {row_indices.dtype} row_indices.")
+ if column_indices.dtype != torch.int16:
+ raise ValueError(
+ f"Expected int16 column_indices. Got {column_indices.dtype} column_indices.")
+ if offsets.dtype != torch.int32:
+ raise ValueError(
+ f"Expected int32 offsets. Got {offsets.dtype} offsets.")
+ return data
+
+def _transpose(size, data: torch.Tensor, row_indices: torch.Tensor, column_indices: torch.Tensor, offsets):
+ block_columns = size[1] // data.shape[1]
+
+ gather_indices = column_indices.argsort()
+ column_indices_t = row_indices.gather(0, gather_indices)
+ block_offsets_t = gather_indices.int()
+
+ column_indices_float = column_indices.float()
+
+ zero = torch.zeros((1,), dtype=torch.int32, device=data.device)
+ nnz_per_column = column_indices_float.histc(block_columns, 0, block_columns)
+ nnz_per_column = nnz_per_column.int()
+ offsets_t = torch.cat([zero, nnz_per_column.cumsum(0, dtype=torch.int32)])
+ return column_indices_t, offsets_t, block_offsets_t
+
+class SBSCMatrix(torch.nn.Module):
+ """Single Block Sparse Column (SBSC) matrix format."""
+ def __init__(
+ self,
+ size,
+ data: torch.Tensor,
+ offset: int,
+ block_size: int
+ ):
+ super().__init__()
+ self.data = data
+ self.offset = offset
+ self.size = size
+ self.num_blocks = data.shape[0]
+ self.col_width = data.shape[2]
+ self.col_block_size = block_size
+
+class Matrix(torch.nn.Module):
+ """A matrix stored in block compressed sparse row (BCSR) format."""
+
+ def __init__(self,
+ size,
+ data: torch.Tensor,
+ row_indices: torch.Tensor,
+ column_indices: torch.Tensor,
+ offsets: torch.Tensor,
+ column_indices_t: torch.Tensor=None,
+ offsets_t: torch.Tensor=None,
+ block_offsets_t: torch.Tensor=None):
+ super().__init__()
+ self._size = size
+ self._data = data
+ self._row_indices = row_indices
+ self._column_indices = column_indices
+ self._offsets = offsets
+
+ if ((column_indices_t is None) or (offsets_t is None) or
+ (block_offsets_t is None)):
+ column_indices_t, offsets_t, block_offsets_t = _transpose(
+ size, data, row_indices, column_indices, offsets)
+ self._column_indices_t = column_indices_t
+ self._offsets_t = offsets_t
+ self._block_offsets_t = block_offsets_t
+
+ self._transposed = False
+
+ max_dim = np.iinfo(np.int16).max * self.blocking
+ if column_indices.dtype == torch.int16:
+ if size[0] > max_dim or size[1] > max_dim:
+ raise ValueError(
+ "Sparse matrix with shape {size} exceeds representable "
+ "size with 16-bit indices.")
+
+ def validate(self):
+ _validate_matrix(self._size,
+ self._data,
+ self._row_indices,
+ self._column_indices,
+ self._offsets)
+
+ def to(self, device):
+ self._data = self._data.to(device)
+ self._row_indices = self._row_indices.to(device)
+ self._column_indices = self._column_indices.to(device)
+ self._offsets = self._offsets.to(device)
+ self._column_indices_t = self._column_indices_t.to(device)
+ self._offsets_t = self._offsets_t.to(device)
+ self._block_offsets_t = self._block_offsets_t.to(device)
+ return self
+
+ def cuda(self):
+ return self.to(torch.cuda.current_device())
+
+ def clone(self):
+ return Matrix(
+ self.size(),
+ self.data.clone(),
+ self.row_indices.clone(),
+ self.column_indices.clone(),
+ self.offsets.clone(),
+ self.column_indices_t.clone(),
+ self.offsets_t.clone(),
+ self.block_offsets_t.clone())
+
+ def t(self):
+ if self.dim() != 2:
+ raise ValueError(
+ "t() expects a tensor with <= 2 dimensions, "
+ f"but self is {self.dim()}D.")
+ out = Matrix(self.size(),
+ self.data,
+ self.row_indices,
+ self.column_indices,
+ self.offsets,
+ self.column_indices_t,
+ self.offsets_t,
+ self.block_offsets_t)
+ out._transposed = not self._transposed
+ out._size = torch.Size((self._size[1], self._size[0]))
+ return out
+
+ def contiguous(self):
+ raise ValueError("Not yet implemented.")
+
+ def is_contiguous(self):
+ return not self._transposed
+
+ @property
+ def is_cuda(self):
+ return self._data.is_cuda
+
+ @property
+ def device(self):
+ return self._data.device
+
+ def size(self):
+ return self._size
+
+ @property
+ def shape(self):
+ return self.size()
+
+ def dim(self):
+ return len(self._size)
+
+ @property
+ def data(self):
+ return self._data
+
+ @property
+ def row_indices(self):
+ return self._row_indices
+
+ @property
+ def column_indices(self):
+ return self._column_indices
+
+ @property
+ def offsets(self):
+ return self._offsets
+
+ @property
+ def offsets_t(self):
+ return self._offsets_t
+
+ @property
+ def column_indices_t(self):
+ return self._column_indices_t
+
+ @property
+ def block_offsets_t(self):
+ return self._block_offsets_t
+
+ @property
+ def dtype(self):
+ return self.data.dtype
+
+ @property
+ def nnz(self):
+ return self.data.numel()
+
+ @property
+ def blocking(self):
+ return self.data.shape[1]
+
+ @property
+ def requires_grad(self):
+ return self.data.requires_grad
+
+ def requires_grad_(self, x):
+ self.data.requires_grad_(x)
+ return self
+
+ def view(self, *shape):
+ assert self.is_contiguous()
+ if shape[-1] != self.size()[-1]:
+ raise ValueError(
+ "Can't change view on compressed dimension. "
+ f"{self.size()[-1]} v. {shape[-1]}.")
+ if np.prod(shape) != np.prod(self.size()):
+ raise ValueError(
+ "Mismatch in numel of Matrix and new shape. "
+ f"{np.prod(self.size())} v. {np.prod(shape)}")
+ return Matrix(shape,
+ self.data,
+ self.row_indices,
+ self.column_indices,
+ self.offsets,
+ self.column_indices_t,
+ self.offsets_t,
+ self.block_offsets_t)
+
+ @property
+ def grad(self):
+ size = self.size()
+ if not self.is_contiguous():
+ size = torch.Size((size[1], size[0]))
+ out = Matrix(size,
+ self.data.grad,
+ self.row_indices,
+ self.column_indices,
+ self.offsets,
+ self.column_indices_t,
+ self.offsets_t,
+ self.block_offsets_t)
+ return out if self.is_contiguous() else out.t()
+
+@torch.no_grad()
+def _expand_for_blocking(idxs, blocking):
+ idxs = torch.reshape(idxs, [idxs.size()[0], 1, 2]).repeat(1, blocking, 1)
+
+ idxs[:, :, 1] *= blocking
+ idxs[:, :, 1] += torch.reshape(torch.arange(blocking, device=idxs.device), [1, blocking])
+
+ idxs = torch.reshape(idxs, [idxs.size()[0], 1, blocking, 2])
+ idxs = idxs.repeat(1, blocking, 1, 1)
+
+ idxs[:, :, :, 0] *= blocking
+ idxs[:, :, :, 0] += torch.reshape(torch.arange(blocking, device=idxs.device), [1, blocking, 1])
+ idxs = torch.reshape(idxs, [-1, 2])
+ return idxs
+
+
+@torch.no_grad()
+def to_dense(x):
+ assert isinstance(x, Matrix)
+
+ shape = (np.prod(x.shape[:-1]), x.shape[-1])
+ row_idxs = x.row_indices.type(torch.int32)
+ col_idxs = x.column_indices.type(torch.int32)
+ indices = _expand_for_blocking(torch.stack([row_idxs, col_idxs], dim=1), x.blocking)
+ indices = (indices[:, 0] * shape[1] + indices[:, 1]).type(torch.int64)
+
+ out = torch.zeros(shape[0] * shape[1], dtype=x.dtype, device=x.device)
+ out.scatter_(0, indices, x.data.flatten())
+ return out.reshape(x.size())
+
+
+@dataclass
+class TritonConfig:
+ BLOCK_M: int = 128
+ BLOCK_N: int = 128
+ BLOCK_K: int = 32
+ BLOCK_SIZE: int = 64
+ NUM_STAGES: int = 4
+ NUM_WARPS: int = 4
+
+@triton.autotune(
+ configs=[
+ triton.Config({}, num_stages=TritonConfig.NUM_STAGES, num_warps=TritonConfig.NUM_WARPS),
+ ],
+ key=['M', 'N', 'K'],
+)
+@triton.jit
+def _dds_kernel(
+ A: tl.tensor, B: tl.tensor, C: tl.tensor, M, N, K, O_M, O_N,
+ stride_ac, stride_am, stride_ak,
+ stride_bk, stride_bn,
+ stride_cc, stride_cm, stride_cn,
+ row_indices: tl.tensor, column_indices: tl.tensor,
+ offsets: tl.tensor, block_offsets_t: tl.tensor,
+ fuse_srgb: tl.constexpr, clamp_output: tl.constexpr,
+ BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr,
+ ):
+
+ pid_c = tl.program_id(0)
+ pid_m = tl.program_id(1)
+ pid_n = tl.program_id(2)
+
+ num_pid_m = tl.num_programs(1)
+ num_pid_n = tl.num_programs(2)
+
+ pid_n, pid_m = tl.swizzle2d(pid_n, pid_m, num_pid_n, num_pid_m, GROUP_M)
+
+ offsets += pid_n
+
+ start_inx = tl.load(offsets)
+ end_inx = tl.load(offsets + 1)
+
+ column_indices += start_inx
+ block_offsets_t += start_inx
+
+ BLOCK_ELEMENTS = BLOCK_SIZE * BLOCK_SIZE
+
+ A_block_ptr = tl.make_block_ptr(
+ base=A + pid_c * stride_ac, shape=(M, K),
+ strides=(stride_am, stride_ak),
+ offsets=(pid_m * BLOCK_M, 0),
+ block_shape=(BLOCK_M, BLOCK_K),
+ order=(0, 1)
+ )
+
+ rn = tl.arange(0, BLOCK_N)
+ rbk = tl.arange(0, BLOCK_K)
+
+ B += (rbk[:, None] * stride_bk + rn[None, :] * stride_bn)
+
+ acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float16)
+ nsub_blocks = tl.cdiv(BLOCK_SIZE, BLOCK_K)
+
+ bk_sub_incr = BLOCK_K * stride_bk
+
+ for block_inx in range(end_inx - start_inx):
+ a_col_idx = tl.load(column_indices + block_inx)
+ ptr_A = tl.advance(A_block_ptr, (0, a_col_idx * BLOCK_SIZE))
+
+ b_block_offset = tl.load(block_offsets_t + block_inx)
+ ptr_B = B + b_block_offset * BLOCK_ELEMENTS
+
+ for sub_block_inx in range(nsub_blocks):
+ a = tl.load(ptr_A)
+ b = tl.load(ptr_B)
+
+ acc = tl.dot(a, b, acc, out_dtype=tl.float16)
+
+ ptr_A = tl.advance(ptr_A, (0, BLOCK_K))
+ ptr_B += bk_sub_incr
+
+ if fuse_srgb:
+ acc = linear_to_srgb_triton(acc)
+
+ if clamp_output:
+ acc = tl.clamp(acc, 0.0, 1.0)
+
+ if fuse_srgb or clamp_output:
+ acc = acc.to(C.dtype.element_ty)
+
+ C_block_ptr = tl.make_block_ptr(
+ base=C + pid_c * stride_cc, shape=(O_M, O_N),
+ strides=(stride_cm, stride_cn),
+ offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
+ block_shape=(BLOCK_M, BLOCK_N),
+ order=(1, 0)
+ )
+
+ tl.store(C_block_ptr, acc, boundary_check=(0, 1))
+
+
+def triton_dds(
+ lhs: torch.Tensor,
+ rhs: Matrix,
+ fuse_srgb: bool = False,
+ clamp_output: bool = False,
+ output_mt: bool = False,
+ output_slice: None | Tuple[int,int] = None
+ ):
+ assert isinstance(lhs, torch.Tensor)
+ assert isinstance(rhs, Matrix)
+ assert lhs.ndim == 3
+ CH = lhs.shape[0]
+ stride_ac = lhs.stride(0)
+
+ M, K = lhs.shape[-2:]
+ N = rhs.shape[-1]
+
+ if output_mt:
+ if output_slice is not None:
+ O_N, O_M = output_slice
+ out = torch.empty(
+ (*lhs.shape[:-2], *output_slice),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ else:
+ O_N, O_M = N, M
+ out = torch.empty(
+ (*lhs.shape[:-2], rhs.shape[1], lhs.shape[-2]),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ stride_cm, stride_cn = out.stride(-1), out.stride(-2)
+ stride_cc = out.stride(-3)
+ else:
+ if output_slice is not None:
+ O_M, O_N = output_slice
+ out = torch.empty(
+ (*lhs.shape[:-2], *output_slice),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ else:
+ O_M, O_N = M, N
+ out = torch.empty(
+ (*lhs.shape[:-1], rhs.shape[1]),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ stride_cm, stride_cn = out.stride(-2), out.stride(-1)
+ stride_cc = out.stride(-3)
+
+ trans_B = not rhs.is_contiguous()
+ trans_A = (lhs.stride(-2) > 1 and lhs.stride(-1) > 1)
+ assert trans_A == False, trans_B == False
+
+ assert lhs.shape[-1] <= rhs.shape[0], "incompatible dimensions"
+
+ stride_am, stride_ak = lhs.stride(-2), lhs.stride(-1)
+
+ if trans_B:
+ stride_bk, stride_bn = rhs.data.stride(2), rhs.data.stride(1)
+ b_column_indices, b_offsets = rhs.column_indices, rhs.offsets
+ else:
+ stride_bk, stride_bn = rhs.data.stride(1), rhs.data.stride(2)
+ b_column_indices, b_offsets = rhs.column_indices_t, rhs.offsets_t
+
+ grid = lambda META: (CH, triton.cdiv(M, META['BLOCK_M']), triton.cdiv(N, META['BLOCK_N']))
+
+ _dds_kernel[grid](
+ lhs, rhs.data, out, M, N, K, O_M, O_N,
+ stride_ac, stride_am, stride_ak,
+ stride_bk, stride_bn,
+ stride_cc, stride_cm, stride_cn,
+ rhs.row_indices, b_column_indices, b_offsets,
+ rhs.block_offsets_t, fuse_srgb, clamp_output,
+ GROUP_M=128, ACC_TYPE=tl.float16, BLOCK_M=min(rhs.data.shape[1], 64),
+ BLOCK_N=rhs.data.shape[1], BLOCK_SIZE=rhs.data.shape[1], BLOCK_K=min(rhs.data.shape[1], 64)
+ )
+ return out
+
+
+@triton.autotune(
+ configs=[
+ triton.Config({}, num_stages=4, num_warps=2),
+ ],
+ key=['BLOCK_SIZE', 'BLOCK_N'],
+)
+@triton.jit
+def _dds_sbsc_kernel(
+ A: tl.tensor, B: tl.tensor, C: tl.tensor, M, N, K, O_M, O_N,
+ stride_ac, stride_am, stride_ak,
+ stride_bb, stride_bk, stride_bn,
+ stride_cc, stride_cm, stride_cn,
+ block_offset: tl.constexpr,
+ fuse_srgb: tl.constexpr, clamp_output: tl.constexpr,
+ BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr, GROUP_M: tl.constexpr, ACC_TYPE: tl.constexpr,
+ ):
+
+ pid_n = tl.program_id(0)
+ pid_m = tl.program_id(1)
+ pid_c = tl.program_id(2)
+
+ nsub_blocks = tl.cdiv(BLOCK_SIZE, BLOCK_K)
+
+ start_row = block_offset * pid_n
+
+ A_block_ptr = tl.make_block_ptr(
+ base=A + pid_c * stride_ac, shape=(M, K),
+ strides=(stride_am, stride_ak),
+ offsets=(pid_m * BLOCK_M, start_row),
+ block_shape=(BLOCK_M, BLOCK_K),
+ order=(0, 1)
+ )
+
+ B_block_ptr = tl.make_block_ptr(
+ base=B + pid_n * stride_bb, shape=(BLOCK_SIZE, BLOCK_N),
+ strides=(stride_bk, stride_bn),
+ offsets=(0, 0),
+ block_shape=(BLOCK_K, BLOCK_N),
+ order=(0, 1)
+ )
+
+ acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+
+ for block_slice in range(nsub_blocks):
+ a = tl.load(A_block_ptr, eviction_policy='evict_first', boundary_check=(0,), padding_option='zero')
+ b = tl.load(B_block_ptr, eviction_policy='evict_last')
+
+ acc = tl.dot(a, b, acc, out_dtype=tl.float32)
+
+ A_block_ptr = A_block_ptr.advance((0, BLOCK_K))
+ B_block_ptr = B_block_ptr.advance((BLOCK_K, 0))
+
+ if fuse_srgb:
+ acc = linear_to_srgb_triton(acc)
+
+ if clamp_output:
+ acc = tl.clamp(acc, 0.0, 1.0)
+
+ acc = acc.to(C.dtype.element_ty)
+
+ C_block_ptr = tl.make_block_ptr(
+ base=C + pid_c * stride_cc, shape=(O_M, O_N),
+ strides=(stride_cm, stride_cn),
+ offsets=(pid_m * BLOCK_M, pid_n * BLOCK_N),
+ block_shape=(BLOCK_M, BLOCK_N),
+ order=(1, 0)
+ )
+
+ tl.store(C_block_ptr, acc, boundary_check=(0, 1), cache_modifier='.cs')
+
+def triton_dds_sbsc(
+ lhs: torch.Tensor,
+ rhs: SBSCMatrix,
+ fuse_srgb: bool = False,
+ clamp_output: bool = False,
+ output_mt: bool = False,
+ output_slice: None | Tuple[int,int] = None
+ ):
+ assert isinstance(lhs, torch.Tensor)
+ assert isinstance(rhs, SBSCMatrix)
+ assert lhs.ndim == 3
+ CH = lhs.shape[0]
+ stride_ac = lhs.stride(0)
+
+ M, K = lhs.shape[-2:]
+ N = rhs.size[-1]
+
+ if output_mt:
+ if output_slice is not None:
+ O_N, O_M = output_slice
+ out = torch.empty(
+ (*lhs.shape[:-2], *output_slice),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ else:
+ O_N, O_M = N, M
+ out = torch.empty(
+ (*lhs.shape[:-2], rhs.size[1], lhs.shape[-2]),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ stride_cm, stride_cn = out.stride(-1), out.stride(-2)
+ stride_cc = out.stride(-3)
+ else:
+ if output_slice is not None:
+ O_M, O_N = output_slice
+ out = torch.empty(
+ (*lhs.shape[:-2], *output_slice),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ else:
+ O_M, O_N = M, N
+ out = torch.empty(
+ (*lhs.shape[:-1], rhs.size[1]),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ stride_cm, stride_cn = out.stride(-2), out.stride(-1)
+ stride_cc = out.stride(-3)
+
+ assert lhs.shape[-1] <= rhs.size[0], f"incompatible dimensions: {lhs.shape[-1]} > {rhs.size[0]}"
+
+ stride_am, stride_ak = lhs.stride(-2), lhs.stride(-1)
+
+ stride_bb, stride_bk, stride_bn = rhs.data.stride(0), rhs.data.stride(1), rhs.data.stride(2)
+
+ grid = lambda META: (triton.cdiv(N, META['BLOCK_N']), triton.cdiv(M, META['BLOCK_M']), CH)
+
+ _dds_sbsc_kernel[grid](
+ lhs, rhs.data, out, M, N, K, O_M, O_N,
+ stride_ac, stride_am, stride_ak,
+ stride_bb, stride_bk, stride_bn,
+ stride_cc, stride_cm, stride_cn,
+ rhs.offset, fuse_srgb, clamp_output,
+ GROUP_M=32, ACC_TYPE=tl.float16, BLOCK_M=32,
+ BLOCK_N=rhs.data.shape[2], BLOCK_SIZE=rhs.data.shape[1], BLOCK_K=rhs.col_block_size
+ )
+ return out
+
+from triton.language.extra import libdevice
+
+@triton.autotune(
+ configs=[
+ triton.Config({}, num_stages=4, num_warps=2),
+ ],
+ key=['BLOCK_SIZE', 'BLOCK_N'],
+)
+@triton.jit
+def _dds_sbsc_zerorhs_kernel(
+ A: tl.tensor, C: tl.tensor, M, N, K, O_M, O_N,
+ stride_ac, stride_am, stride_ak,
+ stride_cc, stride_cm, stride_cn,
+ k, PAD, block_offset: tl.constexpr,
+ fuse_srgb: tl.constexpr, gamma_correction: tl.constexpr, clamp_output: tl.constexpr,
+ BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr,
+ ):
+
+ pid_n = tl.program_id(0)
+ pid_m = tl.program_id(1)
+ pid_c = tl.program_id(2)
+
+ nsub_blocks = triton.cdiv(BLOCK_SIZE, BLOCK_K)
+
+ start_row = block_offset * pid_n
+
+ offs_k = (start_row + tl.arange(0, BLOCK_K)) * stride_ak
+ m_range = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+ A_mask = (m_range < M)[None, :].broadcast_to(BLOCK_K, BLOCK_M)
+
+ A_M_ptr = A + pid_c * stride_ac + stride_am * m_range
+
+ b_k = ((start_row - PAD + tl.arange(0, BLOCK_K)).to(tl.float32) + 0.5) * k
+ b_n = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)).to(tl.float32) + 0.5
+
+ b_base = (b_k[None, :] - b_n[:, None])
+
+ acc = tl.zeros((BLOCK_N, BLOCK_M), dtype=tl.float16)
+
+ for _ in tl.range(nsub_blocks):
+ A_ptr = A_M_ptr[None, :] + tl.minimum(tl.maximum(offs_k, PAD) - PAD, K - 1)[:, None]
+
+ b = magic_kernel_sharp_2021_triton(b_base) * k
+
+ b = b.to(tl.float16)
+
+ a = tl.load(A_ptr, mask=A_mask)
+
+ if fuse_srgb == 'input':
+ if gamma_correction == 'fast':
+ a = libdevice.fast_powf(a, 2.2).to(tl.float16)
+ elif gamma_correction == 'srgb':
+ a = srgb_to_linear_triton(a).to(tl.float16)
+
+ acc = tl.dot(b, a, acc, out_dtype=tl.float16)
+
+ offs_k += BLOCK_K * stride_ak
+ b_base += BLOCK_K * k
+
+ if fuse_srgb == 'output':
+ if gamma_correction == 'fast':
+ acc = libdevice.fast_powf(acc, 1.0/2.2)
+ elif gamma_correction == 'srgb':
+ acc = linear_to_srgb_triton(acc)
+
+ if clamp_output:
+ acc = tl.clamp(acc, 0.0, 1.0)
+
+ if fuse_srgb == 'output' or clamp_output:
+ acc = acc.to(C.dtype.element_ty)
+
+ C_block_ptr = tl.make_block_ptr(
+ base=C + pid_c * stride_cc, shape=(O_N, O_M),
+ strides=(stride_cn, stride_cm),
+ offsets=(pid_n * BLOCK_N, pid_m * BLOCK_M),
+ block_shape=(BLOCK_N, BLOCK_M),
+ order=(1, 0)
+ )
+
+ tl.store(C_block_ptr, acc, boundary_check=(0, 1), cache_modifier='.cs')
+
+import math
+
+
+def triton_dds_zerorhs_sbsc(
+ lhs: torch.Tensor,
+ target_size: int,
+ source_size: int,
+ kernel_window: float,
+ block_specs,
+ fuse_srgb: str = '',
+ gamma_correction: str = 'fast',
+ clamp_output: bool = False,
+ output_mt: bool = False,
+ output_slice: None | Tuple[int,int] = None
+ ):
+ assert isinstance(lhs, torch.Tensor)
+
+ assert fuse_srgb in ['input', 'output', '']
+ assert gamma_correction in ['fast', 'srgb']
+
+ k = target_size / source_size
+
+ PAD = math.ceil((kernel_window - 0.5) / k)
+
+ offset, block_height, num_blocks, col_width = block_specs
+
+ assert lhs.ndim == 3
+ CH = lhs.shape[0]
+ stride_ac = lhs.stride(0)
+
+ M, K = lhs.shape[-2:]
+ N = target_size
+
+ if output_mt:
+ if output_slice is not None:
+ O_N, O_M = output_slice
+ out = torch.empty(
+ (*lhs.shape[:-2], *output_slice),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ else:
+ O_N, O_M = N, M
+ out = torch.empty(
+ (*lhs.shape[:-2], N, M),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ stride_cm, stride_cn = out.stride(-1), out.stride(-2)
+ stride_cc = out.stride(-3)
+ else:
+ if output_slice is not None:
+ O_M, O_N = output_slice
+ out = torch.empty(
+ (*lhs.shape[:-2], *output_slice),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ else:
+ O_M, O_N = M, N
+ out = torch.empty(
+ (*lhs.shape[:-2], M, N),
+ dtype=lhs.dtype,
+ device=lhs.device
+ )
+ stride_cm, stride_cn = out.stride(-2), out.stride(-1)
+ stride_cc = out.stride(-3)
+
+ stride_am, stride_ak = lhs.stride(-2), lhs.stride(-1)
+
+ grid = lambda META: (triton.cdiv(N, META['BLOCK_N']), triton.cdiv(M, META['BLOCK_M']), CH)
+
+ _dds_sbsc_zerorhs_kernel[grid](
+ lhs, out, M, N, K, O_M, O_N,
+ stride_ac, stride_am, stride_ak,
+ stride_cc, stride_cm, stride_cn,
+ k, PAD, offset, fuse_srgb, gamma_correction, clamp_output,
+ BLOCK_M=32, BLOCK_K=16, BLOCK_N=col_width, BLOCK_SIZE=block_height,
+ )
+ return out
diff --git a/modules/sharpfin/transforms.py b/modules/sharpfin/transforms.py
new file mode 100644
index 000000000..95aeb836d
--- /dev/null
+++ b/modules/sharpfin/transforms.py
@@ -0,0 +1,234 @@
+"""Sharpfin transform classes for torchvision integration.
+
+Vendored from https://github.com/drhead/sharpfin (Apache 2.0)
+Imports patched: absolute sharpfin.X -> relative .X, torchvision guarded.
+"""
+
+import torch
+import torch.nn.functional as F
+
+try:
+ from torchvision.transforms.v2 import Transform
+except ImportError:
+ class Transform:
+ _transformed_types = ()
+ def __init__(self):
+ pass
+
+from .util import QuantHandling, ResizeKernel, SharpenKernel, srgb_to_linear, linear_to_srgb
+from . import functional as SFF
+from .cms import apply_srgb
+import math
+from typing import Any, Dict, Tuple
+from PIL import Image
+from .functional import _get_resize_kernel
+from contextlib import nullcontext
+
+try:
+ from .triton_functional import downscale_sparse
+except ImportError:
+ downscale_sparse = None
+
+# from Pytorch >= 2.6
+set_stance = getattr(torch.compiler, "set_stance", None)
+
+__all__ = ["ResizeKernel", "SharpenKernel", "QuantHandling"]
+
+class Scale(Transform):
+ """Rescaling transform supporting multiple algorithms with sRGB linearization."""
+ _transformed_types = (torch.Tensor,)
+ def __init__(self,
+ out_res: tuple[int, int] | int,
+ device: torch.device | str = torch.device('cpu'),
+ dtype: torch.dtype = torch.float32,
+ out_dtype: torch.dtype | None = None,
+ quantization: QuantHandling = QuantHandling.ROUND,
+ generator: torch.Generator | None = None,
+ resize_kernel: ResizeKernel = ResizeKernel.MAGIC_KERNEL_SHARP_2021,
+ sharpen_kernel: SharpenKernel | None = None,
+ do_srgb_conversion: bool = True,
+ use_sparse: bool = False,
+ ):
+ super().__init__()
+ if isinstance(device, str):
+ device = torch.device(device)
+ if not dtype.is_floating_point:
+ raise ValueError("dtype must be a floating point type")
+ if dtype.itemsize == 1:
+ raise ValueError("float8 types are not supported due to severe accuracy issues and limited function support. float16 or float32 is recommended.")
+ if out_dtype is not None and not out_dtype.is_floating_point and out_dtype not in [torch.uint8, torch.uint16, torch.uint32, torch.uint64]:
+ raise ValueError("out_dtype must be a torch float format or a torch unsigned int format")
+ if use_sparse:
+ assert device.type != 'cpu'
+ if resize_kernel != ResizeKernel.MAGIC_KERNEL_SHARP_2021:
+ raise NotImplementedError
+ self.use_sparse = use_sparse
+
+ if isinstance(out_res, int):
+ out_res = (out_res, out_res)
+ self.out_res = out_res
+ self.device = device
+ self.dtype = dtype
+ self.out_dtype = out_dtype if out_dtype is not None else dtype
+ self.do_srgb_conversion = do_srgb_conversion
+
+ if self.out_dtype in [torch.uint8, torch.uint16, torch.uint32, torch.uint64]:
+ match quantization:
+ case QuantHandling.TRUNCATE:
+ self.quantize_function = lambda x: x.mul(torch.iinfo(self.out_dtype).max).to(self.out_dtype)
+ case QuantHandling.ROUND:
+ self.quantize_function = lambda x: x.mul(torch.iinfo(self.out_dtype).max).round().to(self.out_dtype)
+ case QuantHandling.STOCHASTIC_ROUND:
+ if generator is not None:
+ self.generator = torch.Generator(self.device)
+ else:
+ self.generator = generator
+ self.quantize_function = lambda x: SFF.stochastic_round(x, self.out_dtype, self.generator)
+ case QuantHandling.BAYER:
+ self.bayer_matrix = torch.tensor(SFF.generate_bayer_matrix(16), dtype=self.dtype, device=self.device) / 255
+ self.quantize_function = lambda x: self.apply_bayer_matrix(x)
+ case _:
+ raise ValueError(f"Unknown quantization handling type {quantization}")
+ else:
+ self.quantize_function = lambda x: x.to(dtype=out_dtype)
+
+ self.resize_kernel, self.kernel_window = _get_resize_kernel(resize_kernel)
+
+ match sharpen_kernel:
+ case SharpenKernel.SHARP_2013:
+ kernel = torch.tensor([-1, 6, -1], dtype=dtype, device=device) / 4
+ self.sharp_2013_kernel = torch.outer(kernel, kernel).view(1, 1, 3, 3).expand(3, -1, -1, -1)
+ self.sharpen_step = lambda x: SFF.sharpen_conv2d(x, self.sharp_2013_kernel, 1)
+ case SharpenKernel.SHARP_2021:
+ kernel = torch.tensor([-1, 6, -35, 204, -35, 6, -1], dtype=dtype, device=device) / 144
+ self.sharp_2021_kernel = torch.outer(kernel, kernel).view(1, 1, 7, 7).expand(3, -1, -1, -1)
+ self.sharpen_step = lambda x: SFF.sharpen_conv2d(x, self.sharp_2021_kernel, 3)
+ case None:
+ self.sharpen_step = lambda x: x
+ case _:
+ raise ValueError(f"Unknown sharpen kernel {sharpen_kernel}")
+
+ def apply_bayer_matrix(self, x: torch.Tensor):
+ H, W = x.shape[-2:]
+ b = self.bayer_matrix.repeat(1,1,math.ceil(H/16),math.ceil(W/16))[:,:,:H,:W]
+ return (x*255 + b).to(self.out_dtype)
+
+ @torch.compile(disable=False)
+ def downscale(self, image: torch.Tensor, out_res: tuple[int, int]):
+ H, W = out_res
+ image = image.to(dtype=self.dtype)
+ if self.do_srgb_conversion:
+ image = srgb_to_linear(image)
+
+ image = SFF._downscale_axis(image, W, self.kernel_window, self.resize_kernel, self.device, self.dtype)
+ image = SFF._downscale_axis(image.mT, H, self.kernel_window, self.resize_kernel, self.device, self.dtype).mT
+
+ image = self.sharpen_step(image)
+
+ if self.do_srgb_conversion:
+ image = linear_to_srgb(image)
+ image = image.clamp(0,1)
+ image = self.quantize_function(image)
+ return image
+
+ @torch.compile(disable=False)
+ def downscale_sparse(self, image: torch.Tensor, out_res: tuple[int, int]):
+ image = image.to(dtype=self.dtype)
+ if downscale_sparse is not None:
+ image = downscale_sparse(image, out_res)
+ image = self.quantize_function(image)
+ return image
+
+ @torch.compile(disable=False)
+ def upscale(self, image: torch.Tensor, out_res: tuple[int, int]):
+ H, W = out_res
+ image = image.to(dtype=self.dtype)
+ if self.do_srgb_conversion:
+ image = srgb_to_linear(image)
+
+ image = self.sharpen_step(image)
+
+ image = SFF._upscale_axis(image, W, self.kernel_window, self.resize_kernel, self.device, self.dtype)
+ image = SFF._upscale_axis(image.mT, H, self.kernel_window, self.resize_kernel, self.device, self.dtype).mT
+
+ if self.do_srgb_conversion:
+ image = linear_to_srgb(image)
+ image = image.clamp(0,1)
+ image = self.quantize_function(image)
+ return image
+
+ def _transform(self, inpt: torch.Tensor, params: Dict[str, Any]) -> torch.Tensor:
+ image = inpt.to(device=self.device)
+ context_manager = (
+ set_stance("force_eager") if set_stance and self.device.type == "cpu" else nullcontext()
+ )
+ with context_manager:
+ if image.shape[-1] <= self.out_res[-1] and image.shape[-2] <= self.out_res[-2]:
+ return self.upscale(image, self.out_res)
+ elif image.shape[-1] >= self.out_res[-1] and image.shape[-2] >= self.out_res[-2]:
+ if self.use_sparse:
+ return self.downscale_sparse(image, self.out_res)
+ return self.downscale(image, self.out_res)
+ else:
+ raise ValueError("Mixed axis resizing (e.g. scaling one axis up and the other down) is not supported. File a bug report with your use case if needed.")
+
+class ApplyCMS(Transform):
+ """Apply color management to a PIL Image to standardize it to sRGB color space."""
+ _transformed_types = (Image.Image,)
+
+ def _transform(self, inpt: Image.Image, params: Dict[str, Any]) -> Image.Image:
+ if not isinstance(inpt, Image.Image):
+ raise TypeError(f"inpt should be PIL Image. Got {type(inpt)}")
+
+ return apply_srgb(inpt)
+
+class AlphaComposite(Transform):
+ _transformed_types = (Image.Image,)
+ def __init__(
+ self,
+ background: Tuple[int,int,int] = (255, 255, 255)
+ ):
+ super().__init__()
+ self.background = background
+
+ def _transform(self, inpt: Image.Image, params: Dict[str, Any]) -> Image.Image:
+ if not isinstance(inpt, Image.Image):
+ raise TypeError(f"inpt should be PIL Image. Got {type(inpt)}")
+ if not inpt.has_transparency_data:
+ return inpt
+
+ bg = Image.new("RGB", inpt.size, self.background).convert('RGBA')
+
+ return Image.alpha_composite(bg, inpt).convert('RGB')
+
+class AspectRatioCrop(Transform):
+ _transformed_types = (Image.Image,)
+ def __init__(
+ self,
+ width: int,
+ height: int,
+ ):
+ super().__init__()
+ self.ref_width = width
+ self.ref_height = height
+ self.aspect_ratio = width / height
+
+ def _transform(self, inpt: Image.Image, params: Dict[str, Any]) -> Image.Image:
+ if not isinstance(inpt, Image.Image):
+ raise TypeError(f"inpt should be PIL Image. Got {type(inpt)}")
+
+ left, top, right, bottom = 0, 0, inpt.width, inpt.height
+ inpt_ar = inpt.width / inpt.height
+
+ if inpt_ar > self.aspect_ratio:
+ result_width = int(round(inpt.height / self.ref_height * self.ref_width))
+ crop_amt = (inpt.width - result_width) // 2
+ left += crop_amt
+ right -= crop_amt
+ elif inpt_ar < self.aspect_ratio:
+ result_height = int(round(inpt.width / self.ref_width * self.ref_height))
+ crop_amt = (inpt.height - result_height) // 2
+ top += crop_amt
+ bottom -= crop_amt
+
+ return inpt.crop((left, top, right, bottom))
diff --git a/modules/sharpfin/triton_functional.py b/modules/sharpfin/triton_functional.py
new file mode 100644
index 000000000..98986ff75
--- /dev/null
+++ b/modules/sharpfin/triton_functional.py
@@ -0,0 +1,708 @@
+"""Sharpfin Triton-accelerated GPU scaling functions.
+
+Vendored from https://github.com/drhead/sharpfin (Apache 2.0)
+Imports patched: absolute sharpfin.X -> relative .X
+Requires: triton (only available on CUDA platforms)
+"""
+
+import torch
+import math
+import triton
+import triton.language as tl
+
+from .util import ResizeKernel
+from typing import Tuple
+import torch.nn.functional as F
+from triton.language.extra import libdevice
+from .util import linear_to_srgb, srgb_to_linear
+
+# Magic Kernel Sharp with Triton optimizations. Mainly converted to polynomials so that
+# FMA operators can be used.
+@triton.jit
+def magic_kernel_sharp_2021_triton(x: tl.tensor):
+ out = tl.zeros_like(x) # inplace operation doesn't help much.
+ x = tl.abs(x)
+
+ lte_05 = x <= 0.5
+ lte_15 = x <= 1.5
+ lte_25 = x <= 2.5
+ lte_35 = x <= 3.5
+ lte_45 = x <= 4.5
+
+ x_sq = x*x # triton would compile like this anyways but it helps readability
+
+ out = tl.where(lte_05, tl.fma(x_sq, -239/144, 577/576), out)
+ out = tl.where(lte_15 and not lte_05, tl.fma(x_sq, 35/36, tl.fma(x, -379/144, 239/144)), out)
+ out = tl.where(lte_25 and not lte_15, tl.fma(x_sq, -1/6, tl.fma(x, 113/144, -65/72)), out)
+ out = tl.where(lte_35 and not lte_25, tl.fma(x_sq, 1/36, tl.fma(x, -3/16, 5/16)), out)
+ out = tl.where(lte_45 and not lte_35, tl.fma(x_sq, -1/288, tl.fma(x, 1/32, -9/128)), out)
+
+ return out
+
+@triton.jit
+def sinc_triton(x: tl.tensor):
+ y = tl.fma(x, math.pi, 1e-8)
+ return libdevice.fast_sinf(y) / y
+
+@triton.jit
+def lanczos_triton(x: tl.tensor, n: tl.constexpr = 3):
+ return tl.where(
+ tl.abs(x) < n,
+ sinc_triton(x) * sinc_triton(x/n),
+ 0
+ )
+
+# NOTE: there is no reason to use libdevice.pow, its only differences are with subnormals
+@triton.jit
+def linear_to_srgb_triton(x):
+ return tl.where(
+ x <= 0.0031308,
+ x * 12.92,
+ tl.fma(1.055, libdevice.fast_powf(x, 1/2.4), -0.055)
+ )
+
+@triton.jit
+def srgb_to_linear_triton(x):
+ return tl.where(
+ x <= 0.04045,
+ x / 12.92,
+ libdevice.fast_powf(tl.fma(1/1.055, x, 0.055/1.055), 2.4)
+ )
+
+from .sparse_backend import triton_dds, triton_dds_sbsc, triton_dds_zerorhs_sbsc, Matrix, SBSCMatrix
+
+def _get_resize_kernel_triton(k: ResizeKernel):
+ match k:
+ case ResizeKernel.NEAREST:
+ raise NotImplementedError
+ case ResizeKernel.BILINEAR:
+ raise NotImplementedError
+ case ResizeKernel.MITCHELL:
+ raise NotImplementedError
+ case ResizeKernel.CATMULL_ROM:
+ raise NotImplementedError
+ case ResizeKernel.B_SPLINE:
+ raise NotImplementedError
+ case ResizeKernel.LANCZOS2:
+ raise NotImplementedError
+ case ResizeKernel.LANCZOS3:
+ resize_kernel = lanczos_triton
+ kernel_window = 3.
+ case ResizeKernel.MAGIC_KERNEL:
+ raise NotImplementedError
+ case ResizeKernel.MAGIC_KERNEL_SHARP_2013:
+ raise NotImplementedError
+ case ResizeKernel.MAGIC_KERNEL_SHARP_2021:
+ resize_kernel = magic_kernel_sharp_2021_triton
+ kernel_window = 4.5
+ case _:
+ raise ValueError(f"Unknown resize kernel {k}")
+ return resize_kernel, kernel_window
+
+# Sparse Downscale and support functions.
+
+# Amanatides, John and Woo, Andrew -- Fast Voxel Traversal
+def grid_line_tiles(x0, y0, x1, y1, grid_width, grid_height):
+ tiles = set()
+
+ dx = x1 - x0
+ dy = y1 - y0
+
+ x = math.floor(x0)
+ y = math.floor(y0)
+
+ end_x = math.floor(x1)
+ end_y = math.floor(y1)
+
+ step_x = 1 if dx > 0 else -1
+ step_y = 1 if dy > 0 else -1
+
+ t_max_x = ((x + (step_x > 0)) - x0) / dx if dx != 0 else float('inf')
+ t_max_y = ((y + (step_y > 0)) - y0) / dy if dy != 0 else float('inf')
+
+ t_delta_x = abs(1 / dx) if dx != 0 else float('inf')
+ t_delta_y = abs(1 / dy) if dy != 0 else float('inf')
+
+ while True:
+ if 0 <= x < grid_width and 0 <= y < grid_height:
+ tiles.add((y,x))
+ if x == end_x and y == end_y:
+ break
+ if t_max_x < t_max_y:
+ t_max_x += t_delta_x
+ x += step_x
+ else:
+ t_max_y += t_delta_y
+ y += step_y
+
+ return tiles
+
+def tile_mask_function(dest_size, src_size, kernel_window=4.5, tile_size=64):
+ k = dest_size / src_size
+ PAD = math.ceil((kernel_window-0.5) / k)
+
+ grid_size = math.ceil((src_size + 2*PAD)/tile_size), math.ceil(dest_size/tile_size)
+
+ line_1 = 0, 0.5/tile_size, (dest_size)/tile_size, (src_size+0.5)/tile_size
+ line_2 = 0, (2*PAD - 0.5)/tile_size, (dest_size)/tile_size, (src_size + 2*PAD - 0.5)/tile_size
+ lines = line_1, line_2
+
+ mask = torch.zeros(grid_size, dtype=torch.bool)
+
+ tiles = set()
+
+ for (x0, y0, x1, y1) in lines:
+ tiles.update(grid_line_tiles(x0, y0, x1, y1, grid_size[1], grid_size[0]))
+
+ tiles = torch.tensor(list(tiles))
+
+ mask[tiles[:,0], tiles[:,1]] = True
+
+ return mask, tiles
+
+def create_tensor_metadata(
+ tile_mask: torch.Tensor,
+ tiles: torch.Tensor,
+ indices: torch.Tensor,
+ offsets: torch.Tensor,
+ offsets_t: torch.Tensor,
+ ):
+
+ indices[:,:2] = tiles
+
+ torch.argsort(indices[:,1], stable=True, out=indices[:,2]) # block_offsets_t
+ torch.take(indices[:,0], indices[:,2], out=indices[:,3]) # col_indices_t
+
+ # reusing the offsets buffer here helps performance
+ torch.sum(tile_mask, dim=1, out=offsets[1:])
+ torch.sum(tile_mask, dim=0, out=offsets_t[1:])
+ torch.cumsum(offsets, dim=0, out=offsets)
+ torch.cumsum(offsets_t, dim=0, out=offsets_t)
+
+ return indices, offsets, offsets_t
+
+# for isolating the one mandatory graph break
+@torch.compiler.disable
+def _get_nnz_and_buffers(tile_mask):
+ num_sparse_blocks = torch.sum(tile_mask).item()
+
+ return [
+ torch.empty((4, num_sparse_blocks), dtype=torch.int64, pin_memory=True).T, # indices
+ torch.zeros((tile_mask.shape[0] + 1,), dtype=torch.int32, pin_memory=True), # offsets
+ torch.zeros((tile_mask.shape[1] + 1,), dtype=torch.int32, pin_memory=True) # offsets_t
+ ]
+
+
+def generate_sparse_matrix(dest_size, src_size, kernel_window=4.5, tile_size=64):
+ tile_mask, tiles = tile_mask_function(dest_size, src_size, kernel_window, tile_size)
+
+ buffers = _get_nnz_and_buffers(tile_mask)
+ num_sparse_blocks = buffers[0].shape[0]
+
+ indices, offsets, offsets_t = create_tensor_metadata(
+ tile_mask,
+ tiles,
+ *buffers
+ )
+
+ indices = indices.to(device='cuda', dtype=torch.int32, non_blocking=True)
+
+ return Matrix(
+ (tile_mask.shape[0] * tile_size, tile_mask.shape[1] * tile_size),
+ torch.empty(num_sparse_blocks, tile_size, tile_size, dtype=torch.float16, device='cuda'),
+ row_indices=indices[:,0],
+ column_indices=indices[:,1],
+ offsets=offsets.to(device='cuda', non_blocking=True),
+ column_indices_t=indices[:,3],
+ offsets_t=offsets_t.to(device='cuda', non_blocking=True),
+ block_offsets_t=indices[:,2]
+ )
+
+@triton.jit
+def compute_sparse_coord_grid_kernel(
+ coords_source_ptr, coords_dest_ptr, sparse_data_ptr,
+ row_indices_ptr, col_indices_ptr,
+ k: float, M: int, N: int, BLOCK_SIZE: tl.constexpr, SPARSE_BLOCK_SIZE: tl.constexpr
+ ):
+ SPARSE_BLOCK_NUMEL = SPARSE_BLOCK_SIZE * SPARSE_BLOCK_SIZE
+ sparse_block = tl.program_id(0)
+
+ tile_row = tl.program_id(1)
+ tile_col = tl.program_id(2)
+
+ row_offsets = tl.load(row_indices_ptr + sparse_block) * SPARSE_BLOCK_SIZE + tile_row * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ col_offsets = tl.load(col_indices_ptr + sparse_block) * SPARSE_BLOCK_SIZE + tile_col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+ mask_row = row_offsets < M
+ mask_col = col_offsets < N
+
+ coord_source = tl.load(coords_source_ptr + row_offsets, mask=mask_row, other=0.0)
+ coord_dest = tl.load(coords_dest_ptr + col_offsets, mask=mask_col, other=0.0)
+
+ x = tl.cast(coord_source[:, None] - coord_dest[None, :], tl.float16)
+
+ x = magic_kernel_sharp_2021_triton(x)
+
+ x *= k
+
+ sparse_block_ptr = sparse_data_ptr + sparse_block * SPARSE_BLOCK_NUMEL
+
+ local_row_start = tile_row * BLOCK_SIZE
+ local_col_start = tile_col * BLOCK_SIZE
+
+ local_rows = local_row_start + tl.arange(0, BLOCK_SIZE)
+ local_cols = local_col_start + tl.arange(0, BLOCK_SIZE)
+
+ local_rows_2d = local_rows[:, None]
+ local_cols_2d = local_cols[None, :]
+
+ store_offset = local_rows_2d * SPARSE_BLOCK_SIZE + local_cols_2d
+
+ tl.store(sparse_block_ptr + store_offset, x)
+
+def compute_sparse_coord_grid(target_size, source_size, kernel_window, BLOCK_SIZE=32, SPARSE_BLOCK_SIZE=64):
+ assert SPARSE_BLOCK_SIZE % BLOCK_SIZE == 0
+
+ k = target_size / source_size
+ PAD = math.ceil((kernel_window - 0.5) / k)
+
+ coords_source = torch.arange((-PAD + 0.5)*k, (source_size + PAD + 0.5)*k, k, dtype=torch.float32, device='cuda')
+ coords_dest = torch.arange(0.5, target_size + 0.5, 1, dtype=torch.float32, device='cuda')
+
+ M, N = coords_source.shape[0], coords_dest.shape[0]
+ x = generate_sparse_matrix(target_size, source_size, kernel_window, SPARSE_BLOCK_SIZE)
+
+ SPARSE_NUM_BLOCKS = x.data.shape[0]
+
+ grid = lambda meta: (SPARSE_NUM_BLOCKS, triton.cdiv(SPARSE_BLOCK_SIZE, meta['BLOCK_SIZE']), triton.cdiv(SPARSE_BLOCK_SIZE, meta['BLOCK_SIZE']))
+ compute_sparse_coord_grid_kernel[grid](
+ coords_source, coords_dest, x.data,
+ x.row_indices, x.column_indices,
+ k, M, N, BLOCK_SIZE, SPARSE_BLOCK_SIZE)
+ return x
+
+# Dense kernel for downsampling coord_grids
+
+@triton.jit
+def compute_coord_grid_kernel(
+ coords_source_ptr, coords_dest_ptr, coord_grid_ptr, k,
+ M, N, BLOCK_SIZE: tl.constexpr,
+ ):
+ row_offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ col_offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+ mask_row = row_offsets < M
+ mask_col = col_offsets < N
+
+ coord_source = tl.load(coords_source_ptr + row_offsets, mask=mask_row)
+ coord_dest = tl.load(coords_dest_ptr + col_offsets, mask=mask_col)
+
+ x = tl.cast(coord_source[:, None] - coord_dest[None, :], tl.float16)
+
+ x = magic_kernel_sharp_2021_triton(x)
+
+ x *= k
+
+ tl.store(coord_grid_ptr + row_offsets[:, None] * N + col_offsets[None, :], x, mask=mask_row[:, None] & mask_col[None, :])
+
+def compute_coord_grid(target_size, source_size, kernel_window=4.5, BLOCK_SIZE=32):
+ k = target_size / source_size
+ PAD = math.ceil((kernel_window - 0.5) / k)
+
+ coords_source = torch.arange((-PAD + 0.5)*k, (source_size + PAD + 0.5)*k, k, dtype=torch.float32, device='cuda')
+ coords_dest = torch.arange(0.5, target_size + 0.5, 1, dtype=torch.float32, device='cuda')
+
+ M, N = coords_source.shape[0], coords_dest.shape[0]
+ coord_grid = torch.empty((M, N), dtype=torch.float16, device='cuda')
+
+ grid = lambda meta: (triton.cdiv(M, meta['BLOCK_SIZE']), triton.cdiv(N, meta['BLOCK_SIZE']))
+ compute_coord_grid_kernel[grid](coords_source, coords_dest, coord_grid, k, M, N, BLOCK_SIZE)
+ return coord_grid
+
+@triton.jit
+def pad_replicate_kernel(
+ A, B,
+ M_X, N_X,
+ M_Y, N_Y,
+ M_PAD, N_PAD,
+ stride_xc, stride_xm, stride_xn,
+ stride_yc, stride_ym, stride_yn,
+ BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
+ fuse_linrgb: tl.constexpr
+ ):
+ pid_c = tl.program_id(0)
+ pid_m = tl.program_id(1)
+ pid_n = tl.program_id(2)
+
+ offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+ offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+ offs_m_cl = tl.maximum(offs_m, M_PAD) - M_PAD
+ offs_m_cl = tl.minimum(offs_m_cl, M_X - 1)
+ offs_n_cl = tl.maximum(offs_n, N_PAD) - N_PAD
+ offs_n_cl = tl.minimum(offs_n_cl, N_X - 1)
+
+ mask_m = offs_m < M_Y
+ mask_n = offs_n < N_Y
+
+ A_ptr = A + pid_c * stride_xc + offs_m_cl[:, None] * stride_xm + offs_n_cl[None, :] * stride_xn
+ B_ptr = B + pid_c * stride_yc + offs_m[:, None] * stride_ym + offs_n[None, :] * stride_yn
+
+ t = tl.load(A_ptr)
+ if fuse_linrgb:
+ t = srgb_to_linear_triton(t)
+
+ tl.store(B_ptr, t, mask=mask_m[:, None] & mask_n[None, :])
+
+def pad_replicate(
+ img: torch.Tensor,
+ pad_h: int,
+ pad_w: int,
+ sparse_block_size: int = 0,
+ fuse_linrgb: bool = True,
+ ):
+ C = img.shape[0]
+
+ M_PAD = pad_h
+ N_PAD = pad_w
+
+ if sparse_block_size != 0:
+ out_H = img.shape[-2] + M_PAD + (-(img.shape[-2] + M_PAD)) % sparse_block_size
+ out_W = img.shape[-1] + N_PAD + (-(img.shape[-1] + N_PAD)) % sparse_block_size
+ else:
+ out_H = img.shape[-2] + M_PAD + M_PAD
+ out_W = img.shape[-1] + N_PAD + N_PAD
+
+ out = torch.empty(C, out_H, out_W, dtype=img.dtype, device=img.device)
+
+ BLOCK_M = 1
+ BLOCK_N = 512
+
+ grid = lambda META: (
+ C,
+ (out.shape[1] + META['BLOCK_M'] - 1) // META['BLOCK_M'],
+ (out.shape[2] + META['BLOCK_N'] - 1) // META['BLOCK_N'],
+ )
+
+ pad_replicate_kernel[grid](
+ img, out,
+ img.shape[1], img.shape[2],
+ out.shape[1], out.shape[2],
+ M_PAD, N_PAD,
+ img.stride(0), img.stride(1), img.stride(2),
+ out.stride(0), out.stride(1), out.stride(2),
+ BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
+ fuse_linrgb=fuse_linrgb,
+ )
+ return out
+
+def downscale_sparse(
+ image: torch.Tensor,
+ target_size: Tuple[int, int],
+ resize_kernel: ResizeKernel = ResizeKernel.MAGIC_KERNEL_SHARP_2021,
+ do_gamma_handling=True,
+ BLOCK_SIZE: int = 32,
+ SPARSE_BLOCK_SIZE: int = 64,
+ ) -> torch.Tensor:
+ kernel, window = _get_resize_kernel_triton(resize_kernel)
+
+ T_W = target_size[-1]
+ T_H = target_size[-2]
+ S_W = image.shape[-1]
+ S_H = image.shape[-2]
+
+ y_s_w = compute_sparse_coord_grid(T_W, S_W, window, BLOCK_SIZE, SPARSE_BLOCK_SIZE)
+ y_s_h = compute_sparse_coord_grid(T_H, S_H, window, BLOCK_SIZE, SPARSE_BLOCK_SIZE)
+
+ PAD_W = math.ceil((window - 0.5) / (T_W / S_W))
+ PAD_H = math.ceil((window - 0.5) / (T_H / S_H))
+
+ image = pad_replicate(
+ image,
+ PAD_H,
+ PAD_W,
+ SPARSE_BLOCK_SIZE,
+ fuse_linrgb=do_gamma_handling
+ )
+
+ image = triton_dds(
+ image,
+ y_s_w,
+ output_mt=True
+ )
+
+ image = triton_dds(
+ image,
+ y_s_h,
+ fuse_srgb=do_gamma_handling,
+ clamp_output=True,
+ output_mt=True,
+ output_slice=(T_H, T_W)
+ )
+
+ return image
+
+def downscale_triton(
+ image: torch.Tensor,
+ target_size: torch.Size,
+ resize_kernel: ResizeKernel = ResizeKernel.MAGIC_KERNEL_SHARP_2021,
+ do_gamma_handling=True,
+ ) -> torch.Tensor:
+ kernel, window = _get_resize_kernel_triton(resize_kernel)
+
+ y_s_w = compute_coord_grid(target_size[-1], image.shape[-1], window)
+ y_s_h = compute_coord_grid(target_size[-2], image.shape[-2], window)
+
+ PAD_W = math.ceil((window - 0.5) / (target_size[-1] / image.shape[-1]))
+ PAD_H = math.ceil((window - 0.5) / (target_size[-2] / image.shape[-2]))
+
+ image = pad_replicate(image, PAD_H, PAD_W, fuse_linrgb=do_gamma_handling)
+
+ image = image.view(-1, image.shape[-1])
+ image = image @ y_s_w
+ image = image.view(3, -1, image.shape[-1])
+ image = image.mT
+ image = image.reshape(-1, image.shape[-1])
+ image = image @ y_s_h
+ image = image.view(3, -1, image.shape[-1])
+ image = image.mT
+ if do_gamma_handling:
+ image = linear_to_srgb(image[:, :target_size[0], :target_size[1]])
+ image.clamp_(0.,1.)
+ return image
+
+# Single Block Sparse Column implementations.
+
+def evaluate_line(x, x0, y0, x1, y1):
+ """Evaluate the y-coordinate at a given x along a line from (x0, y0) to (x1, y1)."""
+ if x1 == x0:
+ return float('inf')
+ t = (x - x0) / (x1 - x0)
+ return y0 + t * (y1 - y0)
+
+def pad_height_to_multiple(height, multiple):
+ """Pad a height up to the next multiple of 'multiple'."""
+ return int(math.ceil(height / multiple) * multiple)
+
+def generate_sbsc_structure(
+ dest_size,
+ src_size,
+ kernel_window=4.5,
+ tile_size=64,
+ y_tile_size=32
+ ):
+ k = dest_size / src_size
+ PAD = math.ceil((kernel_window - 0.5) / k)
+
+ line1 = (0, 0.5, dest_size, src_size + 0.5)
+ line2 = (0, 2 * PAD - 0.5, dest_size, src_size + 2 * PAD - 0.5)
+
+ y_mins = []
+ y_maxs = []
+ n_blocks = math.ceil(dest_size / tile_size)
+ max_height = 0
+
+ for i in range(n_blocks):
+ x0 = i * tile_size
+ x1 = min(dest_size - 1, x0 + tile_size - 1)
+
+ yt0 = evaluate_line(x0, *line1)
+ yt1 = evaluate_line(x1, *line1)
+ yb0 = evaluate_line(x0, *line2)
+ yb1 = evaluate_line(x1, *line2)
+
+ y_min = min(yt0, yt1)
+ y_max = max(yb0, yb1)
+
+ height = y_max - y_min
+ padded = pad_height_to_multiple(height, y_tile_size)
+
+ y_mins.append(y_min)
+ y_maxs.append(y_max)
+ max_height = max(max_height, padded)
+
+ slope_top = (line1[3] - line1[1]) / (line1[2] - line1[0])
+ ideal_step = slope_top * tile_size
+
+ lower_bounds = []
+ upper_bounds = []
+ for i in range(1, n_blocks):
+ lower_bounds.append((y_maxs[i] - max_height) / i)
+ upper_bounds.append(y_mins[i] / i)
+
+ lower = math.ceil(max(lower_bounds)) if lower_bounds else 0
+ upper = math.floor(min(upper_bounds)) if upper_bounds else int(round(ideal_step))
+
+ fixed_offset = int(round(ideal_step))
+ if fixed_offset < lower:
+ fixed_offset = lower
+ elif fixed_offset > upper:
+ fixed_offset = upper
+
+ return fixed_offset, max_height, n_blocks, tile_size
+
+
+def generate_sbsc_matrix(dest_size, src_size, kernel_window=4.5, tile_size=64, y_tile_size=32):
+ offset, block_height, num_blocks, col_width = generate_sbsc_structure(
+ dest_size, src_size, kernel_window, tile_size, y_tile_size
+ )
+
+ return SBSCMatrix(
+ size=((offset * (num_blocks - 1)) + block_height, dest_size),
+ data=torch.empty((num_blocks, block_height, col_width), dtype=torch.float16, device='cuda'),
+ offset=offset,
+ block_size=y_tile_size
+ )
+
+@triton.jit
+def compute_sbsc_coord_grid_kernel(
+ coords_source_ptr, coords_dest_ptr,
+ sparse_data_ptr, offset: tl.constexpr,
+ stride_xb, stride_xw, stride_xh,
+ k: float, M: int, N: int, BLOCK_SIZE: tl.constexpr, SPARSE_BLOCK_SIZE: tl.constexpr
+ ):
+
+ pid_w = tl.program_id(0)
+ pid_h = tl.program_id(1)
+
+ start_row = offset * pid_w + pid_h * BLOCK_SIZE
+ start_col = pid_w * SPARSE_BLOCK_SIZE
+
+ row_offsets = start_row + tl.arange(0, BLOCK_SIZE)
+ col_offsets = start_col + tl.arange(0, SPARSE_BLOCK_SIZE)
+
+ mask_row = row_offsets < M
+ mask_col = col_offsets < N
+
+ coord_source = tl.load(coords_source_ptr + row_offsets, mask=mask_row, other=0.0)
+ coord_dest = tl.load(coords_dest_ptr + col_offsets, mask=mask_col, other=0.0)
+
+ y = tl.cast(coord_source[:, None] - coord_dest[None, :], tl.float16)
+
+ y = magic_kernel_sharp_2021_triton(y)
+
+ y *= k
+
+ sparse_block_ptr = sparse_data_ptr + pid_w * stride_xb
+
+ local_row_start = pid_h * BLOCK_SIZE
+
+ local_rows = local_row_start + tl.arange(0, BLOCK_SIZE)
+ local_cols = tl.arange(0, SPARSE_BLOCK_SIZE)
+
+ local_rows_2d = local_rows[:, None]
+ local_cols_2d = local_cols[None, :]
+
+ store_offset = local_rows_2d * SPARSE_BLOCK_SIZE + local_cols_2d
+
+ tl.store(sparse_block_ptr + store_offset, y)
+
+def compute_sbsc_coord_grid(target_size, source_size, kernel_window, BLOCK_SIZE=32, SPARSE_BLOCK_SIZE=64):
+ k = target_size / source_size
+ PAD = math.ceil((kernel_window - 0.5) / k)
+
+ coords_source = torch.arange((-PAD + 0.5)*k, (source_size + PAD + 0.5)*k, k, dtype=torch.float32, device='cuda')
+ coords_dest = torch.arange(0.5, target_size + 0.5, 1, dtype=torch.float32, device='cuda')
+
+ M, N = coords_source.shape[0], coords_dest.shape[0]
+ x = generate_sbsc_matrix(target_size, source_size, kernel_window, SPARSE_BLOCK_SIZE, BLOCK_SIZE)
+
+ SPARSE_BLOCKS, BLOCK_HEIGHT, _ = x.data.shape
+ stride_xb, stride_xh, stride_xw = x.data.stride()
+
+ grid = lambda meta: (SPARSE_BLOCKS, triton.cdiv(BLOCK_HEIGHT, meta['BLOCK_SIZE']))
+ compute_sbsc_coord_grid_kernel[grid](
+ coords_source, coords_dest,
+ x.data, x.offset,
+ stride_xb, stride_xh, stride_xw,
+ k, M, N, BLOCK_SIZE, SPARSE_BLOCK_SIZE)
+ return x
+
+
+def downscale_sbsc(
+ image: torch.Tensor,
+ target_size: Tuple[int, int],
+ resize_kernel: ResizeKernel = ResizeKernel.MAGIC_KERNEL_SHARP_2021,
+ do_gamma_handling: bool = True,
+ BLOCK_SIZE: int = 32,
+ SPARSE_BLOCK_SIZE: int = 64,
+ ) -> torch.Tensor:
+ kernel, window = _get_resize_kernel_triton(resize_kernel)
+
+ T_W = target_size[-1]
+ T_H = target_size[-2]
+ S_W = image.shape[-1]
+ S_H = image.shape[-2]
+
+ y_s_w = compute_sbsc_coord_grid(T_W, S_W, window, BLOCK_SIZE, SPARSE_BLOCK_SIZE)
+ y_s_h = compute_sbsc_coord_grid(T_H, S_H, window, BLOCK_SIZE, SPARSE_BLOCK_SIZE)
+
+ PAD_W = math.ceil((window - 0.5) / (T_W / S_W))
+ PAD_H = math.ceil((window - 0.5) / (T_H / S_H))
+
+ image = pad_replicate(
+ image,
+ PAD_H,
+ PAD_W,
+ fuse_linrgb=do_gamma_handling,
+ sparse_block_size=SPARSE_BLOCK_SIZE,
+ )
+
+ image = triton_dds_sbsc(
+ image,
+ y_s_w,
+ output_mt=True
+ )
+
+ image = triton_dds_sbsc(
+ image,
+ y_s_h,
+ fuse_srgb=do_gamma_handling,
+ clamp_output=True,
+ output_mt=True,
+ )
+
+ return image
+
+
+def downscale_sbsc_zerorhs(
+ image: torch.Tensor,
+ target_size: Tuple[int, int],
+ resize_kernel: ResizeKernel = ResizeKernel.MAGIC_KERNEL_SHARP_2021,
+ do_gamma_handling=True,
+ gamma_handling_type: str = 'fast',
+ BLOCK_SIZE: int = 32,
+ SPARSE_BLOCK_SIZE: int = 64,
+ ) -> torch.Tensor:
+ kernel, window = _get_resize_kernel_triton(resize_kernel)
+
+ T_W = target_size[-1]
+ T_H = target_size[-2]
+ S_W = image.shape[-1]
+ S_H = image.shape[-2]
+
+ block_specs_w = generate_sbsc_structure(
+ T_W, S_W, window, BLOCK_SIZE, SPARSE_BLOCK_SIZE
+ )
+
+ block_specs_h = generate_sbsc_structure(
+ T_H, S_H, window, BLOCK_SIZE, SPARSE_BLOCK_SIZE
+ )
+
+ image = triton_dds_zerorhs_sbsc(
+ image,
+ T_W, S_W, window, block_specs_w,
+ fuse_srgb='input' if do_gamma_handling else '',
+ gamma_correction=gamma_handling_type,
+ output_mt=True
+ )
+
+ image = triton_dds_zerorhs_sbsc(
+ image,
+ T_H, S_H, window, block_specs_h,
+ fuse_srgb='output' if do_gamma_handling else '',
+ gamma_correction=gamma_handling_type,
+ clamp_output=True,
+ output_mt=True,
+ )
+
+ return image
diff --git a/modules/sharpfin/util.py b/modules/sharpfin/util.py
new file mode 100644
index 000000000..25340704c
--- /dev/null
+++ b/modules/sharpfin/util.py
@@ -0,0 +1,49 @@
+"""Sharpfin utility types and color space conversion functions.
+
+Vendored from https://github.com/drhead/sharpfin (Apache 2.0)
+"""
+
+from enum import Enum
+import torch
+
+
+def srgb_to_linear(image: torch.Tensor) -> torch.Tensor:
+ return torch.where(
+ image <= 0.04045,
+ image / 12.92,
+ # Clamping is for protection against NaNs during backwards passes.
+ ((torch.clamp(image, min=0.04045) + 0.055) / 1.055) ** 2.4
+ )
+
+
+def linear_to_srgb(image: torch.Tensor) -> torch.Tensor:
+ return torch.where(
+ image <= 0.0031308,
+ image * 12.92,
+ torch.clamp(1.055 * image ** (1 / 2.4) - 0.055, min=0.0, max=1.0)
+ )
+
+
+class ResizeKernel(Enum):
+ NEAREST = "nearest"
+ BILINEAR = "bilinear"
+ CATMULL_ROM = "catmull-rom"
+ MITCHELL = "mitchell"
+ B_SPLINE = "b-spline"
+ LANCZOS2 = "lanczos2"
+ LANCZOS3 = "lanczos3"
+ MAGIC_KERNEL = "magic_kernel"
+ MAGIC_KERNEL_SHARP_2013 = "magic_kernel_sharp_2013"
+ MAGIC_KERNEL_SHARP_2021 = "magic_kernel_sharp_2021"
+
+
+class SharpenKernel(Enum):
+ SHARP_2013 = "sharp_2013"
+ SHARP_2021 = "sharp_2021"
+
+
+class QuantHandling(Enum):
+ TRUNCATE = "truncate"
+ ROUND = "round"
+ STOCHASTIC_ROUND = "stochastic_round"
+ BAYER = "bayer"
diff --git a/modules/upscaler.py b/modules/upscaler.py
index c0c155324..6a7d5e275 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -109,7 +109,8 @@ class Upscaler:
if img.width >= dest_w and img.height >= dest_h:
break
if img.width != dest_w or img.height != dest_h:
- img = img.resize((int(dest_w), int(dest_h)), resample=Image.Resampling.LANCZOS)
+ from modules import images_sharpfin
+ img = images_sharpfin.resize(img, (int(dest_w), int(dest_h)))
shared.state.end(jobid)
return img
diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py
index e1185efbb..5e5e606f3 100644
--- a/modules/upscaler_simple.py
+++ b/modules/upscaler_simple.py
@@ -27,6 +27,8 @@ class UpscalerResize(Upscaler):
UpscalerData("Resize Bilinear", None, self),
UpscalerData("Resize Hamming", None, self),
UpscalerData("Resize Box", None, self),
+ UpscalerData("Resize Sharpfin MKS2021", None, self),
+ UpscalerData("Resize Sharpfin Lanczos3", None, self),
]
def do_upscale(self, img: Image, selected_model=None):
@@ -44,6 +46,12 @@ class UpscalerResize(Upscaler):
return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.HAMMING)
elif selected_model == "Resize Box":
return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=Image.Resampling.BOX)
+ elif selected_model == "Resize Sharpfin MKS2021":
+ from modules import images_sharpfin
+ return images_sharpfin.resize(img, (int(img.width * self.scale), int(img.height * self.scale)), kernel="Sharpfin MKS2021")
+ elif selected_model == "Resize Sharpfin Lanczos3":
+ from modules import images_sharpfin
+ return images_sharpfin.resize(img, (int(img.width * self.scale), int(img.height * self.scale)), kernel="Sharpfin Lanczos3")
else:
return img
diff --git a/modules/upscaler_spandrel.py b/modules/upscaler_spandrel.py
index 975685efe..20e69d558 100644
--- a/modules/upscaler_spandrel.py
+++ b/modules/upscaler_spandrel.py
@@ -25,15 +25,15 @@ class UpscalerSpandrel(Upscaler):
self.scalers.append(scaler)
def process(self, img: Image.Image) -> Image.Image:
- import torchvision.transforms.functional as TF
- tensor = TF.to_tensor(img).unsqueeze(0).to(devices.device)
+ from modules import images_sharpfin
+ tensor = images_sharpfin.to_tensor(img).unsqueeze(0).to(devices.device)
img = img.convert('RGB')
t0 = time.time()
with devices.inference_context():
tensor = self.model(tensor)
tensor = tensor.clamp(0, 1).squeeze(0).cpu()
t1 = time.time()
- upscaled = TF.to_pil_image(tensor)
+ upscaled = images_sharpfin.to_pil(tensor)
log.debug(f'Upscale: name="{self.selected}" input={img.size} output={upscaled.size} time={t1 - t0:.2f}')
return upscaled
diff --git a/modules/upscaler_vae.py b/modules/upscaler_vae.py
index a8c014ee4..fb86702a1 100644
--- a/modules/upscaler_vae.py
+++ b/modules/upscaler_vae.py
@@ -17,9 +17,8 @@ class UpscalerAsymmetricVAE(Upscaler):
def do_upscale(self, img: Image, selected_model=None):
if selected_model is None:
return img
- import torchvision.transforms.functional as F
import diffusers
- from modules import shared, devices
+ from modules import shared, devices, images_sharpfin
if self.vae is None or (selected_model != self.selected):
if 'v1' in selected_model:
repo_id = 'Heasterian/AsymmetricAutoencoderKLUpscaler'
@@ -32,11 +31,11 @@ class UpscalerAsymmetricVAE(Upscaler):
self.selected = selected_model
shared.log.debug(f'Upscaler load: selected="{self.selected}" vae="{repo_id}"')
t0 = time.time()
- img = img.resize((8 * (img.width // 8), 8 * (img.height // 8)), resample=Image.Resampling.LANCZOS).convert('RGB')
- tensor = (F.pil_to_tensor(img).unsqueeze(0) / 255.0).to(device=devices.device, dtype=devices.dtype)
+ img = images_sharpfin.resize(img, (8 * (img.width // 8), 8 * (img.height // 8))).convert('RGB')
+ tensor = images_sharpfin.to_tensor(img).unsqueeze(0).to(device=devices.device, dtype=devices.dtype)
self.vae = self.vae.to(device=devices.device)
tensor = self.vae(tensor).sample
- upscaled = F.to_pil_image(tensor.squeeze().clamp(0.0, 1.0).float().cpu())
+ upscaled = images_sharpfin.to_pil(tensor.squeeze().clamp(0.0, 1.0).float().cpu())
self.vae = self.vae.to(device=devices.cpu)
t1 = time.time()
shared.log.debug(f'Upscale: name="{self.selected}" input={img.size} output={upscaled.size} time={t1 - t0:.2f}')
@@ -57,10 +56,9 @@ class UpscalerWanUpscale(Upscaler):
def do_upscale(self, img: Image, selected_model=None):
if selected_model is None:
return img
- import torchvision.transforms.functional as F
import torch.nn.functional as FN
import diffusers
- from modules import shared, devices
+ from modules import shared, devices, images_sharpfin
if (self.vae_encode is None) or (self.vae_decode is None) or (selected_model != self.selected):
repo_encode = 'Qwen/Qwen-Image-Edit-2509'
subfolder_encode = 'vae'
@@ -79,7 +77,7 @@ class UpscalerWanUpscale(Upscaler):
t0 = time.time()
self.vae_encode = self.vae_encode.to(device=devices.device)
- tensor = (F.pil_to_tensor(img).unsqueeze(0).unsqueeze(2) / 255.0).to(device=devices.device, dtype=devices.dtype)
+ tensor = images_sharpfin.to_tensor(img).unsqueeze(0).unsqueeze(2).to(device=devices.device, dtype=devices.dtype)
tensor = self.vae_encode.encode(tensor).latent_dist.mode()
self.vae_encode.to(device=devices.cpu)
@@ -88,7 +86,7 @@ class UpscalerWanUpscale(Upscaler):
tensor = FN.pixel_shuffle(tensor.movedim(2, 1), upscale_factor=2).movedim(1, 2) # pixel shuffle needs [..., C, H, W] format
self.vae_decode.to(device=devices.cpu)
- upscaled = F.to_pil_image(tensor.squeeze().clamp(0.0, 1.0).float().cpu())
+ upscaled = images_sharpfin.to_pil(tensor.squeeze().clamp(0.0, 1.0).float().cpu())
t1 = time.time()
shared.log.debug(f'Upscale: name="{self.selected}" input={img.size} output={upscaled.size} time={t1 - t0:.2f}')
return upscaled
diff --git a/pipelines/f_lite/pipeline.py b/pipelines/f_lite/pipeline.py
index 18b327087..3cd8b5b09 100644
--- a/pipelines/f_lite/pipeline.py
+++ b/pipelines/f_lite/pipeline.py
@@ -293,10 +293,9 @@ class FLitePipeline(DiffusionPipeline):
raise
# 8. Post-process images
+ from modules import images_sharpfin
images = (decoded_images / 2 + 0.5).clamp(0, 1)
- # Convert to PIL Images
- images = (images * 255).round().clamp(0, 255).to(torch.uint8).cpu()
- pil_images = [Image.fromarray(img.permute(1, 2, 0).numpy()) for img in images]
+ pil_images = [images_sharpfin.to_pil(img) for img in images]
return FLitePipelineOutput(
images=pil_images,
diff --git a/pipelines/model_stablecascade.py b/pipelines/model_stablecascade.py
index 85eea8aef..5a64e4ae0 100644
--- a/pipelines/model_stablecascade.py
+++ b/pipelines/model_stablecascade.py
@@ -332,8 +332,8 @@ class StableCascadeDecoderPipelineFixed(diffusers.StableCascadeDecoderPipeline):
if output_type == "np":
images = images.permute(0, 2, 3, 1).cpu().float().numpy() # float() as bfloat16-> numpy doesnt work
elif output_type == "pil":
- images = images.permute(0, 2, 3, 1).cpu().float().numpy() # float() as bfloat16-> numpy doesnt work
- images = self.numpy_to_pil(images)
+ from modules import images_sharpfin
+ images = [images_sharpfin.to_pil(images[i]) for i in range(images.shape[0])]
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
else:
images = latents
diff --git a/scripts/ctrlx/media.py b/scripts/ctrlx/media.py
index 39086d5d5..dea5a0b76 100644
--- a/scripts/ctrlx/media.py
+++ b/scripts/ctrlx/media.py
@@ -1,6 +1,6 @@
import numpy as np
import torch
-import torchvision.transforms.functional as vF
+from modules import images_sharpfin
import PIL
@@ -13,7 +13,7 @@ def preprocess(image, processor, **kwargs):
elif isinstance(image, np.ndarray):
image = PIL.Image.fromarray(image)
elif isinstance(image, torch.Tensor):
- image = vF.to_pil_image(image)
+ image = images_sharpfin.to_pil(image)
else:
raise TypeError(f"Image must be of type PIL.Image, np.ndarray, or torch.Tensor, got {type(image)} instead.")
diff --git a/scripts/differential_diffusion.py b/scripts/differential_diffusion.py
index f351571a6..fd772a600 100644
--- a/scripts/differential_diffusion.py
+++ b/scripts/differential_diffusion.py
@@ -14,7 +14,6 @@ from packaging import version
import PIL.Image
import numpy as np
import torch
-import torchvision
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
from diffusers.image_processor import VaeImageProcessor
from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
@@ -859,7 +858,8 @@ class StableDiffusionXLDiffImg2ImgPipeline(DiffusionPipeline, FromSingleFileMixi
# 4. Preprocess image
#image = self.image_processor.preprocess(image) #ideally we would have preprocess the image with diffusers, but for this POC we won't --- it throws a deprecated warning
- map = torchvision.transforms.Resize(tuple(s // self.vae_scale_factor for s in original_image.shape[2:]),antialias=None)(map)
+ from modules import images_sharpfin
+ map = images_sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in original_image.shape[2:]), linearize=False)
# 5. Prepare timesteps
def denoising_value_valid(dnv):
return type(denoising_end) == float and 0 < dnv < 1
@@ -1758,7 +1758,8 @@ class StableDiffusionDiffImg2ImgPipeline(DiffusionPipeline):
# 7. Prepare extra step kwargs.
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
- map = torchvision.transforms.Resize(tuple(s // self.vae_scale_factor for s in image.shape[2:]),antialias=None)(map)
+ from modules import images_sharpfin
+ map = images_sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in image.shape[2:]), linearize=False)
# 8. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
@@ -1833,8 +1834,7 @@ class StableDiffusionDiffImg2ImgPipeline(DiffusionPipeline):
import gradio as gr
import diffusers
from PIL import Image, ImageEnhance, ImageOps # pylint: disable=reimported
-from torchvision import transforms
-from modules import errors, shared, devices, scripts_manager, processing, sd_models, images
+from modules import errors, shared, devices, scripts_manager, processing, sd_models, images, images_sharpfin
detector = None
@@ -1888,9 +1888,9 @@ class Script(scripts_manager.Script):
else:
return None, None, None
image_mask = image_map.copy()
- image_map = transforms.ToTensor()(image_map)
+ image_map = images_sharpfin.to_tensor(image_map)
image_map = image_map.to(devices.device)
- image_init = 2 * transforms.ToTensor()(image_init) - 1
+ image_init = 2 * images_sharpfin.to_tensor(image_init) - 1
image_init = image_init.unsqueeze(0)
image_init = image_init.to(devices.device)
return image_init, image_map, image_mask
diff --git a/scripts/lbm_ext.py b/scripts/lbm_ext.py
index d08a72b25..c0f3f7428 100644
--- a/scripts/lbm_ext.py
+++ b/scripts/lbm_ext.py
@@ -84,7 +84,7 @@ class Script(scripts_manager.Script):
from installer import install
install('lpips')
- from torchvision.transforms import ToPILImage, ToTensor
+ from modules import images_sharpfin
from scripts.lbm import get_model, extract_object, resize_and_center_crop # pylint: disable=no-name-in-module
ori_h_bg, ori_w_bg = fg_image.size
@@ -110,7 +110,7 @@ class Script(scripts_manager.Script):
if lbm_method == 'Simple':
output_image = img_pasted
else:
- img_pasted_tensor = ToTensor()(img_pasted).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) * 2 - 1
+ img_pasted_tensor = images_sharpfin.to_tensor(img_pasted).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) * 2 - 1
batch = { "source_image": img_pasted_tensor }
z_source = model.vae.encode(batch[model.source_key])
output_image = model.sample(
@@ -120,7 +120,7 @@ class Script(scripts_manager.Script):
max_samples=1,
)
output_image = (output_image[0].clamp(-1, 1).float().cpu() + 1) / 2
- output_image = ToPILImage()(output_image)
+ output_image = images_sharpfin.to_pil(output_image)
if lbm_composite:
output_image = Image.composite(output_image, bg_image, fg_mask)
diff --git a/scripts/pixelsmith_ext.py b/scripts/pixelsmith_ext.py
index da9686e7b..a87570d0f 100644
--- a/scripts/pixelsmith_ext.py
+++ b/scripts/pixelsmith_ext.py
@@ -26,17 +26,13 @@ class Script(scripts_manager.Script):
def encode(self, p: processing.StableDiffusionProcessing, image: Image.Image):
if image is None:
return None
- import numpy as np
- import torch
+ from modules import images_sharpfin
if p.width is None or p.width == 0:
p.width = int(8 * (image.width * p.scale_by // 8))
if p.height is None or p.height == 0:
p.height = int(8 * (image.height * p.scale_by // 8))
image = images.resize_image(p.resize_mode, image, p.width, p.height, upscaler_name=p.resize_name, context=p.resize_context)
- tensor = np.array(image).astype(np.float16) / 255.0
- tensor = tensor[None].transpose(0, 3, 1, 2)
- # image = image.transpose(0, 3, 1, 2)
- tensor = torch.from_numpy(tensor).to(device=devices.device, dtype=devices.dtype)
+ tensor = images_sharpfin.to_tensor(image).unsqueeze(0).to(device=devices.device, dtype=devices.dtype)
tensor = 2.0 * tensor - 1.0
with devices.inference_context():
latent = shared.sd_model.vae.tiled_encode(tensor)
diff --git a/scripts/softfill.py b/scripts/softfill.py
index 71d22a7b9..51a286dc7 100644
--- a/scripts/softfill.py
+++ b/scripts/softfill.py
@@ -18,7 +18,6 @@ import cv2
import numpy as np
from PIL import Image, ImageFilter
import torch
-import torchvision
from torchvision import transforms
from transformers import (
CLIPImageProcessor,
@@ -1323,7 +1322,8 @@ class StableDiffusionXLSoftFillPipeline(
image.save("noised_image.png")
image = transforms.CenterCrop((image.size[1] // 64 * 64, image.size[0] // 64 * 64))(image)
- image = transforms.ToTensor()(image)
+ from modules import images_sharpfin
+ image = images_sharpfin.to_tensor(image)
image = image * 2 - 1 # Normalize to [-1, 1]
return image.unsqueeze(0)
@@ -1334,7 +1334,8 @@ class StableDiffusionXLSoftFillPipeline(
"""
map = map.convert("L")
map = transforms.CenterCrop((map.size[1] // 64 * 64, map.size[0] // 64 * 64))(map)
- map = transforms.ToTensor()(map)
+ from modules import images_sharpfin
+ map = images_sharpfin.to_tensor(map)
map = (map - 0.05) / (0.95 - 0.05)
map = torch.clamp(map, 0.0, 1.0)
return 1.0 - map
@@ -1349,9 +1350,8 @@ class StableDiffusionXLSoftFillPipeline(
# Prepare mask as rescaled tensor map
map = preprocess_map(mask).to(device)
- map = torchvision.transforms.Resize(
- tuple(s // self.vae_scale_factor for s in original_image_tensor.shape[2:]), antialias=None
- )(map)
+ from modules import images_sharpfin
+ map = images_sharpfin.resize_tensor(map, tuple(s // self.vae_scale_factor for s in original_image_tensor.shape[2:]), linearize=False)
# Generate latent tensor with noise
original_with_noise = self.prepare_latents(