Mirror/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2024-06-07 21:20:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

fp8 for TE

Browse Source
This commit is contained in:
Kohaku-Blueleaf 2023-10-25 11:36:43 +08:00
parent 9c1eba2af3
commit 1df6c8bfec
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
modules/sd_models.py
View File

@ -407,6 +407,13 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
module.to(torch.float8_e4m3fn) module.to(torch.float8_e4m3fn)
timer.record("apply fp8 unet for cpu") timer.record("apply fp8 unet for cpu")
else: else:
if model.is_sdxl:
cond_stage = model.conditioner
else:
cond_stage = model.cond_stage_model
for module in cond_stage.modules():
if isinstance(module, torch.nn.Linear):
module.to(torch.float8_e4m3fn)
model.model.diffusion_model = model.model.diffusion_model.to(torch.float8_e4m3fn) model.model.diffusion_model = model.model.diffusion_model.to(torch.float8_e4m3fn)
timer.record("apply fp8 unet") timer.record("apply fp8 unet")