fix: resolve missing model sizes and improve URL parsing
This commit is contained in:
@@ -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() {
|
||||
|
||||
@@ -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<ModelInfo> 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) {
|
||||
|
||||
@@ -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<ModelInfo> 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<ModelInfo> 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();
|
||||
|
||||
@@ -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<ModelInfo> 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<ModelInfo> 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<ModelInfo> 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<ModelInfo> 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<ModelInfo> 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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user