wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-03-07 02:07:32 +08:00
Code Issues Actions 15 Packages Projects Releases Wiki Activity

Fix Aimdo fallback on probe to not use zero-copy sft (#12634)

Browse Source 
* utils: dont use comfy sft loader in aimdo fallback

This was going to the raw command line switch and should respect main.py
probe of whether aimdo actually loaded successfully.

* ops: dont use deferred linear load in Aimdo fallback

Avoid changes of behaviour on --fast dynamic_vram when aimdo doesnt work.
This commit is contained in:
rattus 2026-02-25 13:49:48 -08:00 committed by GitHub
parent 3ebe1ac22e
commit 4f5b7dbf1f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/ops.py
View File

@ -19,7 +19,7 @@
import torch import torch
import logging import logging
import comfy.model_management import comfy.model_management
from comfy.cli_args import args, PerformanceFeature, enables_dynamic_vram from comfy.cli_args import args, PerformanceFeature
import comfy.float import comfy.float
import json import json
import comfy.memory_management import comfy.memory_management
@ -296,7 +296,7 @@ class disable_weight_init:
class Linear(torch.nn.Linear, CastWeightBiasOp): class Linear(torch.nn.Linear, CastWeightBiasOp):
def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
if not comfy.model_management.WINDOWS or not enables_dynamic_vram(): if not comfy.model_management.WINDOWS or not comfy.memory_management.aimdo_enabled:
super().__init__(in_features, out_features, bias, device, dtype) super().__init__(in_features, out_features, bias, device, dtype)
return return
@ -317,7 +317,7 @@ class disable_weight_init:
def _load_from_state_dict(self, state_dict, prefix, local_metadata, def _load_from_state_dict(self, state_dict, prefix, local_metadata,
strict, missing_keys, unexpected_keys, error_msgs): strict, missing_keys, unexpected_keys, error_msgs):
if not comfy.model_management.WINDOWS or not enables_dynamic_vram(): if not comfy.model_management.WINDOWS or not comfy.memory_management.aimdo_enabled:
return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs) missing_keys, unexpected_keys, error_msgs)
assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False) assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False)

4
comfy/utils.py
View File

@ -29,7 +29,7 @@ import itertools
from torch.nn.functional import interpolate from torch.nn.functional import interpolate
from tqdm.auto import trange from tqdm.auto import trange
from einops import rearrange from einops import rearrange
from comfy.cli_args import args, enables_dynamic_vram from comfy.cli_args import args
import json import json
import time import time
import mmap import mmap
@ -113,7 +113,7 @@ def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
metadata = None metadata = None
if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"): if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"):
try: try:
if enables_dynamic_vram(): if comfy.memory_management.aimdo_enabled:
sd, metadata = load_safetensors(ckpt) sd, metadata = load_safetensors(ckpt)
if not return_metadata: if not return_metadata:
metadata = None metadata = None