diff --git a/app/user_manager.py b/app/user_manager.py index 95ad157c3..05e6805f7 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -7,16 +7,15 @@ import shutil from aiohttp import web from urllib import parse from comfy.cli_args import args -from folder_paths import user_directory +import folder_paths from .app_settings import AppSettings default_user = "default" -users_file = os.path.join(user_directory, "users.json") class UserManager(): def __init__(self): - global user_directory + user_directory = folder_paths.get_user_directory() self.settings = AppSettings(self) if not os.path.exists(user_directory): @@ -26,14 +25,17 @@ class UserManager(): print("****** For multi-user setups add the --multi-user CLI argument to enable multiple user profiles. ******") if args.multi_user: - if os.path.isfile(users_file): - with open(users_file) as f: + if os.path.isfile(self.get_users_file()): + with open(self.get_users_file()) as f: self.users = json.load(f) else: self.users = {} else: self.users = {"default": "default"} + def get_users_file(self): + return os.path.join(folder_paths.get_user_directory(), "users.json") + def get_request_user_id(self, request): user = "default" if args.multi_user and "comfy-user" in request.headers: @@ -45,7 +47,7 @@ class UserManager(): return user def get_request_user_filepath(self, request, file, type="userdata", create_dir=True): - global user_directory + user_directory = folder_paths.get_user_directory() if type == "userdata": root_dir = user_directory @@ -63,7 +65,7 @@ class UserManager(): # Check if filename is url encoded if "%" in file: file = parse.unquote(file) - + # prevent leaving /{type}/{user} path = os.path.abspath(os.path.join(user_root, file)) if os.path.commonpath((user_root, path)) != user_root: @@ -85,8 +87,7 @@ class UserManager(): self.users[user_id] = name - global users_file - with open(users_file, "w") as f: + with open(self.get_users_file(), "w") as f: json.dump(self.users, f) return user_id @@ -119,23 +120,44 @@ class UserManager(): async def listuserdata(request): directory = request.rel_url.query.get('dir', '') if not directory: - return web.Response(status=400) - + return web.Response(status=400, text="Directory not provided") + path = self.get_request_user_filepath(request, directory) if not path: - return web.Response(status=403) - + return web.Response(status=403, text="Invalid directory") + if not os.path.exists(path): - return web.Response(status=404) - + return web.Response(status=404, text="Directory not found") + recurse = request.rel_url.query.get('recurse', '').lower() == "true" - results = glob.glob(os.path.join( - glob.escape(path), '**/*'), recursive=recurse) - results = [os.path.relpath(x, path) for x in results if os.path.isfile(x)] - + full_info = request.rel_url.query.get('full_info', '').lower() == "true" + + # Use different patterns based on whether we're recursing or not + if recurse: + pattern = os.path.join(glob.escape(path), '**', '*') + else: + pattern = os.path.join(glob.escape(path), '*') + + results = glob.glob(pattern, recursive=recurse) + + if full_info: + results = [ + { + 'path': os.path.relpath(x, path).replace(os.sep, '/'), + 'size': os.path.getsize(x), + 'modified': os.path.getmtime(x) + } for x in results if os.path.isfile(x) + ] + else: + results = [ + os.path.relpath(x, path).replace(os.sep, '/') + for x in results + if os.path.isfile(x) + ] + split_path = request.rel_url.query.get('split', '').lower() == "true" - if split_path: - results = [[x] + x.split(os.sep) for x in results] + if split_path and not full_info: + results = [[x] + x.split('/') for x in results] return web.json_response(results) @@ -143,14 +165,14 @@ class UserManager(): file = request.match_info.get(param, None) if not file: return web.Response(status=400) - + path = self.get_request_user_filepath(request, file) if not path: return web.Response(status=403) - + if check_exists and not os.path.exists(path): return web.Response(status=404) - + return path @routes.get("/userdata/{file}") @@ -158,7 +180,7 @@ class UserManager(): path = get_user_data_path(request, check_exists=True) if not isinstance(path, str): return path - + return web.FileResponse(path) @routes.post("/userdata/{file}") @@ -166,7 +188,7 @@ class UserManager(): path = get_user_data_path(request) if not isinstance(path, str): return path - + overwrite = request.query["overwrite"] != "false" if not overwrite and os.path.exists(path): return web.Response(status=409) @@ -175,7 +197,7 @@ class UserManager(): with open(path, "wb") as f: f.write(body) - + resp = os.path.relpath(path, self.get_request_user_filepath(request, None)) return web.json_response(resp) @@ -186,7 +208,7 @@ class UserManager(): return path os.remove(path) - + return web.Response(status=204) @routes.post("/userdata/{file}/move/{dest}") @@ -194,17 +216,17 @@ class UserManager(): source = get_user_data_path(request, check_exists=True) if not isinstance(source, str): return source - + dest = get_user_data_path(request, check_exists=False, param="dest") if not isinstance(source, str): return dest - + overwrite = request.query["overwrite"] != "false" if not overwrite and os.path.exists(dest): return web.Response(status=409) print(f"moving '{source}' -> '{dest}'") shutil.move(source, dest) - + resp = os.path.relpath(dest, self.get_request_user_filepath(request, None)) return web.json_response(resp) diff --git a/comfy/cldm/mmdit.py b/comfy/cldm/mmdit.py index 025c2fb5d..54a58ab83 100644 --- a/comfy/cldm/mmdit.py +++ b/comfy/cldm/mmdit.py @@ -6,6 +6,7 @@ class ControlNet(comfy.ldm.modules.diffusionmodules.mmdit.MMDiT): def __init__( self, num_blocks = None, + control_latent_channels = None, dtype = None, device = None, operations = None, @@ -17,10 +18,13 @@ class ControlNet(comfy.ldm.modules.diffusionmodules.mmdit.MMDiT): for _ in range(len(self.joint_blocks)): self.controlnet_blocks.append(operations.Linear(self.hidden_size, self.hidden_size, device=device, dtype=dtype)) + if control_latent_channels is None: + control_latent_channels = self.in_channels + self.pos_embed_input = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed( None, self.patch_size, - self.in_channels, + control_latent_channels, self.hidden_size, bias=True, strict_img_size=False, diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 65b879907..d0f706432 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -171,6 +171,8 @@ parser.add_argument( help="The local filesystem path to the directory where the frontend is located. Overrides --front-end-version.", ) +parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path.") + if comfy.options.args_parsing: args = parser.parse_args() else: diff --git a/comfy/types.py b/comfy/comfy_types.py similarity index 100% rename from comfy/types.py rename to comfy/comfy_types.py diff --git a/comfy/controlnet.py b/comfy/controlnet.py index c0f9b6511..1ea00eccf 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -79,13 +79,19 @@ class ControlBase: self.previous_controlnet = None self.extra_conds = [] self.strength_type = StrengthType.CONSTANT + self.concat_mask = False + self.extra_concat_orig = [] + self.extra_concat = None - def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None): + def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None, extra_concat=[]): self.cond_hint_original = cond_hint self.strength = strength self.timestep_percent_range = timestep_percent_range if self.latent_format is not None: self.vae = vae + self.extra_concat_orig = extra_concat.copy() + if self.concat_mask and len(self.extra_concat_orig) == 0: + self.extra_concat_orig.append(torch.tensor([[[[1.0]]]])) return self def pre_run(self, model, percent_to_timestep_function): @@ -100,9 +106,9 @@ class ControlBase: def cleanup(self): if self.previous_controlnet is not None: self.previous_controlnet.cleanup() - if self.cond_hint is not None: - del self.cond_hint - self.cond_hint = None + + self.cond_hint = None + self.extra_concat = None self.timestep_range = None def get_models(self): @@ -123,6 +129,8 @@ class ControlBase: c.vae = self.vae c.extra_conds = self.extra_conds.copy() c.strength_type = self.strength_type + c.concat_mask = self.concat_mask + c.extra_concat_orig = self.extra_concat_orig.copy() def inference_memory_requirements(self, dtype): if self.previous_controlnet is not None: @@ -175,7 +183,7 @@ class ControlBase: class ControlNet(ControlBase): - def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT): + def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False): super().__init__(device) self.control_model = control_model self.load_device = load_device @@ -189,6 +197,7 @@ class ControlNet(ControlBase): self.latent_format = latent_format self.extra_conds += extra_conds self.strength_type = strength_type + self.concat_mask = concat_mask def get_control(self, x_noisy, t, cond, batched_number): control_prev = None @@ -220,6 +229,13 @@ class ControlNet(ControlBase): comfy.model_management.load_models_gpu(loaded_models) if self.latent_format is not None: self.cond_hint = self.latent_format.process_in(self.cond_hint) + if len(self.extra_concat_orig) > 0: + to_concat = [] + for c in self.extra_concat_orig: + c = comfy.utils.common_upscale(c, self.cond_hint.shape[3], self.cond_hint.shape[2], self.upscale_algorithm, "center") + to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0])) + self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1) + self.cond_hint = self.cond_hint.to(device=self.device, dtype=dtype) if x_noisy.shape[0] != self.cond_hint.shape[0]: self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number) @@ -410,12 +426,17 @@ def load_controlnet_mmdit(sd): for k in sd: new_sd[k] = sd[k] - control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) + concat_mask = False + control_latent_channels = new_sd.get("pos_embed_input.proj.weight").shape[1] + if control_latent_channels == 17: #inpaint controlnet + concat_mask = True + + control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) control_model = controlnet_load_state_dict(control_model, new_sd) 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, 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) return control @@ -449,12 +470,17 @@ def load_controlnet_flux_instantx(sd): if union_cnet in new_sd: num_union_modes = new_sd[union_cnet].shape[0] - control_model = comfy.ldm.flux.controlnet.ControlNetFlux(latent_input=True, num_union_modes=num_union_modes, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) + control_latent_channels = new_sd.get("pos_embed_input.weight").shape[1] // 4 + concat_mask = False + if control_latent_channels == 17: + concat_mask = True + + control_model = comfy.ldm.flux.controlnet.ControlNetFlux(latent_input=True, num_union_modes=num_union_modes, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) control_model = controlnet_load_state_dict(control_model, new_sd) latent_format = comfy.latent_formats.Flux() extra_conds = ['y', 'guidance'] - 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) + 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) return control def convert_mistoline(sd): diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index 385a25164..ae56aee8b 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -1101,3 +1101,45 @@ def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=No if sigmas[i + 1] > 0: x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up return x +@torch.no_grad() +def sample_dpmpp_2s_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): + """Ancestral sampling with DPM-Solver++(2S) second-order steps.""" + extra_args = {} if extra_args is None else extra_args + noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler + + temp = [0] + def post_cfg_function(args): + temp[0] = args["uncond_denoised"] + return args["denoised"] + + model_options = extra_args.get("model_options", {}).copy() + extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True) + + s_in = x.new_ones([x.shape[0]]) + sigma_fn = lambda t: t.neg().exp() + t_fn = lambda sigma: sigma.log().neg() + + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) + sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) + if sigma_down == 0: + # Euler method + d = to_d(x, sigmas[i], temp[0]) + dt = sigma_down - sigmas[i] + x = denoised + d * sigma_down + else: + # DPM-Solver++(2S) + t, t_next = t_fn(sigmas[i]), t_fn(sigma_down) + # r = torch.sinh(1 + (2 - eta) * (t_next - t) / (t - t_fn(sigma_up))) works only on non-cfgpp, weird + r = 1 / 2 + h = t_next - t + s = t + r * h + x_2 = (sigma_fn(s) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h * r).expm1() * denoised + denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args) + x = (sigma_fn(t_next) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h).expm1() * denoised_2 + # Noise addition + if sigmas[i + 1] > 0: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up + return x diff --git a/comfy/ldm/flux/controlnet.py b/comfy/ldm/flux/controlnet.py index d8b776129..c033dea52 100644 --- a/comfy/ldm/flux/controlnet.py +++ b/comfy/ldm/flux/controlnet.py @@ -52,7 +52,7 @@ class MistolineControlnetBlock(nn.Module): class ControlNetFlux(Flux): - def __init__(self, latent_input=False, num_union_modes=0, mistoline=False, image_model=None, dtype=None, device=None, operations=None, **kwargs): + def __init__(self, latent_input=False, num_union_modes=0, mistoline=False, control_latent_channels=None, image_model=None, dtype=None, device=None, operations=None, **kwargs): super().__init__(final_layer=False, dtype=dtype, device=device, operations=operations, **kwargs) self.main_model_double = 19 @@ -80,7 +80,12 @@ class ControlNetFlux(Flux): self.gradient_checkpointing = False self.latent_input = latent_input - self.pos_embed_input = operations.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device) + if control_latent_channels is None: + control_latent_channels = self.in_channels + else: + control_latent_channels *= 2 * 2 #patch size + + self.pos_embed_input = operations.Linear(control_latent_channels, self.hidden_size, bias=True, dtype=dtype, device=device) if not self.latent_input: if self.mistoline: self.input_cond_block = MistolineCondDownsamplBlock(dtype=dtype, device=device, operations=operations) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 3f5d90273..a89ad4fd7 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -28,7 +28,7 @@ import comfy.utils import comfy.float import comfy.model_management import comfy.lora -from comfy.types import UnetWrapperFunction +from comfy.comfy_types import UnetWrapperFunction def string_to_seed(data): crc = 0xFFFFFFFF diff --git a/comfy/samplers.py b/comfy/samplers.py index b278dae24..a20f65d4a 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -570,7 +570,7 @@ class Sampler: return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", - "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", + "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu", "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "ipndm", "ipndm_v", "deis"] diff --git a/comfy/sd.py b/comfy/sd.py index 1e1a6594b..07310b9d4 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -445,12 +445,8 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_target.tokenizer = comfy.text_encoders.sa_t5.SAT5Tokenizer else: w = clip_data[0].get("text_model.embeddings.position_embedding.weight", None) - if w is not None and w.shape[0] == 248: - clip_target.clip = comfy.text_encoders.long_clipl.LongClipModel - clip_target.tokenizer = comfy.text_encoders.long_clipl.LongClipTokenizer - else: - clip_target.clip = sd1_clip.SD1ClipModel - clip_target.tokenizer = sd1_clip.SD1Tokenizer + clip_target.clip = sd1_clip.SD1ClipModel + clip_target.tokenizer = sd1_clip.SD1Tokenizer elif len(clip_data) == 2: if clip_type == CLIPType.SD3: clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=True, clip_g=True, t5=False) @@ -475,10 +471,12 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer parameters = 0 + tokenizer_data = {} for c in clip_data: parameters += comfy.utils.calculate_parameters(c) + tokenizer_data, model_options = comfy.text_encoders.long_clipl.model_options_long_clip(c, tokenizer_data, model_options) - clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, model_options=model_options) + clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, tokenizer_data=tokenizer_data, model_options=model_options) for c in clip_data: m, u = clip.load_sd(c) if len(m) > 0: diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py index 676653f77..9f0d95b0a 100644 --- a/comfy/sd1_clip.py +++ b/comfy/sd1_clip.py @@ -542,6 +542,7 @@ class SD1Tokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}, clip_name="l", tokenizer=SDTokenizer): self.clip_name = clip_name self.clip = "clip_{}".format(self.clip_name) + tokenizer = tokenizer_data.get("{}_tokenizer_class".format(self.clip), tokenizer) setattr(self, self.clip, tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -570,6 +571,7 @@ class SD1ClipModel(torch.nn.Module): self.clip_name = clip_name self.clip = "clip_{}".format(self.clip_name) + clip_model = model_options.get("{}_class".format(self.clip), clip_model) setattr(self, self.clip, clip_model(device=device, dtype=dtype, model_options=model_options, **kwargs)) self.dtypes = set() diff --git a/comfy/sdxl_clip.py b/comfy/sdxl_clip.py index a0145caa4..4d0a4e8e7 100644 --- a/comfy/sdxl_clip.py +++ b/comfy/sdxl_clip.py @@ -22,7 +22,8 @@ class SDXLClipGTokenizer(sd1_clip.SDTokenizer): class SDXLTokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.clip_g = SDXLClipGTokenizer(embedding_directory=embedding_directory) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -40,7 +41,8 @@ class SDXLTokenizer: class SDXLClipModel(torch.nn.Module): def __init__(self, device="cpu", dtype=None, model_options={}): super().__init__() - self.clip_l = sd1_clip.SDClipModel(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, model_options=model_options) self.clip_g = SDXLClipG(device=device, dtype=dtype, model_options=model_options) self.dtypes = set([dtype]) @@ -57,7 +59,8 @@ class SDXLClipModel(torch.nn.Module): token_weight_pairs_l = token_weight_pairs["l"] g_out, g_pooled = self.clip_g.encode_token_weights(token_weight_pairs_g) l_out, l_pooled = self.clip_l.encode_token_weights(token_weight_pairs_l) - return torch.cat([l_out, g_out], dim=-1), g_pooled + cut_to = min(l_out.shape[1], g_out.shape[1]) + return torch.cat([l_out[:,:cut_to], g_out[:,:cut_to]], dim=-1), g_pooled def load_sd(self, sd): if "text_model.encoder.layers.30.mlp.fc1.weight" in sd: diff --git a/comfy/text_encoders/flux.py b/comfy/text_encoders/flux.py index 91d1f249d..7c75fe64b 100644 --- a/comfy/text_encoders/flux.py +++ b/comfy/text_encoders/flux.py @@ -18,7 +18,8 @@ class T5XXLTokenizer(sd1_clip.SDTokenizer): class FluxTokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.t5xxl = T5XXLTokenizer(embedding_directory=embedding_directory) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -38,7 +39,8 @@ class FluxClipModel(torch.nn.Module): def __init__(self, dtype_t5=None, device="cpu", dtype=None, model_options={}): super().__init__() dtype_t5 = comfy.model_management.pick_weight_dtype(dtype_t5, dtype, device) - self.clip_l = sd1_clip.SDClipModel(device=device, dtype=dtype, return_projected_pooled=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(device=device, dtype=dtype, return_projected_pooled=False, model_options=model_options) self.t5xxl = T5XXLModel(device=device, dtype=dtype_t5, model_options=model_options) self.dtypes = set([dtype, dtype_t5]) diff --git a/comfy/text_encoders/long_clipl.py b/comfy/text_encoders/long_clipl.py index 4677fb3b0..b81912cb3 100644 --- a/comfy/text_encoders/long_clipl.py +++ b/comfy/text_encoders/long_clipl.py @@ -6,9 +6,9 @@ class LongClipTokenizer_(sd1_clip.SDTokenizer): super().__init__(max_length=248, embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) class LongClipModel_(sd1_clip.SDClipModel): - def __init__(self, device="cpu", dtype=None, model_options={}): + def __init__(self, *args, **kwargs): textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "long_clipl.json") - super().__init__(device=device, textmodel_json_config=textmodel_json_config, return_projected_pooled=False, dtype=dtype, model_options=model_options) + super().__init__(*args, textmodel_json_config=textmodel_json_config, **kwargs) class LongClipTokenizer(sd1_clip.SD1Tokenizer): def __init__(self, embedding_directory=None, tokenizer_data={}): @@ -17,3 +17,14 @@ class LongClipTokenizer(sd1_clip.SD1Tokenizer): class LongClipModel(sd1_clip.SD1ClipModel): def __init__(self, device="cpu", dtype=None, model_options={}, **kwargs): super().__init__(device=device, dtype=dtype, model_options=model_options, clip_model=LongClipModel_, **kwargs) + +def model_options_long_clip(sd, tokenizer_data, model_options): + w = sd.get("clip_l.text_model.embeddings.position_embedding.weight", None) + if w is None: + w = sd.get("text_model.embeddings.position_embedding.weight", None) + if w is not None and w.shape[0] == 248: + tokenizer_data = tokenizer_data.copy() + model_options = model_options.copy() + tokenizer_data["clip_l_tokenizer_class"] = LongClipTokenizer_ + model_options["clip_l_class"] = LongClipModel_ + return tokenizer_data, model_options diff --git a/comfy/text_encoders/sd3_clip.py b/comfy/text_encoders/sd3_clip.py index e3832ac24..c54f28856 100644 --- a/comfy/text_encoders/sd3_clip.py +++ b/comfy/text_encoders/sd3_clip.py @@ -20,7 +20,8 @@ class T5XXLTokenizer(sd1_clip.SDTokenizer): class SD3Tokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.clip_g = sdxl_clip.SDXLClipGTokenizer(embedding_directory=embedding_directory) self.t5xxl = T5XXLTokenizer(embedding_directory=embedding_directory) @@ -42,7 +43,8 @@ class SD3ClipModel(torch.nn.Module): super().__init__() self.dtypes = set() if clip_l: - self.clip_l = sd1_clip.SDClipModel(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, return_projected_pooled=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, return_projected_pooled=False, model_options=model_options) self.dtypes.add(dtype) else: self.clip_l = None @@ -95,7 +97,8 @@ class SD3ClipModel(torch.nn.Module): if self.clip_g is not None: g_out, g_pooled = self.clip_g.encode_token_weights(token_weight_pairs_g) if lg_out is not None: - lg_out = torch.cat([lg_out, g_out], dim=-1) + cut_to = min(lg_out.shape[1], g_out.shape[1]) + lg_out = torch.cat([lg_out[:,:cut_to], g_out[:,:cut_to]], dim=-1) else: lg_out = torch.nn.functional.pad(g_out, (768, 0)) else: diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index 62311ed7f..630f280fc 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -1,11 +1,21 @@ import itertools -from typing import Sequence, Mapping +from typing import Sequence, Mapping, Dict from comfy_execution.graph import DynamicPrompt import nodes from comfy_execution.graph_utils import is_link +NODE_CLASS_CONTAINS_UNIQUE_ID: Dict[str, bool] = {} + + +def include_unique_id_in_input(class_type: str) -> bool: + if class_type in NODE_CLASS_CONTAINS_UNIQUE_ID: + return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] + class_def = nodes.NODE_CLASS_MAPPINGS[class_type] + NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] = "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values() + return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] + class CacheKeySet: def __init__(self, dynprompt, node_ids, is_changed_cache): self.keys = {} @@ -98,7 +108,7 @@ class CacheKeySetInputSignature(CacheKeySet): class_type = node["class_type"] class_def = nodes.NODE_CLASS_MAPPINGS[class_type] signature = [class_type, self.is_changed_cache.get(node_id)] - if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values(): + if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or include_unique_id_in_input(class_type): signature.append(node_id) inputs = node["inputs"] for key in sorted(inputs.keys()): diff --git a/comfy_execution/graph.py b/comfy_execution/graph.py index b53e10f3f..0b5bf1899 100644 --- a/comfy_execution/graph.py +++ b/comfy_execution/graph.py @@ -99,30 +99,44 @@ class TopologicalSort: self.add_strong_link(from_node_id, from_socket, to_node_id) def add_strong_link(self, from_node_id, from_socket, to_node_id): - self.add_node(from_node_id) - if to_node_id not in self.blocking[from_node_id]: - self.blocking[from_node_id][to_node_id] = {} - self.blockCount[to_node_id] += 1 - self.blocking[from_node_id][to_node_id][from_socket] = True + if not self.is_cached(from_node_id): + self.add_node(from_node_id) + if to_node_id not in self.blocking[from_node_id]: + self.blocking[from_node_id][to_node_id] = {} + self.blockCount[to_node_id] += 1 + self.blocking[from_node_id][to_node_id][from_socket] = True - def add_node(self, unique_id, include_lazy=False, subgraph_nodes=None): - if unique_id in self.pendingNodes: - return - self.pendingNodes[unique_id] = True - self.blockCount[unique_id] = 0 - self.blocking[unique_id] = {} + def add_node(self, node_unique_id, include_lazy=False, subgraph_nodes=None): + node_ids = [node_unique_id] + links = [] - inputs = self.dynprompt.get_node(unique_id)["inputs"] - for input_name in inputs: - value = inputs[input_name] - if is_link(value): - from_node_id, from_socket = value - if subgraph_nodes is not None and from_node_id not in subgraph_nodes: - continue - input_type, input_category, input_info = self.get_input_info(unique_id, input_name) - is_lazy = input_info is not None and "lazy" in input_info and input_info["lazy"] - if include_lazy or not is_lazy: - self.add_strong_link(from_node_id, from_socket, unique_id) + while len(node_ids) > 0: + unique_id = node_ids.pop() + if unique_id in self.pendingNodes: + continue + + self.pendingNodes[unique_id] = True + self.blockCount[unique_id] = 0 + self.blocking[unique_id] = {} + + inputs = self.dynprompt.get_node(unique_id)["inputs"] + for input_name in inputs: + value = inputs[input_name] + if is_link(value): + from_node_id, from_socket = value + if subgraph_nodes is not None and from_node_id not in subgraph_nodes: + continue + input_type, input_category, input_info = self.get_input_info(unique_id, input_name) + is_lazy = input_info is not None and "lazy" in input_info and input_info["lazy"] + if (include_lazy or not is_lazy) and not self.is_cached(from_node_id): + node_ids.append(from_node_id) + links.append((from_node_id, from_socket, unique_id)) + + for link in links: + self.add_strong_link(*link) + + def is_cached(self, node_id): + return False def get_ready_nodes(self): return [node_id for node_id in self.pendingNodes if self.blockCount[node_id] == 0] @@ -146,11 +160,8 @@ class ExecutionList(TopologicalSort): self.output_cache = output_cache self.staged_node_id = None - def add_strong_link(self, from_node_id, from_socket, to_node_id): - if self.output_cache.get(from_node_id) is not None: - # Nothing to do - return - super().add_strong_link(from_node_id, from_socket, to_node_id) + def is_cached(self, node_id): + return self.output_cache.get(node_id) is not None def stage_node_execution(self): assert self.staged_node_id is None diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py index 762b48279..6990b3f9a 100644 --- a/comfy_extras/nodes_audio.py +++ b/comfy_extras/nodes_audio.py @@ -183,17 +183,10 @@ class PreviewAudio(SaveAudio): } class LoadAudio: - SUPPORTED_FORMATS = ('.wav', '.mp3', '.ogg', '.flac', '.aiff', '.aif') - @classmethod def INPUT_TYPES(s): input_dir = folder_paths.get_input_directory() - files = [ - f for f in os.listdir(input_dir) - if (os.path.isfile(os.path.join(input_dir, f)) - and f.endswith(LoadAudio.SUPPORTED_FORMATS) - ) - ] + files = folder_paths.filter_files_content_types(os.listdir(input_dir), ["audio", "video"]) return {"required": {"audio": (sorted(files), {"audio_upload": True})}} CATEGORY = "audio" diff --git a/comfy_extras/nodes_controlnet.py b/comfy_extras/nodes_controlnet.py index 0773c8a5b..2d20e1fed 100644 --- a/comfy_extras/nodes_controlnet.py +++ b/comfy_extras/nodes_controlnet.py @@ -1,4 +1,6 @@ from comfy.cldm.control_types import UNION_CONTROLNET_TYPES +import nodes +import comfy.utils class SetUnionControlNetType: @classmethod @@ -22,6 +24,37 @@ class SetUnionControlNetType: return (control_net,) +class ControlNetInpaintingAliMamaApply(nodes.ControlNetApplyAdvanced): + @classmethod + def INPUT_TYPES(s): + return {"required": {"positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "control_net": ("CONTROL_NET", ), + "vae": ("VAE", ), + "image": ("IMAGE", ), + "mask": ("MASK", ), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}) + }} + + FUNCTION = "apply_inpaint_controlnet" + + CATEGORY = "conditioning/controlnet" + + def apply_inpaint_controlnet(self, positive, negative, control_net, vae, image, mask, strength, start_percent, end_percent): + extra_concat = [] + if control_net.concat_mask: + mask = 1.0 - mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) + mask_apply = comfy.utils.common_upscale(mask, image.shape[2], image.shape[1], "bilinear", "center").round() + image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3]) + extra_concat = [mask] + + return self.apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent, vae=vae, extra_concat=extra_concat) + + + NODE_CLASS_MAPPINGS = { "SetUnionControlNetType": SetUnionControlNetType, + "ControlNetInpaintingAliMamaApply": ControlNetInpaintingAliMamaApply, } diff --git a/comfy_extras/nodes_perpneg.py b/comfy_extras/nodes_perpneg.py index 546276aa1..762c40220 100644 --- a/comfy_extras/nodes_perpneg.py +++ b/comfy_extras/nodes_perpneg.py @@ -26,6 +26,7 @@ class PerpNeg: FUNCTION = "patch" CATEGORY = "_for_testing" + DEPRECATED = True def patch(self, model, empty_conditioning, neg_scale): m = model.clone() diff --git a/comfy_extras/nodes_torch_compile.py b/comfy_extras/nodes_torch_compile.py new file mode 100644 index 000000000..1d914fa93 --- /dev/null +++ b/comfy_extras/nodes_torch_compile.py @@ -0,0 +1,21 @@ +import torch + +class TorchCompileModel: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model": ("MODEL",), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "_for_testing" + EXPERIMENTAL = True + + def patch(self, model): + m = model.clone() + m.add_object_patch("diffusion_model", torch.compile(model=m.get_model_object("diffusion_model"))) + return (m, ) + +NODE_CLASS_MAPPINGS = { + "TorchCompileModel": TorchCompileModel, +} diff --git a/extra_model_paths.yaml.example b/extra_model_paths.yaml.example index 846d04dbe..d4fd55c87 100644 --- a/extra_model_paths.yaml.example +++ b/extra_model_paths.yaml.example @@ -30,6 +30,9 @@ a111: # clip_vision: models/clip_vision/ # configs: models/configs/ # controlnet: models/controlnet/ +# diffusion_models: | +# models/diffusion_models +# models/unet # embeddings: models/embeddings/ # loras: models/loras/ # upscale_models: models/upscale_models/ diff --git a/folder_paths.py b/folder_paths.py index b154448fc..263704a4c 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -2,7 +2,9 @@ from __future__ import annotations import os import time +import mimetypes import logging +from typing import Set, List, Dict, Tuple, Literal from collections.abc import Collection supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl', '.sft'} @@ -44,6 +46,10 @@ 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]] = {} +extension_mimetypes_cache = { + "webp" : "image", +} + def map_legacy(folder_name: str) -> str: legacy = {"unet": "diffusion_models"} return legacy.get(folder_name, folder_name) @@ -78,6 +84,13 @@ def get_input_directory() -> str: global input_directory return input_directory +def get_user_directory() -> str: + return user_directory + +def set_user_directory(user_dir: str) -> None: + global user_directory + user_directory = user_dir + #NOTE: used in http server so don't put folders that should not be accessed remotely def get_directory_by_type(type_name: str) -> str | None: @@ -89,6 +102,28 @@ def get_directory_by_type(type_name: str) -> str | None: return get_input_directory() return None +def filter_files_content_types(files: List[str], content_types: Literal["image", "video", "audio"]) -> List[str]: + """ + Example: + files = os.listdir(folder_paths.get_input_directory()) + filter_files_content_types(files, ["image", "audio", "video"]) + """ + global extension_mimetypes_cache + result = [] + for file in files: + extension = file.split('.')[-1] + if extension not in extension_mimetypes_cache: + mime_type, _ = mimetypes.guess_type(file, strict=False) + if not mime_type: + continue + content_type = mime_type.split('/')[0] + extension_mimetypes_cache[extension] = content_type + else: + content_type = extension_mimetypes_cache[extension] + + if content_type in content_types: + result.append(file) + return result # determine base_dir rely on annotation if name is 'filename.ext [annotation]' format # otherwise use default_path as base_dir diff --git a/main.py b/main.py index d791a169c..00b49fc4c 100644 --- a/main.py +++ b/main.py @@ -63,6 +63,7 @@ import threading import gc import logging +import utils.extra_config if os.name == "nt": logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage()) @@ -85,7 +86,6 @@ if args.windows_standalone_build: pass import comfy.utils -import yaml import execution import server @@ -180,27 +180,6 @@ def cleanup_temp(): shutil.rmtree(temp_dir, ignore_errors=True) -def load_extra_path_config(yaml_path): - with open(yaml_path, 'r') as stream: - config = yaml.safe_load(stream) - for c in config: - conf = config[c] - if conf is None: - continue - base_path = None - if "base_path" in conf: - base_path = conf.pop("base_path") - for x in conf: - for y in conf[x].split("\n"): - if len(y) == 0: - continue - full_path = y - 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) - - if __name__ == "__main__": if args.temp_directory: temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp") @@ -222,11 +201,11 @@ if __name__ == "__main__": extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") if os.path.isfile(extra_model_paths_config_path): - load_extra_path_config(extra_model_paths_config_path) + utils.extra_config.load_extra_path_config(extra_model_paths_config_path) if args.extra_model_paths_config: for config_path in itertools.chain(*args.extra_model_paths_config): - load_extra_path_config(config_path) + utils.extra_config.load_extra_path_config(config_path) nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes) @@ -253,6 +232,11 @@ if __name__ == "__main__": input_dir = os.path.abspath(args.input_directory) logging.info(f"Setting input directory to: {input_dir}") folder_paths.set_input_directory(input_dir) + + if args.user_directory: + user_dir = os.path.abspath(args.user_directory) + logging.info(f"Setting user directory to: {user_dir}") + folder_paths.set_user_directory(user_dir) if args.quick_test_for_ci: exit(0) diff --git a/nodes.py b/nodes.py index 7378a528c..478d1af43 100644 --- a/nodes.py +++ b/nodes.py @@ -511,6 +511,7 @@ class CheckpointLoader: FUNCTION = "load_checkpoint" CATEGORY = "advanced/loaders" + DEPRECATED = True def load_checkpoint(self, config_name, ckpt_name): config_path = folder_paths.get_full_path("configs", config_name) @@ -823,7 +824,7 @@ class ControlNetApplyAdvanced: CATEGORY = "conditioning/controlnet" - def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None): + def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None, extra_concat=[]): if strength == 0: return (positive, negative) @@ -840,7 +841,7 @@ class ControlNetApplyAdvanced: if prev_cnet in cnets: c_net = cnets[prev_cnet] else: - c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae) + c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae=vae, extra_concat=extra_concat) c_net.set_previous_controlnet(prev_cnet) cnets[prev_cnet] = c_net @@ -2103,6 +2104,7 @@ def init_builtin_extra_nodes(): "nodes_hunyuan.py", "nodes_flux.py", "nodes_lora_extract.py", + "nodes_torch_compile.py", ] import_failed = [] diff --git a/server.py b/server.py index 93147f6c5..a2c078e40 100644 --- a/server.py +++ b/server.py @@ -12,6 +12,8 @@ import json import glob import struct import ssl +import socket +import ipaddress from PIL import Image, ImageOps from PIL.PngImagePlugin import PngInfo from io import BytesIO @@ -80,20 +82,55 @@ def create_cors_middleware(allowed_origin: str): return cors_middleware +def is_loopback(host): + if host is None: + return False + try: + if ipaddress.ip_address(host).is_loopback: + return True + else: + return False + except: + pass + + loopback = False + for family in (socket.AF_INET, socket.AF_INET6): + try: + r = socket.getaddrinfo(host, None, family, socket.SOCK_STREAM) + for family, _, _, _, sockaddr in r: + if not ipaddress.ip_address(sockaddr[0]).is_loopback: + return loopback + else: + loopback = True + except socket.gaierror: + pass + + return loopback + + def create_origin_only_middleware(): @web.middleware async def origin_only_middleware(request: web.Request, handler): + #this code is used to prevent the case where a random website can queue comfy workflows by making a POST to 127.0.0.1 which browsers don't prevent for some dumb reason. + #in that case the Host and Origin hostnames won't match + #I know the proper fix would be to add a cookie but this should take care of the problem in the meantime if 'Host' in request.headers and 'Origin' in request.headers: host = request.headers['Host'] origin = request.headers['Origin'] host_domain = host.lower() parsed = urllib.parse.urlparse(origin) origin_domain = parsed.netloc.lower() - if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers - result = urllib.parse.urlsplit('//' + host_domain) - host_domain = result.hostname + host_domain_parsed = urllib.parse.urlsplit('//' + host_domain) - if len(host_domain) > 0 and len(origin_domain) > 0: + #limit the check to when the host domain is localhost, this makes it slightly less safe but should still prevent the exploit + loopback = is_loopback(host_domain_parsed.hostname) + + if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers, same for host + host_domain = host_domain_parsed.hostname + if host_domain_parsed.port is None: + origin_domain = parsed.hostname + + if loopback and host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0: if host_domain != origin_domain: logging.warning("WARNING: request with non matching host and origin {} != {}, returning 403".format(host_domain, origin_domain)) return web.Response(status=403) diff --git a/tests-unit/README.md b/tests-unit/README.md index 94abd9853..81692b8f1 100644 --- a/tests-unit/README.md +++ b/tests-unit/README.md @@ -2,7 +2,7 @@ ## Install test dependencies -`pip install -r tests-units/requirements.txt` +`pip install -r tests-unit/requirements.txt` ## Run tests -`pytest tests-units/` +`pytest tests-unit/` diff --git a/tests-unit/comfy_test/folder_path_test.py b/tests-unit/comfy_test/folder_path_test.py new file mode 100644 index 000000000..0bbec593b --- /dev/null +++ b/tests-unit/comfy_test/folder_path_test.py @@ -0,0 +1,66 @@ +### 🗻 This file is created through the spirit of Mount Fuji at its peak +# TODO(yoland): clean up this after I get back down +import pytest +import os +import tempfile +from unittest.mock import patch + +import folder_paths + +@pytest.fixture +def temp_dir(): + with tempfile.TemporaryDirectory() as tmpdirname: + yield tmpdirname + + +def test_get_directory_by_type(): + test_dir = "/test/dir" + folder_paths.set_output_directory(test_dir) + assert folder_paths.get_directory_by_type("output") == test_dir + assert folder_paths.get_directory_by_type("invalid") is None + +def test_annotated_filepath(): + assert folder_paths.annotated_filepath("test.txt") == ("test.txt", None) + assert folder_paths.annotated_filepath("test.txt [output]") == ("test.txt", folder_paths.get_output_directory()) + assert folder_paths.annotated_filepath("test.txt [input]") == ("test.txt", folder_paths.get_input_directory()) + assert folder_paths.annotated_filepath("test.txt [temp]") == ("test.txt", folder_paths.get_temp_directory()) + +def test_get_annotated_filepath(): + default_dir = "/default/dir" + assert folder_paths.get_annotated_filepath("test.txt", default_dir) == os.path.join(default_dir, "test.txt") + assert folder_paths.get_annotated_filepath("test.txt [output]") == os.path.join(folder_paths.get_output_directory(), "test.txt") + +def test_add_model_folder_path(): + folder_paths.add_model_folder_path("test_folder", "/test/path") + assert "/test/path" in folder_paths.get_folder_paths("test_folder") + +def test_recursive_search(temp_dir): + os.makedirs(os.path.join(temp_dir, "subdir")) + open(os.path.join(temp_dir, "file1.txt"), "w").close() + open(os.path.join(temp_dir, "subdir", "file2.txt"), "w").close() + + files, dirs = folder_paths.recursive_search(temp_dir) + assert set(files) == {"file1.txt", os.path.join("subdir", "file2.txt")} + assert len(dirs) == 2 # temp_dir and subdir + +def test_filter_files_extensions(): + files = ["file1.txt", "file2.jpg", "file3.png", "file4.txt"] + assert folder_paths.filter_files_extensions(files, [".txt"]) == ["file1.txt", "file4.txt"] + assert folder_paths.filter_files_extensions(files, [".jpg", ".png"]) == ["file2.jpg", "file3.png"] + assert folder_paths.filter_files_extensions(files, []) == files + +@patch("folder_paths.recursive_search") +@patch("folder_paths.folder_names_and_paths") +def test_get_filename_list(mock_folder_names_and_paths, mock_recursive_search): + mock_folder_names_and_paths.__getitem__.return_value = (["/test/path"], {".txt"}) + mock_recursive_search.return_value = (["file1.txt", "file2.jpg"], {}) + assert folder_paths.get_filename_list("test_folder") == ["file1.txt"] + +def test_get_save_image_path(temp_dir): + with patch("folder_paths.output_directory", temp_dir): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("test", temp_dir, 100, 100) + assert os.path.samefile(full_output_folder, temp_dir) + assert filename == "test" + assert counter == 1 + assert subfolder == "" + assert filename_prefix == "test" \ No newline at end of file diff --git a/tests-unit/folder_paths_test/__init__.py b/tests-unit/folder_paths_test/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests-unit/folder_paths_test/filter_by_content_types_test.py b/tests-unit/folder_paths_test/filter_by_content_types_test.py new file mode 100644 index 000000000..5941bfa94 --- /dev/null +++ b/tests-unit/folder_paths_test/filter_by_content_types_test.py @@ -0,0 +1,52 @@ +import pytest +import os +import tempfile +from folder_paths import filter_files_content_types + +@pytest.fixture(scope="module") +def file_extensions(): + return { + 'image': ['bmp', 'cdr', 'gif', 'heif', 'ico', 'jpeg', 'jpg', 'pcx', 'png', 'pnm', 'ppm', 'psd', 'sgi', 'svg', 'tiff', 'webp', 'xbm', 'xcf', 'xpm'], + 'audio': ['aif', 'aifc', 'aiff', 'au', 'awb', 'flac', 'm4a', 'mp2', 'mp3', 'ogg', 'sd2', 'smp', 'snd', 'wav'], + 'video': ['avi', 'flv', 'm2v', 'm4v', 'mj2', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ogv', 'qt', 'webm', 'wmv'] + } + + +@pytest.fixture(scope="module") +def mock_dir(file_extensions): + with tempfile.TemporaryDirectory() as directory: + for content_type, extensions in file_extensions.items(): + for extension in extensions: + with open(f"{directory}/sample_{content_type}.{extension}", "w") as f: + f.write(f"Sample {content_type} file in {extension} format") + yield directory + + +def test_categorizes_all_correctly(mock_dir, file_extensions): + files = os.listdir(mock_dir) + for content_type, extensions in file_extensions.items(): + filtered_files = filter_files_content_types(files, [content_type]) + for extension in extensions: + assert f"sample_{content_type}.{extension}" in filtered_files + + +def test_categorizes_all_uniquely(mock_dir, file_extensions): + files = os.listdir(mock_dir) + for content_type, extensions in file_extensions.items(): + filtered_files = filter_files_content_types(files, [content_type]) + assert len(filtered_files) == len(extensions) + + +def test_handles_bad_extensions(): + files = ["file1.txt", "file2.py", "file3.example", "file4.pdf", "file5.ini", "file6.doc", "file7.md"] + assert filter_files_content_types(files, ["image", "audio", "video"]) == [] + + +def test_handles_no_extension(): + files = ["file1", "file2", "file3", "file4", "file5", "file6", "file7"] + assert filter_files_content_types(files, ["image", "audio", "video"]) == [] + + +def test_handles_no_files(): + files = [] + assert filter_files_content_types(files, ["image", "audio", "video"]) == [] \ No newline at end of file diff --git a/tests-unit/prompt_server_test/user_manager_test.py b/tests-unit/prompt_server_test/user_manager_test.py new file mode 100644 index 000000000..936c6bd27 --- /dev/null +++ b/tests-unit/prompt_server_test/user_manager_test.py @@ -0,0 +1,120 @@ +import pytest +import os +from aiohttp import web +from app.user_manager import UserManager +from unittest.mock import patch + +pytestmark = ( + pytest.mark.asyncio +) # This applies the asyncio mark to all test functions in the module + + +@pytest.fixture +def user_manager(tmp_path): + um = UserManager() + um.get_request_user_filepath = lambda req, file, **kwargs: os.path.join( + tmp_path, file + ) + return um + + +@pytest.fixture +def app(user_manager): + app = web.Application() + routes = web.RouteTableDef() + user_manager.add_routes(routes) + app.add_routes(routes) + return app + + +async def test_listuserdata_empty_directory(aiohttp_client, app, tmp_path): + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir") + assert resp.status == 404 + + +async def test_listuserdata_with_files(aiohttp_client, app, tmp_path): + os.makedirs(tmp_path / "test_dir") + with open(tmp_path / "test_dir" / "file1.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir") + assert resp.status == 200 + assert await resp.json() == ["file1.txt"] + + +async def test_listuserdata_recursive(aiohttp_client, app, tmp_path): + os.makedirs(tmp_path / "test_dir" / "subdir") + with open(tmp_path / "test_dir" / "file1.txt", "w") as f: + f.write("test content") + with open(tmp_path / "test_dir" / "subdir" / "file2.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir&recurse=true") + assert resp.status == 200 + assert set(await resp.json()) == {"file1.txt", "subdir/file2.txt"} + + +async def test_listuserdata_full_info(aiohttp_client, app, tmp_path): + os.makedirs(tmp_path / "test_dir") + with open(tmp_path / "test_dir" / "file1.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir&full_info=true") + assert resp.status == 200 + result = await resp.json() + assert len(result) == 1 + assert result[0]["path"] == "file1.txt" + assert "size" in result[0] + assert "modified" in result[0] + + +async def test_listuserdata_split_path(aiohttp_client, app, tmp_path): + os.makedirs(tmp_path / "test_dir" / "subdir") + with open(tmp_path / "test_dir" / "subdir" / "file1.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir&recurse=true&split=true") + assert resp.status == 200 + assert await resp.json() == [ + ["subdir/file1.txt", "subdir", "file1.txt"] + ] + + +async def test_listuserdata_invalid_directory(aiohttp_client, app): + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=") + assert resp.status == 400 + + +async def test_listuserdata_normalized_separator(aiohttp_client, app, tmp_path): + os_sep = "\\" + with patch("os.sep", os_sep): + with patch("os.path.sep", os_sep): + os.makedirs(tmp_path / "test_dir" / "subdir") + with open(tmp_path / "test_dir" / "subdir" / "file1.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir&recurse=true") + assert resp.status == 200 + result = await resp.json() + assert len(result) == 1 + assert "/" in result[0] # Ensure forward slash is used + assert "\\" not in result[0] # Ensure backslash is not present + assert result[0] == "subdir/file1.txt" + + # Test with full_info + resp = await client.get( + "/userdata?dir=test_dir&recurse=true&full_info=true" + ) + assert resp.status == 200 + result = await resp.json() + assert len(result) == 1 + assert "/" in result[0]["path"] # Ensure forward slash is used + assert "\\" not in result[0]["path"] # Ensure backslash is not present + assert result[0]["path"] == "subdir/file1.txt" diff --git a/tests-unit/utils/extra_config_test.py b/tests-unit/utils/extra_config_test.py new file mode 100644 index 000000000..06349560d --- /dev/null +++ b/tests-unit/utils/extra_config_test.py @@ -0,0 +1,124 @@ +import pytest +import yaml +import os +from unittest.mock import Mock, patch, mock_open + +from utils.extra_config import load_extra_path_config +import folder_paths + +@pytest.fixture +def mock_yaml_content(): + return { + 'test_config': { + 'base_path': '~/App/', + 'checkpoints': 'subfolder1', + } + } + +@pytest.fixture +def mock_expanded_home(): + return '/home/user' + +@pytest.fixture +def yaml_config_with_appdata(): + return """ + test_config: + base_path: '%APPDATA%/ComfyUI' + checkpoints: 'models/checkpoints' + """ + +@pytest.fixture +def mock_yaml_content_appdata(yaml_config_with_appdata): + return yaml.safe_load(yaml_config_with_appdata) + +@pytest.fixture +def mock_expandvars_appdata(): + mock = Mock() + mock.side_effect = lambda path: path.replace('%APPDATA%', 'C:/Users/TestUser/AppData/Roaming') + return mock + +@pytest.fixture +def mock_add_model_folder_path(): + return Mock() + +@pytest.fixture +def mock_expanduser(mock_expanded_home): + def _expanduser(path): + if path.startswith('~/'): + return os.path.join(mock_expanded_home, path[2:]) + return path + return _expanduser + +@pytest.fixture +def mock_yaml_safe_load(mock_yaml_content): + return Mock(return_value=mock_yaml_content) + +@patch('builtins.open', new_callable=mock_open, read_data="dummy file content") +def test_load_extra_model_paths_expands_userpath( + mock_file, + monkeypatch, + mock_add_model_folder_path, + mock_expanduser, + mock_yaml_safe_load, + mock_expanded_home +): + # Attach mocks used by load_extra_path_config + monkeypatch.setattr(folder_paths, 'add_model_folder_path', mock_add_model_folder_path) + monkeypatch.setattr(os.path, 'expanduser', mock_expanduser) + monkeypatch.setattr(yaml, 'safe_load', mock_yaml_safe_load) + + dummy_yaml_file_name = 'dummy_path.yaml' + load_extra_path_config(dummy_yaml_file_name) + + expected_calls = [ + ('checkpoints', os.path.join(mock_expanded_home, 'App', 'subfolder1')), + ] + + assert mock_add_model_folder_path.call_count == len(expected_calls) + + # Check if add_model_folder_path was called with the correct arguments + for actual_call, expected_call in zip(mock_add_model_folder_path.call_args_list, expected_calls): + assert actual_call.args == expected_call + + # Check if yaml.safe_load was called + mock_yaml_safe_load.assert_called_once() + + # Check if open was called with the correct file path + mock_file.assert_called_once_with(dummy_yaml_file_name, 'r') + +@patch('builtins.open', new_callable=mock_open) +def test_load_extra_model_paths_expands_appdata( + mock_file, + monkeypatch, + mock_add_model_folder_path, + mock_expandvars_appdata, + yaml_config_with_appdata, + mock_yaml_content_appdata +): + # Set the mock_file to return yaml with appdata as a variable + mock_file.return_value.read.return_value = yaml_config_with_appdata + + # Attach mocks + monkeypatch.setattr(folder_paths, 'add_model_folder_path', mock_add_model_folder_path) + monkeypatch.setattr(os.path, 'expandvars', mock_expandvars_appdata) + monkeypatch.setattr(yaml, 'safe_load', Mock(return_value=mock_yaml_content_appdata)) + + # Mock expanduser to do nothing (since we're not testing it here) + monkeypatch.setattr(os.path, 'expanduser', lambda x: x) + + dummy_yaml_file_name = 'dummy_path.yaml' + load_extra_path_config(dummy_yaml_file_name) + + expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI' + expected_calls = [ + ('checkpoints', os.path.join(expected_base_path, 'models/checkpoints')), + ] + + assert mock_add_model_folder_path.call_count == len(expected_calls) + + # Check the base path variable was expanded + for actual_call, expected_call in zip(mock_add_model_folder_path.call_args_list, expected_calls): + assert actual_call.args == expected_call + + # Verify that expandvars was called + assert mock_expandvars_appdata.called diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/utils/extra_config.py b/utils/extra_config.py new file mode 100644 index 000000000..7ac3650ac --- /dev/null +++ b/utils/extra_config.py @@ -0,0 +1,25 @@ +import os +import yaml +import folder_paths +import logging + +def load_extra_path_config(yaml_path): + with open(yaml_path, 'r') as stream: + config = yaml.safe_load(stream) + for c in config: + conf = config[c] + if conf is None: + continue + base_path = None + if "base_path" in conf: + base_path = conf.pop("base_path") + base_path = os.path.expandvars(os.path.expanduser(base_path)) + for x in conf: + for y in conf[x].split("\n"): + if len(y) == 0: + continue + full_path = y + 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)