From b24bb1521603aa60444f71d2b9e019ad90df1a32 Mon Sep 17 00:00:00 2001 From: Thomas Date: Tue, 5 May 2026 18:51:51 +0200 Subject: [PATCH] feat: ensure .venv and archive folders are ignored during model scanning and cleanup --- src/main/java/de/tki/comfymodels/Main.java | 4 ++++ .../service/impl/CleanupService.java | 22 ++++++++++++++----- .../service/impl/LocalModelScanner.java | 11 +++++----- .../service/LocalModelScannerTest.java | 12 +++++++--- 4 files changed, 36 insertions(+), 13 deletions(-) diff --git a/src/main/java/de/tki/comfymodels/Main.java b/src/main/java/de/tki/comfymodels/Main.java index e7c4a30..2b00e56 100644 --- a/src/main/java/de/tki/comfymodels/Main.java +++ b/src/main/java/de/tki/comfymodels/Main.java @@ -1298,6 +1298,10 @@ public class Main extends JFrame { List allFiles = Files.walk(root.toPath()) .filter(Files::isRegularFile) + .filter(p -> { + String relPath = root.toPath().relativize(p).toString().toLowerCase(); + return !relPath.contains(".venv") && !relPath.contains("archive"); + }) .filter(p -> { String n = p.getFileName().toString().toLowerCase(); return n.endsWith(".safetensors") || n.endsWith(".sft") || n.endsWith(".ckpt") || n.endsWith(".pth") || n.endsWith(".pt") || n.endsWith(".bin"); diff --git a/src/main/java/de/tki/comfymodels/service/impl/CleanupService.java b/src/main/java/de/tki/comfymodels/service/impl/CleanupService.java index 825b3e1..3e060c7 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/CleanupService.java +++ b/src/main/java/de/tki/comfymodels/service/impl/CleanupService.java @@ -44,6 +44,7 @@ public class CleanupService implements ICleanupService { try (Stream walk = Files.walk(root)) { List candidates = walk.filter(Files::isRegularFile) .filter(this::isSupportedModel) + .filter(p -> !isIgnored(p, root)) .collect(Collectors.toList()); for (Path p : candidates) { @@ -67,6 +68,22 @@ public class CleanupService implements ICleanupService { return unused; } + private boolean isIgnored(Path path, Path rootPath) { + int rootCount = rootPath.getNameCount(); + for (int i = rootCount; i < path.getNameCount(); i++) { + String part = path.getName(i).toString().toLowerCase(); + if (part.equals("archive") || part.equals(".venv")) { + return true; + } + } + return false; + } + + private boolean isSupportedModel(Path p) { + String n = p.getFileName().toString().toLowerCase(); + return n.endsWith(".safetensors") || n.endsWith(".ckpt") || n.endsWith(".pt") || n.endsWith(".pth") || n.endsWith(".bin"); + } + @Override public void archiveModels(List models, BiConsumer onStatusUpdate, Runnable onFinished) { new Thread(() -> { @@ -85,11 +102,6 @@ public class CleanupService implements ICleanupService { }).start(); } - private boolean isSupportedModel(Path p) { - String n = p.getFileName().toString().toLowerCase(); - return n.endsWith(".safetensors") || n.endsWith(".ckpt") || n.endsWith(".pt") || n.endsWith(".pth") || n.endsWith(".bin"); - } - private ModelInfo createInfo(Path p, Path root, BasicFileAttributes attrs) { ModelInfo info = new ModelInfo(); info.setName(p.getFileName().toString()); diff --git a/src/main/java/de/tki/comfymodels/service/impl/LocalModelScanner.java b/src/main/java/de/tki/comfymodels/service/impl/LocalModelScanner.java index 431409d..63eb27d 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/LocalModelScanner.java +++ b/src/main/java/de/tki/comfymodels/service/impl/LocalModelScanner.java @@ -53,7 +53,7 @@ public class LocalModelScanner { List files = walk .filter(Files::isRegularFile) .filter(this::isSupportedModelFile) - .filter(p -> !isArchived(p, rootPath)) + .filter(p -> !isIgnored(p, rootPath)) .collect(Collectors.toList()); for (Path file : files) { @@ -91,7 +91,7 @@ public class LocalModelScanner { try (Stream walk = Files.walk(root)) { List matches = walk.filter(Files::isRegularFile) - .filter(p -> !isArchived(p, root)) + .filter(p -> !isIgnored(p, root)) .filter(p -> p.getFileName().toString().equalsIgnoreCase(filename)) .collect(Collectors.toList()); @@ -119,7 +119,7 @@ public class LocalModelScanner { return name.endsWith(".safetensors") || name.endsWith(".ckpt") || name.endsWith(".pt"); } - private boolean isArchived(Path path, Path rootPath) { + private boolean isIgnored(Path path, Path rootPath) { // 1. If we are already searching inside an archive folder (e.g. E:\Archive), // we should not filter out anything based on the "archive" name. String rootName = rootPath.getFileName() != null ? rootPath.getFileName().toString().toLowerCase() : ""; @@ -138,10 +138,11 @@ public class LocalModelScanner { } } catch (Exception ignored) {} - // 3. Fallback: ignore any subfolder named "archive" + // 3. Fallback: ignore any subfolder named "archive" or ".venv" int rootCount = rootPath.getNameCount(); for (int i = rootCount; i < path.getNameCount(); i++) { - if (path.getName(i).toString().equalsIgnoreCase("archive")) { + String part = path.getName(i).toString().toLowerCase(); + if (part.equals("archive") || part.equals(".venv")) { return true; } } diff --git a/src/test/java/de/tki/comfymodels/service/LocalModelScannerTest.java b/src/test/java/de/tki/comfymodels/service/LocalModelScannerTest.java index 031454f..2aefeda 100644 --- a/src/test/java/de/tki/comfymodels/service/LocalModelScannerTest.java +++ b/src/test/java/de/tki/comfymodels/service/LocalModelScannerTest.java @@ -41,34 +41,40 @@ public class LocalModelScannerTest { // models/checkpoints/test1.safetensors // models/loras/test2.ckpt // models/archive/old.pt (should be ignored) + // models/unet/.venv/hidden.safetensors (should be ignored) Path checkpointsDir = tempDir.resolve("checkpoints"); Path lorasDir = tempDir.resolve("loras"); Path archiveDir = tempDir.resolve("archive"); + Path venvDir = tempDir.resolve("unet").resolve(".venv"); Files.createDirectories(checkpointsDir); Files.createDirectories(lorasDir); Files.createDirectories(archiveDir); + Files.createDirectories(venvDir); Files.writeString(checkpointsDir.resolve("test1.safetensors"), "dummy content"); Files.writeString(lorasDir.resolve("test2.ckpt"), "dummy content"); Files.writeString(archiveDir.resolve("old.pt"), "dummy content"); + Files.writeString(venvDir.resolve("hidden.safetensors"), "dummy content"); // Act List results = scanner.scanLocalModels(); // Assert - assertEquals(2, results.size(), "Should find exactly 2 models (excluding archive)"); + assertEquals(2, results.size(), "Should find exactly 2 models (excluding archive and .venv)"); ModelInfo m1 = results.stream().filter(m -> m.getName().equals("test1.safetensors")).findFirst().orElse(null); assertNotNull(m1); assertEquals("checkpoints", m1.getType()); - assertTrue(m1.getSize().contains("MB")); ModelInfo m2 = results.stream().filter(m -> m.getName().equals("test2.ckpt")).findFirst().orElse(null); assertNotNull(m2); assertEquals("loras", m2.getType()); - assertTrue(m2.getSize().contains("MB")); + + // Ensure nothing from .venv or archive was found + assertFalse(results.stream().anyMatch(m -> m.getName().equals("old.pt"))); + assertFalse(results.stream().anyMatch(m -> m.getName().equals("hidden.safetensors"))); } @Test