andreas/ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Greatly improve lowvram sampling speed by getting rid of accelerate.

Browse Source 
Let me know if this breaks anything.
This commit is contained in:
comfyanonymous
2023-12-22 14:24:04 -05:00
parent 261bcbb0d9
commit 36a7953142
5 changed files with 103 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/model_base.py
View File

@@ -162,11 +162,7 @@ class BaseModel(torch.nn.Module):
def state_dict_for_saving(self, clip_state_dict, vae_state_dict):
clip_state_dict = self.model_config.process_clip_state_dict_for_saving(clip_state_dict)
unet_sd = self.diffusion_model.state_dict()
unet_state_dict = {}
for k in unet_sd:
unet_state_dict[k] = comfy.model_management.resolve_lowvram_weight(unet_sd[k], self.diffusion_model, k)
unet_state_dict = self.diffusion_model.state_dict()
unet_state_dict = self.model_config.process_unet_state_dict_for_saving(unet_state_dict)
vae_state_dict = self.model_config.process_vae_state_dict_for_saving(vae_state_dict)
if self.get_dtype() == torch.float16: