lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-04-28 18:21:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

revise inference dtype

Browse Source
This commit is contained in:
layerdiffusion
2024-08-07 17:08:47 -07:00
parent 89ea50a19d
commit b61bf553ea
5 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
backend/patcher/unet.py
View File

@@ -8,12 +8,12 @@ from backend import memory_management
class UnetPatcher(ModelPatcher):
@classmethod
def from_model(cls, model, diffusers_scheduler, k_predictor=None):
def from_model(cls, model, diffusers_scheduler, config, k_predictor=None):
parameters = memory_management.module_size(model)
unet_dtype = memory_management.unet_dtype(model_params=parameters)
load_device = memory_management.get_torch_device()
initial_load_device = memory_management.unet_inital_load_device(parameters, unet_dtype)
manual_cast_dtype = memory_management.unet_manual_cast(unet_dtype, load_device)
manual_cast_dtype = memory_management.unet_manual_cast(unet_dtype, load_device, supported_dtypes=config.supported_inference_dtypes)
manual_cast_dtype = unet_dtype if manual_cast_dtype is None else manual_cast_dtype
model.to(device=initial_load_device, dtype=unet_dtype)
model = KModel(model=model, diffusers_scheduler=diffusers_scheduler, k_predictor=k_predictor, storage_dtype=unet_dtype, computation_dtype=manual_cast_dtype)