Use Spandrel for upscaling and face restoration architectures#14425
Use Spandrel for upscaling and face restoration architectures#14425AUTOMATIC1111 merged 7 commits intoAUTOMATIC1111:devfrom
Conversation
b1a61e9 to
e61c70b
Compare
ecee8df to
8bfe7cf
Compare
|
|
||
| shared.opts.add_option("SWIN_tile", shared.OptionInfo(192, "Tile size for all SwinIR.", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}, section=('upscaling', "Upscaling"))) | ||
| shared.opts.add_option("SWIN_tile_overlap", shared.OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}, section=('upscaling', "Upscaling"))) | ||
| if int(torch.__version__.split('.')[0]) >= 2 and platform.system() != "Windows": # torch.compile() require pytorch 2.0 or above, and not on Windows |
There was a problem hiding this comment.
We're always torch >= 2.0, and we now just try compile without checking the platform.
.github/workflows/run_tests.yaml
Outdated
| env: | ||
| IGNORE_CMD_ARGS_ERRORS: "1" |
There was a problem hiding this comment.
The new test will fail without this since a module attempts to read the pytest options as regular webui arguments.
| img, | ||
| tile_size=opts.ESRGAN_tile, | ||
| tile_overlap=opts.ESRGAN_tile_overlap, | ||
| # TODO: `outscale`? |
There was a problem hiding this comment.
Do we need to downscale too-large images here according to info.scale? AIUI, there might be some other process that also does that?
modules/upscaler_utils.py
Outdated
| return Image.fromarray(output, 'RGB') | ||
|
|
||
|
|
||
| def upscale_with_model(model, img: Image.Image, *, tile_size: int, tile_overlap: int = 0): |
There was a problem hiding this comment.
This is used by both esrgan and realesrgan.
| def inference( | ||
| img, | ||
| model, | ||
| *, | ||
| tile: int, | ||
| tile_overlap: int, | ||
| window_size: int, | ||
| scale: int, | ||
| device, | ||
| ): |
There was a problem hiding this comment.
This smells like tiled_upscale, but with tile overlap handled by weight scaling.
ce58f5d to
37458e6
Compare
|
Oh yeah, I've been using this PR for a couple days now; it works. |
|
@gel-crabs Thanks for trying it out! I (force-)pushed this branch to update spandrel to a newer version, as well as add experimental support for HAT upscalers, if you want to try that out. (You'll need to bring your own models and put them in |
It works! Admittedly it has issues with deepcache where it adds black splotches to the image during hires fix, but otherwise working. I tried to hack in support for DAT as well by copying hat_model.py and replacing HAT with DAT, but it just made the image go full black. Edit: It actually has nothing to do with deepcache, or any extensions at all. I'm going to try testing with different models. I tried with a different 4x HAT upscaler and it gives full black images, so the HAT support doesn't seem to be working correctly. |
|
I'm generally not pumped about adding new dependencies, but this removes a lot of code we just copy pasted, so that seems nice. Some questions:
|
I think this actually leads to less dependencies in total (I'll run the numbers later). The Spandrel folks seem nice and responsive too. :)
Autogenerated by PyCharm when refactoring code. Will yeet, my bad.
Also accidentally added to this PR (since I was tired of having a gazillion WebUI tabs get auto-opened), my bad. Will yeet.
I can also add an
Since we do I'll add a simple "output image was different" check!
Good question - since Spandrel auto-detects the model arch from the checkpoint, it'd happily load it, and maybe fail with a parameter error down the line when we try to call the architecture with kwargs it doesn't get. I can add
I did, on my machine (Macbook). |
|
Looks like SwinIR x2 is not working now. I get this in any model: |
|
@wcde Thanks, I'll take a peek – what's your SwinIR tile size and overlap setting, and the size of the image you're trying to upscale? |
|
In code hardcoded scale to 4. Second problem - model is loaded with dtype Which give: |
|
@wcde In fairness, I'll take a look at the |
|
I guess it will happen with a lot of extensions after updating. Maybe it should be mentioned in changelog? |
Description
This PR yeets most of the copy-pasted or otherwise vendored model architectures in favor of just using
Spandrel.
Converted models are:
Not converted is LDSR; it doesn't exist in Spandrel.
There's still some more cleanup that could be done – there are multiple implementations of tiled inference right now, for one, and the model loading/downloading/... code is kind of a mess (should continue where I left off with Upscaler model loading cleanup #10823), but I'll hold off on that for this PR.
As an added bonus, this adds (experimental, works-on-my-machine) support for HAT upscaling models.
Screenshots/videos:
No visual changes. This seems to Work On My Machine but it'd be lovely if someone else tried this out too.
Checklist: