Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Anime4K-WebGPU project #221

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add Anime4K-WebGPU project #221

wants to merge 1 commit into from

Conversation

plasmas
Copy link

@plasmas plasmas commented Dec 21, 2023

Hi,

We created a WebGPU implementation for Anime4K. It features cross-platform APIs for integration into any WebGPU code, and can be used for real-time video upscaling, as well as any texture upscaling. We created a live demo and released an NPM package:

The package currently has a selected number of pipelines for each category (deblur, upscale, etc.), which run smoothly with any modern NVIDIA GPUs. It also works in Chrome on Macs, but may experience frame drops with heavy pipelines. Optimizations are still underway as more browser vendors support WebGPU.

I added the project to README.md. Any question is welcome!

@arianaa30
Copy link

@plasmas how easy is it to add other upscale model sizes like M, L, VL as well? It currently has 2xUUL only.

@plasmas
Copy link
Author

plasmas commented Mar 4, 2024

@plasmas how easy is it to add other upscale model sizes like M, L, VL as well? It currently has 2xUUL only.

As of now, conversion from original glsl shader to wgsl shaders is automated, but putting together wgsl shaders into a single pipeline requires knowledge about the model and webGPU.

A good example is Upscale_Shader.glsl. The steps to create its webGPU pipeline are:

  1. Run shader.py on this glsl shader to generate wgsl files, each containing a conv2d layer.
  2. Create UpscaleCNN.ts, which implements Anime4KPipeline. Helper pipelines like Conv2d, Depth2Space, and Overlay are created and chained according the the glsl stages / model arch. Shaders generated in the previous step are used to create Conv2d helper pipelines.

Note that shader.py may generate glitchy wgsl shaders for complex conv2d stages that also involves upscaling (e.g. Anime4K-v4.1-Upscale-GAN-x3-(L)-Conv-4x3x3x24 in Upscale_GAN_x3_L.glsl), and will need additional tweaking.

This way to convert from glsl to wgsl can be buggy. A better way would be converting from tensorflow models (#220) directly. Hope this helps.

@arianaa30
Copy link

arianaa30 commented Mar 4, 2024

@plasmas Yeah I'm only interested in the simple upscalers (no GANs). But one thing: I measured the SSIM of your CNNx2UL (downloaded the Canvas on web demo), and realized it is much lower (0.77) than that of the PyTorch converted model (0.97) in #220 . So aren't they really equivalent?! I actually measured the upscale-VL on PyTorch but I assume UL should even be higher.

@plasmas
Copy link
Author

plasmas commented Mar 4, 2024

@plasmas Yeah I'm only interested in the simple upscalers (no GANs). But one thing: I measured the SSIM of your CNNx2UL (downloaded the Canvas on web demo), and realized it is much lower (0.77) than that of the PyTorch converted model (0.97) in #220 . So aren't they really equivalent?! I actually measured the upscale-VL on PyTorch but I assume UL should even be higher.

Anime4K_Upscale_CNN_x2_UL.glsl was followed closely to create CNNx2UL pipeline, and weights and model architecture should have little difference from the original glsl shader, after reviewing. One possibility is that FP16 is used to store intermediate tensors due to webGPU limits, although computation in each stage is still in FP32.

Further comparison between the recovered PyTorch model and glsl/wgsl might be needed. Please consider creating an issue under Anime4K-WebGPU for more questions.

@arianaa30
Copy link

@plasmas Thanks. I created an issue under the original repo.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants