Merge pull request #12 from Filexor/colorfix_fix

Accept result.images with 4 components
2023-12-04 19:12:35 +08:00 · 2023-12-04 19:12:35 +08:00 · e3a7c69c08
parent 4499d79649 06eb56b4b8
commit e3a7c69c08
1 changed files with 8 additions and 0 deletions
--- a/scripts/stablesr.py
+++ b/scripts/stablesr.py
@ -251,6 +251,14 @@ class Script(scripts.Script):

            fixed_images = []
            # fix the color
+
+            for i in range(len(result.images)):
+                if result.images[i].mode == "RGBA":
+                    # wavelet_color_fix expects shape[1] of input image is 3
+                    result.images[i] = result.images[i].convert(mode = "RGB")
+                    pass
+                pass
+
            color_fix_func = wavelet_color_fix if color_fix == 'Wavelet' else adain_color_fix
            for i in range(len(result.images)):
                try: