AI/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-10-24 02:49:04 +03:00
Code Activity

Fast disk detection to all model loaders. (#16425)

Browse Source
This commit is contained in:
comfyanonymous
2026-09-21 18:57:06 -07:00
committed by GitHub
parent fc584aaa22
commit e638023d54
15 changed files with 49 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
comfy/audio_encoders/audio_encoders.py
View File

@@ -4,6 +4,7 @@ from .sheetsage2 import SheetSage2
from .sheetsage2_abc import events_to_abc
import comfy.model_management
import comfy.ops
import comfy.storage
import comfy.utils
import logging
import comfy.audio
@@ -11,7 +12,7 @@ import torch
class AudioEncoderModel():
def __init__(self, config):
def __init__(self, config, fast_disk=False):
self.load_device = comfy.model_management.text_encoder_device()
offload_device = comfy.model_management.text_encoder_offload_device()
self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
@@ -33,7 +34,7 @@ class AudioEncoderModel():
elif model_type == "sheetsage2":
self.model = SheetSage2(**model_config)
self.model.eval()
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device, fast_disk=fast_disk)
comfy.model_management.archive_model_dtypes(self.model)
def load_sd(self, sd):
@@ -102,7 +103,8 @@ def load_audio_encoder_from_sd(sd, prefix=""):
else:
raise RuntimeError("ERROR: audio encoder not supported.")
audio_encoder = SheetSage2AudioEncoder(config) if config["model_type"] == "sheetsage2" else AudioEncoderModel(config)
fast_disk = comfy.storage.state_dict_fast_disk(sd)
audio_encoder = SheetSage2AudioEncoder(config, fast_disk=fast_disk) if config["model_type"] == "sheetsage2" else AudioEncoderModel(config, fast_disk=fast_disk)
m, u = audio_encoder.load_sd(sd)
if len(m) > 0:
logging.warning("missing audio encoder: {}".format(m))

7
comfy/bg_removal_model.py
View File

@@ -7,6 +7,7 @@ import logging
import comfy.ops
import comfy.model_patcher
import comfy.model_management
import comfy.storage
import comfy.clip_model
import comfy.background_removal.birefnet
@@ -15,7 +16,7 @@ BG_REMOVAL_MODELS = {
}
class BackgroundRemovalModel():
def __init__(self, json_config):
def __init__(self, json_config, fast_disk=False):
with open(json_config) as f:
config = json.load(f)
@@ -32,7 +33,7 @@ class BackgroundRemovalModel():
self.model = model_class(config, self.dtype, offload_device, comfy.ops.manual_cast)
self.model.eval()
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device, fast_disk=fast_disk)
def load_sd(self, sd):
return self.model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
@@ -64,7 +65,7 @@ def load_background_removal_model(sd):
else:
return None
bg_model = BackgroundRemovalModel(json_config)
bg_model = BackgroundRemovalModel(json_config, fast_disk=comfy.storage.state_dict_fast_disk(sd))
m, u = bg_model.load_sd(sd)
if len(m) > 0:
logging.warning("missing background removal: {}".format(m))

9
comfy/clip_vision.py
View File

@@ -6,6 +6,7 @@ import logging
import comfy.ops
import comfy.model_patcher
import comfy.model_management
import comfy.storage
import comfy.utils
import comfy.clip_model
import comfy.image_encoders.dino2
@@ -29,7 +30,7 @@ IMAGE_ENCODERS = {
}
class ClipVisionModel():
def __init__(self, json_config):
def __init__(self, json_config, fast_disk=False):
if isinstance(json_config, dict):
config = json_config
else:
@@ -53,7 +54,7 @@ class ClipVisionModel():
self.model = model_class(config, self.dtype, offload_device, comfy.ops.manual_cast)
self.model.eval()
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device, fast_disk=fast_disk)
self.naf = None
def load_sd(self, sd):
@@ -148,7 +149,7 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
else:
return None
clip = ClipVisionModel(json_config)
clip = ClipVisionModel(json_config, fast_disk=comfy.storage.state_dict_fast_disk(sd))
m, u = clip.load_sd(sd)
if len(m) > 0:
logging.warning("missing clip vision: {}".format(m))
@@ -164,7 +165,7 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
naf = NAF(operations=comfy.ops.manual_cast).eval()
naf.load_state_dict(naf_sd)
naf.to(comfy.model_management.text_encoder_dtype(clip.load_device))
clip.naf = comfy.model_patcher.CoreModelPatcher(naf, load_device=clip.load_device, offload_device=comfy.model_management.text_encoder_offload_device())
clip.naf = comfy.model_patcher.CoreModelPatcher(naf, load_device=clip.load_device, offload_device=comfy.model_management.text_encoder_offload_device(), fast_disk=comfy.storage.state_dict_fast_disk(naf_sd))
return clip
def load(ckpt_path):

22
comfy/controlnet.py
View File

@@ -27,6 +27,7 @@ import comfy.utils
import comfy.model_management
import comfy.model_detection
import comfy.model_patcher
import comfy.storage
import comfy.ops
import comfy.latent_formats
import comfy.model_base
@@ -233,12 +234,12 @@ class ControlBase:
class ControlNet(ControlBase):
def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False, preprocess_image=lambda a: a):
def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False, preprocess_image=lambda a: a, fast_disk=False):
super().__init__()
self.control_model = control_model
self.load_device = load_device
if control_model is not None:
self.control_model_wrapped = comfy.model_patcher.CoreModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device())
self.control_model_wrapped = comfy.model_patcher.CoreModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device(), fast_disk=fast_disk)
self.compression_ratio = compression_ratio
self.global_average_pooling = global_average_pooling
@@ -322,7 +323,7 @@ class ControlNet(ControlBase):
def deepclone_multigpu(self, load_device, autoregister=False):
c = self.copy()
c.control_model = copy.deepcopy(c.control_model)
c.control_model_wrapped = comfy.model_patcher.ModelPatcher(c.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device())
c.control_model_wrapped = comfy.model_patcher.ModelPatcher(c.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device(), fast_disk=self.control_model_wrapped.fast_disk)
if autoregister:
self.multigpu_clones[load_device] = c
return c
@@ -534,7 +535,7 @@ def load_controlnet_mmdit(sd, model_options={}):
latent_format = comfy.latent_formats.SD3()
latent_format.shift_factor = 0 #SD3 controlnet weirdness
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, fast_disk=comfy.storage.state_dict_fast_disk(new_sd))
return control
@@ -609,7 +610,7 @@ def load_controlnet_sd35(sd, model_options={}):
elif depth_cnet:
preprocess_image = lambda a: 1.0 - a
control = ControlNetSD35(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, preprocess_image=preprocess_image)
control = ControlNetSD35(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, preprocess_image=preprocess_image, fast_disk=comfy.storage.state_dict_fast_disk(sd))
return control
@@ -622,7 +623,7 @@ def load_controlnet_hunyuandit(controlnet_data, model_options={}):
latent_format = comfy.latent_formats.SDXL()
extra_conds = ['text_embedding_mask', 'encoder_hidden_states_t5', 'text_embedding_mask_t5', 'image_meta_size', 'style', 'cos_cis_img', 'sin_cis_img']
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, strength_type=StrengthType.CONSTANT)
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, strength_type=StrengthType.CONSTANT, fast_disk=comfy.storage.state_dict_fast_disk(controlnet_data))
return control
def load_controlnet_flux_xlabs_mistoline(sd, mistoline=False, model_options={}):
@@ -631,7 +632,7 @@ def load_controlnet_flux_xlabs_mistoline(sd, mistoline=False, model_options={}):
sd = model_config.process_unet_state_dict(sd)
control_model = controlnet_load_state_dict(control_model, sd)
extra_conds = ['y', 'guidance']
control = ControlNet(control_model, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
control = ControlNet(control_model, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, fast_disk=comfy.storage.state_dict_fast_disk(sd))
return control
def load_controlnet_flux_instantx(sd, model_options={}):
@@ -655,7 +656,7 @@ def load_controlnet_flux_instantx(sd, model_options={}):
latent_format = comfy.latent_formats.Flux()
extra_conds = ['y', 'guidance']
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, fast_disk=comfy.storage.state_dict_fast_disk(new_sd))
return control
def load_controlnet_qwen_instantx(sd, model_options={}):
@@ -671,7 +672,7 @@ def load_controlnet_qwen_instantx(sd, model_options={}):
control_model = controlnet_load_state_dict(control_model, sd)
latent_format = comfy.latent_formats.Wan21()
extra_conds = []
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, fast_disk=comfy.storage.state_dict_fast_disk(sd))
return control
@@ -718,6 +719,7 @@ def load_controlnet_qwen_fun(sd, model_options={}):
load_device=load_device,
manual_cast_dtype=manual_cast_dtype,
extra_conds=[],
fast_disk=comfy.storage.state_dict_fast_disk(sd),
)
return control
@@ -877,7 +879,7 @@ def load_controlnet_state_dict(state_dict, model=None, model_options={}):
logging.debug("unexpected controlnet keys: {}".format(unexpected))
global_average_pooling = model_options.get("global_average_pooling", False)
control = ControlNet(control_model, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
control = ControlNet(control_model, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype, fast_disk=comfy.storage.state_dict_fast_disk(controlnet_data))
return control
def load_controlnet(ckpt_path, model=None, model_options={}):

4
comfy/ldm/hunyuan_video/upsampler.py
View File

@@ -102,14 +102,14 @@ UPSAMPLERS = {
}
class HunyuanVideo15SRModel():
def __init__(self, model_type, config):
def __init__(self, model_type, config, fast_disk=False):
self.load_device = comfy.model_management.vae_device()
offload_device = comfy.model_management.vae_offload_device()
self.dtype = comfy.model_management.vae_dtype(self.load_device)
self.model_class = UPSAMPLERS.get(model_type)
self.model = self.model_class(**config).eval()
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device, fast_disk=fast_disk)
def load_sd(self, sd):
return self.model.load_state_dict(sd, strict=True, assign=self.patcher.is_dynamic())

3
comfy/ldm/moge/model.py
View File

@@ -16,6 +16,7 @@ import torch.nn.functional as F
import comfy.ops
import comfy.model_management
import comfy.model_patcher
import comfy.storage
from comfy.image_encoders.dino2 import Dinov2Model
@@ -342,7 +343,7 @@ class MoGeModel:
self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
self.model = build_from_state_dict(state_dict, dtype=self.dtype, device=offload_device, operations=comfy.ops.manual_cast).eval()
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device, fast_disk=comfy.storage.state_dict_fast_disk(state_dict))
if not hasattr(self.model, "encoder"):
self.version = "v1"
else:

2
comfy/model_patcher.py
View File

@@ -341,7 +341,7 @@ class ModelPatcher:
def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False, fast_disk=False):
self.size = size
self.model = model
self.fast_disk = not comfy.model_management.args.disable_fast_disk and bool(comfy.model_management.args.fast_disk or fast_disk)
self.fast_disk = fast_disk
if not hasattr(self.model, 'device'):
logging.debug("Model doesn't have a device attribute.")
self.model.device = offload_device

2
comfy/sd.py
View File

@@ -2097,7 +2097,7 @@ def load_gligen(ckpt_path):
model = gligen.load_gligen(data)
if model_management.should_use_fp16():
model = model.half()
return comfy.model_patcher.CoreModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
return comfy.model_patcher.CoreModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device(), fast_disk=comfy.storage.state_dict_fast_disk(data))
def model_detection_error_hint(path, state_dict):
filename = os.path.basename(path)

2
comfy_extras/nodes_frame_interpolation.py
View File

@@ -3,6 +3,7 @@ from tqdm import tqdm
from typing_extensions import override
import comfy.model_patcher
import comfy.storage
import comfy.utils
import folder_paths
from comfy import model_management
@@ -41,6 +42,7 @@ class FrameInterpolationModelLoader(io.ComfyNode):
model,
load_device=model_management.get_torch_device(),
offload_device=model_management.unet_offload_device(),
fast_disk=comfy.storage.state_dict_fast_disk(sd),
)
return io.NodeOutput(patcher)

8
comfy_extras/nodes_hunyuan.py
View File

@@ -3,6 +3,7 @@ import node_helpers
import torch
import comfy.model_management
import comfy.model_patcher
import comfy.storage
import comfy.ops
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
@@ -195,6 +196,7 @@ class LatentUpscaleModelLoader(io.ComfyNode):
def execute(cls, model_name) -> io.NodeOutput:
model_path = folder_paths.get_full_path_or_raise("latent_upscale_models", model_name)
sd, metadata = comfy.utils.load_torch_file(model_path, safe_load=True, return_metadata=True)
fast_disk = comfy.storage.state_dict_fast_disk(sd)
if "blocks.0.block.0.conv.weight" in sd:
config = {
@@ -205,7 +207,7 @@ class LatentUpscaleModelLoader(io.ComfyNode):
"global_residual": False,
}
model_type = "720p"
model = HunyuanVideo15SRModel(model_type, config)
model = HunyuanVideo15SRModel(model_type, config, fast_disk=fast_disk)
model.load_sd(sd)
elif "up.0.block.0.conv1.conv.weight" in sd:
sd = {key.replace("nin_shortcut", "nin_shortcut.conv", 1): value for key, value in sd.items()}
@@ -215,13 +217,13 @@ class LatentUpscaleModelLoader(io.ComfyNode):
"block_out_channels": tuple(sd[f"up.{i}.block.0.conv1.conv.weight"].shape[0] for i in range(len([k for k in sd.keys() if k.startswith("up.") and k.endswith(".block.0.conv1.conv.weight")]))),
}
model_type = "1080p"
model = HunyuanVideo15SRModel(model_type, config)
model = HunyuanVideo15SRModel(model_type, config, fast_disk=fast_disk)
model.load_sd(sd)
elif "post_upsample_res_blocks.0.conv2.bias" in sd:
config = json.loads(metadata["config"])
model = LatentUpsampler.from_config(config, operations=comfy.ops.disable_weight_init).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32]))
comfy.model_management.archive_model_dtypes(model)
model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device())
model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device(), fast_disk=fast_disk)
model.load_state_dict(sd, assign=model_patcher.is_dynamic())
model = model_patcher

2
comfy_extras/nodes_mediapipe.py
View File

@@ -19,6 +19,7 @@ from typing_extensions import override
import comfy.model_management
import comfy.model_patcher
import comfy.storage
import comfy.utils
import folder_paths
from comfy_api.latest import ComfyExtension, io
@@ -72,6 +73,7 @@ class FaceLandmarkerModel:
self.patchers[variant] = comfy.model_patcher.CoreModelPatcher(
fl, load_device=self.load_device, offload_device=offload_device,
size=comfy.model_management.module_size(fl),
fast_disk=comfy.storage.state_dict_fast_disk(sub),
)
def detect_batch(self, images, num_faces: int, score_thresh: float, variant: str):

3
comfy_extras/nodes_model_patch.py
View File

@@ -6,6 +6,7 @@ import folder_paths
import comfy.utils
import comfy.ops
import comfy.model_management
import comfy.storage
import comfy.ldm.common_dit
import comfy.latent_formats
import comfy.ldm.lumina.controlnet
@@ -376,7 +377,7 @@ class ModelPatchLoader:
if denoise_encoder_sd:
model.denoise_encoder_sd = denoise_encoder_sd
model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device())
model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device(), fast_disk=comfy.storage.state_dict_fast_disk(sd))
model.load_state_dict(sd, assign=model_patcher.is_dynamic())
return (model_patcher,)

2
comfy_extras/nodes_sam3d_body.py
View File

@@ -11,6 +11,7 @@ from scipy.signal import savgol_coeffs
import comfy.model_management
import comfy.model_patcher
import comfy.storage
import comfy.ops
import comfy.utils
from comfy_api.latest import io, ComfyExtension, Types
@@ -96,6 +97,7 @@ class SAM3DBody_Loader(io.ComfyNode):
load_device=load_device,
offload_device=comfy.model_management.unet_offload_device(),
size=comfy.model_management.module_size(model),
fast_disk=comfy.storage.state_dict_fast_disk(sd),
)
return io.NodeOutput(patcher)

3
comfy_extras/nodes_upscale_model.py
View File

@@ -8,6 +8,7 @@ from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
import comfy.model_management
import comfy.model_patcher
import comfy.storage
try:
from spandrel_extra_arches import EXTRA_REGISTRY
@@ -43,7 +44,7 @@ class UpscaleModelLoader(io.ComfyNode):
if not isinstance(out, ImageModelDescriptor):
raise Exception("Upscale model must be a single-image model.")
out.patcher = comfy.model_patcher.CoreModelPatcher(out.model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
out.patcher = comfy.model_patcher.CoreModelPatcher(out.model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device(), fast_disk=comfy.storage.state_dict_fast_disk(sd))
return io.NodeOutput(out)
load_model = execute # TODO: remove

2
comfy_extras/nodes_void.py
View File

@@ -5,6 +5,7 @@ import torch
import comfy
import comfy.model_management
import comfy.model_patcher
import comfy.storage
import comfy.samplers
import comfy.utils
import folder_paths
@@ -101,6 +102,7 @@ class OpticalFlowLoader(io.ComfyNode):
model,
load_device=comfy.model_management.get_torch_device(),
offload_device=comfy.model_management.unet_offload_device(),
fast_disk=comfy.storage.state_dict_fast_disk(sd),
)
return io.NodeOutput(patcher)