mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2025-12-28 05:35:00 +00:00
273 lines
8.5 KiB
Python
273 lines
8.5 KiB
Python
from modules_forge.initialization import initialize_forge
|
|
|
|
initialize_forge()
|
|
|
|
import sys
|
|
import types
|
|
import os
|
|
import torch
|
|
import inspect
|
|
import functools
|
|
import gradio.oauth
|
|
import gradio.routes
|
|
import contextlib
|
|
|
|
from backend import memory_management
|
|
from backend.operations import DynamicSwapInstaller
|
|
from diffusers.models import modeling_utils as diffusers_modeling_utils
|
|
from transformers import modeling_utils as transformers_modeling_utils
|
|
from backend.attention import AttentionProcessorForge
|
|
from starlette.requests import Request
|
|
|
|
|
|
_original_init = Request.__init__
|
|
|
|
|
|
def patched_init(self, scope, receive=None, send=None):
|
|
if 'session' not in scope:
|
|
scope['session'] = dict()
|
|
_original_init(self, scope, receive, send)
|
|
return
|
|
|
|
|
|
Request.__init__ = patched_init
|
|
gradio.oauth.attach_oauth = lambda x: None
|
|
gradio.routes.attach_oauth = lambda x: None
|
|
|
|
ALWAYS_SWAP = False
|
|
|
|
module_in_gpu: torch.nn.Module = None
|
|
gpu = memory_management.get_torch_device()
|
|
cpu = torch.device('cpu')
|
|
|
|
diffusers_modeling_utils.get_parameter_device = lambda *args, **kwargs: gpu
|
|
transformers_modeling_utils.get_parameter_device = lambda *args, **kwargs: gpu
|
|
|
|
|
|
def unload_module():
|
|
global module_in_gpu
|
|
|
|
if module_in_gpu is None:
|
|
return
|
|
|
|
DynamicSwapInstaller.uninstall_model(module_in_gpu)
|
|
module_in_gpu.to(cpu)
|
|
print(f'Move module to CPU: {type(module_in_gpu).__name__}')
|
|
|
|
module_in_gpu = None
|
|
memory_management.soft_empty_cache()
|
|
return
|
|
|
|
|
|
def greedy_move_to_gpu(model, model_gpu_memory_when_using_cpu_swap):
|
|
mem_counter = 0
|
|
memory_in_swap = 0
|
|
for m in model.modules():
|
|
if hasattr(m, "weight"):
|
|
module_mem = memory_management.module_size(m)
|
|
if mem_counter + module_mem < model_gpu_memory_when_using_cpu_swap:
|
|
m.to(gpu)
|
|
mem_counter += module_mem
|
|
else:
|
|
m.to(cpu)
|
|
memory_in_swap += module_mem
|
|
|
|
print(f"[Memory Management] Loaded to CPU Swap: {memory_in_swap / (1024 * 1024):.2f} MB")
|
|
print(f"[Memory Management] Loaded to GPU: {mem_counter / (1024 * 1024):.2f} MB")
|
|
return
|
|
|
|
|
|
def load_module(m):
|
|
global module_in_gpu
|
|
|
|
if module_in_gpu == m:
|
|
return
|
|
|
|
unload_module()
|
|
|
|
model_memory = memory_management.module_size(m)
|
|
current_free_mem = memory_management.get_free_memory(gpu)
|
|
inference_memory = 1.5 * 1024 * 1024 * 1024 # memory_management.minimum_inference_memory() # TODO: connect to main memory system
|
|
estimated_remaining_memory = current_free_mem - model_memory - inference_memory
|
|
|
|
print(f"[Memory Management] Current Free GPU Memory: {current_free_mem / (1024 * 1024):.2f} MB")
|
|
print(f"[Memory Management] Required Model Memory: {model_memory / (1024 * 1024):.2f} MB")
|
|
print(f"[Memory Management] Required Inference Memory: {inference_memory / (1024 * 1024):.2f} MB")
|
|
print(f"[Memory Management] Estimated Remaining GPU Memory: {estimated_remaining_memory / (1024 * 1024):.2f} MB")
|
|
|
|
is_torch_jit = 'ScriptModule' in type(m).__name__
|
|
|
|
if is_torch_jit:
|
|
print(f'Detected torch jit module: {type(m).__name__}')
|
|
|
|
if (ALWAYS_SWAP or estimated_remaining_memory < 0) and not is_torch_jit:
|
|
print(f'Move module to SWAP: {type(m).__name__}')
|
|
DynamicSwapInstaller.install_model(m, target_device=gpu)
|
|
model_gpu_memory_when_using_cpu_swap = memory_management.compute_model_gpu_memory_when_using_cpu_swap(current_free_mem, inference_memory)
|
|
greedy_move_to_gpu(m, model_gpu_memory_when_using_cpu_swap)
|
|
else:
|
|
print(f'Move module to GPU: {type(m).__name__}')
|
|
m.to(gpu)
|
|
|
|
module_in_gpu = m
|
|
return
|
|
|
|
|
|
class GPUObject:
|
|
def __init__(self):
|
|
self.module_list = []
|
|
|
|
def __enter__(self):
|
|
self.original_init = torch.nn.Module.__init__
|
|
|
|
def patched_init(module, *args, **kwargs):
|
|
self.module_list.append(module)
|
|
return self.original_init(module, *args, **kwargs)
|
|
|
|
torch.nn.Module.__init__ = patched_init
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
torch.nn.Module.__init__ = self.original_init
|
|
self.module_list = list(set(self.module_list))
|
|
self.to(device=torch.device('cpu'))
|
|
memory_management.soft_empty_cache()
|
|
return
|
|
|
|
def to(self, device):
|
|
for module in self.module_list:
|
|
module.to(device)
|
|
print(f'Forge Space: Moved {len(self.module_list)} Modules to {device}')
|
|
return self
|
|
|
|
def gpu(self):
|
|
self.to(device=gpu)
|
|
return self
|
|
|
|
|
|
def capture_gpu_object(capture=True):
|
|
if capture:
|
|
return GPUObject()
|
|
else:
|
|
return contextlib.nullcontext()
|
|
|
|
|
|
def GPU(gpu_objects=None, manual_load=False, **kwargs):
|
|
gpu_objects = gpu_objects or []
|
|
|
|
if not isinstance(gpu_objects, (list, tuple)):
|
|
gpu_objects = [gpu_objects]
|
|
|
|
def decorator(func):
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
print("Entering Forge Space GPU ...")
|
|
memory_management.unload_all_models()
|
|
if not manual_load:
|
|
for o in gpu_objects:
|
|
o.gpu()
|
|
result = func(*args, **kwargs)
|
|
print("Cleaning Forge Space GPU ...")
|
|
unload_module()
|
|
for o in gpu_objects:
|
|
o.to(device=torch.device('cpu'))
|
|
memory_management.soft_empty_cache()
|
|
return result
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
def convert_root_path():
|
|
frame = inspect.currentframe().f_back
|
|
caller_file = frame.f_code.co_filename
|
|
caller_file = os.path.abspath(caller_file)
|
|
result = os.path.join(os.path.dirname(caller_file), 'huggingface_space_mirror')
|
|
return result + '/'
|
|
|
|
|
|
def download_single_file(
|
|
url: str,
|
|
*,
|
|
model_dir: str,
|
|
progress: bool = True,
|
|
file_name: str | None = None,
|
|
hash_prefix: str | None = None,
|
|
) -> str:
|
|
os.makedirs(model_dir, exist_ok=True)
|
|
if not file_name:
|
|
from urllib.parse import urlparse
|
|
parts = urlparse(url)
|
|
file_name = os.path.basename(parts.path)
|
|
cached_file = os.path.abspath(os.path.join(model_dir, file_name))
|
|
if not os.path.exists(cached_file):
|
|
tmp_filename = cached_file + '.tmp'
|
|
print(f'Downloading: "{url}" to {cached_file} using temp file {tmp_filename}\n')
|
|
from torch.hub import download_url_to_file
|
|
download_url_to_file(url, tmp_filename, progress=progress, hash_prefix=hash_prefix)
|
|
os.replace(tmp_filename, cached_file)
|
|
return cached_file
|
|
|
|
|
|
def automatically_move_to_gpu_when_forward(m: torch.nn.Module, target_model: torch.nn.Module = None):
|
|
if target_model is None:
|
|
target_model = m
|
|
|
|
def patch_method(method_name):
|
|
if not hasattr(m, method_name):
|
|
return
|
|
|
|
if not hasattr(m, 'forge_space_hooked_names'):
|
|
m.forge_space_hooked_names = []
|
|
|
|
if method_name in m.forge_space_hooked_names:
|
|
print(f'Already hooked {type(m).__name__}.{method_name}')
|
|
return
|
|
|
|
print(f'Automatic hook: {type(m).__name__}.{method_name}')
|
|
|
|
original_method = getattr(m, method_name)
|
|
|
|
def patched_method(*args, **kwargs):
|
|
load_module(target_model)
|
|
return original_method(*args, **kwargs)
|
|
|
|
setattr(m, method_name, patched_method)
|
|
|
|
m.forge_space_hooked_names.append(method_name)
|
|
return
|
|
|
|
for method_name in ['forward', 'encode', 'decode']:
|
|
patch_method(method_name)
|
|
|
|
return
|
|
|
|
|
|
def automatically_move_pipeline_components(pipe):
|
|
for attr_name in dir(pipe):
|
|
attr_value = getattr(pipe, attr_name, None)
|
|
if isinstance(attr_value, torch.nn.Module):
|
|
automatically_move_to_gpu_when_forward(attr_value)
|
|
return
|
|
|
|
|
|
def change_attention_from_diffusers_to_forge(m):
|
|
m.set_attn_processor(AttentionProcessorForge())
|
|
return
|
|
|
|
|
|
# diffusers version fix
|
|
|
|
import diffusers.models
|
|
|
|
diffusers.models.embeddings.PositionNet = diffusers.models.embeddings.GLIGENTextBoundingboxProjection
|
|
|
|
import diffusers.models.transformers.dual_transformer_2d
|
|
dual_transformer_2d = types.ModuleType('diffusers.models.dual_transformer_2d')
|
|
dual_transformer_2d.__dict__.update(diffusers.models.transformers.dual_transformer_2d.__dict__)
|
|
sys.modules['diffusers.models.dual_transformer_2d'] = dual_transformer_2d
|
|
|
|
import diffusers.models.transformers.transformer_2d
|
|
transformer_2d = types.ModuleType('diffusers.models.transformer_2d')
|
|
transformer_2d.__dict__.update(diffusers.models.transformers.transformer_2d.__dict__)
|
|
sys.modules['diffusers.models.transformer_2d'] = transformer_2d
|