fluxdev/stable-diffusion-webui-forge

1import torch2import numpy as np3import os4import time5import random6import string7import cv28 9from ldm_patched.modules import model_management10 11 12def prepare_free_memory(aggressive=False):13    if aggressive:14        model_management.unload_all_models()15        print('Cleanup all memory.')16        return17 18    model_management.free_memory(memory_required=model_management.minimum_inference_memory(),19                                 device=model_management.get_torch_device())20    print('Cleanup minimal inference memory.')21    return22 23 24def apply_circular_forge(model, tiling_enabled=False):25    if model.tiling_enabled == tiling_enabled:26        return27 28    print(f'Tiling: {tiling_enabled}')29    model.tiling_enabled = tiling_enabled30 31    def flatten(el):32        flattened = [flatten(children) for children in el.children()]33        res = [el]34        for c in flattened:35            res += c36        return res37 38    layers = flatten(model)39 40    for layer in [layer for layer in layers if 'Conv' in type(layer).__name__]:41        layer.padding_mode = 'circular' if tiling_enabled else 'zeros'42    return43 44 45def HWC3(x):46    assert x.dtype == np.uint847    if x.ndim == 2:48        x = x[:, :, None]49    assert x.ndim == 350    H, W, C = x.shape51    assert C == 1 or C == 3 or C == 452    if C == 3:53        return x54    if C == 1:55        return np.concatenate([x, x, x], axis=2)56    if C == 4:57        color = x[:, :, 0:3].astype(np.float32)58        alpha = x[:, :, 3:4].astype(np.float32) / 255.059        y = color * alpha + 255.0 * (1.0 - alpha)60        y = y.clip(0, 255).astype(np.uint8)61        return y62 63 64def generate_random_filename(extension=".txt"):65    timestamp = time.strftime("%Y%m%d-%H%M%S")66    random_string = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))67    filename = f"{timestamp}-{random_string}{extension}"68    return filename69 70 71@torch.no_grad()72@torch.inference_mode()73def pytorch_to_numpy(x):74    return [np.clip(255. * y.cpu().numpy(), 0, 255).astype(np.uint8) for y in x]75 76 77@torch.no_grad()78@torch.inference_mode()79def numpy_to_pytorch(x):80    y = x.astype(np.float32) / 255.081    y = y[None]82    y = np.ascontiguousarray(y.copy())83    y = torch.from_numpy(y).float()84    return y85 86 87def write_images_to_mp4(frame_list: list, filename=None, fps=6):88    from modules.paths_internal import default_output_dir89 90    video_folder = os.path.join(default_output_dir, 'svd')91    os.makedirs(video_folder, exist_ok=True)92 93    if filename is None:94        filename = generate_random_filename('.mp4')95 96    full_path = os.path.join(video_folder, filename)97 98    try:99        import av100    except ImportError:101        from launch import run_pip102        run_pip(103            "install imageio[pyav]",104            "imageio[pyav]",105        )106        import av107 108    options = {109        "crf": str(23)110    }111 112    output = av.open(full_path, "w")113 114    stream = output.add_stream('libx264', fps, options=options)115    stream.width = frame_list[0].shape[1]116    stream.height = frame_list[0].shape[0]117    for img in frame_list:118        frame = av.VideoFrame.from_ndarray(img)119        packet = stream.encode(frame)120        output.mux(packet)121    packet = stream.encode(None)122    output.mux(packet)123    output.close()124 125    return full_path126 127 128def pad64(x):129    return int(np.ceil(float(x) / 64.0) * 64 - x)130 131 132def safer_memory(x):133    # Fix many MAC/AMD problems134    return np.ascontiguousarray(x.copy()).copy()135 136 137def resize_image_with_pad(img, resolution):138    H_raw, W_raw, _ = img.shape139    k = float(resolution) / float(min(H_raw, W_raw))140    interpolation = cv2.INTER_CUBIC if k > 1 else cv2.INTER_AREA141    H_target = int(np.round(float(H_raw) * k))142    W_target = int(np.round(float(W_raw) * k))143    img = cv2.resize(img, (W_target, H_target), interpolation=interpolation)144    H_pad, W_pad = pad64(H_target), pad64(W_target)145    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode='edge')146 147    def remove_pad(x):148        return safer_memory(x[:H_target, :W_target])149 150    return safer_memory(img_padded), remove_pad151 152 153def lazy_memory_management(model):154    required_memory = model_management.module_size(model) + model_management.minimum_inference_memory()155    model_management.free_memory(required_memory, device=model_management.get_torch_device())156    return157