Freak-ppa commited on
Commit
1ce42b4
1 Parent(s): dc5c77a

Upload 16 files

Browse files
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