divisor.xflux1.util

View Source

 1# SPDX-License-Identifier:Apache-2.0
 2# original XFlux code from https://github.com/TencentARC/FluxKits
 3
 4import cv2
 5import numpy as np
 6
 7
 8def c_crop(image):
 9    width, height = image.size
10    new_size = min(width, height)
11    left = (width - new_size) / 2
12    top = (height - new_size) / 2
13    right = (width + new_size) / 2
14    bottom = (height + new_size) / 2
15    return image.crop((left, top, right, bottom))
16
17
18def pad64(x):
19    return int(np.ceil(float(x) / 64.0) * 64 - x)
20
21
22def HWC3(x):
23    assert x.dtype == np.uint8
24    if x.ndim == 2:
25        x = x[:, :, None]
26    assert x.ndim == 3
27    H, W, C = x.shape
28    assert C == 1 or C == 3 or C == 4
29    if C == 3:
30        return x
31    if C == 1:
32        return np.concatenate([x, x, x], axis=2)
33    if C == 4:
34        color = x[:, :, 0:3].astype(np.float32)
35        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
36        y = color * alpha + 255.0 * (1.0 - alpha)
37        y = y.clip(0, 255).astype(np.uint8)
38        return y
39
40
41def safer_memory(x):
42    # Fix many MAC/AMD problems
43    return np.ascontiguousarray(x.copy()).copy()
44
45
46# https://github.com/Mikubill/sd-webui-controlnet/blob/main/scripts/processor.py#L17
47# Added upscale_method, mode params
48def resize_image_with_pad(input_image, resolution, skip_hwc3=False, mode="edge"):
49    if skip_hwc3:
50        img = input_image
51    else:
52        img = HWC3(input_image)
53    H_raw, W_raw, _ = img.shape  # type: ignore
54    if resolution == 0:
55        return img, lambda x: x
56    k = float(resolution) / float(min(H_raw, W_raw))
57    H_target = int(np.round(float(H_raw) * k))
58    W_target = int(np.round(float(W_raw) * k))
59    img = cv2.resize(img, (W_target, H_target), interpolation=cv2.INTER_AREA)  # type: ignore
60    H_pad, W_pad = pad64(H_target), pad64(W_target)
61    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode=mode)  # type: ignore
62
63    def remove_pad(x):
64        return safer_memory(x[:H_target, :W_target, ...])
65
66    return safer_memory(img_padded), remove_pad

def c_crop(image): View Source

 9def c_crop(image):
10    width, height = image.size
11    new_size = min(width, height)
12    left = (width - new_size) / 2
13    top = (height - new_size) / 2
14    right = (width + new_size) / 2
15    bottom = (height + new_size) / 2
16    return image.crop((left, top, right, bottom))

def pad64(x): View Source

19def pad64(x):
20    return int(np.ceil(float(x) / 64.0) * 64 - x)

def HWC3(x): View Source

23def HWC3(x):
24    assert x.dtype == np.uint8
25    if x.ndim == 2:
26        x = x[:, :, None]
27    assert x.ndim == 3
28    H, W, C = x.shape
29    assert C == 1 or C == 3 or C == 4
30    if C == 3:
31        return x
32    if C == 1:
33        return np.concatenate([x, x, x], axis=2)
34    if C == 4:
35        color = x[:, :, 0:3].astype(np.float32)
36        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
37        y = color * alpha + 255.0 * (1.0 - alpha)
38        y = y.clip(0, 255).astype(np.uint8)
39        return y

def safer_memory(x): View Source

42def safer_memory(x):
43    # Fix many MAC/AMD problems
44    return np.ascontiguousarray(x.copy()).copy()

def resize_image_with_pad(input_image, resolution, skip_hwc3=False, mode='edge'): View Source

49def resize_image_with_pad(input_image, resolution, skip_hwc3=False, mode="edge"):
50    if skip_hwc3:
51        img = input_image
52    else:
53        img = HWC3(input_image)
54    H_raw, W_raw, _ = img.shape  # type: ignore
55    if resolution == 0:
56        return img, lambda x: x
57    k = float(resolution) / float(min(H_raw, W_raw))
58    H_target = int(np.round(float(H_raw) * k))
59    W_target = int(np.round(float(W_raw) * k))
60    img = cv2.resize(img, (W_target, H_target), interpolation=cv2.INTER_AREA)  # type: ignore
61    H_pad, W_pad = pad64(H_target), pad64(W_target)
62    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode=mode)  # type: ignore
63
64    def remove_pad(x):
65        return safer_memory(x[:H_target, :W_target, ...])
66
67    return safer_memory(img_padded), remove_pad