fix: resolve missing model sizes and improve URL parsing

This commit is contained in:
Thomas
2026-04-30 21:40:17 +02:00
parent 896daaf7e7
commit f5a665499c
4 changed files with 149 additions and 1 deletions
@@ -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());
}
}