mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
106 lines
3.0 KiB
Python
106 lines
3.0 KiB
Python
import random
|
|
import string
|
|
import time
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
|
|
from backend import memory_management
|
|
|
|
|
|
def prepare_free_memory(aggressive=False):
|
|
if aggressive:
|
|
memory_management.unload_all_models()
|
|
print("Cleanup all memory...")
|
|
return
|
|
|
|
memory_management.free_memory(memory_required=memory_management.minimum_inference_memory(), device=memory_management.get_torch_device())
|
|
print("Cleanup minimal inference memory...")
|
|
|
|
|
|
def apply_circular_forge(model, tiling_enabled=False):
|
|
if not model.is_webui_legacy_model():
|
|
return
|
|
if model.tiling_enabled == tiling_enabled:
|
|
return
|
|
|
|
model.tiling_enabled = tiling_enabled
|
|
|
|
unet: torch.nn.Module = model.forge_objects.unet.model.diffusion_model
|
|
for layer in [layer for layer in unet.modules() if isinstance(layer, torch.nn.Conv2d)]:
|
|
layer.padding_mode = "circular" if tiling_enabled else "zeros"
|
|
|
|
vae: torch.nn.Module = model.forge_objects.vae.first_stage_model
|
|
for layer in [layer for layer in vae.modules() if isinstance(layer, torch.nn.Conv2d)]:
|
|
layer.padding_mode = "circular" if tiling_enabled else "zeros"
|
|
|
|
print(f"Tiling: {tiling_enabled}")
|
|
|
|
|
|
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 generate_random_filename(extension=".txt"):
|
|
timestamp = time.strftime("%Y%m%d-%H%M%S")
|
|
random_string = "".join(random.choices(string.ascii_lowercase + string.digits, k=5))
|
|
filename = f"{timestamp}-{random_string}{extension}"
|
|
return filename
|
|
|
|
|
|
@torch.no_grad()
|
|
@torch.inference_mode()
|
|
def pytorch_to_numpy(x):
|
|
return [np.clip(255.0 * y.cpu().numpy(), 0, 255).astype(np.uint8) for y in x]
|
|
|
|
|
|
@torch.no_grad()
|
|
@torch.inference_mode()
|
|
def numpy_to_pytorch(x):
|
|
y = x.astype(np.float32) / 255.0
|
|
y = y[None]
|
|
y = np.ascontiguousarray(y.copy())
|
|
y = torch.from_numpy(y).float()
|
|
return y
|
|
|
|
|
|
def pad64(x):
|
|
return int(np.ceil(float(x) / 64.0) * 64 - x)
|
|
|
|
|
|
def safer_memory(x):
|
|
# Fix many MAC/AMD problems
|
|
return np.ascontiguousarray(x.copy()).copy()
|
|
|
|
|
|
def resize_image_with_pad(img, resolution):
|
|
H_raw, W_raw, _ = img.shape
|
|
k = float(resolution) / float(min(H_raw, W_raw))
|
|
interpolation = cv2.INTER_CUBIC if k > 1 else cv2.INTER_AREA
|
|
H_target = int(np.round(float(H_raw) * k))
|
|
W_target = int(np.round(float(W_raw) * k))
|
|
img = cv2.resize(img, (W_target, H_target), interpolation=interpolation)
|
|
H_pad, W_pad = pad64(H_target), pad64(W_target)
|
|
img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode="edge")
|
|
|
|
def remove_pad(x):
|
|
return safer_memory(x[:H_target, :W_target])
|
|
|
|
return safer_memory(img_padded), remove_pad
|