From f5a665499c8bdff682aebbc32d1d146fb1b16aed Mon Sep 17 00:00:00 2001 From: Thomas Date: Thu, 30 Apr 2026 21:40:17 +0200 Subject: [PATCH] fix: resolve missing model sizes and improve URL parsing --- src/main/java/de/tki/comfymodels/Main.java | 34 ++++++++ .../service/impl/ComfyModelAnalyzer.java | 3 +- .../comfymodels/ComfyModelAnalyzerTest.java | 29 +++++++ .../MissingSizeReproductionTest.java | 84 +++++++++++++++++++ 4 files changed, 149 insertions(+), 1 deletion(-) create mode 100644 src/test/java/de/tki/comfymodels/MissingSizeReproductionTest.java diff --git a/src/main/java/de/tki/comfymodels/Main.java b/src/main/java/de/tki/comfymodels/Main.java index 691118f..39808a9 100644 --- a/src/main/java/de/tki/comfymodels/Main.java +++ b/src/main/java/de/tki/comfymodels/Main.java @@ -697,10 +697,44 @@ public class Main extends JFrame { String type = info.getType() != null ? info.getType() : "checkpoints"; Path local = Paths.get(base, info.getSave_path() != null ? info.getSave_path() : type, info.getName()); boolean exists = Files.exists(local); + + // Populate size from local file if it exists and currently unknown + if (exists && "Unknown".equals(info.getSize())) { + try { + long bytes = Files.size(local); + info.setSize(searchService.formatSize(bytes)); + } catch (Exception e) {} + } + String status = exists ? "✅ Already exists" : (info.getUrl().equals("MISSING") ? "Idle" : "✅ Known Good"); tableModel.addRow(new Object[]{!exists, info.getType(), info.getName(), info.getSize(), info.getPopularity(), "models/" + type, info.getUrl(), status}); } downloadButton.setEnabled(!modelsToDownload.isEmpty()); + + // NEW: Automatically fetch missing sizes for remote models in background + fetchMissingRemoteSizes(); + } + + private void fetchMissingRemoteSizes() { + if (modelsToDownload == null) return; + new Thread(() -> { + for (int i = 0; i < modelsToDownload.size(); i++) { + final int idx = i; + ModelInfo info = modelsToDownload.get(idx); + if (!info.getUrl().equals("MISSING") && "Unknown".equals(info.getSize())) { + long size = searchService.getRemoteSize(info.getUrl()); + if (size > 0) { + String formatted = searchService.formatSize(size); + info.setSize(formatted); + SwingUtilities.invokeLater(() -> { + if (idx < tableModel.getRowCount()) { + tableModel.setValueAt(formatted, idx, 3); + } + }); + } + } + } + }).start(); } private void searchMissingOnline() { diff --git a/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java b/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java index e358bb7..a4779f0 100644 --- a/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java +++ b/src/main/java/de/tki/comfymodels/service/impl/ComfyModelAnalyzer.java @@ -28,7 +28,7 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { private final Pattern COMP_CLIP_PATTERN = Pattern.compile("(clip|text.encoder|t5|llama|gemma|mistral|qwen|bert|vit|embed)", Pattern.CASE_INSENSITIVE); private final Pattern COMP_UNET_PATTERN = Pattern.compile("(unet|diffusion|transformer|dit|model|base|upscale|esrgan|resnet|sampling)", Pattern.CASE_INSENSITIVE); private final Pattern MD_LINK_PATTERN = Pattern.compile("\\[([^\\]]+)\\]\\((https?://[^\\)]+)\\)", Pattern.CASE_INSENSITIVE); - private final Pattern NAKED_URL_PATTERN = Pattern.compile("(https?://[a-zA-Z0-9\\-\\.\\/\\?\\=\\&\\%]+(?:\\.(?:safetensors|sft|ckpt|pth|pt|bin|onnx|yaml))(?:\\?[^\\s\\\"\\']*)?)", Pattern.CASE_INSENSITIVE); + private final Pattern NAKED_URL_PATTERN = Pattern.compile("(https?://[a-zA-Z0-9\\-\\._\\/\\?\\=\\&\\%]+(?:\\.(?:safetensors|sft|ckpt|pth|pt|bin|onnx|yaml))(?:\\?[^\\s\\\"\\']*)?)", Pattern.CASE_INSENSITIVE); @Override public List analyze(String jsonText, String fileName) { @@ -54,6 +54,7 @@ public class ComfyModelAnalyzer implements IModelAnalyzer { ModelInfo match = listMatch.get(); info.setUrl(match.getUrl()); info.setSave_path(match.getSave_path()); + info.setSize(match.getSize()); info.setPopularity("📂 MODEL LIST MATCH"); } else { if (aiService != null) { diff --git a/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java b/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java index 946fd1e..a9f53a7 100644 --- a/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java +++ b/src/test/java/de/tki/comfymodels/ComfyModelAnalyzerTest.java @@ -40,6 +40,35 @@ public class ComfyModelAnalyzerTest { assertEquals("https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors", models.get(0).getUrl()); } + @Test + public void testAnalyzeWithModelListSize() throws Exception { + ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer(); + + // Mock ModelListService + ModelListService mockListService = new ModelListService() { + @Override + public Optional findByFilename(String filename) { + if ("model_with_size.safetensors".equalsIgnoreCase(filename)) { + ModelInfo info = new ModelInfo("checkpoints", "model_with_size.safetensors", "http://example.com/model"); + info.setSize("1.23 GB"); + return Optional.of(info); + } + return Optional.empty(); + } + }; + + Field field = ComfyModelAnalyzer.class.getDeclaredField("modelListService"); + field.setAccessible(true); + field.set(analyzer, mockListService); + + String json = "{\"nodes\": [{\"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"model_with_size.safetensors\"]}]}"; + List models = analyzer.analyze(json, "test.json"); + + assertFalse(models.isEmpty()); + assertEquals("model_with_size.safetensors", models.get(0).getName()); + assertEquals("1.23 GB", models.get(0).getSize()); + } + @Test public void testInvalidJson() { ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer(); diff --git a/src/test/java/de/tki/comfymodels/MissingSizeReproductionTest.java b/src/test/java/de/tki/comfymodels/MissingSizeReproductionTest.java new file mode 100644 index 0000000..6d2bc7b --- /dev/null +++ b/src/test/java/de/tki/comfymodels/MissingSizeReproductionTest.java @@ -0,0 +1,84 @@ +package de.tki.comfymodels; + +import de.tki.comfymodels.domain.ModelInfo; +import de.tki.comfymodels.service.impl.ComfyModelAnalyzer; +import de.tki.comfymodels.service.impl.ModelListService; +import org.junit.jupiter.api.Test; +import java.lang.reflect.Field; +import java.util.List; +import java.util.Optional; + +import static org.junit.jupiter.api.Assertions.*; + +public class MissingSizeReproductionTest { + + @Test + public void testNakedUrlSizeIsUnknown() throws Exception { + ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer(); + + // No ModelListService match for this one + ModelListService mockListService = new ModelListService() { + @Override + public Optional findByFilename(String filename) { + return Optional.empty(); + } + }; + + Field field = ComfyModelAnalyzer.class.getDeclaredField("modelListService"); + field.setAccessible(true); + field.set(analyzer, mockListService); + + // A workflow containing a naked URL to a model + String json = "{\"nodes\": [{\"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/diffusion_models/ernie-image-turbo.safetensors\"]}]}"; + List models = analyzer.analyze(json, "test.json"); + + assertFalse(models.isEmpty()); + ModelInfo info = models.get(0); + assertEquals("ernie-image-turbo.safetensors", info.getName()); + assertEquals("https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/diffusion_models/ernie-image-turbo.safetensors", info.getUrl()); + + // This reproduces the issue: Size is "Unknown" because analyzer doesn't do network calls + assertEquals("Unknown", info.getSize()); + } + + @Test + public void testUrlWithUnderscoresIsFound() throws Exception { + ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer(); + String json = "{\"nodes\": [{\"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/diffusion_models/ernie_image_turbo.safetensors\"]}]}"; + List models = analyzer.analyze(json, "test.json"); + + assertFalse(models.isEmpty(), "Should find URL with underscores"); + assertEquals("ernie_image_turbo.safetensors", models.get(0).getName()); + } + + @Test + public void testModelListMatchWithoutSize() throws Exception { + ComfyModelAnalyzer analyzer = new ComfyModelAnalyzer(); + + ModelListService mockListService = new ModelListService() { + @Override + public Optional findByFilename(String filename) { + if ("ernie-image-turbo.safetensors".equalsIgnoreCase(filename)) { + // Entry in model list has URL but NO size (or "Unknown") + ModelInfo info = new ModelInfo("diffusion_models", "ernie-image-turbo.safetensors", + "https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/diffusion_models/ernie-image-turbo.safetensors"); + info.setSize("Unknown"); + return Optional.of(info); + } + return Optional.empty(); + } + }; + + Field field = ComfyModelAnalyzer.class.getDeclaredField("modelListService"); + field.setAccessible(true); + field.set(analyzer, mockListService); + + String json = "{\"nodes\": [{\"type\": \"CheckpointLoaderSimple\", \"widgets_values\": [\"ernie-image-turbo.safetensors\"]}]}"; + List models = analyzer.analyze(json, "test.json"); + + assertFalse(models.isEmpty()); + ModelInfo info = models.get(0); + assertEquals("ernie-image-turbo.safetensors", info.getName()); + assertEquals("Unknown", info.getSize()); + } +}