Merge remote-tracking branch 'upstream/master' into pr

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):
+            if os.path.isfile(self.get_users_file()):
-                with open(users_file) as f:
+                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
@@ -85,8 +87,7 @@ class UserManager():
 
         self.users[user_id] = name
 
-        global users_file
+        with open(self.get_users_file(), "w") as f:
-        with open(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(
+            full_info = request.rel_url.query.get('full_info', '').lower() == "true"
-                glob.escape(path), '**/*'), recursive=recurse)
+
-            results = [os.path.relpath(x, path) for x in results if os.path.isfile(x)]
+            # 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:
+            if split_path and not full_info:
-                results = [[x] + x.split(os.sep) for x in results]
+                results = [[x] + x.split('/') for x in results]
 
             return web.json_response(results)
 

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,

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:

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.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):

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

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)

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

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"]
 

comfy/sd.py

@@ -445,10 +445,6 @@ 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
     elif len(clip_data) == 2:
@@ -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:

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()

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:

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])
 

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

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:

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()):

comfy_execution/graph.py

@@ -99,15 +99,22 @@ 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):
+        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):
+    def add_node(self, node_unique_id, include_lazy=False, subgraph_nodes=None):
+        node_ids = [node_unique_id]
+        links = []
+
+        while len(node_ids) > 0:
+            unique_id = node_ids.pop()
             if unique_id in self.pendingNodes:
-            return
+                continue
+
             self.pendingNodes[unique_id] = True
             self.blockCount[unique_id] = 0
             self.blocking[unique_id] = {}
@@ -121,8 +128,15 @@ class TopologicalSort:
                         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:
+                    if (include_lazy or not is_lazy) and not self.is_cached(from_node_id):
-                    self.add_strong_link(from_node_id, from_socket, unique_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):
+    def is_cached(self, node_id):
-        if self.output_cache.get(from_node_id) is not None:
+        return self.output_cache.get(node_id) is not None
-            # Nothing to do
-            return
-        super().add_strong_link(from_node_id, from_socket, to_node_id)
 
     def stage_node_execution(self):
         assert self.staged_node_id is None

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 = [
+        files = folder_paths.filter_files_content_types(os.listdir(input_dir), ["audio", "video"])
-            f for f in os.listdir(input_dir)
-            if (os.path.isfile(os.path.join(input_dir, f))
-                and f.endswith(LoadAudio.SUPPORTED_FORMATS)
-            )
-        ]
         return {"required": {"audio": (sorted(files), {"audio_upload": True})}}
 
     CATEGORY = "audio"

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,
 }

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()

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,
+}

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/

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

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)
 
@@ -254,6 +233,11 @@ if __name__ == "__main__":
         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)
 

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 = []

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
+            host_domain_parsed = urllib.parse.urlsplit('//' + host_domain)
-                result = urllib.parse.urlsplit('//' + host_domain)
-                host_domain = result.hostname
 
-            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)

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/`

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"

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"]) == []

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"

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

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)