Merge 126cdc582b into ff838657fa

Signed-off-by: bigcat88 <bigcat88@icloud.com>

comfy/ldm/lightricks/model.py

@@ -456,9 +456,8 @@ class LTXVModel(torch.nn.Module):
         x = self.patchify_proj(x)
         timestep = timestep * 1000.0
 
-        attention_mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1]))
-        attention_mask = attention_mask.masked_fill(attention_mask.to(torch.bool), float("-inf"))  # not sure about this
-        # attention_mask = (context != 0).any(dim=2).to(dtype=x.dtype)
+        if attention_mask is not None and not torch.is_floating_point(attention_mask):
+            attention_mask = (attention_mask - 1).to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(x.dtype).max
 
         pe = precompute_freqs_cis(indices_grid, dim=self.inner_dim, out_dtype=x.dtype)
 

tests-unit/utils/extra_config_test.py

@@ -1,11 +1,22 @@
 import pytest
 import yaml
 import os
+import sys
 from unittest.mock import Mock, patch, mock_open
 
 from utils.extra_config import load_extra_path_config
 import folder_paths
 
+
+@pytest.fixture()
+def clear_folder_paths():
+    # Clear the global dictionary before each test to ensure isolation
+    original = folder_paths.folder_names_and_paths.copy()
+    folder_paths.folder_names_and_paths.clear()
+    yield
+    folder_paths.folder_names_and_paths = original
+
+
 @pytest.fixture
 def mock_yaml_content():
     return {
@@ -15,10 +26,12 @@ def mock_yaml_content():
         }
     }
 
+
 @pytest.fixture
 def mock_expanded_home():
     return '/home/user'
 
+
 @pytest.fixture
 def yaml_config_with_appdata():
     return """
@@ -27,20 +40,33 @@ def yaml_config_with_appdata():
       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')
+
+    def expandvars(path):
+        if '%APPDATA%' in path:
+            if sys.platform == 'win32':
+                return path.replace('%APPDATA%', 'C:/Users/TestUser/AppData/Roaming')
+            else:
+                return path.replace('%APPDATA%', '/Users/TestUser/AppData/Roaming')
+        return path
+
+    mock.side_effect = expandvars
     return mock
 
+
 @pytest.fixture
 def mock_add_model_folder_path():
     return Mock()
 
+
 @pytest.fixture
 def mock_expanduser(mock_expanded_home):
     def _expanduser(path):
@@ -49,10 +75,12 @@ def mock_expanduser(mock_expanded_home):
         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,
@@ -88,6 +116,7 @@ def test_load_extra_model_paths_expands_userpath(
     # 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,
@@ -111,7 +140,10 @@ def test_load_extra_model_paths_expands_appdata(
     dummy_yaml_file_name = 'dummy_path.yaml'
     load_extra_path_config(dummy_yaml_file_name)
 
-    expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI'
+    if sys.platform == "win32":
+        expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI'
+    else:
+        expected_base_path = '/Users/TestUser/AppData/Roaming/ComfyUI'
     expected_calls = [
         ('checkpoints', os.path.join(expected_base_path, 'models/checkpoints'), False),
     ]
@@ -124,3 +156,148 @@ def test_load_extra_model_paths_expands_appdata(
 
     # Verify that expandvars was called
     assert mock_expandvars_appdata.called
+
+
+@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content")
+@patch("yaml.safe_load")
+def test_load_extra_path_config_relative_base_path(
+    mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path
+):
+    """
+    Test that when 'base_path' is a relative path in the YAML, it is joined to the YAML file directory, and then
+    the items in the config are correctly converted to absolute paths.
+    """
+    sub_folder = "./my_rel_base"
+    config_data = {
+        "some_model_folder": {
+            "base_path": sub_folder,
+            "is_default": True,
+            "checkpoints": "checkpoints",
+            "some_key": "some_value"
+        }
+    }
+    mock_yaml_load.return_value = config_data
+
+    dummy_yaml_name = "dummy_file.yaml"
+
+    def fake_abspath(path):
+        if path == dummy_yaml_name:
+            # If it's the YAML path, treat it like it lives in tmp_path
+            return os.path.join(str(tmp_path), dummy_yaml_name)
+        return os.path.join(str(tmp_path), path)  # Otherwise, do a normal join relative to tmp_path
+
+    def fake_dirname(path):
+        # We expect path to be the result of fake_abspath(dummy_yaml_name)
+        if path.endswith(dummy_yaml_name):
+            return str(tmp_path)
+        return os.path.dirname(path)
+
+    monkeypatch.setattr(os.path, "abspath", fake_abspath)
+    monkeypatch.setattr(os.path, "dirname", fake_dirname)
+
+    load_extra_path_config(dummy_yaml_name)
+
+    expected_checkpoints = os.path.abspath(os.path.join(str(tmp_path), sub_folder, "checkpoints"))
+    expected_some_value = os.path.abspath(os.path.join(str(tmp_path), sub_folder, "some_value"))
+
+    actual_paths = folder_paths.folder_names_and_paths["checkpoints"][0]
+    assert len(actual_paths) == 1, "Should have one path added for 'checkpoints'."
+    assert actual_paths[0] == expected_checkpoints
+
+    actual_paths = folder_paths.folder_names_and_paths["some_key"][0]
+    assert len(actual_paths) == 1, "Should have one path added for 'some_key'."
+    assert actual_paths[0] == expected_some_value
+
+
+@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content")
+@patch("yaml.safe_load")
+def test_load_extra_path_config_absolute_base_path(
+    mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path
+):
+    """
+    Test that when 'base_path' is an absolute path, each subdirectory is joined with that absolute path,
+    rather than being relative to the YAML's directory.
+    """
+    abs_base = os.path.join(str(tmp_path), "abs_base")
+    config_data = {
+        "some_absolute_folder": {
+            "base_path": abs_base,   # <-- absolute
+            "is_default": True,
+            "loras": "loras_folder",
+            "embeddings": "embeddings_folder"
+        }
+    }
+    mock_yaml_load.return_value = config_data
+
+    dummy_yaml_name = "dummy_abs.yaml"
+
+    def fake_abspath(path):
+        if path == dummy_yaml_name:
+            # If it's the YAML path, treat it like it is in tmp_path
+            return os.path.join(str(tmp_path), dummy_yaml_name)
+        return path  # For absolute base, we just return path directly
+
+    def fake_dirname(path):
+        return str(tmp_path) if path.endswith(dummy_yaml_name) else os.path.dirname(path)
+
+    monkeypatch.setattr(os.path, "abspath", fake_abspath)
+    monkeypatch.setattr(os.path, "dirname", fake_dirname)
+
+    load_extra_path_config(dummy_yaml_name)
+
+    # Expect the final paths to be <abs_base>/loras_folder and <abs_base>/embeddings_folder
+    expected_loras = os.path.join(abs_base, "loras_folder")
+    expected_embeddings = os.path.join(abs_base, "embeddings_folder")
+
+    actual_loras = folder_paths.folder_names_and_paths["loras"][0]
+    assert len(actual_loras) == 1, "Should have one path for 'loras'."
+    assert actual_loras[0] == os.path.abspath(expected_loras)
+
+    actual_embeddings = folder_paths.folder_names_and_paths["embeddings"][0]
+    assert len(actual_embeddings) == 1, "Should have one path for 'embeddings'."
+    assert actual_embeddings[0] == os.path.abspath(expected_embeddings)
+
+
+@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content")
+@patch("yaml.safe_load")
+def test_load_extra_path_config_no_base_path(
+    mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path
+):
+    """
+    Test that if 'base_path' is not present, each path is joined
+    with the directory of the YAML file (unless it's already absolute).
+    """
+    config_data = {
+        "some_folder_without_base": {
+            "is_default": True,
+            "text_encoders": "clip",
+            "diffusion_models": "unet"
+        }
+    }
+    mock_yaml_load.return_value = config_data
+
+    dummy_yaml_name = "dummy_no_base.yaml"
+
+    def fake_abspath(path):
+        if path == dummy_yaml_name:
+            return os.path.join(str(tmp_path), dummy_yaml_name)
+        return os.path.join(str(tmp_path), path)
+
+    def fake_dirname(path):
+        return str(tmp_path) if path.endswith(dummy_yaml_name) else os.path.dirname(path)
+
+    monkeypatch.setattr(os.path, "abspath", fake_abspath)
+    monkeypatch.setattr(os.path, "dirname", fake_dirname)
+
+    load_extra_path_config(dummy_yaml_name)
+
+    expected_clip = os.path.join(str(tmp_path), "clip")
+    expected_unet = os.path.join(str(tmp_path), "unet")
+
+    actual_text_encoders = folder_paths.folder_names_and_paths["text_encoders"][0]
+    assert len(actual_text_encoders) == 1, "Should have one path for 'text_encoders'."
+    assert actual_text_encoders[0] == os.path.abspath(expected_clip)
+
+    actual_diffusion = folder_paths.folder_names_and_paths["diffusion_models"][0]
+    assert len(actual_diffusion) == 1, "Should have one path for 'diffusion_models'."
+    assert actual_diffusion[0] == os.path.abspath(expected_unet)

utils/extra_config.py

@@ -6,6 +6,7 @@ import logging
 def load_extra_path_config(yaml_path):
     with open(yaml_path, 'r') as stream:
         config = yaml.safe_load(stream)
+    yaml_dir = os.path.dirname(os.path.abspath(yaml_path))
     for c in config:
         conf = config[c]
         if conf is None:
@@ -14,6 +15,8 @@ def load_extra_path_config(yaml_path):
         if "base_path" in conf:
             base_path = conf.pop("base_path")
             base_path = os.path.expandvars(os.path.expanduser(base_path))
+            if not os.path.isabs(base_path):
+                base_path = os.path.abspath(os.path.join(yaml_dir, base_path))
         is_default = False
         if "is_default" in conf:
             is_default = conf.pop("is_default")
@@ -22,10 +25,9 @@ def load_extra_path_config(yaml_path):
                 if len(y) == 0:
                     continue
                 full_path = y
-                if base_path is not None:
+                if base_path:
                     full_path = os.path.join(base_path, full_path)
                 elif not os.path.isabs(full_path):
-                    yaml_dir = os.path.dirname(os.path.abspath(yaml_path))
                     full_path = os.path.abspath(os.path.join(yaml_dir, y))
                 logging.info("Adding extra search path {} {}".format(x, full_path))
                 folder_paths.add_model_folder_path(x, full_path, is_default)