wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-12-19 19:13:02 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

sd: Free RAM on main model load

Browse Source
This commit is contained in:
Rattus 2025-11-18 09:17:10 +10:00
parent 62a2622591
commit 4a83a9bc0e
2 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/model_base.py
View File

@ -304,6 +304,7 @@ class BaseModel(torch.nn.Module):
to_load[k[len(unet_prefix):]] = sd.pop(k) to_load[k[len(unet_prefix):]] = sd.pop(k)
to_load = self.model_config.process_unet_state_dict(to_load) to_load = self.model_config.process_unet_state_dict(to_load)
comfy.model_management.free_ram(state_dict=to_load)
m, u = self.diffusion_model.load_state_dict(to_load, strict=False) m, u = self.diffusion_model.load_state_dict(to_load, strict=False)
if len(m) > 0: if len(m) > 0:
logging.warning("unet missing: {}".format(m)) logging.warning("unet missing: {}".format(m))

3
comfy/sd.py
View File

@ -260,6 +260,7 @@ class CLIP:
def load_sd(self, sd, full_model=False): def load_sd(self, sd, full_model=False):
if full_model: if full_model:
comfy.model_management.free_ram(state_dict=sd)
return self.cond_stage_model.load_state_dict(sd, strict=False) return self.cond_stage_model.load_state_dict(sd, strict=False)
else: else:
return self.cond_stage_model.load_sd(sd) return self.cond_stage_model.load_sd(sd)
@ -625,6 +626,7 @@ class VAE:
self.first_stage_model = AutoencoderKL(**(config['params'])) self.first_stage_model = AutoencoderKL(**(config['params']))
self.first_stage_model = self.first_stage_model.eval() self.first_stage_model = self.first_stage_model.eval()
comfy.model_management.free_ram(state_dict=sd)
m, u = self.first_stage_model.load_state_dict(sd, strict=False) m, u = self.first_stage_model.load_state_dict(sd, strict=False)
if len(m) > 0: if len(m) > 0:
logging.warning("Missing VAE keys {}".format(m)) logging.warning("Missing VAE keys {}".format(m))
@ -933,6 +935,7 @@ def load_style_model(ckpt_path):
model = comfy.ldm.flux.redux.ReduxImageEncoder() model = comfy.ldm.flux.redux.ReduxImageEncoder()
else: else:
raise Exception("invalid style model {}".format(ckpt_path)) raise Exception("invalid style model {}".format(ckpt_path))
comfy.model_management.free_ram(state_dict=model_data)
model.load_state_dict(model_data) model.load_state_dict(model_data)
return StyleModel(model) return StyleModel(model)