diff --git a/metadata_utils.py b/metadata_utils.py index b3450b1..416b096 100644 --- a/metadata_utils.py +++ b/metadata_utils.py @@ -98,7 +98,9 @@ def find_model_file(model_name: str, model_type: str) -> Optional[Path]: Path to model file or None if not found """ candidate_names = [model_name] - if "." not in Path(model_name).name: + known_model_suffixes = {".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf"} + suffix = Path(model_name).suffix.lower() + if suffix not in known_model_suffixes: candidate_names.append(f"{model_name}.safetensors") # Get search paths @@ -460,8 +462,11 @@ def build_metadata_sources(manual_inputs: dict, extracted: dict, fields: List[st for field in fields: manual_value = manual_inputs.get(field) extracted_value = extracted.get(field) - if manual_value is not None and not (isinstance(manual_value, str) and not manual_value.strip()): - sources[field] = "manual_override" + if manual_value is not None: + if isinstance(manual_value, str) and not manual_value.strip(): + sources[field] = "unknown" + else: + sources[field] = "manual_override" elif extracted_value is not None and not (isinstance(extracted_value, str) and not extracted_value.strip()): sources[field] = "detected" else: @@ -1000,6 +1005,7 @@ def _sanitize_relative_segments(path_text: str) -> Path: normalized = path_text.replace("\\", "/") if _has_drive_or_root(normalized): raise ValueError(f"Absolute paths are not allowed here: {path_text!r}") + normalized = normalized.rstrip("/") raw_segments = normalized.split("/") if any(segment == "" for segment in raw_segments): raise ValueError(f"Empty path segments are not allowed: {path_text!r}") diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index 242d0d8..c2fbbc1 100644 --- a/tests/test_metadata_utils.py +++ b/tests/test_metadata_utils.py @@ -11,6 +11,7 @@ from metadata_utils import ( build_metadata_status, ensure_metahub_save_node, ensure_prompt_in_workflow, + find_model_file, get_next_filename, get_output_directory, get_workflow_json, @@ -58,6 +59,20 @@ def test_output_path_relative_resolves_under_comfy_output(tmp_path, monkeypatch) assert output.exists() +def test_output_path_allows_trailing_separator(tmp_path, monkeypatch): + class FakeFolderPaths: + @staticmethod + def get_output_directory(): + return str(tmp_path / "output") + + monkeypatch.setitem(sys.modules, "folder_paths", FakeFolderPaths) + + output = get_output_directory("project/") + + assert output == tmp_path / "output" / "project" + assert output.exists() + + def test_filename_pattern_supports_subfolders_and_tokens(): params = {"seed": 123, "model_name": "model.safetensors"} @@ -98,6 +113,27 @@ def test_metadata_status_marks_defaulted_fields_as_partial(): assert build_metadata_status(sources, ["positive", "model_name", "seed"]) == "partial" +def test_blank_string_override_is_unknown_not_detected(): + sources = build_metadata_sources( + {"positive": ""}, + {"positive": "extracted prompt"}, + ["positive"], + ) + + assert sources["positive"] == "unknown" + assert build_metadata_status(sources, ["positive"]) == "fallback" + + +def test_find_model_file_tries_safetensors_for_dotted_extensionless_names(tmp_path, monkeypatch): + model_dir = tmp_path / "models" + model_dir.mkdir() + model_path = model_dir / "dreamshaper.v8.safetensors" + model_path.write_bytes(b"model") + monkeypatch.setenv("COMFYUI_CHECKPOINT_PATH", str(model_dir)) + + assert find_model_file("dreamshaper.v8", "checkpoint") == model_path + + def test_ensure_prompt_in_workflow_adds_prompt(): workflow = {"workflow": {"nodes": []}} prompt = {"1": {"class_type": "KSampler", "inputs": {}}} diff --git a/tests/test_video_metadata_utils.py b/tests/test_video_metadata_utils.py index 6961c4d..ec193cb 100644 --- a/tests/test_video_metadata_utils.py +++ b/tests/test_video_metadata_utils.py @@ -44,3 +44,19 @@ def test_video_metadata_nulls_defaulted_canonical_fields(): assert payload["prompt"] is None assert payload["model"] is None assert payload["seed"] is None + + +def test_video_metadata_nulls_unknown_blank_overrides(): + payload = build_video_metahub_metadata( + { + "positive": "", + "width": 640, + "height": 480, + "metadata_sources": { + "positive": "unknown", + }, + }, + {}, + ) + + assert payload["prompt"] is None diff --git a/video_metadata_utils.py b/video_metadata_utils.py index eee7616..b557d9e 100644 --- a/video_metadata_utils.py +++ b/video_metadata_utils.py @@ -209,7 +209,7 @@ def build_video_metahub_metadata(params: dict, workflow_json: dict) -> dict: sources = params.get("metadata_sources") if isinstance(params.get("metadata_sources"), dict) else {} def canonical(field: str, param_key: str, fallback: Any = None) -> Any: - if sources.get(field) == "default": + if sources.get(field) in {"default", "unknown"}: return None return params.get(param_key, fallback)