mirror of
https://github.com/lllyasviel/Fooocus.git
synced 2026-08-16 13:13:16 +02:00
Use weights_only for loading (#3427)
Co-authored-by: Manuel Schmid <9307310+mashb1t@users.noreply.github.com>
This commit is contained in:
co-authored by
Manuel Schmid
parent
1a53e0676a
commit
da3d4d006f
@@ -377,15 +377,15 @@ class VQAutoEncoder(nn.Module):
|
||||
)
|
||||
|
||||
if model_path is not None:
|
||||
chkpt = torch.load(model_path, map_location="cpu")
|
||||
chkpt = torch.load(model_path, map_location="cpu", weights_only=True)
|
||||
if "params_ema" in chkpt:
|
||||
self.load_state_dict(
|
||||
torch.load(model_path, map_location="cpu")["params_ema"]
|
||||
torch.load(model_path, map_location="cpu", weights_only=True)["params_ema"]
|
||||
)
|
||||
logger.info(f"vqgan is loaded from: {model_path} [params_ema]")
|
||||
elif "params" in chkpt:
|
||||
self.load_state_dict(
|
||||
torch.load(model_path, map_location="cpu")["params"]
|
||||
torch.load(model_path, map_location="cpu", weights_only=True)["params"]
|
||||
)
|
||||
logger.info(f"vqgan is loaded from: {model_path} [params]")
|
||||
else:
|
||||
|
||||
Reference in New Issue
Block a user