Skip to content

Commit

Permalink
支持在某种情况下预先使用常规算法放大以减少调用一次Real-ESRGAN
Browse files Browse the repository at this point in the history 
#80
  • Loading branch information
@TransparentLC
TransparentLC committed May 20, 2024
1 parent de07ed2 commit b389353
Show file tree
Hide file tree
Showing 4 changed files with 42 additions and 11 deletions.
4 changes: 4 additions & 0 deletions i18n.ini
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ EnableLossyMode = 使用有损压缩（JPEG/WebP）
LossyModeQuality = 有损压缩质量（0-100）
CustomCommand = 自定义压缩/后期处理命令
EnableIgnoreError = 在批处理过程中忽略错误并继续处理
EnablePreupscale = 尝试预先使用常规算法放大
ViewREGUISource = 查看源代码
ViewRESource = 查看 Real-ESRGAN 介绍
ViewAdditionalModel = 下载附加模型
Expand Down Expand Up @@ -65,6 +66,7 @@ EnableLossyMode = 使用有損壓縮（JPEG/WebP）
LossyModeQuality = 有損壓縮質量（0-100）
CustomCommand = 自定義壓縮/後期處理命令
EnableIgnoreError = 在批處理過程中忽略錯誤並繼續處理
EnablePreupscale = 嘗試預先使用常規算法放大
ViewREGUISource = 查看源代碼
ViewRESource = 查看 Real-ESRGAN 介紹
ViewAdditionalModel = 下載附加模型
Expand Down Expand Up @@ -109,6 +111,7 @@ EnableLossyMode = 使用有損壓縮（JPEG/WebP）
LossyModeQuality = 有損壓縮質量（0-100）
CustomCommand = 自訂壓縮/後期處理命令
EnableIgnoreError = 在批處理過程中忽略錯誤並繼續處理
EnablePreupscale = 嘗試預先使用常規算法放大
ViewREGUISource = 查看原始碼
ViewRESource = 查看 Real-ESRGAN 介紹
ViewAdditionalModel = 下載附加模型
Expand Down Expand Up @@ -153,6 +156,7 @@ EnableLossyMode = Enable lossy compression (JPEG/WebP)
LossyModeQuality = Lossy compression quality (0-100)
CustomCommand = Custom compression/post-processing command
EnableIgnoreError = Ignore error and continue during batch processing
EnablePreupscale = Try to pre-upscale with general algorithm
ViewREGUISource = View source code
ViewRESource = About Real-ESRGAN
ViewAdditionalModel = Download additional models
Expand Down
9 changes: 9 additions & 0 deletions main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,7 @@ def outputPathTraceCallback(var: tk.IntVar | tk.StringVar, index: str, mode: str
self.varboolOptimizeGIF = tk.BooleanVar(value=self.config['Config'].getboolean('OptimizeGIF'))
self.varboolLossyMode = tk.BooleanVar(value=self.config['Config'].getboolean('LossyMode'))
self.varboolIgnoreError = tk.BooleanVar(value=self.config['Config'].getboolean('IgnoreError'))
self.varboolPreupscale = tk.BooleanVar(value=self.config['Config'].getboolean('Preupscale'))
self.varboolProcessing = tk.BooleanVar(value=False)
self.varboolProcessingPaused = tk.BooleanVar(value=False)
self.varstrCustomCommand = tk.StringVar(value=self.config['Config'].get('CustomCommand'))
Expand Down Expand Up @@ -182,6 +183,7 @@ def outputPathTraceCallback(var: tk.IntVar | tk.StringVar, index: str, mode: str
self.varstrLabelGIFOptimizeTransparency = tk.StringVar(value=i18n.getTranslatedString('GIFOptimizeTransparency'))
self.varstrLabelEnableLossyMode = tk.StringVar(value=i18n.getTranslatedString('EnableLossyMode'))
self.varstrLabelEnableIgnoreError = tk.StringVar(value=i18n.getTranslatedString('EnableIgnoreError'))
self.varstrLabelEnablePreupscale = tk.StringVar(value=i18n.getTranslatedString('EnablePreupscale'))
self.varstrLabelViewREGUISource = tk.StringVar(value=i18n.getTranslatedString('ViewREGUISource'))
self.varstrLabelViewRESource = tk.StringVar(value=i18n.getTranslatedString('ViewRESource'))
self.varstrLabelViewAdditionalModel = tk.StringVar(value=i18n.getTranslatedString('ViewAdditionalModel'))
Expand Down Expand Up @@ -294,6 +296,8 @@ def setupWidgets(self):
self.checkLossyMode.pack(padx=10, pady=5, fill=tk.X)
self.checkIgnoreError = ttk.Checkbutton(self.frameAdvancedConfigRight, textvariable=self.varstrLabelEnableIgnoreError, style='Switch.TCheckbutton', variable=self.varboolIgnoreError)
self.checkIgnoreError.pack(padx=10, pady=5, fill=tk.X)
self.checkPreupscale = ttk.Checkbutton(self.frameAdvancedConfigRight, textvariable=self.varstrLabelEnablePreupscale, style='Switch.TCheckbutton', variable=self.varboolPreupscale)
self.checkPreupscale.pack(padx=10, pady=5, fill=tk.X)
self.comboLanguage = ttk.Combobox(self.frameAdvancedConfigRight, state='readonly', values=tuple(i18n.locales_map.keys()))
self.comboLanguage.current(i18n.get_current_locale_display_name())
self.comboLanguage.pack(padx=10, pady=5, fill=tk.X)
Expand Down Expand Up @@ -361,6 +365,7 @@ def change_app_lang(self, event: tk.Event):
self.varstrLabelGIFOptimizeTransparency.set(i18n.getTranslatedString('GIFOptimizeTransparency'))
self.varstrLabelEnableLossyMode.set(i18n.getTranslatedString('EnableLossyMode'))
self.varstrLabelEnableIgnoreError.set(i18n.getTranslatedString('EnableIgnoreError'))
self.varstrLabelEnablePreupscale.set(i18n.getTranslatedString('EnablePreupscale'))
self.varstrLabelViewREGUISource.set(i18n.getTranslatedString('ViewREGUISource'))
self.varstrLabelViewRESource.set(i18n.getTranslatedString('ViewRESource'))
self.varstrLabelViewAdditionalModel.set(i18n.getTranslatedString('ViewAdditionalModel'))
Expand All @@ -385,6 +390,8 @@ def close(self):
'UseTTA': self.varboolUseTTA.get(),
'OptimizeGIF': self.varboolOptimizeGIF.get(),
'LossyMode': self.varboolLossyMode.get(),
'IgnoreError': self.varboolIgnoreError.get(),
'Preupscale': self.varboolPreupscale.get(),
'CustomCommand': self.varstrCustomCommand.get(),
'AppLanguage': i18n.current_language
}
Expand Down Expand Up @@ -621,6 +628,7 @@ def getConfigParams(self) -> param.REConfigParams:
self.tileSize[self.varintTileSizeIndex.get()],
self.varintGPUID.get(),
self.varboolUseTTA.get(),
self.varboolPreupscale.get(),
self.varstrCustomCommand.get().strip(),
)

Expand Down Expand Up @@ -665,6 +673,7 @@ def init_config_and_model_paths() -> tuple[configparser.ConfigParser, list[str]]
'OptimizeGIF': False,
'LossyMode': False,
'IgnoreError': False,
'Preupscale': False,
'CustomCommand': '',
'AppLanguage': locale.getdefaultlocale()[0],
})
Expand Down
1 change: 1 addition & 0 deletions param.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,4 +17,5 @@ class REConfigParams(typing.NamedTuple):
tileSize: int
gpuID: int
useTTA: bool
preupscale: bool
customCommand: str
39 changes: 28 additions & 11 deletions task.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import collections
import secrets
import subprocess
import io
import math
import os
import re
import shlex
Expand All @@ -18,9 +18,6 @@
import define
import param

def buildTempPath(ext: str) -> str:
return os.path.join(tempfile.gettempdir(), secrets.token_urlsafe(12) + ext)

class AbstractTask:
def __init__(self, outputCallback: typing.Callable[[str], None]) -> None:
self.outputCallback = outputCallback
Expand Down Expand Up @@ -52,7 +49,7 @@ def run(self) -> None:
srcWidth, srcHeight = img.size
srcRatio = srcWidth / srcHeight
if img.mode == 'P':
self.inputPath = buildTempPath('.png')
self.inputPath = tempfile.mktemp('.png')
img.convert('RGBA').save(self.inputPath)
self.removeInput = True
match self.config.resizeMode:
Expand All @@ -65,6 +62,26 @@ def run(self) -> None:
case param.ResizeMode.HEIGHT:
dstHeight = self.config.resizeModeValue
dstWidth = round(dstHeight * srcRatio)
inputPathPreupscaled: str = None
if self.config.preupscale:
match self.config.resizeMode:
case param.ResizeMode.RATIO:
scaleRatio = self.config.resizeModeValue
case param.ResizeMode.WIDTH:
scaleRatio = self.config.resizeModeValue / srcWidth
case param.ResizeMode.HEIGHT:
scaleRatio = self.config.resizeModeValue / srcHeight
frac, intg = math.modf(math.log(scaleRatio, self.config.modelFactor))
preWidth = math.ceil(dstWidth / (self.config.modelFactor ** intg))
preHeight = math.ceil(dstHeight / (self.config.modelFactor ** intg))
if frac < .5 and (srcWidth != preWidth or srcHeight != preHeight):
self.outputCallback(f'Pre-upscale from {srcWidth}x{srcHeight} to {preWidth}x{preHeight}.\n')
inputPathPreupscaled = tempfile.mktemp('.webp' if os.path.splitext(self.inputPath)[1] == '.webp' else '.png')
with Image.open(self.inputPath) as img:
resized = img.resize((preWidth, preHeight), Image.LANCZOS)
resized.save(inputPathPreupscaled, lossless=True)
resized.close()
srcWidth, srcHeight = preWidth, preHeight
scalePass = 0
while srcWidth < dstWidth and srcHeight < dstHeight:
scalePass += 1
Expand All @@ -75,7 +92,7 @@ def run(self) -> None:
# input -> temp0 -> output
# input -> temp0 -> temp1 -> output
outputExt = os.path.splitext(self.outputPath)[1]
files = (self.inputPath, *(buildTempPath(outputExt) for _ in range(scalePass)))
files = (inputPathPreupscaled or self.inputPath, *(tempfile.mktemp(outputExt) for _ in range(scalePass)))
for i in range(len(files) - 1):
inputPath, outputPath = files[i:(i + 2)]
alphaOverridePath = None
Expand Down Expand Up @@ -132,7 +149,7 @@ def run(self) -> None:
self.outputCallback(line)
if p.returncode:
raise subprocess.CalledProcessError(p.returncode, cmd)
if i > 0 or self.removeInput:
if i > 0 or inputPath == inputPathPreupscaled or self.removeInput:
os.remove(inputPath)
if alphaOverridePath:
shutil.move(alphaOverridePath, outputPath)
Expand All @@ -146,7 +163,7 @@ def run(self) -> None:
else:
with Image.open(files[-1]) as img:
self.outputCallback(f'Downsample from {img.size[0]}x{img.size[1]} to {dstWidth}x{dstHeight}.\n')
resized: Image.Image = img.resize((dstWidth, dstHeight), self.config.downsample)
resized = img.resize((dstWidth, dstHeight), self.config.downsample)
resized.save(self.outputPath)
resized.close()
if scalePass:
Expand Down Expand Up @@ -233,8 +250,8 @@ def run(self) -> None:
with Image.open(self.inputPath) as img:
for f in ImageSequence.Iterator(img):
f: Image.Image
frameSrcPath = buildTempPath('.png' if self.optimizeTransparency else '.webp')
frameDstPath = buildTempPath('.png' if self.optimizeTransparency else '.webp')
frameSrcPath = tempfile.mktemp('.png' if self.optimizeTransparency else '.webp')
frameDstPath = tempfile.mktemp('.png' if self.optimizeTransparency else '.webp')
d = f.info.get('duration', 0)
if self.optimizeTransparency:
f = f.convert('RGBA')
Expand All @@ -251,7 +268,7 @@ def run(self) -> None:
self.progressValue[2] += 1
self.progressValue[2] -= 1
if self.config.customCommand:
t = buildTempPath('.gif')
t = tempfile.mktemp('.gif')
tasks.append(MergeGIFTask(self.outputCallback, t, frames, durations, self.optimizeTransparency))
tasks.append(CustomCompressTask(self.outputCallback, t, self.outputPath, self.config.customCommand, True))
else:
Expand Down

0 comments on commit b389353

Please sign in to comment.