Replace the GPU Lock with a queue-based approach
This commit is contained in:
parent
9a3101d358
commit
11f50facb9
|
@ -4,7 +4,7 @@ plugins {
|
|||
}
|
||||
|
||||
group 'moe.lemonsh'
|
||||
version '0.1'
|
||||
version '0.2'
|
||||
|
||||
jar {
|
||||
manifest {
|
||||
|
|
|
@ -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;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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;
|
||||
|
|
Reference in a new issue