add hqx and icb interpolations

Signed-off-by: vladmandic <mandic00@live.com>
pull/4626/head
vladmandic 2026-02-05 11:21:17 +01:00
parent 605d87cb2d
commit 59654d68ea
8 changed files with 8159 additions and 167 deletions

View File

@ -4,10 +4,11 @@
Bugfix refresh
- add metadata restore to always-on scripts
- improve wildcard weights parsing, thanks @Tillerz
- fix `anima` model detection
- reorganize `cli` scripts
- add metadata restore to always-on scripts
- improve wildcard weights parsing, thanks @Tillerz
- fix `anima` model detection
- reorganize `cli` scripts
- upscalers: hqx, icbi
## Update for 2026-02-04

View File

@ -41,12 +41,6 @@
TODO: Investigate which models are diffusers-compatible and prioritize!
### Upscalers
- [HQX](https://github.com/uier/py-hqx/blob/main/hqx.py)
- [DCCI](https://every-algorithm.github.io/2024/11/06/directional_cubic_convolution_interpolation.html)
- [ICBI](https://github.com/gyfastas/ICBI/blob/master/icbi.py)
### Image-Base
- [Chroma Zeta](https://huggingface.co/lodestones/Zeta-Chroma): Image and video generator for creative effects and professional filters
- [Chroma Radiance](https://huggingface.co/lodestones/Chroma1-Radiance): Pixel-space model eliminating VAE artifacts for high visual fidelity

7724
modules/postprocess/hqx.py Normal file

File diff suppressed because it is too large Load Diff

225
modules/postprocess/icbi.py Normal file
View File

@ -0,0 +1,225 @@
'''
This is the python implementation of icbi.m
Author: gyf
Begin: 2019-1-16
'''
import numpy as np
import cv2
def icbi(IM,ZK = 1,SZ = 8,PF = 1,ST = 20,TM = 100,TC = 50,SC = 1,TS = 100,AL = 1,BT = -1,GM = 5):
'''
:param IM: Source image
:param ZK: Power of zoom factor (default:1)
:param SZ: Number of image bits per layer (default:8)
:param PF: Potential to be minimized (default:1)
:param ST: Maximum number of iterations (default:20)
:param TM: Maximum edge step (default:100)
:param TC: Edge continuity threshold (deafult:50).
:param SC: Stopping criterion: 1 = change under threshold, 0 = ST iterations (default:1).
:param TS: Threshold on image change for stopping iterations (default:100).
:param AL: Weight for Curvature Continuity energy (default:1.0).
:param BT: Weight for Curvature enhancement energy (default:-1.0).
:param GM: Weight for Isophote smoothing energy (default:5.0).
:return: EI: Enlarged image
'''
H = IM.shape[0]
W = IM.shape[1]
if ZK < 1:
EI = cv2.resize(IM,(H*(2**ZK),W*(2**ZK)))
#check image type
IDIM = np.ndim(IM)
if IDIM == 3:
CL = IM.shape[2] #number of colors
elif IDIM == 2:
IM = np.reshape(IM,(H,W,1))
CL = 1
else:
print('Unrecognized image type, please use RGB or grayscale images')
return 0
#calculate final size
fm = H * (2**ZK) - (2**ZK - 1)
fn = W * (2**ZK) - (2**ZK - 1)
#initialize output image
if SZ>32:
EI = np.zeros([fm,fn,CL],dtype= np.uint64)
elif SZ>16:
EI = np.zeros([fm,fn,CL],dtype= np.uint32)
elif SZ>8:
EI = np.zeros([fm,fn,CL],dtype= np.uint16)
else:
EI = np.zeros([fm,fn,CL],dtype= np.uint8)
#each image color
IMG = IM.copy()
for CID in range(CL):
IMG = IM[:,:,CID]
#The image is enlarged by scaling factor 2**ZK-1 at each cycle
for _ZF in range(ZK):
#size of enlarged image
mm = 2*H - 1
nn = 2*W - 1
#initialize expanded and support matrix
IMGEXP = np.zeros([mm,nn])
D1 = np.zeros([mm,nn])
D2 = np.zeros([mm,nn])
D3 = np.zeros([mm,nn])
C1 = np.zeros([mm,nn])
C2 = np.zeros([mm,nn])
#copy low resolution grid on high resolution grid
IMGEXP[::2,::2] = IMG
#interpolation at borders (average value of 2 neighbors)
for i in range(1,mm-1,2):
#left col
IMGEXP[i,0] = (IMGEXP[i-1,0]+IMGEXP[i+1,0])/2
#right col
IMGEXP[i,nn-1] = (IMGEXP[i-1,nn-1]+IMGEXP[i+1,nn-1])/2
for i in range(1,nn,2):
#top row
IMGEXP[0,i] = (IMGEXP[0,i-1] + IMGEXP[0,i+1])/2
#bottom row
IMGEXP[mm-1,i] = (IMGEXP[mm-1,i-1]+IMGEXP[mm-1,i+1])/2
#Calculate interpolated points in two steps
#s = 0 calculates on diagonal directions
#s = 1 calculates on vertical and horizontal directions
for s in range(2):
#FCBI (Fast Curvature Based Interpolation)
for i in range(1,mm-s,2-s):
for j in range(1+(s*(1-np.mod(i+1,2))),nn-s,2):
v1 = np.abs(IMGEXP[i-1,j-1+s]-IMGEXP[i+1,j+1-s])
v2 = np.abs(IMGEXP[i+1-s,j-1]-IMGEXP[i-1+s,j+1])
p1 = (IMGEXP[i-1,j-1+s]+IMGEXP[i+1,j+1-s])/2
p2 = (IMGEXP[i+1-s,j-1]+IMGEXP[i-1+s,j+1])/2
if (v1<TM) and (v2<TM) and (i>2-s) and i<mm-4-s and j>2-s and j<nn-4-s and (np.abs(p1-p2)<TM):
if np.abs( IMGEXP[i-1-s,j-3+2*s] + IMGEXP[i-3+s,j-1+2*s] + IMGEXP[i+1+s,j+3-2*s] +IMGEXP[i+3-s,j+1-2*s] + 2*p2-6*p1)> np.abs( IMGEXP[i-3+2*s,j+1+s] + IMGEXP[i-1+2*s,j+3-s] + IMGEXP[i+3-2*s,j-1-s] +IMGEXP[i+1-2*s,j-3+s] + 2*p1-6*p2):
IMGEXP[i,j] = p1
else:
IMGEXP[i,j] = p2
else:
if v1<v2:
IMGEXP[i,j] = p1
else:
IMGEXP[i,j] = p2
step = 4.0/(1+s)
#iterative refinement
for g in range(ST):
diff = 0
if g<ST/4 -1:
step = 1
elif g<ST/2 -1:
step = 2
elif g<3*ST/4 -1:
step = 2
#computation of derivatives:
for i in range(3-2*s,mm-3+s):
for j in range(3-2*s+(1-s)*np.mod(i+1,2),nn-3+s,2-s):
C1[i,j] = (IMGEXP[i-1+s,j-1] - IMGEXP[i+1-s,j+1])/2
C2[i,j] = (IMGEXP[i+1-2*s,j-1+s] - IMGEXP[i-1+2*s,j+1-s])/2
D1[i,j] = IMGEXP[i-1+s,j-1] + IMGEXP[i+1-s,j+1] - 2*IMGEXP[i,j]
D2[i,j] = IMGEXP[i+1,j-1+s] + IMGEXP[i-1,j+1-s] - 2*IMGEXP[i,j]
D3[i,j] = (IMGEXP[i-s,j-2+s] - IMGEXP[i-2+s,j+s] + IMGEXP[i+s,j+2-s] - IMGEXP[i+2-s,j-s])/2
for i in range(5-3*s,mm-5+3*s,2-s):
for j in range(5+s*(np.mod(i+1,2)-2),nn-5+3*s,2):
c_1 = 1
c_2 = 1
c_3 = 1
c_4 = 1
if np.abs(IMGEXP[i+1-s,j+1] - IMGEXP[i,j])>TC:
c_1 = 0
if np.abs(IMGEXP[i-1+s,j-1] - IMGEXP[i,j])>TC:
c_2 = 0
if np.abs(IMGEXP[i+1,j-1+s] - IMGEXP[i,j])>TC:
c_3 = 0
if np.abs(IMGEXP[i-1,j+1-s] - IMGEXP[i,j])>TC:
c_4 = 0
EN1 = c_1*np.abs(D1[i,j] - D1[i+1-s,j+1]) + c_2*np.abs(D1[i,j] - D1[i-1+s,j-1])
EN2 = c_3*np.abs(D1[i,j] - D1[i+1,j-1+s]) + c_4*np.abs(D1[i,j] - D1[i-1,j+1-s])
EN3 = c_1*np.abs(D2[i,j] - D2[i+1-s,j+1]) + c_2*np.abs(D2[i,j] - D2[i-1+s,j-1])
EN4 = c_3*np.abs(D2[i,j] - D2[i+1,j-1+s]) + c_4*np.abs(D2[i,j] - D2[i-1,j+1-s])
EN5 = np.abs(IMGEXP[i-2+2*s,j-2] + IMGEXP[i+2-2*s,j+2] - 2*IMGEXP[i,j])
EN6 = np.abs(IMGEXP[i+2,j-2+2*s] + IMGEXP[i-2,j+2-2*s] - 2*IMGEXP[i,j])
EA1 = c_1*np.abs(D1[i,j] - D1[i+1-s,j+1] - 3*step) + c_2*np.abs(D1[i,j] - D1[i-1+s,j-1] - 3*step)
EA2 = c_3*np.abs(D1[i,j] - D1[i+1,j-1+s] - 3*step) + c_4*np.abs(D1[i,j] - D1[i-1,j+1-s] - 3*step)
EA3 = c_1*np.abs(D2[i,j] - D2[i+1-s,j+1] - 3*step) + c_2*np.abs(D2[i,j] - D2[i-1+s,j-1] - 3*step)
EA4 = c_3*np.abs(D2[i,j] - D2[i+1,j-1+s] - 3*step) + c_4*np.abs(D2[i,j] - D2[i-1,j+1-s] - 3*step)
EA5 = np.abs(IMGEXP[i-2+2*s,j-2] + IMGEXP[i+2-2*s,j+2] - 2*IMGEXP[i,j] - 2*step)
EA6 = np.abs(IMGEXP[i+2,j-2+2*s] + IMGEXP[i-2,j+2-2*s] - 2*IMGEXP[i,j] - 2*step)
ES1 = c_1*np.abs(D1[i,j] - D1[i+1-s,j+1] + 3*step) + c_2*np.abs(D1[i,j] - D1[i-1+s,j-1] + 3*step)
ES2 = c_3*np.abs(D1[i,j] - D1[i+1,j-1+s] + 3*step) + c_4*np.abs(D1[i,j] - D1[i-1,j+1-s] + 3*step)
ES3 = c_1*np.abs(D2[i,j] - D2[i+1-s,j+1] + 3*step) + c_2*np.abs(D2[i,j] - D2[i-1+s,j-1] + 3*step)
ES4 = c_3*np.abs(D2[i,j] - D2[i+1,j-1+s] + 3*step) + c_4*np.abs(D2[i,j] - D2[i-1,j+1-s] + 3*step)
ES5 = np.abs(IMGEXP[i-2+2*s,j-2] + IMGEXP[i+2-2*s,j+2] - 2*IMGEXP[i,j] + 2*step)
ES6 = np.abs(IMGEXP[i+2,j-2+2*s] + IMGEXP[i-2,j+2-2*s] - 2*IMGEXP[i,j] + 2*step)
EISO = (C1[i,j]*C1[i,j]*D2[i,j] - 2*C1[i,j]*C2[i,j]*D3[i,j] + C2[i,j]*C2[i,j]*D1[i,j])/(C1[i,j]*C1[i,j]+C2[i,j]*C2[i,j])
if np.abs(EISO) < 0.2:
EISO = 0
if PF==1:
EN = AL*(EN1 + EN2 + EN3 + EN4) + BT*(EN5 + EN6)
EA = AL*(EA1 + EA2 + EA3 + EA4) + BT*(EA5 + EA6)
ES = AL*(ES1 + ES2 + ES3 + ES4) + BT*(ES5 + ES6)
elif PF==2:
EN = AL*(EN1 + EN2 + EN3 + EN4)
EA = AL*(EA1 + EA2 + EA3 + EA4) - GM*np.sign(EISO)
ES = AL*(ES1 + ES2 + ES3 + ES4) - GM*np.sign(EISO)
else:
EN = AL*(EN1 + EN2 + EN3 + EN4) + BT*(EN5 + EN6)
EA = AL*(EA1 + EA2 + EA3 + EA4) + BT*(EA5 + EA6) - GM*np.sign(EISO)
ES = AL*(ES1 + ES2 + ES3 + ES4) + BT*(ES5 + ES6) + GM*np.sign(EISO)
if (EN>EA) and (ES>EA):
IMGEXP[i,j] = IMGEXP[i,j] + step
diff = diff + step
elif (EN>ES) and (EA>ES):
IMGEXP[i,j] = IMGEXP[i,j] - step
diff = diff + step
if (SC==1) and (diff<TS):
break
#assign the expanded image to the current image
IMG = IMGEXP
EI[:,:,CID] = np.round(IMG)
#back to 2D array if gray
if CL ==1:
EI = np.reshape(EI,(fm,fn))
return EI

116
modules/upscaler_algo.py Normal file
View File

@ -0,0 +1,116 @@
import time
from PIL import Image
from modules.upscaler import Upscaler, UpscalerData
from modules.shared import log
class UpscalerDCC(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "DCC Interpolation"
self.vae = None
self.scalers = [
UpscalerData("DCC Interpolation", None, self),
]
def do_upscale(self, img: Image, selected_model=None):
import math
import numpy as np
from modules.postprocess.dcc import DCC
t0 = time.time()
normalized = np.array(img).astype(np.float32) / 255.0
scale = math.ceil(self.scale)
upscaled = DCC(normalized, scale)
upscaled = (upscaled - upscaled.min()) / (upscaled.max() - upscaled.min())
upscaled = (255.0 * upscaled).astype(np.uint8)
upscaled = Image.fromarray(upscaled)
t1 = time.time()
log.debug(f"Upscale: name=DCC input={img.size} output={upscaled.size} time={t1 - t0:.2f}")
return upscaled
class UpscalerVIPS(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "VIPS"
self.scalers = [
UpscalerData("VIPS Lanczos 2", None, self),
UpscalerData("VIPS Lanczos 3", None, self),
UpscalerData("VIPS Mitchell", None, self),
UpscalerData("VIPS MagicKernelSharp 2013", None, self),
UpscalerData("VIPS MagicKernelSharp 2021", None, self),
]
def do_upscale(self, img: Image, selected_model=None):
if selected_model is None:
return img
from installer import install
install('pyvips')
try:
import pyvips
except Exception as e:
log.error(f"Upscaler: vips {e}")
return img
t0 = time.time()
vips_image = pyvips.Image.new_from_array(img)
try:
if selected_model is None:
return img
elif selected_model == "VIPS Lanczos 2":
vips_image = vips_image.resize(2, kernel='lanczos2')
elif selected_model == "VIPS Lanczos 3":
vips_image = vips_image.resize(2, kernel='lanczos3')
elif selected_model == "VIPS Mitchell":
vips_image = vips_image.resize(2, kernel='mitchell')
elif selected_model == "VIPS MagicKernelSharp 2013":
vips_image = vips_image.resize(2, kernel='mks2013')
elif selected_model == "VIPS MagicKernelSharp 2021":
vips_image = vips_image.resize(2, kernel='mks2021')
else:
return img
except Exception as e:
log.error(f"Upscaler: vips {e}")
return img
upscaled = Image.fromarray(vips_image.numpy())
t1 = time.time()
log.debug(f"Upscale: name=VIPS input={img.size} output={upscaled.size} time={t1 - t0:.2f}")
return upscaled
class UpscalerHQX(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "HQX"
self.scalers = [
UpscalerData("HQX Interpolation", None, self),
]
def do_upscale(self, img: Image, selected_model=None):
import numpy as np
from modules.postprocess.hqx import hqx
t0 = time.time()
np_img = np.array(img).astype(np.uint32)
upscaled = hqx(np_img, 2)
upscaled = (upscaled).astype(np.uint8)
upscaled = Image.fromarray(upscaled)
t1 = time.time()
log.debug(f"Upscale: name=HQX input={img.size} output={upscaled.size} time={t1 - t0:.2f}")
return upscaled
class UpscalerICBI(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "ICB"
self.scalers = [
UpscalerData("ICB Interpolation", None, self),
]
def do_upscale(self, img: Image, selected_model=None):
import numpy as np
from modules.postprocess.icbi import icbi
t0 = time.time()
np_img = np.array(img)
upscaled = icbi(np_img)
upscaled = Image.fromarray(upscaled)
t1 = time.time()
log.debug(f"Upscale: name=ICB input={img.size} output={upscaled.size} time={t1 - t0:.2f}")
return upscaled

View File

@ -93,160 +93,3 @@ class UpscalerLatent(Upscaler):
else:
raise log.error(f"Upscale: type=latent model={selected_model} unknown")
return F.interpolate(img, size=(h, w), mode=mode, antialias=antialias)
class UpscalerAsymmetricVAE(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "Asymmetric VAE"
self.vae = None
self.selected = None
self.scalers = [
UpscalerData("Asymmetric VAE v1", None, self),
UpscalerData("Asymmetric VAE v2", None, self),
]
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
if self.vae is None or (selected_model != self.selected):
if 'v1' in selected_model:
repo_id = 'Heasterian/AsymmetricAutoencoderKLUpscaler'
else:
repo_id = 'Heasterian/AsymmetricAutoencoderKLUpscaler_v2'
self.vae = diffusers.AsymmetricAutoencoderKL.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir)
self.vae.requires_grad_(False)
self.vae = self.vae.to(device=devices.device, dtype=devices.dtype)
self.vae.eval()
self.selected = selected_model
shared.log.debug(f'Upscaler load: selected="{self.selected}" vae="{repo_id}"')
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)
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())
self.vae = self.vae.to(device=devices.cpu)
return upscaled
class UpscalerWanUpscale(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "WAN Upscale"
self.vae_encode = None
self.vae_decode = None
self.selected = None
self.scalers = [
UpscalerData("WAN Asymmetric Upscale", None, self),
]
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
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'
self.vae_encode = diffusers.AutoencoderKLWan.from_pretrained(repo_encode, subfolder=subfolder_encode, cache_dir=shared.opts.hfcache_dir)
self.vae_encode.requires_grad_(False)
self.vae_encode = self.vae_encode.to(device=devices.device, dtype=devices.dtype)
self.vae_encode.eval()
repo_decode = 'spacepxl/Wan2.1-VAE-upscale2x'
subfolder_decode = "diffusers/Wan2.1_VAE_upscale2x_imageonly_real_v1"
self.vae_decode = diffusers.AutoencoderKLWan.from_pretrained(repo_decode, subfolder=subfolder_decode, cache_dir=shared.opts.hfcache_dir)
self.vae_decode.requires_grad_(False)
self.vae_decode = self.vae_decode.to(device=devices.device, dtype=devices.dtype)
self.vae_decode.eval()
self.selected = selected_model
shared.log.debug(f'Upscaler load: selected="{self.selected}" encode="{repo_encode}" decode="{repo_decode}"')
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 = self.vae_encode.encode(tensor).latent_dist.mode()
self.vae_encode.to(device=devices.cpu)
self.vae_decode = self.vae_decode.to(device=devices.device)
tensor = self.vae_decode.decode(tensor).sample
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())
return upscaled
class UpscalerDCC(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "DCC Interpolation"
self.vae = None
self.scalers = [
UpscalerData("DCC Interpolation", None, self),
]
def do_upscale(self, img: Image, selected_model=None):
import math
import numpy as np
from modules.postprocess.dcc import DCC
normalized = np.array(img).astype(np.float32) / 255.0
scale = math.ceil(self.scale)
upscaled = DCC(normalized, scale)
upscaled = (upscaled - upscaled.min()) / (upscaled.max() - upscaled.min())
upscaled = (255.0 * upscaled).astype(np.uint8)
upscaled = Image.fromarray(upscaled)
return upscaled
class UpscalerVIPS(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "VIPS"
self.scalers = [
UpscalerData("VIPS Lanczos 2", None, self),
UpscalerData("VIPS Lanczos 3", None, self),
UpscalerData("VIPS Mitchell", None, self),
UpscalerData("VIPS MagicKernelSharp 2013", None, self),
UpscalerData("VIPS MagicKernelSharp 2021", None, self),
]
def do_upscale(self, img: Image, selected_model=None):
if selected_model is None:
return img
from installer import install
install('pyvips')
try:
import pyvips
except Exception as e:
log.error(f"Upscaler: vips {e}")
return img
vips_image = pyvips.Image.new_from_array(img)
# import numpy as np
# np_image = np.array(img)
# h, w, c = np_image.shape
# np_linear = np_image.reshape(w * h * c)
# vips_image = pyvips.Image.new_from_memory(np_linear.data, w, h, c, 'uchar')
try:
if selected_model is None:
return img
elif selected_model == "VIPS Lanczos 2":
vips_image = vips_image.resize(2, kernel='lanczos2')
elif selected_model == "VIPS Lanczos 3":
vips_image = vips_image.resize(2, kernel='lanczos3')
elif selected_model == "VIPS Mitchell":
vips_image = vips_image.resize(2, kernel='mitchell')
elif selected_model == "VIPS MagicKernelSharp 2013":
vips_image = vips_image.resize(2, kernel='mks2013')
elif selected_model == "VIPS MagicKernelSharp 2021":
vips_image = vips_image.resize(2, kernel='mks2021')
else:
return img
except Exception as e:
log.error(f"Upscaler: vips {e}")
return img
upscaled = Image.fromarray(vips_image.numpy())
return upscaled

87
modules/upscaler_vae.py Normal file
View File

@ -0,0 +1,87 @@
from PIL import Image
from modules.upscaler import Upscaler, UpscalerData
class UpscalerAsymmetricVAE(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "Asymmetric VAE"
self.vae = None
self.selected = None
self.scalers = [
UpscalerData("Asymmetric VAE v1", None, self),
UpscalerData("Asymmetric VAE v2", None, self),
]
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
if self.vae is None or (selected_model != self.selected):
if 'v1' in selected_model:
repo_id = 'Heasterian/AsymmetricAutoencoderKLUpscaler'
else:
repo_id = 'Heasterian/AsymmetricAutoencoderKLUpscaler_v2'
self.vae = diffusers.AsymmetricAutoencoderKL.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir)
self.vae.requires_grad_(False)
self.vae = self.vae.to(device=devices.device, dtype=devices.dtype)
self.vae.eval()
self.selected = selected_model
shared.log.debug(f'Upscaler load: selected="{self.selected}" vae="{repo_id}"')
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)
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())
self.vae = self.vae.to(device=devices.cpu)
return upscaled
class UpscalerWanUpscale(Upscaler):
def __init__(self, dirname=None): # pylint: disable=unused-argument
super().__init__(False)
self.name = "WAN Upscale"
self.vae_encode = None
self.vae_decode = None
self.selected = None
self.scalers = [
UpscalerData("WAN Asymmetric Upscale", None, self),
]
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
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'
self.vae_encode = diffusers.AutoencoderKLWan.from_pretrained(repo_encode, subfolder=subfolder_encode, cache_dir=shared.opts.hfcache_dir)
self.vae_encode.requires_grad_(False)
self.vae_encode = self.vae_encode.to(device=devices.device, dtype=devices.dtype)
self.vae_encode.eval()
repo_decode = 'spacepxl/Wan2.1-VAE-upscale2x'
subfolder_decode = "diffusers/Wan2.1_VAE_upscale2x_imageonly_real_v1"
self.vae_decode = diffusers.AutoencoderKLWan.from_pretrained(repo_decode, subfolder=subfolder_decode, cache_dir=shared.opts.hfcache_dir)
self.vae_decode.requires_grad_(False)
self.vae_decode = self.vae_decode.to(device=devices.device, dtype=devices.dtype)
self.vae_decode.eval()
self.selected = selected_model
shared.log.debug(f'Upscaler load: selected="{self.selected}" encode="{repo_encode}" decode="{repo_decode}"')
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 = self.vae_encode.encode(tensor).latent_dist.mode()
self.vae_encode.to(device=devices.cpu)
self.vae_decode = self.vae_decode.to(device=devices.device)
tensor = self.vae_decode.decode(tensor).sample
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())
return upscaled

View File

@ -37,6 +37,8 @@ import modules.txt2img
import modules.img2img
import modules.upscaler
import modules.upscaler_simple
import modules.upscaler_vae
import modules.upscaler_algo
import modules.extra_networks
import modules.ui_extra_networks
import modules.textual_inversion