mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-02-24 16:53:56 +00:00
80 lines
2.2 KiB
Python
80 lines
2.2 KiB
Python
import numpy as np
|
|
import cv2
|
|
import os
|
|
|
|
|
|
def load_model(filename: str, remote_url: str, model_dir: str) -> str:
|
|
"""
|
|
Load the model from the specified filename and remote URL if it doesn't exist locally.
|
|
|
|
Args:
|
|
filename (str): The filename of the model.
|
|
remote_url (str): The remote URL of the model.
|
|
"""
|
|
local_path = os.path.join(model_dir, filename)
|
|
if not os.path.exists(local_path):
|
|
from basicsr.utils.download_util import load_file_from_url
|
|
|
|
load_file_from_url(remote_url, model_dir=model_dir)
|
|
return local_path
|
|
|
|
|
|
def HWC3(x):
|
|
assert x.dtype == np.uint8
|
|
if x.ndim == 2:
|
|
x = x[:, :, None]
|
|
assert x.ndim == 3
|
|
H, W, C = x.shape
|
|
assert C == 1 or C == 3 or C == 4
|
|
if C == 3:
|
|
return x
|
|
if C == 1:
|
|
return np.concatenate([x, x, x], axis=2)
|
|
if C == 4:
|
|
color = x[:, :, 0:3].astype(np.float32)
|
|
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
|
|
y = color * alpha + 255.0 * (1.0 - alpha)
|
|
y = y.clip(0, 255).astype(np.uint8)
|
|
return y
|
|
|
|
|
|
def make_noise_disk(H, W, C, F):
|
|
noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
|
|
noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
|
|
noise = noise[F: F + H, F: F + W]
|
|
noise -= np.min(noise)
|
|
noise /= np.max(noise)
|
|
if C == 1:
|
|
noise = noise[:, :, None]
|
|
return noise
|
|
|
|
|
|
def nms(x, t, s):
|
|
x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
|
|
|
|
f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
|
|
f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
|
|
f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
|
|
f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
|
|
|
|
y = np.zeros_like(x)
|
|
|
|
for f in [f1, f2, f3, f4]:
|
|
np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
|
|
|
|
z = np.zeros_like(y, dtype=np.uint8)
|
|
z[y > t] = 255
|
|
return z
|
|
|
|
|
|
def min_max_norm(x):
|
|
x -= np.min(x)
|
|
x /= np.maximum(np.max(x), 1e-5)
|
|
return x
|
|
|
|
|
|
def safe_step(x, step=2):
|
|
y = x.astype(np.float32) * float(step + 1)
|
|
y = y.astype(np.int32).astype(np.float32) / float(step)
|
|
return y
|