Skip to content

Commit

Permalink
waifu2x: unlimited: Add model option; Support CUNet models
Browse files Browse the repository at this point in the history
  • Loading branch information
@nagadomi
nagadomi committed Feb 7, 2023
1 parent c5d26ff commit 2a1e5e3
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 8 deletions.
10 changes: 10 additions & 0 deletions waifu2x/unlimited_waifu2x/public_html/index.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,15 @@ <h1>unlimited:waifu2x</h1>
<input type="file" id="file" name="file" accept="image/*">
</td>
</tr>
<tr>
<td>Model</td>
<td>
<select name="model">
<option value="swin_unet.art" selected>swin_unet / art (202302)</option>
<option value="cunet.art">cunet / art (201811)</option>
</select>
</td>
</tr>
<tr>
<td>JPEG DeNoise</td>
<td>
Expand All @@ -46,6 +55,7 @@ <h1>unlimited:waifu2x</h1>
<option value="2" selected>2x</option>
<option value="4">4x</option>
</select>
<span id="scale-comment">no 4x support</span>
</td>
</tr>
<tr>
Expand Down
66 changes: 58 additions & 8 deletions waifu2x/unlimited_waifu2x/public_html/script.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,23 @@ function gen_arch_config()
scale4x: {scale: 4, offset: 32},
scale1x: {scale: 1, offset: 8}, // bypass for alpha denoise
};
var base = config["swin_unet"];
for (var i = 0; i < 4; ++i) {
config["swin_unet"]["art"]["noise" + i + "_scale2x"] = {scale: 2, offset: 16};
config["swin_unet"]["art"]["noise" + i + "_scale4x"] = {scale: 4, offset: 32};
config["swin_unet"]["art"]["noise" + i] = {scale: 1, offset: 8};
base["art"]["noise" + i + "_scale2x"] = {scale: 2, offset: 16};
base["art"]["noise" + i + "_scale4x"] = {scale: 4, offset: 32};
base["art"]["noise" + i] = {scale: 1, offset: 8};
}
config["cunet"] = {art: {}}
config["cunet"]["art"] = {
scale2x: {scale: 2, offset: 36},
scale1x: {scale: 1, offset: 28}, // bypass for alpha denoise
};
var base = config["cunet"];
for (var i = 0; i < 4; ++i) {
base["art"]["noise" + i + "_scale2x"] = {scale: 2, offset: 36};
base["art"]["noise" + i] = {scale: 1, offset: 28};
}

return config;
}

Expand All @@ -24,6 +36,23 @@ const CONFIG = {
if ((arch in this.arch) && (style in this.arch[arch]) && (method in this.arch[arch][style])) {
config = this.arch[arch][style][method];
config["path"] = `models/${arch}/${style}/${method}.onnx`;
if (arch == "swin_unet") {
config.calc_tile_size = function (tile_size) {
while (true) {
if ((tile_size - 16) % 12 == 0 && (tile_size - 16) % 16 == 0) {
break;
}
tile_size += 1;
}
return tile_size;
};
} else if (arch == "cunet") {
config.calc_tile_size = function (tile_size) {
tile_size = tile_size + (config.offset - 16) * 2;
tile_size -= tile_size % 4;
return tile_size;
};
}
return config;
} else {
return null;
Expand Down Expand Up @@ -507,7 +536,8 @@ $(function () {
console.log("Already running");
return;
}
var style = "art";
var model_name = $("select[name=model]").val();
var [arch, style] = model_name.split(".");
var scale = parseInt($("select[name=scale]").val());
var noise_level = parseInt($("select[name=noise_level]").val());
var method;
Expand All @@ -517,28 +547,25 @@ $(function () {
return;
}
method = "noise" + noise_level;
arch = "swin_unet";
} else if (scale == 2) {
if (noise_level == -1) {
method = "scale2x";
} else {
method = "noise" + noise_level + "_scale2x";
}
arch = "swin_unet";
} else if (scale == 4) {
if (noise_level == -1) {
method = "scale4x";
} else {
method = "noise" + noise_level + "_scale4x";
}
arch = "swin_unet";
}
const config = CONFIG.get_config(arch, style, method);
if (config == null) {
set_message("(・A・) Model Not found!");
return;
}
const tile_size = parseInt($("select[name=tile_size]").val());
const tile_size = config.calc_tile_size(parseInt($("select[name=tile_size]").val()));
const tile_random = $("input[name=tile_random]").prop("checked");
const tta_level = parseInt($("select[name=tta]").val());

Expand Down Expand Up @@ -712,6 +739,12 @@ $(function () {

function restore_from_cookie()
{
if ($.cookie("model")) {
$("select[name=model]").val($.cookie("model"));
if (!$("select[name=model]").val()) {
$("select[name=model]").val("swin_unet.art");
}
}
if ($.cookie("noise_level")) {
$("select[name=noise_level]").val($.cookie("noise_level"));
}
Expand All @@ -733,6 +766,23 @@ $(function () {
};
restore_from_cookie();

$("select[name=model]").change(() => {
var model = $("select[name=model]").val();
var [arch, style] = model.split(".");
$.cookie("model", model, {expires: g_expires});
if (arch == "swin_unet") {
$("select[name=scale]").children("option[value=4]").show();
$("#scale-comment").hide();
} else {
var scale = $("select[name=scale]").val();
$("select[name=scale]").children("option[value=4]").hide();
$("#scale-comment").show();
if (scale == "4") {
$("select[name=scale]").val("2");
}
}
});
$("select[name=model]").trigger("change");
$("select[name=noise_level]").change(() => {
$.cookie("noise_level", $("select[name=noise_level]").val(), {expires: g_expires});
});
Expand Down
2 changes: 2 additions & 0 deletions waifu2x/unlimited_waifu2x/public_html/style.css
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,5 @@ tr {border: 1px solid;}
td {border: 3px groove; padding: 4px 8px 4px 8px;}
select { padding: 4px; }
input[type=file] { padding: 4px; width: 80%;}
#scale-comment { font-size: 0.8em; color: green; margin-left: 1em; }

0 comments on commit 2a1e5e3

Please sign in to comment.