Upload 16 files
Browse files- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/README.md +35 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/__init__.py +48 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/gradio.py +1 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/__init__.py +0 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/devices.py +2 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/images.py +8 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/processing.py +213 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/scripts.py +2 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/shared.py +24 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/upscaler.py +32 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/nodes.py +234 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/repositories/__init__.py +14 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/repositories/ultimate_sd_upscale/README.md +119 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/repositories/ultimate_sd_upscale/scripts/ultimate-upscale.py +569 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/usdu_patch.py +71 -0
- ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/utils.py +460 -0
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/README.md
ADDED
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# ComfyUI_UltimateSDUpscale
|
2 |
+
|
3 |
+
[ComfyUI](https://github.com/comfyanonymous/ComfyUI) nodes for the [Ultimate Stable Diffusion Upscale script by Coyote-A](https://github.com/Coyote-A/ultimate-upscale-for-automatic1111). This is a wrapper for the script used in the A1111 extension.
|
4 |
+
|
5 |
+
## Installation
|
6 |
+
|
7 |
+
Enter the following command from the commandline starting in ComfyUI/custom_nodes/
|
8 |
+
```
|
9 |
+
git clone https://github.com/ssitu/ComfyUI_UltimateSDUpscale --recursive
|
10 |
+
```
|
11 |
+
|
12 |
+
## Usage
|
13 |
+
|
14 |
+
Nodes can be found in the node menu under `image/upscaling`:
|
15 |
+
|
16 |
+
|Node|Description|
|
17 |
+
| --- | --- |
|
18 |
+
| Ultimate SD Upscale | The primary node that has the most of the inputs as the original extension script. |
|
19 |
+
| Ultimate SD Upscale <br>(No Upscale) | Same as the primary node, but without the upscale inputs and assumes that the input image is already upscaled. Use this if you already have an upscaled image or just want to do the tiled sampling. |
|
20 |
+
| Ultimate SD Upscale <br>(Custom Sample) | Same as the primary node, but has additional inputs for a custom sampler and custom sigmas. Both must be provided if one is used. If neither is provided, the widgets (the settings below the input slots) for the sampler and step/denoise settings will be used, like in the base USDU node. |
|
21 |
+
|
22 |
+
---
|
23 |
+
|
24 |
+
Details about most of the parameters can be found [here](https://github.com/Coyote-A/ultimate-upscale-for-automatic1111/wiki/FAQ#parameters-descriptions).
|
25 |
+
|
26 |
+
Parameters not found in the original repository:
|
27 |
+
|
28 |
+
* `upscale_by` The number to multiply the width and height of the image by. If you want to specify an exact width and height, use the "No Upscale" version of the node and perform the upscaling separately (e.g., ImageUpscaleWithModel -> ImageScale -> UltimateSDUpscaleNoUpscale).
|
29 |
+
* `force_uniform_tiles` If enabled, tiles that would be cut off by the edges of the image will expand the tile using the rest of the image to keep the same tile size determined by `tile_width` and `tile_height`, which is what the A1111 Web UI does. If disabled, the minimal size for tiles will be used, which may make the sampling faster but may cause artifacts due to irregular tile sizes.
|
30 |
+
|
31 |
+
## Examples
|
32 |
+
|
33 |
+
#### Using the ControlNet tile model:
|
34 |
+
|
35 |
+
![image](https://github.com/ssitu/ComfyUI_UltimateSDUpscale/assets/57548627/64f8d3b2-10ae-45ee-9f8a-40b798a51655)
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/__init__.py
ADDED
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import sys
|
2 |
+
import os
|
3 |
+
|
4 |
+
# Remove other custom_node paths from sys.path to avoid conflicts
|
5 |
+
custom_node_paths = [path for path in sys.path if "custom_node" in path]
|
6 |
+
original_sys_path = sys.path.copy()
|
7 |
+
for path in custom_node_paths:
|
8 |
+
sys.path.remove(path)
|
9 |
+
|
10 |
+
# Add this repository's path to sys.path for third-party imports
|
11 |
+
repo_dir = os.path.dirname(os.path.realpath(__file__))
|
12 |
+
sys.path.insert(0, repo_dir)
|
13 |
+
original_modules = sys.modules.copy()
|
14 |
+
|
15 |
+
# Place aside potentially conflicting modules
|
16 |
+
modules_used = [
|
17 |
+
"modules",
|
18 |
+
"modules.devices",
|
19 |
+
"modules.images",
|
20 |
+
"modules.processing",
|
21 |
+
"modules.scripts",
|
22 |
+
"modules.shared",
|
23 |
+
"modules.upscaler",
|
24 |
+
"utils",
|
25 |
+
]
|
26 |
+
original_imported_modules = {}
|
27 |
+
for module in modules_used:
|
28 |
+
if module in sys.modules:
|
29 |
+
original_imported_modules[module] = sys.modules.pop(module)
|
30 |
+
|
31 |
+
# Proceed with node setup
|
32 |
+
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
33 |
+
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
34 |
+
|
35 |
+
# Clean up imports
|
36 |
+
# Remove any new modules
|
37 |
+
modules_to_remove = []
|
38 |
+
for module in sys.modules:
|
39 |
+
if module not in original_modules:
|
40 |
+
modules_to_remove.append(module)
|
41 |
+
for module in modules_to_remove:
|
42 |
+
del sys.modules[module]
|
43 |
+
|
44 |
+
# Restore original modules
|
45 |
+
sys.modules.update(original_imported_modules)
|
46 |
+
|
47 |
+
# Restore original sys.path
|
48 |
+
sys.path = original_sys_path
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/gradio.py
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
# Empty gradio module for the ultimate-upscale.py import because gradio is not needed
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/__init__.py
ADDED
File without changes
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/devices.py
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
def torch_gc():
|
2 |
+
pass
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/images.py
ADDED
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from PIL import Image
|
2 |
+
|
3 |
+
|
4 |
+
def flatten(img, bgcolor):
|
5 |
+
# Replace transparency with bgcolor
|
6 |
+
if img.mode in ("RGB"):
|
7 |
+
return img
|
8 |
+
return Image.alpha_composite(Image.new("RGBA", img.size, bgcolor), img).convert("RGB")
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/processing.py
ADDED
@@ -0,0 +1,213 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from PIL import Image, ImageFilter
|
2 |
+
import torch
|
3 |
+
import math
|
4 |
+
from nodes import common_ksampler, VAEEncode, VAEDecode, VAEDecodeTiled
|
5 |
+
from comfy_extras.nodes_custom_sampler import SamplerCustom
|
6 |
+
from utils import pil_to_tensor, tensor_to_pil, get_crop_region, expand_crop, crop_cond
|
7 |
+
from modules import shared
|
8 |
+
|
9 |
+
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
|
10 |
+
Image.Resampling = Image
|
11 |
+
|
12 |
+
|
13 |
+
class StableDiffusionProcessing:
|
14 |
+
|
15 |
+
def __init__(
|
16 |
+
self,
|
17 |
+
init_img,
|
18 |
+
model,
|
19 |
+
positive,
|
20 |
+
negative,
|
21 |
+
vae,
|
22 |
+
seed,
|
23 |
+
steps,
|
24 |
+
cfg,
|
25 |
+
sampler_name,
|
26 |
+
scheduler,
|
27 |
+
denoise,
|
28 |
+
upscale_by,
|
29 |
+
uniform_tile_mode,
|
30 |
+
tiled_decode,
|
31 |
+
custom_sampler=None,
|
32 |
+
custom_sigmas=None
|
33 |
+
):
|
34 |
+
# Variables used by the USDU script
|
35 |
+
self.init_images = [init_img]
|
36 |
+
self.image_mask = None
|
37 |
+
self.mask_blur = 0
|
38 |
+
self.inpaint_full_res_padding = 0
|
39 |
+
self.width = init_img.width
|
40 |
+
self.height = init_img.height
|
41 |
+
|
42 |
+
# ComfyUI Sampler inputs
|
43 |
+
self.model = model
|
44 |
+
self.positive = positive
|
45 |
+
self.negative = negative
|
46 |
+
self.vae = vae
|
47 |
+
self.seed = seed
|
48 |
+
self.steps = steps
|
49 |
+
self.cfg = cfg
|
50 |
+
self.sampler_name = sampler_name
|
51 |
+
self.scheduler = scheduler
|
52 |
+
self.denoise = denoise
|
53 |
+
|
54 |
+
# Optional custom sampler and sigmas
|
55 |
+
self.custom_sampler = custom_sampler
|
56 |
+
self.custom_sigmas = custom_sigmas
|
57 |
+
|
58 |
+
if (custom_sampler is not None) ^ (custom_sigmas is not None):
|
59 |
+
print("[USDU] Both custom sampler and custom sigmas must be provided, defaulting to widget sampler and sigmas")
|
60 |
+
|
61 |
+
# Variables used only by this script
|
62 |
+
self.init_size = init_img.width, init_img.height
|
63 |
+
self.upscale_by = upscale_by
|
64 |
+
self.uniform_tile_mode = uniform_tile_mode
|
65 |
+
self.tiled_decode = tiled_decode
|
66 |
+
self.vae_decoder = VAEDecode()
|
67 |
+
self.vae_encoder = VAEEncode()
|
68 |
+
self.vae_decoder_tiled = VAEDecodeTiled()
|
69 |
+
|
70 |
+
# Other required A1111 variables for the USDU script that is currently unused in this script
|
71 |
+
self.extra_generation_params = {}
|
72 |
+
|
73 |
+
|
74 |
+
class Processed:
|
75 |
+
|
76 |
+
def __init__(self, p: StableDiffusionProcessing, images: list, seed: int, info: str):
|
77 |
+
self.images = images
|
78 |
+
self.seed = seed
|
79 |
+
self.info = info
|
80 |
+
|
81 |
+
def infotext(self, p: StableDiffusionProcessing, index):
|
82 |
+
return None
|
83 |
+
|
84 |
+
|
85 |
+
def fix_seed(p: StableDiffusionProcessing):
|
86 |
+
pass
|
87 |
+
|
88 |
+
|
89 |
+
def sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise, custom_sampler, custom_sigmas):
|
90 |
+
# Choose way to sample based on given inputs
|
91 |
+
|
92 |
+
# Custom sampler and sigmas
|
93 |
+
if custom_sampler is not None and custom_sigmas is not None:
|
94 |
+
custom_sample = SamplerCustom()
|
95 |
+
(samples, _) = getattr(custom_sample, custom_sample.FUNCTION)(
|
96 |
+
model=model,
|
97 |
+
add_noise=True,
|
98 |
+
noise_seed=seed,
|
99 |
+
cfg=cfg,
|
100 |
+
positive=positive,
|
101 |
+
negative=negative,
|
102 |
+
sampler=custom_sampler,
|
103 |
+
sigmas=custom_sigmas,
|
104 |
+
latent_image=latent
|
105 |
+
)
|
106 |
+
return samples
|
107 |
+
|
108 |
+
# Default
|
109 |
+
(samples,) = common_ksampler(model, seed, steps, cfg, sampler_name,
|
110 |
+
scheduler, positive, negative, latent, denoise=denoise)
|
111 |
+
return samples
|
112 |
+
|
113 |
+
|
114 |
+
def process_images(p: StableDiffusionProcessing) -> Processed:
|
115 |
+
# Where the main image generation happens in A1111
|
116 |
+
|
117 |
+
# Setup
|
118 |
+
image_mask = p.image_mask.convert('L')
|
119 |
+
init_image = p.init_images[0]
|
120 |
+
|
121 |
+
# Locate the white region of the mask outlining the tile and add padding
|
122 |
+
crop_region = get_crop_region(image_mask, p.inpaint_full_res_padding)
|
123 |
+
|
124 |
+
if p.uniform_tile_mode:
|
125 |
+
# Expand the crop region to match the processing size ratio and then resize it to the processing size
|
126 |
+
x1, y1, x2, y2 = crop_region
|
127 |
+
crop_width = x2 - x1
|
128 |
+
crop_height = y2 - y1
|
129 |
+
crop_ratio = crop_width / crop_height
|
130 |
+
p_ratio = p.width / p.height
|
131 |
+
if crop_ratio > p_ratio:
|
132 |
+
target_width = crop_width
|
133 |
+
target_height = round(crop_width / p_ratio)
|
134 |
+
else:
|
135 |
+
target_width = round(crop_height * p_ratio)
|
136 |
+
target_height = crop_height
|
137 |
+
crop_region, _ = expand_crop(crop_region, image_mask.width, image_mask.height, target_width, target_height)
|
138 |
+
tile_size = p.width, p.height
|
139 |
+
else:
|
140 |
+
# Uses the minimal size that can fit the mask, minimizes tile size but may lead to image sizes that the model is not trained on
|
141 |
+
x1, y1, x2, y2 = crop_region
|
142 |
+
crop_width = x2 - x1
|
143 |
+
crop_height = y2 - y1
|
144 |
+
target_width = math.ceil(crop_width / 8) * 8
|
145 |
+
target_height = math.ceil(crop_height / 8) * 8
|
146 |
+
crop_region, tile_size = expand_crop(crop_region, image_mask.width,
|
147 |
+
image_mask.height, target_width, target_height)
|
148 |
+
|
149 |
+
# Blur the mask
|
150 |
+
if p.mask_blur > 0:
|
151 |
+
image_mask = image_mask.filter(ImageFilter.GaussianBlur(p.mask_blur))
|
152 |
+
|
153 |
+
# Crop the images to get the tiles that will be used for generation
|
154 |
+
tiles = [img.crop(crop_region) for img in shared.batch]
|
155 |
+
|
156 |
+
# Assume the same size for all images in the batch
|
157 |
+
initial_tile_size = tiles[0].size
|
158 |
+
|
159 |
+
# Resize if necessary
|
160 |
+
for i, tile in enumerate(tiles):
|
161 |
+
if tile.size != tile_size:
|
162 |
+
tiles[i] = tile.resize(tile_size, Image.Resampling.LANCZOS)
|
163 |
+
|
164 |
+
# Crop conditioning
|
165 |
+
positive_cropped = crop_cond(p.positive, crop_region, p.init_size, init_image.size, tile_size)
|
166 |
+
negative_cropped = crop_cond(p.negative, crop_region, p.init_size, init_image.size, tile_size)
|
167 |
+
|
168 |
+
# Encode the image
|
169 |
+
batched_tiles = torch.cat([pil_to_tensor(tile) for tile in tiles], dim=0)
|
170 |
+
(latent,) = p.vae_encoder.encode(p.vae, batched_tiles)
|
171 |
+
|
172 |
+
# Generate samples
|
173 |
+
samples = sample(p.model, p.seed, p.steps, p.cfg, p.sampler_name, p.scheduler, positive_cropped,
|
174 |
+
negative_cropped, latent, p.denoise, p.custom_sampler, p.custom_sigmas)
|
175 |
+
|
176 |
+
# Decode the sample
|
177 |
+
if not p.tiled_decode:
|
178 |
+
(decoded,) = p.vae_decoder.decode(p.vae, samples)
|
179 |
+
else:
|
180 |
+
print("[USDU] Using tiled decode")
|
181 |
+
(decoded,) = p.vae_decoder_tiled.decode(p.vae, samples, 512) # Default tile size is 512
|
182 |
+
|
183 |
+
# Convert the sample to a PIL image
|
184 |
+
tiles_sampled = [tensor_to_pil(decoded, i) for i in range(len(decoded))]
|
185 |
+
|
186 |
+
for i, tile_sampled in enumerate(tiles_sampled):
|
187 |
+
init_image = shared.batch[i]
|
188 |
+
|
189 |
+
# Resize back to the original size
|
190 |
+
if tile_sampled.size != initial_tile_size:
|
191 |
+
tile_sampled = tile_sampled.resize(initial_tile_size, Image.Resampling.LANCZOS)
|
192 |
+
|
193 |
+
# Put the tile into position
|
194 |
+
image_tile_only = Image.new('RGBA', init_image.size)
|
195 |
+
image_tile_only.paste(tile_sampled, crop_region[:2])
|
196 |
+
|
197 |
+
# Add the mask as an alpha channel
|
198 |
+
# Must make a copy due to the possibility of an edge becoming black
|
199 |
+
temp = image_tile_only.copy()
|
200 |
+
temp.putalpha(image_mask)
|
201 |
+
image_tile_only.paste(temp, image_tile_only)
|
202 |
+
|
203 |
+
# Add back the tile to the initial image according to the mask in the alpha channel
|
204 |
+
result = init_image.convert('RGBA')
|
205 |
+
result.alpha_composite(image_tile_only)
|
206 |
+
|
207 |
+
# Convert back to RGB
|
208 |
+
result = result.convert('RGB')
|
209 |
+
|
210 |
+
shared.batch[i] = result
|
211 |
+
|
212 |
+
processed = Processed(p, [shared.batch[0]], p.seed, None)
|
213 |
+
return processed
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/scripts.py
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
class Script:
|
2 |
+
pass
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/shared.py
ADDED
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
class Options:
|
2 |
+
img2img_background_color = "#ffffff" # Set to white for now
|
3 |
+
|
4 |
+
|
5 |
+
class State:
|
6 |
+
interrupted = False
|
7 |
+
|
8 |
+
def begin(self):
|
9 |
+
pass
|
10 |
+
|
11 |
+
def end(self):
|
12 |
+
pass
|
13 |
+
|
14 |
+
|
15 |
+
opts = Options()
|
16 |
+
state = State()
|
17 |
+
|
18 |
+
# Will only ever hold 1 upscaler
|
19 |
+
sd_upscalers = [None]
|
20 |
+
# The upscaler usable by ComfyUI nodes
|
21 |
+
actual_upscaler = None
|
22 |
+
|
23 |
+
# Batch of images to upscale
|
24 |
+
batch = None
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/modules/upscaler.py
ADDED
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from PIL import Image
|
2 |
+
from utils import tensor_to_pil, pil_to_tensor
|
3 |
+
from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel
|
4 |
+
from modules import shared
|
5 |
+
|
6 |
+
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
|
7 |
+
Image.Resampling = Image
|
8 |
+
|
9 |
+
|
10 |
+
class Upscaler:
|
11 |
+
|
12 |
+
def _upscale(self, img: Image, scale):
|
13 |
+
if scale == 1.0:
|
14 |
+
return img
|
15 |
+
if (shared.actual_upscaler is None):
|
16 |
+
return img.resize((img.width * scale, img.height * scale), Image.Resampling.NEAREST)
|
17 |
+
tensor = pil_to_tensor(img)
|
18 |
+
image_upscale_node = ImageUpscaleWithModel()
|
19 |
+
(upscaled,) = image_upscale_node.upscale(shared.actual_upscaler, tensor)
|
20 |
+
return tensor_to_pil(upscaled)
|
21 |
+
|
22 |
+
def upscale(self, img: Image, scale, selected_model: str = None):
|
23 |
+
shared.batch = [self._upscale(img, scale) for img in shared.batch]
|
24 |
+
return shared.batch[0]
|
25 |
+
|
26 |
+
|
27 |
+
class UpscalerData:
|
28 |
+
name = ""
|
29 |
+
data_path = ""
|
30 |
+
|
31 |
+
def __init__(self):
|
32 |
+
self.scaler = Upscaler()
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/nodes.py
ADDED
@@ -0,0 +1,234 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# ComfyUI Node for Ultimate SD Upscale by Coyote-A: https://github.com/Coyote-A/ultimate-upscale-for-automatic1111
|
2 |
+
|
3 |
+
import logging
|
4 |
+
import torch
|
5 |
+
import comfy
|
6 |
+
from usdu_patch import usdu
|
7 |
+
from utils import tensor_to_pil, pil_to_tensor
|
8 |
+
from modules.processing import StableDiffusionProcessing
|
9 |
+
import modules.shared as shared
|
10 |
+
from modules.upscaler import UpscalerData
|
11 |
+
|
12 |
+
MAX_RESOLUTION = 8192
|
13 |
+
# The modes available for Ultimate SD Upscale
|
14 |
+
MODES = {
|
15 |
+
"Linear": usdu.USDUMode.LINEAR,
|
16 |
+
"Chess": usdu.USDUMode.CHESS,
|
17 |
+
"None": usdu.USDUMode.NONE,
|
18 |
+
}
|
19 |
+
# The seam fix modes
|
20 |
+
SEAM_FIX_MODES = {
|
21 |
+
"None": usdu.USDUSFMode.NONE,
|
22 |
+
"Band Pass": usdu.USDUSFMode.BAND_PASS,
|
23 |
+
"Half Tile": usdu.USDUSFMode.HALF_TILE,
|
24 |
+
"Half Tile + Intersections": usdu.USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS,
|
25 |
+
}
|
26 |
+
|
27 |
+
|
28 |
+
def USDU_base_inputs():
|
29 |
+
required = [
|
30 |
+
("image", ("IMAGE",)),
|
31 |
+
# Sampling Params
|
32 |
+
("model", ("MODEL",)),
|
33 |
+
("positive", ("CONDITIONING",)),
|
34 |
+
("negative", ("CONDITIONING",)),
|
35 |
+
("vae", ("VAE",)),
|
36 |
+
("upscale_by", ("FLOAT", {"default": 2, "min": 0.05, "max": 4, "step": 0.05})),
|
37 |
+
("seed", ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff})),
|
38 |
+
("steps", ("INT", {"default": 20, "min": 1, "max": 10000, "step": 1})),
|
39 |
+
("cfg", ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0})),
|
40 |
+
("sampler_name", (comfy.samplers.KSampler.SAMPLERS,)),
|
41 |
+
("scheduler", (comfy.samplers.KSampler.SCHEDULERS,)),
|
42 |
+
("denoise", ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01})),
|
43 |
+
# Upscale Params
|
44 |
+
("upscale_model", ("UPSCALE_MODEL",)),
|
45 |
+
("mode_type", (list(MODES.keys()),)),
|
46 |
+
("tile_width", ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8})),
|
47 |
+
("tile_height", ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8})),
|
48 |
+
("mask_blur", ("INT", {"default": 8, "min": 0, "max": 64, "step": 1})),
|
49 |
+
("tile_padding", ("INT", {"default": 32, "min": 0, "max": MAX_RESOLUTION, "step": 8})),
|
50 |
+
# Seam fix params
|
51 |
+
("seam_fix_mode", (list(SEAM_FIX_MODES.keys()),)),
|
52 |
+
("seam_fix_denoise", ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})),
|
53 |
+
("seam_fix_width", ("INT", {"default": 64, "min": 0, "max": MAX_RESOLUTION, "step": 8})),
|
54 |
+
("seam_fix_mask_blur", ("INT", {"default": 8, "min": 0, "max": 64, "step": 1})),
|
55 |
+
("seam_fix_padding", ("INT", {"default": 16, "min": 0, "max": MAX_RESOLUTION, "step": 8})),
|
56 |
+
# Misc
|
57 |
+
("force_uniform_tiles", ("BOOLEAN", {"default": True})),
|
58 |
+
("tiled_decode", ("BOOLEAN", {"default": False})),
|
59 |
+
]
|
60 |
+
|
61 |
+
optional = []
|
62 |
+
|
63 |
+
return required, optional
|
64 |
+
|
65 |
+
|
66 |
+
def prepare_inputs(required: list, optional: list = None):
|
67 |
+
inputs = {}
|
68 |
+
if required:
|
69 |
+
inputs["required"] = {}
|
70 |
+
for name, type in required:
|
71 |
+
inputs["required"][name] = type
|
72 |
+
if optional:
|
73 |
+
inputs["optional"] = {}
|
74 |
+
for name, type in optional:
|
75 |
+
inputs["optional"][name] = type
|
76 |
+
return inputs
|
77 |
+
|
78 |
+
|
79 |
+
def remove_input(inputs: list, input_name: str):
|
80 |
+
for i, (n, _) in enumerate(inputs):
|
81 |
+
if n == input_name:
|
82 |
+
del inputs[i]
|
83 |
+
break
|
84 |
+
|
85 |
+
|
86 |
+
def rename_input(inputs: list, old_name: str, new_name: str):
|
87 |
+
for i, (n, t) in enumerate(inputs):
|
88 |
+
if n == old_name:
|
89 |
+
inputs[i] = (new_name, t)
|
90 |
+
break
|
91 |
+
|
92 |
+
|
93 |
+
class UltimateSDUpscale:
|
94 |
+
@classmethod
|
95 |
+
def INPUT_TYPES(s):
|
96 |
+
required, optional = USDU_base_inputs()
|
97 |
+
return prepare_inputs(required, optional)
|
98 |
+
|
99 |
+
RETURN_TYPES = ("IMAGE",)
|
100 |
+
FUNCTION = "upscale"
|
101 |
+
CATEGORY = "image/upscaling"
|
102 |
+
|
103 |
+
def upscale(self, image, model, positive, negative, vae, upscale_by, seed,
|
104 |
+
steps, cfg, sampler_name, scheduler, denoise, upscale_model,
|
105 |
+
mode_type, tile_width, tile_height, mask_blur, tile_padding,
|
106 |
+
seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur,
|
107 |
+
seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode,
|
108 |
+
custom_sampler=None, custom_sigmas=None):
|
109 |
+
# Store params
|
110 |
+
self.tile_width = tile_width
|
111 |
+
self.tile_height = tile_height
|
112 |
+
self.mask_blur = mask_blur
|
113 |
+
self.tile_padding = tile_padding
|
114 |
+
self.seam_fix_width = seam_fix_width
|
115 |
+
self.seam_fix_denoise = seam_fix_denoise
|
116 |
+
self.seam_fix_padding = seam_fix_padding
|
117 |
+
self.seam_fix_mode = seam_fix_mode
|
118 |
+
self.mode_type = mode_type
|
119 |
+
self.upscale_by = upscale_by
|
120 |
+
self.seam_fix_mask_blur = seam_fix_mask_blur
|
121 |
+
|
122 |
+
#
|
123 |
+
# Set up A1111 patches
|
124 |
+
#
|
125 |
+
|
126 |
+
# Upscaler
|
127 |
+
# An object that the script works with
|
128 |
+
shared.sd_upscalers[0] = UpscalerData()
|
129 |
+
# Where the actual upscaler is stored, will be used when the script upscales using the Upscaler in UpscalerData
|
130 |
+
shared.actual_upscaler = upscale_model
|
131 |
+
|
132 |
+
# Set the batch of images
|
133 |
+
shared.batch = [tensor_to_pil(image, i) for i in range(len(image))]
|
134 |
+
|
135 |
+
# Processing
|
136 |
+
self.sdprocessing = StableDiffusionProcessing(
|
137 |
+
tensor_to_pil(image), model, positive, negative, vae,
|
138 |
+
seed, steps, cfg, sampler_name, scheduler, denoise, upscale_by, force_uniform_tiles, tiled_decode,
|
139 |
+
custom_sampler, custom_sigmas
|
140 |
+
)
|
141 |
+
|
142 |
+
# Disable logging
|
143 |
+
logger = logging.getLogger()
|
144 |
+
old_level = logger.getEffectiveLevel()
|
145 |
+
logger.setLevel(logging.CRITICAL + 1)
|
146 |
+
try:
|
147 |
+
#
|
148 |
+
# Running the script
|
149 |
+
#
|
150 |
+
script = usdu.Script()
|
151 |
+
processed = script.run(p=self.sdprocessing, _=None, tile_width=self.tile_width, tile_height=self.tile_height,
|
152 |
+
mask_blur=self.mask_blur, padding=self.tile_padding, seams_fix_width=self.seam_fix_width,
|
153 |
+
seams_fix_denoise=self.seam_fix_denoise, seams_fix_padding=self.seam_fix_padding,
|
154 |
+
upscaler_index=0, save_upscaled_image=False, redraw_mode=MODES[self.mode_type],
|
155 |
+
save_seams_fix_image=False, seams_fix_mask_blur=self.seam_fix_mask_blur,
|
156 |
+
seams_fix_type=SEAM_FIX_MODES[self.seam_fix_mode], target_size_type=2,
|
157 |
+
custom_width=None, custom_height=None, custom_scale=self.upscale_by)
|
158 |
+
|
159 |
+
# Return the resulting images
|
160 |
+
images = [pil_to_tensor(img) for img in shared.batch]
|
161 |
+
tensor = torch.cat(images, dim=0)
|
162 |
+
return (tensor,)
|
163 |
+
finally:
|
164 |
+
# Restore the original logging level
|
165 |
+
logger.setLevel(old_level)
|
166 |
+
|
167 |
+
class UltimateSDUpscaleNoUpscale(UltimateSDUpscale):
|
168 |
+
@classmethod
|
169 |
+
def INPUT_TYPES(s):
|
170 |
+
required, optional = USDU_base_inputs()
|
171 |
+
remove_input(required, "upscale_model")
|
172 |
+
remove_input(required, "upscale_by")
|
173 |
+
rename_input(required, "image", "upscaled_image")
|
174 |
+
return prepare_inputs(required, optional)
|
175 |
+
|
176 |
+
RETURN_TYPES = ("IMAGE",)
|
177 |
+
FUNCTION = "upscale"
|
178 |
+
CATEGORY = "image/upscaling"
|
179 |
+
|
180 |
+
def upscale(self, upscaled_image, model, positive, negative, vae, seed,
|
181 |
+
steps, cfg, sampler_name, scheduler, denoise,
|
182 |
+
mode_type, tile_width, tile_height, mask_blur, tile_padding,
|
183 |
+
seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur,
|
184 |
+
seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode):
|
185 |
+
upscale_by = 1.0
|
186 |
+
return super().upscale(upscaled_image, model, positive, negative, vae, upscale_by, seed,
|
187 |
+
steps, cfg, sampler_name, scheduler, denoise, None,
|
188 |
+
mode_type, tile_width, tile_height, mask_blur, tile_padding,
|
189 |
+
seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur,
|
190 |
+
seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode)
|
191 |
+
|
192 |
+
class UltimateSDUpscaleCustomSample(UltimateSDUpscale):
|
193 |
+
@classmethod
|
194 |
+
def INPUT_TYPES(s):
|
195 |
+
required, optional = USDU_base_inputs()
|
196 |
+
remove_input(required, "upscale_model")
|
197 |
+
optional.append(("upscale_model", ("UPSCALE_MODEL",)))
|
198 |
+
optional.append(("custom_sampler", ("SAMPLER",)))
|
199 |
+
optional.append(("custom_sigmas", ("SIGMAS",)))
|
200 |
+
return prepare_inputs(required, optional)
|
201 |
+
|
202 |
+
RETURN_TYPES = ("IMAGE",)
|
203 |
+
FUNCTION = "upscale"
|
204 |
+
CATEGORY = "image/upscaling"
|
205 |
+
|
206 |
+
def upscale(self, image, model, positive, negative, vae, upscale_by, seed,
|
207 |
+
steps, cfg, sampler_name, scheduler, denoise,
|
208 |
+
mode_type, tile_width, tile_height, mask_blur, tile_padding,
|
209 |
+
seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur,
|
210 |
+
seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode,
|
211 |
+
upscale_model=None,
|
212 |
+
custom_sampler=None, custom_sigmas=None):
|
213 |
+
return super().upscale(image, model, positive, negative, vae, upscale_by, seed,
|
214 |
+
steps, cfg, sampler_name, scheduler, denoise, upscale_model,
|
215 |
+
mode_type, tile_width, tile_height, mask_blur, tile_padding,
|
216 |
+
seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur,
|
217 |
+
seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode,
|
218 |
+
custom_sampler, custom_sigmas)
|
219 |
+
|
220 |
+
|
221 |
+
# A dictionary that contains all nodes you want to export with their names
|
222 |
+
# NOTE: names should be globally unique
|
223 |
+
NODE_CLASS_MAPPINGS = {
|
224 |
+
"UltimateSDUpscale": UltimateSDUpscale,
|
225 |
+
"UltimateSDUpscaleNoUpscale": UltimateSDUpscaleNoUpscale,
|
226 |
+
"UltimateSDUpscaleCustomSample": UltimateSDUpscaleCustomSample
|
227 |
+
}
|
228 |
+
|
229 |
+
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
230 |
+
NODE_DISPLAY_NAME_MAPPINGS = {
|
231 |
+
"UltimateSDUpscale": "Ultimate SD Upscale",
|
232 |
+
"UltimateSDUpscaleNoUpscale": "Ultimate SD Upscale (No Upscale)",
|
233 |
+
"UltimateSDUpscaleCustomSample": "Ultimate SD Upscale (Custom Sample)"
|
234 |
+
}
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/repositories/__init__.py
ADDED
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import sys
|
3 |
+
import importlib.util
|
4 |
+
|
5 |
+
repositories_path = os.path.dirname(os.path.realpath(__file__))
|
6 |
+
|
7 |
+
# Import the script
|
8 |
+
script_name = os.path.join("scripts", "ultimate-upscale")
|
9 |
+
repo_name = "ultimate_sd_upscale"
|
10 |
+
script_path = os.path.join(repositories_path, repo_name, f"{script_name}.py")
|
11 |
+
spec = importlib.util.spec_from_file_location(script_name, script_path)
|
12 |
+
ultimate_upscale = importlib.util.module_from_spec(spec)
|
13 |
+
sys.modules[script_name] = ultimate_upscale
|
14 |
+
spec.loader.exec_module(ultimate_upscale)
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/repositories/ultimate_sd_upscale/README.md
ADDED
@@ -0,0 +1,119 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Ultimate SD Upscale extension for [AUTOMATIC1111 Stable Diffusion web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
|
2 |
+
Now you have the opportunity to use a large denoise (0.3-0.5) and not spawn many artifacts. Works on any video card, since you can use a 512x512 tile size and the image will converge.
|
3 |
+
|
4 |
+
News channel: https://t.me/usdunews
|
5 |
+
|
6 |
+
# Instructions
|
7 |
+
All instructions can be found on the project's [wiki](https://github.com/Coyote-A/ultimate-upscale-for-automatic1111/wiki).
|
8 |
+
|
9 |
+
# Refs
|
10 |
+
|
11 |
+
https://github.com/ssitu/ComfyUI_UltimateSDUpscale - Implementation for ComfyUI
|
12 |
+
|
13 |
+
# Examples
|
14 |
+
More on [wiki page](https://github.com/Coyote-A/ultimate-upscale-for-automatic1111/wiki/Examples)
|
15 |
+
|
16 |
+
<details>
|
17 |
+
<summary>E1</summary>
|
18 |
+
Original image
|
19 |
+
|
20 |
+
![Original](https://i.imgur.com/J8mRYOD.png)
|
21 |
+
|
22 |
+
2k upscaled. **Tile size**: 512, **Padding**: 32, **Mask blur**: 16, **Denoise**: 0.4
|
23 |
+
![2k upscale](https://i.imgur.com/0aKua4r.png)
|
24 |
+
</details>
|
25 |
+
|
26 |
+
<details>
|
27 |
+
<summary>E2</summary>
|
28 |
+
Original image
|
29 |
+
|
30 |
+
![Original](https://i.imgur.com/aALNI2w.png)
|
31 |
+
|
32 |
+
2k upscaled. **Tile size**: 768, **Padding**: 55, **Mask blur**: 20, **Denoise**: 0.35
|
33 |
+
![2k upscale](https://i.imgur.com/B5PHz0J.png)
|
34 |
+
|
35 |
+
4k upscaled. **Tile size**: 768, **Padding**: 55, **Mask blur**: 20, **Denoise**: 0.35
|
36 |
+
![4k upscale](https://i.imgur.com/tIUQ7TJ.jpg)
|
37 |
+
</details>
|
38 |
+
|
39 |
+
<details>
|
40 |
+
<summary>E3</summary>
|
41 |
+
Original image
|
42 |
+
|
43 |
+
![Original](https://i.imgur.com/AGtszA8.png)
|
44 |
+
|
45 |
+
4k upscaled. **Tile size**: 768, **Padding**: 55, **Mask blur**: 20, **Denoise**: 0.4
|
46 |
+
![4k upscale](https://i.imgur.com/LCYLfCs.jpg)
|
47 |
+
</details>
|
48 |
+
|
49 |
+
# API Usage
|
50 |
+
|
51 |
+
```javascript
|
52 |
+
{
|
53 |
+
"script_name" : "ultimate sd upscale",
|
54 |
+
"script_args" : [
|
55 |
+
null, // _ (not used)
|
56 |
+
512, // tile_width
|
57 |
+
512, // tile_height
|
58 |
+
8, // mask_blur
|
59 |
+
32, // padding
|
60 |
+
64, // seams_fix_width
|
61 |
+
0.35, // seams_fix_denoise
|
62 |
+
32, // seams_fix_padding
|
63 |
+
0, // upscaler_index
|
64 |
+
true, // save_upscaled_image a.k.a Upscaled
|
65 |
+
0, // redraw_mode
|
66 |
+
false, // save_seams_fix_image a.k.a Seams fix
|
67 |
+
8, // seams_fix_mask_blur
|
68 |
+
0, // seams_fix_type
|
69 |
+
0, // target_size_type
|
70 |
+
2048, // custom_width
|
71 |
+
2048, // custom_height
|
72 |
+
2 // custom_scale
|
73 |
+
]
|
74 |
+
}
|
75 |
+
```
|
76 |
+
upscaler_index
|
77 |
+
| Value | |
|
78 |
+
|:-------------:| -----:|
|
79 |
+
| 0 | None |
|
80 |
+
| 1 | Lanczos |
|
81 |
+
| 2 | Nearest |
|
82 |
+
| 3 | ESRGAN_4x |
|
83 |
+
| 4 | LDSR |
|
84 |
+
| 5 | R-ESRGAN_4x+ |
|
85 |
+
| 6 | R-ESRGAN 4x+ Anime6B |
|
86 |
+
| 7 | ScuNET GAN |
|
87 |
+
| 8 | ScuNET PSNR |
|
88 |
+
| 9 | SwinIR 4x |
|
89 |
+
|
90 |
+
redraw_mode
|
91 |
+
| Value | |
|
92 |
+
|:-------------:| -----:|
|
93 |
+
| 0 | Linear |
|
94 |
+
| 1 | Chess |
|
95 |
+
| 2 | None |
|
96 |
+
|
97 |
+
seams_fix_mask_blur
|
98 |
+
| Value | |
|
99 |
+
|:-------------:| -----:|
|
100 |
+
| 0 | None |
|
101 |
+
| 1 | BAND_PASS |
|
102 |
+
| 2 | HALF_TILE |
|
103 |
+
| 3 | HALF_TILE_PLUS_INTERSECTIONS |
|
104 |
+
|
105 |
+
seams_fix_type
|
106 |
+
| Value | |
|
107 |
+
|:-------------:| -----:|
|
108 |
+
| 0 | None |
|
109 |
+
| 1 | Band pass |
|
110 |
+
| 2 | Half tile offset pass |
|
111 |
+
| 3 | Half tile offset pass + intersections |
|
112 |
+
|
113 |
+
seams_fix_type
|
114 |
+
| Value | |
|
115 |
+
|:-------------:| -----:|
|
116 |
+
| 0 | From img2img2 settings |
|
117 |
+
| 1 | Custom size |
|
118 |
+
| 2 | Scale from image size |
|
119 |
+
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/repositories/ultimate_sd_upscale/scripts/ultimate-upscale.py
ADDED
@@ -0,0 +1,569 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import math
|
2 |
+
import gradio as gr
|
3 |
+
from PIL import Image, ImageDraw, ImageOps
|
4 |
+
from modules import processing, shared, images, devices, scripts
|
5 |
+
from modules.processing import StableDiffusionProcessing
|
6 |
+
from modules.processing import Processed
|
7 |
+
from modules.shared import opts, state
|
8 |
+
from enum import Enum
|
9 |
+
|
10 |
+
elem_id_prefix = "ultimateupscale"
|
11 |
+
|
12 |
+
class USDUMode(Enum):
|
13 |
+
LINEAR = 0
|
14 |
+
CHESS = 1
|
15 |
+
NONE = 2
|
16 |
+
|
17 |
+
class USDUSFMode(Enum):
|
18 |
+
NONE = 0
|
19 |
+
BAND_PASS = 1
|
20 |
+
HALF_TILE = 2
|
21 |
+
HALF_TILE_PLUS_INTERSECTIONS = 3
|
22 |
+
|
23 |
+
class USDUpscaler():
|
24 |
+
|
25 |
+
def __init__(self, p, image, upscaler_index:int, save_redraw, save_seams_fix, tile_width, tile_height) -> None:
|
26 |
+
self.p:StableDiffusionProcessing = p
|
27 |
+
self.image:Image = image
|
28 |
+
self.scale_factor = math.ceil(max(p.width, p.height) / max(image.width, image.height))
|
29 |
+
self.upscaler = shared.sd_upscalers[upscaler_index]
|
30 |
+
self.redraw = USDURedraw()
|
31 |
+
self.redraw.save = save_redraw
|
32 |
+
self.redraw.tile_width = tile_width if tile_width > 0 else tile_height
|
33 |
+
self.redraw.tile_height = tile_height if tile_height > 0 else tile_width
|
34 |
+
self.seams_fix = USDUSeamsFix()
|
35 |
+
self.seams_fix.save = save_seams_fix
|
36 |
+
self.seams_fix.tile_width = tile_width if tile_width > 0 else tile_height
|
37 |
+
self.seams_fix.tile_height = tile_height if tile_height > 0 else tile_width
|
38 |
+
self.initial_info = None
|
39 |
+
self.rows = math.ceil(self.p.height / self.redraw.tile_height)
|
40 |
+
self.cols = math.ceil(self.p.width / self.redraw.tile_width)
|
41 |
+
|
42 |
+
def get_factor(self, num):
|
43 |
+
# Its just return, don't need elif
|
44 |
+
if num == 1:
|
45 |
+
return 2
|
46 |
+
if num % 4 == 0:
|
47 |
+
return 4
|
48 |
+
if num % 3 == 0:
|
49 |
+
return 3
|
50 |
+
if num % 2 == 0:
|
51 |
+
return 2
|
52 |
+
return 0
|
53 |
+
|
54 |
+
def get_factors(self):
|
55 |
+
scales = []
|
56 |
+
current_scale = 1
|
57 |
+
current_scale_factor = self.get_factor(self.scale_factor)
|
58 |
+
while current_scale_factor == 0:
|
59 |
+
self.scale_factor += 1
|
60 |
+
current_scale_factor = self.get_factor(self.scale_factor)
|
61 |
+
while current_scale < self.scale_factor:
|
62 |
+
current_scale_factor = self.get_factor(self.scale_factor // current_scale)
|
63 |
+
scales.append(current_scale_factor)
|
64 |
+
current_scale = current_scale * current_scale_factor
|
65 |
+
if current_scale_factor == 0:
|
66 |
+
break
|
67 |
+
self.scales = enumerate(scales)
|
68 |
+
|
69 |
+
def upscale(self):
|
70 |
+
# Log info
|
71 |
+
print(f"Canva size: {self.p.width}x{self.p.height}")
|
72 |
+
print(f"Image size: {self.image.width}x{self.image.height}")
|
73 |
+
print(f"Scale factor: {self.scale_factor}")
|
74 |
+
# Check upscaler is not empty
|
75 |
+
if self.upscaler.name == "None":
|
76 |
+
self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS)
|
77 |
+
return
|
78 |
+
# Get list with scale factors
|
79 |
+
self.get_factors()
|
80 |
+
# Upscaling image over all factors
|
81 |
+
for index, value in self.scales:
|
82 |
+
print(f"Upscaling iteration {index+1} with scale factor {value}")
|
83 |
+
self.image = self.upscaler.scaler.upscale(self.image, value, self.upscaler.data_path)
|
84 |
+
# Resize image to set values
|
85 |
+
self.image = self.image.resize((self.p.width, self.p.height), resample=Image.LANCZOS)
|
86 |
+
|
87 |
+
def setup_redraw(self, redraw_mode, padding, mask_blur):
|
88 |
+
self.redraw.mode = USDUMode(redraw_mode)
|
89 |
+
self.redraw.enabled = self.redraw.mode != USDUMode.NONE
|
90 |
+
self.redraw.padding = padding
|
91 |
+
self.p.mask_blur = mask_blur
|
92 |
+
|
93 |
+
def setup_seams_fix(self, padding, denoise, mask_blur, width, mode):
|
94 |
+
self.seams_fix.padding = padding
|
95 |
+
self.seams_fix.denoise = denoise
|
96 |
+
self.seams_fix.mask_blur = mask_blur
|
97 |
+
self.seams_fix.width = width
|
98 |
+
self.seams_fix.mode = USDUSFMode(mode)
|
99 |
+
self.seams_fix.enabled = self.seams_fix.mode != USDUSFMode.NONE
|
100 |
+
|
101 |
+
def save_image(self):
|
102 |
+
if type(self.p.prompt) != list:
|
103 |
+
images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt, opts.samples_format, info=self.initial_info, p=self.p)
|
104 |
+
else:
|
105 |
+
images.save_image(self.image, self.p.outpath_samples, "", self.p.seed, self.p.prompt[0], opts.samples_format, info=self.initial_info, p=self.p)
|
106 |
+
|
107 |
+
def calc_jobs_count(self):
|
108 |
+
redraw_job_count = (self.rows * self.cols) if self.redraw.enabled else 0
|
109 |
+
seams_job_count = 0
|
110 |
+
if self.seams_fix.mode == USDUSFMode.BAND_PASS:
|
111 |
+
seams_job_count = self.rows + self.cols - 2
|
112 |
+
elif self.seams_fix.mode == USDUSFMode.HALF_TILE:
|
113 |
+
seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols
|
114 |
+
elif self.seams_fix.mode == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS:
|
115 |
+
seams_job_count = self.rows * (self.cols - 1) + (self.rows - 1) * self.cols + (self.rows - 1) * (self.cols - 1)
|
116 |
+
|
117 |
+
state.job_count = redraw_job_count + seams_job_count
|
118 |
+
|
119 |
+
def print_info(self):
|
120 |
+
print(f"Tile size: {self.redraw.tile_width}x{self.redraw.tile_height}")
|
121 |
+
print(f"Tiles amount: {self.rows * self.cols}")
|
122 |
+
print(f"Grid: {self.rows}x{self.cols}")
|
123 |
+
print(f"Redraw enabled: {self.redraw.enabled}")
|
124 |
+
print(f"Seams fix mode: {self.seams_fix.mode.name}")
|
125 |
+
|
126 |
+
def add_extra_info(self):
|
127 |
+
self.p.extra_generation_params["Ultimate SD upscale upscaler"] = self.upscaler.name
|
128 |
+
self.p.extra_generation_params["Ultimate SD upscale tile_width"] = self.redraw.tile_width
|
129 |
+
self.p.extra_generation_params["Ultimate SD upscale tile_height"] = self.redraw.tile_height
|
130 |
+
self.p.extra_generation_params["Ultimate SD upscale mask_blur"] = self.p.mask_blur
|
131 |
+
self.p.extra_generation_params["Ultimate SD upscale padding"] = self.redraw.padding
|
132 |
+
|
133 |
+
def process(self):
|
134 |
+
state.begin()
|
135 |
+
self.calc_jobs_count()
|
136 |
+
self.result_images = []
|
137 |
+
if self.redraw.enabled:
|
138 |
+
self.image = self.redraw.start(self.p, self.image, self.rows, self.cols)
|
139 |
+
self.initial_info = self.redraw.initial_info
|
140 |
+
self.result_images.append(self.image)
|
141 |
+
if self.redraw.save:
|
142 |
+
self.save_image()
|
143 |
+
|
144 |
+
if self.seams_fix.enabled:
|
145 |
+
self.image = self.seams_fix.start(self.p, self.image, self.rows, self.cols)
|
146 |
+
self.initial_info = self.seams_fix.initial_info
|
147 |
+
self.result_images.append(self.image)
|
148 |
+
if self.seams_fix.save:
|
149 |
+
self.save_image()
|
150 |
+
state.end()
|
151 |
+
|
152 |
+
class USDURedraw():
|
153 |
+
|
154 |
+
def init_draw(self, p, width, height):
|
155 |
+
p.inpaint_full_res = True
|
156 |
+
p.inpaint_full_res_padding = self.padding
|
157 |
+
p.width = math.ceil((self.tile_width+self.padding) / 64) * 64
|
158 |
+
p.height = math.ceil((self.tile_height+self.padding) / 64) * 64
|
159 |
+
mask = Image.new("L", (width, height), "black")
|
160 |
+
draw = ImageDraw.Draw(mask)
|
161 |
+
return mask, draw
|
162 |
+
|
163 |
+
def calc_rectangle(self, xi, yi):
|
164 |
+
x1 = xi * self.tile_width
|
165 |
+
y1 = yi * self.tile_height
|
166 |
+
x2 = xi * self.tile_width + self.tile_width
|
167 |
+
y2 = yi * self.tile_height + self.tile_height
|
168 |
+
|
169 |
+
return x1, y1, x2, y2
|
170 |
+
|
171 |
+
def linear_process(self, p, image, rows, cols):
|
172 |
+
mask, draw = self.init_draw(p, image.width, image.height)
|
173 |
+
for yi in range(rows):
|
174 |
+
for xi in range(cols):
|
175 |
+
if state.interrupted:
|
176 |
+
break
|
177 |
+
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
178 |
+
p.init_images = [image]
|
179 |
+
p.image_mask = mask
|
180 |
+
processed = processing.process_images(p)
|
181 |
+
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
182 |
+
if (len(processed.images) > 0):
|
183 |
+
image = processed.images[0]
|
184 |
+
|
185 |
+
p.width = image.width
|
186 |
+
p.height = image.height
|
187 |
+
self.initial_info = processed.infotext(p, 0)
|
188 |
+
|
189 |
+
return image
|
190 |
+
|
191 |
+
def chess_process(self, p, image, rows, cols):
|
192 |
+
mask, draw = self.init_draw(p, image.width, image.height)
|
193 |
+
tiles = []
|
194 |
+
# calc tiles colors
|
195 |
+
for yi in range(rows):
|
196 |
+
for xi in range(cols):
|
197 |
+
if state.interrupted:
|
198 |
+
break
|
199 |
+
if xi == 0:
|
200 |
+
tiles.append([])
|
201 |
+
color = xi % 2 == 0
|
202 |
+
if yi > 0 and yi % 2 != 0:
|
203 |
+
color = not color
|
204 |
+
tiles[yi].append(color)
|
205 |
+
|
206 |
+
for yi in range(len(tiles)):
|
207 |
+
for xi in range(len(tiles[yi])):
|
208 |
+
if state.interrupted:
|
209 |
+
break
|
210 |
+
if not tiles[yi][xi]:
|
211 |
+
tiles[yi][xi] = not tiles[yi][xi]
|
212 |
+
continue
|
213 |
+
tiles[yi][xi] = not tiles[yi][xi]
|
214 |
+
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
215 |
+
p.init_images = [image]
|
216 |
+
p.image_mask = mask
|
217 |
+
processed = processing.process_images(p)
|
218 |
+
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
219 |
+
if (len(processed.images) > 0):
|
220 |
+
image = processed.images[0]
|
221 |
+
|
222 |
+
for yi in range(len(tiles)):
|
223 |
+
for xi in range(len(tiles[yi])):
|
224 |
+
if state.interrupted:
|
225 |
+
break
|
226 |
+
if not tiles[yi][xi]:
|
227 |
+
continue
|
228 |
+
draw.rectangle(self.calc_rectangle(xi, yi), fill="white")
|
229 |
+
p.init_images = [image]
|
230 |
+
p.image_mask = mask
|
231 |
+
processed = processing.process_images(p)
|
232 |
+
draw.rectangle(self.calc_rectangle(xi, yi), fill="black")
|
233 |
+
if (len(processed.images) > 0):
|
234 |
+
image = processed.images[0]
|
235 |
+
|
236 |
+
p.width = image.width
|
237 |
+
p.height = image.height
|
238 |
+
self.initial_info = processed.infotext(p, 0)
|
239 |
+
|
240 |
+
return image
|
241 |
+
|
242 |
+
def start(self, p, image, rows, cols):
|
243 |
+
self.initial_info = None
|
244 |
+
if self.mode == USDUMode.LINEAR:
|
245 |
+
return self.linear_process(p, image, rows, cols)
|
246 |
+
if self.mode == USDUMode.CHESS:
|
247 |
+
return self.chess_process(p, image, rows, cols)
|
248 |
+
|
249 |
+
class USDUSeamsFix():
|
250 |
+
|
251 |
+
def init_draw(self, p):
|
252 |
+
self.initial_info = None
|
253 |
+
p.width = math.ceil((self.tile_width+self.padding) / 64) * 64
|
254 |
+
p.height = math.ceil((self.tile_height+self.padding) / 64) * 64
|
255 |
+
|
256 |
+
def half_tile_process(self, p, image, rows, cols):
|
257 |
+
|
258 |
+
self.init_draw(p)
|
259 |
+
processed = None
|
260 |
+
|
261 |
+
gradient = Image.linear_gradient("L")
|
262 |
+
row_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
|
263 |
+
row_gradient.paste(gradient.resize(
|
264 |
+
(self.tile_width, self.tile_height//2), resample=Image.BICUBIC), (0, 0))
|
265 |
+
row_gradient.paste(gradient.rotate(180).resize(
|
266 |
+
(self.tile_width, self.tile_height//2), resample=Image.BICUBIC),
|
267 |
+
(0, self.tile_height//2))
|
268 |
+
col_gradient = Image.new("L", (self.tile_width, self.tile_height), "black")
|
269 |
+
col_gradient.paste(gradient.rotate(90).resize(
|
270 |
+
(self.tile_width//2, self.tile_height), resample=Image.BICUBIC), (0, 0))
|
271 |
+
col_gradient.paste(gradient.rotate(270).resize(
|
272 |
+
(self.tile_width//2, self.tile_height), resample=Image.BICUBIC), (self.tile_width//2, 0))
|
273 |
+
|
274 |
+
p.denoising_strength = self.denoise
|
275 |
+
p.mask_blur = self.mask_blur
|
276 |
+
|
277 |
+
for yi in range(rows-1):
|
278 |
+
for xi in range(cols):
|
279 |
+
if state.interrupted:
|
280 |
+
break
|
281 |
+
p.width = self.tile_width
|
282 |
+
p.height = self.tile_height
|
283 |
+
p.inpaint_full_res = True
|
284 |
+
p.inpaint_full_res_padding = self.padding
|
285 |
+
mask = Image.new("L", (image.width, image.height), "black")
|
286 |
+
mask.paste(row_gradient, (xi*self.tile_width, yi*self.tile_height + self.tile_height//2))
|
287 |
+
|
288 |
+
p.init_images = [image]
|
289 |
+
p.image_mask = mask
|
290 |
+
processed = processing.process_images(p)
|
291 |
+
if (len(processed.images) > 0):
|
292 |
+
image = processed.images[0]
|
293 |
+
|
294 |
+
for yi in range(rows):
|
295 |
+
for xi in range(cols-1):
|
296 |
+
if state.interrupted:
|
297 |
+
break
|
298 |
+
p.width = self.tile_width
|
299 |
+
p.height = self.tile_height
|
300 |
+
p.inpaint_full_res = True
|
301 |
+
p.inpaint_full_res_padding = self.padding
|
302 |
+
mask = Image.new("L", (image.width, image.height), "black")
|
303 |
+
mask.paste(col_gradient, (xi*self.tile_width+self.tile_width//2, yi*self.tile_height))
|
304 |
+
|
305 |
+
p.init_images = [image]
|
306 |
+
p.image_mask = mask
|
307 |
+
processed = processing.process_images(p)
|
308 |
+
if (len(processed.images) > 0):
|
309 |
+
image = processed.images[0]
|
310 |
+
|
311 |
+
p.width = image.width
|
312 |
+
p.height = image.height
|
313 |
+
if processed is not None:
|
314 |
+
self.initial_info = processed.infotext(p, 0)
|
315 |
+
|
316 |
+
return image
|
317 |
+
|
318 |
+
def half_tile_process_corners(self, p, image, rows, cols):
|
319 |
+
fixed_image = self.half_tile_process(p, image, rows, cols)
|
320 |
+
processed = None
|
321 |
+
self.init_draw(p)
|
322 |
+
gradient = Image.radial_gradient("L").resize(
|
323 |
+
(self.tile_width, self.tile_height), resample=Image.BICUBIC)
|
324 |
+
gradient = ImageOps.invert(gradient)
|
325 |
+
p.denoising_strength = self.denoise
|
326 |
+
#p.mask_blur = 0
|
327 |
+
p.mask_blur = self.mask_blur
|
328 |
+
|
329 |
+
for yi in range(rows-1):
|
330 |
+
for xi in range(cols-1):
|
331 |
+
if state.interrupted:
|
332 |
+
break
|
333 |
+
p.width = self.tile_width
|
334 |
+
p.height = self.tile_height
|
335 |
+
p.inpaint_full_res = True
|
336 |
+
p.inpaint_full_res_padding = 0
|
337 |
+
mask = Image.new("L", (fixed_image.width, fixed_image.height), "black")
|
338 |
+
mask.paste(gradient, (xi*self.tile_width + self.tile_width//2,
|
339 |
+
yi*self.tile_height + self.tile_height//2))
|
340 |
+
|
341 |
+
p.init_images = [fixed_image]
|
342 |
+
p.image_mask = mask
|
343 |
+
processed = processing.process_images(p)
|
344 |
+
if (len(processed.images) > 0):
|
345 |
+
fixed_image = processed.images[0]
|
346 |
+
|
347 |
+
p.width = fixed_image.width
|
348 |
+
p.height = fixed_image.height
|
349 |
+
if processed is not None:
|
350 |
+
self.initial_info = processed.infotext(p, 0)
|
351 |
+
|
352 |
+
return fixed_image
|
353 |
+
|
354 |
+
def band_pass_process(self, p, image, cols, rows):
|
355 |
+
|
356 |
+
self.init_draw(p)
|
357 |
+
processed = None
|
358 |
+
|
359 |
+
p.denoising_strength = self.denoise
|
360 |
+
p.mask_blur = 0
|
361 |
+
|
362 |
+
gradient = Image.linear_gradient("L")
|
363 |
+
mirror_gradient = Image.new("L", (256, 256), "black")
|
364 |
+
mirror_gradient.paste(gradient.resize((256, 128), resample=Image.BICUBIC), (0, 0))
|
365 |
+
mirror_gradient.paste(gradient.rotate(180).resize((256, 128), resample=Image.BICUBIC), (0, 128))
|
366 |
+
|
367 |
+
row_gradient = mirror_gradient.resize((image.width, self.width), resample=Image.BICUBIC)
|
368 |
+
col_gradient = mirror_gradient.rotate(90).resize((self.width, image.height), resample=Image.BICUBIC)
|
369 |
+
|
370 |
+
for xi in range(1, rows):
|
371 |
+
if state.interrupted:
|
372 |
+
break
|
373 |
+
p.width = self.width + self.padding * 2
|
374 |
+
p.height = image.height
|
375 |
+
p.inpaint_full_res = True
|
376 |
+
p.inpaint_full_res_padding = self.padding
|
377 |
+
mask = Image.new("L", (image.width, image.height), "black")
|
378 |
+
mask.paste(col_gradient, (xi * self.tile_width - self.width // 2, 0))
|
379 |
+
|
380 |
+
p.init_images = [image]
|
381 |
+
p.image_mask = mask
|
382 |
+
processed = processing.process_images(p)
|
383 |
+
if (len(processed.images) > 0):
|
384 |
+
image = processed.images[0]
|
385 |
+
for yi in range(1, cols):
|
386 |
+
if state.interrupted:
|
387 |
+
break
|
388 |
+
p.width = image.width
|
389 |
+
p.height = self.width + self.padding * 2
|
390 |
+
p.inpaint_full_res = True
|
391 |
+
p.inpaint_full_res_padding = self.padding
|
392 |
+
mask = Image.new("L", (image.width, image.height), "black")
|
393 |
+
mask.paste(row_gradient, (0, yi * self.tile_height - self.width // 2))
|
394 |
+
|
395 |
+
p.init_images = [image]
|
396 |
+
p.image_mask = mask
|
397 |
+
processed = processing.process_images(p)
|
398 |
+
if (len(processed.images) > 0):
|
399 |
+
image = processed.images[0]
|
400 |
+
|
401 |
+
p.width = image.width
|
402 |
+
p.height = image.height
|
403 |
+
if processed is not None:
|
404 |
+
self.initial_info = processed.infotext(p, 0)
|
405 |
+
|
406 |
+
return image
|
407 |
+
|
408 |
+
def start(self, p, image, rows, cols):
|
409 |
+
if USDUSFMode(self.mode) == USDUSFMode.BAND_PASS:
|
410 |
+
return self.band_pass_process(p, image, rows, cols)
|
411 |
+
elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE:
|
412 |
+
return self.half_tile_process(p, image, rows, cols)
|
413 |
+
elif USDUSFMode(self.mode) == USDUSFMode.HALF_TILE_PLUS_INTERSECTIONS:
|
414 |
+
return self.half_tile_process_corners(p, image, rows, cols)
|
415 |
+
else:
|
416 |
+
return image
|
417 |
+
|
418 |
+
class Script(scripts.Script):
|
419 |
+
def title(self):
|
420 |
+
return "Ultimate SD upscale"
|
421 |
+
|
422 |
+
def show(self, is_img2img):
|
423 |
+
return is_img2img
|
424 |
+
|
425 |
+
def ui(self, is_img2img):
|
426 |
+
|
427 |
+
target_size_types = [
|
428 |
+
"From img2img2 settings",
|
429 |
+
"Custom size",
|
430 |
+
"Scale from image size"
|
431 |
+
]
|
432 |
+
|
433 |
+
seams_fix_types = [
|
434 |
+
"None",
|
435 |
+
"Band pass",
|
436 |
+
"Half tile offset pass",
|
437 |
+
"Half tile offset pass + intersections"
|
438 |
+
]
|
439 |
+
|
440 |
+
redrow_modes = [
|
441 |
+
"Linear",
|
442 |
+
"Chess",
|
443 |
+
"None"
|
444 |
+
]
|
445 |
+
|
446 |
+
info = gr.HTML(
|
447 |
+
"<p style=\"margin-bottom:0.75em\">Will upscale the image depending on the selected target size type</p>")
|
448 |
+
|
449 |
+
with gr.Row():
|
450 |
+
target_size_type = gr.Dropdown(label="Target size type", elem_id=f"{elem_id_prefix}_target_size_type", choices=[k for k in target_size_types], type="index",
|
451 |
+
value=next(iter(target_size_types)))
|
452 |
+
|
453 |
+
custom_width = gr.Slider(label='Custom width', elem_id=f"{elem_id_prefix}_custom_width", minimum=64, maximum=8192, step=64, value=2048, visible=False, interactive=True)
|
454 |
+
custom_height = gr.Slider(label='Custom height', elem_id=f"{elem_id_prefix}_custom_height", minimum=64, maximum=8192, step=64, value=2048, visible=False, interactive=True)
|
455 |
+
custom_scale = gr.Slider(label='Scale', elem_id=f"{elem_id_prefix}_custom_scale", minimum=1, maximum=16, step=0.01, value=2, visible=False, interactive=True)
|
456 |
+
|
457 |
+
gr.HTML("<p style=\"margin-bottom:0.75em\">Redraw options:</p>")
|
458 |
+
with gr.Row():
|
459 |
+
upscaler_index = gr.Radio(label='Upscaler', elem_id=f"{elem_id_prefix}_upscaler_index", choices=[x.name for x in shared.sd_upscalers],
|
460 |
+
value=shared.sd_upscalers[0].name, type="index")
|
461 |
+
with gr.Row():
|
462 |
+
redraw_mode = gr.Dropdown(label="Type", elem_id=f"{elem_id_prefix}_redraw_mode", choices=[k for k in redrow_modes], type="index", value=next(iter(redrow_modes)))
|
463 |
+
tile_width = gr.Slider(elem_id=f"{elem_id_prefix}_tile_width", minimum=0, maximum=2048, step=64, label='Tile width', value=512)
|
464 |
+
tile_height = gr.Slider(elem_id=f"{elem_id_prefix}_tile_height", minimum=0, maximum=2048, step=64, label='Tile height', value=0)
|
465 |
+
mask_blur = gr.Slider(elem_id=f"{elem_id_prefix}_mask_blur", label='Mask blur', minimum=0, maximum=64, step=1, value=8)
|
466 |
+
padding = gr.Slider(elem_id=f"{elem_id_prefix}_padding", label='Padding', minimum=0, maximum=512, step=1, value=32)
|
467 |
+
gr.HTML("<p style=\"margin-bottom:0.75em\">Seams fix:</p>")
|
468 |
+
with gr.Row():
|
469 |
+
seams_fix_type = gr.Dropdown(label="Type", elem_id=f"{elem_id_prefix}_seams_fix_type", choices=[k for k in seams_fix_types], type="index", value=next(iter(seams_fix_types)))
|
470 |
+
seams_fix_denoise = gr.Slider(label='Denoise', elem_id=f"{elem_id_prefix}_seams_fix_denoise", minimum=0, maximum=1, step=0.01, value=0.35, visible=False, interactive=True)
|
471 |
+
seams_fix_width = gr.Slider(label='Width', elem_id=f"{elem_id_prefix}_seams_fix_width", minimum=0, maximum=128, step=1, value=64, visible=False, interactive=True)
|
472 |
+
seams_fix_mask_blur = gr.Slider(label='Mask blur', elem_id=f"{elem_id_prefix}_seams_fix_mask_blur", minimum=0, maximum=64, step=1, value=4, visible=False, interactive=True)
|
473 |
+
seams_fix_padding = gr.Slider(label='Padding', elem_id=f"{elem_id_prefix}_seams_fix_padding", minimum=0, maximum=128, step=1, value=16, visible=False, interactive=True)
|
474 |
+
gr.HTML("<p style=\"margin-bottom:0.75em\">Save options:</p>")
|
475 |
+
with gr.Row():
|
476 |
+
save_upscaled_image = gr.Checkbox(label="Upscaled", elem_id=f"{elem_id_prefix}_save_upscaled_image", value=True)
|
477 |
+
save_seams_fix_image = gr.Checkbox(label="Seams fix", elem_id=f"{elem_id_prefix}_save_seams_fix_image", value=False)
|
478 |
+
|
479 |
+
def select_fix_type(fix_index):
|
480 |
+
all_visible = fix_index != 0
|
481 |
+
mask_blur_visible = fix_index == 2 or fix_index == 3
|
482 |
+
width_visible = fix_index == 1
|
483 |
+
|
484 |
+
return [gr.update(visible=all_visible),
|
485 |
+
gr.update(visible=width_visible),
|
486 |
+
gr.update(visible=mask_blur_visible),
|
487 |
+
gr.update(visible=all_visible)]
|
488 |
+
|
489 |
+
seams_fix_type.change(
|
490 |
+
fn=select_fix_type,
|
491 |
+
inputs=seams_fix_type,
|
492 |
+
outputs=[seams_fix_denoise, seams_fix_width, seams_fix_mask_blur, seams_fix_padding]
|
493 |
+
)
|
494 |
+
|
495 |
+
def select_scale_type(scale_index):
|
496 |
+
is_custom_size = scale_index == 1
|
497 |
+
is_custom_scale = scale_index == 2
|
498 |
+
|
499 |
+
return [gr.update(visible=is_custom_size),
|
500 |
+
gr.update(visible=is_custom_size),
|
501 |
+
gr.update(visible=is_custom_scale),
|
502 |
+
]
|
503 |
+
|
504 |
+
target_size_type.change(
|
505 |
+
fn=select_scale_type,
|
506 |
+
inputs=target_size_type,
|
507 |
+
outputs=[custom_width, custom_height, custom_scale]
|
508 |
+
)
|
509 |
+
|
510 |
+
def init_field(scale_name):
|
511 |
+
try:
|
512 |
+
scale_index = target_size_types.index(scale_name)
|
513 |
+
custom_width.visible = custom_height.visible = scale_index == 1
|
514 |
+
custom_scale.visible = scale_index == 2
|
515 |
+
except:
|
516 |
+
pass
|
517 |
+
|
518 |
+
target_size_type.init_field = init_field
|
519 |
+
|
520 |
+
return [info, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding,
|
521 |
+
upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur,
|
522 |
+
seams_fix_type, target_size_type, custom_width, custom_height, custom_scale]
|
523 |
+
|
524 |
+
def run(self, p, _, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding,
|
525 |
+
upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur,
|
526 |
+
seams_fix_type, target_size_type, custom_width, custom_height, custom_scale):
|
527 |
+
|
528 |
+
# Init
|
529 |
+
processing.fix_seed(p)
|
530 |
+
devices.torch_gc()
|
531 |
+
|
532 |
+
p.do_not_save_grid = True
|
533 |
+
p.do_not_save_samples = True
|
534 |
+
p.inpaint_full_res = False
|
535 |
+
|
536 |
+
p.inpainting_fill = 1
|
537 |
+
p.n_iter = 1
|
538 |
+
p.batch_size = 1
|
539 |
+
|
540 |
+
seed = p.seed
|
541 |
+
|
542 |
+
# Init image
|
543 |
+
init_img = p.init_images[0]
|
544 |
+
if init_img == None:
|
545 |
+
return Processed(p, [], seed, "Empty image")
|
546 |
+
init_img = images.flatten(init_img, opts.img2img_background_color)
|
547 |
+
|
548 |
+
#override size
|
549 |
+
if target_size_type == 1:
|
550 |
+
p.width = custom_width
|
551 |
+
p.height = custom_height
|
552 |
+
if target_size_type == 2:
|
553 |
+
p.width = math.ceil((init_img.width * custom_scale) / 64) * 64
|
554 |
+
p.height = math.ceil((init_img.height * custom_scale) / 64) * 64
|
555 |
+
|
556 |
+
# Upscaling
|
557 |
+
upscaler = USDUpscaler(p, init_img, upscaler_index, save_upscaled_image, save_seams_fix_image, tile_width, tile_height)
|
558 |
+
upscaler.upscale()
|
559 |
+
|
560 |
+
# Drawing
|
561 |
+
upscaler.setup_redraw(redraw_mode, padding, mask_blur)
|
562 |
+
upscaler.setup_seams_fix(seams_fix_padding, seams_fix_denoise, seams_fix_mask_blur, seams_fix_width, seams_fix_type)
|
563 |
+
upscaler.print_info()
|
564 |
+
upscaler.add_extra_info()
|
565 |
+
upscaler.process()
|
566 |
+
result_images = upscaler.result_images
|
567 |
+
|
568 |
+
return Processed(p, result_images, seed, upscaler.initial_info if upscaler.initial_info is not None else "")
|
569 |
+
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/usdu_patch.py
ADDED
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Make some patches to the script
|
2 |
+
from repositories import ultimate_upscale as usdu
|
3 |
+
import modules.shared as shared
|
4 |
+
import math
|
5 |
+
from PIL import Image
|
6 |
+
|
7 |
+
|
8 |
+
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
|
9 |
+
Image.Resampling = Image
|
10 |
+
|
11 |
+
#
|
12 |
+
# Instead of using multiples of 64, use multiples of 8
|
13 |
+
#
|
14 |
+
|
15 |
+
|
16 |
+
def round_length(length, multiple=8):
|
17 |
+
return math.ceil(length / multiple) * multiple
|
18 |
+
|
19 |
+
|
20 |
+
# Upscaler
|
21 |
+
old_init = usdu.USDUpscaler.__init__
|
22 |
+
|
23 |
+
|
24 |
+
def new_init(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height):
|
25 |
+
p.width = round_length(image.width * p.upscale_by)
|
26 |
+
p.height = round_length(image.height * p.upscale_by)
|
27 |
+
old_init(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height)
|
28 |
+
|
29 |
+
|
30 |
+
usdu.USDUpscaler.__init__ = new_init
|
31 |
+
|
32 |
+
# Redraw
|
33 |
+
old_setup_redraw = usdu.USDURedraw.init_draw
|
34 |
+
|
35 |
+
|
36 |
+
def new_setup_redraw(self, p, width, height):
|
37 |
+
mask, draw = old_setup_redraw(self, p, width, height)
|
38 |
+
p.width = round_length(self.tile_width + self.padding)
|
39 |
+
p.height = round_length(self.tile_height + self.padding)
|
40 |
+
return mask, draw
|
41 |
+
|
42 |
+
|
43 |
+
usdu.USDURedraw.init_draw = new_setup_redraw
|
44 |
+
|
45 |
+
# Seams fix
|
46 |
+
old_setup_seams_fix = usdu.USDUSeamsFix.init_draw
|
47 |
+
|
48 |
+
|
49 |
+
def new_setup_seams_fix(self, p):
|
50 |
+
old_setup_seams_fix(self, p)
|
51 |
+
p.width = round_length(self.tile_width + self.padding)
|
52 |
+
p.height = round_length(self.tile_height + self.padding)
|
53 |
+
|
54 |
+
|
55 |
+
usdu.USDUSeamsFix.init_draw = new_setup_seams_fix
|
56 |
+
|
57 |
+
|
58 |
+
#
|
59 |
+
# Make the script upscale on a batch of images instead of one image
|
60 |
+
#
|
61 |
+
|
62 |
+
old_upscale = usdu.USDUpscaler.upscale
|
63 |
+
|
64 |
+
|
65 |
+
def new_upscale(self):
|
66 |
+
old_upscale(self)
|
67 |
+
shared.batch = [self.image] + \
|
68 |
+
[img.resize((self.p.width, self.p.height), resample=Image.LANCZOS) for img in shared.batch[1:]]
|
69 |
+
|
70 |
+
|
71 |
+
usdu.USDUpscaler.upscale = new_upscale
|
ComfyUI/custom_nodes/ComfyUI_UltimateSDUpscale/utils.py
ADDED
@@ -0,0 +1,460 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
from PIL import Image, ImageFilter
|
3 |
+
import torch
|
4 |
+
import torch.nn.functional as F
|
5 |
+
from torchvision.transforms import GaussianBlur
|
6 |
+
import math
|
7 |
+
|
8 |
+
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
|
9 |
+
Image.Resampling = Image
|
10 |
+
|
11 |
+
BLUR_KERNEL_SIZE = 15
|
12 |
+
|
13 |
+
|
14 |
+
def tensor_to_pil(img_tensor, batch_index=0):
|
15 |
+
# Takes an image in a batch in the form of a tensor of shape [batch_size, channels, height, width]
|
16 |
+
# and returns an PIL Image with the corresponding mode deduced by the number of channels
|
17 |
+
|
18 |
+
# Take the image in the batch given by batch_index
|
19 |
+
img_tensor = img_tensor[batch_index].unsqueeze(0)
|
20 |
+
i = 255. * img_tensor.cpu().numpy()
|
21 |
+
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8).squeeze())
|
22 |
+
return img
|
23 |
+
|
24 |
+
|
25 |
+
def pil_to_tensor(image):
|
26 |
+
# Takes a PIL image and returns a tensor of shape [1, height, width, channels]
|
27 |
+
image = np.array(image).astype(np.float32) / 255.0
|
28 |
+
image = torch.from_numpy(image).unsqueeze(0)
|
29 |
+
if len(image.shape) == 3: # If the image is grayscale, add a channel dimension
|
30 |
+
image = image.unsqueeze(-1)
|
31 |
+
return image
|
32 |
+
|
33 |
+
|
34 |
+
def controlnet_hint_to_pil(tensor, batch_index=0):
|
35 |
+
return tensor_to_pil(tensor.movedim(1, -1), batch_index)
|
36 |
+
|
37 |
+
|
38 |
+
def pil_to_controlnet_hint(img):
|
39 |
+
return pil_to_tensor(img).movedim(-1, 1)
|
40 |
+
|
41 |
+
|
42 |
+
def crop_tensor(tensor, region):
|
43 |
+
# Takes a tensor of shape [batch_size, height, width, channels] and crops it to the given region
|
44 |
+
x1, y1, x2, y2 = region
|
45 |
+
return tensor[:, y1:y2, x1:x2, :]
|
46 |
+
|
47 |
+
|
48 |
+
def resize_tensor(tensor, size, mode="nearest-exact"):
|
49 |
+
# Takes a tensor of shape [B, C, H, W] and resizes
|
50 |
+
# it to a shape of [B, C, size[0], size[1]] using the given mode
|
51 |
+
return torch.nn.functional.interpolate(tensor, size=size, mode=mode)
|
52 |
+
|
53 |
+
|
54 |
+
def get_crop_region(mask, pad=0):
|
55 |
+
# Takes a black and white PIL image in 'L' mode and returns the coordinates of the white rectangular mask region
|
56 |
+
# Should be equivalent to the get_crop_region function from https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/masking.py
|
57 |
+
coordinates = mask.getbbox()
|
58 |
+
if coordinates is not None:
|
59 |
+
x1, y1, x2, y2 = coordinates
|
60 |
+
else:
|
61 |
+
x1, y1, x2, y2 = mask.width, mask.height, 0, 0
|
62 |
+
# Apply padding
|
63 |
+
x1 = max(x1 - pad, 0)
|
64 |
+
y1 = max(y1 - pad, 0)
|
65 |
+
x2 = min(x2 + pad, mask.width)
|
66 |
+
y2 = min(y2 + pad, mask.height)
|
67 |
+
return fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height))
|
68 |
+
|
69 |
+
|
70 |
+
def fix_crop_region(region, image_size):
|
71 |
+
# Remove the extra pixel added by the get_crop_region function
|
72 |
+
image_width, image_height = image_size
|
73 |
+
x1, y1, x2, y2 = region
|
74 |
+
if x2 < image_width:
|
75 |
+
x2 -= 1
|
76 |
+
if y2 < image_height:
|
77 |
+
y2 -= 1
|
78 |
+
return x1, y1, x2, y2
|
79 |
+
|
80 |
+
|
81 |
+
def expand_crop(region, width, height, target_width, target_height):
|
82 |
+
'''
|
83 |
+
Expands a crop region to a specified target size.
|
84 |
+
:param region: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points
|
85 |
+
of the rectangular region. Expected to have x2 > x1 and y2 > y1.
|
86 |
+
:param width: The width of the image the crop region is from.
|
87 |
+
:param height: The height of the image the crop region is from.
|
88 |
+
:param target_width: The desired width of the crop region.
|
89 |
+
:param target_height: The desired height of the crop region.
|
90 |
+
'''
|
91 |
+
x1, y1, x2, y2 = region
|
92 |
+
actual_width = x2 - x1
|
93 |
+
actual_height = y2 - y1
|
94 |
+
# target_width = math.ceil(actual_width / 8) * 8
|
95 |
+
# target_height = math.ceil(actual_height / 8) * 8
|
96 |
+
|
97 |
+
# Try to expand region to the right of half the difference
|
98 |
+
width_diff = target_width - actual_width
|
99 |
+
x2 = min(x2 + width_diff // 2, width)
|
100 |
+
# Expand region to the left of the difference including the pixels that could not be expanded to the right
|
101 |
+
width_diff = target_width - (x2 - x1)
|
102 |
+
x1 = max(x1 - width_diff, 0)
|
103 |
+
# Try the right again
|
104 |
+
width_diff = target_width - (x2 - x1)
|
105 |
+
x2 = min(x2 + width_diff, width)
|
106 |
+
|
107 |
+
# Try to expand region to the bottom of half the difference
|
108 |
+
height_diff = target_height - actual_height
|
109 |
+
y2 = min(y2 + height_diff // 2, height)
|
110 |
+
# Expand region to the top of the difference including the pixels that could not be expanded to the bottom
|
111 |
+
height_diff = target_height - (y2 - y1)
|
112 |
+
y1 = max(y1 - height_diff, 0)
|
113 |
+
# Try the bottom again
|
114 |
+
height_diff = target_height - (y2 - y1)
|
115 |
+
y2 = min(y2 + height_diff, height)
|
116 |
+
|
117 |
+
return (x1, y1, x2, y2), (target_width, target_height)
|
118 |
+
|
119 |
+
|
120 |
+
def resize_region(region, init_size, resize_size):
|
121 |
+
# Resize a crop so that it fits an image that was resized to the given width and height
|
122 |
+
x1, y1, x2, y2 = region
|
123 |
+
init_width, init_height = init_size
|
124 |
+
resize_width, resize_height = resize_size
|
125 |
+
x1 = math.floor(x1 * resize_width / init_width)
|
126 |
+
x2 = math.ceil(x2 * resize_width / init_width)
|
127 |
+
y1 = math.floor(y1 * resize_height / init_height)
|
128 |
+
y2 = math.ceil(y2 * resize_height / init_height)
|
129 |
+
return (x1, y1, x2, y2)
|
130 |
+
|
131 |
+
|
132 |
+
def pad_image(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
|
133 |
+
'''
|
134 |
+
Pads an image with the given number of pixels on each side and fills the padding with data from the edges.
|
135 |
+
:param image: A PIL image
|
136 |
+
:param left_pad: The number of pixels to pad on the left side
|
137 |
+
:param right_pad: The number of pixels to pad on the right side
|
138 |
+
:param top_pad: The number of pixels to pad on the top side
|
139 |
+
:param bottom_pad: The number of pixels to pad on the bottom side
|
140 |
+
:param blur: Whether to blur the padded edges
|
141 |
+
:return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad)
|
142 |
+
'''
|
143 |
+
left_edge = image.crop((0, 1, 1, image.height - 1))
|
144 |
+
right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1))
|
145 |
+
top_edge = image.crop((1, 0, image.width - 1, 1))
|
146 |
+
bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height))
|
147 |
+
new_width = image.width + left_pad + right_pad
|
148 |
+
new_height = image.height + top_pad + bottom_pad
|
149 |
+
padded_image = Image.new(image.mode, (new_width, new_height))
|
150 |
+
padded_image.paste(image, (left_pad, top_pad))
|
151 |
+
if fill:
|
152 |
+
for i in range(left_pad):
|
153 |
+
edge = left_edge.resize(
|
154 |
+
(1, new_height - i * (top_pad + bottom_pad) // left_pad), resample=Image.Resampling.NEAREST)
|
155 |
+
padded_image.paste(edge, (i, i * top_pad // left_pad))
|
156 |
+
for i in range(right_pad):
|
157 |
+
edge = right_edge.resize(
|
158 |
+
(1, new_height - i * (top_pad + bottom_pad) // right_pad), resample=Image.Resampling.NEAREST)
|
159 |
+
padded_image.paste(edge, (new_width - 1 - i, i * top_pad // right_pad))
|
160 |
+
for i in range(top_pad):
|
161 |
+
edge = top_edge.resize(
|
162 |
+
(new_width - i * (left_pad + right_pad) // top_pad, 1), resample=Image.Resampling.NEAREST)
|
163 |
+
padded_image.paste(edge, (i * left_pad // top_pad, i))
|
164 |
+
for i in range(bottom_pad):
|
165 |
+
edge = bottom_edge.resize(
|
166 |
+
(new_width - i * (left_pad + right_pad) // bottom_pad, 1), resample=Image.Resampling.NEAREST)
|
167 |
+
padded_image.paste(edge, (i * left_pad // bottom_pad, new_height - 1 - i))
|
168 |
+
if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0):
|
169 |
+
padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE))
|
170 |
+
padded_image.paste(image, (left_pad, top_pad))
|
171 |
+
return padded_image
|
172 |
+
|
173 |
+
|
174 |
+
def pad_image2(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
|
175 |
+
'''
|
176 |
+
Pads an image with the given number of pixels on each side and fills the padding with data from the edges.
|
177 |
+
Faster than pad_image, but only pads with edge data in straight lines.
|
178 |
+
:param image: A PIL image
|
179 |
+
:param left_pad: The number of pixels to pad on the left side
|
180 |
+
:param right_pad: The number of pixels to pad on the right side
|
181 |
+
:param top_pad: The number of pixels to pad on the top side
|
182 |
+
:param bottom_pad: The number of pixels to pad on the bottom side
|
183 |
+
:param blur: Whether to blur the padded edges
|
184 |
+
:return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad)
|
185 |
+
'''
|
186 |
+
left_edge = image.crop((0, 1, 1, image.height - 1))
|
187 |
+
right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1))
|
188 |
+
top_edge = image.crop((1, 0, image.width - 1, 1))
|
189 |
+
bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height))
|
190 |
+
new_width = image.width + left_pad + right_pad
|
191 |
+
new_height = image.height + top_pad + bottom_pad
|
192 |
+
padded_image = Image.new(image.mode, (new_width, new_height))
|
193 |
+
padded_image.paste(image, (left_pad, top_pad))
|
194 |
+
if fill:
|
195 |
+
if left_pad > 0:
|
196 |
+
padded_image.paste(left_edge.resize((left_pad, new_height), resample=Image.Resampling.NEAREST), (0, 0))
|
197 |
+
if right_pad > 0:
|
198 |
+
padded_image.paste(right_edge.resize((right_pad, new_height),
|
199 |
+
resample=Image.Resampling.NEAREST), (new_width - right_pad, 0))
|
200 |
+
if top_pad > 0:
|
201 |
+
padded_image.paste(top_edge.resize((new_width, top_pad), resample=Image.Resampling.NEAREST), (0, 0))
|
202 |
+
if bottom_pad > 0:
|
203 |
+
padded_image.paste(bottom_edge.resize((new_width, bottom_pad),
|
204 |
+
resample=Image.Resampling.NEAREST), (0, new_height - bottom_pad))
|
205 |
+
if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0):
|
206 |
+
padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE))
|
207 |
+
padded_image.paste(image, (left_pad, top_pad))
|
208 |
+
return padded_image
|
209 |
+
|
210 |
+
|
211 |
+
def pad_tensor(tensor, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
|
212 |
+
'''
|
213 |
+
Pads an image tensor with the given number of pixels on each side and fills the padding with data from the edges.
|
214 |
+
:param tensor: A tensor of shape [B, H, W, C]
|
215 |
+
:param left_pad: The number of pixels to pad on the left side
|
216 |
+
:param right_pad: The number of pixels to pad on the right side
|
217 |
+
:param top_pad: The number of pixels to pad on the top side
|
218 |
+
:param bottom_pad: The number of pixels to pad on the bottom side
|
219 |
+
:param blur: Whether to blur the padded edges
|
220 |
+
:return: A tensor of shape [B, H + top_pad + bottom_pad, W + left_pad + right_pad, C]
|
221 |
+
'''
|
222 |
+
batch_size, channels, height, width = tensor.shape
|
223 |
+
h_pad = left_pad + right_pad
|
224 |
+
v_pad = top_pad + bottom_pad
|
225 |
+
new_width = width + h_pad
|
226 |
+
new_height = height + v_pad
|
227 |
+
|
228 |
+
# Create empty image
|
229 |
+
padded = torch.zeros((batch_size, channels, new_height, new_width), dtype=tensor.dtype)
|
230 |
+
|
231 |
+
# Copy the original image into the centor of the padded tensor
|
232 |
+
padded[:, :, top_pad:top_pad + height, left_pad:left_pad + width] = tensor
|
233 |
+
|
234 |
+
# Duplicate the edges of the original image into the padding
|
235 |
+
if top_pad > 0:
|
236 |
+
padded[:, :, :top_pad, :] = padded[:, :, top_pad:top_pad + 1, :] # Top edge
|
237 |
+
if bottom_pad > 0:
|
238 |
+
padded[:, :, -bottom_pad:, :] = padded[:, :, -bottom_pad - 1:-bottom_pad, :] # Bottom edge
|
239 |
+
if left_pad > 0:
|
240 |
+
padded[:, :, :, :left_pad] = padded[:, :, :, left_pad:left_pad + 1] # Left edge
|
241 |
+
if right_pad > 0:
|
242 |
+
padded[:, :, :, -right_pad:] = padded[:, :, :, -right_pad - 1:-right_pad] # Right edge
|
243 |
+
|
244 |
+
return padded
|
245 |
+
|
246 |
+
|
247 |
+
def resize_and_pad_image(image, width, height, fill=False, blur=False):
|
248 |
+
'''
|
249 |
+
Resizes an image to the given width and height and pads it to the given width and height.
|
250 |
+
:param image: A PIL image
|
251 |
+
:param width: The width of the resized image
|
252 |
+
:param height: The height of the resized image
|
253 |
+
:param fill: Whether to fill the padding with data from the edges
|
254 |
+
:param blur: Whether to blur the padded edges
|
255 |
+
:return: A PIL image of size (width, height)
|
256 |
+
'''
|
257 |
+
width_ratio = width / image.width
|
258 |
+
height_ratio = height / image.height
|
259 |
+
if height_ratio > width_ratio:
|
260 |
+
resize_ratio = width_ratio
|
261 |
+
else:
|
262 |
+
resize_ratio = height_ratio
|
263 |
+
resize_width = round(image.width * resize_ratio)
|
264 |
+
resize_height = round(image.height * resize_ratio)
|
265 |
+
resized = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS)
|
266 |
+
# Pad the sides of the image to get the image to the desired size that wasn't covered by the resize
|
267 |
+
horizontal_pad = (width - resize_width) // 2
|
268 |
+
vertical_pad = (height - resize_height) // 2
|
269 |
+
result = pad_image2(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur)
|
270 |
+
result = result.resize((width, height), resample=Image.Resampling.LANCZOS)
|
271 |
+
return result, (horizontal_pad, vertical_pad)
|
272 |
+
|
273 |
+
|
274 |
+
def resize_and_pad_tensor(tensor, width, height, fill=False, blur=False):
|
275 |
+
'''
|
276 |
+
Resizes an image tensor to the given width and height and pads it to the given width and height.
|
277 |
+
:param tensor: A tensor of shape [B, H, W, C]
|
278 |
+
:param width: The width of the resized image
|
279 |
+
:param height: The height of the resized image
|
280 |
+
:param fill: Whether to fill the padding with data from the edges
|
281 |
+
:param blur: Whether to blur the padded edges
|
282 |
+
:return: A tensor of shape [B, height, width, C]
|
283 |
+
'''
|
284 |
+
# Resize the image to the closest size that maintains the aspect ratio
|
285 |
+
width_ratio = width / tensor.shape[3]
|
286 |
+
height_ratio = height / tensor.shape[2]
|
287 |
+
if height_ratio > width_ratio:
|
288 |
+
resize_ratio = width_ratio
|
289 |
+
else:
|
290 |
+
resize_ratio = height_ratio
|
291 |
+
resize_width = round(tensor.shape[3] * resize_ratio)
|
292 |
+
resize_height = round(tensor.shape[2] * resize_ratio)
|
293 |
+
resized = F.interpolate(tensor, size=(resize_height, resize_width), mode='nearest-exact')
|
294 |
+
# Pad the sides of the image to get the image to the desired size that wasn't covered by the resize
|
295 |
+
horizontal_pad = (width - resize_width) // 2
|
296 |
+
vertical_pad = (height - resize_height) // 2
|
297 |
+
result = pad_tensor(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur)
|
298 |
+
result = F.interpolate(result, size=(height, width), mode='nearest-exact')
|
299 |
+
return result
|
300 |
+
|
301 |
+
|
302 |
+
def crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
|
303 |
+
if "control" not in cond_dict:
|
304 |
+
return
|
305 |
+
c = cond_dict["control"]
|
306 |
+
controlnet = c.copy()
|
307 |
+
cond_dict["control"] = controlnet
|
308 |
+
while c is not None:
|
309 |
+
# hint is shape (B, C, H, W)
|
310 |
+
hint = controlnet.cond_hint_original
|
311 |
+
resized_crop = resize_region(region, canvas_size, hint.shape[:-3:-1])
|
312 |
+
hint = crop_tensor(hint.movedim(1, -1), resized_crop).movedim(-1, 1)
|
313 |
+
hint = resize_tensor(hint, tile_size[::-1])
|
314 |
+
controlnet.cond_hint_original = hint
|
315 |
+
c = c.previous_controlnet
|
316 |
+
controlnet.set_previous_controlnet(c.copy() if c is not None else None)
|
317 |
+
controlnet = controlnet.previous_controlnet
|
318 |
+
|
319 |
+
|
320 |
+
def region_intersection(region1, region2):
|
321 |
+
"""
|
322 |
+
Returns the coordinates of the intersection of two rectangular regions.
|
323 |
+
:param region1: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points
|
324 |
+
of the first rectangular region. Expected to have x2 > x1 and y2 > y1.
|
325 |
+
:param region2: The second rectangular region with the same format as the first.
|
326 |
+
:return: A tuple of the form (x1, y1, x2, y2) denoting the rectangular intersection.
|
327 |
+
None if there is no intersection.
|
328 |
+
"""
|
329 |
+
x1, y1, x2, y2 = region1
|
330 |
+
x1_, y1_, x2_, y2_ = region2
|
331 |
+
x1 = max(x1, x1_)
|
332 |
+
y1 = max(y1, y1_)
|
333 |
+
x2 = min(x2, x2_)
|
334 |
+
y2 = min(y2, y2_)
|
335 |
+
if x1 >= x2 or y1 >= y2:
|
336 |
+
return None
|
337 |
+
return (x1, y1, x2, y2)
|
338 |
+
|
339 |
+
|
340 |
+
def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
|
341 |
+
if "gligen" not in cond_dict:
|
342 |
+
return
|
343 |
+
type, model, cond = cond_dict["gligen"]
|
344 |
+
if type != "position":
|
345 |
+
from warnings import warn
|
346 |
+
warn(f"Unknown gligen type {type}")
|
347 |
+
return
|
348 |
+
cropped = []
|
349 |
+
for c in cond:
|
350 |
+
emb, h, w, y, x = c
|
351 |
+
# Get the coordinates of the box in the upscaled image
|
352 |
+
x1 = x * 8
|
353 |
+
y1 = y * 8
|
354 |
+
x2 = x1 + w * 8
|
355 |
+
y2 = y1 + h * 8
|
356 |
+
gligen_upscaled_box = resize_region((x1, y1, x2, y2), init_size, canvas_size)
|
357 |
+
|
358 |
+
# Calculate the intersection of the gligen box and the region
|
359 |
+
intersection = region_intersection(gligen_upscaled_box, region)
|
360 |
+
if intersection is None:
|
361 |
+
continue
|
362 |
+
x1, y1, x2, y2 = intersection
|
363 |
+
|
364 |
+
# Offset the gligen box so that the origin is at the top left of the tile region
|
365 |
+
x1 -= region[0]
|
366 |
+
y1 -= region[1]
|
367 |
+
x2 -= region[0]
|
368 |
+
y2 -= region[1]
|
369 |
+
|
370 |
+
# Add the padding
|
371 |
+
x1 += w_pad
|
372 |
+
y1 += h_pad
|
373 |
+
x2 += w_pad
|
374 |
+
y2 += h_pad
|
375 |
+
|
376 |
+
# Set the new position params
|
377 |
+
h = (y2 - y1) // 8
|
378 |
+
w = (x2 - x1) // 8
|
379 |
+
x = x1 // 8
|
380 |
+
y = y1 // 8
|
381 |
+
cropped.append((emb, h, w, y, x))
|
382 |
+
|
383 |
+
cond_dict["gligen"] = (type, model, cropped)
|
384 |
+
|
385 |
+
|
386 |
+
def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
|
387 |
+
if "area" not in cond_dict:
|
388 |
+
return
|
389 |
+
|
390 |
+
# Resize the area conditioning to the canvas size and confine it to the tile region
|
391 |
+
h, w, y, x = cond_dict["area"]
|
392 |
+
w, h, x, y = 8 * w, 8 * h, 8 * x, 8 * y
|
393 |
+
x1, y1, x2, y2 = resize_region((x, y, x + w, y + h), init_size, canvas_size)
|
394 |
+
intersection = region_intersection((x1, y1, x2, y2), region)
|
395 |
+
if intersection is None:
|
396 |
+
del cond_dict["area"]
|
397 |
+
del cond_dict["strength"]
|
398 |
+
return
|
399 |
+
x1, y1, x2, y2 = intersection
|
400 |
+
|
401 |
+
# Offset origin to the top left of the tile
|
402 |
+
x1 -= region[0]
|
403 |
+
y1 -= region[1]
|
404 |
+
x2 -= region[0]
|
405 |
+
y2 -= region[1]
|
406 |
+
|
407 |
+
# Add the padding
|
408 |
+
x1 += w_pad
|
409 |
+
y1 += h_pad
|
410 |
+
x2 += w_pad
|
411 |
+
y2 += h_pad
|
412 |
+
|
413 |
+
# Set the params for tile
|
414 |
+
w, h = (x2 - x1) // 8, (y2 - y1) // 8
|
415 |
+
x, y = x1 // 8, y1 // 8
|
416 |
+
|
417 |
+
cond_dict["area"] = (h, w, y, x)
|
418 |
+
|
419 |
+
|
420 |
+
def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
|
421 |
+
if "mask" not in cond_dict:
|
422 |
+
return
|
423 |
+
mask_tensor = cond_dict["mask"] # (B, H, W)
|
424 |
+
masks = []
|
425 |
+
for i in range(mask_tensor.shape[0]):
|
426 |
+
# Convert to PIL image
|
427 |
+
mask = tensor_to_pil(mask_tensor, i) # W x H
|
428 |
+
|
429 |
+
# Resize the mask to the canvas size
|
430 |
+
mask = mask.resize(canvas_size, Image.Resampling.BICUBIC)
|
431 |
+
|
432 |
+
# Crop the mask to the region
|
433 |
+
mask = mask.crop(region)
|
434 |
+
|
435 |
+
# Add padding
|
436 |
+
mask, _ = resize_and_pad_image(mask, tile_size[0], tile_size[1], fill=True)
|
437 |
+
|
438 |
+
# Resize the mask to the tile size
|
439 |
+
if tile_size != mask.size:
|
440 |
+
mask = mask.resize(tile_size, Image.Resampling.BICUBIC)
|
441 |
+
|
442 |
+
# Convert back to tensor
|
443 |
+
mask = pil_to_tensor(mask) # (1, H, W, 1)
|
444 |
+
mask = mask.squeeze(-1) # (1, H, W)
|
445 |
+
masks.append(mask)
|
446 |
+
|
447 |
+
cond_dict["mask"] = torch.cat(masks, dim=0) # (B, H, W)
|
448 |
+
|
449 |
+
|
450 |
+
def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0):
|
451 |
+
cropped = []
|
452 |
+
for emb, x in cond:
|
453 |
+
cond_dict = x.copy()
|
454 |
+
n = [emb, cond_dict]
|
455 |
+
crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
456 |
+
crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
457 |
+
crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
458 |
+
crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
459 |
+
cropped.append(n)
|
460 |
+
return cropped
|