diff --git a/api_server/routes/internal/internal_routes.py b/api_server/routes/internal/internal_routes.py index 8c46215f0..63704f13a 100644 --- a/api_server/routes/internal/internal_routes.py +++ b/api_server/routes/internal/internal_routes.py @@ -1,6 +1,6 @@ from aiohttp import web from typing import Optional -from folder_paths import models_dir, user_directory, output_directory +from folder_paths import models_dir, user_directory, output_directory, folder_names_and_paths from api_server.services.file_service import FileService import app.logger @@ -36,6 +36,13 @@ class InternalRoutes: async def get_logs(request): return web.json_response(app.logger.get_logs()) + @self.routes.get('/folder_paths') + async def get_folder_paths(request): + response = {} + for key in folder_names_and_paths: + response[key] = folder_names_and_paths[key][0] + return web.json_response(response) + def get_app(self): if self._app is None: self._app = web.Application() diff --git a/app/user_manager.py b/app/user_manager.py index 05e6805f7..208178444 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -118,6 +118,29 @@ class UserManager(): @routes.get("/userdata") async def listuserdata(request): + """ + List user data files in a specified directory. + + This endpoint allows listing files in a user's data directory, with options for recursion, + full file information, and path splitting. + + Query Parameters: + - dir (required): The directory to list files from. + - recurse (optional): If "true", recursively list files in subdirectories. + - full_info (optional): If "true", return detailed file information (path, size, modified time). + - split (optional): If "true", split file paths into components (only applies when full_info is false). + + Returns: + - 400: If 'dir' parameter is missing. + - 403: If the requested path is not allowed. + - 404: If the requested directory does not exist. + - 200: JSON response with the list of files or file information. + + The response format depends on the query parameters: + - Default: List of relative file paths. + - full_info=true: List of dictionaries with file details. + - split=true (and full_info=false): List of lists, each containing path components. + """ directory = request.rel_url.query.get('dir', '') if not directory: return web.Response(status=400, text="Directory not provided") diff --git a/comfy/controlnet.py b/comfy/controlnet.py index 1ea00eccf..9dfd69977 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -88,6 +88,8 @@ class ControlBase: self.strength = strength self.timestep_percent_range = timestep_percent_range if self.latent_format is not None: + if vae is None: + logging.warning("WARNING: no VAE provided to the controlnet apply node when this controlnet requires one.") self.vae = vae self.extra_concat_orig = extra_concat.copy() if self.concat_mask and len(self.extra_concat_orig) == 0: @@ -222,6 +224,9 @@ class ControlNet(ControlBase): compression_ratio = self.compression_ratio if self.vae is not None: compression_ratio *= self.vae.downscale_ratio + else: + if self.latent_format is not None: + raise ValueError("This Controlnet needs a VAE but none was provided, please use a different ControlNetApply node with a VAE input.") self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center") if self.vae is not None: loaded_models = comfy.model_management.loaded_models(only_currently_used=True) @@ -335,7 +340,7 @@ class ControlLoraOps: class ControlLora(ControlNet): - def __init__(self, control_weights, global_average_pooling=False, device=None): + def __init__(self, control_weights, global_average_pooling=False, device=None, model_options={}): #TODO? model_options ControlBase.__init__(self, device) self.control_weights = control_weights self.global_average_pooling = global_average_pooling @@ -392,19 +397,25 @@ class ControlLora(ControlNet): def inference_memory_requirements(self, dtype): return comfy.utils.calculate_parameters(self.control_weights) * comfy.model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype) -def controlnet_config(sd): +def controlnet_config(sd, model_options={}): model_config = comfy.model_detection.model_config_from_unet(sd, "", True) - supported_inference_dtypes = model_config.supported_inference_dtypes + unet_dtype = model_options.get("dtype", None) + if unet_dtype is None: + weight_dtype = comfy.utils.weight_dtype(sd) + + supported_inference_dtypes = list(model_config.supported_inference_dtypes) + if weight_dtype is not None: + supported_inference_dtypes.append(weight_dtype) + + unet_dtype = comfy.model_management.unet_dtype(model_params=-1, supported_dtypes=supported_inference_dtypes) - controlnet_config = model_config.unet_config - unet_dtype = comfy.model_management.unet_dtype(supported_dtypes=supported_inference_dtypes) load_device = comfy.model_management.get_torch_device() manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device) - if manual_cast_dtype is not None: - operations = comfy.ops.manual_cast - else: - operations = comfy.ops.disable_weight_init + + operations = model_options.get("custom_operations", None) + if operations is None: + operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True) offload_device = comfy.model_management.unet_offload_device() return model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device @@ -419,9 +430,9 @@ def controlnet_load_state_dict(control_model, sd): logging.debug("unexpected controlnet keys: {}".format(unexpected)) return control_model -def load_controlnet_mmdit(sd): +def load_controlnet_mmdit(sd, model_options={}): new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "") - model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd) + model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd, model_options=model_options) num_blocks = comfy.model_detection.count_blocks(new_sd, 'joint_blocks.{}.') for k in sd: new_sd[k] = sd[k] @@ -440,8 +451,8 @@ def load_controlnet_mmdit(sd): return control -def load_controlnet_hunyuandit(controlnet_data): - model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(controlnet_data) +def load_controlnet_hunyuandit(controlnet_data, model_options={}): + model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(controlnet_data, model_options=model_options) control_model = comfy.ldm.hydit.controlnet.HunYuanControlNet(operations=operations, device=offload_device, dtype=unet_dtype) control_model = controlnet_load_state_dict(control_model, controlnet_data) @@ -451,17 +462,17 @@ def load_controlnet_hunyuandit(controlnet_data): 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) return control -def load_controlnet_flux_xlabs_mistoline(sd, mistoline=False): - model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(sd) +def load_controlnet_flux_xlabs_mistoline(sd, mistoline=False, model_options={}): + model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(sd, model_options=model_options) control_model = comfy.ldm.flux.controlnet.ControlNetFlux(mistoline=mistoline, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) 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) return control -def load_controlnet_flux_instantx(sd): +def load_controlnet_flux_instantx(sd, model_options={}): new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "") - model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd) + model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd, model_options=model_options) for k in sd: new_sd[k] = sd[k] @@ -487,13 +498,13 @@ def convert_mistoline(sd): return comfy.utils.state_dict_prefix_replace(sd, {"single_controlnet_blocks.": "controlnet_single_blocks."}) -def load_controlnet(ckpt_path, model=None): - controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True) +def load_controlnet_state_dict(state_dict, model=None, model_options={}): + controlnet_data = state_dict if 'after_proj_list.18.bias' in controlnet_data.keys(): #Hunyuan DiT - return load_controlnet_hunyuandit(controlnet_data) + return load_controlnet_hunyuandit(controlnet_data, model_options=model_options) if "lora_controlnet" in controlnet_data: - return ControlLora(controlnet_data) + return ControlLora(controlnet_data, model_options=model_options) controlnet_config = None supported_inference_dtypes = None @@ -550,13 +561,13 @@ def load_controlnet(ckpt_path, model=None): controlnet_data = new_sd elif "controlnet_blocks.0.weight" in controlnet_data: if "double_blocks.0.img_attn.norm.key_norm.scale" in controlnet_data: - return load_controlnet_flux_xlabs_mistoline(controlnet_data) + return load_controlnet_flux_xlabs_mistoline(controlnet_data, model_options=model_options) elif "pos_embed_input.proj.weight" in controlnet_data: - return load_controlnet_mmdit(controlnet_data) #SD3 diffusers controlnet + return load_controlnet_mmdit(controlnet_data, model_options=model_options) #SD3 diffusers controlnet elif "controlnet_x_embedder.weight" in controlnet_data: - return load_controlnet_flux_instantx(controlnet_data) + return load_controlnet_flux_instantx(controlnet_data, model_options=model_options) elif "controlnet_blocks.0.linear.weight" in controlnet_data: #mistoline flux - return load_controlnet_flux_xlabs_mistoline(convert_mistoline(controlnet_data), mistoline=True) + return load_controlnet_flux_xlabs_mistoline(convert_mistoline(controlnet_data), mistoline=True, model_options=model_options) pth_key = 'control_model.zero_convs.0.0.weight' pth = False @@ -568,25 +579,36 @@ def load_controlnet(ckpt_path, model=None): elif key in controlnet_data: prefix = "" else: - net = load_t2i_adapter(controlnet_data) + net = load_t2i_adapter(controlnet_data, model_options=model_options) if net is None: - logging.error("error checkpoint does not contain controlnet or t2i adapter data {}".format(ckpt_path)) + logging.error("error could not detect control model type.") return net if controlnet_config is None: model_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, True) - supported_inference_dtypes = model_config.supported_inference_dtypes + supported_inference_dtypes = list(model_config.supported_inference_dtypes) controlnet_config = model_config.unet_config + unet_dtype = model_options.get("dtype", None) + if unet_dtype is None: + weight_dtype = comfy.utils.weight_dtype(controlnet_data) + + if supported_inference_dtypes is None: + supported_inference_dtypes = [comfy.model_management.unet_dtype()] + + if weight_dtype is not None: + supported_inference_dtypes.append(weight_dtype) + + unet_dtype = comfy.model_management.unet_dtype(model_params=-1, supported_dtypes=supported_inference_dtypes) + load_device = comfy.model_management.get_torch_device() - if supported_inference_dtypes is None: - unet_dtype = comfy.model_management.unet_dtype() - else: - unet_dtype = comfy.model_management.unet_dtype(supported_dtypes=supported_inference_dtypes) manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device) - if manual_cast_dtype is not None: - controlnet_config["operations"] = comfy.ops.manual_cast + operations = model_options.get("custom_operations", None) + if operations is None: + operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype) + + controlnet_config["operations"] = operations controlnet_config["dtype"] = unet_dtype controlnet_config["device"] = comfy.model_management.unet_offload_device() controlnet_config.pop("out_channels") @@ -622,14 +644,21 @@ def load_controlnet(ckpt_path, model=None): if len(unexpected) > 0: logging.debug("unexpected controlnet keys: {}".format(unexpected)) - global_average_pooling = False - filename = os.path.splitext(ckpt_path)[0] - if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling - global_average_pooling = True - + 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) return control +def load_controlnet(ckpt_path, model=None, model_options={}): + if "global_average_pooling" not in model_options: + filename = os.path.splitext(ckpt_path)[0] + if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling + model_options["global_average_pooling"] = True + + cnet = load_controlnet_state_dict(comfy.utils.load_torch_file(ckpt_path, safe_load=True), model=model, model_options=model_options) + if cnet is None: + logging.error("error checkpoint does not contain controlnet or t2i adapter data {}".format(ckpt_path)) + return cnet + class T2IAdapter(ControlBase): def __init__(self, t2i_model, channels_in, compression_ratio, upscale_algorithm, device=None): super().__init__(device) @@ -685,7 +714,7 @@ class T2IAdapter(ControlBase): self.copy_to(c) return c -def load_t2i_adapter(t2i_data): +def load_t2i_adapter(t2i_data, model_options={}): #TODO: model_options compression_ratio = 8 upscale_algorithm = 'nearest-exact' diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index ae56aee8b..e9e4edcc6 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -44,6 +44,17 @@ def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device='cpu'): return append_zero(sigmas) +def get_sigmas_laplace(n, sigma_min, sigma_max, mu=0., beta=0.5, device='cpu'): + """Constructs the noise schedule proposed by Tiankai et al. (2024). """ + epsilon = 1e-5 # avoid log(0) + x = torch.linspace(0, 1, n, device=device) + clamp = lambda x: torch.clamp(x, min=sigma_min, max=sigma_max) + lmb = mu - beta * torch.sign(0.5-x) * torch.log(1 - 2 * torch.abs(0.5-x) + epsilon) + sigmas = clamp(torch.exp(lmb)) + return sigmas + + + def to_d(x, sigma, denoised): """Converts a denoiser output to a Karras ODE derivative.""" return (x - denoised) / utils.append_dims(sigma, x.ndim) diff --git a/comfy/lora.py b/comfy/lora.py index 4c01a37b2..7970d6f70 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -415,7 +415,7 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, ori weight *= strength_model if isinstance(v, list): - v = (calculate_weight(v[1:], v[0].clone(), key, intermediate_dtype=intermediate_dtype), ) + v = (calculate_weight(v[1:], comfy.model_management.cast_to_device(v[0], weight.device, intermediate_dtype, copy=True), key, intermediate_dtype=intermediate_dtype), ) if len(v) == 1: patch_type = "diff" diff --git a/comfy/model_management.py b/comfy/model_management.py index 77c05510e..a97d489d5 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -326,7 +326,7 @@ class LoadedModel: self.model_unload() raise e - if is_intel_xpu() and not args.disable_ipex_optimize and self.real_model is not None: + if is_intel_xpu() and not args.disable_ipex_optimize and 'ipex' in globals() and self.real_model is not None: with torch.no_grad(): self.real_model = ipex.optimize(self.real_model.eval(), inplace=True, graph_mode=True, concat_linear=True) @@ -626,6 +626,8 @@ def maximum_vram_for_weights(device=None): return (get_total_memory(device) * 0.88 - minimum_inference_memory()) def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]): + if model_params < 0: + model_params = 1000000000000000000000 if args.bf16_unet: return torch.bfloat16 if args.fp16_unet: diff --git a/comfy/ops.py b/comfy/ops.py index 43ed55adb..1b386dba7 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -300,10 +300,10 @@ class fp8_ops(manual_cast): return torch.nn.functional.linear(input, weight, bias) -def pick_operations(weight_dtype, compute_dtype, load_device=None): +def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False): if compute_dtype is None or weight_dtype == compute_dtype: return disable_weight_init - if args.fast: + if args.fast and not disable_fast_fp8: if comfy.model_management.supports_fp8_compute(load_device): return fp8_ops return manual_cast diff --git a/comfy/sd.py b/comfy/sd.py index 07310b9d4..99859d241 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -70,14 +70,14 @@ class CLIP: clip = target.clip tokenizer = target.tokenizer - load_device = model_management.text_encoder_device() - offload_device = model_management.text_encoder_offload_device() + load_device = model_options.get("load_device", model_management.text_encoder_device()) + offload_device = model_options.get("offload_device", model_management.text_encoder_offload_device()) dtype = model_options.get("dtype", None) if dtype is None: dtype = model_management.text_encoder_dtype(load_device) params['dtype'] = dtype - params['device'] = model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype)) + params['device'] = model_options.get("initial_device", model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype))) params['model_options'] = model_options self.cond_stage_model = clip(**(params)) @@ -645,7 +645,7 @@ def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffuse manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes) model_config.set_inference_dtype(unet_dtype, manual_cast_dtype) - model_config.custom_operations = model_options.get("custom_operations", None) + model_config.custom_operations = model_options.get("custom_operations", model_config.custom_operations) model = model_config.get_model(new_sd, "") model = model.to(offload_device) model.load_model_weights(new_sd, "") diff --git a/comfy/utils.py b/comfy/utils.py index 1bc35df7a..78f8314fc 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -713,7 +713,9 @@ def common_upscale(samples, width, height, upscale_method, crop): return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method) def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap): - return math.ceil((height / (tile_y - overlap))) * math.ceil((width / (tile_x - overlap))) + rows = 1 if height <= tile_y else math.ceil((height - overlap) / (tile_y - overlap)) + cols = 1 if width <= tile_x else math.ceil((width - overlap) / (tile_x - overlap)) + return rows * cols @torch.inference_mode() def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None): @@ -722,10 +724,20 @@ def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_ for b in range(samples.shape[0]): s = samples[b:b+1] + + # handle entire input fitting in a single tile + if all(s.shape[d+2] <= tile[d] for d in range(dims)): + output[b:b+1] = function(s).to(output_device) + if pbar is not None: + pbar.update(1) + continue + out = torch.zeros([s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), device=output_device) out_div = torch.zeros([s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), device=output_device) - for it in itertools.product(*map(lambda a: range(0, a[0], a[1] - overlap), zip(s.shape[2:], tile))): + positions = [range(0, s.shape[d+2], tile[d] - overlap) if s.shape[d+2] > tile[d] else [0] for d in range(dims)] + + for it in itertools.product(*positions): s_in = s upscaled = [] @@ -734,15 +746,16 @@ def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_ l = min(tile[d], s.shape[d + 2] - pos) s_in = s_in.narrow(d + 2, pos, l) upscaled.append(round(pos * upscale_amount)) + ps = function(s_in).to(output_device) mask = torch.ones_like(ps) feather = round(overlap * upscale_amount) + for t in range(feather): for d in range(2, dims + 2): - m = mask.narrow(d, t, 1) - m *= ((1.0/feather) * (t + 1)) - m = mask.narrow(d, mask.shape[d] -1 -t, 1) - m *= ((1.0/feather) * (t + 1)) + a = (t + 1) / feather + mask.narrow(d, t, 1).mul_(a) + mask.narrow(d, mask.shape[d] - 1 - t, 1).mul_(a) o = out o_d = out_div @@ -750,8 +763,8 @@ def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_ o = o.narrow(d + 2, upscaled[d], mask.shape[d + 2]) o_d = o_d.narrow(d + 2, upscaled[d], mask.shape[d + 2]) - o += ps * mask - o_d += mask + o.add_(ps * mask) + o_d.add_(mask) if pbar is not None: pbar.update(1) diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py index 219975e27..c7ff9a4d8 100644 --- a/comfy_extras/nodes_custom_sampler.py +++ b/comfy_extras/nodes_custom_sampler.py @@ -90,6 +90,27 @@ class PolyexponentialScheduler: sigmas = k_diffusion_sampling.get_sigmas_polyexponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho) return (sigmas, ) +class LaplaceScheduler: + @classmethod + def INPUT_TYPES(s): + return {"required": + {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}), + "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}), + "mu": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step":0.1, "round": False}), + "beta": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 10.0, "step":0.1, "round": False}), + } + } + RETURN_TYPES = ("SIGMAS",) + CATEGORY = "sampling/custom_sampling/schedulers" + + FUNCTION = "get_sigmas" + + def get_sigmas(self, steps, sigma_max, sigma_min, mu, beta): + sigmas = k_diffusion_sampling.get_sigmas_laplace(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, mu=mu, beta=beta) + return (sigmas, ) + + class SDTurboScheduler: @classmethod def INPUT_TYPES(s): @@ -673,6 +694,7 @@ NODE_CLASS_MAPPINGS = { "KarrasScheduler": KarrasScheduler, "ExponentialScheduler": ExponentialScheduler, "PolyexponentialScheduler": PolyexponentialScheduler, + "LaplaceScheduler": LaplaceScheduler, "VPScheduler": VPScheduler, "BetaSamplingScheduler": BetaSamplingScheduler, "SDTurboScheduler": SDTurboScheduler, diff --git a/comfy_extras/nodes_hypernetwork.py b/comfy_extras/nodes_hypernetwork.py index cafafa6ab..665632292 100644 --- a/comfy_extras/nodes_hypernetwork.py +++ b/comfy_extras/nodes_hypernetwork.py @@ -107,7 +107,7 @@ class HypernetworkLoader: CATEGORY = "loaders" def load_hypernetwork(self, model, hypernetwork_name, strength): - hypernetwork_path = folder_paths.get_full_path("hypernetworks", hypernetwork_name) + hypernetwork_path = folder_paths.get_full_path_or_raise("hypernetworks", hypernetwork_name) model_hypernetwork = model.clone() patch = load_hypernetwork_patch(hypernetwork_path, strength) if patch is not None: diff --git a/comfy_extras/nodes_photomaker.py b/comfy_extras/nodes_photomaker.py index 29d127d74..95d24dd22 100644 --- a/comfy_extras/nodes_photomaker.py +++ b/comfy_extras/nodes_photomaker.py @@ -126,7 +126,7 @@ class PhotoMakerLoader: CATEGORY = "_for_testing/photomaker" def load_photomaker_model(self, photomaker_model_name): - photomaker_model_path = folder_paths.get_full_path("photomaker", photomaker_model_name) + photomaker_model_path = folder_paths.get_full_path_or_raise("photomaker", photomaker_model_name) photomaker_model = PhotoMakerIDEncoder() data = comfy.utils.load_torch_file(photomaker_model_path, safe_load=True) if "id_encoder" in data: diff --git a/comfy_extras/nodes_sd3.py b/comfy_extras/nodes_sd3.py index 046096cba..60a85bbdf 100644 --- a/comfy_extras/nodes_sd3.py +++ b/comfy_extras/nodes_sd3.py @@ -15,9 +15,9 @@ class TripleCLIPLoader: CATEGORY = "advanced/loaders" def load_clip(self, clip_name1, clip_name2, clip_name3): - clip_path1 = folder_paths.get_full_path("clip", clip_name1) - clip_path2 = folder_paths.get_full_path("clip", clip_name2) - clip_path3 = folder_paths.get_full_path("clip", clip_name3) + clip_path1 = folder_paths.get_full_path_or_raise("clip", clip_name1) + clip_path2 = folder_paths.get_full_path_or_raise("clip", clip_name2) + clip_path3 = folder_paths.get_full_path_or_raise("clip", clip_name3) clip = comfy.sd.load_clip(ckpt_paths=[clip_path1, clip_path2, clip_path3], embedding_directory=folder_paths.get_folder_paths("embeddings")) return (clip,) @@ -36,7 +36,7 @@ class EmptySD3LatentImage: CATEGORY = "latent/sd3" def generate(self, width, height, batch_size=1): - latent = torch.ones([batch_size, 16, height // 8, width // 8], device=self.device) * 0.0609 + latent = torch.zeros([batch_size, 16, height // 8, width // 8], device=self.device) return ({"samples":latent}, ) class CLIPTextEncodeSD3: @@ -103,5 +103,5 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { # Sampling - "ControlNetApplySD3": "ControlNetApply SD3 and HunyuanDiT", + "ControlNetApplySD3": "Apply Controlnet", } diff --git a/comfy_extras/nodes_upscale_model.py b/comfy_extras/nodes_upscale_model.py index bca79ef2e..6ba3e404f 100644 --- a/comfy_extras/nodes_upscale_model.py +++ b/comfy_extras/nodes_upscale_model.py @@ -25,7 +25,7 @@ class UpscaleModelLoader: CATEGORY = "loaders" def load_model(self, model_name): - model_path = folder_paths.get_full_path("upscale_models", model_name) + model_path = folder_paths.get_full_path_or_raise("upscale_models", model_name) sd = comfy.utils.load_torch_file(model_path, safe_load=True) if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd: sd = comfy.utils.state_dict_prefix_replace(sd, {"module.":""}) diff --git a/comfy_extras/nodes_video_model.py b/comfy_extras/nodes_video_model.py index 1a0189ed4..7f2146d1e 100644 --- a/comfy_extras/nodes_video_model.py +++ b/comfy_extras/nodes_video_model.py @@ -17,7 +17,7 @@ class ImageOnlyCheckpointLoader: CATEGORY = "loaders/video_models" def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True): - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=False, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) return (out[0], out[3], out[2]) diff --git a/extra_model_paths.yaml.example b/extra_model_paths.yaml.example index d4fd55c87..b55913a5a 100644 --- a/extra_model_paths.yaml.example +++ b/extra_model_paths.yaml.example @@ -25,6 +25,8 @@ a111: #comfyui: # base_path: path/to/comfyui/ +# # You can use is_default to mark that these folders should be listed first, and used as the default dirs for eg downloads +# #is_default: true # checkpoints: models/checkpoints/ # clip: models/clip/ # clip_vision: models/clip_vision/ diff --git a/folder_paths.py b/folder_paths.py index 263704a4c..1f03c08d8 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -46,6 +46,36 @@ user_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "user filename_list_cache: dict[str, tuple[list[str], dict[str, float], float]] = {} +class CacheHelper: + """ + Helper class for managing file list cache data. + """ + def __init__(self): + self.cache: dict[str, tuple[list[str], dict[str, float], float]] = {} + self.active = False + + def get(self, key: str, default=None) -> tuple[list[str], dict[str, float], float]: + if not self.active: + return default + return self.cache.get(key, default) + + def set(self, key: str, value: tuple[list[str], dict[str, float], float]) -> None: + if self.active: + self.cache[key] = value + + def clear(self): + self.cache.clear() + + def __enter__(self): + self.active = True + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.active = False + self.clear() + +cache_helper = CacheHelper() + extension_mimetypes_cache = { "webp" : "image", } @@ -165,11 +195,14 @@ def exists_annotated_filepath(name) -> bool: return os.path.exists(filepath) -def add_model_folder_path(folder_name: str, full_folder_path: str) -> None: +def add_model_folder_path(folder_name: str, full_folder_path: str, is_default: bool = False) -> None: global folder_names_and_paths folder_name = map_legacy(folder_name) if folder_name in folder_names_and_paths: - folder_names_and_paths[folder_name][0].append(full_folder_path) + if is_default: + folder_names_and_paths[folder_name][0].insert(0, full_folder_path) + else: + folder_names_and_paths[folder_name][0].append(full_folder_path) else: folder_names_and_paths[folder_name] = ([full_folder_path], set()) @@ -235,6 +268,14 @@ def get_full_path(folder_name: str, filename: str) -> str | None: return None + +def get_full_path_or_raise(folder_name: str, filename: str) -> str: + full_path = get_full_path(folder_name, filename) + if full_path is None: + raise FileNotFoundError(f"Model in folder '{folder_name}' with filename '{filename}' not found.") + return full_path + + def get_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float]: folder_name = map_legacy(folder_name) global folder_names_and_paths @@ -249,6 +290,10 @@ def get_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], f return sorted(list(output_list)), output_folders, time.perf_counter() def cached_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float] | None: + strong_cache = cache_helper.get(folder_name) + if strong_cache is not None: + return strong_cache + global filename_list_cache global folder_names_and_paths folder_name = map_legacy(folder_name) @@ -277,6 +322,7 @@ def get_filename_list(folder_name: str) -> list[str]: out = get_filename_list_(folder_name) global filename_list_cache filename_list_cache[folder_name] = out + cache_helper.set(folder_name, out) return list(out[0]) def get_save_image_path(filename_prefix: str, output_dir: str, image_width=0, image_height=0) -> tuple[str, str, int, str, str]: diff --git a/main.py b/main.py index 00b49fc4c..f5b633e5b 100644 --- a/main.py +++ b/main.py @@ -241,6 +241,7 @@ if __name__ == "__main__": if args.quick_test_for_ci: exit(0) + os.makedirs(folder_paths.get_temp_directory(), exist_ok=True) call_on_start = None if args.auto_launch: def startup_server(scheme, address, port): diff --git a/nodes.py b/nodes.py index 70f52741e..9ec2d6c58 100644 --- a/nodes.py +++ b/nodes.py @@ -515,7 +515,7 @@ class CheckpointLoader: def load_checkpoint(self, config_name, ckpt_name): config_path = folder_paths.get_full_path("configs", config_name) - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) return comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) class CheckpointLoaderSimple: @@ -536,7 +536,7 @@ class CheckpointLoaderSimple: DESCRIPTION = "Loads a diffusion model checkpoint, diffusion models are used to denoise latents." def load_checkpoint(self, ckpt_name): - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) return out[:3] @@ -578,7 +578,7 @@ class unCLIPCheckpointLoader: CATEGORY = "loaders" def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True): - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) return out @@ -625,7 +625,7 @@ class LoraLoader: if strength_model == 0 and strength_clip == 0: return (model, clip) - lora_path = folder_paths.get_full_path("loras", lora_name) + lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) lora = None if self.loaded_lora is not None: if self.loaded_lora[0] == lora_path: @@ -704,11 +704,11 @@ class VAELoader: encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes)) decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes)) - enc = comfy.utils.load_torch_file(folder_paths.get_full_path("vae_approx", encoder)) + enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder)) for k in enc: sd["taesd_encoder.{}".format(k)] = enc[k] - dec = comfy.utils.load_torch_file(folder_paths.get_full_path("vae_approx", decoder)) + dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder)) for k in dec: sd["taesd_decoder.{}".format(k)] = dec[k] @@ -739,7 +739,7 @@ class VAELoader: if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]: sd = self.load_taesd(vae_name) else: - vae_path = folder_paths.get_full_path("vae", vae_name) + vae_path = folder_paths.get_full_path_or_raise("vae", vae_name) sd = comfy.utils.load_torch_file(vae_path) vae = comfy.sd.VAE(sd=sd) return (vae,) @@ -755,7 +755,7 @@ class ControlNetLoader: CATEGORY = "loaders" def load_controlnet(self, control_net_name): - controlnet_path = folder_paths.get_full_path("controlnet", control_net_name) + controlnet_path = folder_paths.get_full_path_or_raise("controlnet", control_net_name) controlnet = comfy.controlnet.load_controlnet(controlnet_path) return (controlnet,) @@ -771,7 +771,7 @@ class DiffControlNetLoader: CATEGORY = "loaders" def load_controlnet(self, model, control_net_name): - controlnet_path = folder_paths.get_full_path("controlnet", control_net_name) + controlnet_path = folder_paths.get_full_path_or_raise("controlnet", control_net_name) controlnet = comfy.controlnet.load_controlnet(controlnet_path, model) return (controlnet,) @@ -871,7 +871,7 @@ class UNETLoader: elif weight_dtype == "fp8_e5m2": model_options["dtype"] = torch.float8_e5m2 - unet_path = folder_paths.get_full_path("diffusion_models", unet_name) + unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name) model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options) return (model,) @@ -896,7 +896,7 @@ class CLIPLoader: else: clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION - clip_path = folder_paths.get_full_path("clip", clip_name) + clip_path = folder_paths.get_full_path_or_raise("clip", clip_name) clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type) return (clip,) @@ -913,8 +913,8 @@ class DualCLIPLoader: CATEGORY = "advanced/loaders" def load_clip(self, clip_name1, clip_name2, type): - clip_path1 = folder_paths.get_full_path("clip", clip_name1) - clip_path2 = folder_paths.get_full_path("clip", clip_name2) + clip_path1 = folder_paths.get_full_path_or_raise("clip", clip_name1) + clip_path2 = folder_paths.get_full_path_or_raise("clip", clip_name2) if type == "sdxl": clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION elif type == "sd3": @@ -936,7 +936,7 @@ class CLIPVisionLoader: CATEGORY = "loaders" def load_clip(self, clip_name): - clip_path = folder_paths.get_full_path("clip_vision", clip_name) + clip_path = folder_paths.get_full_path_or_raise("clip_vision", clip_name) clip_vision = comfy.clip_vision.load(clip_path) return (clip_vision,) @@ -966,7 +966,7 @@ class StyleModelLoader: CATEGORY = "loaders" def load_style_model(self, style_model_name): - style_model_path = folder_paths.get_full_path("style_models", style_model_name) + style_model_path = folder_paths.get_full_path_or_raise("style_models", style_model_name) style_model = comfy.sd.load_style_model(style_model_path) return (style_model,) @@ -1031,7 +1031,7 @@ class GLIGENLoader: CATEGORY = "loaders" def load_gligen(self, gligen_name): - gligen_path = folder_paths.get_full_path("gligen", gligen_name) + gligen_path = folder_paths.get_full_path_or_raise("gligen", gligen_name) gligen = comfy.sd.load_gligen(gligen_path) return (gligen,) @@ -1917,8 +1917,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ConditioningSetArea": "Conditioning (Set Area)", "ConditioningSetAreaPercentage": "Conditioning (Set Area with Percentage)", "ConditioningSetMask": "Conditioning (Set Mask)", - "ControlNetApply": "Apply ControlNet", - "ControlNetApplyAdvanced": "Apply ControlNet (Advanced)", + "ControlNetApply": "Apply ControlNet (OLD)", + "ControlNetApplyAdvanced": "Apply ControlNet (OLD Advanced)", # Latent "VAEEncodeForInpaint": "VAE Encode (for Inpainting)", "SetLatentNoiseMask": "Set Latent Noise Mask", diff --git a/server.py b/server.py index a2c078e40..9321e4d08 100644 --- a/server.py +++ b/server.py @@ -221,6 +221,12 @@ class PromptServer(): def get_embeddings(self): embeddings = folder_paths.get_filename_list("embeddings") return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings))) + + @routes.get("/models") + def list_model_types(request): + model_types = list(folder_paths.folder_names_and_paths.keys()) + + return web.json_response(model_types) @routes.get("/models/{folder}") async def get_models(request): @@ -546,14 +552,15 @@ class PromptServer(): @routes.get("/object_info") async def get_object_info(request): - out = {} - for x in nodes.NODE_CLASS_MAPPINGS: - try: - out[x] = node_info(x) - except Exception as e: - logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.") - logging.error(traceback.format_exc()) - return web.json_response(out) + with folder_paths.cache_helper: + out = {} + for x in nodes.NODE_CLASS_MAPPINGS: + try: + out[x] = node_info(x) + except Exception as e: + logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.") + logging.error(traceback.format_exc()) + return web.json_response(out) @routes.get("/object_info/{node_class}") async def get_object_info_node(request): diff --git a/tests-unit/utils/extra_config_test.py b/tests-unit/utils/extra_config_test.py index 06349560d..0effd89e8 100644 --- a/tests-unit/utils/extra_config_test.py +++ b/tests-unit/utils/extra_config_test.py @@ -71,7 +71,7 @@ def test_load_extra_model_paths_expands_userpath( load_extra_path_config(dummy_yaml_file_name) expected_calls = [ - ('checkpoints', os.path.join(mock_expanded_home, 'App', 'subfolder1')), + ('checkpoints', os.path.join(mock_expanded_home, 'App', 'subfolder1'), False), ] assert mock_add_model_folder_path.call_count == len(expected_calls) @@ -111,7 +111,7 @@ def test_load_extra_model_paths_expands_appdata( expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI' expected_calls = [ - ('checkpoints', os.path.join(expected_base_path, 'models/checkpoints')), + ('checkpoints', os.path.join(expected_base_path, 'models/checkpoints'), False), ] assert mock_add_model_folder_path.call_count == len(expected_calls) diff --git a/utils/extra_config.py b/utils/extra_config.py index 7ac3650ac..908765902 100644 --- a/utils/extra_config.py +++ b/utils/extra_config.py @@ -14,6 +14,9 @@ def load_extra_path_config(yaml_path): if "base_path" in conf: base_path = conf.pop("base_path") base_path = os.path.expandvars(os.path.expanduser(base_path)) + is_default = False + if "is_default" in conf: + is_default = conf.pop("is_default") for x in conf: for y in conf[x].split("\n"): if len(y) == 0: @@ -22,4 +25,4 @@ def load_extra_path_config(yaml_path): if base_path is not None: full_path = os.path.join(base_path, full_path) logging.info("Adding extra search path {} {}".format(x, full_path)) - folder_paths.add_model_folder_path(x, full_path) + folder_paths.add_model_folder_path(x, full_path, is_default)