lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-03-05 04:59:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

move to new backend - part 2

Browse Source
This commit is contained in:
layerdiffusion
2024-08-03 15:10:37 -07:00
parent 8a01b2c5db
commit 4add428e25
9 changed files with 61 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
modules_forge/diffusers_patcher.py
View File

@@ -1,22 +1,21 @@
import torch
import ldm_patched.modules.ops as ops
from backend import operations, memory_management
from backend.patcher.base import ModelPatcher
from ldm_patched.modules.model_patcher import ModelPatcher
from ldm_patched.modules import model_management
from transformers import modeling_utils
class DiffusersModelPatcher:
def __init__(self, pipeline_class, dtype=torch.float16, *args, **kwargs):
load_device = model_management.get_torch_device()
load_device = memory_management.get_torch_device()
offload_device = torch.device("cpu")
if not model_management.should_use_fp16(device=load_device):
if not memory_management.should_use_fp16(device=load_device):
dtype = torch.float32
self.dtype = dtype
with ops.use_patched_ops(ops.manual_cast):
with operations.using_forge_operations():
with modeling_utils.no_init_weights():
self.pipeline = pipeline_class.from_pretrained(*args, **kwargs)
@@ -41,7 +40,7 @@ class DiffusersModelPatcher:
def prepare_memory_before_sampling(self, batchsize, latent_width, latent_height):
area = 2 * batchsize * latent_width * latent_height
inference_memory = (((area * 0.6) / 0.9) + 1024) * (1024 * 1024)
model_management.load_models_gpu(
memory_management.load_models_gpu(
models=[self.patcher],
memory_required=inference_memory
)