lemonsh/waifu2xlemon
Archived
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

Replace the GPU Lock with a queue-based approach

Browse source
This commit is contained in:
lemon-sh 2021-08-20 23:59:13 +02:00
parent 9a3101d358
commit 11f50facb9
3 changed files with 76 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
build.gradle
View file

@ -4,7 +4,7 @@ plugins {
}
group 'moe.lemonsh'
version '0.1'
version '0.2'
jar {
manifest {

30
src/main/java/moe/lemonsh/waifu2xlemon/App.java
View file

@ -6,6 +6,7 @@ import java.io.*;
import java.net.URLConnection;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
import java.util.concurrent.CountDownLatch;
import static spark.Spark.*;
@ -33,7 +34,7 @@ public class App {
}
}
Waifu2x waifu2x = new Waifu2x(envexec);
threadPool(4, 1, -1);
new Thread(waifu2x::runQueue).start();
try {
port(Integer.parseInt(envport));
} catch (NumberFormatException e) {
@ -106,18 +107,29 @@ public class App {
case 2 -> outputExtension = ".webp";
default -> halt(400);
}
Waifu2x.WaifuResult waifuResult;
waifuResult = waifu2x.run(filePart.getInputStream(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0'));
if (waifuResult.success) {
resp.raw().setContentType(URLConnection.guessContentTypeFromName(waifuResult.filename));
resp.raw().setHeader("Content-Disposition","attachment; filename="+waifuResult.filename);
resp.raw().setContentLengthLong(waifuResult.size);
waifuResult.image.transferTo(resp.raw().getOutputStream());
final var latch = new CountDownLatch(1);
var resultRef = new Object() {
Waifu2x.WaifuResult waifuResult;
};
var task = new Waifu2x.WaifuTask(filePart.getInputStream().readAllBytes(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0'), result -> {
resultRef.waifuResult = result;
latch.countDown();
});
switch (waifu2x.addTask(task)) {
case 1 -> { return "Invalid image (too big or in non-supported format)"; }
case 2 -> { return "The processing queue is full. Please try again later."; }
}
latch.await();
if (resultRef.waifuResult.success) {
resp.raw().setContentType(URLConnection.guessContentTypeFromName(resultRef.waifuResult.filename));
resp.raw().setHeader("Content-Disposition","attachment; filename="+resultRef.waifuResult.filename);
resp.raw().setContentLengthLong(resultRef.waifuResult.size);
resultRef.waifuResult.image.transferTo(resp.raw().getOutputStream());
resp.raw().getOutputStream().flush();
resp.raw().getOutputStream().close();
return null;
}
return "Error.\n"+waifuResult.stdout;
return "Error.\n"+resultRef.waifuResult.stdout;
});
}
}

76
src/main/java/moe/lemonsh/waifu2xlemon/Waifu2x.java
View file

@ -3,21 +3,24 @@ package moe.lemonsh.waifu2xlemon;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.imageio.ImageIO;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.concurrent.BlockingDeque;
import java.util.concurrent.LinkedBlockingDeque;
public class Waifu2x {
public static class WaifuResult {
public boolean success;
public String stdout;
public InputStream image;
public String filename;
public long size;
@FunctionalInterface
public interface WaifuCallback {
void exec(WaifuResult result);
}
private final Object GPULock;
public record WaifuTask(byte[] inputImage, String inputExtension, String outputExtension,
char noise, char scale, WaifuCallback callback) {}
private final BlockingDeque<WaifuTask> taskQueue;
private final Path ExecutablePath;
private final Logger log = LoggerFactory.getLogger(Waifu2x.class);
@ -27,11 +30,43 @@ public class Waifu2x {
throw new IllegalArgumentException("'%s' is not a valid path.".formatted(wf2xExecutablePath));
} else {
ExecutablePath = executable.toPath().toAbsolutePath();
GPULock = new Object();
taskQueue = new LinkedBlockingDeque<>(2);
}
}
public WaifuResult run(InputStream inputImage, String inputExtension, String outputExtension, char noise, char scale) throws IOException, InterruptedException {
public synchronized void runQueue() {
while (true) {
try {
WaifuTask task = taskQueue.take();
log.info("New job! %s image with options noise=%c scale=%c output=%s".formatted(task.inputExtension, task.noise, task.scale, task.outputExtension));
var result = run(task.inputImage, task.inputExtension, task.outputExtension, task.noise, task.scale);
log.info("Job completed! Result:\n"+result.stdout);
task.callback.exec(result);
} catch (IOException | InterruptedException e) {
System.exit(10);
}
}
}
private boolean imageIllegal(byte[] fullImage) throws IOException {
// this can most likely be optimized
var image = ImageIO.read(new ByteArrayInputStream(fullImage));
return image == null || image.getWidth() > 2000 || image.getHeight() > 2000;
}
public int addTask(WaifuTask task) throws IOException {
return imageIllegal(task.inputImage) ? 1 : taskQueue.offer(task) ? 0 : 2;
}
static class WaifuResult {
public boolean success;
public String stdout;
public InputStream image;
public String filename;
public long size;
}
private WaifuResult run(byte[] inputImage, String inputExtension, String outputExtension, char noise, char scale) throws IOException, InterruptedException {
var inputFile = new File("wfx_input" + inputExtension).toPath().toAbsolutePath();
var outputFile = new File("wfx_output" + outputExtension);
var outputFilePath = outputFile.toPath().toAbsolutePath();
@ -43,19 +78,16 @@ public class Waifu2x {
return buf;
}
};
synchronized (GPULock) {
log.info("Starting processing %s image with options noise=%c scale=%c output=%s".formatted(inputExtension, noise, scale, outputExtension));
Files.deleteIfExists(inputFile);
Files.copy(inputImage, inputFile);
var waifuProcess = new ProcessBuilder(ExecutablePath.toString(), "-s", Character.toString(scale), "-n",
Character.toString(noise), "-i", inputFile.toString(), "-o", outputFilePath.toString()).start();
waifuProcess.waitFor();
Files.delete(inputFile);
log.info(result.stdout = new String(waifuProcess.getErrorStream().readAllBytes(), StandardCharsets.UTF_8));
if (!(result.success = outputFile.isFile())) return result;
Files.copy(outputFilePath, outputBytes);
Files.delete(outputFilePath);
}
Files.deleteIfExists(inputFile);
Files.write(inputFile, inputImage);
var waifuProcess = new ProcessBuilder(ExecutablePath.toString(), "-s", Character.toString(scale), "-n",
Character.toString(noise), "-i", inputFile.toString(), "-o", outputFilePath.toString()).start();
waifuProcess.waitFor();
Files.delete(inputFile);
result.stdout = new String(waifuProcess.getErrorStream().readAllBytes(), StandardCharsets.UTF_8);
if (!(result.success = outputFile.isFile())) return result;
Files.copy(outputFilePath, outputBytes);
Files.delete(outputFilePath);
result.image = new ByteArrayInputStream(outputBytes.toByteArray(), 0, outputBytes.size());
result.size = outputBytes.size();
return result;